From 0f7c297e36c7a3bbb25967973b4536b944a5b61d Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 22 May 2026 14:19:59 +0000 Subject: [PATCH] Studio: render generated images inline for the Images pill The pill wired the request end of the loop but the response was lost on the client: the backend emits a `tool_end` _toolEvent carrying the base64 PNG on `image_b64` / `image_mime`, but the chat-adapter only read the `result` string and the generic ToolFallback printed the prompt as JSON args with an empty Result block -- the "I see no image" symptom in the chat. - chat-adapter: when the closing `tool_end` is for `image_generation`, repackage `image_b64` + `image_mime` (+ size/quality/background) into a structured result object instead of dropping them. - New `ImageGenerationToolUI` reads that result and renders the image inline via `` with the prompt as a caption. Falls back to a spinner while the request is still running. - Register the component under `image_generation` in thread.tsx's tools.by_name map so it preempts ToolFallback for this tool only. --- .../src/components/assistant-ui/thread.tsx | 2 + .../assistant-ui/tool-ui-image-generation.tsx | 132 ++++++++++++++++++ .../src/features/chat/api/chat-adapter.ts | 32 ++++- 3 files changed, 164 insertions(+), 2 deletions(-) create mode 100644 studio/frontend/src/components/assistant-ui/tool-ui-image-generation.tsx diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx index 5baf0f62da..431e568205 100644 --- a/studio/frontend/src/components/assistant-ui/thread.tsx +++ b/studio/frontend/src/components/assistant-ui/thread.tsx @@ -18,6 +18,7 @@ import { import { ToolFallback } from "@/components/assistant-ui/tool-fallback"; import { ToolGroup } from "@/components/assistant-ui/tool-group"; import { CodeExecutionToolUI } from "@/components/assistant-ui/tool-ui-code-execution"; +import { ImageGenerationToolUI } from "@/components/assistant-ui/tool-ui-image-generation"; import { PythonToolUI } from "@/components/assistant-ui/tool-ui-python"; import { TerminalToolUI } from "@/components/assistant-ui/tool-ui-terminal"; import { WebSearchToolUI } from "@/components/assistant-ui/tool-ui-web-search"; @@ -1105,6 +1106,7 @@ const AssistantMessage: FC = () => { python: PythonToolUI, terminal: TerminalToolUI, code_execution: CodeExecutionToolUI, + image_generation: ImageGenerationToolUI, }, Fallback: ToolFallback, }, diff --git a/studio/frontend/src/components/assistant-ui/tool-ui-image-generation.tsx b/studio/frontend/src/components/assistant-ui/tool-ui-image-generation.tsx new file mode 100644 index 0000000000..246d8b6978 --- /dev/null +++ b/studio/frontend/src/components/assistant-ui/tool-ui-image-generation.tsx @@ -0,0 +1,132 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"use client"; + +import { type ToolCallMessagePartComponent, useAuiState } from "@assistant-ui/react"; +import { ImageIcon, LoaderIcon } from "lucide-react"; +import { memo, useEffect, useState } from "react"; +import { + ToolFallbackContent, + ToolFallbackRoot, + ToolFallbackTrigger, +} from "./tool-fallback"; + +/** + * Renders the synthetic `_toolEvent` chunks emitted by + * `_stream_openai_responses` when OpenAI's Responses-API + * `image_generation` tool fires. The backend stashes the base64 + * PNG/WebP/JPEG (the gpt-image backbone output) on an `image_b64` + * field of the tool_end event so the JSON result stays small, and the + * adapter repackages it into a structured `result` shape: + * + * { + * image_b64: string, + * image_mime: string, // e.g. "image/png" + * size?: string, // "1024x1024" etc + * quality?: string, + * background?: string, + * } + * + * The corresponding `tool_start` carries the prompt as + * `args.prompt` (after gpt-image's revision pass) plus `args.kind: + * "image"`. Without this component the generic ToolFallback would + * print the prompt as JSON args text with an empty Result block -- + * which is exactly the "no image" symptom users hit before this UI + * landed. + */ +interface ImageGenerationArgs { + prompt?: string; + kind?: string; +} + +interface ImageGenerationResult { + image_b64?: string; + image_mime?: string; + size?: string; + quality?: string; + background?: string; +} + +const ImageGenerationToolUIImpl: ToolCallMessagePartComponent = ({ + args, + result, + status, +}) => { + const parsedArgs = (args as ImageGenerationArgs) ?? {}; + const prompt = parsedArgs.prompt ?? ""; + const isRunning = status?.type === "running"; + + const isImageResult = + !!result && + typeof result === "object" && + typeof (result as ImageGenerationResult).image_b64 === "string"; + const imageResult = isImageResult ? (result as ImageGenerationResult) : null; + const mime = imageResult?.image_mime || "image/png"; + const imageSrc = imageResult?.image_b64 + ? `data:${mime};base64,${imageResult.image_b64}` + : null; + + // Collapse the card once the model has resumed streaming prose + // after the image. Mirrors CodeExecutionToolUI so the inline image + // doesn't collapse mid-stream and the user can click to re-expand. + const hasText = useAuiState(({ message }) => + message.content.some( + (p) => + p.type === "text" && + "text" in p && + (p as { text: string }).text.length > 0, + ), + ); + const [open, setOpen] = useState(true); + useEffect(() => { + if (isRunning) { + setOpen(true); + } else if (hasText && !imageSrc) { + setOpen(false); + } + }, [isRunning, hasText, imageSrc]); + + const runningLabel = "Generating image…"; + const completedLabel = prompt + ? prompt.length > 80 + ? `Generated image: ${prompt.slice(0, 80)}…` + : `Generated image: ${prompt}` + : "Generated image"; + + return ( + + + + {isRunning && !imageSrc ? ( +
+ + {runningLabel} +
+ ) : imageSrc ? ( +
+ {prompt + {prompt ? ( +
+ {prompt} +
+ ) : null} +
+ ) : null} + + + ); +}; + +export const ImageGenerationToolUI = memo( + ImageGenerationToolUIImpl, +) as unknown as ToolCallMessagePartComponent; +ImageGenerationToolUI.displayName = "ImageGenerationToolUI"; diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index d2fb4bb220..0521ca2e10 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -1610,8 +1610,36 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { const imgIdx = rawResult.lastIndexOf(imgMarker); let parsedResult: | string - | { text: string; images: string[]; sessionId: string }; - if (imgIdx !== -1) { + | { text: string; images: string[]; sessionId: string } + | { + image_b64: string; + image_mime: string; + size?: string; + quality?: string; + background?: string; + }; + const imageB64 = toolEvent.image_b64 as string | undefined; + if ( + toolCallParts[idx].toolName === "image_generation" && + typeof imageB64 === "string" && + imageB64 + ) { + // OpenAI Responses image_generation_call: the + // backend stashes the base64 PNG/WebP/JPEG on + // separate `image_b64` / `image_mime` fields on + // the synthetic _toolEvent so the JSON result + // string stays small enough to log. Repackage as + // a structured result for the dedicated tool UI. + parsedResult = { + image_b64: imageB64, + image_mime: + (toolEvent.image_mime as string | undefined) ?? + "image/png", + size: toolEvent.size as string | undefined, + quality: toolEvent.quality as string | undefined, + background: toolEvent.background as string | undefined, + }; + } else if (imgIdx !== -1) { const text = rawResult.slice(0, imgIdx); // Fall back to "_default" to match the backend sandbox directory // used when no session_id is provided (see tools.py _get_workdir).