From 7e0ee4a7190867d22de16a2e20131bdca0c0cf97 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 22 May 2026 07:08:42 -0700 Subject: [PATCH] Studio: surface OpenAI image_generation as composer Images pill (#5699) The backend already wires OpenAI's Responses-API image_generation server tool: when `enabled_tools` carries "image_generation" on an OpenAI cloud request, _stream_openai_responses appends `{type: "image_generation"}` to the request's tools array and emits `image_generation_call` output items back to the assistant stream (see backend/core/inference/external_provider.py and backend/tests/test_openai_image_generation.py for the round-trip). This wires the frontend half so a user can actually opt into it from the composer next to the Search and Code pills, instead of the tool sitting dormant. - `providerSupportsBuiltinImageGeneration` gates on OpenAI cloud (`api.openai.com`) + a Responses-API model prefix (gpt-5.x, o3). Mirror of the backend's `is_openai_cloud` guard so the pill is hidden on custom OpenAI-compat backends (ollama / llama.cpp / vLLM) that report `provider_type="openai"` but would 400 on the tool. - New `imageToolsEnabled` flag in chat-runtime-store, persisted under `unsloth_chat_image_tools_enabled` and reset on model change in chat-page exactly like `codeToolsEnabled`. - `chat-adapter` appends "image_generation" to `enabled_tools` and flips `enable_tools: true` when the pill is on, so the existing backend dispatch picks it up. - Composer renders an Images pill (lucide `ImageIcon`) immediately after the Code pill, only when the active model advertises the capability. The in-thread composer (assistant-ui/thread.tsx) gets the matching `ImagesToggle` for parity. --- .../src/components/assistant-ui/thread.tsx | 38 ++++++++++++++++++ .../src/features/chat/api/chat-adapter.ts | 27 ++++++++++++- .../frontend/src/features/chat/chat-page.tsx | 40 +++++++++++++++++-- .../features/chat/provider-capabilities.ts | 38 ++++++++++++++++++ .../src/features/chat/shared-composer.tsx | 34 +++++++++++++++- .../chat/stores/chat-runtime-store.ts | 20 ++++++++++ 6 files changed, 190 insertions(+), 7 deletions(-) diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx index c0d577b12d..5baf0f62da 100644 --- a/studio/frontend/src/components/assistant-ui/thread.tsx +++ b/studio/frontend/src/components/assistant-ui/thread.tsx @@ -68,6 +68,7 @@ import { DownloadIcon, GlobeIcon, HeadphonesIcon, + ImageIcon, LightbulbIcon, LightbulbOffIcon, MicIcon, @@ -870,6 +871,42 @@ const CodeToolsToggle: FC = () => { ); }; +const ImagesToggle: FC = () => { + const modelLoaded = useChatRuntimeStore( + (s) => !!s.params.checkpoint && !s.modelLoading, + ); + // OpenAI cloud Responses-API models advertise image_generation as a + // server-side tool; no local runtime fallback exists. Mirror of + // shared-composer's imageDisabled / showImagePill so the in-thread + // composer surfaces the same control as the empty-state composer. + const supportsBuiltinImageGeneration = useChatRuntimeStore( + (s) => s.supportsBuiltinImageGeneration, + ); + const imageToolsEnabled = useChatRuntimeStore((s) => s.imageToolsEnabled); + const setImageToolsEnabled = useChatRuntimeStore( + (s) => s.setImageToolsEnabled, + ); + if (!supportsBuiltinImageGeneration) { + return null; + } + const disabled = !modelLoaded; + return ( + + ); +}; + const ToolStatusDisplay: FC = () => { const toolStatus = useChatRuntimeStore((s) => s.toolStatus); const isThreadRunning = useAuiState(({ thread }) => thread.isRunning); @@ -941,6 +978,7 @@ const ComposerAction: FC<{ disabled?: boolean; blockSend?: () => boolean }> = ({ +
diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index c214c17c0a..d2fb4bb220 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -23,6 +23,7 @@ import { getExternalReasoningCapabilities, getProviderCapabilities, providerSupportsBuiltinCodeExecution, + providerSupportsBuiltinImageGeneration, providerSupportsBuiltinWebFetch, providerSupportsBuiltinWebSearch, } from "../provider-capabilities"; @@ -830,7 +831,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { // Re-read store after potential auto-load / model ready wait runtime = useChatRuntimeStore.getState(); const { params } = runtime; - const { supportsTools, toolsEnabled, codeToolsEnabled } = runtime; + const { supportsTools, toolsEnabled, codeToolsEnabled, imageToolsEnabled } = runtime; const externalSelection = parseExternalModelId(params.checkpoint); const isExternalRequest = externalSelection !== null; if ( @@ -900,6 +901,20 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { externalProvider && providerSupportsBuiltinWebFetch(externalProvider.providerType), ); + // OpenAI Responses-API image_generation server tool. Pill is + // gated on OpenAI cloud + a Responses-API model id; the backend + // additionally re-checks is_openai_cloud before appending + // {type:"image_generation"} to the request tools array. + const imageGenerationEnabledForThisTurn = Boolean( + externalProvider && + externalSelection && + imageToolsEnabled && + providerSupportsBuiltinImageGeneration( + externalProvider.providerType, + externalSelection.modelId, + externalProvider.baseUrl, + ), + ); const outboundMessages = messages .map(toOpenAIMessage) @@ -1396,7 +1411,8 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { // body["tools"] inside _stream_anthropic. ...(webSearchEnabledForThisTurn || webFetchEnabledForThisTurn || - codeExecEnabledForThisTurn + codeExecEnabledForThisTurn || + imageGenerationEnabledForThisTurn ? { enable_tools: true, enabled_tools: [ @@ -1410,6 +1426,13 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { // tool. There is no separate UI toggle yet. ...(webFetchEnabledForThisTurn ? ["web_fetch"] : []), ...(codeExecEnabledForThisTurn ? ["code_execution"] : []), + // OpenAI Responses-API only: `image_generation` + // returns inline image_generation_call output + // items; the backend's _stream_openai_responses + // path translates them to assistant tool events. + ...(imageGenerationEnabledForThisTurn + ? ["image_generation"] + : []), ], } : {}), diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx index 2f062568dd..0943e26313 100644 --- a/studio/frontend/src/features/chat/chat-page.tsx +++ b/studio/frontend/src/features/chat/chat-page.tsx @@ -56,6 +56,7 @@ import { getExternalReasoningCapabilities, getProviderCapabilities, providerSupportsBuiltinCodeExecution, + providerSupportsBuiltinImageGeneration, providerSupportsBuiltinWebSearch, } from "./provider-capabilities"; import { ChatRuntimeProvider } from "./runtime-provider"; @@ -68,6 +69,7 @@ import { } from "./shared-composer"; import { CHAT_CODE_TOOLS_ENABLED_KEY, + CHAT_IMAGE_TOOLS_ENABLED_KEY, CHAT_TOOLS_ENABLED_KEY, loadOptionalBool, useChatRuntimeStore, @@ -771,6 +773,12 @@ export function ChatPage(): ReactElement { selection.modelId, provider?.baseUrl, ); + const supportsBuiltinImageGeneration = + providerSupportsBuiltinImageGeneration( + provider?.providerType, + selection.modelId, + provider?.baseUrl, + ); // Kimi's k2.6/k2.5 default to thinking enabled on the server side // (per https://platform.kimi.ai/docs/models). Mirror that default // in the UI so the Think pill comes up clicked when the user picks @@ -790,6 +798,9 @@ export function ChatPage(): ReactElement { provider?.providerType === "openai"); const storedToolsEnabled = loadOptionalBool(CHAT_TOOLS_ENABLED_KEY); const storedCodeToolsEnabled = loadOptionalBool(CHAT_CODE_TOOLS_ENABLED_KEY); + const storedImageToolsEnabled = loadOptionalBool( + CHAT_IMAGE_TOOLS_ENABLED_KEY, + ); const nextToolsEnabled = supportsBuiltinWebSearch ? isKimi ? false @@ -811,19 +822,25 @@ export function ChatPage(): ReactElement { : state.reasoningEnabled, supportsPreserveThinking: false, // External models never give us a local tool runtime (no - // python sandbox), so `supportsTools` must be false. The two + // python sandbox), so `supportsTools` must be false. The three // `supportsBuiltin*` flags pick up the slack for providers that // run the tool server-side: `supportsBuiltinWebSearch` lights // up the Search pill (OpenAI / Anthropic / OpenRouter / Kimi), // `supportsBuiltinCodeExecution` lights up the Code pill - // (Anthropic Claude 4.x only, today). + // (Anthropic Claude 4.x and OpenAI gpt-5.5), and + // `supportsBuiltinImageGeneration` lights up the Images pill + // (OpenAI cloud Responses-API models only). supportsTools: false, supportsBuiltinWebSearch, supportsBuiltinCodeExecution, + supportsBuiltinImageGeneration, toolsEnabled: nextToolsEnabled, codeToolsEnabled: supportsBuiltinCodeExecution ? (storedCodeToolsEnabled ?? false) : false, + imageToolsEnabled: supportsBuiltinImageGeneration + ? (storedImageToolsEnabled ?? false) + : false, }); }, [externalProvidersForChat, inferenceParams.checkpoint]); const canCompare = useMemo(() => { @@ -988,6 +1005,12 @@ export function ChatPage(): ReactElement { selectedExternal?.modelId, selectedProvider?.baseUrl, ); + const supportsBuiltinImageGeneration = + providerSupportsBuiltinImageGeneration( + selectedProvider?.providerType, + selectedExternal?.modelId, + selectedProvider?.baseUrl, + ); // See sibling useEffect above: Kimi's k2.x default to thinking // enabled, so the Think pill comes up clicked. Search pill stays // off by default; mutual exclusion flips them via the composer. @@ -1003,6 +1026,9 @@ export function ChatPage(): ReactElement { const storedCodeToolsEnabled = loadOptionalBool( CHAT_CODE_TOOLS_ENABLED_KEY, ); + const storedImageToolsEnabled = loadOptionalBool( + CHAT_IMAGE_TOOLS_ENABLED_KEY, + ); const nextToolsEnabled = supportsBuiltinWebSearch ? isKimi ? false @@ -1029,18 +1055,24 @@ export function ChatPage(): ReactElement { : store.reasoningEnabled, supportsPreserveThinking: false, // External models have no local tool runtime → supportsTools - // stays false. The two supportsBuiltin* flags carry the + // stays false. The three supportsBuiltin* flags carry the // server-side capability info for each pill: // - Search → providerSupportsBuiltinWebSearch // - Code → providerSupportsBuiltinCodeExecution - // (Anthropic Claude 4.x only, today) + // (Anthropic Claude 4.x + OpenAI gpt-5.5) + // - Images → providerSupportsBuiltinImageGeneration + // (OpenAI cloud Responses-API models) supportsTools: false, supportsBuiltinWebSearch, supportsBuiltinCodeExecution, + supportsBuiltinImageGeneration, toolsEnabled: nextToolsEnabled, codeToolsEnabled: supportsBuiltinCodeExecution ? (storedCodeToolsEnabled ?? false) : false, + imageToolsEnabled: supportsBuiltinImageGeneration + ? (storedImageToolsEnabled ?? false) + : false, ...(stillOnOpenRouterFree ? {} : { lastOpenRouterChosenModel: null }), }); return; diff --git a/studio/frontend/src/features/chat/provider-capabilities.ts b/studio/frontend/src/features/chat/provider-capabilities.ts index 8b16c2c85f..da1d6e3431 100644 --- a/studio/frontend/src/features/chat/provider-capabilities.ts +++ b/studio/frontend/src/features/chat/provider-capabilities.ts @@ -217,6 +217,44 @@ export function providerSupportsBuiltinCodeExecution( return false; } +/** + * Whether the selected external provider/model exposes OpenAI's + * Responses-API server-side image_generation tool. Lit on for OpenAI + * cloud (`api.openai.com`) when the picked model is a Responses-API + * family id (gpt-5.x today). The backend additionally gates on + * `is_openai_cloud`; mirror that here so the pill is hidden on custom + * OpenAI-compat backends (ollama / llama.cpp / vLLM) that report + * `provider_type="openai"` but would 400 on a `{type:"image_generation"}` + * tool. See backend/core/inference/external_provider.py near line 2770 + * for the dispatch and backend/tests/test_openai_image_generation.py + * for the round-trip coverage. + */ +const OPENAI_IMAGE_GENERATION_MODEL_PREFIXES = [ + "gpt-5.5-pro", + "gpt-5.5", + "gpt-5.4-pro", + "gpt-5.4", + "gpt-5.3", + "gpt-5.2", + "gpt-5.1", + "gpt-5", + "o3", +] as const; + +export function providerSupportsBuiltinImageGeneration( + providerType: string | null | undefined, + modelId: string | null | undefined, + baseUrl?: string | null, +): boolean { + if (providerType !== "openai") return false; + if (!isOpenAICloudBaseUrl(baseUrl)) return false; + const normalized = modelId?.trim().toLowerCase() ?? ""; + if (!normalized) return false; + return OPENAI_IMAGE_GENERATION_MODEL_PREFIXES.some((prefix) => + normalized.startsWith(prefix), + ); +} + /** * Per-provider minimum on the outbound max_tokens. Kimi's docs require * `max_tokens >= 16000` whenever a thinking model is in use so the diff --git a/studio/frontend/src/features/chat/shared-composer.tsx b/studio/frontend/src/features/chat/shared-composer.tsx index 57a8a132d7..246fd81510 100644 --- a/studio/frontend/src/features/chat/shared-composer.tsx +++ b/studio/frontend/src/features/chat/shared-composer.tsx @@ -21,7 +21,7 @@ import { isTauri } from "@/lib/api-base"; import { isMultimodalResponse } from "./types/api"; import { getImageInputUnavailableReason } from "./utils/image-input-support"; import { useAui } from "@assistant-ui/react"; -import { ArrowUpIcon, GlobeIcon, HeadphonesIcon, LightbulbIcon, LightbulbOffIcon, MicIcon, PlusIcon, SquareIcon, XIcon } from "lucide-react"; +import { ArrowUpIcon, GlobeIcon, HeadphonesIcon, ImageIcon, LightbulbIcon, LightbulbOffIcon, MicIcon, PlusIcon, SquareIcon, XIcon } from "lucide-react"; import { toast } from "@/lib/toast"; import { loadModel, validateModel } from "./api/chat-api"; import { parseExternalModelId, providerTypeSupportsVision } from "./external-providers"; @@ -33,6 +33,7 @@ import { import { getExternalReasoningCapabilities, providerSupportsBuiltinCodeExecution, + providerSupportsBuiltinImageGeneration, } from "./provider-capabilities"; import { type CompositionEvent, @@ -331,6 +332,10 @@ export function SharedComposer({ const setToolsEnabled = useChatRuntimeStore((s) => s.setToolsEnabled); const codeToolsEnabled = useChatRuntimeStore((s) => s.codeToolsEnabled); const setCodeToolsEnabled = useChatRuntimeStore((s) => s.setCodeToolsEnabled); + const imageToolsEnabled = useChatRuntimeStore((s) => s.imageToolsEnabled); + const setImageToolsEnabled = useChatRuntimeStore( + (s) => s.setImageToolsEnabled, + ); const lastOpenRouterChosenModel = useChatRuntimeStore( (s) => s.lastOpenRouterChosenModel, ); @@ -416,10 +421,22 @@ export function SharedComposer({ effectiveExternalModelId, selectedExternalProvider?.baseUrl, ); + const supportsBuiltinImageGeneration = providerSupportsBuiltinImageGeneration( + selectedExternalProvider?.providerType, + effectiveExternalModelId, + selectedExternalProvider?.baseUrl, + ); const searchDisabled = !modelLoaded || !(supportsTools || supportsBuiltinWebSearch); const codeDisabled = !modelLoaded || !(supportsTools || supportsBuiltinCodeExecution); + // Images pill is only ever lit on OpenAI cloud's Responses-API models. + // No local tool runtime fallback because the only image-generation + // server tool we wire today is OpenAI's; local models cannot dispatch + // it. Hidden entirely when the active model does not advertise it so + // the pill row stays compact for providers without the capability. + const imageDisabled = !modelLoaded || !supportsBuiltinImageGeneration; + const showImagePill = supportsBuiltinImageGeneration; // Backwards-compatible alias for any other call site that may still // reference `toolsDisabled` (rare; both pills used it before). const toolsDisabled = codeDisabled; @@ -1074,6 +1091,21 @@ export function SharedComposer({ Code + {showImagePill && ( + + )}
{dictationSupported && ( diff --git a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts index 1f7442ccd2..a2cc547548 100644 --- a/studio/frontend/src/features/chat/stores/chat-runtime-store.ts +++ b/studio/frontend/src/features/chat/stores/chat-runtime-store.ts @@ -24,6 +24,8 @@ const HF_TOKEN_KEY = "unsloth_hf_token"; export const CHAT_REASONING_ENABLED_KEY = "unsloth_chat_reasoning_enabled"; export const CHAT_TOOLS_ENABLED_KEY = "unsloth_chat_tools_enabled"; export const CHAT_CODE_TOOLS_ENABLED_KEY = "unsloth_chat_code_tools_enabled"; +export const CHAT_IMAGE_TOOLS_ENABLED_KEY = "unsloth_chat_image_tools_enabled"; + // External provider selection is encoded into `params.checkpoint` as // `external::::`. PersistedChatSettings deliberately // Omits `checkpoint` because the local-model side is mirrored by the @@ -253,8 +255,16 @@ type ChatRuntimeStore = { * execution server-side. Read by both composers' Code pill gate. */ supportsBuiltinCodeExecution: boolean; + /** + * Whether the active external provider exposes a server-side + * image-generation tool (OpenAI's Responses-API `image_generation` + * today). Gates the chat composer's Images pill. Local models never + * receive the tool because their runtime cannot dispatch it. + */ + supportsBuiltinImageGeneration: boolean; toolsEnabled: boolean; codeToolsEnabled: boolean; + imageToolsEnabled: boolean; toolStatus: string | null; generatingStatus: string | null; autoHealToolCalls: boolean; @@ -313,6 +323,7 @@ type ChatRuntimeStore = { setPreserveThinking: (value: boolean) => void; setToolsEnabled: (enabled: boolean, options?: { persist?: boolean }) => void; setCodeToolsEnabled: (enabled: boolean) => void; + setImageToolsEnabled: (enabled: boolean) => void; setToolStatus: (status: string | null) => void; setGeneratingStatus: (status: string | null) => void; setAutoHealToolCalls: (enabled: boolean) => void; @@ -552,8 +563,10 @@ export const useChatRuntimeStore = create((set, get) => ({ supportsTools: false, supportsBuiltinWebSearch: false, supportsBuiltinCodeExecution: false, + supportsBuiltinImageGeneration: false, toolsEnabled: loadBool(CHAT_TOOLS_ENABLED_KEY, false), codeToolsEnabled: loadBool(CHAT_CODE_TOOLS_ENABLED_KEY, false), + imageToolsEnabled: loadBool(CHAT_IMAGE_TOOLS_ENABLED_KEY, false), toolStatus: null, generatingStatus: null, autoHealToolCalls: true, @@ -741,8 +754,10 @@ export const useChatRuntimeStore = create((set, get) => ({ supportsTools: false, supportsBuiltinWebSearch: false, supportsBuiltinCodeExecution: false, + supportsBuiltinImageGeneration: false, toolsEnabled: false, codeToolsEnabled: false, + imageToolsEnabled: false, toolStatus: null, kvCacheDtype: null, loadedKvCacheDtype: null, @@ -797,6 +812,11 @@ export const useChatRuntimeStore = create((set, get) => ({ saveBool(CHAT_CODE_TOOLS_ENABLED_KEY, codeToolsEnabled); return { codeToolsEnabled }; }), + setImageToolsEnabled: (imageToolsEnabled) => + set(() => { + saveBool(CHAT_IMAGE_TOOLS_ENABLED_KEY, imageToolsEnabled); + return { imageToolsEnabled }; + }), setToolStatus: (toolStatus) => set({ toolStatus }), setGeneratingStatus: (generatingStatus) => set({ generatingStatus }), setAutoHealToolCalls: (autoHealToolCalls) =>