feat(server): refactor provider interface (#11665)

fix AI-4
fix AI-18

better provider/model choose to allow fallback to similar models (e.g., self-hosted) when the provider is not fully configured
split functions of different output types

packages/backend/server/src/__tests__/__snapshots__/copilot.spec.ts.md

@@ -261,3 +261,40 @@ Generated by [AVA](https://avajs.dev).
         role: 'assistant',
       },
     ]
+
+## should be able to run image executor
+
+> should generate image stream
+
+    [
+      {
+        params: {
+          key: [
+            'https://example.com/test-image.jpg',
+            'tag1, tag2, tag3, tag4, tag5, ',
+          ],
+        },
+        type: 2,
+      },
+    ]
+
+> should render the prompt with params array
+
+    [
+      {
+        modelId: 'test-image',
+      },
+      [
+        {
+          content: 'tag1, tag2, tag3, tag4, tag5, ',
+          params: {
+            tags: [
+              'tag4',
+              'tag5',
+            ],
+          },
+          role: 'user',
+        },
+      ],
+      {},
+    ]

packages/backend/server/src/__tests__/copilot-provider.spec.ts

@@ -342,7 +342,7 @@ const actions = [
         TranscriptionResponseSchema.parse(JSON.parse(result));
       });
     },
-    type: 'text' as const,
+    type: 'structured' as const,
   },
   {
     name: 'Should transcribe middle audio',
@@ -364,7 +364,7 @@ const actions = [
         TranscriptionResponseSchema.parse(JSON.parse(result));
       });
     },
-    type: 'text' as const,
+    type: 'structured' as const,
   },
   {
     name: 'Should transcribe long audio',
@@ -386,7 +386,7 @@ const actions = [
         TranscriptionResponseSchema.parse(JSON.parse(result));
       });
     },
-    type: 'text' as const,
+    type: 'structured' as const,
   },
   {
     promptName: [
@@ -564,43 +564,71 @@ for (const { name, promptName, messages, verifier, type, config } of actions) {
         const provider = (await factory.getProviderByModel(prompt.model))!;
         t.truthy(provider, 'should have provider');
         await retry(`action: ${promptName}`, t, async t => {
-          if (type === 'text' && 'generateText' in provider) {
-            const result = await provider.generateText(
-              [
-                ...prompt.finish(
-                  messages.reduce(
-                    // @ts-expect-error
-                    (acc, m) => Object.assign(acc, m.params),
-                    {}
-                  )
-                ),
-                ...messages,
-              ],
-              prompt.model,
-              Object.assign({}, prompt.config, config)
-            );
-            t.truthy(result, 'should return result');
-            verifier?.(t, result);
-          } else if (type === 'image' && 'generateImages' in provider) {
-            const result = await provider.generateImages(
-              [
-                ...prompt.finish(
-                  messages.reduce(
-                    // @ts-expect-error
-                    (acc, m) => Object.assign(acc, m.params),
-                    {}
-                  )
-                ),
-                ...messages,
-              ],
-              prompt.model
-            );
-            t.truthy(result.length, 'should return result');
-            for (const r of result) {
-              verifier?.(t, r);
+          switch (type) {
+            case 'text': {
+              const result = await provider.text(
+                { modelId: prompt.model },
+                [
+                  ...prompt.finish(
+                    messages.reduce(
+                      // @ts-expect-error
+                      (acc, m) => Object.assign(acc, m.params),
+                      {}
+                    )
+                  ),
+                  ...messages,
+                ],
+                Object.assign({}, prompt.config, config)
+              );
+              t.truthy(result, 'should return result');
+              verifier?.(t, result);
+              break;
+            }
+            case 'structured': {
+              const result = await provider.structure(
+                { modelId: prompt.model },
+                [
+                  ...prompt.finish(
+                    messages.reduce(
+                      (acc, m) => Object.assign(acc, m.params),
+                      {}
+                    )
+                  ),
+                  ...messages,
+                ],
+                Object.assign({}, prompt.config, config)
+              );
+              t.truthy(result, 'should return result');
+              verifier?.(t, result);
+              break;
+            }
+            case 'image': {
+              const stream = provider.streamImages({ modelId: prompt.model }, [
+                ...prompt.finish(
+                  messages.reduce(
+                    // @ts-expect-error
+                    (acc, m) => Object.assign(acc, m.params),
+                    {}
+                  )
+                ),
+                ...messages,
+              ]);
+
+              const result = [];
+              for await (const attachment of stream) {
+                result.push(attachment);
+              }
+
+              t.truthy(result.length, 'should return result');
+              for (const r of result) {
+                verifier?.(t, r);
+              }
+              break;
+            }
+            default: {
+              t.fail('unsupported provider type');
+              break;
             }
-          } else {
-            t.fail('unsupported provider type');
           }
         });
       }

packages/backend/server/src/__tests__/copilot.e2e.ts

@@ -121,6 +121,7 @@ test.before(async t => {
 });
 
 const textPromptName = 'prompt';
+const imagePromptName = 'prompt-image';
 test.beforeEach(async t => {
   Sinon.restore();
   const { app, prompt } = t.context;
@@ -131,6 +132,10 @@ test.beforeEach(async t => {
   await prompt.set(textPromptName, 'test', [
     { role: 'system', content: 'hello {{word}}' },
   ]);
+
+  await prompt.set(imagePromptName, 'test-image', [
+    { role: 'system', content: 'hello {{word}}' },
+  ]);
 });
 
 test.after.always(async t => {
@@ -441,33 +446,44 @@ test('should be able to chat with api', async t => {
   Sinon.stub(storage, 'handleRemoteLink').resolvesArg(2);
 
   const { id } = await createWorkspace(app);
-  const sessionId = await createCopilotSession(
-    app,
-    id,
-    randomUUID(),
-    textPromptName
-  );
-  const messageId = await createCopilotMessage(app, sessionId);
-  const ret = await chatWithText(app, sessionId, messageId);
-  t.is(ret, 'generate text to text', 'should be able to chat with text');
+  {
+    const sessionId = await createCopilotSession(
+      app,
+      id,
+      randomUUID(),
+      textPromptName
+    );
+    const messageId = await createCopilotMessage(app, sessionId);
+    const ret = await chatWithText(app, sessionId, messageId);
+    t.is(ret, 'generate text to text', 'should be able to chat with text');
 
-  const ret2 = await chatWithTextStream(app, sessionId, messageId);
-  t.is(
-    ret2,
-    textToEventStream('generate text to text stream', messageId),
-    'should be able to chat with text stream'
-  );
+    const ret2 = await chatWithTextStream(app, sessionId, messageId);
+    t.is(
+      ret2,
+      textToEventStream('generate text to text stream', messageId),
+      'should be able to chat with text stream'
+    );
+  }
 
-  const ret3 = await chatWithImages(app, sessionId, messageId);
-  t.is(
-    array2sse(sse2array(ret3).filter(e => e.event !== 'event')),
-    textToEventStream(
-      ['https://example.com/test.jpg', 'hello '],
-      messageId,
-      'attachment'
-    ),
-    'should be able to chat with images'
-  );
+  {
+    const sessionId = await createCopilotSession(
+      app,
+      id,
+      randomUUID(),
+      imagePromptName
+    );
+    const messageId = await createCopilotMessage(app, sessionId);
+    const ret3 = await chatWithImages(app, sessionId, messageId);
+    t.is(
+      array2sse(sse2array(ret3).filter(e => e.event !== 'event')),
+      textToEventStream(
+        ['https://example.com/test-image.jpg', 'hello '],
+        messageId,
+        'attachment'
+      ),
+      'should be able to chat with images'
+    );
+  }
 
   Sinon.restore();
 });
@@ -918,7 +934,10 @@ test('should be able to transcript', async t => {
 
   const { id: workspaceId } = await createWorkspace(app);
 
-  Sinon.stub(app.get(GeminiProvider), 'generateText').resolves(
+  Sinon.stub(app.get(GeminiProvider), 'structure').resolves(
+    '[{"a":"A","s":30,"e":45,"t":"Hello, everyone."},{"a":"B","s":46,"e":70,"t":"Hi, thank you for joining the meeting today."}]'
+  );
+  Sinon.stub(app.get(GeminiProvider), 'text').resolves(
     '[{"a":"A","s":30,"e":45,"t":"Hello, everyone."},{"a":"B","s":46,"e":70,"t":"Hi, thank you for joining the meeting today."}]'
   );
 

packages/backend/server/src/__tests__/copilot.spec.ts

@@ -20,9 +20,10 @@ import {
 import { MockEmbeddingClient } from '../plugins/copilot/context/embedding';
 import { prompts, PromptService } from '../plugins/copilot/prompt';
 import {
-  CopilotCapability,
   CopilotProviderFactory,
   CopilotProviderType,
+  ModelInputType,
+  ModelOutputType,
   OpenAIProvider,
 } from '../plugins/copilot/providers';
 import { CitationParser } from '../plugins/copilot/providers/utils';
@@ -756,9 +757,7 @@ test('should be able to get provider', async t => {
   const { factory } = t.context;
 
   {
-    const p = await factory.getProviderByCapability(
-      CopilotCapability.TextToText
-    );
+    const p = await factory.getProvider({ outputType: ModelOutputType.Text });
     t.is(
       p?.type.toString(),
       'openai',
@@ -767,36 +766,41 @@ test('should be able to get provider', async t => {
   }
 
   {
-    const p = await factory.getProviderByCapability(
-      CopilotCapability.ImageToImage,
-      { model: 'lora/image-to-image' }
-    );
+    const p = await factory.getProvider({
+      outputType: ModelOutputType.Image,
+      inputTypes: [ModelInputType.Image],
+      modelId: 'lora/image-to-image',
+    });
     t.is(
       p?.type.toString(),
       'fal',
-      'should get provider support text-to-embedding'
+      'should get provider supporting image output'
     );
   }
 
   {
-    const p = await factory.getProviderByCapability(
-      CopilotCapability.ImageToText,
+    const p = await factory.getProvider(
+      {
+        outputType: ModelOutputType.Image,
+        inputTypes: [ModelInputType.Image],
+      },
       { prefer: CopilotProviderType.FAL }
     );
     t.is(
       p?.type.toString(),
       'fal',
-      'should get provider support text-to-embedding'
+      'should get provider supporting text output with image input'
     );
   }
 
   // if a model is not defined and not available in online api
   // it should return null
   {
-    const p = await factory.getProviderByCapability(
-      CopilotCapability.ImageToText,
-      { model: 'gpt-4-not-exist' }
-    );
+    const p = await factory.getProvider({
+      outputType: ModelOutputType.Text,
+      inputTypes: [ModelInputType.Text],
+      modelId: 'gpt-4-not-exist',
+    });
     t.falsy(p, 'should not get provider');
   }
 });
@@ -987,10 +991,9 @@ test('should be able to run text executor', async t => {
     { role: 'system', content: 'hello {{word}}' },
   ]);
   // mock provider
-  const testProvider =
-    (await factory.getProviderByModel<CopilotCapability.TextToText>('test'))!;
-  const text = Sinon.spy(testProvider, 'generateText');
-  const textStream = Sinon.spy(testProvider, 'generateTextStream');
+  const testProvider = (await factory.getProviderByModel('test'))!;
+  const text = Sinon.spy(testProvider, 'text');
+  const textStream = Sinon.spy(testProvider, 'streamText');
 
   const nodeData: WorkflowNodeData = {
     id: 'basic',
@@ -1013,7 +1016,7 @@ test('should be able to run text executor', async t => {
       },
     ]);
     t.deepEqual(
-      text.lastCall.args[0][0].content,
+      text.lastCall.args[1][0].content,
       'hello world',
       'should render the prompt with params'
     );
@@ -1036,7 +1039,7 @@ test('should be able to run text executor', async t => {
       }))
     );
     t.deepEqual(
-      textStream.lastCall.args[0][0].params?.attachments,
+      textStream.lastCall.args[1][0].params?.attachments,
       ['https://affine.pro/example.jpg'],
       'should pass attachments to provider'
     );
@@ -1050,14 +1053,13 @@ test('should be able to run image executor', async t => {
 
   executors.image.register();
   const executor = getWorkflowExecutor(executors.image.type);
-  await prompt.set('test', 'test', [
+  await prompt.set('test', 'test-image', [
     { role: 'user', content: 'tag1, tag2, tag3, {{#tags}}{{.}}, {{/tags}}' },
   ]);
   // mock provider
-  const testProvider =
-    (await factory.getProviderByModel<CopilotCapability.TextToImage>('test'))!;
-  const image = Sinon.spy(testProvider, 'generateImages');
-  const imageStream = Sinon.spy(testProvider, 'generateImagesStream');
+  const testProvider = (await factory.getProviderByModel('test'))!;
+
+  const imageStream = Sinon.spy(testProvider, 'streamImages');
 
   const nodeData: WorkflowNodeData = {
     id: 'basic',
@@ -1076,20 +1078,9 @@ test('should be able to run image executor', async t => {
       )
     );
 
-    t.deepEqual(ret, [
-      {
-        type: NodeExecuteState.Params,
-        params: {
-          key: [
-            'https://example.com/test.jpg',
-            'tag1, tag2, tag3, tag4, tag5, ',
-          ],
-        },
-      },
-    ]);
-    t.deepEqual(
-      image.lastCall.args[0][0].content,
-      'tag1, tag2, tag3, tag4, tag5, ',
+    t.snapshot(ret, 'should generate image stream');
+    t.snapshot(
+      imageStream.lastCall.args,
       'should render the prompt with params array'
     );
   }
@@ -1104,16 +1095,17 @@ test('should be able to run image executor', async t => {
 
     t.deepEqual(
       ret,
-      Array.from(['https://example.com/test.jpg', 'tag1, tag2, tag3, ']).map(
-        t => ({
-          attachment: t,
-          nodeId: 'basic',
-          type: NodeExecuteState.Attachment,
-        })
-      )
+      Array.from([
+        'https://example.com/test-image.jpg',
+        'tag1, tag2, tag3, ',
+      ]).map(t => ({
+        attachment: t,
+        nodeId: 'basic',
+        type: NodeExecuteState.Attachment,
+      }))
     );
     t.deepEqual(
-      imageStream.lastCall.args[0][0].params?.attachments,
+      imageStream.lastCall.args[1][0].params?.attachments,
       ['https://affine.pro/example.jpg'],
       'should pass attachments to provider'
     );

packages/backend/server/src/__tests__/mocks/copilot.mock.ts

@@ -1,9 +1,13 @@
 import { randomBytes } from 'node:crypto';
 
 import {
-  CopilotCapability,
   CopilotChatOptions,
   CopilotEmbeddingOptions,
+  CopilotImageOptions,
+  CopilotStructuredOptions,
+  ModelConditions,
+  ModelInputType,
+  ModelOutputType,
   PromptMessage,
 } from '../../plugins/copilot/providers';
 import {
@@ -14,49 +18,135 @@ import { sleep } from '../utils/utils';
 
 export class MockCopilotProvider extends OpenAIProvider {
   override readonly models = [
-    'test',
-    'gpt-4o',
-    'gpt-4o-2024-08-06',
-    'gpt-4.1',
-    'gpt-4.1-2025-04-14',
-    'gpt-4.1-mini',
-    'fast-sdxl/image-to-image',
-    'lcm-sd15-i2i',
-    'clarity-upscaler',
-    'imageutils/rembg',
-    'gemini-2.5-pro-preview-03-25',
+    {
+      id: 'test',
+      capabilities: [
+        {
+          input: [ModelInputType.Text],
+          output: [ModelOutputType.Text],
+          defaultForOutputType: true,
+        },
+      ],
+    },
+    {
+      id: 'test-image',
+      capabilities: [
+        {
+          input: [ModelInputType.Text],
+          output: [ModelOutputType.Image],
+          defaultForOutputType: true,
+        },
+      ],
+    },
+    {
+      id: 'gpt-4o',
+      capabilities: [
+        {
+          input: [ModelInputType.Text, ModelInputType.Image],
+          output: [ModelOutputType.Text],
+        },
+      ],
+    },
+    {
+      id: 'gpt-4o-2024-08-06',
+      capabilities: [
+        {
+          input: [ModelInputType.Text, ModelInputType.Image],
+          output: [ModelOutputType.Text],
+        },
+      ],
+    },
+    {
+      id: 'gpt-4.1',
+      capabilities: [
+        {
+          input: [ModelInputType.Text, ModelInputType.Image],
+          output: [ModelOutputType.Text],
+        },
+      ],
+    },
+    {
+      id: 'gpt-4.1-2025-04-14',
+      capabilities: [
+        {
+          input: [ModelInputType.Text, ModelInputType.Image],
+          output: [ModelOutputType.Text],
+        },
+      ],
+    },
+    {
+      id: 'gpt-4.1-mini',
+      capabilities: [
+        {
+          input: [ModelInputType.Text, ModelInputType.Image],
+          output: [ModelOutputType.Text],
+        },
+      ],
+    },
+    {
+      id: 'lcm-sd15-i2i',
+      capabilities: [
+        {
+          input: [ModelInputType.Image],
+          output: [ModelOutputType.Image],
+        },
+      ],
+    },
+    {
+      id: 'clarity-upscaler',
+      capabilities: [
+        {
+          input: [ModelInputType.Image],
+          output: [ModelOutputType.Image],
+        },
+      ],
+    },
+    {
+      id: 'imageutils/rembg',
+      capabilities: [
+        {
+          input: [ModelInputType.Image],
+          output: [ModelOutputType.Image],
+        },
+      ],
+    },
+    {
+      id: 'gemini-2.5-pro-preview-03-25',
+      capabilities: [
+        {
+          input: [ModelInputType.Text, ModelInputType.Image],
+          output: [ModelOutputType.Text, ModelOutputType.Structured],
+        },
+      ],
+    },
   ];
 
-  override readonly capabilities = [
-    CopilotCapability.TextToText,
-    CopilotCapability.TextToEmbedding,
-    CopilotCapability.TextToImage,
-    CopilotCapability.ImageToImage,
-    CopilotCapability.ImageToText,
-  ];
-
-  // ====== text to text ======
-
-  override async generateText(
+  override async text(
+    cond: ModelConditions,
     messages: PromptMessage[],
-    model: string = 'test',
     options: CopilotChatOptions = {}
   ): Promise<string> {
-    await this.checkParams({ messages, model, options });
+    const fullCond = {
+      ...cond,
+      outputType: ModelOutputType.Text,
+    };
+    await this.checkParams({ messages, cond: fullCond, options });
     // make some time gap for history test case
     await sleep(100);
     return 'generate text to text';
   }
 
-  override async *generateTextStream(
+  override async *streamText(
+    cond: ModelConditions,
     messages: PromptMessage[],
-    model: string = 'gpt-4.1-mini',
     options: CopilotChatOptions = {}
   ): AsyncIterable<string> {
-    await this.checkParams({ messages, model, options });
+    const fullCond = { ...cond, outputType: ModelOutputType.Text };
+    await this.checkParams({ messages, cond: fullCond, options });
 
     // make some time gap for history test case
     await sleep(100);
+
     const result = 'generate text to text stream';
     for (const message of result) {
       yield message;
@@ -66,52 +156,60 @@ export class MockCopilotProvider extends OpenAIProvider {
     }
   }
 
+  override async structure(
+    cond: ModelConditions,
+    messages: PromptMessage[],
+    options: CopilotStructuredOptions = {}
+  ): Promise<string> {
+    const fullCond = { ...cond, outputType: ModelOutputType.Structured };
+    await this.checkParams({ messages, cond: fullCond, options });
+
+    // make some time gap for history test case
+    await sleep(100);
+    return 'generate text to text';
+  }
+
+  override async *streamImages(
+    cond: ModelConditions,
+    messages: PromptMessage[],
+    options: CopilotImageOptions = {}
+  ) {
+    const fullCond = { ...cond, outputType: ModelOutputType.Image };
+    await this.checkParams({ messages, cond: fullCond, options });
+
+    // make some time gap for history test case
+    await sleep(100);
+
+    const { content: prompt } = [...messages].pop() || {};
+    if (!prompt) throw new Error('Prompt is required');
+
+    const imageUrls = [
+      `https://example.com/${cond.modelId || 'test'}.jpg`,
+      prompt,
+    ];
+
+    for (const imageUrl of imageUrls) {
+      yield imageUrl;
+      if (options.signal?.aborted) {
+        break;
+      }
+    }
+    return;
+  }
+
   // ====== text to embedding ======
 
-  override async generateEmbedding(
+  override async embedding(
+    cond: ModelConditions,
     messages: string | string[],
-    model: string,
     options: CopilotEmbeddingOptions = { dimensions: DEFAULT_DIMENSIONS }
   ): Promise<number[][]> {
     messages = Array.isArray(messages) ? messages : [messages];
-    await this.checkParams({ embeddings: messages, model, options });
+    const fullCond = { ...cond, outputType: ModelOutputType.Embedding };
+    await this.checkParams({ embeddings: messages, cond: fullCond, options });
 
     // make some time gap for history test case
     await sleep(100);
     return [Array.from(randomBytes(options.dimensions)).map(v => v % 128)];
   }
-
-  // ====== text to image ======
-  override async generateImages(
-    messages: PromptMessage[],
-    model: string = 'test',
-    _options: {
-      signal?: AbortSignal;
-      user?: string;
-    } = {}
-  ): Promise<Array<string>> {
-    const { content: prompt } = messages[0] || {};
-    if (!prompt) {
-      throw new Error('Prompt is required');
-    }
-
-    // make some time gap for history test case
-    await sleep(100);
-    // just let test case can easily verify the final prompt
-    return [`https://example.com/${model}.jpg`, prompt];
-  }
-
-  override async *generateImagesStream(
-    messages: PromptMessage[],
-    model: string = 'dall-e-3',
-    options: {
-      signal?: AbortSignal;
-      user?: string;
-    } = {}
-  ): AsyncIterable<string> {
-    const ret = await this.generateImages(messages, model, options);
-    for (const url of ret) {
-      yield url;
-    }
-  }
 }