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)

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'
+      ? PROMPT_NAME_NETWORK_SEARCH
-      : 'Chat With AFFiNE AI';
+      : PROMPT_NAME_AFFINE_AI;
   }
 
   private async _updatePromptName() {

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

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) => {
+          @keydown=${this._handleKeyDown}
-            if (evt.key === 'Enter' && !evt.shiftKey && !evt.isComposing) {
+          @input=${this._handleInput}
-              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';
-            }
-          }}
           @focus=${() => {
             this._focused = true;
           }}
           @blur=${() => {
             this._focused = false;
           }}
-          @paste=${(event: ClipboardEvent) => {
+          @paste=${this._handlePaste}
-            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]);
-              }
-            }
-          }}
           data-testid="chat-block-input"
         ></textarea>
         <div class="chat-panel-input-actions">
-          <div
+          <div class=${cleanButtonClasses} @click=${this._handleCleanup}>
-            class=${cleanButtonClasses}
-            @click=${async () => {
-              if (disableCleanUp) {
-                return;
-              }
-              await this.cleanupHistories();
-            }}
-          >
             ${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
+            ? html`<div class="image-upload" @click=${this._handleImageUpload}>
-                class="image-upload"
+                ${ImageIcon()}
-                @click=${async () => {
-                  const images = await openFileOrFiles({
-                    acceptType: 'Images',
-                    multiple: true,
-                  });
-                  if (!images) return;
-                  this._addImages(images);
-                }}
-              >
-                ${ImageIcon}
               </div>`
             : nothing}
           ${status === 'transmitting'
-            ? html`<div
+            ? html`<div @click=${this._handleAbort}>${ChatAbortIcon}</div>`
-                @click=${() => {
-                  this.chatContext.abortController?.abort();
-                  this.updateContext({ status: 'success' });
-                  reportResponse('aborted:stop');
-                }}
-              >
-                ${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,

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) => {
+    return html`${repeat(
-      const { status, error } = this.chatContext;
+      currentMessages,
-      const isAssistantMessage = message.role === 'assistant';
+      message => message.id || message.createdAt,
-      const isLastReply =
+      (message, idx) => {
-        idx === currentMessages.length - 1 && isAssistantMessage;
+        const { status, error } = this.chatContext;
-      const messageState =
+        const isAssistantMessage = message.role === 'assistant';
-        isLastReply && (status === 'transmitting' || status === 'loading')
+        const isLastReply =
-          ? 'generating'
+          idx === currentMessages.length - 1 && isAssistantMessage;
-          : 'finished';
+        const messageState =
-      const shouldRenderError = isLastReply && status === 'error' && !!error;
+          isLastReply && (status === 'transmitting' || status === 'loading')
-      const isNotReady = status === 'transmitting' || status === 'loading';
+            ? 'generating'
-      const shouldRenderCopyMore =
+            : 'finished';
-        isAssistantMessage && !(isLastReply && isNotReady);
+        const shouldRenderError = isLastReply && status === 'error' && !!error;
-      const shouldRenderActions =
+        const isNotReady = status === 'transmitting' || status === 'loading';
-        isLastReply && !!message.content && !isNotReady;
+        const shouldRenderCopyMore =
+          isAssistantMessage && !(isLastReply && isNotReady);
+        const shouldRenderActions =
+          isLastReply && !!message.content && !isNotReady;
 
-      const messageClasses = classMap({
+        const messageClasses = classMap({
-        'assistant-message-container': isAssistantMessage,
+          'assistant-message-container': isAssistantMessage,
-      });
+        });
 
-      const { attachments, role, content } = message;
+        const { attachments, role, content, userId, userName, avatarUrl } =
-      const userInfo = {
+          message;
-        userId: message.userId,
-        userName: message.userName,
-        avatarUrl: message.avatarUrl,
-      };
-      const textRendererOptions: TextRendererOptions = {
-        extensions: this.previewSpecBuilder.value,
-      };
 
-      return html`<div class=${messageClasses}>
+        return html`<div class=${messageClasses}>
-        <ai-chat-message
+          <ai-chat-message
-          .host=${host}
+            .host=${host}
-          .state=${messageState}
+            .state=${messageState}
-          .content=${content}
+            .content=${content}
-          .attachments=${attachments}
+            .attachments=${attachments}
-          .messageRole=${role}
+            .messageRole=${role}
-          .userInfo=${userInfo}
+            .userId=${userId}
-          .textRendererOptions=${textRendererOptions}
+            .userName=${userName}
-        ></ai-chat-message>
+            .avatarUrl=${avatarUrl}
-        ${shouldRenderError ? AIChatErrorRenderer(host, error) : nothing}
+            .textRendererOptions=${this._textRendererOptions}
-        ${shouldRenderCopyMore
+          ></ai-chat-message>
-          ? html` <chat-copy-more
+          ${shouldRenderError ? AIChatErrorRenderer(host, error) : nothing}
-              .host=${host}
+          ${shouldRenderCopyMore
-              .actions=${actions}
+            ? html` <chat-copy-more
-              .content=${message.content}
+                .host=${host}
-              .isLast=${isLastReply}
+                .actions=${actions}
-              .chatSessionId=${this.chatContext.currentSessionId ?? undefined}
+                .content=${message.content}
-              .messageId=${message.id ?? undefined}
+                .isLast=${isLastReply}
-              .retry=${() => this.retry()}
+                .chatSessionId=${this.chatContext.currentSessionId ?? undefined}
-            ></chat-copy-more>`
+                .messageId=${message.id ?? undefined}
-          : nothing}
+                .retry=${() => this.retry()}
-        ${shouldRenderActions
+              ></chat-copy-more>`
-          ? html`<chat-action-list
+            : nothing}
-              .host=${host}
+          ${shouldRenderActions
-              .actions=${actions}
+            ? html`<chat-action-list
-              .content=${message.content}
+                .host=${host}
-              .chatSessionId=${this.chatContext.currentSessionId ?? undefined}
+                .actions=${actions}
-              .messageId=${message.id ?? undefined}
+                .content=${message.content}
-              .layoutDirection=${'horizontal'}
+                .chatSessionId=${this.chatContext.currentSessionId ?? undefined}
-            ></chat-action-list>`
+                .messageId=${message.id ?? undefined}
-          : nothing}
+                .layoutDirection=${'horizontal'}
-      </div>`;
+              ></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>`;
 };

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 {
+import type { ChatMessage, MessageRole } from '../../../../blocks';
-  ChatMessage,
-  MessageRole,
-  MessageUserInfo,
-} 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 { attachments, role, content, userId, userName, avatarUrl } =
-          const userInfo = {
+            message;
-            userId: message.userId,
-            userName: message.userName,
-            avatarUrl: message.avatarUrl,
-          };
           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>
           `;
         }

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