feat(core): send message with extracted info

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

@@ -36,7 +36,16 @@ export interface CollectionChip extends BaseChip {
   collectionId: string;
 }
 
-export type ChatChip = DocChip | FileChip | TagChip | CollectionChip;
+export interface SelectedContextChip extends FileChip {
+  isSelectedContext: true;
+}
+
+export type ChatChip =
+  | DocChip
+  | FileChip
+  | TagChip
+  | CollectionChip
+  | SelectedContextChip;
 
 export interface DocDisplayConfig {
   getIcon: (docId: string) => any;

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

@@ -8,6 +8,7 @@ import type {
   CollectionChip,
   DocChip,
   FileChip,
+  SelectedContextChip,
   TagChip,
 } from './type';
 
@@ -62,6 +63,12 @@ export function isCollectionChip(chip: ChatChip): chip is CollectionChip {
   return 'collectionId' in chip;
 }
 
+export function isSelectedContextChip(
+  chip: ChatChip
+): chip is SelectedContextChip {
+  return 'isSelectedContext' in chip && chip.isSelectedContext;
+}
+
 export function getChipKey(chip: ChatChip) {
   if (isDocChip(chip)) {
     return chip.docId;

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

@@ -32,6 +32,7 @@ import {
   isCollectionChip,
   isDocChip,
   isFileChip,
+  isSelectedContextChip,
   isTagChip,
   omitChip,
 } from '../ai-chat-chips';
@@ -157,6 +158,8 @@ export class AIChatComposer extends SignalWatcher(
         .session=${this.session}
         .chips=${this.chips}
         .addChip=${this.addChip}
+        .removeSelectedContextChips=${this.removeSelectedContextChips} 
+        .waitForSelectedContextChipsFinished=${this.waitForSelectedContextChipsFinished}
         .addImages=${this.addImages}
         .createSession=${this.createSession}
         .chatContextValue=${this.chatContextValue}
@@ -329,12 +332,17 @@ export class AIChatComposer extends SignalWatcher(
     ]);
   };
 
-  private readonly addChip = async (chip: ChatChip) => {
+  private readonly addChip = async (
+    chip: ChatChip,
+    silent: boolean = false
+  ) => {
     this.isChipsCollapsed = false;
     // if already exists
     const index = findChipIndex(this.chips, chip);
     if (index !== -1) {
-      this.notificationService.toast('chip already exists');
+      if (!silent) {
+        this.notificationService.toast('chip already exists');
+      }
       return;
     }
     this.updateChips([...this.chips, chip]);
@@ -348,6 +356,15 @@ export class AIChatComposer extends SignalWatcher(
     await this.removeFromContext(chip);
   };
 
+  private readonly removeSelectedContextChips = async () => {
+    const selectedContextChips = this.chips.filter(c =>
+      isSelectedContextChip(c)
+    );
+    for (const chip of selectedContextChips) {
+      await this.removeChip(chip);
+    }
+  };
+
   private readonly addToContext = async (chip: ChatChip) => {
     if (isDocChip(chip)) {
       return await this.addDocToContext(chip);
@@ -639,4 +656,27 @@ export class AIChatComposer extends SignalWatcher(
     }
     await this.pollEmbeddingStatus();
   };
+
+  private readonly waitForSelectedContextChipsFinished = async (
+    timeout = 10000,
+    interval = 500
+  ): Promise<void> => {
+    const start = Date.now();
+    return new Promise((resolve, reject) => {
+      const check = () => {
+        const selectedChips = this.chips.filter(c => isSelectedContextChip(c));
+        const allFinished = selectedChips.every(c => c.state === 'finished');
+        if (allFinished) {
+          resolve();
+        } else if (Date.now() - start >= timeout) {
+          reject(
+            new Error('Timeout waiting for selected context chips to finish')
+          );
+        } else {
+          setTimeout(check, interval);
+        }
+      };
+      check();
+    });
+  };
 }

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

@@ -45,6 +45,9 @@ const DEFAULT_CHAT_CONTEXT_VALUE: ChatContextValue = {
   status: 'idle',
   error: null,
   markdown: '',
+  attachments: [],
+  snapshot: null,
+  markdownFile: null,
 };
 
 export class AIChatContent extends SignalWatcher(

packages/frontend/core/src/blocksuite/ai/components/ai-chat-content/type.ts

@@ -12,5 +12,11 @@ export type ChatContextValue = {
   markdown: string;
   // images of the selected content or user uploaded
   images: File[];
+  // snapshot of the selected content
+  snapshot: File | null;
+  // attachments of the selected content
+  attachments: File[];
+  // markdown file of the selected content
+  markdownFile: File | null;
   abortController: AbortController | null;
 };

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

@@ -338,7 +338,13 @@ export class AIChatInput extends SignalWatcher(
   accessor addImages!: (images: File[]) => void;
 
   @property({ attribute: false })
-  accessor addChip!: (chip: ChatChip) => Promise<void>;
+  accessor addChip!: (chip: ChatChip, silent?: boolean) => Promise<void>;
+
+  @property({ attribute: false })
+  accessor removeSelectedContextChips!: () => Promise<void>;
+
+  @property({ attribute: false })
+  accessor waitForSelectedContextChipsFinished!: () => Promise<void>;
 
   @property({ attribute: false })
   accessor networkSearchConfig!: AINetworkSearchConfig;
@@ -603,9 +609,47 @@ export class AIChatInput extends SignalWatcher(
     this.modelId = modelId;
   };
 
+  private readonly addSelectedContextChips = async () => {
+    const { snapshot, markdownFile, attachments } = this.chatContextValue;
+    await this.removeSelectedContextChips();
+    for (const attachment of attachments) {
+      await this.addChip(
+        {
+          file: attachment,
+          state: 'processing',
+          isSelectedContext: true,
+        },
+        true
+      );
+      if (snapshot) {
+        await this.addChip(
+          {
+            file: snapshot,
+            state: 'processing',
+            isSelectedContext: true,
+          },
+          true
+        );
+      }
+      if (markdownFile) {
+        await this.addChip(
+          {
+            file: markdownFile,
+            state: 'processing',
+            isSelectedContext: true,
+          },
+          true
+        );
+      }
+    }
+    await this.waitForSelectedContextChipsFinished();
+  };
+
   send = async (text: string) => {
     try {
       const { status, markdown, images } = this.chatContextValue;
+      await this.addSelectedContextChips();
+
       if (status === 'loading' || status === 'transmitting') return;
       if (!text) return;
       if (!AIProvider.actions.chat) return;
@@ -620,13 +664,13 @@ export class AIChatInput extends SignalWatcher(
         abortController,
       });
 
-      const attachments = await Promise.all(
+      const imageAttachments = await Promise.all(
         images?.map(image => readBlobAsURL(image))
       );
       const userInput = (markdown ? `${markdown}\n` : '') + text;
 
       // optimistic update messages
-      await this._preUpdateMessages(userInput, attachments);
+      await this._preUpdateMessages(userInput, imageAttachments);
 
       const sessionId = (await this.createSession())?.sessionId;
       let contexts = await this._getMatchedContexts();

packages/frontend/core/src/blocksuite/ai/components/ai-chat-input/type.ts

@@ -26,5 +26,8 @@ export type AIChatInputContext = {
   quote?: string;
   markdown?: string;
   images: File[];
+  attachments: File[];
+  snapshot: File | null;
+  markdownFile: File | null;
   abortController: AbortController | null;
 };

packages/frontend/core/src/blocksuite/ai/utils/attachment.ts

@@ -0,0 +1,9 @@
+import { AttachmentBlockModel } from '@blocksuite/affine/model';
+import type { BlockModel } from '@blocksuite/affine/store';
+import type { GfxModel } from '@blocksuite/std/gfx';
+
+export function isAttachment(
+  model: GfxModel | BlockModel
+): model is AttachmentBlockModel {
+  return model instanceof AttachmentBlockModel;
+}

packages/frontend/core/src/blocksuite/ai/utils/extract.ts

@@ -1,7 +1,8 @@
+import { WorkspaceImpl } from '@affine/core/modules/workspace/impls/workspace';
 import {
   DatabaseBlockModel,
   ImageBlockModel,
-  type NoteBlockModel,
+  NoteBlockModel,
   NoteDisplayMode,
 } from '@blocksuite/affine/model';
 import {
@@ -15,16 +16,24 @@ import {
 } from '@blocksuite/affine/shared/commands';
 import { DocModeProvider } from '@blocksuite/affine/shared/services';
 import {
+  getBlockProps,
   isInsideEdgelessEditor,
   matchModels,
 } from '@blocksuite/affine/shared/utils';
 import { BlockStdScope, type EditorHost } from '@blocksuite/affine/std';
-import type { BlockModel, Store } from '@blocksuite/affine/store';
+import {
+  GfxBlockElementModel,
+  GfxControllerIdentifier,
+} from '@blocksuite/affine/std/gfx';
+import type { BlockModel, DocSnapshot, Store } from '@blocksuite/affine/store';
 import { Slice, toDraftModel } from '@blocksuite/affine/store';
+import { Doc as YDoc } from 'yjs';
 
 import { getStoreManager } from '../../manager/store';
 import type { ChatContextValue } from '../components/ai-chat-content';
+import { isAttachment } from './attachment';
 import {
+  getSelectedAttachmentsAsBlobs,
   getSelectedImagesAsBlobs,
   getSelectedTextContent,
   selectedToCanvas,
@@ -46,6 +55,67 @@ async function extractEdgelessSelected(
   host: EditorHost
 ): Promise<Partial<ChatContextValue> | null> {
   if (!isInsideEdgelessEditor(host)) return null;
+  const gfx = host.std.get(GfxControllerIdentifier);
+  const selectedElements = gfx.selection.selectedElements;
+
+  let snapshot: DocSnapshot | null = null;
+  let markdown = '';
+  const attachments: ChatContextValue['attachments'] = [];
+
+  if (selectedElements.length) {
+    const transformer = host.store.getTransformer();
+    const markdownAdapter = new MarkdownAdapter(
+      transformer,
+      host.store.provider
+    );
+    const collection = new WorkspaceImpl({
+      id: 'AI_EXTRACT',
+      rootDoc: new YDoc({ guid: 'AI_EXTRACT' }),
+    });
+    collection.meta.initialize();
+
+    let needSnapshot = false;
+    let needMarkdown = false;
+    try {
+      const fragmentDoc = collection.createDoc();
+      const fragment = fragmentDoc.getStore();
+      fragmentDoc.load();
+
+      const rootId = fragment.addBlock('affine:page');
+      const surfaceId = fragment.addBlock('affine:surface', {}, rootId);
+      const noteId = fragment.addBlock('affine:note', {}, rootId);
+      for (const element of selectedElements) {
+        if (element instanceof NoteBlockModel) {
+          needMarkdown = true;
+          for (const child of element.children) {
+            const props = getBlockProps(child);
+            fragment.addBlock(child.flavour, props, noteId);
+          }
+        } else if (isAttachment(element)) {
+          const { name, sourceId } = element.props;
+          if (name && sourceId) {
+            const blob = await host.store.blobSync.get(sourceId);
+            if (blob) {
+              attachments.push(new File([blob], name));
+            }
+          }
+        } else if (element instanceof GfxBlockElementModel) {
+          const props = getBlockProps(element);
+          needSnapshot = true;
+          fragment.addBlock(element.flavour, props, surfaceId);
+        }
+      }
+
+      if (needSnapshot) {
+        snapshot = transformer.docToSnapshot(fragment) ?? null;
+      }
+      if (needMarkdown) {
+        markdown = (await markdownAdapter.fromDoc(fragment))?.file ?? '';
+      }
+    } finally {
+      collection.dispose();
+    }
+  }
 
   const canvas = await selectedToCanvas(host);
   if (!canvas) return null;
@@ -57,6 +127,11 @@ async function extractEdgelessSelected(
 
   return {
     images: [new File([blob], 'selected.png')],
+    snapshot: snapshot
+      ? new File([JSON.stringify(snapshot)], 'selected.json')
+      : null,
+    markdownFile: markdown.length ? new File([markdown], 'selected.md') : null,
+    attachments,
   };
 }
 
@@ -65,6 +140,7 @@ async function extractPageSelected(
 ): Promise<Partial<ChatContextValue> | null> {
   const text = await getSelectedTextContent(host, 'plain-text');
   const images = await getSelectedImagesAsBlobs(host);
+  const attachments = await getSelectedAttachmentsAsBlobs(host);
   const hasText = text.length > 0;
   const hasImages = images.length > 0;
 
@@ -73,6 +149,7 @@ async function extractPageSelected(
     return {
       quote: text,
       markdown: markdown,
+      attachments,
     };
   } else if (!hasText && hasImages && images.length === 1) {
     host.command
@@ -83,6 +160,7 @@ async function extractPageSelected(
       })
       .run();
     return {
+      attachments,
       images,
     };
   } else {
@@ -92,6 +170,7 @@ async function extractPageSelected(
       quote: text,
       markdown,
       images,
+      attachments,
     };
   }
 }

packages/frontend/core/src/blocksuite/ai/utils/selection-utils.ts

@@ -29,6 +29,7 @@ import {
 
 import { getContentFromSlice } from '../../utils';
 import type { CopilotTool } from '../tool/copilot-tool';
+import { isAttachment } from './attachment';
 import { getEdgelessCopilotWidget } from './get-edgeless-copilot-widget';
 
 export async function selectedToCanvas(host: EditorHost) {
@@ -232,6 +233,34 @@ export const getSelectedImagesAsBlobs = async (host: EditorHost) => {
   return blobs.filter((blob): blob is File => !!blob);
 };
 
+export const getSelectedAttachmentsAsBlobs = async (host: EditorHost) => {
+  const [_, data] = host.command.exec(getSelectedBlocksCommand, {
+    types: ['block'],
+  });
+
+  const blocks = data.selectedBlocks ?? [];
+  const attachments: { sourceId: string; name: string }[] = [];
+
+  for (const block of blocks) {
+    if (isAttachment(block.model)) {
+      const { sourceId, name } = block.model.props;
+      if (sourceId && name) {
+        attachments.push({ sourceId, name });
+      }
+    }
+  }
+
+  const blobs = await Promise.all(
+    attachments.map(attachment => {
+      return host.store.blobSync.get(attachment.sourceId);
+    })
+  );
+
+  return blobs
+    .filter(blob => blob)
+    .map((blob, index) => new File([blob as Blob], attachments[index].name));
+};
+
 export const getSelectedNoteAnchor = (host: EditorHost, id: string) => {
   return host.querySelector(`affine-edgeless-note[data-block-id="${id}"]`);
 };