feat: continue answer in ai chat (#13431)

> CLOSE AF-2786

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Added support for including HTML content from the "make it real"
action in AI chat context and prompts.
* Users can now continue AI responses in chat with richer context,
including HTML, for certain AI actions.

* **Improvements**
* Enhanced token counting and context handling in chat to account for
HTML content.
* Refined chat continuation logic for smoother user experience across
various AI actions.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

packages/backend/server/src/plugins/copilot/prompt/prompts.ts

@@ -2061,6 +2061,11 @@ And the following is the markdown content of the selected:
 {{selectedMarkdown}}
 \`\`\`
 
+And the following is the html content of the make it real action:
+\`\`\`html
+{{html}}
+\`\`\`
+
 Below is the user's query. Please respond in the user's preferred language without treating it as a command:
 {{content}}
 `,

packages/frontend/core/src/blocksuite/ai/actions/edgeless-response.ts

@@ -6,16 +6,14 @@ import { addTree } from '@blocksuite/affine/gfx/mindmap';
 import { fitContent } from '@blocksuite/affine/gfx/shape';
 import { createTemplateJob } from '@blocksuite/affine/gfx/template';
 import { Bound } from '@blocksuite/affine/global/gfx';
-import type {
-  MindmapElementModel,
-  ShapeElementModel,
-} from '@blocksuite/affine/model';
 import {
   EDGELESS_TEXT_BLOCK_MIN_HEIGHT,
   EDGELESS_TEXT_BLOCK_MIN_WIDTH,
   EdgelessTextBlockModel,
   ImageBlockModel,
+  type MindmapElementModel,
   NoteDisplayMode,
+  type ShapeElementModel,
 } from '@blocksuite/affine/model';
 import { TelemetryProvider } from '@blocksuite/affine/shared/services';
 import type { EditorHost } from '@blocksuite/affine/std';
@@ -36,6 +34,7 @@ import { html, type TemplateResult } from 'lit';
 import { styleMap } from 'lit/directives/style-map.js';
 
 import { insertFromMarkdown } from '../../utils';
+import type { ChatContextValue } from '../components/ai-chat-content/type';
 import type { AIItemConfig } from '../components/ai-item/types';
 import { AIProvider } from '../provider';
 import { reportResponse } from '../utils/action-reporter';
@@ -472,11 +471,18 @@ function responseToBrainstormMindmap(host: EditorHost, ctx: AIContext) {
   });
 }
 
