feat(server): refactor copilot (#14892)

#### PR Dependency Tree


* **PR #14892** 👈

This tree was auto-generated by
[Charcoal](https://github.com/danerwilliams/charcoal)

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

@@ -152,7 +152,6 @@ declare global {
         selectedMarkdown?: string;
         html?: string;
       };
-      postfix?: (text: string) => string;
     }
 
     interface TranslateOptions extends AITextActionOptions {

packages/frontend/core/src/blocksuite/ai/provider/copilot-client.spec.ts

@@ -0,0 +1,34 @@
+/**
+ * @vitest-environment happy-dom
+ */
+import { describe, expect, test, vi } from 'vitest';
+
+import { CopilotClient, Endpoint } from './copilot-client';
+
+describe('CopilotClient action streams', () => {
+  test('routes action endpoint outside the deprecated workflow path', () => {
+    const eventSource = vi.fn(
+      () =>
+        ({
+          close: vi.fn(),
+        }) as unknown as EventSource
+    );
+    const client = new CopilotClient(vi.fn(), eventSource);
+
+    client.chatTextStream(
+      {
+        sessionId: 'session-1',
+        messageId: 'message-1',
+        actionId: 'mindmap.generate',
+        actionVersion: 'v1',
+        retry: true,
+        runId: 'run-1',
+      },
+      Endpoint.Action
+    );
+
+    expect(eventSource).toHaveBeenCalledWith(
+      '/api/copilot/actions/session-1/stream?messageId=message-1&actionId=mindmap.generate&actionVersion=v1&runId=run-1&retry=true'
+    );
+  });
+});

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

@@ -42,8 +42,8 @@ import {
 } from './error';
 
 export enum Endpoint {
+  Action = 'action',
   StreamObject = 'stream-object',
-  Workflow = 'workflow',
   Images = 'images',
 }
 
@@ -469,21 +469,36 @@ export class CopilotClient {
       reasoning,
       modelId,
       toolsConfig,
+      actionId,
+      actionVersion,
+      runId,
+      retry,
     }: {
       sessionId: string;
       messageId?: string;
       reasoning?: boolean;
       modelId?: string;
       toolsConfig?: AIToolsConfig;
+      actionId?: string;
+      actionVersion?: string;
+      runId?: string;
+      retry?: boolean;
     },
     endpoint = Endpoint.StreamObject
   ) {
-    let url = `/api/copilot/chat/${sessionId}/${endpoint}`;
+    let url =
+      endpoint === Endpoint.Action
+        ? `/api/copilot/actions/${sessionId}/stream`
+        : `/api/copilot/chat/${sessionId}/${endpoint}`;
     const queryString = this.paramsToQueryString({
       messageId,
       reasoning,
       modelId,
       toolsConfig,
+      actionId,
+      actionVersion,
+      runId,
+      retry,
     });
     if (queryString) {
       url += `?${queryString}`;

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

@@ -39,9 +39,13 @@ export const promptKeys = [
   'Convert to sticker',
   'Upscale image',
   'Remove background',
-  // workflows
-  'workflow:presentation',
-  'workflow:brainstorm',
+  // actions
+  'mindmap.generate',
+  'slides.outline',
+  'image.filter.sketch',
+  'image.filter.clay',
+  'image.filter.anime',
+  'image.filter.pixel',
 ] as const;
 
 export type PromptKey = (typeof promptKeys)[number];

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

@@ -18,8 +18,10 @@ export type TextToTextOptions = {
   signal?: AbortSignal;
   retry?: boolean;
   endpoint?: Endpoint;
+  actionId?: string;
+  actionVersion?: string;
+  runId?: string;
   isRootSession?: boolean;
-  postfix?: (text: string) => string;
   reasoning?: boolean;
   modelId?: string;
   toolsConfig?: AIToolsConfig;
@@ -120,7 +122,9 @@ export function textToText({
   timeout = TIMEOUT,
   retry = false,
   endpoint = Endpoint.StreamObject,
-  postfix,
+  actionId,
+  actionVersion,
+  runId,
   reasoning,
   modelId,
   toolsConfig,
@@ -148,6 +152,10 @@ export function textToText({
             reasoning,
             modelId,
             toolsConfig,
+            actionId,
+            actionVersion,
+            runId,
+            retry,
           },
           endpoint
         );
@@ -166,25 +174,12 @@ export function textToText({
             signal.addEventListener('abort', onAbort, { once: true });
           }
 
-          if (postfix) {
-            const messages: string[] = [];
-            for await (const event of toTextStream(eventSource, {
-              timeout,
-              signal,
-            })) {
-              if (event.type === 'message') {
-                messages.push(event.data);
-              }
-            }
-            yield postfix(messages.join(''));
-          } else {
-            for await (const event of toTextStream(eventSource, {
-              timeout,
-              signal,
-            })) {
-              if (event.type === 'message') {
-                yield event.data;
-              }
+          for await (const event of toTextStream(eventSource, {
+            timeout,
+            signal,
+          })) {
+            if (event.type === 'message') {
+              yield event.data;
             }
           }
         } finally {
@@ -215,6 +210,10 @@ export function textToText({
           reasoning,
           modelId,
           toolsConfig,
+          actionId,
+          actionVersion,
+          runId,
+          retry,
         },
         endpoint
       );
@@ -244,7 +243,7 @@ export function textToText({
         }
 
         const result = messages.join('');
-        return postfix ? postfix(result) : result;
+        return result;
       } finally {
         eventSource.close();
         if (signal && onAbort) {
@@ -266,6 +265,9 @@ export function toImage({
   timeout = TIMEOUT,
   retry = false,
   endpoint,
+  actionId,
+  actionVersion,
+  runId,
   client,
 }: ToImageOptions) {
   let messageId: string | undefined;
@@ -282,12 +284,20 @@ export function toImage({
           signal,
         });
       }
-      const eventSource = client.imagesStream(
-        sessionId,
-        messageId,
-        seed,
-        endpoint
-      );
+      const eventSource =
+        endpoint === Endpoint.Action
+          ? client.chatTextStream(
+              {
+                sessionId,
+                messageId,
+                actionId,
+                actionVersion,
+                runId,
+                retry,
+              },
+              Endpoint.Action
+            )
+          : client.imagesStream(sessionId, messageId, seed, endpoint);
       AIProvider.LAST_ACTION_SESSIONID = sessionId;
 
       for await (const event of toTextStream(eventSource, {

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

@@ -0,0 +1,120 @@
+/**
+ * @vitest-environment happy-dom
+ */
+import { BehaviorSubject } from 'rxjs';
+import { describe, expect, test, vi } from 'vitest';
+
+import { AIProvider } from './ai-provider';
+import { CopilotClient, Endpoint } from './copilot-client';
+import { setupAIProvider } from './setup-provider';
+
+Object.defineProperty(globalThis, 'EventSource', {
+  configurable: true,
+  value: {
+    CLOSED: 2,
+  },
+});
+
+type SetupAIProviderArgs = Parameters<typeof setupAIProvider>;
+type ActionInput<T extends keyof BlockSuitePresets.AIActions> = Parameters<
+  NonNullable<BlockSuitePresets.AIActions[T]>
+>[0];
+
+async function drain(stream: AsyncIterable<unknown>) {
+  for await (const chunk of stream) {
+    void chunk;
+  }
+}
+
+async function drainActionResult(
+  stream: string | AsyncIterable<unknown> | undefined
+) {
+  expect(stream).toBeDefined();
+  expect(typeof stream).not.toBe('string');
+  await drain(stream as AsyncIterable<unknown>);
+}
+
+function createClosedEventSource(): EventSource {
+  return {
+    readyState: EventSource.CLOSED,
+    addEventListener: vi.fn(),
+    close: vi.fn(),
+  } as unknown as EventSource;
+}
+
+describe('setupAIProvider action migrations', () => {
+  test('routes mindmap, slides and image filter through action API', async () => {
+    const createdSessions: unknown[] = [];
+    const textStreams: unknown[] = [];
+    const client = new CopilotClient(
+      vi.fn(),
+      vi.fn(() => createClosedEventSource())
+    );
+    vi.spyOn(client, 'createSession').mockImplementation(async options => {
+      createdSessions.push(options);
+      return `session:${options.promptName}`;
+    });
+    vi.spyOn(client, 'createMessage').mockResolvedValue('message-1');
+    vi.spyOn(client, 'chatTextStream').mockImplementation(
+      (options, endpoint) => {
+        textStreams.push({ options, endpoint });
+        return createClosedEventSource();
+      }
+    );
+    vi.spyOn(client, 'imagesStream').mockReturnValue(createClosedEventSource());
+
+    setupAIProvider(
+      client,
+      { open: vi.fn() } as unknown as SetupAIProviderArgs[1],
+      {
+        session: {
+          account$: new BehaviorSubject(null),
+        },
+      } as unknown as SetupAIProviderArgs[2]
+    );
+
+    await drainActionResult(
+      await AIProvider.actions.brainstormMindmap?.({
+        workspaceId: 'workspace-1',
+        input: 'make a map',
+        stream: true,
+      } satisfies ActionInput<'brainstormMindmap'>)
+    );
+    await drainActionResult(
+      await AIProvider.actions.createSlides?.({
+        workspaceId: 'workspace-1',
+        input: 'make slides',
+        stream: true,
+      } satisfies ActionInput<'createSlides'>)
+    );
+    await drainActionResult(
+      await AIProvider.actions.filterImage?.({
+        workspaceId: 'workspace-1',
+        input: 'convert',
+        attachments: ['blob-1'],
+        style: 'Sketch style',
+      } satisfies ActionInput<'filterImage'>)
+    );
+
+    expect(createdSessions).toEqual([
+      expect.objectContaining({ promptName: 'mindmap.generate' }),
+      expect.objectContaining({ promptName: 'slides.outline' }),
+      expect.objectContaining({ promptName: 'image.filter.sketch' }),
+    ]);
+    expect(textStreams).toEqual([
+      expect.objectContaining({
+        endpoint: Endpoint.Action,
+        options: expect.objectContaining({ actionId: 'mindmap.generate' }),
+      }),
+      expect.objectContaining({
+        endpoint: Endpoint.Action,
+        options: expect.objectContaining({ actionId: 'slides.outline' }),
+      }),
+      expect.objectContaining({
+        endpoint: Endpoint.Action,
+        options: expect.objectContaining({ actionId: 'image.filter.sketch' }),
+      }),
+    ]);
+    expect(client.imagesStream).not.toHaveBeenCalled();
+  });
+});

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

@@ -10,7 +10,6 @@ import {
   type RequestOptions,
   type UpdateChatSessionInput,
 } from '@affine/graphql';
-import { z } from 'zod';
 
 import { AIProvider } from './ai-provider';
 import { type CopilotClient, Endpoint } from './copilot-client';
@@ -30,10 +29,10 @@ function toAIUserInfo(account: AuthAccountInfo | null) {
 
 const filterStyleToPromptName = new Map<string, PromptKey>(
   Object.entries({
-    'Clay style': 'Convert to Clay style',
-    'Pixel style': 'Convert to Pixel style',
-    'Sketch style': 'Convert to Sketch style',
-    'Anime style': 'Convert to Anime style',
+    'Clay style': 'image.filter.clay',
+    'Pixel style': 'image.filter.pixel',
+    'Sketch style': 'image.filter.sketch',
+    'Anime style': 'image.filter.anime',
   })
 );
 
@@ -350,7 +349,7 @@ export function setupAIProvider(
 
   AIProvider.provide('brainstormMindmap', async options => {
     const sessionId = await createSession({
-      promptName: 'workflow:brainstorm',
+      promptName: 'mindmap.generate',
       ...options,
     });
     return textToText({
@@ -360,7 +359,9 @@ export function setupAIProvider(
       content: options.input,
       // 3 minutes
       timeout: 180000,
-      endpoint: Endpoint.Workflow,
+      endpoint: Endpoint.Action,
+      actionId: 'mindmap.generate',
+      actionVersion: 'v1',
     });
   });
 
@@ -439,44 +440,8 @@ Could you make a new website based on these notes and send back just the html fi
   });
 
   AIProvider.provide('createSlides', async options => {
-    const SlideSchema = z.object({
-      page: z.number(),
-      type: z.enum(['name', 'title', 'content']),
-      content: z.string(),
-    });
-    type Slide = z.infer<typeof SlideSchema>;
-    const parseJson = (json: string) => {
-      try {
-        return SlideSchema.parse(JSON.parse(json));
-      } catch {
-        return null;
-      }
-    };
-    // TODO(@darkskygit): move this to backend's workflow after workflow support custom code action
-    const postfix = (text: string): string => {
-      const slides = text
-        .split('\n')
-        .map(parseJson)
-        .filter((v): v is Slide => !!v);
-      return slides
-        .map(slide => {
-          if (slide.type === 'name') {
-            return `- ${slide.content}`;
-          } else if (slide.type === 'title') {
-            return `  - ${slide.content}`;
-          } else if (slide.content.includes('\n')) {
-            return slide.content
-              .split('\n')
-              .map(c => `    - ${c}`)
-              .join('\n');
-          } else {
-            return `    - ${slide.content}`;
-          }
-        })
-        .join('\n');
-    };
     const sessionId = await createSession({
-      promptName: 'workflow:presentation',
+      promptName: 'slides.outline',
       ...options,
     });
     return textToText({
@@ -486,8 +451,9 @@ Could you make a new website based on these notes and send back just the html fi
       content: options.input,
       // 3 minutes
       timeout: 180000,
-      endpoint: Endpoint.Workflow,
-      postfix,
+      endpoint: Endpoint.Action,
+      actionId: 'slides.outline',
+      actionVersion: 'v1',
     });
   });
 
@@ -521,14 +487,15 @@ Could you make a new website based on these notes and send back just the html fi
       promptName,
       ...options,
     });
-    const isWorkflow = !!promptName?.startsWith('workflow:');
     return toImage({
       ...options,
       client,
       sessionId,
       content: options.input,
       timeout: 180000,
-      endpoint: isWorkflow ? Endpoint.Workflow : Endpoint.Images,
+      endpoint: Endpoint.Action,
+      actionId: promptName,
+      actionVersion: 'v1',
     });
   });