From 9fe8d432dd82d389273fbe671a50307f869dbb0c Mon Sep 17 00:00:00 2001
From: yoyoyohamapi <472285740@qq.com>
Date: Tue, 22 Jul 2025 11:14:21 +0800
Subject: [PATCH] feat(core): extract attachments from page selected

---
 .../core/src/blocksuite/ai/utils/extract.ts   |  5 +++
 .../blocksuite/ai/utils/selection-utils.ts    | 34 ++++++++++++++++++-
 2 files changed, 38 insertions(+), 1 deletion(-)

diff --git a/packages/frontend/core/src/blocksuite/ai/utils/extract.ts b/packages/frontend/core/src/blocksuite/ai/utils/extract.ts
index c338c8cf6f..3f3ef39484 100644
--- a/packages/frontend/core/src/blocksuite/ai/utils/extract.ts
+++ b/packages/frontend/core/src/blocksuite/ai/utils/extract.ts
@@ -33,6 +33,7 @@ import { Doc as YDoc } from 'yjs';
 import { getStoreManager } from '../../manager/store';
 import type { ChatContextValue } from '../components/ai-chat-content';
 import {
+  getSelectedAttachmentsAsBlobs,
   getSelectedImagesAsBlobs,
   getSelectedTextContent,
   selectedToCanvas,
@@ -135,6 +136,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;
 
@@ -143,6 +145,7 @@ async function extractPageSelected(
     return {
       quote: text,
       markdown: markdown,
+      attachments,
     };
   } else if (!hasText && hasImages && images.length === 1) {
     host.command
@@ -153,6 +156,7 @@ async function extractPageSelected(
       })
       .run();
     return {
+      attachments,
       images,
     };
   } else {
@@ -162,6 +166,7 @@ async function extractPageSelected(
       quote: text,
       markdown,
       images,
+      attachments,
     };
   }
 }
diff --git a/packages/frontend/core/src/blocksuite/ai/utils/selection-utils.ts b/packages/frontend/core/src/blocksuite/ai/utils/selection-utils.ts
index 2ae7670d36..b782affe0e 100644
--- a/packages/frontend/core/src/blocksuite/ai/utils/selection-utils.ts
+++ b/packages/frontend/core/src/blocksuite/ai/utils/selection-utils.ts
@@ -6,7 +6,11 @@ import {
   getSurfaceBlock,
   type SurfaceBlockComponent,
 } from '@blocksuite/affine/blocks/surface';
-import { DatabaseBlockModel, ImageBlockModel } from '@blocksuite/affine/model';
+import {
+  AttachmentBlockModel,
+  DatabaseBlockModel,
+  ImageBlockModel,
+} from '@blocksuite/affine/model';
 import {
   getBlockSelectionsCommand,
   getImageSelectionsCommand,
@@ -232,6 +236,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 (block.model instanceof AttachmentBlockModel) {
+      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}"]`);
 };