feat(core): suggest reference docs as candidate chips (#11050)

Close [BS-2464](https://linear.app/affine-design/issue/BS-2464).

![截屏2025-03-20 23.46.42.png](https://graphite-user-uploaded-assets-prod.s3.amazonaws.com/sJGviKxfE3Ap685cl5bj/c52b32de-1beb-4194-8af6-a175e5e6b986.png)

packages/frontend/core/src/blocksuite/ai/chat-panel/chat-config.ts

@@ -31,6 +31,15 @@ export interface DocDisplayConfig {
     cleanup: () => void;
   };
   getDoc: (docId: string) => Store | null;
+  getReferenceDocs: (docIds: string[]) => {
+    signal: Signal<
+      Array<{
+        docId: string;
+        title: string;
+      }>
+    >;
+    cleanup: () => void;
+  };
 }
 
 export interface SearchMenuConfig {

packages/frontend/core/src/blocksuite/ai/chat-panel/chat-panel-chips.ts

@@ -3,12 +3,14 @@ import {
   ShadowlessElement,
 } from '@blocksuite/affine/block-std';
 import { createLitPortal } from '@blocksuite/affine/components/portal';
-import { WithDisposable } from '@blocksuite/affine/global/lit';
+import { SignalWatcher, WithDisposable } from '@blocksuite/affine/global/lit';
 import { PlusIcon } from '@blocksuite/icons/lit';
 import { flip, offset } from '@floating-ui/dom';
+import { type Signal, signal } from '@preact/signals-core';
 import { css, html, nothing, type PropertyValues } from 'lit';
 import { property, query, state } from 'lit/decorators.js';
 import { repeat } from 'lit/directives/repeat.js';
+import { isEqual } from 'lodash-es';
 
 import { AIProvider } from '../provider';
 import type { DocDisplayConfig, SearchMenuConfig } from './chat-config';
@@ -28,7 +30,9 @@ import {
 // 100k tokens limit for the docs context
 const MAX_TOKEN_COUNT = 100000;
 
-export class ChatPanelChips extends WithDisposable(ShadowlessElement) {
+export class ChatPanelChips extends SignalWatcher(
+  WithDisposable(ShadowlessElement)
+) {
   static override styles = css`
     .chips-wrapper {
       display: flex;
@@ -83,15 +87,26 @@ export class ChatPanelChips extends WithDisposable(ShadowlessElement) {
   @state()
   accessor isCollapsed = false;
 
-  override render() {
+  @state()
-    const isCollapsed =
+  accessor referenceDocs: Signal<
-      this.isCollapsed &&
+    Array<{
-      this.chatContextValue.chips.filter(c => c.state !== 'candidate').length >
+      docId: string;
-        1;
+      title: string;
+    }>
+  > = signal([]);
 
-    const chips = isCollapsed
+  private _cleanup: (() => void) | null = null;
-      ? this.chatContextValue.chips.slice(0, 1)
+
-      : this.chatContextValue.chips;
+  private _docIds: string[] = [];
+
+  override render() {
+    const candidates: DocChip[] = this.referenceDocs.value.map(doc => ({
+      docId: doc.docId,
+      state: 'candidate',
+    }));
+    const allChips = this.chatContextValue.chips.concat(candidates);
+    const isCollapsed = this.isCollapsed && allChips.length > 1;
+    const chips = isCollapsed ? allChips.slice(0, 1) : allChips;
 
     return html` <div class="chips-wrapper">
       <div class="add-button" @click=${this._toggleAddDocMenu}>
@@ -123,7 +138,7 @@ export class ChatPanelChips extends WithDisposable(ShadowlessElement) {
       )}
       ${isCollapsed
         ? html`<div class="collapse-button" @click=${this._toggleCollapse}>
-            +${this.chatContextValue.chips.length - 1}
+            +${allChips.length - 1}
           </div>`
         : nothing}
     </div>`;
@@ -137,6 +152,16 @@ export class ChatPanelChips extends WithDisposable(ShadowlessElement) {
     ) {
       this.isCollapsed = true;
     }
+
+    // TODO only update when the chips are changed
+    if (_changedProperties.has('chatContextValue')) {
+      this._updateReferenceDocs();
+    }
+  }
+
+  override disconnectedCallback(): void {
+    super.disconnectedCallback();
+    this._cleanup?.();
   }
 
   private readonly _toggleCollapse = () => {
@@ -179,16 +204,7 @@ export class ChatPanelChips extends WithDisposable(ShadowlessElement) {
 
   private readonly _addChip = async (chip: ChatChip) => {
     this.isCollapsed = false;
-    if (
+
-      this.chatContextValue.chips.length === 1 &&
-      this.chatContextValue.chips[0].state === 'candidate'
-    ) {
-      this.updateContext({
-        chips: [chip],
-      });
-      await this._addToContext(chip);
-      return;
-    }
     // remove the chip if it already exists
     const chips = this.chatContextValue.chips.filter(item => {
       if (isDocChip(chip)) {
@@ -233,25 +249,20 @@ export class ChatPanelChips extends WithDisposable(ShadowlessElement) {
 
   private readonly _removeChip = async (chip: ChatChip) => {
     if (isDocChip(chip)) {
-      const removed = await this._removeFromContext(chip);
+      this.updateContext({
-      if (removed) {
+        chips: this.chatContextValue.chips.filter(item => {
-        this.updateContext({
+          return !isDocChip(item) || item.docId !== chip.docId;
-          chips: this.chatContextValue.chips.filter(item => {
+        }),
-            return !isDocChip(item) || item.docId !== chip.docId;
+      });
-          }),
-        });
-      }
     }
     if (isFileChip(chip)) {
-      const removed = await this._removeFromContext(chip);
+      this.updateContext({
-      if (removed) {
+        chips: this.chatContextValue.chips.filter(item => {
-        this.updateContext({
+          return !isFileChip(item) || item.file !== chip.file;
-          chips: this.chatContextValue.chips.filter(item => {
+        }),
-            return !isFileChip(item) || item.file !== chip.file;
+      });
-          }),
-        });
-      }
     }
+    await this._removeFromContext(chip);
   };
 
   private readonly _addToContext = async (chip: ChatChip) => {
@@ -328,4 +339,20 @@ export class ChatPanelChips extends WithDisposable(ShadowlessElement) {
     }, 0);
     return estimatedTokens <= MAX_TOKEN_COUNT;
   };
+
+  private readonly _updateReferenceDocs = () => {
+    const docIds = this.chatContextValue.chips
+      .filter(isDocChip)
+      .filter(chip => chip.state !== 'candidate')
+      .map(chip => chip.docId);
+    if (isEqual(this._docIds, docIds)) {
+      return;
+    }
+
+    this._cleanup?.();
+    this._docIds = docIds;
+    const { signal, cleanup } = this.docDisplayConfig.getReferenceDocs(docIds);
+    this.referenceDocs = signal;
+    this._cleanup = cleanup;
+  };
 }

packages/frontend/core/src/blocksuite/ai/chat-panel/components/chip.ts

@@ -1,6 +1,6 @@
 import { ShadowlessElement } from '@blocksuite/affine/block-std';
 import { SignalWatcher, WithDisposable } from '@blocksuite/affine/global/lit';
-import { CloseIcon } from '@blocksuite/icons/lit';
+import { CloseIcon, PlusIcon } from '@blocksuite/icons/lit';
 import { css, html, type TemplateResult } from 'lit';
 import { property } from 'lit/decorators.js';
 
@@ -18,14 +18,18 @@ export class ChatPanelChip extends SignalWatcher(
       margin: 4px;
       padding: 0 4px;
       border-radius: 4px;
-      border: 1px solid var(--affine-border-color);
+      border: 0.5px solid var(--affine-border-color);
       background: var(--affine-background-primary-color);
       box-sizing: border-box;
     }
     .chip-card[data-state='candidate'] {
-      border-width: 0.5px;
+      border-width: 1px;
       border-style: dashed;
-      background: var(--affine-background-secondary-color);
+      background: var(--affine-tag-white);
+      color: var(--affine-icon-secondary);
+    }
+    .chip-card[data-state='candidate'] svg {
+      color: var(--affine-icon-secondary);
     }
     .chip-card[data-state='failed'] {
       color: var(--affine-error-color);
@@ -91,6 +95,7 @@ export class ChatPanelChip extends SignalWatcher(
   accessor onChipClick: () => void = () => {};
 
   override render() {
+    const isCandidate = this.state === 'candidate';
     return html`
       <div
         class="chip-card"
@@ -104,13 +109,15 @@ export class ChatPanelChip extends SignalWatcher(
           </span>
           <affine-tooltip>${this.tooltip}</affine-tooltip>
         </div>
-        ${this.closeable
+        ${isCandidate
-          ? html`
+          ? html`${PlusIcon()}`
-              <div class="chip-card-close" @click=${this.onChipDelete}>
+          : this.closeable
-                ${CloseIcon()}
+            ? html`
-              </div>
+                <div class="chip-card-close" @click=${this.onChipDelete}>
-            `
+                  ${CloseIcon()}
-          : ''}
+                </div>
+              `
+            : ''}
       </div>
     `;
   }

packages/frontend/core/src/blocksuite/ai/chat-panel/index.ts

@@ -177,7 +177,7 @@ export class ChatPanel extends SignalWatcher(
   };
 
   private readonly _initChips = async () => {
-    // context not initialized, show candidate chip
+    // context not initialized
     if (!this._chatSessionId || !this._chatContextId) {
       return;
     }

packages/frontend/core/src/desktop/pages/workspace/detail-page/tabs/chat.tsx

@@ -3,6 +3,7 @@ import type { AffineEditorContainer } from '@affine/core/blocksuite/block-suite-
 import { enableFootnoteConfigExtension } from '@affine/core/blocksuite/extensions';
 import { AINetworkSearchService } from '@affine/core/modules/ai-button/services/network-search';
 import { DocDisplayMetaService } from '@affine/core/modules/doc-display-meta';
+import { DocsSearchService } from '@affine/core/modules/docs-search';
 import { SearchMenuService } from '@affine/core/modules/search-menu/services';
 import { WorkbenchService } from '@affine/core/modules/workbench';
 import { WorkspaceService } from '@affine/core/modules/workspace';
@@ -54,11 +55,14 @@ export const EditorChatPanel = forwardRef(function EditorChatPanel(
       chatPanelRef.current = new ChatPanel();
       chatPanelRef.current.host = editor.host;
       chatPanelRef.current.doc = editor.doc;
+
       const searchService = framework.get(AINetworkSearchService);
       const docDisplayMetaService = framework.get(DocDisplayMetaService);
       const workspaceService = framework.get(WorkspaceService);
       const searchMenuService = framework.get(SearchMenuService);
       const workbench = framework.get(WorkbenchService).workbench;
+      const docsSearchService = framework.get(DocsSearchService);
+
       chatPanelRef.current.appSidebarConfig = {
         getWidth: () => {
           const width$ = workbench.sidebarWidth$;
@@ -69,11 +73,13 @@ export const EditorChatPanel = forwardRef(function EditorChatPanel(
           return createSignalFromObservable(open$, true);
         },
       };
+
       chatPanelRef.current.networkSearchConfig = {
         visible: searchService.visible,
         enabled: searchService.enabled,
         setEnabled: searchService.setEnabled,
       };
+
       chatPanelRef.current.docDisplayConfig = {
         getIcon: (docId: string) => {
           return docDisplayMetaService.icon$(docId, { type: 'lit' }).value;
@@ -86,7 +92,12 @@ export const EditorChatPanel = forwardRef(function EditorChatPanel(
           const doc = workspaceService.workspace.docCollection.getDoc(docId);
           return doc;
         },
+        getReferenceDocs: (docIds: string[]) => {
+          const docs$ = docsSearchService.watchRefsFrom(docIds);
+          return createSignalFromObservable(docs$, []);
+        },
       };
+
       chatPanelRef.current.searchMenuConfig = {
         getDocMenuGroup: (query, action, abortSignal) => {
           return searchMenuService.getDocMenuGroup(query, action, abortSignal);
@@ -102,10 +113,11 @@ export const EditorChatPanel = forwardRef(function EditorChatPanel(
           );
         },
       };
-      const previewSpecBuilder = enableFootnoteConfigExtension(
+
+      chatPanelRef.current.previewSpecBuilder = enableFootnoteConfigExtension(
         SpecProvider._.getSpec('preview:page')
       );
-      chatPanelRef.current.previewSpecBuilder = previewSpecBuilder;
+
       containerRef.current?.append(chatPanelRef.current);
     } else {
       chatPanelRef.current.host = editor.host;