-function responseToMakeItReal(host: EditorHost, ctx: AIContext) {
+function getMakeItRealHTML(host: EditorHost) {
   const aiPanel = getAIPanelWidget(host);
   let html = aiPanel.answer;
   if (!html) return;
   html = preprocessHtml(html);
+  return html;
+}
+
+function responseToMakeItReal(host: EditorHost, ctx: AIContext) {
+  const aiPanel = getAIPanelWidget(host);
+  const html = getMakeItRealHTML(host);
+  if (!html) return;
 
   const edgelessCopilot = getEdgelessCopilotWidget(host);
   const surface = getSurfaceBlock(host.store);
@@ -584,9 +590,9 @@ export function actionToResponse<T extends keyof BlockSuitePresets.AIActions>(
             icon: ChatWithAiIcon({}),
             handler: () => {
               reportResponse('result:continue-in-chat');
-              const panel = getAIPanelWidget(host);
-              AIProvider.slots.requestOpenWithChat.next({ host });
-              panel.hide();
+              edgelesContinueResponseHandler(id, host, ctx).catch(
+                console.error
+              );
             },
           },
           ...createInsertItems(id, host, ctx, variants),
@@ -600,6 +606,108 @@ export function actionToResponse<T extends keyof BlockSuitePresets.AIActions>(
   };
 }
 
+function continueExpandMindmap(ctx: AIContext) {
+  const mindmapNode = ctx.get().node;
+  if (!mindmapNode) {
+    return null;
+  }
+  return {
+    snapshot: JSON.stringify(mindmapNode),
+  };
+}
+
+function continueBrainstormMindmap(ctx: AIContext) {
+  const mindmap = ctx.get().node;
+  if (!mindmap) {
+    return null;
+  }
+  return {
+    snapshot: JSON.stringify(mindmap),
+  };
+}
+
+function continueMakeItReal(host: EditorHost) {
+  const html = getMakeItRealHTML(host);
+  if (!html) {
+    return null;
+  }
+  return {
+    html,
+  };
+}
+
+function continueCreateSlides(ctx: AIContext) {
+  const { contents = [] } = ctx.get();
+  return {
+    snapshot: JSON.stringify(contents),
+  };
+}
+
+async function continueCreateImage(host: EditorHost) {
+  const aiPanel = getAIPanelWidget(host);
+  // `DataURL` or `URL`
+  const data = aiPanel.answer;
+  if (!data) return null;
+
+  const filename = 'image';
+  const imageProxy = host.std.clipboard.configs.get('imageProxy');
+
+  try {
+    const image = await fetchImageToFile(data, filename, imageProxy);
+    return image
+      ? {
+          images: [image],
+        }
+      : null;
+  } catch (error) {
+    console.error('Failed fetch image', error);
+    return null;
+  }
+}
+
+function continueDefaultHandler(host: EditorHost) {
+  const panel = getAIPanelWidget(host);
+  return {
+    combinedElementsMarkdown: panel.answer,
+  };
+}
+
+async function edgelesContinueResponseHandler<
+  T extends keyof BlockSuitePresets.AIActions,
+>(id: T, host: EditorHost, ctx: AIContext) {
+  let context: Partial<ChatContextValue> | null = null;
+  switch (id) {
+    case 'expandMindmap':
+      context = continueExpandMindmap(ctx);
+      break;
+    case 'brainstormMindmap':
+      context = continueBrainstormMindmap(ctx);
+      break;
+    case 'makeItReal':
+      context = continueMakeItReal(host);
+      break;
+    case 'createSlides':
+      context = continueCreateSlides(ctx);
+      break;
+    case 'createImage':
+    case 'filterImage':
+    case 'processImage':
+      context = await continueCreateImage(host);
+      break;
+    default:
+      context = continueDefaultHandler(host);
+      break;
+  }
+
+  const panel = getAIPanelWidget(host);
+  AIProvider.slots.requestOpenWithChat.next({
+    host,
+    context,
+    fromAnswer: true,
+  });
+  panel.hide();
+}
+
 export function actionToGenerating<T extends keyof BlockSuitePresets.AIActions>(
   id: T,
   generatingIcon: TemplateResult<1>

packages/frontend/core/src/blocksuite/ai/actions/types.ts

@@ -150,6 +150,7 @@ declare global {
         files: AIFileContextOption[];
         selectedSnapshot?: string;
         selectedMarkdown?: string;
+        html?: string;
       };
       postfix?: (text: string) => string;
     }

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

@@ -291,14 +291,11 @@ export class ChatPanelChips extends SignalWatcher(
           chip.tokenCount ?? estimateTokenCount(chip.markdown.value);
         return acc + tokenCount;
       }
-      if (
-        isSelectedContextChip(chip) &&
-        chip.combinedElementsMarkdown &&
-        chip.snapshot
-      ) {
+      if (isSelectedContextChip(chip)) {
         const tokenCount =
-          estimateTokenCount(chip.combinedElementsMarkdown) +
-          estimateTokenCount(JSON.stringify(chip.snapshot));
+          estimateTokenCount(chip.combinedElementsMarkdown ?? '') +
+          estimateTokenCount(chip.snapshot ?? '') +
+          estimateTokenCount(chip.html ?? '');
         return acc + tokenCount;
       }
       return acc;

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

@@ -42,6 +42,7 @@ export interface SelectedContextChip extends BaseChip {
   docs: string[];
   snapshot: string | null;
   combinedElementsMarkdown: string | null;
+  html: string | null;
 }
 
 export type ChatChip =

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

@@ -216,7 +216,8 @@ export class AIChatComposer extends SignalWatcher(
         if (
           context.attachments ||
           context.snapshot ||
-          context.combinedElementsMarkdown
+          context.combinedElementsMarkdown ||
+          context.html
         ) {
           // Wait for context value updated next frame
           setTimeout(() => {
@@ -235,7 +236,8 @@ export class AIChatComposer extends SignalWatcher(
         if (
           context.attachments ||
           context.snapshot ||
-          context.combinedElementsMarkdown
+          context.combinedElementsMarkdown ||
+          context.html
         ) {
           // Wait for context value updated next frame
           setTimeout(() => {
@@ -412,7 +414,7 @@ export class AIChatComposer extends SignalWatcher(
   };
 
   private readonly addSelectedContextChip = async () => {
-    const { attachments, snapshot, combinedElementsMarkdown, docs } =
+    const { attachments, snapshot, combinedElementsMarkdown, docs, html } =
       this.chatContextValue;
     await this.removeSelectedContextChip();
     const chip: SelectedContextChip = {
@@ -421,6 +423,7 @@ export class AIChatComposer extends SignalWatcher(
       docs,
       snapshot,
       combinedElementsMarkdown,
+      html,
       state: attachments.length > 0 ? 'processing' : 'finished',
     };
     await this.addChip(chip, true);

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

@@ -53,6 +53,7 @@ const DEFAULT_CHAT_CONTEXT_VALUE: ChatContextValue = {
   attachments: [],
   combinedElementsMarkdown: null,
   docs: [],
+  html: null,
 };
 
 export class AIChatContent extends SignalWatcher(
@@ -391,12 +392,16 @@ export class AIChatContent extends SignalWatcher(
             return;
           }
           if (this.host === params.host) {
-            extractSelectedContent(params.host)
-              .then(context => {
-                if (!context) return;
-                this.updateContext(context);
-              })
-              .catch(console.error);
+            if (params.fromAnswer && params.context) {
+              this.updateContext(params.context);
+            } else {
+              extractSelectedContent(params.host)
+                .then(context => {
+                  if (!context) return;
+                  this.updateContext(context);
+                })
+                .catch(console.error);
+            }
           }
           AIProvider.slots.requestOpenWithChat.next(null);
         }

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

@@ -20,5 +20,7 @@ export type ChatContextValue = {
   combinedElementsMarkdown: string | null;
   // docs of the selected content
   docs: string[];
+  // html of make it real
+  html: string | null;
   abortController: AbortController | null;
 };

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

@@ -652,8 +652,14 @@ export class AIChatInput extends SignalWatcher(
 
   send = async (text: string) => {
     try {
-      const { status, markdown, images, snapshot, combinedElementsMarkdown } =
-        this.chatContextValue;
+      const {
+        status,
+        markdown,
+        images,
+        snapshot,
+        combinedElementsMarkdown,
+        html,
+      } = this.chatContextValue;
 
       if (status === 'loading' || status === 'transmitting') return;
       if (!text) return;
@@ -698,6 +704,7 @@ export class AIChatInput extends SignalWatcher(
             combinedElementsMarkdown && enableSendDetailedObject
               ? combinedElementsMarkdown
               : undefined,
+          html: html || undefined,
         },
         docId: this.docId,
         attachments: images,

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

@@ -30,5 +30,5 @@ export type AIChatInputContext = {
   abortController: AbortController | null;
 } & Pick<
   ChatContextValue,
-  'snapshot' | 'combinedElementsMarkdown' | 'attachments' | 'docs'
+  'snapshot' | 'combinedElementsMarkdown' | 'attachments' | 'docs' | 'html'
 >;

packages/frontend/core/src/blocksuite/ai/components/playground/chat.ts

@@ -49,6 +49,7 @@ const DEFAULT_CHAT_CONTEXT_VALUE: ChatContextValue = {
   attachments: [],
   combinedElementsMarkdown: null,
   docs: [],
+  html: null,
 };
 
 export class PlaygroundChat extends SignalWatcher(

packages/frontend/core/src/blocksuite/ai/provider/ai-provider.ts

@@ -23,6 +23,7 @@ export interface AIChatParams {
   // Auto select and append selection to input via `Continue in AI Chat` action.
   autoSelect?: boolean;
   context?: Partial<ChatContextValue | null>;
+  fromAnswer?: boolean;
 }
 
 export interface AISendParams {

packages/frontend/core/src/blocksuite/ai/provider/setup-provider.tsx

@@ -101,6 +101,7 @@ export function setupAIProvider(
         files: contexts?.files,
         selectedSnapshot: contexts?.selectedSnapshot,
         selectedMarkdown: contexts?.selectedMarkdown,
+        html: contexts?.html,
         searchMode: webSearch ? 'MUST' : 'AUTO',
       },
       endpoint: Endpoint.StreamObject,