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)
2026-02-12 04:18:54 +00:00 · 2025-03-21 05:09:20 +00:00
parent 321dfaa881
commit a5b975ac46
5 changed files with 106 additions and 51 deletions
--- a/packages/frontend/core/src/blocksuite/ai/chat-panel/chat-config.ts
+++ b/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 {
--- a/packages/frontend/core/src/blocksuite/ai/chat-panel/chat-panel-chips.ts
+++ b/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() {
-    const isCollapsed =
-      this.isCollapsed &&
-      this.chatContextValue.chips.filter(c => c.state !== 'candidate').length >
-        1;
+  @state()
+  accessor referenceDocs: Signal<
+    Array<{
+      docId: string;
+      title: string;
+    }>
+  > = signal([]);

-    const chips = isCollapsed
-      ? this.chatContextValue.chips.slice(0, 1)
-      : this.chatContextValue.chips;
+  private _cleanup: (() => void) | null = null;
+
+  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);
-      if (removed) {
-        this.updateContext({
-          chips: this.chatContextValue.chips.filter(item => {
-            return !isDocChip(item) || item.docId !== chip.docId;
-          }),
-        });
-      }
+      this.updateContext({
+        chips: this.chatContextValue.chips.filter(item => {
+          return !isDocChip(item) || item.docId !== chip.docId;
+        }),
+      });
    }
    if (isFileChip(chip)) {
-      const removed = await this._removeFromContext(chip);
-      if (removed) {
-        this.updateContext({
-          chips: this.chatContextValue.chips.filter(item => {
-            return !isFileChip(item) || item.file !== chip.file;
-          }),
-        });
-      }
+      this.updateContext({
+        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;
+  };
 }
--- a/packages/frontend/core/src/blocksuite/ai/chat-panel/components/chip.ts
+++ b/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
-          ? html`
-              <div class="chip-card-close" @click=${this.onChipDelete}>
-                ${CloseIcon()}
-              </div>
-            `
-          : ''}
+        ${isCandidate
+          ? html`${PlusIcon()}`
+          : this.closeable
+            ? html`
+                <div class="chip-card-close" @click=${this.onChipDelete}>
+                  ${CloseIcon()}
+                </div>
+              `
+            : ''}
      </div>
    `;
  }
--- a/packages/frontend/core/src/blocksuite/ai/chat-panel/index.ts
+++ b/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;
    }
--- a/packages/frontend/core/src/desktop/pages/workspace/detail-page/tabs/chat.tsx
+++ b/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;