From d00879acebf3d1e28ccebd8a6e7c08f05e7c2b85 Mon Sep 17 00:00:00 2001
From: fundon <fundon@pindash.io>
Date: Mon, 15 Apr 2024 09:32:21 +0000
Subject: [PATCH] feat: ai create an image (#6538)

Related to https://github.com/toeverything/blocksuite/pull/6746
---
 .../block-suite-editor/ai/provider.ts         | 17 ++++-
 .../block-suite-editor/ai/request.ts          | 76 +++++++++++++++----
 2 files changed, 76 insertions(+), 17 deletions(-)

diff --git a/packages/frontend/core/src/components/blocksuite/block-suite-editor/ai/provider.ts b/packages/frontend/core/src/components/blocksuite/block-suite-editor/ai/provider.ts
index c610d4ccb4..defbc1dea6 100644
--- a/packages/frontend/core/src/components/blocksuite/block-suite-editor/ai/provider.ts
+++ b/packages/frontend/core/src/components/blocksuite/block-suite-editor/ai/provider.ts
@@ -1,7 +1,12 @@
 import { assertExists } from '@blocksuite/global/utils';
 import { AIProvider } from '@blocksuite/presets';
 
-import { createChatSession, listHistories, textToText } from './request';
+import {
+  createChatSession,
+  listHistories,
+  textToText,
+  toImage,
+} from './request';
 
 export function setupAIProvider() {
   // a single workspace should have only a single chat session
@@ -235,4 +240,14 @@ export function setupAIProvider() {
       return (await listHistories(workspaceId, docId)) ?? [];
     },
   });
+
+  AIProvider.provide('createImage', options => {
+    // const promptName = 'debug:action:fal-sd15';
+    const promptName = 'debug:action:dalle3';
+    return toImage({
+      ...options,
+      promptName,
+      forceToImage: true,
+    });
+  });
 }
diff --git 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
index 1e4541b7ab..aaee0c7f5e 100644
--- 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
@@ -32,6 +32,7 @@ export type TextToTextOptions = {
   params?: Record<string, string>;
   timeout?: number;
   stream?: boolean;
+  forceToImage?: boolean; // force to image
 };
 
 export function createChatSession({
@@ -56,33 +57,43 @@ async function createSessionMessage({
   sessionId: providedSessionId,
   attachments,
   params,
+  forceToImage,
 }: TextToTextOptions) {
-  const hasAttachments = attachments && attachments.length > 0;
   if (!promptName && !providedSessionId) {
     throw new Error('promptName or sessionId is required');
   }
+  const hasAttachments = attachments && attachments.length > 0;
   const sessionId = await (providedSessionId ??
     client.createSession({
       workspaceId,
       docId,
       promptName: promptName as string,
     }));
-  if (hasAttachments) {
-    const normalizedAttachments = await Promise.all(
-      attachments.map(async attachment => {
-        if (typeof attachment === 'string') {
-          return attachment;
-        }
-        const url = await readBlobAsURL(attachment);
-        return url;
-      })
-    );
-    const messageId = await client.createMessage({
-      sessionId: sessionId,
+
+  if (forceToImage || hasAttachments) {
+    const options = {
+      sessionId,
       content,
-      attachments: normalizedAttachments,
       params,
-    });
+    } as {
+      sessionId: string;
+      content?: string;
+      params?: Record<string, string>;
+      attachments?: string[];
+    };
+    if (hasAttachments) {
+      const normalizedAttachments = await Promise.all(
+        attachments.map(async attachment => {
+          if (typeof attachment === 'string') {
+            return attachment;
+          }
+          const url = await readBlobAsURL(attachment);
+          return url;
+        })
+      );
+      options.attachments = normalizedAttachments;
+    }
+    const messageId = await client.createMessage(options);
     return {
       messageId,
       sessionId,
@@ -126,7 +137,7 @@ export function textToText({
           messageId: message.messageId,
           message: message.message,
         });
-        yield* toTextStream(eventSource, { timeout: timeout });
+        yield* toTextStream(eventSource, { timeout });
       },
     };
   } else {
@@ -158,3 +169,36 @@ export function textToText({
 }
 
 export const listHistories = client.getHistories;
+
+// Only one image is currently being processed
+export function toImage({
+  docId,
+  workspaceId,
+  promptName,
+  content,
+  attachments,
+  params,
+  forceToImage,
+  timeout = TIMEOUT,
+}: TextToTextOptions) {
+  return {
+    [Symbol.asyncIterator]: async function* () {
+      const { messageId, sessionId } = await createSessionMessage({
+        docId,
+        workspaceId,
+        promptName,
+        content,
+        attachments,
+        params,
+        forceToImage,
+      });
+
+      const eventSource = client.imagesStream(
+        // @ts-expect-error: messageId should exist
+        messageId,
+        sessionId
+      );
+      yield* toTextStream(eventSource, { timeout, type: 'attachment' });
+    },
+  };
+}