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.
2026-02-12 04:18:54 +00:00 · 2025-03-22 15:15:42 +00:00
parent 1f0fc9d47a
commit 331dd67e69
17 changed files with 498 additions and 128 deletions
--- a/packages/common/graphql/src/graphql/copilot-context-list-object.gql
+++ b/packages/common/graphql/src/graphql/copilot-context-list-object.gql
@@ -22,18 +22,22 @@ query listContextObject(
          createdAt
        }
        tags {
+          type
          id
          docs {
            id
            status
+            createdAt
          }
          createdAt
        }
        collections {
+          type
          id
          docs {
            id
            status
+            createdAt
          }
          createdAt
        }
--- a/packages/common/graphql/src/graphql/index.ts
+++ b/packages/common/graphql/src/graphql/index.ts
@@ -255,18 +255,22 @@ export const listContextObjectQuery = {
          createdAt
        }
        tags {
+          type
          id
          docs {
            id
            status
+            createdAt
          }
          createdAt
        }
        collections {
+          type
          id
          docs {
            id
            status
+            createdAt
          }
          createdAt
        }
--- a/packages/common/graphql/src/schema.ts
+++ b/packages/common/graphql/src/schema.ts
@@ -2593,22 +2593,26 @@ export type ListContextObjectQuery = {
        }>;
        tags: Array<{
          __typename?: 'CopilotContextCategory';
+          type: ContextCategories;
          id: string;
          createdAt: number;
          docs: Array<{
            __typename?: 'CopilotDocType';
            id: string;
            status: ContextEmbedStatus | null;
+            createdAt: number;
          }>;
        }>;
        collections: Array<{
          __typename?: 'CopilotContextCategory';
+          type: ContextCategories;
          id: string;
          createdAt: number;
          docs: Array<{
            __typename?: 'CopilotDocType';
            id: string;
            status: ContextEmbedStatus | null;
+            createdAt: number;
          }>;
        }>;
      }>;
--- a/packages/frontend/core/src/blocksuite/ai/actions/types.ts
+++ b/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,
--- a/packages/frontend/core/src/blocksuite/ai/chat-panel/chat-config.ts
+++ b/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 {
--- a/packages/frontend/core/src/blocksuite/ai/chat-panel/chat-context.ts
+++ b/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;
--- 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
@@ -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;
+    });
+  };
 }
--- a/packages/frontend/core/src/blocksuite/ai/chat-panel/components/add-popover.ts
+++ b/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();
  };

--- a/packages/frontend/core/src/blocksuite/ai/chat-panel/components/collection-chip.ts
+++ b/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);
+  };
 }
--- a/packages/frontend/core/src/blocksuite/ai/chat-panel/components/doc-chip.ts
+++ b/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!: (
--- a/packages/frontend/core/src/blocksuite/ai/chat-panel/components/file-chip.ts
+++ b/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;
--- a/packages/frontend/core/src/blocksuite/ai/chat-panel/components/tag-chip.ts
+++ b/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);
+  };
 }
--- a/packages/frontend/core/src/blocksuite/ai/chat-panel/components/utils.ts
+++ b/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;
 }
--- a/packages/frontend/core/src/blocksuite/ai/chat-panel/index.ts
+++ b/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();
    }
  };
--- a/packages/frontend/core/src/blocksuite/ai/provider/copilot-client.ts
+++ b/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,
--- a/packages/frontend/core/src/blocksuite/ai/provider/setup-provider.tsx
+++ b/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,
--- 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
@@ -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 = {