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

Studio: wire OpenAI image_generation tool (#5688)

Browse source 
* Studio: wire OpenAI image_generation tool

OpenAI's Responses API exposes server-side image generation as a
tool entry (`{type: "image_generation"}`); the result comes back as
an `image_generation_call` output item with the base64 image on
`result`, the actual prompt used on `revised_prompt`, plus `size`,
`quality`, `output_format`, `background`. The model decides when to
call the tool based on the user's request; rendering uses one of
the gpt-image-* backbones server-side.

Available on every gpt-5.x family member plus gpt-4.1, gpt-4o, o3,
o4-mini per the docs.

Changes:

- Append `{type:"image_generation"}` to the Responses request tools
  array when `enabled_tools` carries `image_generation` AND the base
  URL points at cloud OpenAI. Non-cloud bases (ollama, llama.cpp,
  "custom" presets that collapse to provider="openai") silently drop
  the tool to avoid 400s.
- Mirror the same logic in `_build_body` (the post-expiry retry
  builder) so retries carry the same tool set as the original
  attempt.
- Handle `image_generation_call` items in
  `response.output_item.done`: emit `tool_start` with
  `arguments:{kind:"image", prompt:<revised_prompt>}` and `tool_end`
  with `image_b64`, `image_mime`, `size`, `quality`, `background`
  so the chat adapter can render an inline preview. Image bytes go
  on the tool_end chunk; no extra fields on the chat-completions
  envelope so the OpenAI SDK shape stays clean.
- Add `import time` (used for synthesised tool_call_id fallback).
- Add `test_openai_image_generation.py` with 5 cases: tool entry on
  cloud OpenAI, combined with web_search + code_execution
  (verifies all three coexist), non-cloud drop, omitted pill leaves
  body untouched, output item translation produces the expected
  tool_start + tool_end chunks.

Live verified end-to-end: `gpt-5.4-mini` with `image_generation`
tool returned an `image_generation_call` carrying ~1MB of base64
PNG plus the gpt-image backbone's revised prompt.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Use time.time_ns() for synthesised image_generation tool_call_id

Gemini medium on PR #5688: `int(time.time() * 1000)` has 1ms
resolution; two image generations resolving in the same millisecond
would collide on the synthesised id. Bump to nanoseconds.

(In practice the upstream `image_generation_call` item always carries
its own `id`; the synthesised fallback only fires when OpenAI omits
it -- rare, but cheap to harden.)

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Daniel Han 2026-05-22 06:03:38 -07:00 • committed by GitHub
commit 5b41872e8b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 289 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

71
studio/backend/core/inference/external_provider.py
View file

@ -10,6 +10,7 @@ Anthropic uses native Messages API with translation in this client.
import json as _json
import re
import time
from typing import Any, AsyncGenerator, Literal, NamedTuple, Optional
import httpx
@ -2235,6 +2236,18 @@ class ExternalProviderClient:
code_execution_enabled_openai = bool(
enabled_tools and "code_execution" in enabled_tools and is_openai_cloud
)
# OpenAI's image_generation tool is a Responses-API server tool.
# See https://developers.openai.com/api/docs/guides/tools-image-generation
# The model picks size / quality / background server-side and
# delegates rendering to a gpt-image-* family model; the result
# comes back inline as an `image_generation_call` output item
# with a base64 image. Available on every gpt-5.x family member
# plus gpt-4.1 / gpt-4o / o3 per the docs; restrict to cloud
# OpenAI because the local llama.cpp / ollama backends don't
# implement it and would 400.
image_generation_enabled_openai = bool(
enabled_tools and "image_generation" in enabled_tools and is_openai_cloud
)
if enabled_tools:
tools_array: list[dict[str, Any]] = []
if "web_search" in enabled_tools:
@ -2261,6 +2274,8 @@ class ExternalProviderClient:
else:
shell_env = {"type": "container_auto"}
tools_array.append({"type": "shell", "environment": shell_env})
if image_generation_enabled_openai:
tools_array.append({"type": "image_generation"})
if tools_array:
body["tools"] = tools_array
@ -2291,6 +2306,8 @@ class ExternalProviderClient:
tools_array_attempt.append(
{"type": "shell", "environment": env_attempt}
)
if image_generation_enabled_openai:
tools_array_attempt.append({"type": "image_generation"})
if tools_array_attempt:
attempt_body["tools"] = tools_array_attempt
else:
@ -2730,6 +2747,60 @@ class ExternalProviderClient:
"result": result_text,
}
)
elif item.get("type") == "image_generation_call":
# OpenAI's image_generation tool returns
# a single output item with the base64
# PNG/WebP/JPEG on `result` (sometimes
# `b64_json` depending on output_format).
# `revised_prompt` is what the gpt-image
# backbone actually used after refinement
# of the assistant's request. Emit
# tool_start + tool_end so the chat card
# renders the prompt + the generated
# image inline. The frontend chat-adapter
# decides how to render the base64 blob
# (likely an <img src="data:image/...">)
# based on the `kind: "image"` hint we
# set on tool_start arguments.
# `time_ns()` (nanoseconds) instead of
# millisecond resolution so synthesised
# ids stay unique even when two image
# generations resolve in the same ms.
item_id = item.get("id", "") or (
f"img_{time.time_ns()}"
)
prompt_in = (
item.get("revised_prompt")
or item.get("prompt")
or ""
)
yield _emit_tool_event(
{
"type": "tool_start",
"tool_name": "image_generation",
"tool_call_id": item_id,
"arguments": {
"kind": "image",
"prompt": prompt_in,
},
}
)
b64 = (
item.get("result") or item.get("b64_json") or ""
)
output_format = item.get("output_format") or "png"
yield _emit_tool_event(
{
"type": "tool_end",
"tool_call_id": item_id,
"result": "",
"image_b64": b64,
"image_mime": (f"image/{output_format}"),
"size": item.get("size"),
"quality": item.get("quality"),
"background": item.get("background"),
}
)
elif (
isinstance(event_type, str)

218
studio/backend/tests/test_openai_image_generation.py Normal file
View file

@ -0,0 +1,218 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Unit tests for OpenAI Responses API image_generation tool wiring.
The image_generation tool is a server-side Responses-API tool:
``{type: "image_generation"}`` in the request's tools array, and the
result comes back as an ``image_generation_call`` output item carrying
the base64 image on ``result``. Studio translates the output item
into ``_toolEvent`` chunks (``tool_start`` with `kind:"image"`,
``tool_end`` with ``image_b64`` + ``image_mime``) so the chat adapter
can render the image inline.
These tests pin: the tool is added to the outbound body only when the
caller asks for it on a cloud OpenAI base; the SSE output_item.done
for ``image_generation_call`` produces the expected _toolEvent chunks;
non-cloud bases drop the tool silently.
"""
import asyncio
import json
import httpx
from core.inference import external_provider as ep_mod
from core.inference.external_provider import ExternalProviderClient
def _drive(coro):
return asyncio.new_event_loop().run_until_complete(coro)
def _capture_body(monkeypatch, *, base_url: str, enabled_tools) -> dict:
captured: dict = {}
def handler(request: httpx.Request) -> httpx.Response:
captured["body"] = json.loads(request.content.decode("utf-8"))
return httpx.Response(
200,
content = (
b"event: response.completed\n"
b'data: {"type":"response.completed",'
b'"response":{"output":[],"usage":{"input_tokens":0,'
b'"output_tokens":0}}}\n\n'
),
headers = {"content-type": "text/event-stream"},
)
monkeypatch.setattr(
ep_mod,
"_http_client",
httpx.AsyncClient(transport = httpx.MockTransport(handler)),
)
async def run():
client = ExternalProviderClient(
provider_type = "openai",
base_url = base_url,
api_key = "sk-test",
)
async for _ in client.stream_chat_completion(
messages = [{"role": "user", "content": "draw a cat"}],
model = "gpt-5.5",
temperature = 0.7,
top_p = 0.95,
max_tokens = 32,
reasoning_effort = "medium",
enabled_tools = enabled_tools,
):
pass
await client.close()
_drive(run())
return captured
def _collect_tool_events(monkeypatch) -> list[dict]:
"""Drive a Responses stream that emits one image_generation_call done
event and return the parsed _toolEvent chunks."""
sse = (
b"event: response.output_item.done\n"
b'data: {"type":"response.output_item.done",'
b'"item":{"type":"image_generation_call",'
b'"id":"img_abc",'
b'"revised_prompt":"A photorealistic cat sitting",'
b'"result":"AAAA",'
b'"output_format":"png",'
b'"size":"1024x1024",'
b'"quality":"high",'
b'"background":"opaque"}}\n\n'
b"event: response.completed\n"
b'data: {"type":"response.completed",'
b'"response":{"output":[],"usage":{"input_tokens":0,'
b'"output_tokens":0}}}\n\n'
)
def handler(request: httpx.Request) -> httpx.Response:
return httpx.Response(
200,
content = sse,
headers = {"content-type": "text/event-stream"},
)
monkeypatch.setattr(
ep_mod,
"_http_client",
httpx.AsyncClient(transport = httpx.MockTransport(handler)),
)
events: list[dict] = []
async def run():
client = ExternalProviderClient(
provider_type = "openai",
base_url = "https://api.openai.com/v1",
api_key = "sk-test",
)
async for line in client.stream_chat_completion(
messages = [{"role": "user", "content": "draw a cat"}],
model = "gpt-5.5",
temperature = 0.7,
top_p = 0.95,
max_tokens = 32,
reasoning_effort = "medium",
enabled_tools = ["image_generation"],
):
if not line or not line.startswith("data:"):
continue
payload = line[5:].strip()
if payload == "[DONE]":
continue
try:
obj = json.loads(payload)
except json.JSONDecodeError:
continue
if "_toolEvent" in obj:
events.append(obj["_toolEvent"])
await client.close()
_drive(run())
return events
# ── tool entry appended to outbound body on cloud OpenAI ─────────────
def test_cloud_openai_appends_image_generation_tool(monkeypatch):
captured = _capture_body(
monkeypatch,
base_url = "https://api.openai.com/v1",
enabled_tools = ["image_generation"],
)
tools = captured["body"].get("tools") or []
assert {"type": "image_generation"} in tools, tools
def test_combined_with_web_search_and_code_execution(monkeypatch):
captured = _capture_body(
monkeypatch,
base_url = "https://api.openai.com/v1",
enabled_tools = ["web_search", "code_execution", "image_generation"],
)
tools = captured["body"].get("tools") or []
tool_types = {t["type"] for t in tools if isinstance(t, dict)}
assert tool_types == {"web_search", "shell", "image_generation"}, tools
# ── non-cloud base silently drops the tool ──────────────────────────
def test_non_cloud_base_drops_image_generation(monkeypatch):
captured = _capture_body(
monkeypatch,
base_url = "http://127.0.0.1:11434/v1",
enabled_tools = ["image_generation"],
)
tools = captured["body"].get("tools") or []
assert {"type": "image_generation"} not in tools, tools
# ── omitted pill leaves body untouched ──────────────────────────────
def test_omitted_image_generation_pill_no_tool(monkeypatch):
captured = _capture_body(
monkeypatch,
base_url = "https://api.openai.com/v1",
enabled_tools = ["web_search"],
)
tools = captured["body"].get("tools") or []
assert all(t.get("type") != "image_generation" for t in tools)
# ── output translation surfaces tool_start + tool_end ────────────────
def test_image_generation_done_emits_tool_event_chunks(monkeypatch):
events = _collect_tool_events(monkeypatch)
image_events = [
e
for e in events
if e.get("tool_name") == "image_generation"
or (e.get("type") == "tool_end" and e.get("image_b64"))
]
starts = [e for e in image_events if e.get("type") == "tool_start"]
ends = [e for e in image_events if e.get("type") == "tool_end"]
assert len(starts) == 1, image_events
assert len(ends) == 1, image_events
assert starts[0]["arguments"] == {
"kind": "image",
"prompt": "A photorealistic cat sitting",
}
assert ends[0]["image_b64"] == "AAAA"
assert ends[0]["image_mime"] == "image/png"
assert ends[0]["size"] == "1024x1024"
assert ends[0]["quality"] == "high"
assert ends[0]["background"] == "opaque"