diff --git a/tests/tools/test_hunyuan_image.py b/tests/tools/test_hunyuan_image.py new file mode 100644 index 00000000..efb1f9d5 --- /dev/null +++ b/tests/tools/test_hunyuan_image.py @@ -0,0 +1,488 @@ +"""Unit tests for hunyuan_image — TokenHub 混元生图 3.0 tool.""" + +from __future__ import annotations + +import json +import sys +from pathlib import Path +from unittest.mock import patch + +import pytest + +PROJECT_ROOT = Path(__file__).resolve().parent.parent.parent +sys.path.insert(0, str(PROJECT_ROOT)) + +from tools.base_tool import ToolStatus + +# --------------------------------------------------------------------------- +# Tool discovery & metadata +# --------------------------------------------------------------------------- + + +def test_hunyuan_image_is_discovered_by_registry(): + from tools.tool_registry import ToolRegistry + + registry = ToolRegistry() + registry.discover() + + tool = registry.get("hunyuan_image") + assert tool is not None + assert tool.provider == "hunyuan_cloud" + assert tool.capability == "image_generation" + assert tool.name == "hunyuan_image" + + +def test_hunyuan_image_metadata(): + from tools.graphics.hunyuan_image import HunyuanImage + + tool = HunyuanImage() + info = tool.get_info() + + assert info["tier"] == "generate" + assert info["stability"] == "experimental" + assert info["runtime"] == "api" + assert "text_to_image" in info["capabilities"] + assert info["supports"]["seed"] is True + assert info["supports"]["reference_image"] is True + assert info["supports"]["prompt_rewrite"] is True + assert info["supports"]["negative_prompt"] is False + + +# --------------------------------------------------------------------------- +# Status reporting +# --------------------------------------------------------------------------- + + +def test_status_unavailable_when_no_api_key(monkeypatch): + from tools.graphics.hunyuan_image import HunyuanImage + + monkeypatch.delenv("TENCENT_TOKENHUB_API_KEY", raising=False) + assert HunyuanImage().get_status() == ToolStatus.UNAVAILABLE + + +def test_status_available_when_api_key_set(monkeypatch): + from tools.graphics.hunyuan_image import HunyuanImage + + monkeypatch.setenv("TENCENT_TOKENHUB_API_KEY", "test-key-123") + assert HunyuanImage().get_status() == ToolStatus.AVAILABLE + + +def test_api_key_filters_comment_like_values(monkeypatch): + from tools.graphics.hunyuan_image import HunyuanImage + + monkeypatch.setenv("TENCENT_TOKENHUB_API_KEY", "# this-is-a-comment") + assert HunyuanImage()._api_key() is None + assert HunyuanImage().get_status() == ToolStatus.UNAVAILABLE + + +def test_api_key_strips_whitespace(monkeypatch): + from tools.graphics.hunyuan_image import HunyuanImage + + monkeypatch.setenv("TENCENT_TOKENHUB_API_KEY", " my-key ") + assert HunyuanImage()._api_key() == "my-key" + + +# --------------------------------------------------------------------------- +# Cost & runtime estimation +# --------------------------------------------------------------------------- + + +def test_estimate_cost(): + from tools.graphics.hunyuan_image import HunyuanImage + + tool = HunyuanImage() + cost = tool.estimate_cost({}) + assert cost > 0 + assert cost == pytest.approx(0.08, rel=0.1) + + +def test_estimate_runtime(): + from tools.graphics.hunyuan_image import HunyuanImage + + tool = HunyuanImage() + runtime = tool.estimate_runtime({}) + assert runtime == 120.0 + + +# --------------------------------------------------------------------------- +# Payload construction (_build_payload) +# --------------------------------------------------------------------------- + + +def test_build_payload_minimal(): + from tools.graphics.hunyuan_image import HunyuanImage + + tool = HunyuanImage() + payload = tool._build_payload({"prompt": "a cat"}) + assert payload == {"prompt": "a cat"} + + +def test_build_payload_all_params(): + from tools.graphics.hunyuan_image import HunyuanImage + + tool = HunyuanImage() + payload = tool._build_payload({ + "prompt": "a cat", + "resolution": "768:1024", + "seed": 42, + "revise": 0, + "logo_add": 1, + "images": ["https://example.com/ref.jpg"], + }) + assert payload["prompt"] == "a cat" + assert payload["resolution"] == "768:1024" + assert payload["seed"] == 42 + assert payload["revise"] == 0 + assert payload["logo_add"] == 1 + assert payload["images"] == ["https://example.com/ref.jpg"] + + +def test_build_payload_with_logo_param(): + from tools.graphics.hunyuan_image import HunyuanImage + + tool = HunyuanImage() + payload = tool._build_payload({ + "prompt": "a cat", + "logo_param": {"logo_url": "https://example.com/wm.png"}, + }) + assert payload["logo_param"] == {"logo_url": "https://example.com/wm.png"} + + +def test_build_payload_logo_param_skips_empty(): + from tools.graphics.hunyuan_image import HunyuanImage + + tool = HunyuanImage() + payload = tool._build_payload({ + "prompt": "a cat", + "logo_param": {}, + }) + assert "logo_param" not in payload + + +def test_build_payload_omits_none_seed(): + from tools.graphics.hunyuan_image import HunyuanImage + + tool = HunyuanImage() + payload = tool._build_payload({"prompt": "a cat", "seed": None}) + assert "seed" not in payload + + +# --------------------------------------------------------------------------- +# Image resolution (_resolve_images) +# --------------------------------------------------------------------------- + + +def test_resolve_images_passes_urls_through(): + from tools.graphics.hunyuan_image import HunyuanImage + + refs = [ + "https://example.com/a.jpg", + "data:image/png;base64,abc123", + ] + resolved = HunyuanImage._resolve_images(refs) + assert resolved == refs + + +def test_resolve_images_encodes_local_file(tmp_path): + from tools.graphics.hunyuan_image import HunyuanImage + + img = tmp_path / "test.png" + img.write_bytes(b"fake-png-data") + + resolved = HunyuanImage._resolve_images([str(img)]) + assert len(resolved) == 1 + assert resolved[0].startswith("data:image/png;base64,") + + +def test_resolve_images_raises_on_missing_file(): + from tools.graphics.hunyuan_image import HunyuanImage + + with pytest.raises(FileNotFoundError): + HunyuanImage._resolve_images(["/nonexistent/path.jpg"]) + + +def test_resolve_images_raises_on_oversized_file(tmp_path): + from tools.graphics.hunyuan_image import HunyuanImage + + big = tmp_path / "big.jpg" + big.write_bytes(b"x" * (7 * 1024 * 1024)) # 7MB > 6MB limit + + with pytest.raises(ValueError, match="too large"): + HunyuanImage._resolve_images([str(big)]) + + +def test_resolve_images_detects_mime_from_extension(tmp_path): + from tools.graphics.hunyuan_image import HunyuanImage + + cases = [ + ("ref.jpg", "image/jpeg"), + ("ref.jpeg", "image/jpeg"), + ("ref.png", "image/png"), + ("ref.bmp", "image/bmp"), + ("ref.tiff", "image/tiff"), + ("ref.tif", "image/tiff"), + ("ref.webp", "image/webp"), + ("ref.unknown", "image/png"), # fallback + ] + for filename, expected_mime in cases: + f = tmp_path / filename + f.write_bytes(b"data") + resolved = HunyuanImage._resolve_images([str(f)]) + assert resolved[0].startswith(f"data:{expected_mime};base64,") + + +# --------------------------------------------------------------------------- +# Output path resolution (_resolve_output_paths) +# --------------------------------------------------------------------------- + + +def test_resolve_output_paths_single(): + from tools.graphics.hunyuan_image import HunyuanImage + + paths = HunyuanImage._resolve_output_paths("/out/img.png", 1) + assert len(paths) == 1 + assert paths[0] == Path("/out/img.png") + + +def test_resolve_output_paths_multi(): + from tools.graphics.hunyuan_image import HunyuanImage + + paths = HunyuanImage._resolve_output_paths("/out/img.png", 3) + assert len(paths) == 3 + assert [p.name for p in paths] == ["img_1.png", "img_2.png", "img_3.png"] + assert len(set(paths)) == 3 + + +# --------------------------------------------------------------------------- +# Auth headers +# --------------------------------------------------------------------------- + + +def test_auth_headers(): + from tools.graphics.hunyuan_image import HunyuanImage + + headers = HunyuanImage._auth_headers("my-api-key") + assert headers["Authorization"] == "Bearer my-api-key" + assert headers["Content-Type"] == "application/json" + + +# --------------------------------------------------------------------------- +# JSON error handling +# --------------------------------------------------------------------------- + + +def test_json_or_raise_parses_valid_json(): + from tools.graphics.hunyuan_image import HunyuanImage + + class FakeResp: + status_code = 200 + + def json(self): + return {"status": "ok"} + + assert HunyuanImage._json_or_raise(FakeResp()) == {"status": "ok"} + + +def test_json_or_raise_raises_on_invalid_json(): + from tools.graphics.hunyuan_image import HunyuanImage + + class FakeResp: + status_code = 500 + + def json(self): + raise ValueError("not json") + + with pytest.raises(RuntimeError, match="Non-JSON response"): + HunyuanImage._json_or_raise(FakeResp()) + + +def test_check_response_passes_clean_payload(): + from tools.graphics.hunyuan_image import HunyuanImage + + HunyuanImage._check_response({"status": "completed"}) # no error -> no raise + + +def test_check_response_raises_on_error_field(): + from tools.graphics.hunyuan_image import HunyuanImage + + with pytest.raises(RuntimeError, match="TokenHub API error"): + HunyuanImage._check_response({ + "error": {"code": "AUTH_FAILED", "message": "invalid key"}, + }) + + +def test_check_response_raises_on_error_without_code(): + from tools.graphics.hunyuan_image import HunyuanImage + + with pytest.raises(RuntimeError, match="TokenHub API error"): + HunyuanImage._check_response({"error": {"message": "something broke"}}) + + +# --------------------------------------------------------------------------- +# Safe error redaction +# --------------------------------------------------------------------------- + + +def test_safe_error_redacts_api_key(monkeypatch): + from tools.graphics.hunyuan_image import HunyuanImage + + monkeypatch.setenv("TENCENT_TOKENHUB_API_KEY", "secret-key-abc") + msg = HunyuanImage._safe_error(Exception("failed with secret-key-abc")) + assert "secret-key-abc" not in msg + assert "[redacted]" in msg + + +def test_safe_error_preserves_other_text(monkeypatch): + from tools.graphics.hunyuan_image import HunyuanImage + + monkeypatch.setenv("TENCENT_TOKENHUB_API_KEY", "sk-123") + msg = HunyuanImage._safe_error(Exception("network timeout: connection refused")) + assert "network timeout" in msg + assert "sk-123" not in msg + + +# --------------------------------------------------------------------------- +# Execute guards +# --------------------------------------------------------------------------- + + +def test_execute_returns_error_without_api_key(monkeypatch): + from tools.graphics.hunyuan_image import HunyuanImage + + monkeypatch.delenv("TENCENT_TOKENHUB_API_KEY", raising=False) + result = HunyuanImage().execute({"prompt": "a cat"}) + assert not result.success + assert "TENCENT_TOKENHUB_API_KEY" in result.error + + +# --------------------------------------------------------------------------- +# Dry run +# --------------------------------------------------------------------------- + + +def test_dry_run_no_side_effects(monkeypatch): + from tools.graphics.hunyuan_image import HunyuanImage + + monkeypatch.setenv("TENCENT_TOKENHUB_API_KEY", "test-key") + tool = HunyuanImage() + info = tool.dry_run({"prompt": "a cat"}) + assert info["tool"] == "hunyuan_image" + assert info["estimated_cost_usd"] > 0 + assert info["would_execute"] is True + + +# --------------------------------------------------------------------------- +# End-to-end with mocked API +# --------------------------------------------------------------------------- + + +def test_execute_full_flow_with_mocked_api(monkeypatch, tmp_path): + """Simulate the full submit → poll → download flow.""" + from tools.graphics.hunyuan_image import HunyuanImage, _MODEL, _HOST + + monkeypatch.setenv("TENCENT_TOKENHUB_API_KEY", "test-key") + + # Track calls across all mocked endpoints + api_calls = [] + poll_count = [0] # mutable counter for poll iteration + + class _FakeResp: + status_code = 200 + def __init__(self, data, content=None): + self._data = data + self.content = content + def json(self): + return self._data + def raise_for_status(self): + pass + + submit_url = f"https://{_HOST}/v1/api/image/submit" + query_url = f"https://{_HOST}/v1/api/image/query" + + # Safety sentinel: if a real network call slips through the mock, + # this flag is flipped and we fail-fast instead of leaking files. + mock_active = [False] + + def fake_post(url, *, json, headers, timeout): + mock_active[0] = True + api_calls.append(("post", url, json)) + if "submit" in url: + return _FakeResp({"id": "job-001", "status": "queued"}) + elif "query" in url: + poll_count[0] += 1 + if poll_count[0] == 1: + return _FakeResp({"status": "running"}) + return _FakeResp({ + "status": "completed", + "data": [{"url": "https://example.com/result.png"}], + }) + raise RuntimeError(f"Unexpected URL: {url}") + + def fake_get(url, timeout): + mock_active[0] = True + api_calls.append(("get", url)) + return _FakeResp({}, content=b"fake-image-data") + + with ( + patch("requests.post", side_effect=fake_post), + patch("requests.get", side_effect=fake_get), + ): + out = tmp_path / "gen.png" + result = HunyuanImage().execute({ + "prompt": "a programmer coding", + "resolution": "1024:1024", + "seed": 12345, + "revise": 1, + "logo_add": 0, + "output_path": str(out), + }) + + # Guard: mock must have been exercised — if not, a real API call leaked + assert mock_active[0], ( + "Mock was never triggered — a real API call may have leaked. " + "Check that requests.post / requests.get patching is effective." + ) + + assert result.success, result.error + assert result.data["provider"] == "hunyuan_cloud" + assert result.data["model"] == _MODEL + assert result.data["task_id"] == "job-001" + assert result.data["resolution"] == "1024:1024" + assert result.data["images_generated"] == 1 + assert result.artifacts == [str(out)] + assert out.read_bytes() == b"fake-image-data" + + # Verify submit payload was correct + submit_calls = [c for c in api_calls if "submit" in c[1]] + assert len(submit_calls) == 1 + _, _, submit_body = submit_calls[0] + assert submit_body["prompt"] == "a programmer coding" + assert submit_body["resolution"] == "1024:1024" + assert submit_body["seed"] == 12345 + assert submit_body["revise"] == 1 + assert submit_body["logo_add"] == 0 + assert submit_body["model"] == _MODEL + + # Verify polling happened (submit + at least 1 query + download) + assert any("query" in c[1] for c in api_calls) + assert any(c[0] == "get" for c in api_calls) + + +def test_execute_with_local_reference_images(monkeypatch, tmp_path): + """Reference images from local paths should be base64-encoded in payload.""" + from tools.graphics.hunyuan_image import HunyuanImage + + monkeypatch.setenv("TENCENT_TOKENHUB_API_KEY", "test-key") + + ref_img = tmp_path / "ref.png" + ref_img.write_bytes(b"reference-data") + + tool = HunyuanImage() + payload = tool._build_payload({ + "prompt": "enhance this", + "images": [str(ref_img)], + }) + + assert "images" in payload + assert len(payload["images"]) == 1 + assert payload["images"][0].startswith("data:image/png;base64,") diff --git a/tools/graphics/hunyuan_image.py b/tools/graphics/hunyuan_image.py new file mode 100644 index 00000000..65a52c8c --- /dev/null +++ b/tools/graphics/hunyuan_image.py @@ -0,0 +1,579 @@ +"""Tencent Hunyuan (腾讯混元) cloud image generation (3.0) via TokenHub API. + +Calls the Tencent TokenHub API (tokenhub.tencentmaas.com) using simple Bearer +token authentication. This is the OpenAI-compatible API gateway for Tencent +Hunyuan image models — no TC3-HMAC-SHA256 signing required. + +API flow: POST /v1/api/image/submit -> poll /v1/api/image/query -> +download data[].url. + +Authentication uses a TokenHub API key obtained from the Tencent Cloud +TokenHub console (https://console.cloud.tencent.com/tokenhub). +""" + +from __future__ import annotations + +import os +import time +from pathlib import Path +from typing import Any + +from tools.base_tool import ( + BaseTool, + Determinism, + ExecutionMode, + ResourceProfile, + RetryPolicy, + ToolResult, + ToolRuntime, + ToolStability, + ToolStatus, + ToolTier, +) + +_HOST = "tokenhub.tencentmaas.com" +_SUBMIT_PATH = "/v1/api/image/submit" +_QUERY_PATH = "/v1/api/image/query" + +# TokenHub model identifier for 混元生图 3.0 +_MODEL = "hy-image-v3.0" + + +class HunyuanImage(BaseTool): + """Tencent Hunyuan cloud image generation (3.0) via TokenHub API.""" + + name = "hunyuan_image" + version = "0.1.0" + tier = ToolTier.GENERATE + capability = "image_generation" + provider = "hunyuan_cloud" + stability = ToolStability.EXPERIMENTAL + execution_mode = ExecutionMode.ASYNC + determinism = Determinism.SEEDED + runtime = ToolRuntime.API + + dependencies = [] + install_instructions = ( + "Set TENCENT_TOKENHUB_API_KEY to your Tencent Cloud TokenHub API key.\n" + " Get it at https://console.cloud.tencent.com/tokenhub" + ) + agent_skills = [] + + capabilities = ["generate_image", "text_to_image"] + supports = { + "negative_prompt": False, + "seed": True, + "custom_size": True, + "reference_image": True, + "prompt_rewrite": True, + } + best_for = [ + "Hunyuan text-to-image via Tencent TokenHub API", + "simple Bearer-token auth (no TC3 signing required)", + "direct Tencent Cloud quota usage (not through a third-party gateway)", + "Chinese-language prompt understanding", + ] + not_good_for = [ + "offline generation or air-gapped environments", + "users without Tencent Cloud account and real-name verification", + ] + fallback_tools = ["dashscope_image", "flux_image", "openai_image", "recraft_image"] + + input_schema = { + "type": "object", + "required": ["prompt"], + "properties": { + "prompt": { + "type": "string", + "maxLength": 8192, + "description": ( + "Image description. Max 8192 UTF-8 characters. " + "Supports Chinese and English. Be specific about subject, " + "composition, style, and mood." + ), + }, + "images": { + "type": "array", + "items": {"type": "string"}, + "maxItems": 3, + "description": ( + "Reference images per upstream Images.N param (max 3). " + "Each entry is a publicly accessible URL or a local file path " + "(auto-encoded to base64 data URI). " + "Single image: 50-5000px per side, base64 < 6MB. " + "Formats: jpg/png/jpeg/webp/bmp/tiff." + ), + }, + "resolution": { + "type": "string", + "default": "1024:1024", + "description": ( + 'Image resolution as "W:H" (colon separator, per upstream ' + "Resolution param). W, H in [512, 2048], product (W*H) <= " + '1024x1024 pixels. Examples: "1024:1024", "768:1024", ' + '"1024:576".' + ), + }, + "seed": { + "type": "integer", + "minimum": 1, + "maximum": 4294967295, + "description": ( + "Random seed in [1, 4294967295]. " + "Note: seed is ignored when revise is enabled (default)." + ), + }, + "revise": { + "type": "integer", + "enum": [0, 1], + "default": 1, + "description": ( + "Prompt auto-rewrite toggle per upstream Revise param. " + "1 = enabled (default, adds ~20s processing), 0 = disabled. " + "When disabled, caller should handle prompt rewriting." + ), + }, + "logo_add": { + "type": "integer", + "enum": [0, 1], + "default": 1, + "description": ( + "Add 'AI-generated' watermark per upstream LogoAdd param. " + "1 = add watermark (default), 0 = no watermark. " + "Values other than 0 or 1 are treated as 1." + ), + }, + "logo_param": { + "type": "object", + "properties": { + "logo_url": { + "type": "string", + "description": "Custom watermark image URL.", + }, + "logo_image": { + "type": "string", + "description": "Custom watermark image as base64-encoded string.", + }, + }, + "description": ( + "Custom watermark settings per upstream LogoParam. " + "Default: \"图片由 AI 生成\" at bottom-right. " + "Note: custom watermarks may not be supported by the engine " + "(error: InvalidParameterValue.LogoParamErr)." + ), + }, + "output_path": { + "type": "string", + "description": "Output file path for the generated image (PNG).", + }, + "poll_interval_seconds": { + "type": "number", + "minimum": 2, + "default": 5.0, + "description": "Seconds between status polls.", + }, + "timeout_seconds": { + "type": "integer", + "minimum": 60, + "default": 600, + "description": "Maximum seconds to wait for generation.", + }, + }, + } + + resource_profile = ResourceProfile( + cpu_cores=1, ram_mb=512, vram_mb=0, disk_mb=100, network_required=True, + ) + retry_policy = RetryPolicy( + max_retries=2, + backoff_seconds=2.0, + retryable_errors=["rate_limit", "timeout"], + ) + idempotency_key_fields = [ + "prompt", + "resolution", + "images", + "seed", + "revise", + "logo_add", + ] + side_effects = [ + "writes image file to output_path", + "calls Tencent TokenHub API (Bearer-token submit + poll + download)", + ] + user_visible_verification = [ + "Inspect generated image for quality and prompt adherence", + "Check for watermark if logo_add=0 was requested", + ] + + # ------------------------------------------------------------------ + # Credential helpers + # ------------------------------------------------------------------ + + @staticmethod + def _api_key() -> str | None: + val = os.environ.get("TENCENT_TOKENHUB_API_KEY", "") + if val and not val.strip().startswith("#"): + return val.strip() + return None + + # ------------------------------------------------------------------ + # Tool contract methods + # ------------------------------------------------------------------ + + def get_status(self) -> ToolStatus: + if self._api_key(): + return ToolStatus.AVAILABLE + return ToolStatus.UNAVAILABLE + + def estimate_cost(self, inputs: dict[str, Any]) -> float: + """Estimate cost in USD. + + Tencent TokenHub credit-based pricing (1 credit = 1.2 RMB ≈ $0.167 USD): + - hy-image-v3.0: ~0.5 credits/image → ~$0.08 + + Source: https://cloud.tencent.com.cn/document/product/1823/130054 + """ + _CREDIT_TO_USD = 1.2 / 7.2 # 1 credit = 1.2 RMB, ~7.2 RMB/USD + credits = 0.5 + return round(credits * _CREDIT_TO_USD, 2) + + def estimate_runtime(self, inputs: dict[str, Any]) -> float: + """Estimate wall-clock time in seconds. + + Per upstream docs, prompt rewrite (revise=1) adds ~20s. Including + queuing and download, 120s is a safe upper-bound. + """ + return 120.0 + + # ------------------------------------------------------------------ + # Main execution + # ------------------------------------------------------------------ + + def execute(self, inputs: dict[str, Any]) -> ToolResult: + api_key = self._api_key() + if not api_key: + return ToolResult( + success=False, + error="TENCENT_TOKENHUB_API_KEY not set. " + self.install_instructions, + ) + + start = time.time() + try: + result = self._generate(inputs, api_key=api_key) + except Exception as exc: + return ToolResult( + success=False, + error=f"Hunyuan TokenHub image generation failed: {self._safe_error(exc)}", + ) + + result.duration_seconds = round(time.time() - start, 2) + return result + + # ------------------------------------------------------------------ + # Generation pipeline + # ------------------------------------------------------------------ + + def _generate( + self, inputs: dict[str, Any], *, api_key: str, + ) -> ToolResult: + import requests + + # Guard: refuse to make paid API calls without an explicit output_path. + # A CWD-relative default would leak files into the project root when + # called by selectors or other automated tooling. + if not inputs.get("output_path"): + return ToolResult( + success=False, + error="output_path is required for hunyuan_image generation.", + ) + + payload = self._build_payload(inputs) + task_id = self._submit_task(payload, model=_MODEL, api_key=api_key) + image_urls = self._poll_task( + task_id, + model=_MODEL, + api_key=api_key, + poll_interval=float(inputs.get("poll_interval_seconds", 5.0)), + timeout_seconds=int(inputs.get("timeout_seconds", 600)), + ) + + output_paths = self._resolve_output_paths( + inputs["output_path"], + count=len(image_urls), + ) + for path, url in zip(output_paths, image_urls): + download = requests.get(url, timeout=120) + download.raise_for_status() + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(download.content) + + return ToolResult( + success=True, + data={ + "provider": "hunyuan_cloud", + "route": "tokenhub", + "model": _MODEL, + "prompt": inputs["prompt"], + "resolution": payload.get("resolution", "1024:1024"), + "revise": payload.get("revise", 1), + "logo_add": payload.get("logo_add", 1), + "task_id": task_id, + "output": str(output_paths[0]), + "outputs": [str(p) for p in output_paths], + "images_generated": len(output_paths), + }, + artifacts=[str(p) for p in output_paths], + cost_usd=self.estimate_cost(inputs), + model=_MODEL, + ) + + # ------------------------------------------------------------------ + # Payload construction + # ------------------------------------------------------------------ + + def _build_payload(self, inputs: dict[str, Any]) -> dict[str, Any]: + """Build the request body for TokenHub image submit. + + TokenHub async endpoints use snake_case versions of the upstream + SubmitTextToImageJob parameter names (e.g. Resolution -> resolution, + Images -> images). + """ + payload: dict[str, Any] = { + "prompt": inputs["prompt"], + } + + # Optional parameters (snake_case of upstream SubmitTextToImageJob params) + if inputs.get("resolution"): + payload["resolution"] = inputs["resolution"] + if inputs.get("seed") is not None: + payload["seed"] = int(inputs["seed"]) + if "revise" in inputs: + payload["revise"] = int(inputs["revise"]) + if "logo_add" in inputs: + payload["logo_add"] = int(inputs["logo_add"]) + if inputs.get("logo_param"): + logo_param: dict[str, str] = {} + lp = inputs["logo_param"] + if lp.get("logo_url"): + logo_param["logo_url"] = lp["logo_url"] + if lp.get("logo_image"): + logo_param["logo_image"] = lp["logo_image"] + if logo_param: + payload["logo_param"] = logo_param + + # Reference images — maps to upstream Images.N + # TokenHub accepts URLs or base64 data URIs in the images array + image_refs = inputs.get("images") + if image_refs: + payload["images"] = self._resolve_images(image_refs) + + return payload + + @staticmethod + def _resolve_images(refs: list[str]) -> list[str]: + """Resolve reference images to strings for the TokenHub API. + + Each entry may be: + - An HTTP(S) URL → passed through unchanged + - A data URI (``data:...``) → passed through unchanged + - A local file path → base64-encoded as a data URI + + Per upstream docs: single image 50-5000px per side, base64 < 6MB. + Formats: jpg/jpeg/png/bmp/tiff/webp. + """ + import base64 + + resolved: list[str] = [] + for ref in refs: + if ref.startswith("data:") or ref.startswith("http://") or ref.startswith("https://"): + resolved.append(ref) + continue + + image_path = Path(ref) + if not image_path.is_file(): + raise FileNotFoundError(f"Reference image not found: {ref}") + + raw = image_path.read_bytes() + max_raw = 6 * 1024 * 1024 # 6MB per upstream limit + if len(raw) > max_raw: + raise ValueError( + f"Image too large ({len(raw)} bytes). Max ~6MB raw." + ) + + suffix = image_path.suffix.lower() + mime_map = { + ".jpg": "image/jpeg", + ".jpeg": "image/jpeg", + ".png": "image/png", + ".bmp": "image/bmp", + ".tiff": "image/tiff", + ".tif": "image/tiff", + ".webp": "image/webp", + } + mime = mime_map.get(suffix, "image/png") + data = base64.b64encode(raw).decode("ascii") + resolved.append(f"data:{mime};base64,{data}") + return resolved + + # ------------------------------------------------------------------ + # API communication (TokenHub OpenAI-compatible) + # ------------------------------------------------------------------ + + @staticmethod + def _auth_headers(api_key: str) -> dict[str, str]: + """Build common request headers for TokenHub API calls.""" + return { + "Authorization": f"Bearer {api_key}", + "Content-Type": "application/json", + } + + def _submit_task( + self, payload: dict[str, Any], *, model: str, api_key: str, + ) -> str: + """Submit an image generation task and return the task ID. + + POST /v1/api/image/submit + Request: {"model": "hy-image-v3.0", "prompt": "...", ...} + Response: {"id": "...", "status": "queued", ...} + """ + import requests + + body = { + "model": model, + **payload, + } + url = f"https://{_HOST}{_SUBMIT_PATH}" + resp = requests.post( + url, + json=body, + headers=self._auth_headers(api_key), + timeout=30, + ) + data = self._json_or_raise(resp) + self._check_response(data) + + task_id = data.get("id") + if not task_id: + raise RuntimeError( + f"TokenHub submit returned no task id: {data}" + ) + return task_id + + def _poll_task( + self, + task_id: str, + *, + model: str, + api_key: str, + poll_interval: float, + timeout_seconds: int, + ) -> list[str]: + """Poll /v1/api/image/query until completion, return image download URLs. + + Response when completed: + {"status": "completed", "data": [{"url": "...", "revised_prompt": "..."}]} + + The data array may contain multiple images. Each URL is valid for ~1 hour. + """ + import requests + + url = f"https://{_HOST}{_QUERY_PATH}" + + deadline = time.time() + timeout_seconds + while time.time() < deadline: + time.sleep(poll_interval) + + resp = requests.post( + url, + json={"model": model, "id": task_id}, + headers=self._auth_headers(api_key), + timeout=30, + ) + data = self._json_or_raise(resp) + self._check_response(data) + + status = data.get("status", "") + + if status == "completed": + result_data = data.get("data") or [] + urls = [item.get("url") for item in result_data if item.get("url")] + if not urls: + raise RuntimeError( + f"TokenHub task {task_id} completed but no data[].url: {data}" + ) + return urls + + if status == "failed": + error_info = data.get("error") or {} + error_msg = error_info.get("message", "unknown error") + raise RuntimeError( + f"TokenHub task {task_id} failed: {error_msg}" + ) + + # queued / running / in_progress — continue polling + if status not in ("queued", "running", "in_progress"): + raise RuntimeError( + f"TokenHub task {task_id} returned unknown status: {status}" + ) + + raise TimeoutError( + f"TokenHub task {task_id} did not finish within {timeout_seconds}s" + ) + + # ------------------------------------------------------------------ + # Error handling helpers + # ------------------------------------------------------------------ + + @staticmethod + def _safe_error(exc: Exception) -> str: + """Redact secret values from exception messages.""" + msg = str(exc) + for var in ("TENCENT_TOKENHUB_API_KEY",): + val = os.environ.get(var, "") + if val: + msg = msg.replace(val, "[redacted]") + return msg + + @staticmethod + def _json_or_raise(response: Any) -> dict[str, Any]: + """Parse JSON response body or raise with HTTP status.""" + try: + return response.json() + except ValueError as exc: + raise RuntimeError( + f"Non-JSON response from TokenHub API: HTTP {response.status_code}" + ) from exc + + @staticmethod + def _check_response(payload: dict[str, Any]) -> None: + """Check the TokenHub API response for errors. + + TokenHub returns errors at the top level with an ``error`` field. + """ + error = payload.get("error") + if error: + message = error.get("message", "unknown error") + code = error.get("code", error.get("type", "unknown")) + raise RuntimeError( + f"TokenHub API error: code={code}, message={message}" + ) + + # ------------------------------------------------------------------ + # Output helpers + # ------------------------------------------------------------------ + + @staticmethod + def _resolve_output_paths(base: str, count: int) -> list[Path]: + """Derive distinct paths for ``count`` images. + + Single image keeps the base path unchanged; multiple images insert an + index before the extension (foo.png -> foo_1.png, foo_2.png, ...). + """ + base_path = Path(base) + if count <= 1: + return [base_path] + stem = base_path.stem + suffix = base_path.suffix + parent = base_path.parent + return [parent / f"{stem}_{i}{suffix}" for i in range(1, count + 1)]