From 16096978d86fa4767b108b915a1d0f185a3928f9 Mon Sep 17 00:00:00 2001
From: L-Sun <zover.v@gmail.com>
Date: Wed, 26 Jun 2024 04:05:19 +0000
Subject: [PATCH] fix(core): add synced doc content to edgeless ai context
 (#7327)

Closes: [BS-633](https://linear.app/affine-design/issue/BS-633/synced-block-%E7%99%BD%E6%9D%BF%E6%94%AF%E6%8C%81%EF%BC%9A%E4%BD%9C%E4%B8%BA%E4%B8%8A%E4%B8%8B%E6%96%87%E6%8F%90%E4%BE%9B%E7%BB%99-ai#comment-f336c6d5)
Related PR: https://github.com/toeverything/blocksuite/pull/7417

Read content from synced doc model as AI context in edgeless model.
---
 .../presets/ai/actions/edgeless-handler.ts    | 90 ++++++++++++-------
 1 file changed, 58 insertions(+), 32 deletions(-)

diff --git a/packages/frontend/core/src/blocksuite/presets/ai/actions/edgeless-handler.ts b/packages/frontend/core/src/blocksuite/presets/ai/actions/edgeless-handler.ts
index 59d1cdc597..238e4ad74a 100644
--- a/packages/frontend/core/src/blocksuite/presets/ai/actions/edgeless-handler.ts
+++ b/packages/frontend/core/src/blocksuite/presets/ai/actions/edgeless-handler.ts
@@ -8,6 +8,7 @@ import type {
 import {
   BlocksUtils,
   EdgelessTextBlockModel,
+  EmbedSyncedDocModel,
   ImageBlockModel,
   NoteBlockModel,
   ShapeElementModel,
@@ -96,6 +97,14 @@ function actionToRenderer<T extends keyof BlockSuitePresets.AIActions>(
   return createTextRenderer(host, { maxHeight: 320 });
 }
 
+async function getContentFromEmbedSyncedDocModel(
+  host: EditorHost,
+  models: EmbedSyncedDocModel[]
+) {
+  const slice = Slice.fromModels(host.doc, models);
+  return (await getContentFromSlice(host, slice)).trim();
+}
+
 async function getContentFromHubBlockModel(
   host: EditorHost,
   models: EdgelessTextBlockModel[] | NoteBlockModel[]
@@ -132,39 +141,55 @@ export async function getContentFromSelected(
     return el.caption !== undefined && el.caption.length !== 0;
   }
 
-  const { notes, texts, shapes, images, edgelessTexts } = selected.reduce<{
-    notes: NoteBlockModel[];
-    texts: TextElementModel[];
-    shapes: RemoveUndefinedKey<ShapeElementModel, 'text'>[];
-    images: RemoveUndefinedKey<ImageBlockModel, 'caption'>[];
-    edgelessTexts: EdgelessTextBlockModel[];
-  }>(
-    (pre, cur) => {
-      if (cur instanceof NoteBlockModel) {
-        pre.notes.push(cur);
-      } else if (cur instanceof TextElementModel) {
-        pre.texts.push(cur);
-      } else if (cur instanceof ShapeElementModel && isShapeWithText(cur)) {
-        pre.shapes.push(cur);
-      } else if (cur instanceof ImageBlockModel && isImageWithCaption(cur)) {
-        pre.images.push(cur);
-      } else if (cur instanceof EdgelessTextBlockModel) {
-        pre.edgelessTexts.push(cur);
-      }
+  const { notes, texts, shapes, images, edgelessTexts, embedSyncedDocs } =
+    selected.reduce<{
+      notes: NoteBlockModel[];
+      texts: TextElementModel[];
+      shapes: RemoveUndefinedKey<ShapeElementModel, 'text'>[];
+      images: RemoveUndefinedKey<ImageBlockModel, 'caption'>[];
+      edgelessTexts: EdgelessTextBlockModel[];
+      embedSyncedDocs: EmbedSyncedDocModel[];
+    }>(
+      (pre, cur) => {
+        if (cur instanceof NoteBlockModel) {
+          pre.notes.push(cur);
+        } else if (cur instanceof TextElementModel) {
+          pre.texts.push(cur);
+        } else if (cur instanceof ShapeElementModel && isShapeWithText(cur)) {
+          pre.shapes.push(cur);
+        } else if (cur instanceof ImageBlockModel && isImageWithCaption(cur)) {
+          pre.images.push(cur);
+        } else if (cur instanceof EdgelessTextBlockModel) {
+          pre.edgelessTexts.push(cur);
+        } else if (cur instanceof EmbedSyncedDocModel) {
+          pre.embedSyncedDocs.push(cur);
+        }
 
-      return pre;
-    },
-    { notes: [], texts: [], shapes: [], images: [], edgelessTexts: [] }
-  );
+        return pre;
+      },
+      {
+        notes: [],
+        texts: [],
+        shapes: [],
+        images: [],
+        edgelessTexts: [],
+        embedSyncedDocs: [],
+      }
+    );
 
   const noteContent = await getContentFromHubBlockModel(host, notes);
   const edgelessTextContent = await getContentFromHubBlockModel(
     host,
     edgelessTexts
   );
+  const syncedDocsContent = await getContentFromEmbedSyncedDocModel(
+    host,
+    embedSyncedDocs
+  );
 
   return `${noteContent.join('\n')}
-  ${edgelessTextContent.join('\n')}
+${edgelessTextContent.join('\n')}
+${syncedDocsContent}
 ${texts.map(text => text.text.toString()).join('\n')}
 ${shapes.map(shape => shape.text.toString()).join('\n')}
 ${images.map(image => image.caption.toString()).join('\n')}
@@ -449,14 +474,15 @@ export function actionToHandler<T extends keyof BlockSuitePresets.AIActions>(
     if (isCreateImageAction || isMakeItRealAction) {
       togglePanel = async () => {
         if (isEmpty) return true;
-        const {
-          notes,
-          shapes,
-          images,
-          edgelessTexts,
-          frames: _,
-        } = BlocksUtils.splitElements(selectedElements);
-        const blocks = [...notes, ...shapes, ...images, ...edgelessTexts];
+        const { notes, shapes, images, edgelessTexts, embedSyncedDocs } =
+          BlocksUtils.splitElements(selectedElements);
+        const blocks = [
+          ...notes,
+          ...shapes,
+          ...images,
+          ...edgelessTexts,
+          ...embedSyncedDocs,
+        ];
         if (blocks.length === 0) return true;
         const content = await getContentFromSelected(host, blocks);
         ctx.set({