diff --git a/studio/backend/core/inference/external_provider.py b/studio/backend/core/inference/external_provider.py index 44e8720fed..b9e224bd66 100644 --- a/studio/backend/core/inference/external_provider.py +++ b/studio/backend/core/inference/external_provider.py @@ -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 ) + # 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) diff --git a/studio/backend/tests/test_openai_image_generation.py b/studio/backend/tests/test_openai_image_generation.py new file mode 100644 index 0000000000..1f9c2710e4 --- /dev/null +++ b/studio/backend/tests/test_openai_image_generation.py @@ -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"