feat: support save chat to block (#7481)

2024-12-23 09:22:38 +03:00 · 2024-07-24 09:46:36 +00:00 · 2024-07-24 09:46:36 +00:00 · 9d446469f8
commit 9d446469f8
parent 98281a6394
18 changed files with 463 additions and 44 deletions
--- a/packages/common/infra/package.json
+++ b/packages/common/infra/package.json
@ -16,6 +16,7 @@
    "@affine/templates": "workspace:*",
    "@blocksuite/blocks": "0.16.0-canary-202407230727-128fc57",
    "@blocksuite/global": "0.16.0-canary-202407230727-128fc57",
+    "@blocksuite/presets": "0.16.0-canary-202407230727-128fc57",
    "@blocksuite/store": "0.16.0-canary-202407230727-128fc57",
    "@datastructures-js/binary-search-tree": "^5.3.2",
    "foxact": "^0.2.33",
--- a/packages/common/infra/src/modules/workspace/global-schema.ts
+++ b/packages/common/infra/src/modules/workspace/global-schema.ts
@ -1,6 +1,8 @@
 import { AffineSchemas } from '@blocksuite/blocks/schemas';
+import { AIChatBlockSchema } from '@blocksuite/presets';
 import { Schema } from '@blocksuite/store';

 export const globalBlockSuiteSchema = new Schema();

-globalBlockSuiteSchema.register(AffineSchemas);
+const schemas = [...AffineSchemas, AIChatBlockSchema];
+globalBlockSuiteSchema.register(schemas);
--- a/packages/frontend/core/src/blocksuite/presets/ai/_common/icons.ts
+++ b/packages/frontend/core/src/blocksuite/presets/ai/_common/icons.ts
@ -1178,3 +1178,17 @@ export const CommentIcon = html`<svg
    d="M10.4167 3.125C6.84987 3.125 3.95837 6.01649 3.95837 9.58333C3.95837 10.5626 4.1759 11.4893 4.56469 12.3193C4.62513 12.4484 4.64581 12.5938 4.62202 12.7365L4.09372 15.9063L7.26351 15.378C7.40628 15.3542 7.55167 15.3749 7.68071 15.4354C8.51071 15.8241 9.43744 16.0417 10.4167 16.0417C13.9835 16.0417 16.875 13.1502 16.875 9.58333C16.875 6.01649 13.9835 3.125 10.4167 3.125ZM2.70837 9.58333C2.70837 5.32614 6.15951 1.875 10.4167 1.875C14.6739 1.875 18.125 5.32614 18.125 9.58333C18.125 13.8405 14.6739 17.2917 10.4167 17.2917C9.31104 17.2917 8.25828 17.0585 7.30619 16.6382L3.5512 17.264C3.07179 17.3439 2.65615 16.9283 2.73606 16.4488L3.36189 12.6939C2.94154 11.7418 2.70837 10.689 2.70837 9.58333Z"
  />
 </svg>`;
+
+export const BlockIcon = html`<svg
+  width="20"
+  height="20"
+  viewBox="0 0 20 20"
+  fill="currentColor"
+  xmlns="http://www.w3.org/2000/svg"
+>
+  <path
+    fill-rule="evenodd"
+    clip-rule="evenodd"
+    d="M9.7495 2.34391C9.9092 2.27404 10.0908 2.27404 10.2505 2.34391L16.9172 5.26057C17.1447 5.3601 17.2917 5.58486 17.2917 5.83317V14.1665C17.2917 14.4148 17.1447 14.6396 16.9172 14.7391L10.2505 17.6558C10.0908 17.7256 9.9092 17.7256 9.7495 17.6558L3.08283 14.7391C2.85534 14.6396 2.70834 14.4148 2.70834 14.1665V5.83317C2.70834 5.58486 2.85534 5.3601 3.08283 5.26057L9.7495 2.34391ZM3.95834 6.78881V13.7577L9.37501 16.1275V9.1586L3.95834 6.78881ZM10.625 9.1586V16.1275L16.0417 13.7577V6.78881L10.625 9.1586ZM15.1074 5.83317L10 8.06764L4.89265 5.83317L10 3.5987L15.1074 5.83317Z"
+  />
+</svg>`;
--- a/packages/frontend/core/src/blocksuite/presets/ai/actions/edgeless-response.ts
+++ b/packages/frontend/core/src/blocksuite/presets/ai/actions/edgeless-response.ts
@ -15,7 +15,6 @@ import {
  EDGELESS_TEXT_BLOCK_MIN_HEIGHT,
  EDGELESS_TEXT_BLOCK_MIN_WIDTH,
  EdgelessTextBlockModel,
-  EmbedHtmlBlockSpec,
  fitContent,
  ImageBlockModel,
  InsertBelowIcon,
@ -463,7 +462,7 @@ export const responses: {

    host.doc.transact(() => {
      edgelessRoot.doc.addBlock(
-        EmbedHtmlBlockSpec.schema.model.flavour as 'affine:embed-html',
+        'affine:embed-html',
        {
          html,
          design: 'ai:makeItReal', // as tag
--- a/packages/frontend/core/src/blocksuite/presets/ai/actions/types.ts
+++ b/packages/frontend/core/src/blocksuite/presets/ai/actions/types.ts
@ -1,3 +1,4 @@
+import type { getCopilotHistoriesQuery, RequestOptions } from '@affine/graphql';
 import type { EditorHost } from '@blocksuite/block-std';
 import type { BlockModel } from '@blocksuite/store';

@ -74,6 +75,13 @@ declare global {
      where: TrackerWhere;
    }

+    interface AIForkChatSessionOptions {
+      docId: string;
+      workspaceId: string;
+      sessionId: string;
+      latestMessageId: string;
+    }
+
    interface AIImageActionOptions extends AITextActionOptions {
      content?: string;
      seed?: string;
@ -222,12 +230,22 @@ declare global {
      action: string;
      createdAt: string;
      messages: {
+        id: string; // message id
        content: string;
        createdAt: string;
-        role: 'user' | 'assistant';
+        role: MessageRole;
      }[];
    }

+    type MessageRole = 'user' | 'assistant';
+
+    type AIHistoryIds = Pick<AIHistory, 'sessionId' | 'messages'> & {
+      messages: Pick<
+        AIHistory['messages'][number],
+        'id' | 'createdAt' | 'role'
+      >[];
+    };
+
    interface AIHistoryService {
      // non chat histories
      actions: (
@ -236,13 +254,23 @@ declare global {
      ) => Promise<AIHistory[] | undefined>;
      chats: (
        workspaceId: string,
-        docId?: string
+        docId?: string,
+        options?: RequestOptions<
+          typeof getCopilotHistoriesQuery
+        >['variables']['options']
      ) => Promise<AIHistory[] | undefined>;
      cleanup: (
        workspaceId: string,
        docId: string,
        sessionIds: string[]
      ) => Promise<void>;
+      ids: (
+        workspaceId: string,
+        docId?: string,
+        options?: RequestOptions<
+          typeof getCopilotHistoriesQuery
+        >['variables']['options']
+      ) => Promise<AIHistoryIds[] | undefined>;
    }

    interface AIPhotoEngineService {
--- a/packages/frontend/core/src/blocksuite/presets/ai/chat-panel/actions/actions-handle.ts
+++ b/packages/frontend/core/src/blocksuite/presets/ai/chat-panel/actions/actions-handle.ts
@ -1,9 +1,15 @@
+import { ChatHistoryOrder } from '@affine/graphql';
 import type {
  BlockSelection,
  EditorHost,
  TextSelection,
 } from '@blocksuite/block-std';
-import type { EdgelessRootService, ImageSelection } from '@blocksuite/blocks';
+import type {
+  DocMode,
+  EdgelessRootService,
+  ImageSelection,
+  PageRootService,
+} from '@blocksuite/blocks';
 import {
  BlocksUtils,
  getElementsBound,
@ -11,25 +17,158 @@ import {
 } from '@blocksuite/blocks';
 import type { SerializedXYWH } from '@blocksuite/global/utils';
 import { Bound } from '@blocksuite/global/utils';
+import type { Doc } from '@blocksuite/store';
+import type { TemplateResult } from 'lit';

-import { CreateIcon, InsertBelowIcon, ReplaceIcon } from '../../_common/icons';
+import {
+  BlockIcon,
+  CreateIcon,
+  InsertBelowIcon,
+  ReplaceIcon,
+} from '../../_common/icons';
+import { AIProvider } from '../../provider';
 import { reportResponse } from '../../utils/action-reporter';
 import { insertBelow, replace } from '../../utils/editor-actions';
 import { insertFromMarkdown } from '../../utils/markdown-utils';
+import type { ChatBlockMessage, ChatContextValue } from '../chat-context';

 const { matchFlavours } = BlocksUtils;

-const CommonActions = [
+type Selections = {
+  text?: TextSelection;
+  blocks?: BlockSelection[];
+  images?: ImageSelection[];
+};
+
+type ChatAction = {
+  icon: TemplateResult<1>;
+  title: string;
+  toast: string;
+  showWhen: (host: EditorHost) => boolean;
+  handler: (
+    host: EditorHost,
+    content: string,
+    currentSelections: Selections,
+    chatContext?: ChatContextValue,
+    messageId?: string
+  ) => Promise<boolean>;
+};
+
+async function constructChatBlockMessages(doc: Doc, forkSessionId: string) {
+  const userInfo = await AIProvider.userInfo;
+  // Get fork session messages
+  const histories = await AIProvider.histories?.chats(
+    doc.collection.id,
+    doc.id,
+    {
+      sessionId: forkSessionId,
+      messageOrder: ChatHistoryOrder.asc,
+    }
+  );
+
+  if (!histories || !histories.length) {
+    return [];
+  }
+
+  const messages = histories[0].messages.map(message => {
+    const { role, id, content, createdAt } = message;
+    const isUser = role === 'user';
+    const userInfoProps = isUser
+      ? {
+          userId: userInfo?.id,
+          userName: userInfo?.name,
+          avatarUrl: userInfo?.avatarUrl ?? undefined,
+        }
+      : {};
+    return {
+      id,
+      role,
+      content,
+      createdAt,
+      attachments: [],
+      ...userInfoProps,
+    };
+  });
+  return messages;
+}
+
+function getViewportCenter(
+  mode: DocMode,
+  rootService: PageRootService | EdgelessRootService
+) {
+  const center = { x: 0, y: 0 };
+  if (mode === 'page') {
+    const viewport = rootService.editPropsStore.getStorage('viewport');
+    if (viewport) {
+      if ('xywh' in viewport) {
+        const bound = Bound.deserialize(viewport.xywh);
+        center.x = bound.x + bound.w / 2;
+        center.y = bound.y + bound.h / 2;
+      } else {
+        center.x = viewport.centerX;
+        center.y = viewport.centerY;
+      }
+    }
+  } else {
+    // Else we should get latest viewport center from the edgeless root service
+    const edgelessService = rootService as EdgelessRootService;
+    center.x = edgelessService.viewport.centerX;
+    center.y = edgelessService.viewport.centerY;
+  }
+
+  return center;
+}
+
+// Add AI chat block and focus on it
+function addAIChatBlock(
+  doc: Doc,
+  messages: ChatBlockMessage[],
+  sessionId: string,
+  viewportCenter: { x: number; y: number }
+) {
+  if (!messages.length || !sessionId) {
+    return;
+  }
+
+  const surfaceBlock = doc
+    .getBlocks()
+    .find(block => block.flavour === 'affine:surface');
+  if (!surfaceBlock) {
+    return;
+  }
+
+  // Add AI chat block to the center of the viewport
+  const width = 300; // AI_CHAT_BLOCK_WIDTH = 300
+  const height = 320; // AI_CHAT_BLOCK_HEIGHT = 320
+  const x = viewportCenter.x - width / 2;
+  const y = viewportCenter.y - height / 2;
+  const bound = new Bound(x, y, width, height);
+  const aiChatBlockId = doc.addBlock(
+    'affine:embed-ai-chat' as keyof BlockSuite.BlockModels,
+    {
+      xywh: bound.serialize(),
+      messages: JSON.stringify(messages),
+      sessionId,
+    },
+    surfaceBlock.id
+  );
+
+  return aiChatBlockId;
+}
+
+const CommonActions: ChatAction[] = [
  {
    icon: ReplaceIcon,
    title: 'Replace selection',
+    showWhen: () => true,
    toast: 'Successfully replaced',
    handler: async (
      host: EditorHost,
      content: string,
-      currentTextSelection?: TextSelection,
-      currentBlockSelections?: BlockSelection[]
+      currentSelections: Selections
    ) => {
+      const currentTextSelection = currentSelections.text;
+      const currentBlockSelections = currentSelections.blocks;
      const [_, data] = host.command
        .chain()
        .getSelectedBlocks({
@ -67,14 +206,16 @@ const CommonActions = [
  {
    icon: InsertBelowIcon,
    title: 'Insert below',
+    showWhen: () => true,
    toast: 'Successfully inserted',
    handler: async (
      host: EditorHost,
      content: string,
-      currentTextSelection?: TextSelection,
-      currentBlockSelections?: BlockSelection[],
-      currentImageSelections?: ImageSelection[]
+      currentSelections: Selections
    ) => {
+      const currentTextSelection = currentSelections.text;
+      const currentBlockSelections = currentSelections.blocks;
+      const currentImageSelections = currentSelections.images;
      const [_, data] = host.command
        .chain()
        .getSelectedBlocks({
@ -95,11 +236,90 @@ const CommonActions = [
  },
 ];

+const SAVE_CHAT_TO_BLOCK_ACTION: ChatAction = {
+  icon: BlockIcon,
+  title: 'Save chat to block',
+  toast: 'Successfully saved chat to a block',
+  showWhen: (host: EditorHost) =>
+    !!host.doc.awarenessStore.getFlag('enable_ai_chat_block'),
+  handler: async (
+    host: EditorHost,
+    _,
+    __,
+    chatContext?: ChatContextValue,
+    messageId?: string
+  ) => {
+    // The chat session id and the latest message id are required to fork the chat session
+    const parentSessionId = chatContext?.chatSessionId;
+    if (!messageId || !parentSessionId) {
+      return false;
+    }
+
+    const rootService = host.spec.getService('affine:page');
+    if (!rootService) return false;
+
+    const { docModeService, notificationService } = rootService;
+    const curMode = docModeService.getMode();
+    const viewportCenter = getViewportCenter(curMode, rootService);
+    // If current mode is not edgeless, switch to edgeless mode first
+    if (curMode !== 'edgeless') {
+      // Set mode to edgeless
+      docModeService.setMode('edgeless');
+      // Notify user to switch to edgeless mode
+      notificationService?.notify({
+        title: 'Save chat to a block',
+        accent: 'info',
+        message:
+          'This feature is not available in the page editor. Switch to edgeless mode.',
+        onClose: function (): void {},
+      });
+    }
+
+    try {
+      const newSessionId = await AIProvider.forkChat?.({
+        workspaceId: host.doc.collection.id,
+        docId: host.doc.id,
+        sessionId: parentSessionId,
+        latestMessageId: messageId,
+      });
+
+      if (!newSessionId) {
+        return false;
+      }
+
+      // Construct chat block messages from the forked chat session
+      const messages = await constructChatBlockMessages(host.doc, newSessionId);
+
+      // After switching to edgeless mode, the user can save the chat to a block
+      const blockId = addAIChatBlock(
+        host.doc,
+        messages,
+        newSessionId,
+        viewportCenter
+      );
+      if (!blockId) {
+        return false;
+      }
+
+      return true;
+    } catch (err) {
+      console.error(err);
+      notificationService?.notify({
+        title: 'Failed to save chat to a block',
+        accent: 'error',
+        onClose: function (): void {},
+      });
+      return false;
+    }
+  },
+};
+
 export const PageEditorActions = [
  ...CommonActions,
  {
    icon: CreateIcon,
    title: 'Create as a doc',
+    showWhen: () => true,
    toast: 'New doc created',
    handler: (host: EditorHost, content: string) => {
      reportResponse('result:add-page');
@ -128,6 +348,7 @@ export const PageEditorActions = [
      return true;
    },
  },
+  SAVE_CHAT_TO_BLOCK_ACTION,
 ];

 export const EdgelessEditorActions = [
@ -135,6 +356,7 @@ export const EdgelessEditorActions = [
  {
    icon: CreateIcon,
    title: 'Add to edgeless as note',
+    showWhen: () => true,
    toast: 'New note created',
    handler: async (host: EditorHost, content: string) => {
      reportResponse('result:add-note');
@ -166,4 +388,5 @@ export const EdgelessEditorActions = [
      return true;
    },
  },
+  SAVE_CHAT_TO_BLOCK_ACTION,
 ];
--- a/packages/frontend/core/src/blocksuite/presets/ai/chat-panel/actions/copy-more.ts
+++ b/packages/frontend/core/src/blocksuite/presets/ai/chat-panel/actions/copy-more.ts
@ -91,6 +91,9 @@ export class ChatCopyMore extends WithDisposable(LitElement) {
  @property({ attribute: false })
  accessor content!: string;

+  @property({ attribute: false })
+  accessor messageId!: string;
+
  @property({ attribute: false })
  accessor isLast!: boolean;

@ -182,7 +185,7 @@ export class ChatCopyMore extends WithDisposable(LitElement) {
  }

  override render() {
-    const { host, content, isLast } = this;
+    const { host, content, isLast, messageId } = this;
    return html`<style>
        .more-menu {
          padding: ${this._showMoreMenu ? '8px' : '0px'};
@ -218,16 +221,21 @@ export class ChatCopyMore extends WithDisposable(LitElement) {
      <div class="more-menu">
        ${this._showMoreMenu
          ? repeat(
-              PageEditorActions,
+              PageEditorActions.filter(action => action.showWhen(host)),
              action => action.title,
              action => {
+                const currentSelections = {
+                  text: this.curTextSelection,
+                  blocks: this.curBlockSelections,
+                };
                return html`<div
                  @click=${async () => {
                    const success = await action.handler(
                      host,
                      content,
-                      this.curTextSelection,
-                      this.curBlockSelections
+                      currentSelections,
+                      this.chatContextValue,
+                      messageId ?? undefined
                    );

                    if (success) {
--- a/packages/frontend/core/src/blocksuite/presets/ai/chat-panel/chat-context.ts
+++ b/packages/frontend/core/src/blocksuite/presets/ai/chat-panel/chat-context.ts
@ -1,6 +1,7 @@
 import type { AIError } from '@blocksuite/blocks';

 export type ChatMessage = {
+  id: string;
  content: string;
  role: 'user' | 'assistant';
  attachments?: string[];
@ -31,4 +32,11 @@ export type ChatContextValue = {
  markdown: string;
  images: File[];
  abortController: AbortController | null;
+  chatSessionId: string | null;
+};
+
+export type ChatBlockMessage = ChatMessage & {
+  userId?: string;
+  userName?: string;
+  avatarUrl?: string;
 };
--- a/packages/frontend/core/src/blocksuite/presets/ai/chat-panel/chat-panel-input.ts
+++ b/packages/frontend/core/src/blocksuite/presets/ai/chat-panel/chat-panel-input.ts
@ -428,16 +428,23 @@ export class ChatPanelInput extends WithDisposable(LitElement) {

    const content = (markdown ? `${markdown}\n` : '') + text;

+    // TODO: Should update message id especially for the assistant message
    this.updateContext({
      items: [
        ...this.chatContextValue.items,
        {
+          id: '',
          role: 'user',
          content: content,
          createdAt: new Date().toISOString(),
          attachments,
        },
-        { role: 'assistant', content: '', createdAt: new Date().toISOString() },
+        {
+          id: '',
+          role: 'assistant',
+          content: '',
+          createdAt: new Date().toISOString(),
+        },
      ],
    });

@ -466,6 +473,24 @@ export class ChatPanelInput extends WithDisposable(LitElement) {
        }

        this.updateContext({ status: 'success' });
+
+        if (!this.chatContextValue.chatSessionId) {
+          this.updateContext({
+            chatSessionId: AIProvider.LAST_ACTION_SESSIONID,
+          });
+        }
+
+        const { items } = this.chatContextValue;
+        const last = items[items.length - 1] as ChatMessage;
+        if (!last.id) {
+          const historyIds = await AIProvider.histories?.ids(
+            doc.collection.id,
+            doc.id,
+            { sessionId: this.chatContextValue.chatSessionId }
+          );
+          if (!historyIds || !historyIds[0]) return;
+          last.id = historyIds[0].messages.at(-1)?.id ?? '';
+        }
      }
    } catch (error) {
      this.updateContext({ status: 'error', error: error as AIError });
--- a/packages/frontend/core/src/blocksuite/presets/ai/chat-panel/chat-panel-messages.ts
+++ b/packages/frontend/core/src/blocksuite/presets/ai/chat-panel/chat-panel-messages.ts
@ -59,6 +59,10 @@ export class ChatPanelMessages extends WithDisposable(ShadowlessElement) {
    return this._selectionValue.filter(v => v.type === 'image');
  }

+  private get _rootService() {
+    return this.host.spec.getService('affine:page');
+  }
+
  static override styles = css`
    chat-panel-messages {
      position: relative;
@ -459,7 +463,7 @@ export class ChatPanelMessages extends WithDisposable(ShadowlessElement) {
      return nothing;

    const { host } = this;
-    const { content } = item;
+    const { content, id: messageId } = item;
    const actions = isInsidePageEditor(host)
      ? PageEditorActions
      : EdgelessEditorActions;
@ -505,6 +509,7 @@ export class ChatPanelMessages extends WithDisposable(ShadowlessElement) {
        .host=${host}
        .content=${content}
        .isLast=${isLast}
+        .messageId=${messageId}
        .curTextSelection=${this._currentTextSelection}
        .curBlockSelections=${this._currentBlockSelections}
        .chatContextValue=${this.chatContextValue}
@ -513,19 +518,21 @@ export class ChatPanelMessages extends WithDisposable(ShadowlessElement) {
      ${isLast
        ? html`<div class="actions-container">
            ${repeat(
-              actions.filter(action => {
-                if (!content) return false;
+              actions
+                .filter(action => action.showWhen(host))
+                .filter(action => {
+                  if (!content) return false;

-                if (
-                  action.title === 'Replace selection' &&
-                  (!this._currentTextSelection ||
-                    this._currentTextSelection.from.length === 0) &&
-                  this._currentBlockSelections?.length === 0
-                ) {
-                  return false;
-                }
-                return true;
-              }),
+                  if (
+                    action.title === 'Replace selection' &&
+                    (!this._currentTextSelection ||
+                      this._currentTextSelection.from.length === 0) &&
+                    this._currentBlockSelections?.length === 0
+                  ) {
+                    return false;
+                  }
+                  return true;
+                }),
              action => action.title,
              action => {
                return html`<div class="action">
@ -545,17 +552,21 @@ export class ChatPanelMessages extends WithDisposable(ShadowlessElement) {
                        return;
                      }

+                      const currentSelections = {
+                        text: this._currentTextSelection,
+                        blocks: this._currentBlockSelections,
+                        images: this._currentImageSelections,
+                      };
+
                      const success = await action.handler(
                        host,
                        content,
-                        this._currentTextSelection,
-                        this._currentBlockSelections,
-                        this._currentImageSelections
+                        currentSelections,
+                        this.chatContextValue,
+                        messageId ?? undefined
                      );
-                      const rootService = host.spec.getService('affine:page');
-                      const { notificationService } = rootService;
                      if (success) {
-                        notificationService?.notify({
+                        this._rootService.notificationService?.notify({
                          title: action.toast,
                          accent: 'success',
                          onClose: function (): void {},
--- a/packages/frontend/core/src/blocksuite/presets/ai/chat-panel/index.ts
+++ b/packages/frontend/core/src/blocksuite/presets/ai/chat-panel/index.ts
@ -108,7 +108,7 @@ export class ChatPanel extends WithDisposable(ShadowlessElement) {
      const { doc } = this;

      const [histories, actions] = await Promise.all([
-        AIProvider.histories?.chats(doc.collection.id, doc.id),
+        AIProvider.histories?.chats(doc.collection.id, doc.id, { fork: false }),
        AIProvider.histories?.actions(doc.collection.id, doc.id),
      ]);

@ -116,9 +116,12 @@ export class ChatPanel extends WithDisposable(ShadowlessElement) {

      const items: ChatItem[] = actions ? [...actions] : [];

-      if (histories?.[0]) {
-        this._chatSessionId = histories[0].sessionId;
-        items.push(...histories[0].messages);
+      if (histories?.at(-1)) {
+        const history = histories.at(-1);
+        if (!history) return;
+        this._chatSessionId = history.sessionId;
+        this.chatContextValue.chatSessionId = history.sessionId;
+        items.push(...history.messages);
      }

      this.chatContextValue = {
@ -153,6 +156,7 @@ export class ChatPanel extends WithDisposable(ShadowlessElement) {
    status: 'idle',
    error: null,
    markdown: '',
+    chatSessionId: null,
  };

  private readonly _cleanupHistories = async () => {
--- a/packages/frontend/core/src/blocksuite/presets/ai/provider.ts
+++ b/packages/frontend/core/src/blocksuite/presets/ai/provider.ts
@ -70,6 +70,10 @@ export class AIProvider {
    return AIProvider.instance.toggleGeneralAIOnboarding;
  }

+  static get forkChat() {
+    return AIProvider.instance.forkChat;
+  }
+
  private static readonly instance = new AIProvider();

  static LAST_ACTION_SESSIONID = '';
@ -84,6 +88,12 @@ export class AIProvider {

  private toggleGeneralAIOnboarding: ((value: boolean) => void) | null = null;

+  private forkChat:
+    | ((
+        options: BlockSuitePresets.AIForkChatSessionOptions
+      ) => string | Promise<string>)
+    | null = null;
+
  private readonly slots = {
    // use case: when user selects "continue in chat" in an ask ai result panel
    // do we need to pass the context to the chat panel?
@ -236,6 +246,13 @@ export class AIProvider {
    engine: BlockSuitePresets.AIPhotoEngineService
  ): void;

+  static provide(
+    id: 'forkChat',
+    fn: (
+      options: BlockSuitePresets.AIForkChatSessionOptions
+    ) => string | Promise<string>
+  ): void;
+
  static provide(id: 'onboarding', fn: (value: boolean) => void): void;

  // actions:
@ -259,6 +276,10 @@ export class AIProvider {
      AIProvider.instance.toggleGeneralAIOnboarding = action as (
        value: boolean
      ) => void;
+    } else if (id === 'forkChat') {
+      AIProvider.instance.forkChat = action as (
+        options: BlockSuitePresets.AIForkChatSessionOptions
+      ) => string | Promise<string>;
    } else {
      // eslint-disable-next-line @typescript-eslint/no-explicit-any
      AIProvider.instance.provideAction(id as any, action as any);
--- a/packages/frontend/core/src/components/affine/setting-modal/general-setting/experimental-features/index.tsx
+++ b/packages/frontend/core/src/components/affine/setting-modal/general-setting/experimental-features/index.tsx
@ -128,6 +128,7 @@ const blocksuiteFeatureFlags: Partial<Record<keyof BlockSuiteFlags, string>> = {
  enable_database_statistics: 'Enable Database Block Statistics',
  enable_block_query: 'Enable Todo Block Query',
  enable_ai_onboarding: 'Enable AI Onboarding',
+  enable_ai_chat_block: 'Enable AI Chat Block',
 };

 const BlocksuiteFeatureFlagSettings = () => {
--- a/packages/frontend/core/src/components/blocksuite/block-suite-editor/ai/copilot-client.ts
+++ b/packages/frontend/core/src/components/blocksuite/block-suite-editor/ai/copilot-client.ts
@ -4,8 +4,10 @@ import {
  createCopilotMessageMutation,
  createCopilotSessionMutation,
  fetcher as defaultFetcher,
+  forkCopilotSessionMutation,
  getBaseUrl,
  getCopilotHistoriesQuery,
+  getCopilotHistoryIdsQuery,
  getCopilotSessionsQuery,
  GraphQLError,
  type GraphQLQuery,
@ -76,6 +78,16 @@ export class CopilotClient {
    return res.createCopilotSession;
  }

+  async forkSession(options: OptionsField<typeof forkCopilotSessionMutation>) {
+    const res = await fetcher({
+      query: forkCopilotSessionMutation,
+      variables: {
+        options,
+      },
+    });
+    return res.forkCopilotSession;
+  }
+
  async createMessage(
    options: OptionsField<typeof createCopilotMessageMutation>
  ) {
@ -117,6 +129,25 @@ export class CopilotClient {
    return res.currentUser?.copilot?.histories;
  }

+  async getHistoryIds(
+    workspaceId: string,
+    docId?: string,
+    options?: RequestOptions<
+      typeof getCopilotHistoriesQuery
+    >['variables']['options']
+  ) {
+    const res = await fetcher({
+      query: getCopilotHistoryIdsQuery,
+      variables: {
+        workspaceId,
+        docId,
+        options,
+      },
+    });
+
+    return res.currentUser?.copilot?.histories;
+  }
+
  async cleanupSessions(input: {
    workspaceId: string;
    docId: string;
--- a/packages/frontend/core/src/components/blocksuite/block-suite-editor/ai/request.ts
+++ b/packages/frontend/core/src/components/blocksuite/block-suite-editor/ai/request.ts
@ -1,4 +1,5 @@
 import { AIProvider } from '@affine/core/blocksuite/presets/ai';
+import type { ForkChatSessionInput } from '@affine/graphql';
 import { assertExists } from '@blocksuite/global/utils';
 import { partition } from 'lodash-es';

@ -44,6 +45,10 @@ export function createChatSession({
  });
 }

+export function forkCopilotSession(forkChatSessionInput: ForkChatSessionInput) {
+  return client.forkSession(forkChatSessionInput);
+}
+
 async function resizeImage(blob: Blob | File): Promise<Blob | null> {
  let src = '';
  try {
@ -256,6 +261,8 @@ export function textToText({

 export const listHistories = client.getHistories;

+export const listHistoryIds = client.getHistoryIds;
+
 // Only one image is currently being processed
 export function toImage({
  docId,
--- a/packages/frontend/core/src/components/blocksuite/block-suite-editor/ai/setup-provider.tsx
+++ b/packages/frontend/core/src/components/blocksuite/block-suite-editor/ai/setup-provider.tsx
@ -3,7 +3,11 @@ import { authAtom, openSettingModalAtom } from '@affine/core/atoms';
 import { AIProvider } from '@affine/core/blocksuite/presets/ai';
 import { toggleGeneralAIOnboarding } from '@affine/core/components/affine/ai-onboarding/apis';
 import { mixpanel } from '@affine/core/utils';
-import { getBaseUrl } from '@affine/graphql';
+import {
+  getBaseUrl,
+  type getCopilotHistoriesQuery,
+  type RequestOptions,
+} from '@affine/graphql';
 import { Trans } from '@affine/i18n';
 import { UnauthorizedError } from '@blocksuite/blocks';
 import { assertExists } from '@blocksuite/global/utils';
@ -14,6 +18,7 @@ import type { PromptKey } from './prompt';
 import {
  cleanupSessions,
  createChatSession,
+  forkCopilotSession,
  listHistories,
  textToText,
  toImage,
@ -404,10 +409,13 @@ Could you make a new website based on these notes and send back just the html fi
    },
    chats: async (
      workspaceId: string,
-      docId?: string
+      docId?: string,
+      options?: RequestOptions<
+        typeof getCopilotHistoriesQuery
+      >['variables']['options']
    ): Promise<BlockSuitePresets.AIHistory[]> => {
      // @ts-expect-error - 'action' is missing in server impl
-      return (await listHistories(workspaceId, docId)) ?? [];
+      return (await listHistories(workspaceId, docId, options)) ?? [];
    },
    cleanup: async (
      workspaceId: string,
@ -416,6 +424,16 @@ Could you make a new website based on these notes and send back just the html fi
    ) => {
      await cleanupSessions({ workspaceId, docId, sessionIds });
    },
+    ids: async (
+      workspaceId: string,
+      docId?: string,
+      options?: RequestOptions<
+        typeof getCopilotHistoriesQuery
+      >['variables']['options']
+    ): Promise<BlockSuitePresets.AIHistoryIds[]> => {
+      // @ts-expect-error - 'role' is missing type in server impl
+      return await listHistories(workspaceId, docId, options);
+    },
  });

  AIProvider.provide('photoEngine', {
@ -443,6 +461,10 @@ Could you make a new website based on these notes and send back just the html fi

  AIProvider.provide('onboarding', toggleGeneralAIOnboarding);

+  AIProvider.provide('forkChat', options => {
+    return forkCopilotSession(options);
+  });
+
  AIProvider.slots.requestUpgradePlan.on(() => {
    getCurrentStore().set(openSettingModalAtom, {
      activeTab: 'billing',
--- a/packages/frontend/core/src/components/blocksuite/block-suite-editor/specs/common.ts
+++ b/packages/frontend/core/src/components/blocksuite/block-suite-editor/specs/common.ts
@ -19,6 +19,7 @@ import {
  ListBlockSpec,
  NoteBlockSpec,
 } from '@blocksuite/blocks';
+import { AIChatBlockSpec, EdgelessAIChatBlockSpec } from '@blocksuite/presets';

 import { CustomAttachmentBlockSpec } from './custom/attachment-block';

@ -41,4 +42,6 @@ export const CommonBlockSpecs: BlockSpec[] = [
  AICodeBlockSpec,
  AIImageBlockSpec,
  AIParagraphBlockSpec,
+  AIChatBlockSpec,
+  EdgelessAIChatBlockSpec,
 ];
--- a/scripts/setup/vi-mock.ts
+++ b/scripts/setup/vi-mock.ts
@ -11,6 +11,17 @@ vi.mock('@blocksuite/presets', () => ({
  DocTitle: vi.fn(),
  EdgelessEditor: vi.fn(),
  PageEditor: vi.fn(),
+  AIChatBlockSchema: {
+    version: 1,
+    model: {
+      version: 1,
+      flavour: 'affine:embed-ai-chat',
+      role: 'content',
+      children: [],
+    },
+  },
+  AIChatBlockSpec: {},
+  EdgelessAIChatBlockSpec: {},
 }));

 vi.mock('@blocksuite/presets/ai', () => ({