feat(core): add copilot tags and collections graphql apis (#11076)

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

### What Changed?
- Add `addContextCategoryMutation` and `removeContextCategoryMutation` graphql apis.
- Provide tag and collection apis for front-end components.

packages/frontend/core/src/blocksuite/ai/actions/types.ts

@@ -1,6 +1,7 @@
 import type {
   ChatHistoryOrder,
   ContextMatchedFileChunk,
+  CopilotContextCategory,
   CopilotContextDoc,
   CopilotContextFile,
   CopilotSessionType,
@@ -245,6 +246,8 @@ declare global {
     type AIDocsAndFilesContext = {
       docs: CopilotContextDoc[];
       files: CopilotContextFile[];
+      tags: CopilotContextCategory[];
+      collections: CopilotContextCategory[];
     };
 
     interface AIContextService {
@@ -275,6 +278,24 @@ declare global {
         contextId: string;
         fileId: string;
       }) => Promise<boolean>;
+      addContextTag: (options: {
+        contextId: string;
+        tagId: string;
+        docIds: string[];
+      }) => Promise<CopilotContextCategory>;
+      removeContextTag: (options: {
+        contextId: string;
+        tagId: string;
+      }) => Promise<boolean>;
+      addContextCollection: (options: {
+        contextId: string;
+        collectionId: string;
+        docIds: string[];
+      }) => Promise<CopilotContextCategory>;
+      removeContextCollection: (options: {
+        contextId: string;
+        collectionId: string;
+      }) => Promise<boolean>;
       getContextDocsAndFiles: (
         workspaceId: string,
         sessionId: string,

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

@@ -1,8 +1,10 @@
+import type { TagMeta } from '@affine/core/components/page-list';
 import type {
   SearchCollectionMenuAction,
   SearchDocMenuAction,
   SearchTagMenuAction,
 } from '@affine/core/modules/search-menu/services';
+import type { Collection } from '@affine/env/filter';
 import type { LinkedMenuGroup } from '@blocksuite/affine/blocks/root';
 import type { Store } from '@blocksuite/affine/store';
 import type { Signal } from '@preact/signals-core';
@@ -40,6 +42,15 @@ export interface DocDisplayConfig {
     >;
     cleanup: () => void;
   };
+  getTags: () => {
+    signal: Signal<TagMeta[]>;
+    cleanup: () => void;
+  };
+  getTagPageIds: (tagId: string) => string[];
+  getCollections: () => {
+    signal: Signal<Collection[]>;
+    cleanup: () => void;
+  };
 }
 
 export interface SearchMenuConfig {

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

@@ -83,6 +83,7 @@ export interface BaseChip {
    */
   state: ChipState;
   tooltip?: string | null;
+  createdAt?: number | null;
 }
 
 export interface DocChip extends BaseChip {
@@ -99,13 +100,10 @@ export interface FileChip extends BaseChip {
 
 export interface TagChip extends BaseChip {
   tagId: string;
-  tagName: string;
-  tagColor: string;
 }
 
 export interface CollectionChip extends BaseChip {
   collectionId: string;
-  collectionName: string;
 }
 
-export type ChatChip = DocChip | FileChip;
+export type ChatChip = DocChip | FileChip | TagChip | CollectionChip;

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

@@ -1,3 +1,5 @@
+import type { TagMeta } from '@affine/core/components/page-list';
+import type { Collection } from '@affine/env/filter';
 import {
   type EditorHost,
   ShadowlessElement,
@@ -17,14 +19,18 @@ import type { DocDisplayConfig, SearchMenuConfig } from './chat-config';
 import type {
   ChatChip,
   ChatContextValue,
+  CollectionChip,
   DocChip,
   FileChip,
+  TagChip,
 } from './chat-context';
 import {
   estimateTokenCount,
   getChipKey,
+  isCollectionChip,
   isDocChip,
   isFileChip,
+  isTagChip,
 } from './components/utils';
 
 // 100k tokens limit for the docs context
@@ -95,6 +101,10 @@ export class ChatPanelChips extends SignalWatcher(
     }>
   > = signal([]);
 
+  private _tags: Signal<TagMeta[]> = signal([]);
+
+  private _collections: Signal<Collection[]> = signal([]);
+
   private _cleanup: (() => void) | null = null;
 
   private _docIds: string[] = [];
@@ -133,6 +143,30 @@ export class ChatPanelChips extends SignalWatcher(
               .removeChip=${this._removeChip}
             ></chat-panel-file-chip>`;
           }
+          if (isTagChip(chip)) {
+            const tag = this._tags.value.find(tag => tag.id === chip.tagId);
+            if (!tag) {
+              return null;
+            }
+            return html`<chat-panel-tag-chip
+              .chip=${chip}
+              .tag=${tag}
+              .removeChip=${this._removeChip}
+            ></chat-panel-tag-chip>`;
+          }
+          if (isCollectionChip(chip)) {
+            const collection = this._collections.value.find(
+              collection => collection.id === chip.collectionId
+            );
+            if (!collection) {
+              return null;
+            }
+            return html`<chat-panel-collection-chip
+              .chip=${chip}
+              .collection=${collection}
+              .removeChip=${this._removeChip}
+            ></chat-panel-collection-chip>`;
+          }
           return null;
         }
       )}
@@ -144,6 +178,17 @@ export class ChatPanelChips extends SignalWatcher(
     </div>`;
   }
 
+  override connectedCallback(): void {
+    super.connectedCallback();
+    const tags = this.docDisplayConfig.getTags();
+    this._tags = tags.signal;
+    this._disposables.add(tags.cleanup);
+
+    const collections = this.docDisplayConfig.getCollections();
+    this._collections = collections.signal;
+    this._disposables.add(collections.cleanup);
+  }
+
   protected override updated(_changedProperties: PropertyValues): void {
     if (
       _changedProperties.has('chatContextValue') &&
@@ -204,15 +249,8 @@ export class ChatPanelChips extends SignalWatcher(
 
   private readonly _addChip = async (chip: ChatChip) => {
     this.isCollapsed = false;
-
     // remove the chip if it already exists
-    const chips = this.chatContextValue.chips.filter(item => {
-      if (isDocChip(chip)) {
-        return !isDocChip(item) || item.docId !== chip.docId;
-      } else {
-        return !isFileChip(item) || item.file !== chip.file;
-      }
-    });
+    const chips = this._omitChip(this.chatContextValue.chips, chip);
     this.updateContext({
       chips: [...chips, chip],
     });
@@ -227,13 +265,10 @@ export class ChatPanelChips extends SignalWatcher(
     chip: ChatChip,
     options: Partial<DocChip | FileChip>
   ) => {
-    const index = this.chatContextValue.chips.findIndex(item => {
-      if (isDocChip(chip)) {
-        return isDocChip(item) && item.docId === chip.docId;
-      } else {
-        return isFileChip(item) && item.file === chip.file;
-      }
-    });
+    const index = this._findChipIndex(this.chatContextValue.chips, chip);
+    if (index === -1) {
+      return;
+    }
     const nextChip: ChatChip = {
       ...chip,
       ...options,
@@ -248,21 +283,13 @@ export class ChatPanelChips extends SignalWatcher(
   };
 
   private readonly _removeChip = async (chip: ChatChip) => {
-    if (isDocChip(chip)) {
-      this.updateContext({
-        chips: this.chatContextValue.chips.filter(item => {
-          return !isDocChip(item) || item.docId !== chip.docId;
-        }),
-      });
+    const chips = this._omitChip(this.chatContextValue.chips, chip);
+    this.updateContext({
+      chips,
+    });
+    if (chips.length < this.chatContextValue.chips.length) {
+      await this._removeFromContext(chip);
     }
-    if (isFileChip(chip)) {
-      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) => {
@@ -271,29 +298,111 @@ export class ChatPanelChips extends SignalWatcher(
       return;
     }
     if (isDocChip(chip)) {
+      return await this._addDocToContext(chip);
+    }
+    if (isFileChip(chip)) {
+      return await this._addFileToContext(chip);
+    }
+    if (isTagChip(chip)) {
+      return await this._addTagToContext(chip);
+    }
+    if (isCollectionChip(chip)) {
+      return await this._addCollectionToContext(chip);
+    }
+    return null;
+  };
+
+  private readonly _addDocToContext = async (chip: DocChip) => {
+    const contextId = await this.getContextId();
+    if (!contextId || !AIProvider.context) {
+      return;
+    }
+    try {
       await AIProvider.context.addContextDoc({
         contextId,
         docId: chip.docId,
       });
+    } catch (e) {
+      this._updateChip(chip, {
+        state: 'failed',
+        tooltip: e instanceof Error ? e.message : 'Add context doc error',
+      });
     }
-    if (isFileChip(chip)) {
-      try {
-        const blobId = await this.host.doc.blobSync.set(chip.file);
-        const contextFile = await AIProvider.context.addContextFile(chip.file, {
-          contextId,
-          blobId,
-        });
-        this._updateChip(chip, {
-          state: contextFile.status,
-          blobId: contextFile.blobId,
-          fileId: contextFile.id,
-        });
-      } catch (e) {
-        this._updateChip(chip, {
-          state: 'failed',
-          tooltip: e instanceof Error ? e.message : 'Add context file error',
-        });
-      }
+  };
+
+  private readonly _addFileToContext = async (chip: FileChip) => {
+    const contextId = await this.getContextId();
+    if (!contextId || !AIProvider.context) {
+      return;
+    }
+    try {
+      const blobId = await this.host.doc.blobSync.set(chip.file);
+      const contextFile = await AIProvider.context.addContextFile(chip.file, {
+        contextId,
+        blobId,
+      });
+      this._updateChip(chip, {
+        state: contextFile.status,
+        blobId: contextFile.blobId,
+        fileId: contextFile.id,
+      });
+    } catch (e) {
+      this._updateChip(chip, {
+        state: 'failed',
+        tooltip: e instanceof Error ? e.message : 'Add context file error',
+      });
+    }
+  };
+
+  private readonly _addTagToContext = async (chip: TagChip) => {
+    const contextId = await this.getContextId();
+    if (!contextId || !AIProvider.context) {
+      return;
+    }
+    try {
+      // TODO: server side docIds calculation
+      const docIds = this.docDisplayConfig.getTagPageIds(chip.tagId);
+      await AIProvider.context.addContextTag({
+        contextId,
+        tagId: chip.tagId,
+        docIds,
+      });
+      this._updateChip(chip, {
+        state: 'finished',
+      });
+    } catch (e) {
+      this._updateChip(chip, {
+        state: 'failed',
+        tooltip: e instanceof Error ? e.message : 'Add context tag error',
+      });
+    }
+  };
+
+  private readonly _addCollectionToContext = async (chip: CollectionChip) => {
+    const contextId = await this.getContextId();
+    if (!contextId || !AIProvider.context) {
+      return;
+    }
+    try {
+      const collection = this._collections.value.find(
+        collection => collection.id === chip.collectionId
+      );
+      // TODO: server side docIds calculation
+      const docIds = collection?.allowList ?? [];
+      await AIProvider.context.addContextCollection({
+        contextId,
+        collectionId: chip.collectionId,
+        docIds,
+      });
+      this._updateChip(chip, {
+        state: 'finished',
+      });
+    } catch (e) {
+      this._updateChip(chip, {
+        state: 'failed',
+        tooltip:
+          e instanceof Error ? e.message : 'Add context collection error',
+      });
     }
   };
 
@@ -316,6 +425,18 @@ export class ChatPanelChips extends SignalWatcher(
         fileId: chip.fileId,
       });
     }
+    if (isTagChip(chip)) {
+      return await AIProvider.context.removeContextTag({
+        contextId,
+        tagId: chip.tagId,
+      });
+    }
+    if (isCollectionChip(chip)) {
+      return await AIProvider.context.removeContextCollection({
+        contextId,
+        collectionId: chip.collectionId,
+      });
+    }
     return true;
   };
 
@@ -324,7 +445,7 @@ export class ChatPanelChips extends SignalWatcher(
     newTokenCount: number
   ) => {
     const estimatedTokens = this.chatContextValue.chips.reduce((acc, chip) => {
-      if (isFileChip(chip)) {
+      if (isFileChip(chip) || isTagChip(chip) || isCollectionChip(chip)) {
         return acc;
       }
       if (chip.docId === newChip.docId) {
@@ -355,4 +476,44 @@ export class ChatPanelChips extends SignalWatcher(
     this.referenceDocs = signal;
     this._cleanup = cleanup;
   };
+
+  private readonly _omitChip = (chips: ChatChip[], chip: ChatChip) => {
+    return chips.filter(item => {
+      if (isDocChip(chip)) {
+        return !isDocChip(item) || item.docId !== chip.docId;
+      }
+      if (isFileChip(chip)) {
+        return !isFileChip(item) || item.file !== chip.file;
+      }
+      if (isTagChip(chip)) {
+        return !isTagChip(item) || item.tagId !== chip.tagId;
+      }
+      if (isCollectionChip(chip)) {
+        return (
+          !isCollectionChip(item) || item.collectionId !== chip.collectionId
+        );
+      }
+      return true;
+    });
+  };
+
+  private readonly _findChipIndex = (chips: ChatChip[], chip: ChatChip) => {
+    return chips.findIndex(item => {
+      if (isDocChip(chip)) {
+        return isDocChip(item) && item.docId === chip.docId;
+      }
+      if (isFileChip(chip)) {
+        return isFileChip(item) && item.file === chip.file;
+      }
+      if (isTagChip(chip)) {
+        return isTagChip(item) && item.tagId === chip.tagId;
+      }
+      if (isCollectionChip(chip)) {
+        return (
+          isCollectionChip(item) && item.collectionId === chip.collectionId
+        );
+      }
+      return -1;
+    });
+  };
 }

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

@@ -453,11 +453,19 @@ export class ChatPanelAddPopover extends SignalWatcher(
     this.abortController.abort();
   };
 
-  private readonly _addTagChip = (_tag: TagMeta) => {
+  private readonly _addTagChip = (tag: TagMeta) => {
+    this.addChip({
+      tagId: tag.id,
+      state: 'processing',
+    });
     this.abortController.abort();
   };
 
-  private readonly _addCollectionChip = (_collection: CollectionMeta) => {
+  private readonly _addCollectionChip = (collection: CollectionMeta) => {
+    this.addChip({
+      collectionId: collection.id,
+      state: 'processing',
+    });
     this.abortController.abort();
   };
 

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

@@ -1,3 +1,4 @@
+import type { Collection } from '@affine/env/filter';
 import { ShadowlessElement } from '@blocksuite/affine/block-std';
 import { SignalWatcher, WithDisposable } from '@blocksuite/affine/global/lit';
 import { CollectionsIcon } from '@blocksuite/icons/lit';
@@ -13,19 +14,31 @@ export class ChatPanelCollectionChip extends SignalWatcher(
   @property({ attribute: false })
   accessor chip!: CollectionChip;
 
+  @property({ attribute: false })
+  accessor removeChip!: (chip: CollectionChip) => void;
+
+  @property({ attribute: false })
+  accessor collection!: Collection;
+
   override render() {
-    const { state, collectionName } = this.chip;
+    const { state } = this.chip;
+    const { name } = this.collection;
     const isLoading = state === 'processing';
-    const tooltip = getChipTooltip(state, collectionName, this.chip.tooltip);
+    const tooltip = getChipTooltip(state, name, this.chip.tooltip);
     const collectionIcon = CollectionsIcon();
     const icon = getChipIcon(state, collectionIcon);
 
     return html`<chat-panel-chip
       .state=${state}
-      .name=${collectionName}
+      .name=${name}
       .tooltip=${tooltip}
       .icon=${icon}
       .closeable=${!isLoading}
+      .onChipDelete=${this.onChipDelete}
     ></chat-panel-chip>`;
   }
+
+  private readonly onChipDelete = () => {
+    this.removeChip(this.chip);
+  };
 }

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

@@ -10,7 +10,7 @@ import throttle from 'lodash-es/throttle';
 
 import { extractMarkdownFromDoc } from '../../utils/extract';
 import type { DocDisplayConfig } from '../chat-config';
-import type { ChatChip, DocChip } from '../chat-context';
+import type { DocChip } from '../chat-context';
 import { estimateTokenCount, getChipIcon, getChipTooltip } from './utils';
 
 const EXTRACT_DOC_THROTTLE = 1000;
@@ -22,13 +22,13 @@ export class ChatPanelDocChip extends SignalWatcher(
   accessor chip!: DocChip;
 
   @property({ attribute: false })
-  accessor addChip!: (chip: ChatChip) => void;
+  accessor addChip!: (chip: DocChip) => void;
 
   @property({ attribute: false })
-  accessor updateChip!: (chip: ChatChip, options: Partial<DocChip>) => void;
+  accessor updateChip!: (chip: DocChip, options: Partial<DocChip>) => void;
 
   @property({ attribute: false })
-  accessor removeChip!: (chip: ChatChip) => void;
+  accessor removeChip!: (chip: DocChip) => void;
 
   @property({ attribute: false })
   accessor checkTokenLimit!: (

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

@@ -4,7 +4,7 @@ import { SignalWatcher, WithDisposable } from '@blocksuite/affine/global/lit';
 import { html } from 'lit';
 import { property } from 'lit/decorators.js';
 
-import type { ChatChip, FileChip } from '../chat-context';
+import type { FileChip } from '../chat-context';
 import { getChipIcon, getChipTooltip } from './utils';
 
 export class ChatPanelFileChip extends SignalWatcher(
@@ -14,7 +14,7 @@ export class ChatPanelFileChip extends SignalWatcher(
   accessor chip!: FileChip;
 
   @property({ attribute: false })
-  accessor removeChip!: (chip: ChatChip) => void;
+  accessor removeChip!: (chip: FileChip) => void;
 
   override render() {
     const { state, file } = this.chip;

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

@@ -1,3 +1,4 @@
+import type { TagMeta } from '@affine/core/components/page-list';
 import { ShadowlessElement } from '@blocksuite/affine/block-std';
 import { SignalWatcher, WithDisposable } from '@blocksuite/affine/global/lit';
 import { css, html } from 'lit';
@@ -27,23 +28,35 @@ export class ChatPanelTagChip extends SignalWatcher(
   @property({ attribute: false })
   accessor chip!: TagChip;
 
+  @property({ attribute: false })
+  accessor removeChip!: (chip: TagChip) => void;
+
+  @property({ attribute: false })
+  accessor tag!: TagMeta;
+
   override render() {
-    const { state, tagName, tagColor } = this.chip;
+    const { state } = this.chip;
+    const { title, color } = this.tag;
     const isLoading = state === 'processing';
-    const tooltip = getChipTooltip(state, tagName, this.chip.tooltip);
+    const tooltip = getChipTooltip(state, title, this.chip.tooltip);
     const tagIcon = html`
       <div class="tag-icon-container">
-        <div class="tag-icon" style="background-color: ${tagColor};"></div>
+        <div class="tag-icon" style="background-color: ${color};"></div>
       </div>
     `;
     const icon = getChipIcon(state, tagIcon);
 
     return html`<chat-panel-chip
       .state=${state}
-      .name=${tagName}
+      .name=${title}
       .tooltip=${tooltip}
       .icon=${icon}
       .closeable=${!isLoading}
+      .onChipDelete=${this.onChipDelete}
     ></chat-panel-chip>`;
   }
+
+  private readonly onChipDelete = () => {
+    this.removeChip(this.chip);
+  };
 }

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

@@ -1,9 +1,15 @@
-import type { CopilotContextDoc, CopilotContextFile } from '@affine/graphql';
 import { LoadingIcon } from '@blocksuite/affine/blocks/image';
 import { WarningIcon } from '@blocksuite/icons/lit';
 import { type TemplateResult } from 'lit';
 
-import type { ChatChip, ChipState, DocChip, FileChip } from '../chat-context';
+import type {
+  ChatChip,
+  ChipState,
+  CollectionChip,
+  DocChip,
+  FileChip,
+  TagChip,
+} from '../chat-context';
 
 export function getChipTooltip(
   state: ChipState,
@@ -48,16 +54,12 @@ export function isFileChip(chip: ChatChip): chip is FileChip {
   return 'file' in chip && chip.file instanceof File;
 }
 
-export function isDocContext(
-  context: CopilotContextDoc | CopilotContextFile
-): context is CopilotContextDoc {
-  return !('blobId' in context);
+export function isTagChip(chip: ChatChip): chip is TagChip {
+  return 'tagId' in chip;
 }
 
-export function isFileContext(
-  context: CopilotContextDoc | CopilotContextFile
-): context is CopilotContextFile {
-  return 'blobId' in context;
+export function isCollectionChip(chip: ChatChip): chip is CollectionChip {
+  return 'collectionId' in chip;
 }
 
 export function getChipKey(chip: ChatChip) {
@@ -65,7 +67,13 @@ export function getChipKey(chip: ChatChip) {
     return chip.docId;
   }
   if (isFileChip(chip)) {
-    return chip.fileId;
+    return chip.file.name;
+  }
+  if (isTagChip(chip)) {
+    return chip.tagId;
+  }
+  if (isCollectionChip(chip)) {
+    return chip.collectionId;
   }
   return null;
 }

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

@@ -1,7 +1,12 @@
 import './chat-panel-input';
 import './chat-panel-messages';
 
-import type { CopilotContextDoc, CopilotContextFile } from '@affine/graphql';
+import type {
+  ContextEmbedStatus,
+  CopilotContextDoc,
+  CopilotContextFile,
+  CopilotDocType,
+} from '@affine/graphql';
 import type { EditorHost } from '@blocksuite/affine/block-std';
 import { ShadowlessElement } from '@blocksuite/affine/block-std';
 import { SignalWatcher, WithDisposable } from '@blocksuite/affine/global/lit';
@@ -32,11 +37,13 @@ import type {
   ChatChip,
   ChatContextValue,
   ChatItem,
+  CollectionChip,
   DocChip,
   FileChip,
+  TagChip,
 } from './chat-context';
 import type { ChatPanelMessages } from './chat-panel-messages';
-import { isDocChip, isDocContext } from './components/utils';
+import { isCollectionChip, isDocChip, isTagChip } from './components/utils';
 
 const DEFAULT_CHAT_CONTEXT_VALUE: ChatContextValue = {
   quote: '',
@@ -183,44 +190,61 @@ export class ChatPanel extends SignalWatcher(
     }
 
     // context initialized, show the chips
-    const { docs = [], files = [] } =
-      (await AIProvider.context?.getContextDocsAndFiles(
-        this.doc.workspace.id,
-        this._chatSessionId,
-        this._chatContextId
-      )) || {};
-    const list = [...docs, ...files].sort(
-      (a, b) =>
-        new Date(a.createdAt).getTime() - new Date(b.createdAt).getTime()
-    );
-    const chips: ChatChip[] = await Promise.all(
-      list.map(async item => {
-        if (isDocContext(item)) {
-          return {
-            docId: item.id,
-            state: item.status || 'processing',
-          } as DocChip;
-        }
-        const file = await this.host.doc.blobSync.get(item.blobId);
-        if (!file) {
-          return {
-            blobId: item.id,
-            file: new File([], item.name),
-            state: 'failed',
-            tooltip: 'File not found in blob storage',
-          } as FileChip;
-        } else {
-          return {
-            file: new File([file], item.name),
-            blobId: item.blobId,
-            fileId: item.id,
-            state: item.status,
-            tooltip: item.error,
-          } as FileChip;
-        }
+    const {
+      docs = [],
+      files = [],
+      tags = [],
+      collections = [],
+    } = (await AIProvider.context?.getContextDocsAndFiles(
+      this.doc.workspace.id,
+      this._chatSessionId,
+      this._chatContextId
+    )) || {};
+
+    const docChips: DocChip[] = docs.map(doc => ({
+      docId: doc.id,
+      state: doc.status || 'processing',
+      tooltip: doc.error,
+      createdAt: doc.createdAt,
+    }));
+
+    const fileChips: FileChip[] = await Promise.all(
+      files.map(async file => {
+        const blob = await this.host.doc.blobSync.get(file.blobId);
+        return {
+          file: new File(blob ? [blob] : [], file.name),
+          blobId: file.blobId,
+          fileId: file.id,
+          state: blob ? file.status : 'failed',
+          tooltip: blob ? file.error : 'File not found in blob storage',
+          createdAt: file.createdAt,
+        };
       })
     );
 
+    const tagChips: TagChip[] = tags.map(tag => ({
+      tagId: tag.id,
+      state: 'finished',
+      createdAt: tag.createdAt,
+    }));
+
+    const collectionChips: CollectionChip[] = collections.map(collection => ({
+      collectionId: collection.id,
+      state: 'finished',
+      createdAt: collection.createdAt,
+    }));
+
+    const chips: ChatChip[] = [
+      ...docChips,
+      ...fileChips,
+      ...tagChips,
+      ...collectionChips,
+    ].sort((a, b) => {
+      const aTime = a.createdAt ?? Date.now();
+      const bTime = b.createdAt ?? Date.now();
+      return aTime - bTime;
+    });
+
     this.chatContextValue = {
       ...this.chatContextValue,
       chips,
@@ -385,38 +409,55 @@ export class ChatPanel extends SignalWatcher(
       this._abortPoll();
       return;
     }
-    const { docs = [], files = [] } = result;
-    const hashMap = new Map<string, CopilotContextDoc | CopilotContextFile>();
-    const totalCount = docs.length + files.length;
-    let processingCount = 0;
+    const {
+      docs: sDocs = [],
+      files = [],
+      tags = [],
+      collections = [],
+    } = result;
+    const docs = [
+      ...sDocs,
+      ...tags.flatMap(tag => tag.docs),
+      ...collections.flatMap(collection => collection.docs),
+    ];
+    const hashMap = new Map<
+      string,
+      CopilotContextDoc | CopilotDocType | CopilotContextFile
+    >();
+    const count: Record<ContextEmbedStatus, number> = {
+      finished: 0,
+      processing: 0,
+      failed: 0,
+    };
     docs.forEach(doc => {
       hashMap.set(doc.id, doc);
-      if (doc.status === 'processing') {
-        processingCount++;
-      }
+      doc.status && count[doc.status]++;
     });
     files.forEach(file => {
       hashMap.set(file.id, file);
-      if (file.status === 'processing') {
-        processingCount++;
-      }
+      file.status && count[file.status]++;
     });
     const nextChips = this.chatContextValue.chips.map(chip => {
+      if (isTagChip(chip) || isCollectionChip(chip)) {
+        return chip;
+      }
       const id = isDocChip(chip) ? chip.docId : chip.fileId;
       const item = id && hashMap.get(id);
       if (item && item.status) {
         return {
           ...chip,
           state: item.status,
+          tooltip: 'error' in item ? item.error : undefined,
         };
       }
       return chip;
     });
+    const total = count.finished + count.processing + count.failed;
     this.updateContext({
       chips: nextChips,
-      embeddingProgress: [totalCount - processingCount, totalCount],
+      embeddingProgress: [count.finished, total],
     });
-    if (processingCount === 0) {
+    if (count.processing === 0) {
       this._abortPoll();
     }
   };

packages/frontend/core/src/blocksuite/ai/provider/copilot-client.ts

@@ -6,6 +6,7 @@ import {
 import { showAILoginRequiredAtom } from '@affine/core/components/affine/auth/ai-login-required';
 import type { UserFriendlyError } from '@affine/error';
 import {
+  addContextCategoryMutation,
   addContextDocMutation,
   addContextFileMutation,
   cleanupCopilotSessionMutation,
@@ -22,6 +23,7 @@ import {
   matchContextQuery,
   type QueryOptions,
   type QueryResponse,
+  removeContextCategoryMutation,
   removeContextDocMutation,
   removeContextFileMutation,
   type RequestOptions,
@@ -290,6 +292,30 @@ export class CopilotClient {
     return res.removeContextFile;
   }
 
+  async addContextCategory(
+    options: OptionsField<typeof addContextCategoryMutation>
+  ) {
+    const res = await this.gql({
+      query: addContextCategoryMutation,
+      variables: {
+        options,
+      },
+    });
+    return res.addContextCategory;
+  }
+
+  async removeContextCategory(
+    options: OptionsField<typeof removeContextCategoryMutation>
+  ) {
+    const res = await this.gql({
+      query: removeContextCategoryMutation,
+      variables: {
+        options,
+      },
+    });
+    return res.removeContextCategory;
+  }
+
   async getContextDocsAndFiles(
     workspaceId: string,
     sessionId: string,

packages/frontend/core/src/blocksuite/ai/provider/setup-provider.tsx

@@ -2,6 +2,7 @@ import { toggleGeneralAIOnboarding } from '@affine/core/components/affine/ai-onb
 import type { GlobalDialogService } from '@affine/core/modules/dialogs';
 import {
   type ChatHistoryOrder,
+  ContextCategories,
   type getCopilotHistoriesQuery,
   type RequestOptions,
 } from '@affine/graphql';
@@ -444,6 +445,47 @@ Could you make a new website based on these notes and send back just the html fi
     }) => {
       return client.removeContextFile(options);
     },
+    addContextTag: async (options: {
+      contextId: string;
+      tagId: string;
+      docIds: string[];
+    }) => {
+      return client.addContextCategory({
+        contextId: options.contextId,
+        type: ContextCategories.Tag,
+        categoryId: options.tagId,
+        docs: options.docIds,
+      });
+    },
+    removeContextTag: async (options: { contextId: string; tagId: string }) => {
+      return client.removeContextCategory({
+        contextId: options.contextId,
+        type: ContextCategories.Tag,
+        categoryId: options.tagId,
+      });
+    },
+    addContextCollection: async (options: {
+      contextId: string;
+      collectionId: string;
+      docIds: string[];
+    }) => {
+      return client.addContextCategory({
+        contextId: options.contextId,
+        type: ContextCategories.Collection,
+        categoryId: options.collectionId,
+        docs: options.docIds,
+      });
+    },
+    removeContextCollection: async (options: {
+      contextId: string;
+      collectionId: string;
+    }) => {
+      return client.removeContextCategory({
+        contextId: options.contextId,
+        type: ContextCategories.Collection,
+        categoryId: options.collectionId,
+      });
+    },
     getContextDocsAndFiles: async (
       workspaceId: string,
       sessionId: string,

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

@@ -2,9 +2,11 @@ import { ChatPanel } from '@affine/core/blocksuite/ai';
 import type { AffineEditorContainer } from '@affine/core/blocksuite/block-suite-editor';
 import { enableFootnoteConfigExtension } from '@affine/core/blocksuite/extensions';
 import { AINetworkSearchService } from '@affine/core/modules/ai-button/services/network-search';
+import { CollectionService } from '@affine/core/modules/collection';
 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 { TagService } from '@affine/core/modules/tag';
 import { WorkbenchService } from '@affine/core/modules/workbench';
 import { WorkspaceService } from '@affine/core/modules/workspace';
 import { RefNodeSlotsProvider } from '@blocksuite/affine/inlines/reference';
@@ -62,7 +64,8 @@ export const EditorChatPanel = forwardRef(function EditorChatPanel(
       const searchMenuService = framework.get(SearchMenuService);
       const workbench = framework.get(WorkbenchService).workbench;
       const docsSearchService = framework.get(DocsSearchService);
-
+      const tagService = framework.get(TagService);
+      const collectionService = framework.get(CollectionService);
       chatPanelRef.current.appSidebarConfig = {
         getWidth: () => {
           const width$ = workbench.sidebarWidth$;
@@ -96,6 +99,19 @@ export const EditorChatPanel = forwardRef(function EditorChatPanel(
           const docs$ = docsSearchService.watchRefsFrom(docIds);
           return createSignalFromObservable(docs$, []);
         },
+        getTags: () => {
+          const tagMetas$ = tagService.tagList.tagMetas$;
+          return createSignalFromObservable(tagMetas$, []);
+        },
+        getTagPageIds: (tagId: string) => {
+          const tag$ = tagService.tagList.tagByTagId$(tagId);
+          if (!tag$) return [];
+          return tag$.value?.pageIds$.value ?? [];
+        },
+        getCollections: () => {
+          const collections$ = collectionService.collections$;
+          return createSignalFromObservable(collections$, []);
+        },
       };
 
       chatPanelRef.current.searchMenuConfig = {