diff --git a/frontend/src/components/WorkspaceChat/ChatContainer/PromptInput/ToolsMenu/Tabs/AgentSkills/index.jsx b/frontend/src/components/WorkspaceChat/ChatContainer/PromptInput/ToolsMenu/Tabs/AgentSkills/index.jsx index 6b87501c649..fd4b67203f8 100644 --- a/frontend/src/components/WorkspaceChat/ChatContainer/PromptInput/ToolsMenu/Tabs/AgentSkills/index.jsx +++ b/frontend/src/components/WorkspaceChat/ChatContainer/PromptInput/ToolsMenu/Tabs/AgentSkills/index.jsx @@ -33,6 +33,7 @@ export default function AgentSkillsTab({ // All skill state management const { fileSystemAgentAvailable, + imageGenerationAvailable, importedSkills, flows, mcpServers, @@ -51,6 +52,7 @@ export default function AgentSkillsTab({ const configurableSkills = getConfigurableSkills(t, { fileSystemAgentAvailable, + imageGenerationAvailable, }); // UI state diff --git a/frontend/src/components/WorkspaceChat/ChatContainer/PromptInput/ToolsMenu/Tabs/AgentSkills/useAgentSkillsState.js b/frontend/src/components/WorkspaceChat/ChatContainer/PromptInput/ToolsMenu/Tabs/AgentSkills/useAgentSkillsState.js index 861a5930543..80a44067ca4 100644 --- a/frontend/src/components/WorkspaceChat/ChatContainer/PromptInput/ToolsMenu/Tabs/AgentSkills/useAgentSkillsState.js +++ b/frontend/src/components/WorkspaceChat/ChatContainer/PromptInput/ToolsMenu/Tabs/AgentSkills/useAgentSkillsState.js @@ -16,6 +16,8 @@ export default function useAgentSkillsState(defaultSkills) { // Core skill state const [fileSystemAgentAvailable, setFileSystemAgentAvailable] = useState(false); + const [imageGenerationAvailable, setImageGenerationAvailable] = + useState(false); const [isMultiUser, setIsMultiUser] = useState(false); const [disabledDefaults, setDisabledDefaults] = useState([]); const [enabledConfigurable, setEnabledConfigurable] = useState([]); @@ -37,7 +39,7 @@ export default function useAgentSkillsState(defaultSkills) { async function fetchSkillSettings() { try { const subSkillPrefKeys = getSubSkillPreferenceKeys(); - const [prefs, flowsRes, fsAgentAvailable, multiUserMode] = + const [prefs, flowsRes, fsAgentAvailable, multiUserMode, keys] = await Promise.all([ Admin.systemPreferencesByFields([ "disabled_agent_skills", @@ -48,6 +50,7 @@ export default function useAgentSkillsState(defaultSkills) { AgentFlows.listFlows(), System.isFileSystemAgentAvailable(), System.isMultiUserMode(), + System.keys(), ]); if (prefs?.settings) { @@ -58,6 +61,7 @@ export default function useAgentSkillsState(defaultSkills) { } if (flowsRes?.flows) setFlows(flowsRes.flows); setFileSystemAgentAvailable(fsAgentAvailable); + setImageGenerationAvailable(!!keys?.ImageGenerationProvider); setIsMultiUser(!!multiUserMode); } catch (e) { console.error(e); @@ -172,6 +176,7 @@ export default function useAgentSkillsState(defaultSkills) { return { // State fileSystemAgentAvailable, + imageGenerationAvailable, isMultiUser, disabledDefaults, enabledConfigurable, diff --git a/frontend/src/locales/en/common.js b/frontend/src/locales/en/common.js index 088156c6482..4b6d4cdac6a 100644 --- a/frontend/src/locales/en/common.js +++ b/frontend/src/locales/en/common.js @@ -320,6 +320,11 @@ const TRANSLATIONS = { description: "Enable the default agent to generate various types of charts from data provided or given in chat.", }, + generateImage: { + title: "Generate images", + description: + "Allow the agent to generate images from chat, or edit images attached to the conversation, using your configured image generation provider.", + }, web: { title: "Web Search", description: diff --git a/frontend/src/media/agents/generate-image.png b/frontend/src/media/agents/generate-image.png new file mode 100644 index 00000000000..1708e53a984 Binary files /dev/null and b/frontend/src/media/agents/generate-image.png differ diff --git a/frontend/src/pages/Admin/Agents/skills.jsx b/frontend/src/pages/Admin/Agents/skills.jsx index 4fa99b5d970..0a1e0c4bb7c 100644 --- a/frontend/src/pages/Admin/Agents/skills.jsx +++ b/frontend/src/pages/Admin/Agents/skills.jsx @@ -15,11 +15,13 @@ import { FolderOpen, FilePlus, CalendarCheck, + ImageSquare, } from "@phosphor-icons/react"; import RAGImage from "@/media/agents/rag-memory.png"; import SummarizeImage from "@/media/agents/view-summarize.png"; import ScrapeWebsitesImage from "@/media/agents/scrape-websites.png"; import GenerateChartsImage from "@/media/agents/generate-charts.png"; +import GenerateImageImage from "@/media/agents/generate-image.png"; import GenerateSaveImages from "@/media/agents/generate-save-files.png"; import FileSystemImage from "@/media/agents/file-system.png"; import GMailIcon from "./GMailSkillPanel/gmail.png"; @@ -60,11 +62,18 @@ export const getDefaultSkills = (t) => ({ * @param {object} options - The options for the configurable skills. * @param {boolean} options.fileSystemAgentAvailable - Whether the file system agent is available. * @param {boolean} options.createFilesAgentAvailable - Whether the create files agent is available. + * @param {boolean} options.imageGenerationAvailable - Whether an image generation provider is + * configured. Only the in-chat skills menu passes this - the admin page always lists the skill so + * it can be enabled before a provider is set up. * @returns {object} The configurable skills. */ export const getConfigurableSkills = ( t, - { fileSystemAgentAvailable = true, createFilesAgentAvailable = true } = {} + { + fileSystemAgentAvailable = true, + createFilesAgentAvailable = true, + imageGenerationAvailable = true, + } = {} ) => ({ ...(fileSystemAgentAvailable && { "filesystem-agent": { @@ -94,6 +103,16 @@ export const getConfigurableSkills = ( icon: ChartBar, image: GenerateChartsImage, }, + ...(imageGenerationAvailable && { + "generate-image": { + title: t("agent.skill.generateImage.title"), + description: t("agent.skill.generateImage.description"), + component: GenericSkillPanel, + skill: "generate-image", + icon: ImageSquare, + image: GenerateImageImage, + }, + }), "web-browsing": { title: t("agent.skill.web.title"), description: t("agent.skill.web.description"), diff --git a/frontend/src/utils/chat/agent.js b/frontend/src/utils/chat/agent.js index 7c3c8d3147f..ce63611eb5d 100644 --- a/frontend/src/utils/chat/agent.js +++ b/frontend/src/utils/chat/agent.js @@ -65,6 +65,7 @@ const handledEvents = [ "statusResponse", "fileDownloadCard", "imageGenerationCard", + "imageGenerationPending", "scheduledJobCreated", "awaitingFeedback", "wssFailure", @@ -314,16 +315,41 @@ export default function handleSocketResponse(socket, event, setChatHistory) { }); } - if (data.type === "imageGenerationCard") { + if (data.type === "imageGenerationPending") { return setChatHistory((prev) => { return [ ...prev.filter((msg) => !!msg.content), + { + uuid: data.content.pendingId, + type: "imageGenerationPending", + content: data.content.prompt, + role: "assistant", + sources: [], + closed: false, + error: null, + animate: false, + pending: true, + metrics: {}, + }, + ]; + }); + } + + if (data.type === "imageGenerationCard") { + return setChatHistory((prev) => { + // Drops the placeholder card this result belongs to, if there was one. + const history = prev.filter( + (msg) => !!msg.content && msg.uuid !== data.content.pendingId + ); + if (data.content.failed) return history; + return [ + ...history, { uuid: v4(), type: "textResponse", - content: data.content, - outputs: data.outputs || [], - chatId: data.chatId || null, + content: data.content.text, + outputs: data.content.outputs || [], + chatId: data.content.chatId || null, role: "assistant", sources: [], closed: true, diff --git a/server/__tests__/utils/agents/aibitat/plugins/generate-image.test.js b/server/__tests__/utils/agents/aibitat/plugins/generate-image.test.js new file mode 100644 index 00000000000..56cbcdf94fd --- /dev/null +++ b/server/__tests__/utils/agents/aibitat/plugins/generate-image.test.js @@ -0,0 +1,188 @@ +/* eslint-env jest */ +const fs = require("fs"); +const os = require("os"); +const path = require("path"); + +// `utils/files` resolves its storage paths from STORAGE_DIR at require time. +process.env.STORAGE_DIR = fs.mkdtempSync( + path.join(os.tmpdir(), "generate-image-test-") +); +jest.mock("../../../../../utils/ImageGenerators", () => ({ + generateImageForWorkspace: jest.fn(), + editImageForWorkspace: jest.fn(), +})); +jest.mock("../../../../../utils/helpers", () => ({ + getImageGeneratorProvider: jest.fn(), +})); +jest.mock("../../../../../models/workspaceChats", () => ({ + WorkspaceChats: { _update: jest.fn() }, +})); + +const { + generateImageForWorkspace, + editImageForWorkspace, +} = require("../../../../../utils/ImageGenerators"); +const { WorkspaceChats } = require("../../../../../models/workspaceChats"); +const { + generateImage, +} = require("../../../../../utils/agents/aibitat/plugins/generate-image.js"); + +const SAVED_IMAGE = { + storageFilename: "img-11111111-2222-3333-4444-555555555555.png", + filename: "a-fox.png", + fileSize: 100, + buffer: Buffer.from("image-bytes"), +}; + +beforeAll(() => { + const dir = path.join(process.env.STORAGE_DIR, "generated-images"); + fs.mkdirSync(dir, { recursive: true }); + fs.writeFileSync( + path.join(dir, SAVED_IMAGE.storageFilename), + SAVED_IMAGE.buffer + ); +}); + +afterAll(() => + fs.rmSync(process.env.STORAGE_DIR, { recursive: true, force: true }) +); + +function setupPlugin(chats = []) { + const aibitat = { + chats, + introspect: jest.fn(), + handlerProps: { log: jest.fn() }, + socket: { send: jest.fn() }, + function: (config) => (aibitat._fn = config), + }; + generateImage.plugin.call(generateImage).setup(aibitat); + return { aibitat, handler: aibitat._fn.handler.bind(aibitat._fn) }; +} + +beforeEach(() => jest.clearAllMocks()); + +describe("generate-image agent skill", () => { + test("generates an image and registers the card for live + historical render", async () => { + generateImageForWorkspace.mockResolvedValue(SAVED_IMAGE); + const { aibitat, handler } = setupPlugin(); + + await handler({ prompt: "a red fox in the snow", size: "512x512" }); + + expect(generateImageForWorkspace).toHaveBeenCalledWith( + expect.objectContaining({ + prompt: "a red fox in the snow", + size: "512x512", + }) + ); + const expectedOutput = { + type: "imageGenerationCard", + payload: { + storageFilename: SAVED_IMAGE.storageFilename, + filename: SAVED_IMAGE.filename, + fileSize: SAVED_IMAGE.fileSize, + prompt: "a red fox in the snow", + }, + }; + // The card must only reference the stored file - no image bytes travel over + // the socket or into the chat record. + expect(aibitat._pendingOutputs).toEqual([expectedOutput]); + const [, card] = aibitat.socket.send.mock.calls.find( + ([type]) => type === "imageGenerationCard" + ); + expect(card.outputs).toEqual([expectedOutput]); + }); + + test("writes the image reference to the reserved chat before showing the card", async () => { + generateImageForWorkspace.mockResolvedValue(SAVED_IMAGE); + const { aibitat, handler } = setupPlugin(); + aibitat.trackedChatId = 42; + + await handler({ prompt: "a fox" }); + + // The serve endpoint only authorizes files referenced by a chat record, so + // the reference has to land before the frontend requests the image. + expect(WorkspaceChats._update).toHaveBeenCalledWith(42, { + response: JSON.stringify({ outputs: aibitat._pendingOutputs }), + }); + expect(WorkspaceChats._update.mock.invocationCallOrder[0]).toBeLessThan( + aibitat.socket.send.mock.invocationCallOrder[1] + ); + }); + + test("shows a placeholder card and swaps it for the result", async () => { + generateImageForWorkspace.mockResolvedValue(SAVED_IMAGE); + const { aibitat, handler } = setupPlugin(); + + await handler({ prompt: "a fox" }); + + const [pending, card] = aibitat.socket.send.mock.calls; + expect(pending[0]).toBe("imageGenerationPending"); + expect(card[0]).toBe("imageGenerationCard"); + expect(card[1].pendingId).toBe(pending[1].pendingId); + }); + + test("clears the placeholder card when generation fails", async () => { + generateImageForWorkspace.mockRejectedValue(new Error("provider is down")); + const { aibitat, handler } = setupPlugin(); + + const reply = await handler({ prompt: "a fox" }); + + const [pending, failure] = aibitat.socket.send.mock.calls; + expect(failure[0]).toBe("imageGenerationCard"); + expect(failure[1]).toEqual({ + pendingId: pending[1].pendingId, + text: "provider is down", + failed: true, + }); + expect(reply).toContain("provider is down"); + }); + + test("drops a size the model invented instead of failing the call", async () => { + generateImageForWorkspace.mockResolvedValue(SAVED_IMAGE); + const { handler } = setupPlugin(); + + await handler({ prompt: "a fox", size: "large" }); + + expect(generateImageForWorkspace).toHaveBeenCalledWith( + expect.objectContaining({ size: null }) + ); + }); + + test("edits using the images attached to the last user message", async () => { + editImageForWorkspace.mockResolvedValue(SAVED_IMAGE); + const { handler } = setupPlugin([ + { from: "USER", content: "make it blue" }, + { + from: "USER", + content: "here is my photo", + attachments: [ + { + mime: "image/png", + contentString: `data:image/png;base64,${Buffer.from("source").toString("base64")}`, + }, + ], + }, + ]); + + await handler({ prompt: "make it blue", edit: true }); + + expect(generateImageForWorkspace).not.toHaveBeenCalled(); + expect(editImageForWorkspace).toHaveBeenCalledWith( + expect.objectContaining({ images: [Buffer.from("source")] }) + ); + }); + + test("falls back to the image generated earlier in the session when editing", async () => { + generateImageForWorkspace.mockResolvedValue(SAVED_IMAGE); + editImageForWorkspace.mockResolvedValue(SAVED_IMAGE); + const { aibitat, handler } = setupPlugin(); + + await handler({ prompt: "a fox" }); + await handler({ prompt: "now make it blue", edit: true }); + + expect(aibitat._lastGeneratedImage).toBe(SAVED_IMAGE.storageFilename); + expect(editImageForWorkspace).toHaveBeenCalledWith( + expect.objectContaining({ images: [SAVED_IMAGE.buffer] }) + ); + }); +}); diff --git a/server/endpoints/agentFileServer.js b/server/endpoints/agentFileServer.js index 6e6416b83dd..51ab4123b79 100644 --- a/server/endpoints/agentFileServer.js +++ b/server/endpoints/agentFileServer.js @@ -182,12 +182,16 @@ async function findInWorkspaceChats(storageFilename, { user, isMultiUser }) { // DB-level filter so we don't load every chat into memory. const chats = await WorkspaceChats.where({ workspaceId: { in: workspaceIds }, - include: true, response: { contains: storageFilename }, }); for (const chat of chats) { - const { outputs = [] } = safeJsonParse(chat.response, { outputs: [] }); + const { outputs = [], text } = safeJsonParse(chat.response, { + outputs: [], + }); + // A hidden chat is either an agent reply still being written (no text yet) + // or history the user cleared - only the former can still serve its files. + if (!chat.include && !!text) continue; const output = outputs.find( (o) => o?.payload?.storageFilename === storageFilename ); diff --git a/server/utils/agents/aibitat/plugins/generate-image.js b/server/utils/agents/aibitat/plugins/generate-image.js new file mode 100644 index 00000000000..76859ecc1fc --- /dev/null +++ b/server/utils/agents/aibitat/plugins/generate-image.js @@ -0,0 +1,170 @@ +const { v4: uuidv4 } = require("uuid"); +const { + generateImageForWorkspace, + editImageForWorkspace, +} = require("../../../ImageGenerators"); +const { resolveImageBuffers } = require("../../../chats/commands/img"); +const { WorkspaceChats } = require("../../../../models/workspaceChats"); +const { safeJSONStringify } = require("../../../helpers/chat/responses"); + +/** + * Collects the image buffers an edit request should transform. Prefers the + * images attached to the most recent user message and falls back to the image + * this session generated last, which is not yet in the persisted chat history. + * @param {object} aibitat + * @returns {Buffer[]} + */ +function sourceImages(aibitat) { + const lastUserMessage = [...aibitat.chats] + .reverse() + .find((chat) => chat.from === "USER"); + const images = resolveImageBuffers(lastUserMessage?.attachments || []); + if (images.length || !aibitat._lastGeneratedImage) return images; + return resolveImageBuffers([ + { mime: "image/png", storageFilename: aibitat._lastGeneratedImage }, + ]); +} + +/** + * Writes the outputs collected so far onto the chat row reserved for this reply. + * The image serve endpoint authorizes a request by finding a chat that references + * the file, so that reference has to exist before the card asks for the image - + * the reply itself is not persisted until the agent finishes its turn. + * @param {object} aibitat + */ +async function persistOutputs(aibitat) { + if (!aibitat.trackedChatId) return; + await WorkspaceChats._update(aibitat.trackedChatId, { + response: safeJSONStringify({ outputs: aibitat._pendingOutputs }), + }); +} + +const generateImage = { + name: "generate-image", + startupConfig: { + params: {}, + }, + plugin: function () { + return { + name: this.name, + setup(aibitat) { + aibitat.function({ + super: aibitat, + name: this.name, + description: + "Generate an image from a text prompt, or edit an image already in the conversation. Use for any request to draw, create, render, or modify a picture.", + examples: [ + { + prompt: "Generate an image of a red fox in the snow", + call: JSON.stringify({ prompt: "a red fox in the snow" }), + }, + { + prompt: "Make that image black and white", + call: JSON.stringify({ + prompt: "make the image black and white", + edit: true, + }), + }, + ], + parameters: { + $schema: "http://json-schema.org/draft-07/schema#", + type: "object", + properties: { + prompt: { + type: "string", + description: + "Detailed description of the image to create, or of the change to apply when editing.", + }, + size: { + type: "string", + description: + "Image dimensions as WIDTHxHEIGHT (eg: 512x512, 1024x1024). Omit to use the system default.", + }, + edit: { + type: "boolean", + description: + "Set true to edit an image already in the conversation instead of creating a new one.", + }, + }, + additionalProperties: false, + }, + required: ["prompt"], + handler: async function ({ prompt, size = null, edit = false }) { + const { getImageGeneratorProvider } = require("../../../helpers"); + try { + getImageGeneratorProvider(); + } catch { + this.super.introspect( + `${this.caller}: No image generation provider is configured.` + ); + return "No image generation provider is set up on this instance, so no image could be created. Tell the user to configure one in Settings > Image Generation (an admin must do this) and then try again."; + } + + // Ties the placeholder card to the result so the frontend can swap + // or drop it once generation settles. + const pendingId = uuidv4(); + try { + const images = edit ? sourceImages(this.super) : []; + this.super.introspect( + `${this.caller}: ${images.length ? "Editing image" : "Generating image"} - "${prompt}"` + ); + this.super.socket?.send?.("imageGenerationPending", { + pendingId, + prompt, + }); + + // Models like to invent sizes ("large", "square"), so anything + // that is not WIDTHxHEIGHT falls back to the system default. + if (!/^\d+x\d+$/.test(String(size))) size = null; + + const signal = this.super.abortController?.signal ?? null; + const { storageFilename, filename, fileSize, notice } = + images.length > 0 + ? await editImageForWorkspace({ + prompt, + images, + size, + signal, + }) + : await generateImageForWorkspace({ prompt, size, signal }); + + // Register the card as a pending output so it is saved with the + // reply and re-renders when the chat is reloaded. + const output = { + type: "imageGenerationCard", + payload: { storageFilename, filename, fileSize, prompt }, + }; + if (!Array.isArray(this.super._pendingOutputs)) + this.super._pendingOutputs = []; + this.super._pendingOutputs.push(output); + this.super._lastGeneratedImage = storageFilename; + await persistOutputs(this.super); + + this.super.socket?.send?.("imageGenerationCard", { + pendingId, + text: `Generated an image for: "${prompt}"`, + outputs: [output], + }); + + return `The image was generated and is already displayed to the user.${notice ? ` Note: ${notice}.` : ""} Confirm it is ready in one short sentence - do not describe the image or repeat the prompt.`; + } catch (error) { + this.super.socket?.send?.("imageGenerationCard", { + pendingId, + text: error.message, + failed: true, + }); + const { isAbortError } = require("../../../helpers/abortSignals"); + if (isAbortError(error)) return "Image generation was cancelled."; + this.super.handlerProps.log( + `generate-image raised an error. ${error.message}` + ); + return `Let the user know the image could not be generated. ${error.message}`; + } + }, + }); + }, + }; + }, +}; + +module.exports = { generateImage }; diff --git a/server/utils/agents/aibitat/plugins/index.js b/server/utils/agents/aibitat/plugins/index.js index c5e7f70a94e..ad44cf1af0e 100644 --- a/server/utils/agents/aibitat/plugins/index.js +++ b/server/utils/agents/aibitat/plugins/index.js @@ -5,6 +5,7 @@ const { docSummarizer } = require("./summarize.js"); const { chatHistory } = require("./chat-history.js"); const { memory } = require("./memory.js"); const { rechart } = require("./rechart.js"); +const { generateImage } = require("./generate-image.js"); const { sqlAgent } = require("./sql-agent/index.js"); const { filesystemAgent } = require("./filesystem/index.js"); const { createFilesAgent } = require("./create-files/index.js"); @@ -23,6 +24,7 @@ module.exports = { chatHistory, memory, rechart, + generateImage, sqlAgent, filesystemAgent, createFilesAgent, @@ -41,6 +43,7 @@ module.exports = { [chatHistory.name]: chatHistory, [memory.name]: memory, [rechart.name]: rechart, + [generateImage.name]: generateImage, [sqlAgent.name]: sqlAgent, [filesystemAgent.name]: filesystemAgent, [createFilesAgent.name]: createFilesAgent, diff --git a/server/utils/agents/aibitat/plugins/websocket.js b/server/utils/agents/aibitat/plugins/websocket.js index 0fcdc93ec78..511fcf39a56 100644 --- a/server/utils/agents/aibitat/plugins/websocket.js +++ b/server/utils/agents/aibitat/plugins/websocket.js @@ -90,9 +90,11 @@ async function handleImageCommand({ aibitat, socket, message }) { socket.send( JSON.stringify({ type: "imageGenerationCard", - content: result.textResponse, - outputs: result.outputs || [], - chatId: result.chatId || null, + content: { + text: result.textResponse, + outputs: result.outputs || [], + chatId: result.chatId || null, + }, }) ); diff --git a/server/utils/agents/ephemeral.js b/server/utils/agents/ephemeral.js index e9c63db5671..bec4cbd8b4e 100644 --- a/server/utils/agents/ephemeral.js +++ b/server/utils/agents/ephemeral.js @@ -681,6 +681,14 @@ class EphemeralEventListener extends EventEmitter { continue; } + // Generated images are collected from the agent's pending outputs, so the + // live card event is not a text response and is ignored here. + if ( + msg.type === "imageGenerationCard" || + msg.type === "imageGenerationPending" + ) + continue; + if (msg.type === "reportStreamEvent") { const inner = msg.content; if (inner?.type === "textResponseChunk" && inner?.content) @@ -745,6 +753,14 @@ class EphemeralEventListener extends EventEmitter { }); } + // Generated images travel back in the final response `outputs`, so the + // live card event has nothing to stream. + if ( + data.type === "imageGenerationCard" || + data.type === "imageGenerationPending" + ) + return; + if (data.type === "reportStreamEvent") { const inner = data.content; if (!inner?.type) return; diff --git a/server/utils/chats/commands/img.js b/server/utils/chats/commands/img.js index 8e4811e920e..63caad3a44c 100644 --- a/server/utils/chats/commands/img.js +++ b/server/utils/chats/commands/img.js @@ -190,4 +190,4 @@ function resolveImageBuffers(attachments = []) { return buffers; } -module.exports = { generateImage }; +module.exports = { generateImage, resolveImageBuffers };