ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Studio: surface OpenAI image_generation as composer Images pill (#5699)

Browse source 
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.
This commit is contained in:
Daniel Han 2026-05-22 07:08:42 -07:00 • committed by GitHub
commit 7e0ee4a719
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 190 additions and 7 deletions
Download patch file Download diff file Expand all files Collapse all files

38
studio/frontend/src/components/assistant-ui/thread.tsx
View file

@ -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 (
<button
type="button"
disabled={disabled}
onClick={() => setImageToolsEnabled(!imageToolsEnabled)}
className="composer-pill-btn"
data-active={imageToolsEnabled && !disabled ? "true" : "false"}
aria-label={
imageToolsEnabled ? "Disable image generation" : "Enable image generation"
}
>
<ImageIcon className="size-3.5" />
<span>Images</span>
</button>
);
};
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 }> = ({
<PreserveThinkingToggle />
<WebSearchToggle />
<CodeToolsToggle />
<ImagesToggle />
</div>
<div className="flex items-center gap-1">
<ComposerPrimitive.If dictation={false}>

27
studio/frontend/src/features/chat/api/chat-adapter.ts
View file

@ -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"]
: []),
],
}
: {}),

40
studio/frontend/src/features/chat/chat-page.tsx
View file

@ -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;

38
studio/frontend/src/features/chat/provider-capabilities.ts
View file

@ -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

34
studio/frontend/src/features/chat/shared-composer.tsx
View file

@ -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({
<CodeToggleIcon className="size-3.5" />
<span>Code</span>
</button>
{showImagePill && (
<button
type="button"
disabled={imageDisabled}
onClick={() => setImageToolsEnabled(!imageToolsEnabled)}
className="composer-pill-btn"
data-active={imageToolsEnabled && !imageDisabled ? "true" : "false"}
aria-label={
imageToolsEnabled ? "Disable image generation" : "Enable image generation"
}
>
<ImageIcon className="size-3.5" />
<span>Images</span>
</button>
)}
</div>
<div className="flex items-center gap-1">
{dictationSupported && (

20
studio/frontend/src/features/chat/stores/chat-runtime-store.ts
View file

@ -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::<providerId>::<modelId>`. 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<ChatRuntimeStore>((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<ChatRuntimeStore>((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<ChatRuntimeStore>((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) =>