diff --git a/DOCKER.md b/DOCKER.md index caf2b22..f94ad4d 100644 --- a/DOCKER.md +++ b/DOCKER.md @@ -34,6 +34,7 @@ Set options in `.env` or pass environment variables: - `CHATGPT_LOCAL_ENABLE_WEB_SEARCH`: `true|false` to enable default web search tool - `CHATGPT_LOCAL_MODEL_SYNC`: `true|false` to discover account models automatically (default `true`) - `CHATGPT_LOCAL_MODEL_REFRESH_INTERVAL`: model catalog refresh interval in seconds (default `3600`) +- `CHATGPT_LOCAL_IMAGE_MODEL`: model that orchestrates image requests (default `gpt-5.4-mini`) ## Logs Set `VERBOSE=true` to include extra logging for troubleshooting upstream or chat app requests. Please include and use these logs when submitting bug reports. diff --git a/FORK.md b/FORK.md new file mode 100644 index 0000000..e0a7c24 --- /dev/null +++ b/FORK.md @@ -0,0 +1,124 @@ +# Fork: geração de imagem no ChatMock + +Fork de [RayBytes/ChatMock](https://github.com/RayBytes/ChatMock) (v1.40) com um +objetivo só: **fazer o ChatMock devolver imagens**. + +O backend do Codex já sabia desenhar — quem não sabia receber era o ChatMock. + +## O que foi descoberto + +O endpoint `https://chatgpt.com/backend-api/codex/responses`, o mesmo que o +ChatMock já usava para texto, aceita o tool `image_generation` e devolve a imagem +em base64. Medido, não deduzido: + +- modelo que desenha: `gpt-image-2-codex` (não dá para trocar); +- `tool_choice: {"type": "image_generation"}` funciona e força a geração; +- o consumo aparece em `tool_usage.image_gen` da resposta — ou seja, **sai da + cota do plano ChatGPT/Codex**, não de crédito de API; +- `n` no tool é recusado (`Unknown parameter: 'tools[0].n'`): uma imagem por + requisição; +- `size` e `quality` são aceitos mas ignorados — o echo volta sempre `"auto"`. + A resolução é escolhida a partir do prompt: pedindo 1024x1024 veio 1254x1254, + pedindo 1536x1024 veio 1536x1024 na mosca. Ou seja, dá para pedir — só não dá + para garantir. + +## O bug que estava por baixo + +`response.completed` do backend do Codex vem com `output: []`. Sempre. Os itens +de saída só existem nos eventos `response.output_item.done`. + +Como `aggregate_response_from_sse()` devolvia o objeto do evento final, **todo +`POST /v1/responses` com `stream: false` respondia com output vazio** — inclusive +para texto puro, sem imagem nenhuma. A resposta parecia bem-sucedida (`status: +completed`, `usage` preenchido) e não tinha conteúdo. + +Corrigido em `chatmock/responses_api.py`: os itens são acumulados por +`output_index` e remontados no objeto final quando ele vem vazio. + +## O que mudou + +| Arquivo | Mudança | +|---|---| +| `chatmock/images_api.py` | **novo** — monta o payload e converte para o formato da Images API. Sem import interno de propósito: `utils.py` importa este módulo, e `model_catalog -> utils` fecharia o ciclo | +| `chatmock/responses_api.py` | remonta `output` a partir dos `output_item.done` (o bug acima); `collect_images_from_sse()`, que fica aqui porque é este o módulo que já lê SSE | +| `chatmock/routes_openai.py` | rota `/v1/images/generations`; `image_generation` liberado em `responses_tools`; imagem vira data-url no `/v1/chat/completions` | +| `chatmock/utils.py` | imagem como delta de conteúdo no chat streaming, **depois** de fechar o `` — o fechamento virou o helper `_close_think_tag()`, usado nos três pontos que antes repetiam o mesmo bloco | +| `chatmock/app.py`, `chatmock/cli.py` | flag `--image-model` / env `CHATGPT_LOCAL_IMAGE_MODEL` | +| `tests/test_routes.py` | 8 testes novos | + +## Uso + +```bash +python chatmock.py serve --port 8000 +``` + +### `POST /v1/images/generations` + +Compatível com a Images API da OpenAI, então SDK oficial, n8n e afins falam com +ela sem adaptação: + +```bash +curl http://127.0.0.1:8000/v1/images/generations \ + -H "Content-Type: application/json" \ + -d '{"prompt":"logo minimalista de uma raposa geometrica","size":"1024x1024"}' +``` + +```json +{ "created": 1786806013, + "data": [{ "b64_json": "...", "revised_prompt": "...", "size": "1254x1254" }], + "usage": { "input_tokens": 52, "output_tokens": 915, "total_tokens": 967 } } +``` + +Parâmetros: + +- `prompt` — obrigatório. +- `n` — de 1 a 4. Cada unidade é **uma requisição a mais** na sua cota, porque o + backend não aceita `n` no tool. +- `size` — vira instrução em texto ("quadrada", "horizontal", "vertical" + + os pixels pedidos), já que o backend ignora o parâmetro. Trate como pedido, + não como garantia: a proporção costuma sair certa, a resolução exata não. +- `quality`, `output_format`, `output_compression`, `background`, `moderation` — + repassados ao tool. Se o backend recusar algum com `unknown_parameter`, a + requisição é refeita uma vez com o tool pelado, para o cliente receber a imagem + em vez de um 400. +- `image` — data-url (ou lista delas) para usar como referência. Não é a + `/v1/images/edits` multipart da OpenAI, é um atalho em JSON. +- `chat_model` — troca o modelo que orquestra só nesta chamada. +- `response_format: "url"` responde 400: a imagem vem em base64 e o ChatMock não + hospeda arquivo — igual ao comportamento do `gpt-image-1` na API real. + +### Pelo `/v1/chat/completions` + +```json +{ "model": "gpt-5.4-mini", + "messages": [{"role": "user", "content": "gere uma imagem de um cubo azul"}], + "responses_tools": [{"type": "image_generation"}] } +``` + +A imagem chega embutida no `content` como `![alt](data:image/png;base64,...)`, +em streaming ou não. Serve para UIs de chat que renderizam markdown. + +Com `--reasoning-compat think-tags` (o padrão), a imagem sai **depois** do +``. Sem isso ela cairia dentro do bloco de raciocínio e sumiria em +qualquer cliente que esconde o `` — que é justamente o motivo do modo +existir. + +## Limites conhecidos + +- **Peso.** Uma imagem passa de 2,5 MB em base64. Não ligue `--verbose` nessas + rotas: o log imprime o corpo inteiro. +- Rotas Ollama (`/api/chat`) não foram tocadas — continuam só texto. +- O modelo de imagem é escolhido pelo backend. `model` no corpo da requisição não + troca nada; existe para não quebrar cliente que sempre manda `gpt-image-1`. + +## Rebase + +As mudanças em arquivos que já existiam são pequenas e localizadas — o grosso +está em `images_api.py`, que é arquivo novo. Ao subir de versão, os pontos de +atrito são `aggregate_response_from_sse()`, os dois trechos de +`response.output_item.done` (em `routes_openai.py` e `utils.py`) e o +`_close_think_tag()`, que substituiu duas cópias de um bloco que o upstream +repetia dentro de `sse_translate_chat()`. + +Comentários e mensagens de erro do código estão em inglês, como o resto do +repositório — só este arquivo está em português. diff --git a/README.md b/README.md index 94b78e6..5373138 100644 --- a/README.md +++ b/README.md @@ -104,6 +104,7 @@ account. The current catalog commonly includes: - Tool / function calling - Vision / image input +- Image generation (`/v1/images/generations`, or as a tool in a chat request) - Thinking summaries (via think tags) - Configurable thinking effort - Fast mode for supported models @@ -128,6 +129,7 @@ All flags go after `chatmock serve`. These can also be set as environment variab | `--expose-reasoning-models` | `CHATGPT_LOCAL_EXPOSE_REASONING_MODELS` | true/false | false | List each reasoning level as its own model | | `--model-sync` | `CHATGPT_LOCAL_MODEL_SYNC` | true/false | true | Discover account models automatically | | `--model-refresh-interval` | `CHATGPT_LOCAL_MODEL_REFRESH_INTERVAL` | seconds | 3600 | Refresh interval for model discovery | +| `--image-model` | `CHATGPT_LOCAL_IMAGE_MODEL` | model slug | gpt-5.4-mini | Model that orchestrates image requests |
Web search in a request @@ -143,6 +145,53 @@ All flags go after `chatmock serve`. These can also be set as environment variab
+
+Generating an image + +`/v1/images/generations` mirrors the OpenAI Images API, so existing clients work +unchanged: + +```bash +curl http://127.0.0.1:8000/v1/images/generations \ + -H "Content-Type: application/json" \ + -d '{"prompt": "a minimalist geometric fox logo", "size": "1024x1024"}' +``` + +```json +{ "created": 1786806013, + "data": [{ "b64_json": "...", "revised_prompt": "...", "size": "1254x1254" }], + "usage": { "input_tokens": 52, "output_tokens": 915, "total_tokens": 967 } } +``` + +The same tool works inside a chat request on `/v1/chat/completions` and +`/api/chat`, where the image comes back embedded in the message content as +`![alt](data:image/png;base64,...)`: + +```json +{ + "model": "gpt-5.4-mini", + "messages": [{"role": "user", "content": "draw a blue cube"}], + "responses_tools": [{"type": "image_generation"}] +} +``` + +Worth knowing: + +- The picture is always drawn by the backend's own image model. `model` in the + request body changes nothing; it is accepted so clients that always send + `gpt-image-1` keep working. +- `n` goes up to 4, and each unit is a separate upstream request, because the + backend refuses `n` inside the tool. +- `size` is passed along but the backend decides the final resolution from the + prompt, so it is written into the instructions as well. Treat it as a request, + not a guarantee. +- `response_format: "url"` is rejected: the backend returns base64 and ChatMock + hosts no files. +- One image is a couple of megabytes of base64. `--verbose` prints request bodies, + so leave it off when passing reference images. + +
+
Fast mode in a request diff --git a/chatmock/app.py b/chatmock/app.py index f50eae0..94c9e7a 100644 --- a/chatmock/app.py +++ b/chatmock/app.py @@ -6,6 +6,7 @@ from flask_sock import Sock from .http import build_cors_headers +from .images_api import DEFAULT_IMAGE_ORCHESTRATOR_MODEL from .model_catalog import DEFAULT_REFRESH_INTERVAL_SECONDS, ModelCatalog from .routes_openai import openai_bp from .routes_ollama import ollama_bp @@ -24,8 +25,13 @@ def create_app( default_web_search: bool = False, model_sync: bool | None = None, model_refresh_interval: float | None = None, + image_orchestrator_model: str | None = None, ) -> Flask: app = Flask(__name__) + if not (isinstance(image_orchestrator_model, str) and image_orchestrator_model.strip()): + image_orchestrator_model = ( + os.getenv("CHATGPT_LOCAL_IMAGE_MODEL") or DEFAULT_IMAGE_ORCHESTRATOR_MODEL + ) if model_sync is None: model_sync = (os.getenv("CHATGPT_LOCAL_MODEL_SYNC") or "true").strip().lower() in ( "1", @@ -53,6 +59,7 @@ def create_app( DEFAULT_WEB_SEARCH=bool(default_web_search), MODEL_SYNC=bool(model_sync), MODEL_REFRESH_INTERVAL=float(model_refresh_interval), + IMAGE_ORCHESTRATOR_MODEL=image_orchestrator_model.strip(), ) app.extensions["chatmock_model_catalog"] = ModelCatalog( enabled=bool(model_sync), diff --git a/chatmock/cli.py b/chatmock/cli.py index 1383bf9..04aeb45 100644 --- a/chatmock/cli.py +++ b/chatmock/cli.py @@ -10,6 +10,7 @@ from .app import create_app from .config import CLIENT_ID_DEFAULT +from .images_api import DEFAULT_IMAGE_ORCHESTRATOR_MODEL from .limits import RateLimitWindow, compute_reset_at, load_rate_limit_snapshot from .oauth import OAuthHTTPServer, OAuthHandler, REQUIRED_PORT, URL_BASE, run_device_code_login from .utils import eprint, get_home_dir, load_chatgpt_tokens, parse_jwt_claims, read_auth_file @@ -242,6 +243,7 @@ def cmd_serve( default_web_search: bool, model_sync: bool = True, model_refresh_interval: float = 3600, + image_orchestrator_model: str | None = None, ) -> int: app = create_app( verbose=verbose, @@ -255,6 +257,7 @@ def cmd_serve( default_web_search=default_web_search, model_sync=model_sync, model_refresh_interval=model_refresh_interval, + image_orchestrator_model=image_orchestrator_model, ) app.run(host=host, use_reloader=False, port=port, threaded=True) @@ -345,6 +348,16 @@ def main() -> None: help="Refresh the ChatGPT model catalog after this many seconds (default: 3600).", ) + p_serve.add_argument( + "--image-model", + default=os.getenv("CHATGPT_LOCAL_IMAGE_MODEL", DEFAULT_IMAGE_ORCHESTRATOR_MODEL), + metavar="MODEL", + help=( + "Model that orchestrates /v1/images/generations. The picture itself is always drawn " + f"by the backend's image model, so the cheapest one does (default: {DEFAULT_IMAGE_ORCHESTRATOR_MODEL})." + ), + ) + p_info = sub.add_parser("info", help="Print current stored tokens and derived account id") p_info.add_argument("--json", action="store_true", help="Output raw auth.json contents") @@ -368,6 +381,7 @@ def main() -> None: default_web_search=args.enable_web_search, model_sync=args.model_sync, model_refresh_interval=args.model_refresh_interval, + image_orchestrator_model=args.image_model, ) ) elif args.command == "info": diff --git a/chatmock/images_api.py b/chatmock/images_api.py new file mode 100644 index 0000000..9bda999 --- /dev/null +++ b/chatmock/images_api.py @@ -0,0 +1,148 @@ +from __future__ import annotations + +from typing import Any, Dict, List + +# Deliberately free of intra-package imports: utils.py imports this module, and +# model_catalog -> utils would close the cycle. SSE parsing lives in +# responses_api.py, which already owns it. + + +IMAGE_TOOL_TYPE = "image_generation" + +# Model that orchestrates the call. The picture itself is always drawn by the +# backend's own image model, so the cheapest chat model does the job. +DEFAULT_IMAGE_ORCHESTRATOR_MODEL = "gpt-5.4-mini" + +# Tool parameters the backend accepts today. 'n' shows up in the echoed tool +# config but is refused on input ("Unknown parameter: 'tools[0].n'"), so several +# images means several requests. +IMAGE_TOOL_PARAM_KEYS = ( + "size", + "quality", + "output_format", + "output_compression", + "background", + "moderation", +) + +MAX_IMAGES_PER_REQUEST = 4 + + +def build_image_tool(params: Dict[str, Any] | None) -> Dict[str, Any]: + tool: Dict[str, Any] = {"type": IMAGE_TOOL_TYPE} + if isinstance(params, dict): + for key in IMAGE_TOOL_PARAM_KEYS: + value = params.get(key) + if value is None: + continue + if isinstance(value, str) and not value.strip(): + continue + tool[key] = value + return tool + + +def size_hint_instruction(size: Any) -> str | None: + """Turn a requested size into plain text. + + The backend ignores the `size` parameter (the echo always reads "auto") and + picks the aspect ratio from the prompt, so the request only lands if it is + written out for the model to read. + """ + if not isinstance(size, str): + return None + raw = size.strip().lower() + if not raw or raw == "auto" or "x" not in raw: + return None + left, _, right = raw.partition("x") + try: + width = int(left) + height = int(right) + except ValueError: + return None + if width <= 0 or height <= 0: + return None + if width == height: + shape = "square (1:1)" + elif width > height: + shape = "landscape" + else: + shape = "portrait" + return f"The image must be {shape}, as close as possible to {width}x{height} pixels." + + +def build_image_request_payload( + prompt: str, + *, + model: str, + tool_params: Dict[str, Any] | None = None, + input_images: List[str] | None = None, +) -> Dict[str, Any]: + content: List[Dict[str, Any]] = [{"type": "input_text", "text": prompt}] + for url in input_images or []: + if isinstance(url, str) and url.strip(): + content.append({"type": "input_image", "image_url": url}) + + instructions = [ + "You generate images. Call the image tool exactly once with the user's " + "request and write nothing else." + ] + hint = size_hint_instruction((tool_params or {}).get("size")) + if hint: + instructions.append(hint) + + return { + "model": model, + "instructions": " ".join(instructions), + "input": [{"type": "message", "role": "user", "content": content}], + "tools": [build_image_tool(tool_params)], + "tool_choice": {"type": IMAGE_TOOL_TYPE}, + "parallel_tool_calls": False, + "store": False, + "stream": True, + } + + +def image_item_to_openai(item: Dict[str, Any]) -> Dict[str, Any]: + out: Dict[str, Any] = {"b64_json": item.get("result") or ""} + revised = item.get("revised_prompt") + if isinstance(revised, str) and revised.strip(): + out["revised_prompt"] = revised + # Not part of the OpenAI schema, but the backend decides these on its own and + # the caller has no other way to learn what it actually got. + for key in ("size", "output_format", "quality", "background"): + value = item.get(key) + if value is not None: + out[key] = value + return out + + +def data_url_for_item(item: Dict[str, Any]) -> str | None: + b64 = item.get("result") + if not isinstance(b64, str) or not b64: + return None + fmt = item.get("output_format") if isinstance(item.get("output_format"), str) else "png" + return f"data:image/{fmt};base64,{b64}" + + +def image_markdown_for_item(item: Dict[str, Any]) -> str: + url = data_url_for_item(item) + if not url: + return "" + alt = item.get("revised_prompt") + alt = alt.replace("\n", " ").strip()[:120] if isinstance(alt, str) else "" + return f"\n\n![{alt or 'image'}]({url})" + + +def usage_to_openai(usage: Dict[str, Any] | None) -> Dict[str, Any] | None: + if not isinstance(usage, dict): + return None + try: + input_tokens = int(usage.get("input_tokens") or 0) + output_tokens = int(usage.get("output_tokens") or 0) + except (TypeError, ValueError): + return None + return { + "input_tokens": input_tokens, + "output_tokens": output_tokens, + "total_tokens": int(usage.get("total_tokens") or (input_tokens + output_tokens)), + } diff --git a/chatmock/responses_api.py b/chatmock/responses_api.py index ab66803..6b245bf 100644 --- a/chatmock/responses_api.py +++ b/chatmock/responses_api.py @@ -171,6 +171,12 @@ def aggregate_response_from_sse( ) -> tuple[Dict[str, Any] | None, Dict[str, Any] | None]: response_obj: Dict[str, Any] | None = None error_obj: Dict[str, Any] | None = None + # The Codex backend sends `response.completed` with `output: []` — the items + # (message, function_call, image_generation_call) only ever arrive in the + # `response.output_item.done` events. Without rebuilding them here, every + # non-streaming response comes out empty. + done_items: Dict[int, Dict[str, Any]] = {} + fallback_order = 0 try: for evt in iter_sse_event_payloads(upstream): if callable(on_event): @@ -182,6 +188,14 @@ def aggregate_response_from_sse( if isinstance(response, dict): response_obj = response kind = evt.get("type") + if kind == "response.output_item.done": + item = evt.get("item") + if isinstance(item, dict): + index = evt.get("output_index") + if not isinstance(index, int): + index = 1_000_000 + fallback_order + fallback_order += 1 + done_items[index] = item if kind == "response.failed": if isinstance(response, dict) and isinstance(response.get("error"), dict): error_obj = {"error": response.get("error")} @@ -192,9 +206,64 @@ def aggregate_response_from_sse( break finally: upstream.close() + if isinstance(response_obj, dict) and done_items: + existing = response_obj.get("output") + if not isinstance(existing, list) or not existing: + response_obj = dict(response_obj) + response_obj["output"] = [done_items[key] for key in sorted(done_items)] return response_obj, error_obj +def collect_images_from_sse( + upstream: Any, +) -> tuple[List[Dict[str, Any]], Dict[str, Any] | None, Dict[str, Any] | None, str]: + """Read the stream to the end and return (image items, usage, error, text). + + Same reason as above: `response.completed` carries no output, so the image + items only ever exist in the `response.output_item.done` events. + + The text is collected for one reason: when the model finishes without + drawing, that text is the only place the reason exists. It used to be + dropped, and the caller was left guessing — the error even said the refusal + "usually" came from moderation, which nobody had ever confirmed. + """ + images: List[Dict[str, Any]] = [] + usage: Dict[str, Any] | None = None + error: Dict[str, Any] | None = None + said: List[str] = [] + try: + for evt in iter_sse_event_payloads(upstream): + kind = evt.get("type") + if kind == "response.output_item.done": + item = evt.get("item") + if isinstance(item, dict) and item.get("type") == "image_generation_call": + images.append(item) + elif isinstance(item, dict) and item.get("type") == "message": + for part in item.get("content") or []: + if isinstance(part, dict) and isinstance(part.get("text"), str): + said.append(part["text"]) + elif kind == "response.failed": + response = evt.get("response") + if isinstance(response, dict) and isinstance(response.get("error"), dict): + error = response["error"] + else: + error = {"message": "response.failed"} + break + elif kind == "error": + error = evt.get("error") if isinstance(evt.get("error"), dict) else {"message": "upstream error"} + break + elif kind == "response.completed": + response = evt.get("response") + if isinstance(response, dict): + tool_usage = response.get("tool_usage") + if isinstance(tool_usage, dict) and isinstance(tool_usage.get("image_gen"), dict): + usage = tool_usage["image_gen"] + break + finally: + upstream.close() + return images, usage, error, " ".join(said).strip() + + def stream_upstream_bytes( upstream: Any, *, diff --git a/chatmock/routes_ollama.py b/chatmock/routes_ollama.py index 70981ee..a7930bc 100644 --- a/chatmock/routes_ollama.py +++ b/chatmock/routes_ollama.py @@ -8,6 +8,7 @@ from flask import Blueprint, Response, current_app, jsonify, make_response, request, stream_with_context from .fast_mode import resolve_service_tier +from .images_api import IMAGE_TOOL_TYPE, image_markdown_for_item from .limits import record_rate_limits_from_response from .http import build_cors_headers from .model_registry import list_public_models @@ -206,8 +207,13 @@ def ollama_chat() -> Response: for _t in rt_payload: if not (isinstance(_t, dict) and isinstance(_t.get("type"), str)): continue - if _t.get("type") not in ("web_search", "web_search_preview"): - err = {"error": "Only web_search/web_search_preview are supported in responses_tools"} + if _t.get("type") not in ("web_search", "web_search_preview", IMAGE_TOOL_TYPE): + err = { + "error": ( + "Only web_search/web_search_preview/image_generation are supported " + "in responses_tools" + ) + } if verbose: _log_json("OUT POST /api/chat", err) return jsonify(err), 400 @@ -437,6 +443,41 @@ def _gen(): full_parts.append(delta_txt) else: pass + elif kind == "response.output_item.done" and ( + (evt.get("item") or {}).get("type") == "image_generation_call" + ): + markdown = image_markdown_for_item(evt.get("item") or {}) + if markdown: + # Same reason as the text branch below: close the + # reasoning block first, or the image is hidden along + # with it. + if compat == "think-tags" and think_open and not think_closed: + yield ( + json.dumps( + { + "model": model_out, + "created_at": created_at, + "message": {"role": "assistant", "content": ""}, + "done": False, + } + ) + + "\n" + ) + full_parts.append("") + think_open = False + think_closed = True + yield ( + json.dumps( + { + "model": model_out, + "created_at": created_at, + "message": {"role": "assistant", "content": markdown}, + "done": False, + } + ) + + "\n" + ) + full_parts.append(markdown) elif kind == "response.output_text.delta": delta = evt.get("delta") or "" if compat == "think-tags" and think_open and not think_closed: @@ -534,7 +575,11 @@ def _gen(): reasoning_full_text += evt.get("delta") or "" elif kind == "response.output_item.done": item = evt.get("item") or {} - if isinstance(item, dict) and item.get("type") == "function_call": + if isinstance(item, dict) and item.get("type") == "image_generation_call": + # Appended, not prepended: the think block is added in front + # of full_text further down. + full_text += image_markdown_for_item(item) + elif isinstance(item, dict) and item.get("type") == "function_call": call_id = item.get("call_id") or item.get("id") or "" name = item.get("name") or "" args = item.get("arguments") or "" diff --git a/chatmock/routes_openai.py b/chatmock/routes_openai.py index 673e22f..4dfde08 100644 --- a/chatmock/routes_openai.py +++ b/chatmock/routes_openai.py @@ -7,12 +7,22 @@ from flask import Blueprint, Response, current_app, jsonify, make_response, request from .fast_mode import resolve_service_tier +from .images_api import ( + DEFAULT_IMAGE_ORCHESTRATOR_MODEL, + IMAGE_TOOL_TYPE, + MAX_IMAGES_PER_REQUEST, + build_image_request_payload, + image_item_to_openai, + image_markdown_for_item, + usage_to_openai, +) from .limits import record_rate_limits_from_response from .http import build_cors_headers from .model_registry import list_public_models from .responses_api import ( ResponsesRequestError, aggregate_response_from_sse, + collect_images_from_sse, extract_client_session_id, normalize_responses_payload, stream_upstream_bytes, @@ -156,10 +166,13 @@ def chat_completions() -> Response: for _t in responses_tools_payload: if not (isinstance(_t, dict) and isinstance(_t.get("type"), str)): continue - if _t.get("type") not in ("web_search", "web_search_preview"): + if _t.get("type") not in ("web_search", "web_search_preview", IMAGE_TOOL_TYPE): err = { "error": { - "message": "Only web_search/web_search_preview are supported in responses_tools", + "message": ( + "Only web_search/web_search_preview/image_generation are supported " + "in responses_tools" + ), "code": "RESPONSES_TOOL_UNSUPPORTED", } } @@ -349,7 +362,9 @@ def _extract_usage(evt: Dict[str, Any]) -> Dict[str, int] | None: reasoning_full_text += evt.get("delta") or "" elif kind == "response.output_item.done": item = evt.get("item") or {} - if isinstance(item, dict) and item.get("type") == "function_call": + if isinstance(item, dict) and item.get("type") == "image_generation_call": + full_text += image_markdown_for_item(item) + elif isinstance(item, dict) and item.get("type") == "function_call": call_id = item.get("call_id") or item.get("id") or "" name = item.get("name") or "" args = item.get("arguments") or "" @@ -717,6 +732,189 @@ def responses_create() -> Response: return resp +def _images_error(message: str, status: int, code: str | None = None) -> Response: + body: Dict[str, Any] = {"error": {"message": message}} + if code: + body["error"]["code"] = code + resp = make_response(jsonify(body), status) + for k, v in build_cors_headers().items(): + resp.headers.setdefault(k, v) + return resp + + +def _run_image_request( + upstream_payload: Dict[str, Any], + *, + session_id: str | None, + verbose: bool, +) -> tuple[List[Dict[str, Any]], Dict[str, Any] | None, Response | None, str]: + upstream, error_resp = start_upstream_raw_request( + upstream_payload, + session_id=session_id, + stream=True, + ) + if error_resp is not None: + return [], None, error_resp, "" + + record_rate_limits_from_response(upstream) + + if upstream.status_code >= 400: + try: + raw = upstream.content + err_body = json.loads(raw.decode("utf-8", errors="ignore")) if raw else {} + except Exception: + err_body = {} + finally: + upstream.close() + + err_info = err_body.get("error") if isinstance(err_body.get("error"), dict) else {} + param = err_info.get("param") if isinstance(err_info.get("param"), str) else "" + tools = upstream_payload.get("tools") + tool = tools[0] if isinstance(tools, list) and tools and isinstance(tools[0], dict) else {} + # The backend refuses image parameters it does not know (that is what + # happens with 'n'). Rather than forwarding the 400, retry once with a + # bare tool: the caller gets the image, just without the tweak it asked + # for. + if err_info.get("code") == "unknown_parameter" and param.startswith("tools[") and len(tool) > 1: + if verbose: + print(f"[Images] Backend rejected '{param}'; retrying without the tool parameters") + retry_payload = dict(upstream_payload) + retry_payload["tools"] = [{"type": IMAGE_TOOL_TYPE}] + return _run_image_request(retry_payload, session_id=session_id, verbose=verbose) + + message = err_info.get("message") or "Upstream error" + return [], None, _images_error(str(message), upstream.status_code), "" + + images, usage, error, said = collect_images_from_sse(upstream) + if error is not None: + message = error.get("message") if isinstance(error, dict) else None + return [], None, _images_error(str(message or "Upstream error"), 502), said + return images, usage, None, said + + +@openai_bp.route("/v1/images/generations", methods=["POST"]) +def images_generations() -> Response: + verbose = bool(current_app.config.get("VERBOSE")) + raw = request.get_data(cache=True, as_text=True) or "" + if verbose: + try: + # Truncated on purpose: an 'image' reference arrives here as a data + # URL and would otherwise dump megabytes into the log. + print("IN POST /v1/images/generations\n" + raw[:2000]) + except Exception: + pass + + try: + payload = json.loads(raw) if raw else {} + except Exception: + return _images_error("Invalid JSON body", 400) + if not isinstance(payload, dict): + return _images_error("Request body must be a JSON object", 400) + + prompt = payload.get("prompt") + if not isinstance(prompt, str) or not prompt.strip(): + return _images_error("Missing required parameter: 'prompt'", 400) + + response_format = payload.get("response_format") + if isinstance(response_format, str) and response_format.strip().lower() == "url": + return _images_error( + "response_format 'url' is not available: the backend returns base64 and ChatMock " + "hosts no files. Use 'b64_json'.", + 400, + code="unsupported_value", + ) + + try: + n = int(payload.get("n") or 1) + except (TypeError, ValueError): + return _images_error("Parameter 'n' must be an integer", 400) + if n < 1: + return _images_error("Parameter 'n' must be >= 1", 400) + if n > MAX_IMAGES_PER_REQUEST: + return _images_error( + f"Parameter 'n' must be <= {MAX_IMAGES_PER_REQUEST}: the backend draws one image per " + "call, so every extra unit is another request against your quota.", + 400, + ) + + input_images = payload.get("image") + if isinstance(input_images, str): + input_images = [input_images] + elif not isinstance(input_images, list): + input_images = [] + + orchestrator = payload.get("chat_model") + if not isinstance(orchestrator, str) or not orchestrator.strip(): + orchestrator = current_app.config.get("IMAGE_ORCHESTRATOR_MODEL") or DEFAULT_IMAGE_ORCHESTRATOR_MODEL + orchestrator = normalize_model_name(orchestrator, current_app.config.get("DEBUG_MODEL")) + + upstream_payload = build_image_request_payload( + prompt.strip(), + model=orchestrator, + tool_params=payload, + input_images=[img for img in input_images if isinstance(img, str)], + ) + + session_id = extract_client_session_id(request.headers) + collected: List[Dict[str, Any]] = [] + usage_totals: Dict[str, int] = {"input_tokens": 0, "output_tokens": 0, "total_tokens": 0} + saw_usage = False + # Kept across rounds: with n > 1 the last round is the one that explains an + # empty result, and each round overwrites the previous. + last_text = "" + for _ in range(n): + images, usage, error_resp, said = _run_image_request( + upstream_payload, + session_id=session_id, + verbose=verbose, + ) + if error_resp is not None: + return error_resp + if said: + last_text = said + collected.extend(images) + normalized_usage = usage_to_openai(usage) + if normalized_usage: + saw_usage = True + for key in usage_totals: + usage_totals[key] += int(normalized_usage.get(key) or 0) + + if not collected: + # Say what the model said, instead of guessing why it stopped. This + # error used to assert the prompt "usually" hit moderation, which was + # never verified: the text that carries the actual reason was being + # thrown away one function up. + spoken = last_text.strip() + if spoken: + detail = f' The model answered instead of drawing: "{spoken[:400]}"' + else: + detail = ( + " It returned no text either, so the request was most likely stopped before the " + "model started working." + ) + return _images_error( + "The model finished without generating an image." + detail, + 502, + code="no_image_returned", + ) + + body: Dict[str, Any] = { + "created": int(time.time()), + "data": [image_item_to_openai(item) for item in collected], + } + if saw_usage: + body["usage"] = usage_totals + if verbose: + # Not _log_json: a single image is a couple of megabytes of base64, and + # printing the body would drown the log. + print(f"OUT POST /v1/images/generations ({len(collected)} image(s))") + + resp = make_response(jsonify(body), 200) + for k, v in build_cors_headers().items(): + resp.headers.setdefault(k, v) + return resp + + @openai_bp.route("/v1/models", methods=["GET"]) def list_models() -> Response: expose_variants = bool(current_app.config.get("EXPOSE_REASONING_MODELS")) diff --git a/chatmock/utils.py b/chatmock/utils.py index 96dd314..844a987 100644 --- a/chatmock/utils.py +++ b/chatmock/utils.py @@ -14,6 +14,7 @@ import requests from .config import CLIENT_ID_DEFAULT, OAUTH_TOKEN_URL +from .images_api import image_markdown_for_item from .version import __version__ @@ -471,7 +472,24 @@ def _serialize_tool_args(eff_args: Any) -> str: return json.dumps({"query": eff_args}) else: return "{}" - + + def _close_think_tag(): + """Close the reasoning block before emitting real output, so the content + never lands inside and gets hidden by the client.""" + nonlocal think_open, think_closed + if compat != "think-tags" or not think_open or think_closed: + return + close_chunk = { + "id": response_id, + "object": "chat.completion.chunk", + "created": created, + "model": model, + "choices": [{"index": 0, "delta": {"content": ""}, "finish_reason": None}], + } + yield f"data: {json.dumps(close_chunk)}\n\n".encode("utf-8") + think_open = False + think_closed = True + def _extract_usage(evt: Dict[str, Any]) -> Dict[str, int] | None: try: usage = (evt.get("response") or {}).get("usage") @@ -594,17 +612,7 @@ def _merge_from(src): if kind == "response.output_text.delta": delta = evt.get("delta") or "" - if compat == "think-tags" and think_open and not think_closed: - close_chunk = { - "id": response_id, - "object": "chat.completion.chunk", - "created": created, - "model": model, - "choices": [{"index": 0, "delta": {"content": ""}, "finish_reason": None}], - } - yield f"data: {json.dumps(close_chunk)}\n\n".encode("utf-8") - think_open = False - think_closed = True + yield from _close_think_tag() saw_output = True chunk = { "id": response_id, @@ -616,7 +624,20 @@ def _merge_from(src): yield f"data: {json.dumps(chunk)}\n\n".encode("utf-8") elif kind == "response.output_item.done": item = evt.get("item") or {} - if isinstance(item, dict) and (item.get("type") == "function_call" or item.get("type") == "web_search_call"): + if isinstance(item, dict) and item.get("type") == "image_generation_call": + markdown = image_markdown_for_item(item) + if markdown: + yield from _close_think_tag() + saw_output = True + image_chunk = { + "id": response_id, + "object": "chat.completion.chunk", + "created": created, + "model": model, + "choices": [{"index": 0, "delta": {"content": markdown}, "finish_reason": None}], + } + yield f"data: {json.dumps(image_chunk)}\n\n".encode("utf-8") + elif isinstance(item, dict) and (item.get("type") == "function_call" or item.get("type") == "web_search_call"): call_id = item.get("call_id") or item.get("id") or "" name = item.get("name") or ("web_search" if item.get("type") == "web_search_call" else "") raw_args = item.get("arguments") or item.get("parameters") @@ -782,17 +803,7 @@ def _merge_from(src): m = _extract_usage(evt) if m: upstream_usage = m - if compat == "think-tags" and think_open and not think_closed: - close_chunk = { - "id": response_id, - "object": "chat.completion.chunk", - "created": created, - "model": model, - "choices": [{"index": 0, "delta": {"content": ""}, "finish_reason": None}], - } - yield f"data: {json.dumps(close_chunk)}\n\n".encode("utf-8") - think_open = False - think_closed = True + yield from _close_think_tag() if not sent_stop_chunk: finish_reason = "tool_calls" if saw_function_call else "stop" chunk = { diff --git a/tests/test_routes.py b/tests/test_routes.py index a490670..21d457c 100644 --- a/tests/test_routes.py +++ b/tests/test_routes.py @@ -660,5 +660,248 @@ def close(self) -> None: ) +IMAGE_ITEM = { + "id": "ig_1", + "type": "image_generation_call", + "status": "completed", + "output_format": "png", + "size": "1024x1024", + "result": "QUJD", + "revised_prompt": "a blue cube", +} + + +def image_sse_events(usage_tokens: int = 10) -> list[dict[str, object]]: + return [ + {"type": "response.output_item.done", "output_index": 0, "item": IMAGE_ITEM}, + { + "type": "response.completed", + "response": { + "id": "resp_img", + "status": "completed", + "output": [], + "tool_usage": { + "image_gen": { + "input_tokens": 1, + "output_tokens": usage_tokens, + "total_tokens": usage_tokens + 1, + } + }, + }, + }, + ] + + +class ImageRouteTests(unittest.TestCase): + def setUp(self) -> None: + reset_session_state() + self.app = create_app(model_sync=False) + self.client = self.app.test_client() + + @patch("chatmock.routes_openai.start_upstream_raw_request") + def test_images_generations_returns_b64(self, mock_start) -> None: + mock_start.return_value = (FakeUpstream(image_sse_events()), None) + response = self.client.post( + "/v1/images/generations", + json={"prompt": "a blue cube", "size": "1024x1024", "quality": "low"}, + ) + body = response.get_json() + self.assertEqual(response.status_code, 200) + self.assertEqual(len(body["data"]), 1) + self.assertEqual(body["data"][0]["b64_json"], "QUJD") + self.assertEqual(body["data"][0]["revised_prompt"], "a blue cube") + self.assertEqual(body["usage"]["output_tokens"], 10) + + sent = mock_start.call_args.args[0] + self.assertEqual(sent["tool_choice"], {"type": "image_generation"}) + self.assertEqual(sent["tools"][0]["size"], "1024x1024") + self.assertNotIn("n", sent["tools"][0]) + self.assertIn("1024x1024", sent["instructions"]) + + @patch("chatmock.routes_openai.start_upstream_raw_request") + def test_images_generations_repeats_request_for_n(self, mock_start) -> None: + mock_start.side_effect = [ + (FakeUpstream(image_sse_events(4)), None), + (FakeUpstream(image_sse_events(6)), None), + ] + response = self.client.post("/v1/images/generations", json={"prompt": "x", "n": 2}) + body = response.get_json() + self.assertEqual(response.status_code, 200) + self.assertEqual(mock_start.call_count, 2) + self.assertEqual(len(body["data"]), 2) + self.assertEqual(body["usage"]["output_tokens"], 10) + + @patch("chatmock.routes_openai.start_upstream_raw_request") + def test_images_generations_retries_without_rejected_tool_params(self, mock_start) -> None: + rejection = json.dumps( + { + "error": { + "code": "unknown_parameter", + "message": "Unknown parameter: 'tools[0].quality'.", + "param": "tools[0].quality", + } + } + ).encode("utf-8") + mock_start.side_effect = [ + (FakeUpstream(status_code=400, content=rejection), None), + (FakeUpstream(image_sse_events()), None), + ] + response = self.client.post( + "/v1/images/generations", json={"prompt": "x", "quality": "low"} + ) + self.assertEqual(response.status_code, 200) + self.assertEqual(mock_start.call_count, 2) + self.assertEqual(mock_start.call_args_list[1].args[0]["tools"], [{"type": "image_generation"}]) + + @patch("chatmock.routes_openai.start_upstream_raw_request") + def test_images_generations_reports_empty_result(self, mock_start) -> None: + mock_start.return_value = ( + FakeUpstream([{"type": "response.completed", "response": {"id": "r", "output": []}}]), + None, + ) + response = self.client.post("/v1/images/generations", json={"prompt": "x"}) + self.assertEqual(response.status_code, 502) + self.assertEqual(response.get_json()["error"]["code"], "no_image_returned") + + @patch("chatmock.routes_openai.start_upstream_raw_request") + def test_images_generations_reports_what_the_model_said(self, mock_start) -> None: + """An empty result must carry the model's own words. + + The error used to assert the prompt "usually" hit moderation, which + nobody had verified: the text that holds the real reason was collected + by the stream reader and then dropped. Whoever hits this needs the + reason, not a guess. + """ + mock_start.return_value = ( + FakeUpstream( + [ + { + "type": "response.output_item.done", + "item": { + "type": "message", + "content": [{"type": "output_text", "text": "I can't draw that logo."}], + }, + }, + {"type": "response.completed", "response": {"id": "r", "output": []}}, + ] + ), + None, + ) + response = self.client.post("/v1/images/generations", json={"prompt": "x"}) + self.assertEqual(response.status_code, 502) + body = response.get_json()["error"] + self.assertEqual(body["code"], "no_image_returned") + self.assertIn("draw that logo", body["message"]) + + @patch("chatmock.routes_openai.start_upstream_raw_request") + def test_images_generations_says_when_nothing_came_back(self, mock_start) -> None: + """No image and no text is a different failure, and says so.""" + mock_start.return_value = ( + FakeUpstream([{"type": "response.completed", "response": {"id": "r", "output": []}}]), + None, + ) + response = self.client.post("/v1/images/generations", json={"prompt": "x"}) + message = response.get_json()["error"]["message"] + self.assertIn("no text either", message) + + def test_images_generations_rejects_url_response_format(self) -> None: + response = self.client.post( + "/v1/images/generations", json={"prompt": "x", "response_format": "url"} + ) + self.assertEqual(response.status_code, 400) + self.assertEqual(response.get_json()["error"]["code"], "unsupported_value") + + @patch("chatmock.routes_openai.start_upstream_request") + def test_chat_completions_embeds_image_as_data_url(self, mock_start) -> None: + mock_start.return_value = (FakeUpstream(image_sse_events()), None) + response = self.client.post( + "/v1/chat/completions", + json={ + "model": "gpt-5.4-mini", + "messages": [{"role": "user", "content": "draw a cube"}], + "responses_tools": [{"type": "image_generation"}], + }, + ) + body = response.get_json() + self.assertEqual(response.status_code, 200) + self.assertIn("data:image/png;base64,QUJD", body["choices"][0]["message"]["content"]) + + @patch("chatmock.routes_openai.start_upstream_request") + def test_chat_stream_closes_think_tag_before_image(self, mock_start) -> None: + mock_start.return_value = ( + FakeUpstream( + [ + {"type": "response.reasoning_summary_text.delta", "delta": "planning"}, + {"type": "response.output_item.done", "output_index": 0, "item": IMAGE_ITEM}, + {"type": "response.completed", "response": {"id": "resp_img", "output": []}}, + ] + ), + None, + ) + response = self.client.post( + "/v1/chat/completions", + json={ + "model": "gpt-5.4-mini", + "messages": [{"role": "user", "content": "draw a cube"}], + "responses_tools": [{"type": "image_generation"}], + "stream": True, + }, + ) + content = "" + for line in response.get_data(as_text=True).splitlines(): + if not line.startswith("data: ") or line[6:].strip() == "[DONE]": + continue + event = json.loads(line[6:]) + content += (event.get("choices") or [{}])[0].get("delta", {}).get("content") or "" + image_at = content.find("data:image/") + self.assertGreater(image_at, -1) + # The image is output, not reasoning: it must land after the reasoning + # block is closed, or clients that hide hide the image with it. + self.assertIn("", content[:image_at]) + + @patch("chatmock.routes_ollama.start_upstream_request") + def test_ollama_chat_embeds_image_as_data_url(self, mock_start) -> None: + mock_start.return_value = (FakeUpstream(image_sse_events()), None) + response = self.client.post( + "/api/chat", + json={ + "model": "gpt-5.4-mini", + "messages": [{"role": "user", "content": "draw a cube"}], + "responses_tools": [{"type": "image_generation"}], + "stream": False, + }, + ) + body = response.get_json() + self.assertEqual(response.status_code, 200) + self.assertIn("data:image/png;base64,QUJD", body["message"]["content"]) + + @patch("chatmock.routes_openai.start_upstream_raw_request") + def test_responses_route_rebuilds_output_from_done_items(self, mock_start) -> None: + mock_start.return_value = ( + FakeUpstream( + [ + { + "type": "response.output_item.done", + "output_index": 1, + "item": {"type": "message", "role": "assistant", "content": []}, + }, + {"type": "response.output_item.done", "output_index": 0, "item": {"type": "reasoning"}}, + { + "type": "response.completed", + "response": {"id": "resp_x", "status": "completed", "output": []}, + }, + ] + ), + None, + ) + response = self.client.post( + "/v1/responses", + json={"model": "gpt-5.4-mini", "input": "hi", "stream": False}, + ) + body = response.get_json() + self.assertEqual(response.status_code, 200) + self.assertEqual([item["type"] for item in body["output"]], ["reasoning", "message"]) + + if __name__ == "__main__": unittest.main()