refactor(core): display selected doc & attachment chip (#13443)

<img width="1275" height="997" alt="截屏2025-08-08 15 13 59" src="https://github.com/user-attachments/assets/b429239d-84dc-490d-ad1e-957652e3caba" />  ## Summary by CodeRabbit * **New Features** * Introduced support for attachment chips in AI chat, allowing individual attachments to be displayed, added, and removed as separate chips. * Added a new visual component for displaying attachment chips in the chat interface. * **Improvements** * Enhanced chat composer to handle attachments and document chips separately, improving clarity and control over shared content. * Expanded criteria for triggering chat actions to include both document and attachment selections. * **Refactor** * Updated context management to process attachments individually rather than in batches, streamlining the addition and removal of context items.
2026-02-04 08:38:34 +00:00 · 2025-08-08 15:34:04 +08:00
parent 009288dee2
commit 5fd7dfc8aa
9 changed files with 188 additions and 121 deletions
--- a/packages/frontend/core/src/blocksuite/ai/actions/types.ts
+++ b/packages/frontend/core/src/blocksuite/ai/actions/types.ts
@@ -361,12 +361,12 @@ declare global {
        op: string,
        updates: string
      ) => Promise<string>;
-      addContextBlobs: (options: {
-        blobIds: string[];
+      addContextBlob: (options: {
+        blobId: string;
        contextId: string;
-      }) => Promise<CopilotContextBlob[]>;
-      removeContextBlobs: (options: {
-        blobIds: string[];
+      }) => Promise<CopilotContextBlob>;
+      removeContextBlob: (options: {
+        blobId: string;
        contextId: string;
      }) => Promise<boolean>;
    }
--- a/packages/frontend/core/src/blocksuite/ai/components/ai-chat-chips/attachment-chip.ts
+++ b/packages/frontend/core/src/blocksuite/ai/components/ai-chat-chips/attachment-chip.ts
@@ -0,0 +1,40 @@
+import { getAttachmentFileIcon } from '@blocksuite/affine/components/icons';
+import { SignalWatcher, WithDisposable } from '@blocksuite/affine/global/lit';
+import { ShadowlessElement } from '@blocksuite/affine/std';
+import { html } from 'lit';
+import { property } from 'lit/decorators.js';
+
+import type { AttachmentChip } from './type';
+import { getChipIcon, getChipTooltip } from './utils';
+
+export class ChatPanelAttachmentChip extends SignalWatcher(
+  WithDisposable(ShadowlessElement)
+) {
+  @property({ attribute: false })
+  accessor chip!: AttachmentChip;
+
+  @property({ attribute: false })
+  accessor removeChip!: (chip: AttachmentChip) => void;
+
+  override render() {
+    const { state, name } = this.chip;
+    const isLoading = state === 'processing';
+    const tooltip = getChipTooltip(state, name, this.chip.tooltip);
+    const fileType = name.split('.').pop() ?? '';
+    const fileIcon = getAttachmentFileIcon(fileType);
+    const icon = getChipIcon(state, fileIcon);
+
+    return html`<chat-panel-chip
+      .state=${state}
+      .name=${name}
+      .tooltip=${tooltip}
+      .icon=${icon}
+      .closeable=${!isLoading}
+      .onChipDelete=${this.onChipDelete}
+    ></chat-panel-chip>`;
+  }
+
+  private readonly onChipDelete = () => {
+    this.removeChip(this.chip);
+  };
+}
--- a/packages/frontend/core/src/blocksuite/ai/components/ai-chat-chips/chat-panel-chips.ts
+++ b/packages/frontend/core/src/blocksuite/ai/components/ai-chat-chips/chat-panel-chips.ts
@@ -15,6 +15,7 @@ import type { ChatChip, DocChip, DocDisplayConfig, FileChip } from './type';
 import {
  estimateTokenCount,
  getChipKey,
+  isAttachmentChip,
  isCollectionChip,
  isDocChip,
  isFileChip,
@@ -160,6 +161,12 @@ export class ChatPanelChips extends SignalWatcher(
              .removeChip=${this.removeChip}
            ></chat-panel-file-chip>`;
          }
+          if (isAttachmentChip(chip)) {
+            return html`<chat-panel-attachment-chip
+              .chip=${chip}
+              .removeChip=${this.removeChip}
+            ></chat-panel-attachment-chip>`;
+          }
          if (isTagChip(chip)) {
            const tag = this._tags.value.find(tag => tag.id === chip.tagId);
            if (!tag) {
--- a/packages/frontend/core/src/blocksuite/ai/components/ai-chat-chips/type.ts
+++ b/packages/frontend/core/src/blocksuite/ai/components/ai-chat-chips/type.ts
@@ -36,10 +36,13 @@ export interface CollectionChip extends BaseChip {
  collectionId: string;
 }

+export interface AttachmentChip extends BaseChip {
+  sourceId: string;
+  name: string;
+}
+
 export interface SelectedContextChip extends BaseChip {
  uuid: string;
-  attachments: { sourceId: string; name: string }[];
-  docs: string[];
  snapshot: string | null;
  combinedElementsMarkdown: string | null;
  html: string | null;
@@ -50,6 +53,7 @@ export type ChatChip =
  | FileChip
  | TagChip
  | CollectionChip
+  | AttachmentChip
  | SelectedContextChip;

 export interface DocDisplayConfig {
--- a/packages/frontend/core/src/blocksuite/ai/components/ai-chat-chips/utils.ts
+++ b/packages/frontend/core/src/blocksuite/ai/components/ai-chat-chips/utils.ts
@@ -3,6 +3,7 @@ import { WarningIcon } from '@blocksuite/icons/lit';
 import { type TemplateResult } from 'lit';

 import type {
+  AttachmentChip,
  ChatChip,
  ChipState,
  CollectionChip,
@@ -66,11 +67,11 @@ export function isCollectionChip(chip: ChatChip): chip is CollectionChip {
 export function isSelectedContextChip(
  chip: ChatChip
 ): chip is SelectedContextChip {
-  return (
-    'attachments' in chip &&
-    'snapshot' in chip &&
-    'combinedElementsMarkdown' in chip
-  );
+  return 'snapshot' in chip && 'combinedElementsMarkdown' in chip;
+}
+
+export function isAttachmentChip(chip: ChatChip): chip is AttachmentChip {
+  return 'sourceId' in chip && 'name' in chip;
 }

 export function getChipKey(chip: ChatChip) {
@@ -130,6 +131,9 @@ export function findChipIndex(chips: ChatChip[], chip: ChatChip) {
    if (isSelectedContextChip(chip)) {
      return isSelectedContextChip(item) && item.uuid === chip.uuid;
    }
+    if (isAttachmentChip(chip)) {
+      return isAttachmentChip(item) && item.sourceId === chip.sourceId;
+    }
    return -1;
  });
 }
--- a/packages/frontend/core/src/blocksuite/ai/components/ai-chat-composer/ai-chat-composer.ts
+++ b/packages/frontend/core/src/blocksuite/ai/components/ai-chat-composer/ai-chat-composer.ts
@@ -6,12 +6,13 @@ import type {
 } from '@affine/core/modules/ai-button';
 import type { WorkspaceDialogService } from '@affine/core/modules/dialogs';
 import type {
+  ContextEmbedStatus,
  ContextWorkspaceEmbeddingStatus,
  CopilotChatHistoryFragment,
+  CopilotContextBlob,
  CopilotContextDoc,
  CopilotContextFile,
 } from '@affine/graphql';
-import { ContextEmbedStatus } from '@affine/graphql';
 import { SignalWatcher, WithDisposable } from '@blocksuite/affine/global/lit';
 import type { EditorHost } from '@blocksuite/affine/std';
 import { ShadowlessElement } from '@blocksuite/affine/std';
@@ -23,9 +24,14 @@ import type {
 import { css, html, type PropertyValues } from 'lit';
 import { property, state } from 'lit/decorators.js';

-import { AIProvider } from '../../provider';
+import {
+  type AIChatParams,
+  AIProvider,
+  type AISendParams,
+} from '../../provider';
 import type { SearchMenuConfig } from '../ai-chat-add-context';
 import type {
+  AttachmentChip,
  ChatChip,
  CollectionChip,
  DocChip,
@@ -36,6 +42,7 @@ import type {
 } from '../ai-chat-chips';
 import {
  findChipIndex,
+  isAttachmentChip,
  isCollectionChip,
  isDocChip,
  isFileChip,
@@ -207,44 +214,10 @@ export class AIChatComposer extends SignalWatcher(
  override connectedCallback() {
    super.connectedCallback();
    this._disposables.add(
-      AIProvider.slots.requestOpenWithChat.subscribe(params => {
-        if (!params) return;
-
-        const { context, host } = params;
-        if (this.host !== host || !context) return;
-
-        if (
-          context.attachments ||
-          context.snapshot ||
-          context.combinedElementsMarkdown ||
-          context.html
-        ) {
-          // Wait for context value updated next frame
-          setTimeout(() => {
-            this.addSelectedContextChip().catch(console.error);
-          }, 0);
-        }
-      })
+      AIProvider.slots.requestOpenWithChat.subscribe(this.beforeChatContextSend)
    );
    this._disposables.add(
-      AIProvider.slots.requestSendWithChat.subscribe(params => {
-        if (!params) return;
-
-        const { context, host } = params;
-        if (this.host !== host || !context) return;
-
-        if (
-          context.attachments ||
-          context.snapshot ||
-          context.combinedElementsMarkdown ||
-          context.html
-        ) {
-          // Wait for context value updated next frame
-          setTimeout(() => {
-            this.addSelectedContextChip().catch(console.error);
-          }, 0);
-        }
-      })
+      AIProvider.slots.requestSendWithChat.subscribe(this.beforeChatContextSend)
    );
    this.initComposer().catch(console.error);
  }
@@ -266,6 +239,31 @@ export class AIChatComposer extends SignalWatcher(
    }
  }

+  private readonly beforeChatContextSend = (
+    params: AISendParams | AIChatParams | null
+  ) => {
+    if (!params) return;
+
+    const { context, host } = params;
+    if (this.host !== host || !context) return;
+
+    if (context) {
+      this.updateContext(context);
+    }
+    if (
+      context.docs ||
+      context.attachments ||
+      context.snapshot ||
+      context.combinedElementsMarkdown ||
+      context.html
+    ) {
+      // Wait for context value updated next frame
+      setTimeout(() => {
+        this.addSelectedContextChip().catch(console.error);
+      }, 0);
+    }
+  };
+
  private get isContextProcessing() {
    return this.chips.some(chip => chip.state === 'processing');
  }
@@ -414,19 +412,45 @@ export class AIChatComposer extends SignalWatcher(
  };

  private readonly addSelectedContextChip = async () => {
-    const { attachments, snapshot, combinedElementsMarkdown, docs, html } =
-      this.chatContextValue;
+    const {
+      attachments = [],
+      snapshot,
+      combinedElementsMarkdown,
+      docs = [],
+      html,
+    } = this.chatContextValue;
    await this.removeSelectedContextChip();
    const chip: SelectedContextChip = {
      uuid: uuidv4(),
-      attachments,
-      docs,
      snapshot,
      combinedElementsMarkdown,
      html,
-      state: attachments.length > 0 ? 'processing' : 'finished',
+      state: 'finished',
    };
    await this.addChip(chip, true);
+    await Promise.all(
+      docs.map(docId =>
+        this.addChip(
+          {
+            docId,
+            state: 'processing',
+          },
+          true
+        )
+      )
+    );
+    await Promise.all(
+      attachments.map(attachment =>
+        this.addChip(
+          {
+            sourceId: attachment.sourceId,
+            name: attachment.name,
+            state: 'processing',
+          },
+          true
+        )
+      )
+    );
  };

  private readonly removeSelectedContextChip = async () => {
@@ -449,8 +473,8 @@ export class AIChatComposer extends SignalWatcher(
    if (isCollectionChip(chip)) {
      return await this.addCollectionToContext(chip);
    }
-    if (isSelectedContextChip(chip)) {
-      return await this.addSelectedContextChipToContext(chip);
+    if (isAttachmentChip(chip)) {
+      return await this.addAttachmentChipToContext(chip);
    }
    return null;
  };
@@ -546,26 +570,29 @@ export class AIChatComposer extends SignalWatcher(
    }
  };

-  private readonly addSelectedContextChipToContext = async (
-    chip: SelectedContextChip
+  private readonly addAttachmentChipToContext = async (
+    chip: AttachmentChip
  ) => {
-    const { attachments, docs } = chip;
    const contextId = await this.createContextId();
    if (!contextId || !AIProvider.context) {
      throw new Error('Context not found');
    }
-    await AIProvider.context.addContextBlobs({
-      blobIds: attachments.map(attachment => attachment.sourceId),
-      contextId,
-    });
-    await Promise.all(
-      docs.map(docId =>
-        AIProvider.context?.addContextDoc({
-          contextId,
-          docId,
-        })
-      )
-    );
+    try {
+      const contextBlob = await AIProvider.context.addContextBlob({
+        blobId: chip.sourceId,
+        contextId,
+      });
+      this.updateChip(chip, {
+        state: contextBlob.status || 'processing',
+        blobId: chip.sourceId,
+      });
+    } catch (e) {
+      this.updateChip(chip, {
+        state: 'failed',
+        tooltip:
+          e instanceof Error ? e.message : 'Add context attachment error',
+      });
+    }
  };

  private readonly removeFromContext = async (
@@ -600,20 +627,18 @@ export class AIChatComposer extends SignalWatcher(
          collectionId: chip.collectionId,
        });
      }
-      if (isSelectedContextChip(chip)) {
-        const { attachments, docs } = chip;
-        await AIProvider.context.removeContextBlobs({
+      if (isAttachmentChip(chip)) {
+        return await AIProvider.context.removeContextBlob({
          contextId,
-          blobIds: attachments.map(attachment => attachment.sourceId),
+          blobId: chip.sourceId,
+        });
+      }
+      if (isSelectedContextChip(chip)) {
+        this.updateContext({
+          ...this.chatContextValue,
+          snapshot: null,
+          combinedElementsMarkdown: null,
        });
-        await Promise.all(
-          docs.map(docId =>
-            AIProvider.context?.removeContextDoc({
-              contextId,
-              docId,
-            })
-          )
-        );
      }
      return true;
    } catch {
@@ -707,7 +732,7 @@ export class AIChatComposer extends SignalWatcher(
    ];
    const hashMap = new Map<
      string,
-      CopilotContextDoc | CopilotContextFile | { status: ContextEmbedStatus }
+      CopilotContextDoc | CopilotContextFile | CopilotContextBlob
    >();
    const count: Record<ContextEmbedStatus, number> = {
      finished: 0,
@@ -722,18 +747,10 @@ export class AIChatComposer extends SignalWatcher(
      hashMap.set(file.id, file);
      file.status && count[file.status]++;
    });
-    const selectedChip = this.chips.find(c => isSelectedContextChip(c));
-    if (selectedChip) {
-      const status: ContextEmbedStatus = blobs.every(
-        blob => blob.status === 'finished'
-      )
-        ? ContextEmbedStatus.finished
-        : ContextEmbedStatus.processing;
-      hashMap.set(selectedChip.uuid, {
-        status,
-      });
-      count[status]++;
-    }
+    blobs.forEach(blob => {
+      hashMap.set(blob.id, blob);
+      blob.status && count[blob.status]++;
+    });
    const nextChips = this.chips.map(chip => {
      if (isTagChip(chip) || isCollectionChip(chip)) {
        return chip;
@@ -742,7 +759,11 @@ export class AIChatComposer extends SignalWatcher(
        ? chip.docId
        : isFileChip(chip)
          ? chip.fileId
-          : chip.uuid;
+          : isAttachmentChip(chip)
+            ? chip.sourceId
+            : isSelectedContextChip(chip)
+              ? chip.uuid
+              : undefined;
      const item = id && hashMap.get(id);
      if (item && item.status) {
        return {
--- a/packages/frontend/core/src/blocksuite/ai/effects.ts
+++ b/packages/frontend/core/src/blocksuite/ai/effects.ts
@@ -30,6 +30,7 @@ import { ChatPanelSplitView } from './chat-panel/split-view';
 import { ArtifactSkeleton } from './components/ai-artifact-skeleton';
 import { AIChatAddContext } from './components/ai-chat-add-context';
 import { ChatPanelAddPopover } from './components/ai-chat-chips/add-popover';
+import { ChatPanelAttachmentChip } from './components/ai-chat-chips/attachment-chip';
 import { ChatPanelCandidatesPopover } from './components/ai-chat-chips/candidates-popover';
 import { ChatPanelChips } from './components/ai-chat-chips/chat-panel-chips';
 import { ChatPanelChip } from './components/ai-chat-chips/chip';
@@ -166,6 +167,7 @@ export function registerAIEffects() {
  customElements.define('chat-panel-tag-chip', ChatPanelTagChip);
  customElements.define('chat-panel-collection-chip', ChatPanelCollectionChip);
  customElements.define('chat-panel-selected-chip', ChatPanelSelectedChip);
+  customElements.define('chat-panel-attachment-chip', ChatPanelAttachmentChip);
  customElements.define('chat-panel-chip', ChatPanelChip);
  customElements.define('ai-error-wrapper', AIErrorWrapper);
  customElements.define('ai-slides-renderer', AISlidesRenderer);
--- a/packages/frontend/core/src/blocksuite/ai/entries/edgeless/index.ts
+++ b/packages/frontend/core/src/blocksuite/ai/entries/edgeless/index.ts
@@ -57,7 +57,7 @@ export function edgelessToolbarAIEntryConfig(): ToolbarModuleConfig {
                aiPanel.hide();
                extractSelectedContent(host)
                  .then(context => {
-                    if (context?.attachments?.length) {
+                    if (context?.attachments?.length || context?.docs?.length) {
                      AIProvider.slots.requestOpenWithChat.next({
                        input,
                        host,
--- a/packages/frontend/core/src/blocksuite/ai/provider/setup-provider.tsx
+++ b/packages/frontend/core/src/blocksuite/ai/provider/setup-provider.tsx
@@ -748,31 +748,20 @@ Could you make a new website based on these notes and send back just the html fi
    ) => {
      return client.applyDocUpdates(workspaceId, docId, op, updates);
    },
-    addContextBlobs: async (options: {
-      blobIds: string[];
-      contextId: string;
-    }) => {
-      return Promise.all(
-        options.blobIds.map(blobId =>
-          client.addContextBlob({
-            contextId: options.contextId,
-            blobId,
-          })
-        )
-      );
+    addContextBlob: async (options: { blobId: string; contextId: string }) => {
+      return client.addContextBlob({
+        contextId: options.contextId,
+        blobId: options.blobId,
+      });
    },
-    removeContextBlobs: async (options: {
-      blobIds: string[];
+    removeContextBlob: async (options: {
+      blobId: string;
      contextId: string;
    }) => {
-      return Promise.all(
-        options.blobIds.map(blobId =>
-          client.removeContextBlob({
-            contextId: options.contextId,
-            blobId,
-          })
-        )
-      ).then(results => results.every(Boolean));
+      return client.removeContextBlob({
+        contextId: options.contextId,
+        blobId: options.blobId,
+      });
    },
  });