diff --git a/src/renderer/src/providers/GeminiProvider.ts b/src/renderer/src/providers/GeminiProvider.ts index 66dea24d..449e405c 100644 --- a/src/renderer/src/providers/GeminiProvider.ts +++ b/src/renderer/src/providers/GeminiProvider.ts @@ -731,6 +731,7 @@ export default class GeminiProvider extends BaseProvider { onChunk({ text, generateImage: { + type: 'base64', images }, usage: { diff --git a/src/renderer/src/providers/OpenAIProvider.ts b/src/renderer/src/providers/OpenAIProvider.ts index fb145b83..787f26c4 100644 --- a/src/renderer/src/providers/OpenAIProvider.ts +++ b/src/renderer/src/providers/OpenAIProvider.ts @@ -27,6 +27,7 @@ import { Suggestion } from '@renderer/types' import { removeSpecialCharactersForTopicName } from '@renderer/utils' +import { addImageFileToContents } from '@renderer/utils/formats' import { callMCPTool, mcpToolsToOpenAITools, @@ -354,7 +355,7 @@ export default class OpenAIProvider extends BaseProvider { const defaultModel = getDefaultModel() const model = assistant.model || defaultModel const { contextCount, maxTokens, streamOutput } = getAssistantSettings(assistant) - + messages = addImageFileToContents(messages) let systemMessage = assistant.prompt ? { role: 'system', content: assistant.prompt } : undefined if (isOpenAIoSeries(model)) { diff --git a/src/renderer/src/services/ApiService.ts b/src/renderer/src/services/ApiService.ts index 5cecf713..db211552 100644 --- a/src/renderer/src/services/ApiService.ts +++ b/src/renderer/src/services/ApiService.ts @@ -5,6 +5,7 @@ import store from '@renderer/store' import { setGenerating } from '@renderer/store/runtime' import { Assistant, MCPTool, Message, Model, Provider, Suggestion } from '@renderer/types' import { formatMessageError, isAbortError } from '@renderer/utils/error' +import { withGenerateImage } from '@renderer/utils/formats' import { cloneDeep, findLast, isEmpty } from 'lodash' import AiProvider from '../providers/AiProvider' @@ -156,6 +157,7 @@ export async function fetchChatCompletion({ }) message.status = 'success' + message = withGenerateImage(message) if (!message.usage || !message?.usage?.completion_tokens) { message.usage = await estimateMessagesUsage({ @@ -191,7 +193,6 @@ export async function fetchChatCompletion({ // Reset generating state store.dispatch(setGenerating(false)) - return message } diff --git a/src/renderer/src/types/index.ts b/src/renderer/src/types/index.ts index 84b23661..94039fef 100644 --- a/src/renderer/src/types/index.ts +++ b/src/renderer/src/types/index.ts @@ -308,6 +308,7 @@ export type GenerateImageParams = { } export type GenerateImageResponse = { + type: 'url' | 'base64' images: string[] } diff --git a/src/renderer/src/utils/formats.ts b/src/renderer/src/utils/formats.ts index 4bf0afb8..f08bb131 100644 --- a/src/renderer/src/utils/formats.ts +++ b/src/renderer/src/utils/formats.ts @@ -178,3 +178,56 @@ export function withMessageThought(message: Message) { return message } + +export function withGenerateImage(message: Message) { + const imagePattern = new RegExp(`!\\[[^\\]]*\\]\\((.*?)\\s*("(?:.*[^"])")?\\s*\\)`) + const imageMatches = message.content.match(imagePattern) + + if (!imageMatches || imageMatches[1] === null) { + return message + } + + const cleanImgContent = message.content + .replace(imagePattern, '') + .replace(/\n\s*\n/g, '\n') + .trim() + + const downloadPattern = new RegExp(`\\[[^\\]]*\\]\\((.*?)\\s*("(?:.*[^"])")?\\s*\\)`) + const downloadMatches = cleanImgContent.match(downloadPattern) + + let cleanContent = cleanImgContent + if (downloadMatches) { + cleanContent = cleanImgContent + .replace(downloadPattern, '') + .replace(/\n\s*\n/g, '\n') + .trim() + } + + message = { + ...message, + content: cleanContent, + metadata: { + ...message.metadata, + generateImage: { + type: 'url', + images: [imageMatches[1]] + } + } + } + return message +} + +export function addImageFileToContents(messages: Message[]) { + const lastAssistantMessage = messages.findLast((m) => m.role === 'assistant') + if (!lastAssistantMessage || !lastAssistantMessage.metadata || !lastAssistantMessage.metadata.generateImage) { + return messages + } + + const imageFiles = lastAssistantMessage.metadata.generateImage.images + const updatedAssistantMessage = { + ...lastAssistantMessage, + images: imageFiles + } + + return messages.map((message) => (message.role === 'assistant' ? updatedAssistantMessage : message)) +}