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/blocksuite/presets/ai/ai-panel.ts
+++ b/packages/frontend/core/src/blocksuite/presets/ai/ai-panel.ts
@@ -1,3 +1,4 @@
+import { AINetworkSearchService } from '@affine/core/modules/ai-button/services/network-search';
 import type { EditorHost } from '@blocksuite/affine/block-std';
 import {
  type AffineAIPanelWidget,
@@ -9,6 +10,7 @@ import {
  NoteDisplayMode,
 } from '@blocksuite/affine/blocks';
 import { assertExists, Bound } from '@blocksuite/affine/global/utils';
+import type { FrameworkProvider } from '@toeverything/infra';
 import type { TemplateResult } from 'lit';

 import { createTextRenderer, insertFromMarkdown } from '../_common';
@@ -287,14 +289,21 @@ export function buildCopyConfig(panel: AffineAIPanelWidget) {
 }

 export function buildAIPanelConfig(
-  panel: AffineAIPanelWidget
+  panel: AffineAIPanelWidget,
+  framework: FrameworkProvider
 ): AffineAIPanelWidgetConfig {
  const ctx = new AIContext();
+  const searchService = framework.get(AINetworkSearchService);
  return {
    answerRenderer: createTextRenderer(panel.host, { maxHeight: 320 }),
    finishStateConfig: buildFinishConfig(panel, 'chat', ctx),
    generatingStateConfig: buildGeneratingConfig(),
    errorStateConfig: buildErrorConfig(panel),
    copy: buildCopyConfig(panel),
+    networkSearchConfig: {
+      visible: searchService.visible,
+      enabled: searchService.enabled,
+      setEnabled: searchService.setEnabled,
+    },
  };
 }
--- a/packages/frontend/core/src/blocksuite/presets/ai/ai-spec.ts
+++ b/packages/frontend/core/src/blocksuite/presets/ai/ai-spec.ts
@@ -23,6 +23,7 @@ import {
 } from '@blocksuite/affine/blocks';
 import { assertInstanceOf } from '@blocksuite/affine/global/utils';
 import type { ExtensionType } from '@blocksuite/affine/store';
+import type { FrameworkProvider } from '@toeverything/infra';
 import { literal, unsafeStatic } from 'lit/static-html.js';

 import { buildAIPanelConfig } from './ai-panel';
@@ -36,96 +37,110 @@ import { setupImageToolbarAIEntry } from './entries/image-toolbar/setup-image-to
 import { setupSlashMenuAIEntry } from './entries/slash-menu/setup-slash-menu';
 import { setupSpaceAIEntry } from './entries/space/setup-space';

-class AIPageRootWatcher extends BlockServiceWatcher {
-  static override readonly flavour = 'affine:page';
+function getAIPageRootWatcher(framework: FrameworkProvider) {
+  class AIPageRootWatcher extends BlockServiceWatcher {
+    static override readonly flavour = 'affine:page';

-  override mounted() {
-    super.mounted();
-    this.blockService.specSlots.widgetConnected.on(view => {
-      if (view.component instanceof AffineAIPanelWidget) {
-        view.component.style.width = '630px';
-        view.component.config = buildAIPanelConfig(view.component);
-        setupSpaceAIEntry(view.component);
-      }
+    override mounted() {
+      super.mounted();
+      this.blockService.specSlots.widgetConnected.on(view => {
+        if (view.component instanceof AffineAIPanelWidget) {
+          view.component.style.width = '630px';
+          view.component.config = buildAIPanelConfig(view.component, framework);
+          setupSpaceAIEntry(view.component);
+        }

-      if (view.component instanceof AffineFormatBarWidget) {
-        setupFormatBarAIEntry(view.component);
-      }
+        if (view.component instanceof AffineFormatBarWidget) {
+          setupFormatBarAIEntry(view.component);
+        }

-      if (view.component instanceof AffineSlashMenuWidget) {
-        setupSlashMenuAIEntry(view.component);
-      }
-    });
+        if (view.component instanceof AffineSlashMenuWidget) {
+          setupSlashMenuAIEntry(view.component);
+        }
+      });
+    }
  }
+  return AIPageRootWatcher;
 }

-export const AIPageRootBlockSpec: ExtensionType[] = [
-  ...PageRootBlockSpec,
-  AIPageRootWatcher,
-  {
-    setup: di => {
-      di.override(WidgetViewMapIdentifier('affine:page'), () => {
-        return {
-          ...pageRootWidgetViewMap,
-          [AFFINE_AI_PANEL_WIDGET]: literal`${unsafeStatic(
-            AFFINE_AI_PANEL_WIDGET
-          )}`,
-        };
-      });
+export function createAIPageRootBlockSpec(
+  framework: FrameworkProvider
+): ExtensionType[] {
+  return [
+    ...PageRootBlockSpec,
+    getAIPageRootWatcher(framework),
+    {
+      setup: di => {
+        di.override(WidgetViewMapIdentifier('affine:page'), () => {
+          return {
+            ...pageRootWidgetViewMap,
+            [AFFINE_AI_PANEL_WIDGET]: literal`${unsafeStatic(
+              AFFINE_AI_PANEL_WIDGET
+            )}`,
+          };
+        });
+      },
    },
-  },
-];
-
-class AIEdgelessRootWatcher extends BlockServiceWatcher {
-  static override readonly flavour = 'affine:page';
-
-  override mounted() {
-    super.mounted();
-    this.blockService.specSlots.widgetConnected.on(view => {
-      if (view.component instanceof AffineAIPanelWidget) {
-        view.component.style.width = '430px';
-        view.component.config = buildAIPanelConfig(view.component);
-        setupSpaceAIEntry(view.component);
-      }
-
-      if (view.component instanceof EdgelessCopilotWidget) {
-        setupEdgelessCopilot(view.component);
-      }
-
-      if (view.component instanceof EdgelessElementToolbarWidget) {
-        setupEdgelessElementToolbarAIEntry(view.component);
-      }
-
-      if (view.component instanceof AffineFormatBarWidget) {
-        setupFormatBarAIEntry(view.component);
-      }
-
-      if (view.component instanceof AffineSlashMenuWidget) {
-        setupSlashMenuAIEntry(view.component);
-      }
-    });
-  }
+  ];
 }

-export const AIEdgelessRootBlockSpec: ExtensionType[] = [
-  ...EdgelessRootBlockSpec,
-  AIEdgelessRootWatcher,
-  {
-    setup: di => {
-      di.override(WidgetViewMapIdentifier('affine:page'), () => {
-        return {
-          ...edgelessRootWidgetViewMap,
-          [AFFINE_EDGELESS_COPILOT_WIDGET]: literal`${unsafeStatic(
-            AFFINE_EDGELESS_COPILOT_WIDGET
-          )}`,
-          [AFFINE_AI_PANEL_WIDGET]: literal`${unsafeStatic(
-            AFFINE_AI_PANEL_WIDGET
-          )}`,
-        };
+function getAIEdgelessRootWatcher(framework: FrameworkProvider) {
+  class AIEdgelessRootWatcher extends BlockServiceWatcher {
+    static override readonly flavour = 'affine:page';
+
+    override mounted() {
+      super.mounted();
+      this.blockService.specSlots.widgetConnected.on(view => {
+        if (view.component instanceof AffineAIPanelWidget) {
+          view.component.style.width = '430px';
+          view.component.config = buildAIPanelConfig(view.component, framework);
+          setupSpaceAIEntry(view.component);
+        }
+
+        if (view.component instanceof EdgelessCopilotWidget) {
+          setupEdgelessCopilot(view.component);
+        }
+
+        if (view.component instanceof EdgelessElementToolbarWidget) {
+          setupEdgelessElementToolbarAIEntry(view.component);
+        }
+
+        if (view.component instanceof AffineFormatBarWidget) {
+          setupFormatBarAIEntry(view.component);
+        }
+
+        if (view.component instanceof AffineSlashMenuWidget) {
+          setupSlashMenuAIEntry(view.component);
+        }
      });
+    }
+  }
+  return AIEdgelessRootWatcher;
+}
+
+export function createAIEdgelessRootBlockSpec(
+  framework: FrameworkProvider
+): ExtensionType[] {
+  return [
+    ...EdgelessRootBlockSpec,
+    getAIEdgelessRootWatcher(framework),
+    {
+      setup: di => {
+        di.override(WidgetViewMapIdentifier('affine:page'), () => {
+          return {
+            ...edgelessRootWidgetViewMap,
+            [AFFINE_EDGELESS_COPILOT_WIDGET]: literal`${unsafeStatic(
+              AFFINE_EDGELESS_COPILOT_WIDGET
+            )}`,
+            [AFFINE_AI_PANEL_WIDGET]: literal`${unsafeStatic(
+              AFFINE_AI_PANEL_WIDGET
+            )}`,
+          };
+        });
+      },
    },
-  },
-];
+  ];
+}

 class AIParagraphBlockWatcher extends BlockServiceWatcher {
  static override readonly flavour = 'affine:paragraph';
--- a/packages/frontend/core/src/blocksuite/presets/ai/chat-panel/chat-panel-input.ts
+++ b/packages/frontend/core/src/blocksuite/presets/ai/chat-panel/chat-panel-input.ts
@@ -1,6 +1,14 @@
+import { stopPropagation } from '@affine/core/utils';
 import type { EditorHost } from '@blocksuite/affine/block-std';
 import { type AIError, openFileOrFiles } from '@blocksuite/affine/blocks';
-import { assertExists, WithDisposable } from '@blocksuite/affine/global/utils';
+import {
+  assertExists,
+  SignalWatcher,
+  WithDisposable,
+} from '@blocksuite/affine/global/utils';
+import { unsafeCSSVarV2 } from '@blocksuite/affine-shared/theme';
+import { ImageIcon, PublishIcon } from '@blocksuite/icons/lit';
+import type { Signal } from '@preact/signals-core';
 import { css, html, LitElement, nothing } from 'lit';
 import { property, query, state } from 'lit/decorators.js';
 import { repeat } from 'lit/directives/repeat.js';
@@ -10,7 +18,6 @@ import {
  ChatClearIcon,
  ChatSendIcon,
  CloseIcon,
-  ImageIcon,
 } from '../_common/icons';
 import { AIProvider } from '../provider';
 import { reportResponse } from '../utils/action-reporter';
@@ -24,7 +31,13 @@ function getFirstTwoLines(text: string) {
  return lines.slice(0, 2);
 }

-export class ChatPanelInput extends WithDisposable(LitElement) {
+export interface AINetworkSearchConfig {
+  visible: Signal<boolean | undefined>;
+  enabled: Signal<boolean | undefined>;
+  setEnabled: (state: boolean) => void;
+}
+
+export class ChatPanelInput extends SignalWatcher(WithDisposable(LitElement)) {
  static override styles = css`
    .chat-panel-input {
      display: flex;
@@ -104,10 +117,28 @@ export class ChatPanelInput extends WithDisposable(LitElement) {
        margin-left: auto;
      }

-      .image-upload {
+      .image-upload,
+      .chat-network-search {
        display: flex;
        justify-content: center;
        align-items: center;
+        svg {
+          width: 20px;
+          height: 20px;
+          color: ${unsafeCSSVarV2('icon/primary')};
+        }
+      }
+      .chat-network-search[data-active='true'] svg {
+        color: ${unsafeCSSVarV2('icon/activated')};
+      }
+
+      .image-upload[aria-disabled='true'],
+      .chat-network-search[aria-disabled='true'] {
+        cursor: not-allowed;
+      }
+      .image-upload[aria-disabled='true'] svg,
+      .chat-network-search[aria-disabled='true'] svg {
+        color: var(--affine-text-disable-color) !important;
      }
    }

@@ -235,6 +266,9 @@ export class ChatPanelInput extends WithDisposable(LitElement) {
  @property({ attribute: false })
  accessor cleanupHistories!: () => Promise<void>;

+  @property({ attribute: false })
+  accessor networkSearchConfig!: AINetworkSearchConfig;
+
  private _addImages(images: File[]) {
    const oldImages = this.chatContextValue.images;
    this.updateContext({
@@ -296,6 +330,23 @@ export class ChatPanelInput extends WithDisposable(LitElement) {
    `;
  }

+  private readonly _toggleNetworkSearch = (e: MouseEvent) => {
+    e.preventDefault();
+    e.stopPropagation();
+
+    const enable = this.networkSearchConfig.enabled.value;
+    this.networkSearchConfig.setEnabled(!enable);
+  };
+
+  private readonly _uploadImageFiles = async (_e: MouseEvent) => {
+    const images = await openFileOrFiles({
+      acceptType: 'Images',
+      multiple: true,
+    });
+    if (!images) return;
+    this._addImages(images);
+  };
+
  override connectedCallback() {
    super.connectedCallback();

@@ -305,7 +356,7 @@ export class ChatPanelInput extends WithDisposable(LitElement) {
          if (this.host === host) {
            context && this.updateContext(context);
            await this.updateComplete;
-            await this.send(input);
+            input && (await this.send(input));
          }
        }
      )
@@ -316,7 +367,9 @@ export class ChatPanelInput extends WithDisposable(LitElement) {
    const { images, status } = this.chatContextValue;
    const hasImages = images.length > 0;
    const maxHeight = hasImages ? 272 + 2 : 200 + 2;
-
+    const networkDisabled = !!this.chatContextValue.images.length;
+    const networkActive = !!this.networkSearchConfig.enabled.value;
+    const uploadDisabled = networkActive && !networkDisabled;
    return html`<style>
        .chat-panel-input {
          border-color: ${this.focused
@@ -364,8 +417,7 @@ export class ChatPanelInput extends WithDisposable(LitElement) {
          }}
          @keydown=${async (evt: KeyboardEvent) => {
            if (evt.key === 'Enter' && !evt.shiftKey && !evt.isComposing) {
-              evt.preventDefault();
-              await this.send();
+              this._onTextareaSend(evt);
            }
          }}
          @focus=${() => {
@@ -399,19 +451,29 @@ export class ChatPanelInput extends WithDisposable(LitElement) {
          >
            ${ChatClearIcon}
          </div>
+          ${this.networkSearchConfig.visible.value
+            ? html`
+                <div
+                  class="chat-network-search"
+                  data-testid="chat-network-search"
+                  aria-disabled=${networkDisabled}
+                  data-active=${networkActive}
+                  @click=${networkDisabled
+                    ? undefined
+                    : this._toggleNetworkSearch}
+                  @pointerdown=${stopPropagation}
+                >
+                  ${PublishIcon()}
+                </div>
+              `
+            : nothing}
          ${images.length < MaximumImageCount
            ? html`<div
                class="image-upload"
-                @click=${async () => {
-                  const images = await openFileOrFiles({
-                    acceptType: 'Images',
-                    multiple: true,
-                  });
-                  if (!images) return;
-                  this._addImages(images);
-                }}
+                aria-disabled=${uploadDisabled}
+                @click=${uploadDisabled ? undefined : this._uploadImageFiles}
              >
-                ${ImageIcon}
+                ${ImageIcon()}
              </div>`
            : nothing}
          ${status === 'transmitting'
@@ -425,7 +487,7 @@ export class ChatPanelInput extends WithDisposable(LitElement) {
                ${ChatAbortIcon}
              </div>`
            : html`<div
-                @click="${this.send}"
+                @click="${this._onTextareaSend}"
                class="chat-panel-send"
                aria-disabled=${this.isInputEmpty}
                data-testid="chat-panel-send"
@@ -436,19 +498,30 @@ export class ChatPanelInput extends WithDisposable(LitElement) {
      </div>`;
  }

-  send = async (input?: string) => {
+  private readonly _onTextareaSend = (e: MouseEvent | KeyboardEvent) => {
+    e.preventDefault();
+    e.stopPropagation();
+
+    const value = this.textarea.value.trim();
+    if (value.length === 0) return;
+
+    this.textarea.value = '';
+    this.isInputEmpty = true;
+    this.textarea.style.height = 'unset';
+
+    this.send(value).catch(console.error);
+  };
+
+  send = async (text: string) => {
    const { status, markdown } = this.chatContextValue;
    if (status === 'loading' || status === 'transmitting') return;

-    const text = input || this.textarea.value;
    const { images } = this.chatContextValue;
    if (!text && images.length === 0) {
      return;
    }
    const { doc } = this.host;
-    this.textarea.value = '';
-    this.isInputEmpty = true;
-    this.textarea.style.height = 'unset';
+
    this.updateContext({
      images: [],
      status: 'loading',
--- a/packages/frontend/core/src/blocksuite/presets/ai/chat-panel/index.ts
+++ b/packages/frontend/core/src/blocksuite/presets/ai/chat-panel/index.ts
@@ -17,6 +17,7 @@ import {
  getSelectedTextContent,
 } from '../utils/selection-utils';
 import type { ChatAction, ChatContextValue, ChatItem } from './chat-context';
+import type { AINetworkSearchConfig } from './chat-panel-input';
 import type { ChatPanelMessages } from './chat-panel-messages';

 export class ChatPanel extends WithDisposable(ShadowlessElement) {
@@ -143,6 +144,9 @@ export class ChatPanel extends WithDisposable(ShadowlessElement) {
  @property({ attribute: false })
  accessor doc!: Store;

+  @property({ attribute: false })
+  accessor networkSearchConfig!: AINetworkSearchConfig;
+
  @state()
  accessor isLoading = false;

@@ -280,6 +284,7 @@ export class ChatPanel extends WithDisposable(ShadowlessElement) {
      ></chat-panel-messages>
      <chat-panel-input
        .chatContextValue=${this.chatContextValue}
+        .networkSearchConfig=${this.networkSearchConfig}
        .updateContext=${this.updateContext}
        .host=${this.host}
        .cleanupHistories=${this._cleanupHistories}
--- 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();
--- a/packages/frontend/core/src/desktop/dialogs/setting/general-setting/experimental-features/index.tsx
+++ b/packages/frontend/core/src/desktop/dialogs/setting/general-setting/experimental-features/index.tsx
@@ -105,7 +105,13 @@ const feedbackLink: Record<NonNullable<Flag['feedbackType']>, string> = {
  github: 'https://github.com/toeverything/AFFiNE/issues',
 };

-const ExperimentalFeaturesItem = ({ flag }: { flag: Flag }) => {
+const ExperimentalFeaturesItem = ({
+  flag,
+  flagKey,
+}: {
+  flag: Flag;
+  flagKey: string;
+}) => {
  const value = useLiveData(flag.$);
  const t = useI18n();
  const onChange = useCallback(
@@ -128,7 +134,7 @@ const ExperimentalFeaturesItem = ({ flag }: { flag: Flag }) => {
    <div className={styles.rowContainer}>
      <div className={styles.switchRow}>
        {t[flag.displayName]()}
-        <Switch checked={value} onChange={onChange} />
+        <Switch data-testid={flagKey} checked={value} onChange={onChange} />
      </div>
      {!!flag.description && (
        <Tooltip content={t[flag.description]()}>
@@ -175,6 +181,7 @@ const ExperimentalFeaturesMain = () => {
        {Object.keys(AFFINE_FLAGS).map(key => (
          <ExperimentalFeaturesItem
            key={key}
+            flagKey={key}
            flag={featureFlagService.flags[key as keyof AFFINE_FLAGS]}
          />
        ))}
--- a/packages/frontend/core/src/desktop/pages/workspace/detail-page/tabs/chat.tsx
+++ b/packages/frontend/core/src/desktop/pages/workspace/detail-page/tabs/chat.tsx
@@ -1,9 +1,11 @@
 import { ChatPanel } from '@affine/core/blocksuite/presets/ai';
+import { AINetworkSearchService } from '@affine/core/modules/ai-button/services/network-search';
 import {
  DocModeProvider,
  RefNodeSlotsProvider,
 } from '@blocksuite/affine/blocks';
 import type { AffineEditorContainer } from '@blocksuite/affine/presets';
+import { useFramework } from '@toeverything/infra';
 import { forwardRef, useEffect, useRef } from 'react';

 import * as styles from './chat.css';
@@ -20,6 +22,7 @@ export const EditorChatPanel = forwardRef(function EditorChatPanel(
 ) {
  const chatPanelRef = useRef<ChatPanel | null>(null);
  const containerRef = useRef<HTMLDivElement | null>(null);
+  const framework = useFramework();

  useEffect(() => {
    if (onLoad && chatPanelRef.current) {
@@ -45,6 +48,13 @@ export const EditorChatPanel = forwardRef(function EditorChatPanel(
      chatPanelRef.current.host = editor.host;
      chatPanelRef.current.doc = editor.doc;
      containerRef.current?.append(chatPanelRef.current);
+      const searchService = framework.get(AINetworkSearchService);
+      const networkSearchConfig = {
+        visible: searchService.visible,
+        enabled: searchService.enabled,
+        setEnabled: searchService.setEnabled,
+      };
+      chatPanelRef.current.networkSearchConfig = networkSearchConfig;
    } else {
      chatPanelRef.current.host = editor.host;
      chatPanelRef.current.doc = editor.doc;
@@ -63,7 +73,7 @@ export const EditorChatPanel = forwardRef(function EditorChatPanel(
    ];

    return () => disposable.forEach(d => d?.dispose());
-  }, [editor]);
+  }, [editor, framework]);

  return <div className={styles.root} ref={containerRef} />;
 });
--- a/packages/frontend/core/src/mobile/dialogs/setting/experimental/index.tsx
+++ b/packages/frontend/core/src/mobile/dialogs/setting/experimental/index.tsx
@@ -46,6 +46,7 @@ const ExperimentalFeatureList = () => {
      {Object.keys(AFFINE_FLAGS).map(key => (
        <ExperimentalFeaturesItem
          key={key}
+          flagKey={key}
          flag={featureFlagService.flags[key as keyof AFFINE_FLAGS]}
        />
      ))}
@@ -53,7 +54,13 @@ const ExperimentalFeatureList = () => {
  );
 };

-const ExperimentalFeaturesItem = ({ flag }: { flag: Flag }) => {
+const ExperimentalFeaturesItem = ({
+  flag,
+  flagKey,
+}: {
+  flag: Flag;
+  flagKey: string;
+}) => {
  const t = useI18n();
  const value = useLiveData(flag.$);

@@ -72,7 +79,7 @@ const ExperimentalFeaturesItem = ({ flag }: { flag: Flag }) => {
    <li>
      <div className={styles.itemBlock}>
        {t[flag.displayName]()}
-        <Switch checked={value} onChange={onChange} />
+        <Switch data-testid={flagKey} checked={value} onChange={onChange} />
      </div>
      {flag.description ? (
        <div className={styles.itemDescription}>{t[flag.description]()}</div>
--- a/packages/frontend/core/src/modules/ai-button/index.ts
+++ b/packages/frontend/core/src/modules/ai-button/index.ts
@@ -3,11 +3,21 @@ export { AIButtonService } from './services/ai-button';

 import type { Framework } from '@toeverything/infra';

+import { FeatureFlagService } from '../feature-flag';
+import { GlobalStateService } from '../storage';
 import { AIButtonProvider } from './provider/ai-button';
 import { AIButtonService } from './services/ai-button';
+import { AINetworkSearchService } from './services/network-search';

 export const configureAIButtonModule = (framework: Framework) => {
  framework.service(AIButtonService, container => {
    return new AIButtonService(container.getOptional(AIButtonProvider));
  });
 };
+
+export function configureAINetworkSearchModule(framework: Framework) {
+  framework.service(AINetworkSearchService, [
+    GlobalStateService,
+    FeatureFlagService,
+  ]);
+}
--- a/packages/frontend/core/src/modules/ai-button/services/network-search.ts
+++ b/packages/frontend/core/src/modules/ai-button/services/network-search.ts
@@ -0,0 +1,45 @@
+import {
+  createSignalFromObservable,
+  type Signal,
+} from '@blocksuite/affine-shared/utils';
+import { LiveData, Service } from '@toeverything/infra';
+
+import type { FeatureFlagService } from '../../feature-flag';
+import type { GlobalStateService } from '../../storage';
+
+const AI_NETWORK_SEARCH_KEY = 'AINetworkSearch';
+
+export class AINetworkSearchService extends Service {
+  constructor(
+    private readonly globalStateService: GlobalStateService,
+    private readonly featureFlagService: FeatureFlagService
+  ) {
+    super();
+
+    const { signal: enabled, cleanup: enabledCleanup } =
+      createSignalFromObservable<boolean | undefined>(this._enabled$, false);
+    this.enabled = enabled;
+    this.disposables.push(enabledCleanup);
+
+    const { signal: visible, cleanup: visibleCleanup } =
+      createSignalFromObservable<boolean | undefined>(this._visible$, false);
+    this.visible = visible;
+    this.disposables.push(visibleCleanup);
+  }
+
+  visible: Signal<boolean | undefined>;
+
+  enabled: Signal<boolean | undefined>;
+
+  private readonly _visible$ =
+    this.featureFlagService.flags.enable_ai_network_search.$;
+
+  private readonly _enabled$ = LiveData.from(
+    this.globalStateService.globalState.watch<boolean>(AI_NETWORK_SEARCH_KEY),
+    false
+  );
+
+  setEnabled = (enabled: boolean) => {
+    this.globalStateService.globalState.set(AI_NETWORK_SEARCH_KEY, enabled);
+  };
+}
--- a/packages/frontend/core/src/modules/feature-flag/constant.ts
+++ b/packages/frontend/core/src/modules/feature-flag/constant.ts
@@ -16,6 +16,15 @@ export const AFFINE_FLAGS = {
    configurable: true,
    defaultState: true,
  },
+  enable_ai_network_search: {
+    category: 'affine',
+    displayName:
+      'com.affine.settings.workspace.experimental-features.enable-ai-network-search.name',
+    description:
+      'com.affine.settings.workspace.experimental-features.enable-ai-network-search.description',
+    configurable: true,
+    defaultState: false,
+  },
  enable_database_full_width: {
    category: 'blocksuite',
    bsFlag: 'enable_database_full_width',
--- a/packages/frontend/core/src/modules/index.ts
+++ b/packages/frontend/core/src/modules/index.ts
@@ -1,7 +1,10 @@
 import { configureQuotaModule } from '@affine/core/modules/quota';
 import { type Framework } from '@toeverything/infra';

-import { configureAIButtonModule } from './ai-button';
+import {
+  configureAIButtonModule,
+  configureAINetworkSearchModule,
+} from './ai-button';
 import { configureAppSidebarModule } from './app-sidebar';
 import { configAtMenuConfigModule } from './at-menu-config';
 import { configureCloudModule } from './cloud';
@@ -89,5 +92,6 @@ export function configureCommonModules(framework: Framework) {
  configAtMenuConfigModule(framework);
  configureDndModule(framework);
  configureCommonGlobalStorageImpls(framework);
+  configureAINetworkSearchModule(framework);
  configureAIButtonModule(framework);
 }