diff --git a/plugins/image_gen/_common.py b/plugins/image_gen/_common.py index efc8b8f4ad..aea9d21cda 100644 --- a/plugins/image_gen/_common.py +++ b/plugins/image_gen/_common.py @@ -242,6 +242,25 @@ class HttpFailure: response: Any = None +def record_token_usage(usage: Any, *, model: str, provider: str, base_url: Optional[str] = None) -> None: + """Record a token-billed image call against the ambient session as task ``image_generation``. + + Token-metered image models (OpenRouter chat-image and Image API models, OpenAI ``gpt-image``) + bill exactly like a chat completion, so they get a ``session_model_usage`` row through the + same chokepoint as auxiliary calls. ``usage`` is the response's usage block (dict or SDK + object). Per-image backends (FAL, xAI, Krea, ...) return no token usage and never call this; + a body without tokens is a no-op inside ``record_aux_usage``, as is running outside a turn. + """ + if not usage: + return + from types import SimpleNamespace + + from agent.aux_accounting import record_aux_usage + + record_aux_usage( + SimpleNamespace(model=model, usage=usage), "image_generation", provider=provider, base_url=base_url) + + def post_json( url: str, *, headers: Dict[str, str], payload: Dict[str, Any], timeout: Any, label: str, error_message: Callable[[Any, Exception], str] = requests_error_message, diff --git a/plugins/image_gen/openai/__init__.py b/plugins/image_gen/openai/__init__.py index 2c0742ee8a..260412c758 100644 --- a/plugins/image_gen/openai/__init__.py +++ b/plugins/image_gen/openai/__init__.py @@ -14,7 +14,7 @@ from agent.image_gen_provider import DEFAULT_ASPECT_RATIO, resolve_aspect_ratio, from plugins.image_gen._common import ( GPT_IMAGE_2_API_MODEL as API_MODEL, GPT_IMAGE_2_DEFAULT as DEFAULT_MODEL, GPT_IMAGE_2_TIERS, StaticImageGenProvider, collect_source_images, error_factory, import_openai, materialize_image, - openai_importable, prompt_required_error, resolve_static_model, size_for) + openai_importable, prompt_required_error, record_token_usage, resolve_static_model, size_for) logger = logging.getLogger(__name__) @@ -153,6 +153,8 @@ class OpenAIImageGenProvider(StaticImageGenProvider): model=tier_id, prompt=prompt, aspect=aspect, log=logger) if err: return err + # gpt-image bills per text/image token; the tier id is a Hermes label, the API model prices. + record_token_usage(getattr(response, "usage", None), model=meta["api_model"], provider="openai") extra: Dict[str, Any] = {"size": size, "quality": meta["quality"]} if getattr(first, "revised_prompt", None): extra["revised_prompt"] = first.revised_prompt diff --git a/plugins/image_gen/openrouter/__init__.py b/plugins/image_gen/openrouter/__init__.py index 3a607d1a38..b7b4eb9a62 100644 --- a/plugins/image_gen/openrouter/__init__.py +++ b/plugins/image_gen/openrouter/__init__.py @@ -22,7 +22,7 @@ from typing import Any, Dict, List, Optional, Tuple from agent.image_gen_provider import ( DEFAULT_ASPECT_RATIO, ImageGenProvider, error_response, resolve_aspect_ratio, save_b64_image, save_url_image, success_response) -from plugins.image_gen._common import error_factory, load_image_gen_config, post_json +from plugins.image_gen._common import error_factory, load_image_gen_config, post_json, record_token_usage logger = logging.getLogger(__name__) @@ -692,7 +692,7 @@ class OpenRouterCompatImageProvider(ImageGenProvider): return _fail(f"Could not save generated image: {exc}", "io_error") if not saved: return _fail(f"{self._display} response carried neither b64_json nor url.", "empty_response") - _record_image_api_usage(body, model_id, provider=self._name) + record_token_usage(_dict_at(body, "usage"), model=model_id, provider=self._name, base_url=base_url) return success_response( image=saved[0], model=model_id, prompt=prompt, aspect_ratio=semantic_aspect, provider=self._name, modality="image" if usable_refs else "text", @@ -741,6 +741,7 @@ class OpenRouterCompatImageProvider(ImageGenProvider): saved_path = save_url_image(first, prefix=f"{self._name}_gen") except Exception as exc: # noqa: BLE001 return fail(f"Could not save generated image: {exc}", "io_error"), None + record_token_usage(_dict_at(result, "usage"), model=model_id, provider=self._name, base_url=base_url) return success_response( image=str(saved_path), model=model_id, prompt=prompt, aspect_ratio=aspect, provider=self._name, ), None @@ -803,30 +804,6 @@ class OpenRouterCompatImageProvider(ImageGenProvider): model=model_chain[-1] if model_chain else "", prompt=prompt, aspect_ratio=aspect) -def _record_image_api_usage(body: Any, model_id: str, *, provider: str) -> None: - """Record a token-billed Image API call against the ambient session (#114324). - - Best-effort: accounting must never break image generation. No-ops outside - an agent turn or when the response carries no token usage. - """ - try: - usage = _dict_at(body, "usage") - if not isinstance(usage.get("total_tokens"), int): - return - from agent.aux_accounting import record_aux_usage - from types import SimpleNamespace - - record_aux_usage( - SimpleNamespace(model=model_id, usage=SimpleNamespace( - prompt_tokens=usage.get("prompt_tokens", 0), - completion_tokens=usage.get("completion_tokens", 0), - total_tokens=usage.get("total_tokens", 0), - )), - "image_gen", provider=provider) - except Exception: # noqa: BLE001 - logger.debug("%s: image usage recording failed (non-fatal)", provider, exc_info=True) - - def _build_providers() -> List[OpenRouterCompatImageProvider]: return [ OpenRouterCompatImageProvider( diff --git a/tests/plugins/image_gen/test_openai_provider.py b/tests/plugins/image_gen/test_openai_provider.py index 93ad401321..9a315cbb07 100644 --- a/tests/plugins/image_gen/test_openai_provider.py +++ b/tests/plugins/image_gen/test_openai_provider.py @@ -172,6 +172,34 @@ class TestGenerate: # gpt-image-2 rejects response_format — we must NOT send it. assert "response_format" not in call_kwargs + def test_token_usage_reaches_session_accounting(self, provider): + """gpt-image bills per token: the Images API ``usage`` block lands as one + ``image_generation`` row keyed on the API model, not the Hermes tier label.""" + from agent import aux_accounting + + recorded = [] + + class _DB: + def record_auxiliary_usage(self, *args, **kwargs): + recorded.append((args, kwargs)) + + response = _fake_response(b64=_b64_png()) + response.usage = SimpleNamespace(input_tokens=23, output_tokens=1056, total_tokens=1079) + fake_client = MagicMock() + fake_client.images.generate.return_value = response + token = aux_accounting.set_accounting_context(_DB(), "sess-1") + try: + with _patched_openai(fake_client): + result = provider.generate("a cat", aspect_ratio="landscape") + finally: + aux_accounting.reset_accounting_context(token) + + assert result["success"] is True + ((session_id, task), kwargs), = recorded + assert (session_id, task) == ("sess-1", "image_generation") + assert (kwargs["model"], kwargs["billing_provider"]) == ("gpt-image-2", "openai") + assert (kwargs["input_tokens"], kwargs["output_tokens"]) == (23, 1056) + @pytest.mark.parametrize("api_model,quality", [ ("gpt-image-2", quality) for quality in ("low", "medium", "high") ] + [ diff --git a/tests/plugins/image_gen/test_openrouter_compat_provider.py b/tests/plugins/image_gen/test_openrouter_compat_provider.py index 5e36083325..ec6a1aa5a0 100644 --- a/tests/plugins/image_gen/test_openrouter_compat_provider.py +++ b/tests/plugins/image_gen/test_openrouter_compat_provider.py @@ -703,58 +703,48 @@ class TestImageApiSurface: assert result["exact_aspect_ratio"] == "9:16" assert result["image"] == "/tmp/i.png" - def test_token_billed_call_records_session_usage(self): - """#114324: an OpenRouter token-billed call must reach session_model_usage.""" + _USAGE = {"prompt_tokens": 1000, "completion_tokens": 128, "total_tokens": 1128} + + @pytest.mark.parametrize("surface, model, usage", [ + ("chat", "openai/gpt-5.4-image-2", _USAGE), # default chain: token-billed via /chat/completions + ("images", "krea/krea-2-medium", _USAGE), # curated Image API model + ("images", "krea/krea-2-medium", None), # flat-fee body without usage: no write + ]) + def test_token_usage_reaches_session_accounting(self, surface, model, usage): + """A response carrying token usage records one ``image_generation`` row on the ambient + session; a body without usage records nothing.""" from agent import aux_accounting recorded = [] class _DB: def record_auxiliary_usage(self, *args, **kwargs): - recorded.append(kwargs) + recorded.append((args, kwargs)) + if surface == "chat": + response = _mock_chat_response([_PNG_DATA_URI]) + response.json.return_value["usage"] = dict(usage) + else: + response = _mock_image_api_response(usage=usage) token = aux_accounting.set_accounting_context(_DB(), "sess-1") try: with patch(_RUNTIME, return_value=_runtime_ok()), \ - patch("requests.post", return_value=_mock_image_api_response( - usage={"total_tokens": 1128, "prompt_tokens": 1000, - "completion_tokens": 128})), \ + patch("requests.post", return_value=response), \ patch("plugins.image_gen.openrouter.save_b64_image", return_value=Path("/tmp/i.png")): - result = _openrouter_image_api().generate( - prompt="p", aspect_ratio="portrait", model="krea/krea-2-medium" - ) + result = _openrouter_image_api().generate(prompt="p", aspect_ratio="portrait", model=model) finally: aux_accounting.reset_accounting_context(token) assert result["success"] is True - assert len(recorded) == 1 - assert recorded[0]["input_tokens"] == 1000 - assert recorded[0]["output_tokens"] == 128 - assert recorded[0]["model"] == "krea/krea-2-medium" + if usage is None: + assert recorded == [] + return + ((session_id, task), kwargs), = recorded + assert (session_id, task) == ("sess-1", "image_generation") + assert kwargs["model"] == model + assert kwargs["billing_provider"] == "openrouter" + assert (kwargs["input_tokens"], kwargs["output_tokens"]) == (1000, 128) - def test_untokened_call_records_nothing(self): - """No token usage in the body: no session write (flat-fee image models).""" - from agent import aux_accounting - - recorded = [] - - class _DB: - def record_auxiliary_usage(self, *args, **kwargs): - recorded.append(kwargs) - - token = aux_accounting.set_accounting_context(_DB(), "sess-1") - try: - with patch(_RUNTIME, return_value=_runtime_ok()), \ - patch("requests.post", return_value=_mock_image_api_response()), \ - patch("plugins.image_gen.openrouter.save_b64_image", return_value=Path("/tmp/i.png")): - result = _openrouter_image_api().generate( - prompt="p", aspect_ratio="portrait", model="krea/krea-2-medium" - ) - finally: - aux_accounting.reset_accounting_context(token) - - assert result["success"] is True - assert recorded == [] def test_multiple_images_land_in_additional_images(self): entries = [ diff --git a/website/docs/user-guide/features/image-generation.md b/website/docs/user-guide/features/image-generation.md index f184308054..f3390d6bb7 100644 --- a/website/docs/user-guide/features/image-generation.md +++ b/website/docs/user-guide/features/image-generation.md @@ -319,6 +319,7 @@ If upscaling fails (network issue, rate limit), the original image is returned a 3. **Submission** — `_submit_fal_request()` routes via direct FAL credentials or the managed Nous gateway, according to the stored `image_gen.provider` selection. 4. **Upscaling** — runs only when the agent passed `upscale: true`; every model's catalog default is off. 5. **Delivery** — final image URL returned to the agent, which emits a `MEDIA:` tag that platform adapters convert to native media. +6. **Usage accounting** — token-billed image models (OpenRouter chat-image and Image API models such as `google/gemini-3.1-flash-lite-image`, OpenAI `gpt-image`) return real token counts, so each call is recorded in `session_model_usage` as task `image_generation` under the billing provider and model, and shows up in `hermes insights` and the dashboard's Usage analytics alongside other model calls. Per-image backends (FAL, xAI, Krea, ...) return no token usage and are not recorded there. ## Debugging