feat(core): support network search in chat block center peek (#10186)

[BS-2582](https://linear.app/affine-design/issue/BS-2582/chat-block-center-peek-支持-network-search)
2026-02-12 12:28:42 +00:00 · 2025-02-14 12:43:31 +00:00
parent b6f8027e1b
commit e6b570e613
6 changed files with 276 additions and 148 deletions
--- 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
@@ -27,6 +27,7 @@ import { readBlobAsURL } from '../utils/image';
 import type { AINetworkSearchConfig } from './chat-config';
 import type { ChatContextValue, ChatMessage, DocContext } from './chat-context';
 import { isDocChip } from './components/utils';
+import { PROMPT_NAME_AFFINE_AI, PROMPT_NAME_NETWORK_SEARCH } from './const';

 const MaximumImageCount = 32;

@@ -289,11 +290,11 @@ export class ChatPanelInput extends SignalWatcher(WithDisposable(LitElement)) {

  private get _promptName() {
    if (this._isNetworkDisabled) {
-      return 'Chat With AFFiNE AI';
+      return PROMPT_NAME_AFFINE_AI;
    }
    return this._isNetworkActive
-      ? 'Search With AFFiNE AI'
-      : 'Chat With AFFiNE AI';
+      ? PROMPT_NAME_NETWORK_SEARCH
+      : PROMPT_NAME_AFFINE_AI;
  }

  private async _updatePromptName() {
--- a/packages/frontend/core/src/blocksuite/presets/ai/chat-panel/const.ts
+++ b/packages/frontend/core/src/blocksuite/presets/ai/chat-panel/const.ts
@@ -8,3 +8,6 @@ export const HISTORY_IMAGE_ACTIONS = [
  'Remove background',
  'Convert to sticker',
 ];
+
+export const PROMPT_NAME_AFFINE_AI = 'Chat With AFFiNE AI';
+export const PROMPT_NAME_NETWORK_SEARCH = 'Search With AFFiNE AI';
--- a/packages/frontend/core/src/blocksuite/presets/ai/peek-view/chat-block-input.ts
+++ b/packages/frontend/core/src/blocksuite/presets/ai/peek-view/chat-block-input.ts
@@ -1,5 +1,11 @@
 import type { EditorHost } from '@blocksuite/affine/block-std';
-import { type AIError, openFileOrFiles } from '@blocksuite/affine/blocks';
+import {
+  type AIError,
+  openFileOrFiles,
+  unsafeCSSVarV2,
+} from '@blocksuite/affine/blocks';
+import { SignalWatcher } from '@blocksuite/affine/global/utils';
+import { ImageIcon, PublishIcon } from '@blocksuite/icons/lit';
 import { css, html, LitElement, nothing } from 'lit';
 import { property, query, state } from 'lit/decorators.js';
 import { classMap } from 'lit/directives/class-map.js';
@@ -11,16 +17,21 @@ import {
  ChatClearIcon,
  ChatSendIcon,
  CloseIcon,
-  ImageIcon,
 } from '../_common/icons';
+import type { AINetworkSearchConfig } from '../chat-panel/chat-config';
+import {
+  PROMPT_NAME_AFFINE_AI,
+  PROMPT_NAME_NETWORK_SEARCH,
+} from '../chat-panel/const';
 import { AIProvider } from '../provider';
 import { reportResponse } from '../utils/action-reporter';
 import { readBlobAsURL } from '../utils/image';
+import { stopPropagation } from '../utils/selection-utils';
 import type { ChatContext } from './types';

 const MaximumImageCount = 8;

-export class ChatBlockInput extends LitElement {
+export class ChatBlockInput extends SignalWatcher(LitElement) {
  static override styles = css`
    :host {
      width: 100%;
@@ -126,6 +137,7 @@ export class ChatBlockInput extends LitElement {
      display: flex;
      gap: 8px;
      align-items: center;
+
      div {
        width: 24px;
        height: 24px;
@@ -134,6 +146,28 @@ export class ChatBlockInput extends LitElement {
      div:nth-child(2) {
        margin-left: auto;
      }
+
+      .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')};
+      }
+
+      .chat-network-search[aria-disabled='true'] {
+        cursor: not-allowed;
+      }
+      .chat-network-search[aria-disabled='true'] svg {
+        color: var(--affine-text-disable-color) !important;
+      }
    }

    .chat-history-clear.disabled {
@@ -168,79 +202,44 @@ export class ChatBlockInput extends LitElement {
        <textarea
          rows="1"
          placeholder="What are your thoughts?"
-          @keydown=${async (evt: KeyboardEvent) => {
-            if (evt.key === 'Enter' && !evt.shiftKey && !evt.isComposing) {
-              evt.preventDefault();
-              await this._send();
-            }
-          }}
-          @input=${() => {
-            const { textarea } = this;
-            this._isInputEmpty = !textarea.value.trim();
-            textarea.style.height = 'auto';
-            textarea.style.height = textarea.scrollHeight + 'px';
-            if (this.scrollHeight >= 202) {
-              textarea.style.height = '168px';
-              textarea.style.overflowY = 'scroll';
-            }
-          }}
+          @keydown=${this._handleKeyDown}
+          @input=${this._handleInput}
          @focus=${() => {
            this._focused = true;
          }}
          @blur=${() => {
            this._focused = false;
          }}
-          @paste=${(event: ClipboardEvent) => {
-            const items = event.clipboardData?.items;
-            if (!items) return;
-            for (const index in items) {
-              const item = items[index];
-              if (item.kind === 'file' && item.type.indexOf('image') >= 0) {
-                const blob = item.getAsFile();
-                if (!blob) continue;
-                this._addImages([blob]);
-              }
-            }
-          }}
+          @paste=${this._handlePaste}
          data-testid="chat-block-input"
        ></textarea>
        <div class="chat-panel-input-actions">
-          <div
-            class=${cleanButtonClasses}
-            @click=${async () => {
-              if (disableCleanUp) {
-                return;
-              }
-              await this.cleanupHistories();
-            }}
-          >
+          <div class=${cleanButtonClasses} @click=${this._handleCleanup}>
            ${ChatClearIcon}
          </div>
+          ${this.networkSearchConfig.visible.value
+            ? html`
+                <div
+                  class="chat-network-search"
+                  data-testid="chat-network-search"
+                  aria-disabled=${this._isNetworkDisabled}
+                  data-active=${this._isNetworkActive}
+                  @click=${this._isNetworkDisabled
+                    ? 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);
-                }}
-              >
-                ${ImageIcon}
+            ? html`<div class="image-upload" @click=${this._handleImageUpload}>
+                ${ImageIcon()}
              </div>`
            : nothing}
          ${status === 'transmitting'
-            ? html`<div
-                @click=${() => {
-                  this.chatContext.abortController?.abort();
-                  this.updateContext({ status: 'success' });
-                  reportResponse('aborted:stop');
-                }}
-              >
-                ${ChatAbortIcon}
-              </div>`
+            ? html`<div @click=${this._handleAbort}>${ChatAbortIcon}</div>`
            : html`<div
                @click="${this._send}"
                class="chat-panel-send"
@@ -261,6 +260,9 @@ export class ChatBlockInput extends LitElement {
  @property({ attribute: false })
  accessor host!: EditorHost;

+  @property({ attribute: false })
+  accessor networkSearchConfig!: AINetworkSearchConfig;
+
  @property({ attribute: false })
  accessor updateChatBlock!: () => Promise<void>;

@@ -291,6 +293,41 @@ export class ChatBlockInput extends LitElement {
  @state()
  accessor _curIndex = -1;

+  private _lastPromptName: string | null = null;
+
+  private get _isNetworkActive() {
+    return (
+      !!this.networkSearchConfig.visible.value &&
+      !!this.networkSearchConfig.enabled.value
+    );
+  }
+
+  private get _isNetworkDisabled() {
+    return !!this.chatContext.images.length;
+  }
+
+  private get _promptName() {
+    if (this._isNetworkDisabled) {
+      return PROMPT_NAME_AFFINE_AI;
+    }
+    return this._isNetworkActive
+      ? PROMPT_NAME_NETWORK_SEARCH
+      : PROMPT_NAME_AFFINE_AI;
+  }
+
+  private async _updatePromptName() {
+    if (this._lastPromptName !== this._promptName) {
+      this._lastPromptName = this._promptName;
+      const { currentSessionId } = this.chatContext;
+      if (currentSessionId) {
+        await AIProvider.session?.updateSession(
+          currentSessionId,
+          this._promptName
+        );
+      }
+    }
+  }
+
  private readonly _addImages = (images: File[]) => {
    const oldImages = this.chatContext.images;
    this.updateContext({
@@ -298,6 +335,71 @@ export class ChatBlockInput extends LitElement {
    });
  };

+  private readonly _toggleNetworkSearch = (e: MouseEvent) => {
+    e.preventDefault();
+    e.stopPropagation();
+
+    const enable = this.networkSearchConfig.enabled.value;
+    this.networkSearchConfig.setEnabled(!enable);
+  };
+
+  private readonly _handleKeyDown = async (evt: KeyboardEvent) => {
+    if (evt.key === 'Enter' && !evt.shiftKey && !evt.isComposing) {
+      evt.preventDefault();
+      await this._send();
+    }
+  };
+
+  private readonly _handleInput = () => {
+    const { textarea } = this;
+    this._isInputEmpty = !textarea.value.trim();
+    textarea.style.height = 'auto';
+    textarea.style.height = textarea.scrollHeight + 'px';
+    if (this.scrollHeight >= 202) {
+      textarea.style.height = '168px';
+      textarea.style.overflowY = 'scroll';
+    }
+  };
+
+  private readonly _handlePaste = (event: ClipboardEvent) => {
+    const items = event.clipboardData?.items;
+    if (!items) return;
+    for (const index in items) {
+      const item = items[index];
+      if (item.kind === 'file' && item.type.indexOf('image') >= 0) {
+        const blob = item.getAsFile();
+        if (!blob) continue;
+        this._addImages([blob]);
+      }
+    }
+  };
+
+  private readonly _handleCleanup = async () => {
+    if (
+      this.chatContext.status === 'loading' ||
+      this.chatContext.status === 'transmitting' ||
+      !this.chatContext.messages.length
+    ) {
+      return;
+    }
+    await this.cleanupHistories();
+  };
+
+  private readonly _handleImageUpload = async () => {
+    const images = await openFileOrFiles({
+      acceptType: 'Images',
+      multiple: true,
+    });
+    if (!images) return;
+    this._addImages(images);
+  };
+
+  private readonly _handleAbort = () => {
+    this.chatContext.abortController?.abort();
+    this.updateContext({ status: 'success' });
+    reportResponse('aborted:stop');
+  };
+
  private _renderImages(images: File[]) {
    return html`
      <div
@@ -416,6 +518,8 @@ export class ChatBlockInput extends LitElement {
        chatSessionId = forkSessionId;
      }

+      await this._updatePromptName();
+
      const abortController = new AbortController();
      const stream = AIProvider.actions.chat?.({
        input: text,
--- a/packages/frontend/core/src/blocksuite/presets/ai/peek-view/chat-block-peek-view.ts
+++ b/packages/frontend/core/src/blocksuite/presets/ai/peek-view/chat-block-peek-view.ts
@@ -27,6 +27,7 @@ import {
  queryHistoryMessages,
 } from '../_common/chat-actions-handle';
 import { SmallHintIcon } from '../_common/icons';
+import type { AINetworkSearchConfig } from '../chat-panel/chat-config';
 import { AIChatErrorRenderer } from '../messages/error';
 import { AIProvider } from '../provider';
 import { PeekViewStyles } from './styles';
@@ -60,6 +61,8 @@ export class AIChatBlockPeekView extends LitElement {
    return this.parentModel.rootWorkspaceId;
  }

+  private _textRendererOptions: TextRendererOptions = {};
+
  private readonly _deserializeHistoryChatMessages = (
    historyMessagesString: string
  ) => {
@@ -360,74 +363,76 @@ export class AIChatBlockPeekView extends LitElement {
    const { host } = this;
    const actions = ChatBlockPeekViewActions;

-    return html`${repeat(currentMessages, (message, idx) => {
-      const { status, error } = this.chatContext;
-      const isAssistantMessage = message.role === 'assistant';
-      const isLastReply =
-        idx === currentMessages.length - 1 && isAssistantMessage;
-      const messageState =
-        isLastReply && (status === 'transmitting' || status === 'loading')
-          ? 'generating'
-          : 'finished';
-      const shouldRenderError = isLastReply && status === 'error' && !!error;
-      const isNotReady = status === 'transmitting' || status === 'loading';
-      const shouldRenderCopyMore =
-        isAssistantMessage && !(isLastReply && isNotReady);
-      const shouldRenderActions =
-        isLastReply && !!message.content && !isNotReady;
+    return html`${repeat(
+      currentMessages,
+      message => message.id || message.createdAt,
+      (message, idx) => {
+        const { status, error } = this.chatContext;
+        const isAssistantMessage = message.role === 'assistant';
+        const isLastReply =
+          idx === currentMessages.length - 1 && isAssistantMessage;
+        const messageState =
+          isLastReply && (status === 'transmitting' || status === 'loading')
+            ? 'generating'
+            : 'finished';
+        const shouldRenderError = isLastReply && status === 'error' && !!error;
+        const isNotReady = status === 'transmitting' || status === 'loading';
+        const shouldRenderCopyMore =
+          isAssistantMessage && !(isLastReply && isNotReady);
+        const shouldRenderActions =
+          isLastReply && !!message.content && !isNotReady;

-      const messageClasses = classMap({
-        'assistant-message-container': isAssistantMessage,
-      });
+        const messageClasses = classMap({
+          'assistant-message-container': isAssistantMessage,
+        });

-      const { attachments, role, content } = message;
-      const userInfo = {
-        userId: message.userId,
-        userName: message.userName,
-        avatarUrl: message.avatarUrl,
-      };
-      const textRendererOptions: TextRendererOptions = {
-        extensions: this.previewSpecBuilder.value,
-      };
+        const { attachments, role, content, userId, userName, avatarUrl } =
+          message;

-      return html`<div class=${messageClasses}>
-        <ai-chat-message
-          .host=${host}
-          .state=${messageState}
-          .content=${content}
-          .attachments=${attachments}
-          .messageRole=${role}
-          .userInfo=${userInfo}
-          .textRendererOptions=${textRendererOptions}
-        ></ai-chat-message>
-        ${shouldRenderError ? AIChatErrorRenderer(host, error) : nothing}
-        ${shouldRenderCopyMore
-          ? html` <chat-copy-more
-              .host=${host}
-              .actions=${actions}
-              .content=${message.content}
-              .isLast=${isLastReply}
-              .chatSessionId=${this.chatContext.currentSessionId ?? undefined}
-              .messageId=${message.id ?? undefined}
-              .retry=${() => this.retry()}
-            ></chat-copy-more>`
-          : nothing}
-        ${shouldRenderActions
-          ? html`<chat-action-list
-              .host=${host}
-              .actions=${actions}
-              .content=${message.content}
-              .chatSessionId=${this.chatContext.currentSessionId ?? undefined}
-              .messageId=${message.id ?? undefined}
-              .layoutDirection=${'horizontal'}
-            ></chat-action-list>`
-          : nothing}
-      </div>`;
-    })}`;
+        return html`<div class=${messageClasses}>
+          <ai-chat-message
+            .host=${host}
+            .state=${messageState}
+            .content=${content}
+            .attachments=${attachments}
+            .messageRole=${role}
+            .userId=${userId}
+            .userName=${userName}
+            .avatarUrl=${avatarUrl}
+            .textRendererOptions=${this._textRendererOptions}
+          ></ai-chat-message>
+          ${shouldRenderError ? AIChatErrorRenderer(host, error) : nothing}
+          ${shouldRenderCopyMore
+            ? html` <chat-copy-more
+                .host=${host}
+                .actions=${actions}
+                .content=${message.content}
+                .isLast=${isLastReply}
+                .chatSessionId=${this.chatContext.currentSessionId ?? undefined}
+                .messageId=${message.id ?? undefined}
+                .retry=${() => this.retry()}
+              ></chat-copy-more>`
+            : nothing}
+          ${shouldRenderActions
+            ? html`<chat-action-list
+                .host=${host}
+                .actions=${actions}
+                .content=${message.content}
+                .chatSessionId=${this.chatContext.currentSessionId ?? undefined}
+                .messageId=${message.id ?? undefined}
+                .layoutDirection=${'horizontal'}
+              ></chat-action-list>`
+            : nothing}
+        </div>`;
+      }
+    )}`;
  };

  override connectedCallback() {
    super.connectedCallback();
+    this._textRendererOptions = {
+      extensions: this.previewSpecBuilder.value,
+    };
    this._historyMessages = this._deserializeHistoryChatMessages(
      this.historyMessagesString
    );
@@ -476,19 +481,18 @@ export class AIChatBlockPeekView extends LitElement {
      cleanCurrentChatHistories,
      chatContext,
      updateContext,
+      networkSearchConfig,
+      _textRendererOptions,
    } = this;

    const { messages: currentChatMessages } = chatContext;
-    const textRendererOptions: TextRendererOptions = {
-      extensions: this.previewSpecBuilder.value,
-    };

    return html`<div class="ai-chat-block-peek-view-container">
      <div class="ai-chat-messages-container">
        <ai-chat-messages
          .host=${host}
          .messages=${_historyMessages}
-          .textRendererOptions=${textRendererOptions}
+          .textRendererOptions=${_textRendererOptions}
        ></ai-chat-messages>
        <date-time .date=${latestMessageCreatedAt}></date-time>
        <div class="new-chat-messages-container">
@@ -504,6 +508,7 @@ export class AIChatBlockPeekView extends LitElement {
        .cleanupHistories=${cleanCurrentChatHistories}
        .chatContext=${chatContext}
        .updateContext=${updateContext}
+        .networkSearchConfig=${networkSearchConfig}
      ></chat-block-input>
      <div class="peek-view-footer">
        ${SmallHintIcon}
@@ -524,6 +529,9 @@ export class AIChatBlockPeekView extends LitElement {
  @property({ attribute: false })
  accessor previewSpecBuilder!: SpecBuilder;

+  @property({ attribute: false })
+  accessor networkSearchConfig!: AINetworkSearchConfig;
+
  @state()
  accessor _historyMessages: ChatMessage[] = [];

@@ -548,11 +556,13 @@ declare global {
 export const AIChatBlockPeekViewTemplate = (
  parentModel: AIChatBlockModel,
  host: EditorHost,
-  previewSpecBuilder: SpecBuilder
+  previewSpecBuilder: SpecBuilder,
+  networkSearchConfig: AINetworkSearchConfig
 ) => {
  return html`<ai-chat-block-peek-view
    .parentModel=${parentModel}
    .host=${host}
    .previewSpecBuilder=${previewSpecBuilder}
+    .networkSearchConfig=${networkSearchConfig}
  ></ai-chat-block-peek-view>`;
 };
--- a/packages/frontend/core/src/blocksuite/presets/blocks/ai-chat-block/components/ai-chat-messages.ts
+++ b/packages/frontend/core/src/blocksuite/presets/blocks/ai-chat-block/components/ai-chat-messages.ts
@@ -5,11 +5,7 @@ import { property } from 'lit/decorators.js';
 import { classMap } from 'lit/directives/class-map.js';
 import { repeat } from 'lit/directives/repeat.js';

-import type {
-  ChatMessage,
-  MessageRole,
-  MessageUserInfo,
-} from '../../../../blocks';
+import type { ChatMessage, MessageRole } from '../../../../blocks';
 import type { TextRendererOptions } from '../../../_common/components/text-renderer';
 import { UserInfoTemplate } from './user-info';

@@ -43,7 +39,9 @@ export class AIChatMessage extends LitElement {
      content,
      attachments,
      messageRole,
-      userInfo,
+      userId,
+      userName,
+      avatarUrl,
    } = this;
    const withAttachments = !!attachments && attachments.length > 0;

@@ -53,7 +51,7 @@ export class AIChatMessage extends LitElement {

    return html`
      <div class="ai-chat-message">
-        ${UserInfoTemplate(userInfo, messageRole)}
+        ${UserInfoTemplate({ userId, userName, avatarUrl }, messageRole)}
        <div class="ai-chat-content">
          <chat-images .attachments=${attachments}></chat-images>
          <div class=${messageClasses}>
@@ -88,7 +86,13 @@ export class AIChatMessage extends LitElement {
  accessor textRendererOptions: TextRendererOptions = {};

  @property({ attribute: false })
-  accessor userInfo: MessageUserInfo = {};
+  accessor userId: string | undefined = undefined;
+
+  @property({ attribute: false })
+  accessor userName: string | undefined = undefined;
+
+  @property({ attribute: false })
+  accessor avatarUrl: string | undefined = undefined;
 }

 export class AIChatMessages extends LitElement {
@@ -112,14 +116,10 @@ export class AIChatMessages extends LitElement {
    return html`<div class="ai-chat-messages">
      ${repeat(
        this.messages,
-        message => message.id,
+        message => message.id || message.createdAt,
        message => {
-          const { attachments, role, content } = message;
-          const userInfo = {
-            userId: message.userId,
-            userName: message.userName,
-            avatarUrl: message.avatarUrl,
-          };
+          const { attachments, role, content, userId, userName, avatarUrl } =
+            message;
          return html`
            <ai-chat-message
              .host=${this.host}
@@ -127,7 +127,9 @@ export class AIChatMessages extends LitElement {
              .content=${content}
              .attachments=${attachments}
              .messageRole=${role}
-              .userInfo=${userInfo}
+              .userId=${userId}
+              .userName=${userName}
+              .avatarUrl=${avatarUrl}
            ></ai-chat-message>
          `;
        }
--- a/packages/frontend/core/src/modules/peek-view/view/ai-chat-block-peek-view/index.tsx
+++ b/packages/frontend/core/src/modules/peek-view/view/ai-chat-block-peek-view/index.tsx
@@ -1,5 +1,6 @@
 import { toReactNode } from '@affine/component';
 import { AIChatBlockPeekViewTemplate } from '@affine/core/blocksuite/presets/ai';
+import { AINetworkSearchService } from '@affine/core/modules/ai-button/services/network-search';
 import type { EditorHost } from '@blocksuite/affine/block-std';
 import { useFramework } from '@toeverything/infra';
 import { useMemo } from 'react';
@@ -17,13 +18,20 @@ export const AIChatBlockPeekView = ({
  host,
 }: AIChatBlockPeekViewProps) => {
  const framework = useFramework();
+  const searchService = framework.get(AINetworkSearchService);
  return useMemo(() => {
    const previewSpecBuilder = createPageModePreviewSpecs(framework);
+    const networkSearchConfig = {
+      visible: searchService.visible,
+      enabled: searchService.enabled,
+      setEnabled: searchService.setEnabled,
+    };
    const template = AIChatBlockPeekViewTemplate(
      model,
      host,
-      previewSpecBuilder
+      previewSpecBuilder,
+      networkSearchConfig
    );
    return toReactNode(template);
-  }, [framework, model, host]);
+  }, [framework, model, host, searchService]);
 };