From afc4158f479fcfe8414fb5f604db444a39304cca Mon Sep 17 00:00:00 2001
From: yoyoyohamapi <8338436+yoyoyohamapi@users.noreply.github.com>
Date: Wed, 19 Mar 2025 10:28:56 +0000
Subject: [PATCH] refactor(core): separate rendering logic for user and
 assistant messages (#10909)

### TL;DR

Separate rendering logic for user and assistant messages.

> CLOSE AF-2323

### What Changed
- Rendering user message with `<chat-panel-message-user />` component.
- Rendering assistant message with `<chat-panel-message-assistant />` Component
---
 .../chat-panel-assistant-message.ts           | 198 ++++++++++++++
 .../ai/chat-panel/chat-panel-messages.ts      | 256 +++---------------
 .../ai/chat-panel/chat-panel-user-message.ts  |  82 ++++++
 .../core/src/blocksuite/ai/effects.ts         |   7 +
 4 files changed, 330 insertions(+), 213 deletions(-)
 create mode 100644 packages/frontend/core/src/blocksuite/ai/chat-panel/chat-panel-assistant-message.ts
 create mode 100644 packages/frontend/core/src/blocksuite/ai/chat-panel/chat-panel-user-message.ts

diff --git a/packages/frontend/core/src/blocksuite/ai/chat-panel/chat-panel-assistant-message.ts b/packages/frontend/core/src/blocksuite/ai/chat-panel/chat-panel-assistant-message.ts
new file mode 100644
index 0000000000..bd1b14ed84
--- /dev/null
+++ b/packages/frontend/core/src/blocksuite/ai/chat-panel/chat-panel-assistant-message.ts
@@ -0,0 +1,198 @@
+import type { EditorHost } from '@blocksuite/affine/block-std';
+import { ShadowlessElement } from '@blocksuite/affine/block-std';
+import { WithDisposable } from '@blocksuite/affine/global/lit';
+import { isInsidePageEditor } from '@blocksuite/affine/shared/utils';
+import { AiIcon } from '@blocksuite/icons/lit';
+import { css, html, nothing } from 'lit';
+import { property } from 'lit/decorators.js';
+
+import {
+  EdgelessEditorActions,
+  PageEditorActions,
+} from '../_common/chat-actions-handle';
+import { type AIError } from '../components/ai-item/types';
+import { AIChatErrorRenderer } from '../messages/error';
+import { isChatMessage } from './chat-context';
+import { type ChatItem } from './chat-context';
+import { HISTORY_IMAGE_ACTIONS } from './const';
+
+const AffineAvatarIcon = AiIcon({
+  width: '20px',
+  height: '20px',
+  style: 'color: var(--affine-primary-color)',
+});
+
+export class ChatPanelAssistantMessage extends WithDisposable(
+  ShadowlessElement
+) {
+  static override styles = css`
+    .message-info {
+      color: var(--affine-placeholder-color);
+      font-size: var(--affine-font-xs);
+      font-weight: 400;
+    }
+  `;
+
+  @property({ attribute: false })
+  accessor host!: EditorHost;
+
+  @property({ attribute: false })
+  accessor item!: ChatItem;
+
+  @property({ attribute: false })
+  accessor isLast: boolean = false;
+
+  @property({ attribute: false })
+  accessor status: string = 'idle';
+
+  @property({ attribute: false })
+  accessor error: AIError | null = null;
+
+  @property({ attribute: false })
+  accessor previewSpecBuilder: any;
+
+  @property({ attribute: false })
+  accessor getSessionId!: () => Promise<string | undefined>;
+
+  @property({ attribute: false })
+  accessor retry!: () => void;
+
+  renderAvatar() {
+    const isAssistant =
+      isChatMessage(this.item) && this.item.role === 'assistant';
+    const isWithDocs =
+      isAssistant &&
+      'content' in this.item &&
+      this.item.content &&
+      this.item.content.includes('[^') &&
+      /\[\^\d+\]:{"type":"doc","docId":"[^"]+"}/.test(this.item.content);
+
+    return html`<div class="user-info">
+      ${AffineAvatarIcon} AFFiNE AI
+      ${isWithDocs
+        ? html`<span class="message-info">with your docs</span>`
+        : nothing}
+    </div>`;
+  }
+
+  renderContent() {
+    const { host, item, isLast, status, error } = this;
+
+    if (isLast && status === 'loading') {
+      return html`<ai-loading></ai-loading>`;
+    }
+
+    if (isChatMessage(item)) {
+      const state = isLast
+        ? status !== 'loading' && status !== 'transmitting'
+          ? 'finished'
+          : 'generating'
+        : 'finished';
+      const shouldRenderError = isLast && status === 'error' && !!error;
+      return html`<chat-text
+          .host=${host}
+          .attachments=${item.attachments}
+          .text=${item.content}
+          .state=${state}
+          .previewSpecBuilder=${this.previewSpecBuilder}
+        ></chat-text>
+        ${shouldRenderError ? AIChatErrorRenderer(host, error) : nothing}
+        ${this.renderEditorActions()}`;
+    } else {
+      switch (item.action) {
+        case 'Create a presentation':
+          return html`<action-slides
+            .host=${host}
+            .item=${item}
+          ></action-slides>`;
+        case 'Make it real':
+          return html`<action-make-real
+            .host=${host}
+            .item=${item}
+          ></action-make-real>`;
+        case 'Brainstorm mindmap':
+          return html`<action-mindmap
+            .host=${host}
+            .item=${item}
+          ></action-mindmap>`;
+        case 'Explain this image':
+        case 'Generate a caption':
+          return html`<action-image-to-text
+            .host=${host}
+            .item=${item}
+          ></action-image-to-text>`;
+        default:
+          if (HISTORY_IMAGE_ACTIONS.includes(item.action)) {
+            return html`<action-image
+              .host=${host}
+              .item=${item}
+            ></action-image>`;
+          }
+
+          return html`<action-text
+            .item=${item}
+            .host=${host}
+            .isCode=${item.action === 'Explain this code' ||
+            item.action === 'Check code error'}
+          ></action-text>`;
+      }
+    }
+  }
+
+  renderEditorActions() {
+    const { item, isLast, status } = this;
+
+    if (!isChatMessage(item) || item.role !== 'assistant') return nothing;
+
+    if (
+      isLast &&
+      status !== 'success' &&
+      status !== 'idle' &&
+      status !== 'error'
+    )
+      return nothing;
+
+    const { host } = this;
+    const { content, id: messageId } = item;
+
+    const actions = isInsidePageEditor(host)
+      ? PageEditorActions
+      : EdgelessEditorActions;
+
+    return html`
+      <chat-copy-more
+        .host=${host}
+        .actions=${actions}
+        .content=${content}
+        .isLast=${isLast}
+        .getSessionId=${this.getSessionId}
+        .messageId=${messageId}
+        .withMargin=${true}
+        .retry=${() => this.retry()}
+      ></chat-copy-more>
+      ${isLast && !!content
+        ? html`<chat-action-list
+            .actions=${actions}
+            .host=${host}
+            .content=${content}
+            .getSessionId=${this.getSessionId}
+            .messageId=${messageId ?? undefined}
+            .withMargin=${true}
+          ></chat-action-list>`
+        : nothing}
+    `;
+  }
+
+  protected override render() {
+    return html`
+      ${this.renderAvatar()}
+      <div class="item-wrapper">${this.renderContent()}</div>
+    `;
+  }
+}
+
+declare global {
+  interface HTMLElementTagNameMap {
+    'chat-panel-assistant-message': ChatPanelAssistantMessage;
+  }
+}
diff --git a/packages/frontend/core/src/blocksuite/ai/chat-panel/chat-panel-messages.ts b/packages/frontend/core/src/blocksuite/ai/chat-panel/chat-panel-messages.ts
index 6b98435a33..6ccd8f1d26 100644
--- a/packages/frontend/core/src/blocksuite/ai/chat-panel/chat-panel-messages.ts
+++ b/packages/frontend/core/src/blocksuite/ai/chat-panel/chat-panel-messages.ts
@@ -5,47 +5,25 @@ import {
   DocModeProvider,
   FeatureFlagService,
 } from '@blocksuite/affine/shared/services';
-import {
-  isInsidePageEditor,
-  type SpecBuilder,
-} from '@blocksuite/affine/shared/utils';
+import { type SpecBuilder } from '@blocksuite/affine/shared/utils';
 import type { BaseSelection } from '@blocksuite/affine/store';
-import {
-  AiIcon,
-  ArrowDownBigIcon as ArrowDownIcon,
-} from '@blocksuite/icons/lit';
+import { ArrowDownBigIcon as ArrowDownIcon } from '@blocksuite/icons/lit';
 import { css, html, nothing, type PropertyValues } from 'lit';
 import { property, query, state } from 'lit/decorators.js';
 import { repeat } from 'lit/directives/repeat.js';
 import { debounce } from 'lodash-es';
 
-import {
-  EdgelessEditorActions,
-  PageEditorActions,
-} from '../_common/chat-actions-handle';
 import { AffineIcon } from '../_common/icons';
-import {
-  type AIError,
-  PaymentRequiredError,
-  UnauthorizedError,
-} from '../components/ai-item/types';
-import { AIChatErrorRenderer } from '../messages/error';
+import { type AIError, UnauthorizedError } from '../components/ai-item/types';
 import { AIProvider } from '../provider';
 import {
   type ChatContextValue,
-  type ChatItem,
   type ChatMessage,
   isChatMessage,
 } from './chat-context';
 import { HISTORY_IMAGE_ACTIONS } from './const';
 import { AIPreloadConfig } from './preload-config';
 
-const AffineAvatarIcon = AiIcon({
-  width: '20px',
-  height: '20px',
-  style: 'color: var(--affine-primary-color)',
-});
-
 export class ChatPanelMessages extends WithDisposable(ShadowlessElement) {
   static override styles = css`
     chat-panel-messages {
@@ -61,6 +39,26 @@ export class ChatPanelMessages extends WithDisposable(ShadowlessElement) {
       overflow-y: auto;
     }
 
+    chat-panel-assistant-message,
+    chat-panel-user-message {
+      display: contents;
+    }
+
+    .user-info {
+      display: flex;
+      align-items: center;
+      gap: 10px;
+      margin-bottom: 4px;
+      color: var(--affine-text-primary-color);
+      font-size: var(--affine-font-sm);
+      font-weight: 500;
+      user-select: none;
+    }
+
+    .item-wrapper {
+      margin-left: 32px;
+    }
+
     .messages-placeholder {
       width: 100%;
       position: absolute;
@@ -116,44 +114,8 @@ export class ChatPanelMessages extends WithDisposable(ShadowlessElement) {
       white-space: nowrap;
     }
 
-    .item-wrapper {
-      margin-left: 32px;
-    }
-
-    .user-info {
-      display: flex;
-      align-items: center;
-      gap: 10px;
-      margin-bottom: 4px;
-      color: var(--affine-text-primary-color);
-      font-size: var(--affine-font-sm);
-      font-weight: 500;
-      user-select: none;
-    }
-
-    .message-info {
-      color: var(--affine-placeholder-color);
-      font-size: var(--affine-font-xs);
-      font-weight: 400;
-    }
-
-    .avatar-container {
-      width: 24px;
-      height: 24px;
-    }
-
-    .avatar {
-      width: 100%;
-      height: 100%;
-      border-radius: 50%;
-      background-color: var(--affine-primary-color);
-    }
-
-    .avatar-container img {
-      width: 100%;
-      height: 100%;
-      border-radius: 50%;
-      object-fit: cover;
+    .message {
+      display: contents;
     }
 
     .down-indicator {
@@ -247,7 +209,7 @@ export class ChatPanelMessages extends WithDisposable(ShadowlessElement) {
   };
 
   protected override render() {
-    const { items } = this.chatContextValue;
+    const { items, status, error } = this.chatContextValue;
     const { isLoading } = this;
     const filteredItems = items.filter(item => {
       return (
@@ -288,10 +250,23 @@ export class ChatPanelMessages extends WithDisposable(ShadowlessElement) {
               (item, index) => {
                 const isLast = index === filteredItems.length - 1;
                 return html`<div class="message">
-                  ${this.renderAvatar(item)}
-                  <div class="item-wrapper">
-                    ${this.renderItem(item, isLast)}
-                  </div>
+                  ${isChatMessage(item) && item.role === 'user'
+                    ? html`<chat-panel-user-message
+                        .host=${this.host}
+                        .item=${item}
+                        .avatarUrl=${this.avatarUrl}
+                        .previewSpecBuilder=${this.previewSpecBuilder}
+                      ></chat-panel-user-message>`
+                    : html`<chat-panel-assistant-message
+                        .host=${this.host}
+                        .item=${item}
+                        .isLast=${isLast}
+                        .status=${status}
+                        .error=${error}
+                        .previewSpecBuilder=${this.previewSpecBuilder}
+                        .getSessionId=${this.getSessionId}
+                        .retry=${() => this.retry()}
+                      ></chat-panel-assistant-message>`}
                 </div>`;
               }
             )}
@@ -347,108 +322,6 @@ export class ChatPanelMessages extends WithDisposable(ShadowlessElement) {
     }
   }
 
-  renderItem(item: ChatItem, isLast: boolean) {
-    const { status, error } = this.chatContextValue;
-    const { host } = this;
-
-    if (isLast && status === 'loading') {
-      return this.renderLoading();
-    }
-
-    if (
-      isLast &&
-      status === 'error' &&
-      (error instanceof PaymentRequiredError ||
-        error instanceof UnauthorizedError)
-    ) {
-      return AIChatErrorRenderer(host, error);
-    }
-
-    if (isChatMessage(item)) {
-      const state = isLast
-        ? status !== 'loading' && status !== 'transmitting'
-          ? 'finished'
-          : 'generating'
-        : 'finished';
-      const shouldRenderError = isLast && status === 'error' && !!error;
-      return html`<chat-text
-          .host=${host}
-          .attachments=${item.attachments}
-          .text=${item.content}
-          .state=${state}
-          .previewSpecBuilder=${this.previewSpecBuilder}
-        ></chat-text>
-        ${shouldRenderError ? AIChatErrorRenderer(host, error) : nothing}
-        ${this.renderEditorActions(item, isLast)}`;
-    } else {
-      switch (item.action) {
-        case 'Create a presentation':
-          return html`<action-slides
-            .host=${host}
-            .item=${item}
-          ></action-slides>`;
-        case 'Make it real':
-          return html`<action-make-real
-            .host=${host}
-            .item=${item}
-          ></action-make-real>`;
-        case 'Brainstorm mindmap':
-          return html`<action-mindmap
-            .host=${host}
-            .item=${item}
-          ></action-mindmap>`;
-        case 'Explain this image':
-        case 'Generate a caption':
-          return html`<action-image-to-text
-            .host=${host}
-            .item=${item}
-          ></action-image-to-text>`;
-        default:
-          if (HISTORY_IMAGE_ACTIONS.includes(item.action)) {
-            return html`<action-image
-              .host=${host}
-              .item=${item}
-            ></action-image>`;
-          }
-
-          return html`<action-text
-            .item=${item}
-            .host=${host}
-            .isCode=${item.action === 'Explain this code' ||
-            item.action === 'Check code error'}
-          ></action-text>`;
-      }
-    }
-  }
-
-  renderAvatar(item: ChatItem) {
-    const isUser = isChatMessage(item) && item.role === 'user';
-    const isAssistant = isChatMessage(item) && item.role === 'assistant';
-    const isWithDocs =
-      isAssistant &&
-      item.content &&
-      item.content.includes('[^') &&
-      /\[\^\d+\]:{"type":"doc","docId":"[^"]+"}/.test(item.content);
-
-    return html`<div class="user-info">
-      ${isUser
-        ? html`<div class="avatar-container">
-            ${this.avatarUrl
-              ? html`<img .src=${this.avatarUrl} />`
-              : html`<div class="avatar"></div>`}
-          </div>`
-        : AffineAvatarIcon}
-      ${isUser ? 'You' : 'AFFiNE AI'}
-      ${isWithDocs
-        ? html`<span class="message-info">with your docs</span>`
-        : nothing}
-    </div>`;
-  }
-
-  renderLoading() {
-    return html` <ai-loading></ai-loading>`;
-  }
-
   scrollToEnd() {
     requestAnimationFrame(() => {
       if (!this.messagesContainer) return;
@@ -504,49 +377,6 @@ export class ChatPanelMessages extends WithDisposable(ShadowlessElement) {
       this.updateContext({ abortController: null });
     }
   };
-
-  renderEditorActions(item: ChatMessage, isLast: boolean) {
-    const { status } = this.chatContextValue;
-
-    if (item.role !== 'assistant') return nothing;
-
-    if (
-      isLast &&
-      status !== 'success' &&
-      status !== 'idle' &&
-      status !== 'error'
-    )
-      return nothing;
-
-    const { host } = this;
-    const { content, id: messageId } = item;
-    const actions = isInsidePageEditor(host)
-      ? PageEditorActions
-      : EdgelessEditorActions;
-
-    return html`
-      <chat-copy-more
-        .host=${host}
-        .actions=${actions}
-        .content=${content}
-        .isLast=${isLast}
-        .getSessionId=${this.getSessionId}
-        .messageId=${messageId}
-        .withMargin=${true}
-        .retry=${() => this.retry()}
-      ></chat-copy-more>
-      ${isLast && !!content
-        ? html`<chat-action-list
-            .actions=${actions}
-            .host=${host}
-            .content=${content}
-            .getSessionId=${this.getSessionId}
-            .messageId=${messageId ?? undefined}
-            .withMargin=${true}
-          ></chat-action-list>`
-        : nothing}
-    `;
-  }
 }
 
 declare global {
diff --git a/packages/frontend/core/src/blocksuite/ai/chat-panel/chat-panel-user-message.ts b/packages/frontend/core/src/blocksuite/ai/chat-panel/chat-panel-user-message.ts
new file mode 100644
index 0000000000..784a93d6c4
--- /dev/null
+++ b/packages/frontend/core/src/blocksuite/ai/chat-panel/chat-panel-user-message.ts
@@ -0,0 +1,82 @@
+import type { EditorHost } from '@blocksuite/affine/block-std';
+import { ShadowlessElement } from '@blocksuite/affine/block-std';
+import { WithDisposable } from '@blocksuite/affine/global/lit';
+import { css, html, nothing } from 'lit';
+import { property } from 'lit/decorators.js';
+
+import { type ChatItem, isChatMessage } from './chat-context';
+
+export class ChatPanelUserMessage extends WithDisposable(ShadowlessElement) {
+  static override styles = css`
+    .avatar-container {
+      width: 24px;
+      height: 24px;
+    }
+
+    .avatar {
+      width: 100%;
+      height: 100%;
+      border-radius: 50%;
+      background-color: var(--affine-primary-color);
+    }
+
+    .avatar-container img {
+      width: 100%;
+      height: 100%;
+      border-radius: 50%;
+      object-fit: cover;
+    }
+  `;
+
+  @property({ attribute: false })
+  accessor host!: EditorHost;
+
+  @property({ attribute: false })
+  accessor item!: ChatItem;
+
+  @property({ attribute: false })
+  accessor avatarUrl: string = '';
+
+  @property({ attribute: false })
+  accessor previewSpecBuilder: any;
+
+  renderAvatar() {
+    return html`<div class="user-info">
+      <div class="avatar-container">
+        ${this.avatarUrl
+          ? html`<img .src=${this.avatarUrl} />`
+          : html`<div class="avatar"></div>`}
+      </div>
+      You
+    </div>`;
+  }
+
+  renderContent() {
+    const { host, item } = this;
+
+    if (isChatMessage(item)) {
+      return html`<chat-text
+        .host=${host}
+        .attachments=${item.attachments}
+        .text=${item.content}
+        .state=${'finished'}
+        .previewSpecBuilder=${this.previewSpecBuilder}
+      ></chat-text>`;
+    }
+
+    return nothing;
+  }
+
+  protected override render() {
+    return html`
+      ${this.renderAvatar()}
+      <div class="item-wrapper">${this.renderContent()}</div>
+    `;
+  }
+}
+
+declare global {
+  interface HTMLElementTagNameMap {
+    'chat-panel-user-message': ChatPanelUserMessage;
+  }
+}
diff --git a/packages/frontend/core/src/blocksuite/ai/effects.ts b/packages/frontend/core/src/blocksuite/ai/effects.ts
index ab93fee852..f154a7ffda 100644
--- a/packages/frontend/core/src/blocksuite/ai/effects.ts
+++ b/packages/frontend/core/src/blocksuite/ai/effects.ts
@@ -23,9 +23,11 @@ import { ActionMindmap } from './chat-panel/actions/mindmap';
 import { ActionSlides } from './chat-panel/actions/slides';
 import { ActionText } from './chat-panel/actions/text';
 import { AILoading } from './chat-panel/ai-loading';
+import { ChatPanelAssistantMessage } from './chat-panel/chat-panel-assistant-message';
 import { ChatPanelChips } from './chat-panel/chat-panel-chips';
 import { ChatPanelInput } from './chat-panel/chat-panel-input';
 import { ChatPanelMessages } from './chat-panel/chat-panel-messages';
+import { ChatPanelUserMessage } from './chat-panel/chat-panel-user-message';
 import { ChatPanelAddPopover } from './chat-panel/components/add-popover';
 import { ChatPanelChip } from './chat-panel/components/chip';
 import { ChatPanelCollectionChip } from './chat-panel/components/collection-chip';
@@ -89,6 +91,11 @@ export function registerAIEffects() {
   customElements.define('action-slides', ActionSlides);
   customElements.define('action-text', ActionText);
   customElements.define('ai-loading', AILoading);
+  customElements.define(
+    'chat-panel-assistant-message',
+    ChatPanelAssistantMessage
+  );
+  customElements.define('chat-panel-user-message', ChatPanelUserMessage);
   customElements.define('chat-panel-input', ChatPanelInput);
   customElements.define('chat-panel-messages', ChatPanelMessages);
   customElements.define('chat-panel', ChatPanel);