feat(core): support ai network search (#9357)

### What Changed? - Add `PerplexityProvider` in backend. - Update session prompt name if user toggle network search mode in chat panel. - Add experimental flag for AI network search feature. - Add unit tests and e2e tests. Search results are streamed and appear word for word: <div class='graphite__hidden'> <div>🎥 Video uploaded on Graphite:</div> <a href="https://app.graphite.dev/media/video/sJGviKxfE3Ap685cl5bj/56f6ec7b-4b21-405f-9612-43e083f6fb84.mov"> <img src="https://app.graphite.dev/api/v1/graphite/video/thumbnail/sJGviKxfE3Ap685cl5bj/56f6ec7b-4b21-405f-9612-43e083f6fb84.mov"> </a> </div> <video src="https://graphite-user-uploaded-assets-prod.s3.amazonaws.com/sJGviKxfE3Ap685cl5bj/56f6ec7b-4b21-405f-9612-43e083f6fb84.mov">录屏2024-12-27 18.58.40.mov</video> Click the little globe icon to manually turn on/off Internet search: <div class='graphite__hidden'> <div>🎥 Video uploaded on Graphite:</div> <a href="https://app.graphite.dev/media/video/sJGviKxfE3Ap685cl5bj/778f1406-bf29-498e-a90d-7dad813392d1.mov"> <img src="https://app.graphite.dev/api/v1/graphite/video/thumbnail/sJGviKxfE3Ap685cl5bj/778f1406-bf29-498e-a90d-7dad813392d1.mov"> </a> </div> <video src="https://graphite-user-uploaded-assets-prod.s3.amazonaws.com/sJGviKxfE3Ap685cl5bj/778f1406-bf29-498e-a90d-7dad813392d1.mov">录屏2024-12-27 19.01.16.mov</video> When there is an image, it will automatically switch to the openai model: <div class='graphite__hidden'> <div>🎥 Video uploaded on Graphite:</div> <a href="https://app.graphite.dev/media/video/sJGviKxfE3Ap685cl5bj/56431d8e-75e1-4d84-ab4a-b6636042cc6a.mov"> <img src="https://app.graphite.dev/api/v1/graphite/video/thumbnail/sJGviKxfE3Ap685cl5bj/56431d8e-75e1-4d84-ab4a-b6636042cc6a.mov"> </a> </div> <video src="https://graphite-user-uploaded-assets-prod.s3.amazonaws.com/sJGviKxfE3Ap685cl5bj/56431d8e-75e1-4d84-ab4a-b6636042cc6a.mov">录屏2024-12-27 19.02.13.mov</video>
2026-02-13 21:05:19 +00:00 · 2025-01-09 04:00:58 +00:00
parent 4f10457815
commit 58ce86533e
49 changed files with 1274 additions and 169 deletions
--- a/packages/frontend/core/src/components/blocksuite/block-suite-editor/ai/copilot-client.ts
+++ b/packages/frontend/core/src/components/blocksuite/block-suite-editor/ai/copilot-client.ts
@@ -12,6 +12,7 @@ import {
  type QueryOptions,
  type QueryResponse,
  type RequestOptions,
+  updateCopilotSessionMutation,
  UserFriendlyError,
 } from '@affine/graphql';
 import {
@@ -80,6 +81,18 @@ export class CopilotClient {
    return res.createCopilotSession;
  }

+  async updateSession(
+    options: OptionsField<typeof updateCopilotSessionMutation>
+  ) {
+    const res = await this.gql({
+      query: updateCopilotSessionMutation,
+      variables: {
+        options,
+      },
+    });
+    return res.updateCopilotSession;
+  }
+
  async forkSession(options: OptionsField<typeof forkCopilotSessionMutation>) {
    const res = await this.gql({
      query: forkCopilotSessionMutation,
--- a/packages/frontend/core/src/components/blocksuite/block-suite-editor/ai/prompt.ts
+++ b/packages/frontend/core/src/components/blocksuite/block-suite-editor/ai/prompt.ts
@@ -8,6 +8,7 @@ export const promptKeys = [
  'debug:action:fal-remove-bg',
  'debug:action:fal-face-to-sticker',
  'Chat With AFFiNE AI',
+  'Search With AFFiNE AI',
  'Summary',
  'Generate a caption',
  'Summary the webpage',
--- a/packages/frontend/core/src/components/blocksuite/block-suite-editor/ai/request.ts
+++ b/packages/frontend/core/src/components/blocksuite/block-suite-editor/ai/request.ts
@@ -35,15 +35,32 @@ export function createChatSession({
  client,
  workspaceId,
  docId,
+  promptName,
 }: {
  client: CopilotClient;
  workspaceId: string;
  docId: string;
+  promptName: string;
 }) {
  return client.createSession({
    workspaceId,
    docId,
-    promptName: 'Chat With AFFiNE AI',
+    promptName,
+  });
+}
+
+export function updateChatSession({
+  client,
+  sessionId,
+  promptName,
+}: {
+  client: CopilotClient;
+  sessionId: string;
+  promptName: string;
+}) {
+  return client.updateSession({
+    sessionId,
+    promptName,
  });
 }

--- a/packages/frontend/core/src/components/blocksuite/block-suite-editor/ai/setup-provider.tsx
+++ b/packages/frontend/core/src/components/blocksuite/block-suite-editor/ai/setup-provider.tsx
@@ -1,5 +1,6 @@
 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,
@@ -17,6 +18,7 @@ import {
  forkCopilotSession,
  textToText,
  toImage,
+  updateChatSession,
 } from './request';
 import { setupTracker } from './tracker';

@@ -39,13 +41,30 @@ 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, Promise<string>>();
+const chatSessions = new Map<
+  string,
+  { getSessionId: Promise<string>; promptName: string }
+>();

 export function setupAIProvider(
  client: CopilotClient,
-  globalDialogService: GlobalDialogService
+  globalDialogService: GlobalDialogService,
+  networkSearchService: AINetworkSearchService
 ) {
-  async function getChatSessionId(workspaceId: string, docId: string) {
+  function getChatPrompt(attachments?: (string | File | Blob)[]) {
+    if (attachments?.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(
+    workspaceId: string,
+    docId: string,
+    attachments?: (string | File | Blob)[]
+  ) {
    const userId = (await AIProvider.userInfo)?.id;

    if (!userId) {
@@ -53,19 +72,32 @@ export function setupAIProvider(
    }

    const storeKey = `${userId}:${workspaceId}:${docId}`;
+    const promptName = getChatPrompt(attachments);
    if (!chatSessions.has(storeKey)) {
-      chatSessions.set(
-        storeKey,
-        createChatSession({
+      chatSessions.set(storeKey, {
+        getSessionId: createChatSession({
          client,
          workspaceId,
          docId,
-        })
-      );
+          promptName,
+        }),
+        promptName,
+      });
    }
    try {
-      const sessionId = await chatSessions.get(storeKey);
-      assertExists(sessionId);
+      /* 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
@@ -77,7 +109,8 @@ export function setupAIProvider(
  //#region actions
  AIProvider.provide('chat', options => {
    const sessionId =
-      options.sessionId ?? getChatSessionId(options.workspaceId, options.docId);
+      options.sessionId ??
+      getChatSessionId(options.workspaceId, options.docId, options.attachments);
    return textToText({
      ...options,
      client,
--- a/packages/frontend/core/src/components/blocksuite/block-suite-editor/specs/edgeless.ts
+++ b/packages/frontend/core/src/components/blocksuite/block-suite-editor/specs/edgeless.ts
@@ -1,4 +1,4 @@
-import { AIEdgelessRootBlockSpec } from '@affine/core/blocksuite/presets/ai';
+import { createAIEdgelessRootBlockSpec } from '@affine/core/blocksuite/presets/ai';
 import { FeatureFlagService } from '@affine/core/modules/feature-flag';
 import { builtInTemplates as builtInEdgelessTemplates } from '@affine/templates/edgeless';
 import { builtInTemplates as builtInStickersTemplates } from '@affine/templates/stickers';
@@ -22,7 +22,10 @@ export function createEdgelessModeSpecs(
  enableAffineExtension(framework, edgelessSpec);
  if (enableAI) {
    enableAIExtension(edgelessSpec);
-    edgelessSpec.replace(EdgelessRootBlockSpec, AIEdgelessRootBlockSpec);
+    edgelessSpec.replace(
+      EdgelessRootBlockSpec,
+      createAIEdgelessRootBlockSpec(framework)
+    );
  }

  return edgelessSpec.value;
--- a/packages/frontend/core/src/components/blocksuite/block-suite-editor/specs/page.ts
+++ b/packages/frontend/core/src/components/blocksuite/block-suite-editor/specs/page.ts
@@ -1,4 +1,4 @@
-import { AIPageRootBlockSpec } from '@affine/core/blocksuite/presets/ai';
+import { createAIPageRootBlockSpec } from '@affine/core/blocksuite/presets/ai';
 import { FeatureFlagService } from '@affine/core/modules/feature-flag';
 import { PageRootBlockSpec, SpecProvider } from '@blocksuite/affine/blocks';
 import type { ExtensionType } from '@blocksuite/affine/store';
@@ -16,7 +16,7 @@ export function createPageModeSpecs(
  enableAffineExtension(framework, pageSpec);
  if (enableAI) {
    enableAIExtension(pageSpec);
-    pageSpec.replace(PageRootBlockSpec, AIPageRootBlockSpec);
+    pageSpec.replace(PageRootBlockSpec, createAIPageRootBlockSpec(framework));
  }
  return pageSpec.value;
 }
--- a/packages/frontend/core/src/components/providers/workspace-side-effects.tsx
+++ b/packages/frontend/core/src/components/providers/workspace-side-effects.tsx
@@ -8,6 +8,7 @@ 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,
@@ -139,6 +140,7 @@ export const WorkspaceSideEffects = () => {
  const graphqlService = useService(GraphQLService);
  const eventSourceService = useService(EventSourceService);
  const fetchService = useService(FetchService);
+  const networkSearchService = useService(AINetworkSearchService);

  useEffect(() => {
    const dispose = setupAIProvider(
@@ -147,12 +149,19 @@ export const WorkspaceSideEffects = () => {
        fetchService.fetch,
        eventSourceService.eventSource
      ),
-      globalDialogService
+      globalDialogService,
+      networkSearchService
    );
    return () => {
      dispose();
    };
-  }, [eventSourceService, fetchService, globalDialogService, graphqlService]);
+  }, [
+    eventSourceService,
+    fetchService,
+    globalDialogService,
+    graphqlService,
+    networkSearchService,
+  ]);

  useRegisterWorkspaceCommands();
  useRegisterNavigationCommands();