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>

.github/actions/copilot-test/action.yml

@@ -26,6 +26,7 @@ runs:
         DEV_SERVER_URL: http://localhost:8080
         COPILOT_OPENAI_API_KEY: ${{ inputs.openai-key }}
         COPILOT_FAL_API_KEY: ${{ inputs.fal-key }}
+        COPILOT_PERPLEXITY_API_KEY: ${{ inputs.perplexity-key }}
 
     - name: Upload test results
       if: ${{ failure() }}

.github/actions/deploy/deploy.mjs

@@ -17,6 +17,7 @@ const {
   METRICS_CUSTOMER_IO_TOKEN,
   COPILOT_OPENAI_API_KEY,
   COPILOT_FAL_API_KEY,
+  COPILOT_PERPLEXITY_API_KEY,
   COPILOT_UNSPLASH_API_KEY,
   MAILER_SENDER,
   MAILER_USER,
@@ -147,6 +148,7 @@ const createHelmCommand = ({ isDryRun }) => {
     `--set        graphql.app.copilot.enabled=true`,
     `--set-string graphql.app.copilot.openai.key="${COPILOT_OPENAI_API_KEY}"`,
     `--set-string graphql.app.copilot.fal.key="${COPILOT_FAL_API_KEY}"`,
+    `--set-string graphql.app.copilot.perplexity.key="${COPILOT_PERPLEXITY_API_KEY}"`,
     `--set-string graphql.app.copilot.unsplash.key="${COPILOT_UNSPLASH_API_KEY}"`,
     `--set-string graphql.app.mailer.sender="${MAILER_SENDER}"`,
     `--set-string graphql.app.mailer.user="${MAILER_USER}"`,

.github/helm/affine/charts/graphql/templates/deployment.yaml

@@ -157,6 +157,11 @@ spec:
               secretKeyRef:
                 name: "{{ .Values.app.copilot.secretName }}"
                 key: falSecret
+          - name: COPILOT_PERPLEXITY_API_KEY
+            valueFrom:
+              secretKeyRef:
+                name: "{{ .Values.app.copilot.secretName }}"
+                key: perplexitySecret
           - name: COPILOT_UNSPLASH_API_KEY
             valueFrom:
               secretKeyRef:

.github/workflows/build-test.yml

@@ -531,6 +531,7 @@ jobs:
           CARGO_TARGET_DIR: '${{ github.workspace }}/target'
           COPILOT_OPENAI_API_KEY: ${{ secrets.COPILOT_OPENAI_API_KEY }}
           COPILOT_FAL_API_KEY: ${{ secrets.COPILOT_FAL_API_KEY }}
+          COPILOT_PERPLEXITY_API_KEY: ${{ secrets.COPILOT_PERPLEXITY_API_KEY }}
 
       - name: Upload server test coverage results
         if: ${{ steps.check-blocksuite-update.outputs.skip != 'true' || steps.apifilter.outputs.changed == 'true' }}
@@ -619,6 +620,7 @@ jobs:
           script: yarn affine @affine-test/affine-cloud-copilot e2e --forbid-only --shard=${{ matrix.shardIndex }}/${{ matrix.shardTotal }}
           openai-key: ${{ secrets.COPILOT_OPENAI_API_KEY }}
           fal-key: ${{ secrets.COPILOT_FAL_API_KEY }}
+          perplexity-key: ${{ secrets.COPILOT_PERPLEXITY_API_KEY }}
 
   server-e2e-test:
     name: ${{ matrix.tests.name }}
@@ -703,6 +705,7 @@ jobs:
           DEV_SERVER_URL: http://localhost:8080
           COPILOT_OPENAI_API_KEY: 1
           COPILOT_FAL_API_KEY: 1
+          COPILOT_PERPLEXITY_API_KEY: 1
 
       - name: Upload test results
         if: ${{ failure() }}

.github/workflows/copilot-test.yml

@@ -84,6 +84,7 @@ jobs:
           CARGO_TARGET_DIR: '${{ github.workspace }}/target'
           COPILOT_OPENAI_API_KEY: ${{ secrets.COPILOT_OPENAI_API_KEY }}
           COPILOT_FAL_API_KEY: ${{ secrets.COPILOT_FAL_API_KEY }}
+          COPILOT_PERPLEXITY_API_KEY: ${{ secrets.COPILOT_PERPLEXITY_API_KEY }}
 
       - name: Upload server test coverage results
         uses: codecov/codecov-action@v5
@@ -147,6 +148,7 @@ jobs:
           script: yarn affine @affine-test/affine-cloud-copilot e2e --forbid-only --shard=${{ matrix.shardIndex }}/${{ matrix.shardTotal }}
           openai-key: ${{ secrets.COPILOT_OPENAI_API_KEY }}
           fal-key: ${{ secrets.COPILOT_FAL_API_KEY }}
+          perplexity-key: ${{ secrets.COPILOT_PERPLEXITY_API_KEY }}
 
   test-done:
     needs:

.github/workflows/deploy.yml

@@ -98,6 +98,7 @@ jobs:
           CAPTCHA_TURNSTILE_SECRET: ${{ secrets.CAPTCHA_TURNSTILE_SECRET }}
           COPILOT_OPENAI_API_KEY: ${{ secrets.COPILOT_OPENAI_API_KEY }}
           COPILOT_FAL_API_KEY: ${{ secrets.COPILOT_FAL_API_KEY }}
+          COPILOT_PERPLEXITY_API_KEY: ${{ secrets.COPILOT_PERPLEXITY_API_KEY }}
           COPILOT_UNSPLASH_API_KEY: ${{ secrets.COPILOT_UNSPLASH_API_KEY }}
           METRICS_CUSTOMER_IO_TOKEN: ${{ secrets.METRICS_CUSTOMER_IO_TOKEN }}
           MAILER_SENDER: ${{ secrets.OAUTH_EMAIL_SENDER }}

blocksuite/blocks/src/root-block/widgets/ai-panel/ai-panel.ts

@@ -63,7 +63,7 @@ export class AffineAIPanelWidget extends WidgetComponent {
       box-sizing: border-box;
       width: 100%;
       height: fit-content;
-      padding: 8px 0;
+      padding: 10px 0;
     }
 
     .ai-panel-container:not(:has(ai-panel-generating)) {
@@ -474,6 +474,7 @@ export class AffineAIPanelWidget extends WidgetComponent {
             .onBlur=${this.discard}
             .onFinish=${this._inputFinish}
             .onInput=${this.onInput}
+            .networkSearchConfig=${config.networkSearchConfig}
           ></ai-panel-input>`,
       ],
       [

blocksuite/blocks/src/root-block/widgets/ai-panel/components/state/input.ts

@@ -1,11 +1,14 @@
 import { AIStarIcon } from '@blocksuite/affine-components/icons';
+import { unsafeCSSVarV2 } from '@blocksuite/affine-shared/theme';
 import { stopPropagation } from '@blocksuite/affine-shared/utils';
-import { WithDisposable } from '@blocksuite/global/utils';
+import { SignalWatcher, WithDisposable } from '@blocksuite/global/utils';
-import { SendIcon } from '@blocksuite/icons/lit';
+import { PublishIcon, SendIcon } from '@blocksuite/icons/lit';
 import { css, html, LitElement, nothing } from 'lit';
 import { property, query, state } from 'lit/decorators.js';
 
-export class AIPanelInput extends WithDisposable(LitElement) {
+import type { AINetworkSearchConfig } from '../../type';
+
+export class AIPanelInput extends SignalWatcher(WithDisposable(LitElement)) {
   static override styles = css`
     :host {
       width: 100%;
@@ -20,8 +23,9 @@ export class AIPanelInput extends WithDisposable(LitElement) {
       background: var(--affine-background-overlay-panel-color);
     }
 
-    .icon {
+    .star {
       display: flex;
+      padding: 2px;
       align-items: center;
     }
 
@@ -66,22 +70,36 @@ export class AIPanelInput extends WithDisposable(LitElement) {
       display: flex;
       align-items: center;
       padding: 2px;
-      gap: 10px;
+      gap: 4px;
       border-radius: 4px;
-      background: var(--affine-black-10, rgba(0, 0, 0, 0.1));
+      background: ${unsafeCSSVarV2('icon/disable')};
-
       svg {
-        width: 16px;
+        width: 20px;
-        height: 16px;
+        height: 20px;
-        color: var(--affine-pure-white, #fff);
+        color: ${unsafeCSSVarV2('button/pureWhiteText')};
       }
     }
     .arrow[data-active] {
-      background: var(--affine-brand-color, #1e96eb);
+      background: ${unsafeCSSVarV2('icon/activated')};
     }
     .arrow[data-active]:hover {
       cursor: pointer;
     }
+    .network {
+      display: flex;
+      align-items: center;
+      padding: 2px;
+      gap: 4px;
+      cursor: pointer;
+      svg {
+        width: 20px;
+        height: 20px;
+        color: ${unsafeCSSVarV2('icon/primary')};
+      }
+    }
+    .network[data-active='true'] svg {
+      color: ${unsafeCSSVarV2('icon/activated')};
+    }
   `;
 
   private readonly _onInput = () => {
@@ -101,12 +119,14 @@ export class AIPanelInput extends WithDisposable(LitElement) {
 
   private readonly _onKeyDown = (e: KeyboardEvent) => {
     if (e.key === 'Enter' && !e.shiftKey && !e.isComposing) {
-      e.preventDefault();
+      this._sendToAI(e);
-      this._sendToAI();
     }
   };
 
-  private readonly _sendToAI = () => {
+  private readonly _sendToAI = (e: MouseEvent | KeyboardEvent) => {
+    e.preventDefault();
+    e.stopPropagation();
+
     const value = this.textarea.value.trim();
     if (value.length === 0) return;
 
@@ -114,9 +134,17 @@ export class AIPanelInput extends WithDisposable(LitElement) {
     this.remove();
   };
 
+  private readonly _toggleNetworkSearch = (e: MouseEvent) => {
+    e.preventDefault();
+    e.stopPropagation();
+
+    const enable = this.networkSearchConfig.enabled.value;
+    this.networkSearchConfig.setEnabled(!enable);
+  };
+
   override render() {
     return html`<div class="root">
-      <div class="icon">${AIStarIcon}</div>
+      <div class="star">${AIStarIcon}</div>
       <div class="textarea-container">
         <textarea
           placeholder="What are your thoughts?"
@@ -131,6 +159,21 @@ export class AIPanelInput extends WithDisposable(LitElement) {
           @paste=${stopPropagation}
           @keyup=${stopPropagation}
         ></textarea>
+        ${this.networkSearchConfig.visible.value
+          ? html`
+              <div
+                class="network"
+                data-active=${!!this.networkSearchConfig.enabled.value}
+                @click=${this._toggleNetworkSearch}
+                @pointerdown=${stopPropagation}
+              >
+                ${PublishIcon()}
+                <affine-tooltip .offset=${12}
+                  >Toggle Network Search</affine-tooltip
+                >
+              </div>
+            `
+          : nothing}
         <div
           class="arrow"
           @click=${this._sendToAI}
@@ -157,6 +200,9 @@ export class AIPanelInput extends WithDisposable(LitElement) {
   @state()
   private accessor _hasContent = false;
 
+  @property({ attribute: false })
+  accessor networkSearchConfig!: AINetworkSearchConfig;
+
   @property({ attribute: false })
   accessor onFinish: ((input: string) => void) | undefined = undefined;
 

blocksuite/blocks/src/root-block/widgets/ai-panel/type.ts

@@ -2,6 +2,7 @@ import type {
   AIError,
   AIItemGroupConfig,
 } from '@blocksuite/affine-components/ai-item';
+import type { Signal } from '@preact/signals-core';
 import type { nothing, TemplateResult } from 'lit';
 
 export interface CopyConfig {
@@ -28,6 +29,12 @@ export interface AIPanelGeneratingConfig {
   stages?: string[];
 }
 
+export interface AINetworkSearchConfig {
+  visible: Signal<boolean | undefined>;
+  enabled: Signal<boolean | undefined>;
+  setEnabled: (state: boolean) => void;
+}
+
 export interface AffineAIPanelWidgetConfig {
   answerRenderer: (
     answer: string,
@@ -44,10 +51,10 @@ export interface AffineAIPanelWidgetConfig {
   finishStateConfig: AIPanelAnswerConfig;
   generatingStateConfig: AIPanelGeneratingConfig;
   errorStateConfig: AIPanelErrorConfig;
+  networkSearchConfig: AINetworkSearchConfig;
   hideCallback?: () => void;
   discardCallback?: () => void;
   inputCallback?: (input: string) => void;
-
   copy?: CopyConfig;
 }
 

packages/backend/server/.env.example

@@ -2,6 +2,7 @@
 # REDIS_SERVER_HOST=localhost
 # COPILOT_FAL_API_KEY=YOUR_KEY
 # COPILOT_OPENAI_API_KEY=YOUR_KEY
+# COPILOT_PERPLEXITY_API_KEY=YOUR_KEY
 
 # MAILER_HOST=127.0.0.1
 # MAILER_PORT=1025

packages/backend/server/package.json

@@ -59,6 +59,7 @@
     "@socket.io/redis-adapter": "^8.3.0",
     "cookie-parser": "^1.4.7",
     "dotenv": "^16.4.7",
+    "eventsource-parser": "^3.0.0",
     "express": "^4.21.2",
     "fast-xml-parser": "^4.5.0",
     "get-stream": "^9.0.1",

packages/backend/server/src/config/affine.env.ts

@@ -28,6 +28,7 @@ AFFiNE.ENV_MAP = {
   CAPTCHA_TURNSTILE_SECRET: ['plugins.captcha.turnstile.secret', 'string'],
   COPILOT_OPENAI_API_KEY: 'plugins.copilot.openai.apiKey',
   COPILOT_FAL_API_KEY: 'plugins.copilot.fal.apiKey',
+  COPILOT_PERPLEXITY_API_KEY: 'plugins.copilot.perplexity.apiKey',
   COPILOT_UNSPLASH_API_KEY: 'plugins.copilot.unsplashKey',
   REDIS_SERVER_HOST: 'redis.host',
   REDIS_SERVER_PORT: ['redis.port', 'int'],

packages/backend/server/src/plugins/copilot/config.ts

@@ -3,10 +3,12 @@ import type { ClientOptions as OpenAIClientOptions } from 'openai';
 import { defineStartupConfig, ModuleConfig } from '../../base/config';
 import { StorageConfig } from '../../base/storage/config';
 import type { FalConfig } from './providers/fal';
+import { PerplexityConfig } from './providers/perplexity';
 
 export interface CopilotStartupConfigurations {
   openai?: OpenAIClientOptions;
   fal?: FalConfig;
+  perplexity?: PerplexityConfig;
   test?: never;
   unsplashKey?: string;
   storage: StorageConfig;

packages/backend/server/src/plugins/copilot/index.ts

@@ -13,6 +13,7 @@ import {
   CopilotProviderService,
   FalProvider,
   OpenAIProvider,
+  PerplexityProvider,
   registerCopilotProvider,
 } from './providers';
 import {
@@ -26,6 +27,7 @@ import { CopilotWorkflowExecutors, CopilotWorkflowService } from './workflow';
 
 registerCopilotProvider(FalProvider);
 registerCopilotProvider(OpenAIProvider);
+registerCopilotProvider(PerplexityProvider);
 
 @Plugin({
   name: 'copilot',

packages/backend/server/src/plugins/copilot/prompt/prompts.ts

@@ -952,6 +952,11 @@ const chat: Prompt[] = [
       },
     ],
   },
+  {
+    name: 'Search With AFFiNE AI',
+    model: 'llama-3.1-sonar-small-128k-online',
+    messages: [],
+  },
   // use for believer plan
   {
     name: 'Chat With AFFiNE AI - Believer',

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

@@ -124,9 +124,7 @@ export class CopilotProviderService {
     if (!this.cachedProviders.has(provider)) {
       this.cachedProviders.set(provider, this.create(provider));
     }
-
+    return this.cachedProviders.get(provider) as CopilotProvider;
-    // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
-    return this.cachedProviders.get(provider)!;
   }
 
   async getProviderByCapability<C extends CopilotCapability>(
@@ -196,3 +194,4 @@ export class CopilotProviderService {
 
 export { FalProvider } from './fal';
 export { OpenAIProvider } from './openai';
+export { PerplexityProvider } from './perplexity';

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

@@ -0,0 +1,325 @@
+import assert from 'node:assert';
+
+import { EventSourceParserStream } from 'eventsource-parser/stream';
+import { z } from 'zod';
+
+import {
+  CopilotPromptInvalid,
+  CopilotProviderSideError,
+  metrics,
+} from '../../../base';
+import {
+  CopilotCapability,
+  CopilotChatOptions,
+  CopilotProviderType,
+  CopilotTextToTextProvider,
+  PromptMessage,
+} from '../types';
+
+export type PerplexityConfig = {
+  apiKey: string;
+  endpoint?: string;
+};
+
+const PerplexityErrorSchema = z.object({
+  detail: z.array(
+    z.object({
+      loc: z.array(z.string()),
+      msg: z.string(),
+      type: z.string(),
+    })
+  ),
+});
+
+const PerplexityDataSchema = z.object({
+  citations: z.array(z.string()),
+  choices: z.array(
+    z.object({
+      message: z.object({
+        content: z.string(),
+        role: z.literal('assistant'),
+      }),
+      delta: z.object({
+        content: z.string(),
+        role: z.literal('assistant'),
+      }),
+      finish_reason: z.union([z.literal('stop'), z.literal(null)]),
+    })
+  ),
+});
+
+const PerplexitySchema = z.union([PerplexityDataSchema, PerplexityErrorSchema]);
+
+export class CitationParser {
+  private readonly SQUARE_BRACKET_OPEN = '[';
+
+  private readonly SQUARE_BRACKET_CLOSE = ']';
+
+  private readonly PARENTHESES_OPEN = '(';
+
+  private startToken: string[] = [];
+
+  private endToken: string[] = [];
+
+  private numberToken: string[] = [];
+
+  public parse(content: string, citations: string[]) {
+    let result = '';
+    const contentArray = content.split('');
+    for (const [index, char] of contentArray.entries()) {
+      if (char === this.SQUARE_BRACKET_OPEN) {
+        if (this.numberToken.length === 0) {
+          this.startToken.push(char);
+        } else {
+          result += this.flush() + char;
+        }
+        continue;
+      }
+
+      if (char === this.SQUARE_BRACKET_CLOSE) {
+        this.endToken.push(char);
+        if (this.startToken.length === this.endToken.length) {
+          const cIndex = Number(this.numberToken.join('').trim());
+          if (
+            cIndex > 0 &&
+            cIndex <= citations.length &&
+            contentArray[index + 1] !== this.PARENTHESES_OPEN
+          ) {
+            const content = `[[${cIndex}](${citations[cIndex - 1]})]`;
+            result += content;
+            this.resetToken();
+          } else {
+            result += this.flush();
+          }
+        } else if (this.startToken.length < this.endToken.length) {
+          result += this.flush();
+        }
+        continue;
+      }
+
+      if (this.isNumeric(char)) {
+        if (this.startToken.length > 0) {
+          this.numberToken.push(char);
+        } else {
+          result += this.flush() + char;
+        }
+        continue;
+      }
+
+      if (this.startToken.length > 0) {
+        result += this.flush() + char;
+      } else {
+        result += char;
+      }
+    }
+
+    return result;
+  }
+
+  public flush() {
+    const content = this.getFullContent();
+    this.resetToken();
+    return content;
+  }
+
+  private getFullContent() {
+    return this.startToken.concat(this.numberToken, this.endToken).join('');
+  }
+
+  private resetToken() {
+    this.startToken = [];
+    this.endToken = [];
+    this.numberToken = [];
+  }
+
+  private isNumeric(str: string) {
+    return !isNaN(Number(str)) && str.trim() !== '';
+  }
+}
+
+export class PerplexityProvider implements CopilotTextToTextProvider {
+  static readonly type = CopilotProviderType.Perplexity;
+
+  static readonly capabilities = [CopilotCapability.TextToText];
+
+  static assetsConfig(config: PerplexityConfig) {
+    return !!config.apiKey;
+  }
+
+  constructor(private readonly config: PerplexityConfig) {
+    assert(PerplexityProvider.assetsConfig(config));
+  }
+
+  readonly availableModels = [
+    'llama-3.1-sonar-small-128k-online',
+    'llama-3.1-sonar-large-128k-online',
+    'llama-3.1-sonar-huge-128k-online',
+  ];
+
+  get type(): CopilotProviderType {
+    return PerplexityProvider.type;
+  }
+
+  getCapabilities(): CopilotCapability[] {
+    return PerplexityProvider.capabilities;
+  }
+
+  async isModelAvailable(model: string): Promise<boolean> {
+    return this.availableModels.includes(model);
+  }
+
+  async generateText(
+    messages: PromptMessage[],
+    model: string = 'llama-3.1-sonar-small-128k-online',
+    options: CopilotChatOptions = {}
+  ): Promise<string> {
+    await this.checkParams({ messages, model, options });
+    try {
+      metrics.ai.counter('chat_text_calls').add(1, { model });
+      const sMessages = messages
+        .map(({ content, role }) => ({ content, role }))
+        .filter(({ content }) => typeof content === 'string');
+
+      const params = {
+        method: 'POST',
+        headers: {
+          Authorization: `Bearer ${this.config.apiKey}`,
+          'Content-Type': 'application/json',
+        },
+        body: JSON.stringify({
+          model,
+          messages: sMessages,
+          max_tokens: options.maxTokens || 4096,
+        }),
+      };
+      const response = await fetch(
+        this.config.endpoint || 'https://api.perplexity.ai/chat/completions',
+        params
+      );
+      const data = PerplexitySchema.parse(await response.json());
+      if ('detail' in data) {
+        throw new CopilotProviderSideError({
+          provider: this.type,
+          kind: 'unexpected_response',
+          message: data.detail[0].msg || 'Unexpected perplexity response',
+        });
+      } else {
+        const parser = new CitationParser();
+        const { content } = data.choices[0].message;
+        const { citations } = data;
+        let result = parser.parse(content, citations);
+        result += parser.flush();
+        return result;
+      }
+    } catch (e: any) {
+      metrics.ai.counter('chat_text_errors').add(1, { model });
+      throw this.handleError(e);
+    }
+  }
+
+  async *generateTextStream(
+    messages: PromptMessage[],
+    model: string = 'llama-3.1-sonar-small-128k-online',
+    options: CopilotChatOptions = {}
+  ): AsyncIterable<string> {
+    await this.checkParams({ messages, model, options });
+    try {
+      metrics.ai.counter('chat_text_stream_calls').add(1, { model });
+      const sMessages = messages
+        .map(({ content, role }) => ({ content, role }))
+        .filter(({ content }) => typeof content === 'string');
+
+      const params = {
+        method: 'POST',
+        headers: {
+          Authorization: `Bearer ${this.config.apiKey}`,
+          'Content-Type': 'application/json',
+        },
+        body: JSON.stringify({
+          model,
+          messages: sMessages,
+          max_tokens: options.maxTokens || 4096,
+          stream: true,
+        }),
+      };
+      const response = await fetch(
+        this.config.endpoint || 'https://api.perplexity.ai/chat/completions',
+        params
+      );
+      if (response.body) {
+        const parser = new CitationParser();
+        const provider = this.type;
+        const eventStream = response.body
+          .pipeThrough(new TextDecoderStream())
+          .pipeThrough(new EventSourceParserStream())
+          .pipeThrough(
+            new TransformStream({
+              transform(chunk, controller) {
+                if (options.signal?.aborted) {
+                  controller.enqueue(null);
+                  return;
+                }
+                const json = JSON.parse(chunk.data);
+                if (json) {
+                  const data = PerplexitySchema.parse(json);
+                  if ('detail' in data) {
+                    throw new CopilotProviderSideError({
+                      provider,
+                      kind: 'unexpected_response',
+                      message:
+                        data.detail[0].msg || 'Unexpected perplexity response',
+                    });
+                  }
+                  const { content } = data.choices[0].delta;
+                  const { citations } = data;
+                  const result = parser.parse(content, citations);
+                  controller.enqueue(result);
+                }
+              },
+              flush(controller) {
+                controller.enqueue(parser.flush());
+                controller.enqueue(null);
+              },
+            })
+          );
+
+        const reader = eventStream.getReader();
+        while (true) {
+          const { done, value } = await reader.read();
+          if (done) break;
+          yield value;
+        }
+      } else {
+        const result = await this.generateText(messages, model, options);
+        yield result;
+      }
+    } catch (e) {
+      metrics.ai.counter('chat_text_stream_errors').add(1, { model });
+      throw e;
+    }
+  }
+
+  protected async checkParams({
+    model,
+  }: {
+    messages?: PromptMessage[];
+    embeddings?: string[];
+    model: string;
+    options: CopilotChatOptions;
+  }) {
+    if (!(await this.isModelAvailable(model))) {
+      throw new CopilotPromptInvalid(`Invalid model: ${model}`);
+    }
+  }
+
+  private handleError(e: any) {
+    if (e instanceof CopilotProviderSideError) {
+      return e;
+    }
+    return new CopilotProviderSideError({
+      provider: this.type,
+      kind: 'unexpected_response',
+      message: e?.message || 'Unexpected perplexity response',
+    });
+  }
+}

packages/backend/server/src/plugins/copilot/resolver.ts

@@ -22,6 +22,7 @@ import GraphQLUpload from 'graphql-upload/GraphQLUpload.mjs';
 import {
   CallMetric,
   CopilotFailedToCreateMessage,
+  CopilotSessionNotFound,
   FileUpload,
   RequestMutex,
   Throttle,
@@ -62,6 +63,17 @@ class CreateChatSessionInput {
   promptName!: string;
 }
 
+@InputType()
+class UpdateChatSessionInput {
+  @Field(() => String)
+  sessionId!: string;
+
+  @Field(() => String, {
+    description: 'The prompt name to use for the session',
+  })
+  promptName!: string;
+}
+
 @InputType()
 class ForkChatSessionInput {
   @Field(() => String)
@@ -372,6 +384,38 @@ export class CopilotResolver {
     });
   }
 
+  @Mutation(() => String, {
+    description: 'Update a chat session',
+  })
+  @CallMetric('ai', 'chat_session_update')
+  async updateCopilotSession(
+    @CurrentUser() user: CurrentUser,
+    @Args({ name: 'options', type: () => UpdateChatSessionInput })
+    options: UpdateChatSessionInput
+  ) {
+    const session = await this.chatSession.get(options.sessionId);
+    if (!session) {
+      throw new CopilotSessionNotFound();
+    }
+    const { workspaceId, docId } = session.config;
+    await this.permissions.checkCloudPagePermission(
+      workspaceId,
+      docId,
+      user.id
+    );
+    const lockFlag = `${COPILOT_LOCKER}:session:${user.id}:${workspaceId}`;
+    await using lock = await this.mutex.acquire(lockFlag);
+    if (!lock) {
+      return new TooManyRequest('Server is busy');
+    }
+
+    await this.chatSession.checkQuota(user.id);
+    return await this.chatSession.updateSessionPrompt({
+      ...options,
+      userId: user.id,
+    });
+  }
+
   @Mutation(() => String, {
     description: 'Create a chat session',
   })

packages/backend/server/src/plugins/copilot/session.ts

@@ -10,6 +10,7 @@ import {
   CopilotQuotaExceeded,
   CopilotSessionDeleted,
   CopilotSessionNotFound,
+  PrismaTransaction,
 } from '../../base';
 import { FeatureManagementService } from '../../core/features';
 import { QuotaService } from '../../core/quota';
@@ -22,6 +23,7 @@ import {
   ChatMessageSchema,
   ChatSessionForkOptions,
   ChatSessionOptions,
+  ChatSessionPromptUpdateOptions,
   ChatSessionState,
   getTokenEncoder,
   ListHistoriesOptions,
@@ -198,6 +200,22 @@ export class ChatSessionService {
     private readonly prompt: PromptService
   ) {}
 
+  private async haveSession(
+    sessionId: string,
+    userId: string,
+    tx?: PrismaTransaction
+  ) {
+    const executor = tx ?? this.db;
+    return await executor.aiSession
+      .count({
+        where: {
+          id: sessionId,
+          userId,
+        },
+      })
+      .then(c => c > 0);
+  }
+
   private async setSession(state: ChatSessionState): Promise<string> {
     return await this.db.$transaction(async tx => {
       let sessionId = state.sessionId;
@@ -226,15 +244,7 @@ export class ChatSessionService {
         if (id) sessionId = id;
       }
 
-      const haveSession = await tx.aiSession
+      const haveSession = await this.haveSession(sessionId, state.userId, tx);
-        .count({
-          where: {
-            id: sessionId,
-            userId: state.userId,
-          },
-        })
-        .then(c => c > 0);
-
       if (haveSession) {
         // message will only exists when setSession call by session.save
         if (state.messages.length) {
@@ -570,6 +580,27 @@ export class ChatSessionService {
     });
   }
 
+  async updateSessionPrompt(
+    options: ChatSessionPromptUpdateOptions
+  ): Promise<string> {
+    const prompt = await this.prompt.get(options.promptName);
+    if (!prompt) {
+      this.logger.error(`Prompt not found: ${options.promptName}`);
+      throw new CopilotPromptNotFound({ name: options.promptName });
+    }
+    return await this.db.$transaction(async tx => {
+      let sessionId = options.sessionId;
+      const haveSession = await this.haveSession(sessionId, options.userId, tx);
+      if (haveSession) {
+        await tx.aiSession.update({
+          where: { id: sessionId },
+          data: { promptName: prompt.name },
+        });
+      }
+      return sessionId;
+    });
+  }
+
   async fork(options: ChatSessionForkOptions): Promise<string> {
     const state = await this.getSession(options.sessionId);
     if (!state) {

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

@@ -123,6 +123,11 @@ export interface ChatSessionOptions {
   promptName: string;
 }
 
+export interface ChatSessionPromptUpdateOptions
+  extends Pick<ChatSessionState, 'sessionId' | 'userId'> {
+  promptName: string;
+}
+
 export interface ChatSessionForkOptions
   extends Omit<ChatSessionOptions, 'promptName'> {
   sessionId: string;
@@ -154,6 +159,7 @@ export type ListHistoriesOptions = {
 export enum CopilotProviderType {
   FAL = 'fal',
   OpenAI = 'openai',
+  Perplexity = 'perplexity',
   // only for test
   Test = 'test',
 }

packages/backend/server/src/schema.gql

@@ -551,6 +551,9 @@ type Mutation {
 
   """Update a copilot prompt"""
   updateCopilotPrompt(messages: [CopilotPromptMessageInput!]!, name: String!): CopilotPromptType!
+
+  """Update a chat session"""
+  updateCopilotSession(options: UpdateChatSessionInput!): String!
   updateProfile(input: UpdateUserInput!): UserType!
 
   """update server runtime configurable setting"""
@@ -865,6 +868,12 @@ type UnsupportedSubscriptionPlanDataType {
   plan: String!
 }
 
+input UpdateChatSessionInput {
+  """The prompt name to use for the session"""
+  promptName: String!
+  sessionId: String!
+}
+
 input UpdateUserInput {
   """User name"""
   name: String

packages/backend/server/tests/copilot-provider.spec.ts

@@ -13,6 +13,7 @@ import {
   CopilotProviderService,
   FalProvider,
   OpenAIProvider,
+  PerplexityProvider,
   registerCopilotProvider,
   unregisterCopilotProvider,
 } from '../src/plugins/copilot/providers';
@@ -47,8 +48,10 @@ const test = ava as TestFn<Tester>;
 const isCopilotConfigured =
   !!process.env.COPILOT_OPENAI_API_KEY &&
   !!process.env.COPILOT_FAL_API_KEY &&
+  !!process.env.COPILOT_PERPLEXITY_API_KEY &&
   process.env.COPILOT_OPENAI_API_KEY !== '1' &&
-  process.env.COPILOT_FAL_API_KEY !== '1';
+  process.env.COPILOT_FAL_API_KEY !== '1' &&
+  process.env.COPILOT_PERPLEXITY_API_KEY !== '1';
 const runIfCopilotConfigured = test.macro(
   async (
     t,
@@ -75,6 +78,9 @@ test.serial.before(async t => {
             fal: {
               apiKey: process.env.COPILOT_FAL_API_KEY,
             },
+            perplexity: {
+              apiKey: process.env.COPILOT_PERPLEXITY_API_KEY,
+            },
           },
         },
       }),
@@ -111,6 +117,7 @@ test.serial.before(async t => {
 
   registerCopilotProvider(OpenAIProvider);
   registerCopilotProvider(FalProvider);
+  registerCopilotProvider(PerplexityProvider);
 
   for (const name of await prompt.listNames()) {
     await prompt.delete(name);
@@ -124,6 +131,7 @@ test.serial.before(async t => {
 test.after(async _ => {
   unregisterCopilotProvider(OpenAIProvider.type);
   unregisterCopilotProvider(FalProvider.type);
+  unregisterCopilotProvider(PerplexityProvider.type);
 });
 
 test.after(async t => {
@@ -152,7 +160,6 @@ const checkMDList = (text: string) => {
       return false;
     }
 
-    // eslint-disable-next-line @typescript-eslint/no-non-null-asserted-optional-chain
     const currentIndent = line.match(/^( *)/)?.[0].length!;
     if (Number.isNaN(currentIndent) || currentIndent % 2 !== 0) {
       return false;
@@ -282,6 +289,8 @@ const actions = [
       'Make it longer',
       'Make it shorter',
       'Continue writing',
+      'Chat With AFFiNE AI',
+      'Search With AFFiNE AI',
     ],
     messages: [{ role: 'user' as const, content: TestAssets.SSOT }],
     verifier: (t: ExecutionContext<Tester>, result: string) => {

packages/backend/server/tests/copilot.e2e.ts

@@ -16,6 +16,7 @@ import {
   CopilotProviderService,
   FalProvider,
   OpenAIProvider,
+  PerplexityProvider,
   registerCopilotProvider,
   unregisterCopilotProvider,
 } from '../src/plugins/copilot/providers';
@@ -41,6 +42,7 @@ import {
   sse2array,
   textToEventStream,
   unsplashSearch,
+  updateCopilotSession,
 } from './utils/copilot';
 
 const test = ava as TestFn<{
@@ -63,6 +65,9 @@ test.beforeEach(async t => {
             fal: {
               apiKey: '1',
             },
+            perplexity: {
+              apiKey: '1',
+            },
             unsplashKey: process.env.UNSPLASH_ACCESS_KEY || '1',
           },
         },
@@ -91,6 +96,7 @@ test.beforeEach(async t => {
 
   unregisterCopilotProvider(OpenAIProvider.type);
   unregisterCopilotProvider(FalProvider.type);
+  unregisterCopilotProvider(PerplexityProvider.type);
   registerCopilotProvider(MockCopilotTestProvider);
 
   await prompt.set(promptName, 'test', [
@@ -156,6 +162,85 @@ test('should create session correctly', async t => {
   }
 });
 
+test('should update session correctly', async t => {
+  const { app } = t.context;
+
+  const assertUpdateSession = async (
+    sessionId: string,
+    error: string,
+    asserter = async (x: any) => {
+      t.truthy(await x, error);
+    }
+  ) => {
+    await asserter(updateCopilotSession(app, token, sessionId, promptName));
+  };
+
+  {
+    const { id: workspaceId } = await createWorkspace(app, token);
+    const docId = randomUUID();
+    const sessionId = await createCopilotSession(
+      app,
+      token,
+      workspaceId,
+      docId,
+      promptName
+    );
+    await assertUpdateSession(
+      sessionId,
+      'should be able to update session with cloud workspace that user can access'
+    );
+  }
+
+  {
+    const sessionId = await createCopilotSession(
+      app,
+      token,
+      randomUUID(),
+      randomUUID(),
+      promptName
+    );
+    await assertUpdateSession(
+      sessionId,
+      'should be able to update session with local workspace'
+    );
+  }
+
+  {
+    const aToken = (await signUp(app, 'test', 'test@affine.pro', '123456'))
+      .token.token;
+    const { id: workspaceId } = await createWorkspace(app, aToken);
+    const inviteId = await inviteUser(
+      app,
+      aToken,
+      workspaceId,
+      'darksky@affine.pro'
+    );
+    await acceptInviteById(app, workspaceId, inviteId, false);
+    const sessionId = await createCopilotSession(
+      app,
+      token,
+      workspaceId,
+      randomUUID(),
+      promptName
+    );
+    await assertUpdateSession(
+      sessionId,
+      'should able to update session after user have permission'
+    );
+  }
+
+  {
+    const sessionId = '123456';
+    await assertUpdateSession(sessionId, '', async x => {
+      await t.throwsAsync(
+        x,
+        { instanceOf: Error },
+        'should not able to update invalid session id'
+      );
+    });
+  }
+});
+
 test('should fork session correctly', async t => {
   const { app } = t.context;
 

packages/backend/server/tests/copilot.spec.ts

@@ -16,6 +16,7 @@ import {
   registerCopilotProvider,
   unregisterCopilotProvider,
 } from '../src/plugins/copilot/providers';
+import { CitationParser } from '../src/plugins/copilot/providers/perplexity';
 import { ChatSessionService } from '../src/plugins/copilot/session';
 import {
   CopilotCapability,
@@ -68,7 +69,10 @@ test.beforeEach(async t => {
               apiKey: process.env.COPILOT_OPENAI_API_KEY ?? '1',
             },
             fal: {
-              apiKey: '1',
+              apiKey: process.env.COPILOT_FAL_API_KEY ?? '1',
+            },
+            perplexity: {
+              apiKey: process.env.COPILOT_PERPLEXITY_API_KEY ?? '1',
             },
           },
         },
@@ -274,6 +278,41 @@ test('should be able to manage chat session', async t => {
   }
 });
 
+test('should be able to update chat session prompt', async t => {
+  const { prompt, session } = t.context;
+
+  // Set up a prompt to be used in the session
+  await prompt.set('prompt', 'model', [
+    { role: 'system', content: 'hello {{word}}' },
+  ]);
+
+  // Create a session
+  const sessionId = await session.create({
+    promptName: 'prompt',
+    docId: 'test',
+    workspaceId: 'test',
+    userId,
+  });
+  t.truthy(sessionId, 'should create session');
+
+  // Update the session
+  const updatedSessionId = await session.updateSessionPrompt({
+    sessionId,
+    promptName: 'Search With AFFiNE AI',
+    userId,
+  });
+  t.is(updatedSessionId, sessionId, 'should update session with same id');
+
+  // Verify the session was updated
+  const updatedSession = await session.get(sessionId);
+  t.truthy(updatedSession, 'should retrieve updated session');
+  t.is(
+    updatedSession?.config.promptName,
+    'Search With AFFiNE AI',
+    'should have updated prompt name'
+  );
+});
+
 test('should be able to fork chat session', async t => {
   const { auth, prompt, session } = t.context;
 
@@ -1050,3 +1089,88 @@ test('should be able to run image executor', async t => {
   unregisterCopilotProvider(MockCopilotTestProvider.type);
   registerCopilotProvider(OpenAIProvider);
 });
+
+test('CitationParser should replace citation placeholders with URLs', t => {
+  const content =
+    'This is [a] test sentence with [citations [1]] and [[2]] and [3].';
+  const citations = ['https://example1.com', 'https://example2.com'];
+
+  const parser = new CitationParser();
+  const result = parser.parse(content, citations);
+
+  const expected =
+    'This is [a] test sentence with [citations [[1](https://example1.com)]] and [[2](https://example2.com)] and [3].';
+  t.is(result, expected);
+});
+
+test('CitationParser should replace chunks of citation placeholders with URLs', t => {
+  const contents = [
+    '[[]]',
+    'This is [',
+    'a] test sentence ',
+    'with citations [1',
+    '] and [',
+    '[2]] and [[',
+    '3]] and [[4',
+    ']] and [[5]',
+    '] and [[6]]',
+    ' and [7',
+  ];
+  const citations = [
+    'https://example1.com',
+    'https://example2.com',
+    'https://example3.com',
+    'https://example4.com',
+    'https://example5.com',
+    'https://example6.com',
+    'https://example7.com',
+  ];
+
+  const parser = new CitationParser();
+  let result = contents.reduce((acc, current) => {
+    return acc + parser.parse(current, citations);
+  }, '');
+  result += parser.flush();
+
+  const expected =
+    '[[]]This is [a] test sentence with citations [[1](https://example1.com)] and [[2](https://example2.com)] and [[3](https://example3.com)] and [[4](https://example4.com)] and [[5](https://example5.com)] and [[6](https://example6.com)] and [7';
+  t.is(result, expected);
+});
+
+test('CitationParser should not replace citation already with URLs', t => {
+  const content =
+    'This is [a] test sentence with citations [1](https://example1.com) and [[2]](https://example2.com) and [[3](https://example3.com)].';
+  const citations = [
+    'https://example4.com',
+    'https://example5.com',
+    'https://example6.com',
+  ];
+
+  const parser = new CitationParser();
+  const result = parser.parse(content, citations);
+
+  const expected = content;
+  t.is(result, expected);
+});
+
+test('CitationParser should not replace chunks of citation already with URLs', t => {
+  const contents = [
+    'This is [a] test sentence with citations [1',
+    '](https://example1.com) and [[2]',
+    '](https://example2.com) and [[3](https://example3.com)].',
+  ];
+  const citations = [
+    'https://example4.com',
+    'https://example5.com',
+    'https://example6.com',
+  ];
+
+  const parser = new CitationParser();
+  let result = contents.reduce((acc, current) => {
+    return acc + parser.parse(current, citations);
+  }, '');
+  result += parser.flush();
+
+  const expected = contents.join('');
+  t.is(result, expected);
+});

packages/backend/server/tests/utils/copilot.ts

@@ -184,6 +184,31 @@ export async function createCopilotSession(
   return res.body.data.createCopilotSession;
 }
 
+export async function updateCopilotSession(
+  app: INestApplication,
+  userToken: string,
+  sessionId: string,
+  promptName: string
+): Promise<string> {
+  const res = await request(app.getHttpServer())
+    .post(gql)
+    .auth(userToken, { type: 'bearer' })
+    .set({ 'x-request-id': 'test', 'x-operation-name': 'test' })
+    .send({
+      query: `
+        mutation updateCopilotSession($options: UpdateChatSessionInput!) {
+          updateCopilotSession(options: $options)
+        }
+      `,
+      variables: { options: { sessionId, promptName } },
+    })
+    .expect(200);
+
+  handleGraphQLError(res);
+
+  return res.body.data.updateCopilotSession;
+}
+
 export async function forkCopilotSession(
   app: INestApplication,
   userToken: string,

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,
+    },
   };
 }

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,6 +37,7 @@ 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';
 
+function getAIPageRootWatcher(framework: FrameworkProvider) {
   class AIPageRootWatcher extends BlockServiceWatcher {
     static override readonly flavour = 'affine:page';
 
@@ -44,7 +46,7 @@ class AIPageRootWatcher extends BlockServiceWatcher {
       this.blockService.specSlots.widgetConnected.on(view => {
         if (view.component instanceof AffineAIPanelWidget) {
           view.component.style.width = '630px';
-        view.component.config = buildAIPanelConfig(view.component);
+          view.component.config = buildAIPanelConfig(view.component, framework);
           setupSpaceAIEntry(view.component);
         }
 
@@ -58,10 +60,15 @@ class AIPageRootWatcher extends BlockServiceWatcher {
       });
     }
   }
+  return AIPageRootWatcher;
+}
 
-export const AIPageRootBlockSpec: ExtensionType[] = [
+export function createAIPageRootBlockSpec(
+  framework: FrameworkProvider
+): ExtensionType[] {
+  return [
     ...PageRootBlockSpec,
-  AIPageRootWatcher,
+    getAIPageRootWatcher(framework),
     {
       setup: di => {
         di.override(WidgetViewMapIdentifier('affine:page'), () => {
@@ -75,7 +82,9 @@ export const AIPageRootBlockSpec: ExtensionType[] = [
       },
     },
   ];
+}
 
+function getAIEdgelessRootWatcher(framework: FrameworkProvider) {
   class AIEdgelessRootWatcher extends BlockServiceWatcher {
     static override readonly flavour = 'affine:page';
 
@@ -84,7 +93,7 @@ class AIEdgelessRootWatcher extends BlockServiceWatcher {
       this.blockService.specSlots.widgetConnected.on(view => {
         if (view.component instanceof AffineAIPanelWidget) {
           view.component.style.width = '430px';
-        view.component.config = buildAIPanelConfig(view.component);
+          view.component.config = buildAIPanelConfig(view.component, framework);
           setupSpaceAIEntry(view.component);
         }
 
@@ -106,10 +115,15 @@ class AIEdgelessRootWatcher extends BlockServiceWatcher {
       });
     }
   }
+  return AIEdgelessRootWatcher;
+}
 
-export const AIEdgelessRootBlockSpec: ExtensionType[] = [
+export function createAIEdgelessRootBlockSpec(
+  framework: FrameworkProvider
+): ExtensionType[] {
+  return [
     ...EdgelessRootBlockSpec,
-  AIEdgelessRootWatcher,
+    getAIEdgelessRootWatcher(framework),
     {
       setup: di => {
         di.override(WidgetViewMapIdentifier('affine:page'), () => {
@@ -126,6 +140,7 @@ export const AIEdgelessRootBlockSpec: ExtensionType[] = [
       },
     },
   ];
+}
 
 class AIParagraphBlockWatcher extends BlockServiceWatcher {
   static override readonly flavour = 'affine:paragraph';

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();
+              this._onTextareaSend(evt);
-              await this.send();
             }
           }}
           @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 () => {
+                aria-disabled=${uploadDisabled}
-                  const images = await openFileOrFiles({
+                @click=${uploadDisabled ? undefined : this._uploadImageFiles}
-                    acceptType: 'Images',
-                    multiple: true,
-                  });
-                  if (!images) return;
-                  this._addImages(images);
-                }}
               >
-                ${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',

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}

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,

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',

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

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
+) {
+  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)[]
   ) {
-  async function getChatSessionId(workspaceId: string, docId: string) {
     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(
+      chatSessions.set(storeKey, {
-        storeKey,
+        getSessionId: createChatSession({
-        createChatSession({
           client,
           workspaceId,
           docId,
-        })
+          promptName,
-      );
+        }),
+        promptName,
+      });
     }
     try {
-      const sessionId = await chatSessions.get(storeKey);
+      /* oxlint-disable @typescript-eslint/no-non-null-assertion */
-      assertExists(sessionId);
+      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,

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;

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;
 }

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();

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]}
           />
         ))}

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} />;
 });

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>

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,
+  ]);
+}

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);
+  };
+}

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',

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

packages/frontend/graphql/src/graphql/index.ts

@@ -206,6 +206,17 @@ mutation createCopilotSession($options: CreateChatSessionInput!) {
 }`,
 };
 
+export const updateCopilotSessionMutation = {
+  id: 'updateCopilotSessionMutation' as const,
+  operationName: 'updateCopilotSession',
+  definitionName: 'updateCopilotSession',
+  containsFile: false,
+  query: `
+mutation updateCopilotSession($options: UpdateChatSessionInput!) {
+  updateCopilotSession(options: $options)
+}`,
+};
+
 export const createCustomerPortalMutation = {
   id: 'createCustomerPortalMutation' as const,
   operationName: 'createCustomerPortal',

packages/frontend/graphql/src/graphql/update-copilot-session.gql

@@ -0,0 +1,3 @@
+mutation updateCopilotSession($options: UpdateChatSessionInput!) {
+  updateCopilotSession(options: $options)
+}

packages/frontend/graphql/src/schema.ts

@@ -188,6 +188,11 @@ export interface CreateChatSessionInput {
   workspaceId: Scalars['String']['input'];
 }
 
+export interface UpdateChatSessionInput {
+  sessionId: Scalars['String']['input'];
+  promptName: Scalars['String']['input'];
+}
+
 export interface CreateCheckoutSessionInput {
   args?: InputMaybe<Scalars['JSONObject']['input']>;
   coupon?: InputMaybe<Scalars['String']['input']>;
@@ -688,6 +693,10 @@ export interface MutationCreateCopilotSessionArgs {
   options: CreateChatSessionInput;
 }
 
+export interface MutationUpdateCopilotSessionArgs {
+  options: UpdateChatSessionInput;
+}
+
 export interface MutationCreateInviteLinkArgs {
   expireTime: WorkspaceInviteLinkExpireTime;
   workspaceId: Scalars['String']['input'];
@@ -1596,6 +1605,15 @@ export type CreateCopilotSessionMutation = {
   createCopilotSession: string;
 };
 
+export type UpdateCopilotSessionMutationVariables = Exact<{
+  options: UpdateChatSessionInput;
+}>;
+
+export type UpdateCopilotSessionMutation = {
+  __typename?: 'Mutation';
+  updateCopilotSession: string;
+};
+
 export type CreateCustomerPortalMutationVariables = Exact<{
   [key: string]: never;
 }>;
@@ -3070,6 +3088,11 @@ export type Mutations =
       variables: CreateCopilotSessionMutationVariables;
       response: CreateCopilotSessionMutation;
     }
+  | {
+    name: 'updateCopilotSessionMutation';
+    variables: UpdateCopilotSessionMutationVariables;
+    response: UpdateCopilotSessionMutation;
+  }
   | {
       name: 'createCustomerPortalMutation';
       variables: CreateCustomerPortalMutationVariables;

packages/frontend/i18n/src/resources/en.json

@@ -1281,6 +1281,8 @@
   "com.affine.settings.workspace.experimental-features.prompt-warning-title": "WARNING MESSAGE",
   "com.affine.settings.workspace.experimental-features.enable-ai.name": "Enable AI",
   "com.affine.settings.workspace.experimental-features.enable-ai.description": "Enable or disable ALL AI features.",
+  "com.affine.settings.workspace.experimental-features.enable-ai-network-search.name": "Enable AI Network Search",
+  "com.affine.settings.workspace.experimental-features.enable-ai-network-search.description": "Enable or disable AI Network Search feature.",
   "com.affine.settings.workspace.experimental-features.enable-database-full-width.name": "Database Full Width",
   "com.affine.settings.workspace.experimental-features.enable-database-full-width.description": "The database will be displayed in full-width mode.",
   "com.affine.settings.workspace.experimental-features.enable-database-attachment-note.name": "Database Attachment Note",

tests/affine-cloud-copilot/e2e/copilot.spec.ts

@@ -10,7 +10,10 @@ import {
   getBlockSuiteEditorTitle,
   waitForEditorLoad,
 } from '@affine-test/kit/utils/page-logic';
-import { clickSideBarAllPageButton } from '@affine-test/kit/utils/sidebar';
+import {
+  clickSideBarAllPageButton,
+  clickSideBarUseAvatar,
+} from '@affine-test/kit/utils/sidebar';
 import { createLocalWorkspace } from '@affine-test/kit/utils/workspace';
 import { expect, type Page } from '@playwright/test';
 
@@ -48,14 +51,15 @@ function getUser() {
   };
 }
 
-test.skip(
+const isCopilotConfigured =
-  () =>
+  !!process.env.COPILOT_OPENAI_API_KEY &&
-    !process.env.COPILOT_OPENAI_API_KEY ||
+  !!process.env.COPILOT_FAL_API_KEY &&
-    !process.env.COPILOT_FAL_API_KEY ||
+  !!process.env.COPILOT_PERPLEXITY_API_KEY &&
-    process.env.COPILOT_OPENAI_API_KEY === '1' ||
+  process.env.COPILOT_OPENAI_API_KEY !== '1' &&
-    process.env.COPILOT_FAL_API_KEY === '1',
+  process.env.COPILOT_FAL_API_KEY !== '1' &&
-  'skip test if no copilot api key'
+  process.env.COPILOT_PERPLEXITY_API_KEY !== '1';
-);
+
+test.skip(() => !isCopilotConfigured, 'skip test if no copilot api key');
 
 test('can open chat side panel', async ({ page }) => {
   await openHomePage(page);
@@ -68,13 +72,17 @@ test('can open chat side panel', async ({ page }) => {
   await expect(page.getByTestId('sidebar-tab-content-chat')).toBeVisible();
 });
 
-const makeChat = async (page: Page, content: string) => {
+const openChat = async (page: Page) => {
   if (await page.getByTestId('sidebar-tab-chat').isHidden()) {
     await page.getByTestId('right-sidebar-toggle').click({
       delay: 200,
     });
   }
   await page.getByTestId('sidebar-tab-chat').click();
+};
+
+const makeChat = async (page: Page, content: string) => {
+  await openChat(page);
   await page.getByTestId('chat-panel-input').focus();
   await page.keyboard.type(content);
   await page.keyboard.press('Enter');
@@ -83,6 +91,7 @@ const makeChat = async (page: Page, content: string) => {
 const clearChat = async (page: Page) => {
   await page.getByTestId('chat-panel-clear').click();
   await page.getByTestId('confirm-modal-confirm').click();
+  await page.waitForTimeout(500);
 };
 
 const collectChat = async (page: Page) => {
@@ -343,6 +352,48 @@ test.describe('chat panel', () => {
       expect(editorContent).toBe('');
     }
   });
+
+  test('can open and close network search', async ({ page }) => {
+    await page.reload();
+    await clickSideBarAllPageButton(page);
+    await page.waitForTimeout(200);
+    await createLocalWorkspace({ name: 'test' }, page);
+    await clickNewPageButton(page);
+    await clickSideBarUseAvatar(page);
+    await page.getByTestId('workspace-modal-account-settings-option').click();
+    await page.getByTestId('experimental-features-trigger').click();
+    await page
+      .getByTestId('experimental-prompt')
+      .getByTestId('affine-checkbox')
+      .click();
+    await page.getByTestId('experimental-confirm-button').click();
+    await page.getByTestId('enable_ai_network_search').click();
+    await page.getByTestId('modal-close-button').click();
+    await openChat(page);
+    await page.getByTestId('chat-network-search').click();
+    await makeChat(page, 'hello');
+    let history = await collectChat(page);
+    expect(history[0]).toEqual({
+      name: 'You',
+      content: 'hello',
+    });
+    expect(history[1].name).toBe('AFFiNE AI');
+    expect(
+      await page.locator('chat-panel affine-link').count()
+    ).toBeGreaterThan(0);
+
+    await clearChat(page);
+    expect((await collectChat(page)).length).toBe(0);
+    await page.getByTestId('chat-network-search').click();
+    await makeChat(page, 'hello');
+    history = await collectChat(page);
+    expect(history[0]).toEqual({
+      name: 'You',
+      content: 'hello',
+    });
+    expect(history[1].name).toBe('AFFiNE AI');
+    expect(await page.locator('chat-panel affine-link').count()).toBe(0);
+  });
 });
 
 test.describe('chat with block', () => {

yarn.lock

@@ -808,6 +808,7 @@ __metadata:
     cookie-parser: "npm:^1.4.7"
     cross-env: "npm:^7.0.3"
     dotenv: "npm:^16.4.7"
+    eventsource-parser: "npm:^3.0.0"
     express: "npm:^4.21.2"
     fast-xml-parser: "npm:^4.5.0"
     get-stream: "npm:^9.0.1"
@@ -21165,6 +21166,13 @@ __metadata:
   languageName: node
   linkType: hard
 
+"eventsource-parser@npm:^3.0.0":
+  version: 3.0.0
+  resolution: "eventsource-parser@npm:3.0.0"
+  checksum: 10/8215adf5d8404105ecd0658030b0407e06987ceb9aadcea28a38d69bacf02e5d0fc8bba5fa7c3954552c89509c8ef5e1fa3895e000c061411c055b4bbc26f4b0
+  languageName: node
+  linkType: hard
+
 "execa@npm:^1.0.0":
   version: 1.0.0
   resolution: "execa@npm:1.0.0"