diff --git a/packages/backend/server/src/__tests__/copilot-provider.spec.ts b/packages/backend/server/src/__tests__/copilot-provider.spec.ts index 6ca04e26d..8cfe21f3e 100644 --- a/packages/backend/server/src/__tests__/copilot-provider.spec.ts +++ b/packages/backend/server/src/__tests__/copilot-provider.spec.ts @@ -518,12 +518,7 @@ const actions = [ type: 'text' as const, }, { - promptName: [ - 'debug:action:fal-face-to-sticker', - 'debug:action:fal-remove-bg', - 'debug:action:fal-sd15', - 'debug:action:fal-upscaler', - ], + promptName: ['Convert to sticker', 'Remove background', 'Upscale image'], messages: [ { role: 'user' as const, @@ -590,6 +585,8 @@ for (const { }))!; t.truthy(provider, 'should have provider'); await retry(`action: ${promptName}`, t, async t => { + const finalConfig = Object.assign({}, prompt.config, config); + switch (type) { case 'text': { const result = await provider.text( @@ -604,7 +601,7 @@ for (const { ), ...messages, ], - Object.assign({}, prompt.config, config) + finalConfig ); t.truthy(result, 'should return result'); verifier?.(t, result); @@ -622,23 +619,39 @@ for (const { ), ...messages, ], - Object.assign({}, prompt.config, config) + finalConfig ); 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 finalMessage = [...messages]; + const params = {}; + if (finalMessage.length === 1) { + const latestMessage = finalMessage.pop()!; + Object.assign(params, { + content: latestMessage.content, + attachments: + 'attachments' in latestMessage + ? latestMessage.attachments + : undefined, + }); + } + const stream = provider.streamImages( + { modelId: prompt.model }, + [ + ...prompt.finish( + finalMessage.reduce( + // @ts-expect-error + (acc, m) => Object.assign(acc, m.params), + params + ) + ), + ...finalMessage, + ], + finalConfig + ); const result = []; for await (const attachment of stream) { diff --git a/packages/backend/server/src/__tests__/copilot.e2e.ts b/packages/backend/server/src/__tests__/copilot.e2e.ts index 68cb33541..3ee870cb7 100644 --- a/packages/backend/server/src/__tests__/copilot.e2e.ts +++ b/packages/backend/server/src/__tests__/copilot.e2e.ts @@ -543,12 +543,19 @@ test('should be able to chat with special image model', async t => { ); }; - await testWithModel('debug:action:fal-sd15', 'some-tag'); + await testWithModel('Generate image', 'some-tag'); await testWithModel( - 'debug:action:fal-upscaler', - 'best quality, 8K resolution, highres, clarity, some-tag' + 'Convert to sticker', + 'convert this image to sticker. you need to identify the subject matter and warp a circle of white stroke around the subject matter and with transparent background. some-tag' + ); + await testWithModel( + 'Upscale image', + 'make the image more detailed. some-tag' + ); + await testWithModel( + 'Remove background', + 'Keep the subject and remove other non-subject items. Transparent background. some-tag' ); - await testWithModel('debug:action:fal-remove-bg', 'some-tag'); Sinon.restore(); }); diff --git a/packages/backend/server/src/__tests__/mocks/copilot.mock.ts b/packages/backend/server/src/__tests__/mocks/copilot.mock.ts index 52bf3a5d2..39ecca31c 100644 --- a/packages/backend/server/src/__tests__/mocks/copilot.mock.ts +++ b/packages/backend/server/src/__tests__/mocks/copilot.mock.ts @@ -84,29 +84,12 @@ export class MockCopilotProvider extends OpenAIProvider { ], }, { - id: 'lcm-sd15-i2i', + id: 'gpt-image-1', capabilities: [ { - input: [ModelInputType.Image], - output: [ModelOutputType.Image], - }, - ], - }, - { - id: 'clarity-upscaler', - capabilities: [ - { - input: [ModelInputType.Image], - output: [ModelOutputType.Image], - }, - ], - }, - { - id: 'imageutils/rembg', - capabilities: [ - { - input: [ModelInputType.Image], + input: [ModelInputType.Text, ModelInputType.Image], output: [ModelOutputType.Image], + defaultForOutputType: true, }, ], }, diff --git a/packages/backend/server/src/plugins/copilot/prompt/prompts.ts b/packages/backend/server/src/plugins/copilot/prompt/prompts.ts index 9b5af3fe9..13610a475 100644 --- a/packages/backend/server/src/plugins/copilot/prompt/prompts.ts +++ b/packages/backend/server/src/plugins/copilot/prompt/prompts.ts @@ -20,12 +20,6 @@ type Prompt = Omit< }; const workflows: Prompt[] = [ - { - name: 'debug:action:fal-teed', - action: 'fal-teed', - model: 'workflowutils/teed', - messages: [{ role: 'user', content: '{{content}}' }], - }, { name: 'workflow:presentation', action: 'workflow:presentation', @@ -305,48 +299,7 @@ const workflows: Prompt[] = [ }, ]; -const actions: Prompt[] = [ - { - name: 'debug:action:dalle3', - action: 'image', - model: 'dall-e-3', - messages: [], - }, - { - name: 'debug:action:gpt-image-1', - action: 'image', - model: 'gpt-image-1', - messages: [], - }, - { - name: 'debug:action:fal-sd15', - action: 'image', - model: 'lcm-sd15-i2i', - messages: [], - }, - { - name: 'debug:action:fal-upscaler', - action: 'Clearer', - model: 'clarity-upscaler', - messages: [ - { - role: 'user', - content: 'best quality, 8K resolution, highres, clarity, {{content}}', - }, - ], - }, - { - name: 'debug:action:fal-remove-bg', - action: 'Remove background', - model: 'imageutils/rembg', - messages: [], - }, - { - name: 'debug:action:fal-face-to-sticker', - action: 'Convert to sticker', - model: 'face-to-sticker', - messages: [], - }, +const textActions: Prompt[] = [ { name: 'Transcript audio', action: 'Transcript audio', @@ -1449,6 +1402,161 @@ When sent new notes, respond ONLY with the contents of the html file.`, }, ]; +const imageActions: Prompt[] = [ + { + name: 'Generate image', + action: 'image', + model: 'gpt-image-1', + messages: [ + { + role: 'user', + content: '{{content}}', + }, + ], + }, + { + name: 'Convert to Clay style', + action: 'Convert to Clay style', + model: 'gpt-image-1', + messages: [ + { + role: 'user', + content: + 'Migration style. Migrates the style from the first image to the second. turn to clay/claymation style. {{content}}', + }, + ], + }, + { + name: 'Convert to Sketch style', + action: 'Convert to Sketch style', + model: 'gpt-image-1', + messages: [ + { + role: 'user', + content: 'turn to mono-color sketch style. {{content}}', + }, + ], + }, + { + name: 'Convert to Anime style', + action: 'Convert to Anime style', + model: 'gpt-image-1', + messages: [ + { + role: 'user', + content: 'turn to Suzume style like anime style. {{content}}', + }, + ], + }, + { + name: 'Convert to Pixel style', + action: 'Convert to Pixel style', + model: 'gpt-image-1', + messages: [ + { + role: 'user', + content: 'turn to kairosoft pixel art. {{content}}', + }, + ], + }, + { + name: 'Convert to sticker', + action: 'Convert to sticker', + model: 'gpt-image-1', + messages: [ + { + role: 'user', + content: + 'convert this image to sticker. you need to identify the subject matter and warp a circle of white stroke around the subject matter and with transparent background. {{content}}', + }, + ], + }, + { + name: 'Upscale image', + action: 'Upscale image', + model: 'gpt-image-1', + messages: [ + { + role: 'user', + content: 'make the image more detailed. {{content}}', + }, + ], + }, + { + name: 'Remove background', + action: 'Remove background', + model: 'gpt-image-1', + messages: [ + { + role: 'user', + content: + 'Keep the subject and remove other non-subject items. Transparent background. {{content}}', + }, + ], + }, + // TODO(@darkskygit): deprecated, remove it after <0.22 version is outdated + { + name: 'debug:action:fal-remove-bg', + action: 'Remove background', + model: 'imageutils/rembg', + messages: [], + }, + { + name: 'debug:action:fal-face-to-sticker', + action: 'Convert to sticker', + model: 'face-to-sticker', + messages: [], + }, + { + name: 'debug:action:fal-teed', + action: 'fal-teed', + model: 'workflowutils/teed', + messages: [{ role: 'user', content: '{{content}}' }], + }, + { + name: 'debug:action:dalle3', + action: 'image', + model: 'dall-e-3', + messages: [ + { + role: 'user', + content: '{{content}}', + }, + ], + }, + { + name: 'debug:action:gpt-image-1', + action: 'image', + model: 'gpt-image-1', + messages: [ + { + role: 'user', + content: '{{content}}', + }, + ], + config: { + requireContent: false, + }, + }, + { + name: 'debug:action:fal-sd15', + action: 'image', + model: 'lcm-sd15-i2i', + messages: [], + }, + { + name: 'debug:action:fal-upscaler', + action: 'Clearer', + model: 'clarity-upscaler', + messages: [ + { + role: 'user', + content: 'best quality, 8K resolution, highres, clarity, {{content}}', + }, + ], + }, +]; + const CHAT_PROMPT: Omit = { model: 'gpt-4.1', optionalModels: [ @@ -1622,7 +1730,12 @@ const chat: Prompt[] = [ }, ]; -export const prompts: Prompt[] = [...actions, ...chat, ...workflows]; +export const prompts: Prompt[] = [ + ...textActions, + ...imageActions, + ...chat, + ...workflows, +]; export async function refreshPrompts(db: PrismaClient) { const needToSkip = await db.aiPrompt diff --git a/packages/backend/server/src/plugins/copilot/providers/openai.ts b/packages/backend/server/src/plugins/copilot/providers/openai.ts index 5760d46bb..f5497573e 100644 --- a/packages/backend/server/src/plugins/copilot/providers/openai.ts +++ b/packages/backend/server/src/plugins/copilot/providers/openai.ts @@ -13,6 +13,7 @@ import { streamText, ToolSet, } from 'ai'; +import { z } from 'zod'; import { CopilotPromptInvalid, @@ -40,6 +41,20 @@ export type OpenAIConfig = { baseUrl?: string; }; +const ImageResponseSchema = z.union([ + z.object({ + data: z.array(z.object({ b64_json: z.string() })), + }), + z.object({ + error: z.object({ + message: z.string(), + type: z.string().nullish(), + param: z.any().nullish(), + code: z.union([z.string(), z.number()]).nullish(), + }), + }), +]); + export class OpenAIProvider extends CopilotProvider { readonly type = CopilotProviderType.OpenAI; @@ -389,6 +404,63 @@ export class OpenAIProvider extends CopilotProvider { } } + // ====== text to image ====== + private async *generateImageWithAttachments( + model: string, + prompt: string, + attachments: NonNullable + ): AsyncGenerator { + const form = new FormData(); + form.set('model', model); + form.set('prompt', prompt); + form.set('output_format', 'webp'); + + for (const [idx, entry] of attachments.entries()) { + const url = typeof entry === 'string' ? entry : entry.attachment; + const resp = await fetch(url); + if (resp.ok) { + const type = resp.headers.get('content-type'); + if (type && type.startsWith('image/')) { + const buffer = new Uint8Array(await resp.arrayBuffer()); + const file = new File([buffer], `${idx}.png`, { type }); + form.append('image[]', file); + } + } + } + + if (!form.getAll('image[]').length) { + throw new CopilotPromptInvalid( + 'No valid image attachments found. Please attach images.' + ); + } + + const url = `${this.config.baseUrl || 'https://api.openai.com'}/v1/images/edits`; + const res = await fetch(url, { + method: 'POST', + headers: { Authorization: `Bearer ${this.config.apiKey}` }, + body: form, + }); + + if (!res.ok) { + throw new Error(`OpenAI API error ${res.status}: ${await res.text()}`); + } + + const json = await res.json(); + const imageResponse = ImageResponseSchema.safeParse(json); + if (imageResponse.success) { + const data = imageResponse.data; + if ('error' in data) { + throw new Error(data.error.message); + } else { + for (const image of data.data) { + yield `data:image/webp;base64,${image.b64_json}`; + } + } + } else { + throw new Error(imageResponse.error.message); + } + } + override async *streamImages( cond: ModelConditions, messages: PromptMessage[], @@ -402,30 +474,33 @@ export class OpenAIProvider extends CopilotProvider { .counter('generate_images_stream_calls') .add(1, { model: model.id }); - const { content: prompt } = [...messages].pop() || {}; + const { content: prompt, attachments } = [...messages].pop() || {}; if (!prompt) throw new CopilotPromptInvalid('Prompt is required'); try { - const modelInstance = this.#instance.image(model.id); - - const result = await generateImage({ - model: modelInstance, - prompt, - providerOptions: { - openai: { - quality: options.quality || null, + if (attachments && attachments.length > 0) { + yield* this.generateImageWithAttachments(model.id, prompt, attachments); + } else { + const modelInstance = this.#instance.image(model.id); + const result = await generateImage({ + model: modelInstance, + prompt, + providerOptions: { + openai: { + quality: options.quality || null, + }, }, - }, - }); + }); - const imageUrls = result.images.map( - image => `data:image/png;base64,${image.base64}` - ); + const imageUrls = result.images.map( + image => `data:image/png;base64,${image.base64}` + ); - for (const imageUrl of imageUrls) { - yield imageUrl; - if (options.signal?.aborted) { - break; + for (const imageUrl of imageUrls) { + yield imageUrl; + if (options.signal?.aborted) { + break; + } } } return; diff --git a/packages/backend/server/src/plugins/copilot/providers/utils.ts b/packages/backend/server/src/plugins/copilot/providers/utils.ts index 02dc4cac1..b5c805580 100644 --- a/packages/backend/server/src/plugins/copilot/providers/utils.ts +++ b/packages/backend/server/src/plugins/copilot/providers/utils.ts @@ -39,7 +39,7 @@ const FORMAT_INFER_MAP: Record = { flv: 'video/flv', }; -async function inferMimeType(url: string) { +export async function inferMimeType(url: string) { if (url.startsWith('data:')) { return url.split(';')[0].split(':')[1]; } diff --git a/packages/backend/server/src/plugins/copilot/session.ts b/packages/backend/server/src/plugins/copilot/session.ts index 08ce26e17..e160098b2 100644 --- a/packages/backend/server/src/plugins/copilot/session.ts +++ b/packages/backend/server/src/plugins/copilot/session.ts @@ -141,16 +141,13 @@ export class ChatSession implements AsyncDisposable { return ret; } - finish(params: PromptParams): PromptMessage[] { - const messages = this.takeMessages(); + private mergeUserContent(params: PromptParams) { + const messages = this.stashMessages; const firstMessage = messages.at(0); - // TODO: refactor this {{content}} keyword agreement - // if the message in prompt config contains {{content}}, - // we should combine it with the user message in the prompt if ( - messages.length === 1 && - firstMessage && - this.state.prompt.paramKeys.includes('content') + this.state.prompt.paramKeys.includes('content') && + !messages.some(m => m.role === AiPromptRole.assistant) && + firstMessage ) { const normalizedParams = { ...params, @@ -178,7 +175,18 @@ export class ChatSession implements AsyncDisposable { return finished; } + return; + } + finish(params: PromptParams): PromptMessage[] { + // if the message in prompt config contains {{content}}, + // we should combine it with the user message in the prompt + const mergedMessage = this.mergeUserContent(params); + if (mergedMessage) { + return mergedMessage; + } + + const messages = this.takeMessages(); const lastMessage = messages.at(-1); return [ ...this.state.prompt.finish( diff --git a/packages/frontend/core/src/blocksuite/ai/chat-panel/actions/image.ts b/packages/frontend/core/src/blocksuite/ai/chat-panel/actions/image.ts index d3c977e77..89be5d5a7 100644 --- a/packages/frontend/core/src/blocksuite/ai/chat-panel/actions/image.ts +++ b/packages/frontend/core/src/blocksuite/ai/chat-panel/actions/image.ts @@ -21,7 +21,8 @@ export class ActionImage extends WithDisposable(ShadowlessElement) { accessor testId = 'action-image'; protected override render() { - const images = this.item.messages[0].attachments; + const images = + this.item.messages[1]?.attachments ?? this.item.messages[0].attachments; return html`
diff --git a/packages/frontend/core/src/blocksuite/ai/provider/prompt.ts b/packages/frontend/core/src/blocksuite/ai/provider/prompt.ts index ed9041f63..08933dbcc 100644 --- a/packages/frontend/core/src/blocksuite/ai/provider/prompt.ts +++ b/packages/frontend/core/src/blocksuite/ai/provider/prompt.ts @@ -1,12 +1,7 @@ // manually synced with packages/backend/server/src/data/migrations/utils/prompts.ts // TODO(@Peng): automate this export const promptKeys = [ - 'debug:action:dalle3', - 'debug:action:gpt-image-1', - 'debug:action:fal-sd15', - 'debug:action:fal-upscaler', - 'debug:action:fal-remove-bg', - 'debug:action:fal-face-to-sticker', + // text actions 'Chat With AFFiNE AI', 'Search With AFFiNE AI', 'Summary', @@ -36,12 +31,18 @@ export const promptKeys = [ 'Make it longer', 'Make it shorter', 'Continue writing', + // image actions + 'Generate image', + 'Convert to Anime style', + 'Convert to Clay style', + 'Convert to Pixel style', + 'Convert to Sketch style', + 'Convert to sticker', + 'Upscale image', + 'Remove background', + // workflows 'workflow:presentation', 'workflow:brainstorm', - 'workflow:image-sketch', - 'workflow:image-clay', - 'workflow:image-anime', - 'workflow:image-pixel', ] as const; export type PromptKey = (typeof promptKeys)[number]; diff --git a/packages/frontend/core/src/blocksuite/ai/provider/setup-provider.tsx b/packages/frontend/core/src/blocksuite/ai/provider/setup-provider.tsx index f2ea6ef12..3263d26aa 100644 --- a/packages/frontend/core/src/blocksuite/ai/provider/setup-provider.tsx +++ b/packages/frontend/core/src/blocksuite/ai/provider/setup-provider.tsx @@ -27,18 +27,18 @@ function toAIUserInfo(account: AuthAccountInfo | null) { const filterStyleToPromptName = new Map( Object.entries({ - 'Clay style': 'workflow:image-clay', - 'Pixel style': 'workflow:image-pixel', - 'Sketch style': 'workflow:image-sketch', - 'Anime style': 'workflow:image-anime', + 'Clay style': 'Convert to Clay style', + 'Pixel style': 'Convert to Pixel style', + 'Sketch style': 'Convert to Sketch style', + 'Anime style': 'Convert to Anime style', }) ); const processTypeToPromptName = new Map( Object.entries({ - Clearer: 'debug:action:fal-upscaler', - 'Remove background': 'debug:action:fal-remove-bg', - 'Convert to sticker': 'debug:action:fal-face-to-sticker', + Clearer: 'Upscale image', + 'Remove background': 'Remove background', + 'Convert to sticker': 'Convert to sticker', }) ); @@ -486,22 +486,18 @@ Could you make a new website based on these notes and send back just the html fi }); AIProvider.provide('createImage', async options => { - // test to image - let promptName: PromptKey = 'debug:action:gpt-image-1'; - // image to image - if (options.attachments?.length) { - promptName = 'debug:action:fal-sd15'; - } - const sessionId = await createSession({ - promptName, + promptName: 'Generate image', ...options, }); return toImage({ ...options, client, sessionId, - content: options.input, + content: + !options.input && options.attachments + ? 'Make the image more detailed.' + : options.input, // 5 minutes timeout: 300000, });