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 `<img src="data:image/...;base64,...">` 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.
This commit is contained in:
Daniel Han
parent
18ffbf15f6
commit
0f7c297e36
3 changed files with 164 additions and 2 deletions
|
|
@ -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,
|
||||
},
|
||||
|
|
|
|||
|
|
@ -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 (
|
||||
<ToolFallbackRoot open={open} onOpenChange={setOpen}>
|
||||
<ToolFallbackTrigger
|
||||
toolName={isRunning ? runningLabel : completedLabel}
|
||||
status={status}
|
||||
icon={ImageIcon}
|
||||
/>
|
||||
<ToolFallbackContent>
|
||||
{isRunning && !imageSrc ? (
|
||||
<div className="flex items-center gap-2 text-sm text-muted-foreground">
|
||||
<LoaderIcon className="size-3.5 animate-spin" />
|
||||
<span>{runningLabel}</span>
|
||||
</div>
|
||||
) : imageSrc ? (
|
||||
<figure className="m-0 flex flex-col gap-1.5">
|
||||
<img
|
||||
src={imageSrc}
|
||||
alt={prompt || "Generated image"}
|
||||
className="max-w-full rounded-md border border-border/60"
|
||||
/>
|
||||
{prompt ? (
|
||||
<figcaption className="text-xs leading-snug text-muted-foreground">
|
||||
{prompt}
|
||||
</figcaption>
|
||||
) : null}
|
||||
</figure>
|
||||
) : null}
|
||||
</ToolFallbackContent>
|
||||
</ToolFallbackRoot>
|
||||
);
|
||||
};
|
||||
|
||||
export const ImageGenerationToolUI = memo(
|
||||
ImageGenerationToolUIImpl,
|
||||
) as unknown as ToolCallMessagePartComponent;
|
||||
ImageGenerationToolUI.displayName = "ImageGenerationToolUI";
|
||||
|
|
@ -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).
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue