test(tencent): cover Seedream image generation, edits, and pricing through TokenHub

tests/unit/llms/tencent/image_edit/test_tencent_image_edit_transformation.py

@@ -0,0 +1,91 @@
+import base64
+import io
+import json
+
+import httpx
+import pytest
+
+import litellm
+from litellm.llms.custom_httpx.http_handler import HTTPHandler
+
+PNG_BYTES = b"\x89PNG\r\n\x1a\n" + b"\x00" * 32
+TOKENHUB_RESPONSE = {
+    "created": 1792551838,
+    "data": [{"output_format": "jpeg", "size": "2048x2048", "url": "https://cos.example/edited.jpg"}],
+    "model": "doubao-seedream-5-0-lite",
+    "tokenhub_usage": {"total_tokens": 21875},
+    "request_id": "4f795d48-b227-4c35-af2c-1a0ba288b6b6",
+}
+
+
+def _client(captured: dict) -> HTTPHandler:
+    def respond(request: httpx.Request) -> httpx.Response:
+        captured["url"] = str(request.url)
+        captured["content_type"] = request.headers.get("content-type")
+        captured["authorization"] = request.headers.get("authorization")
+        captured["body"] = json.loads(request.content)
+        return httpx.Response(200, json=TOKENHUB_RESPONSE)
+
+    return HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(respond)))
+
+
+def test_edit_sends_uploaded_images_as_data_urls_in_a_json_body():
+    captured = {}
+    png_data_url = f"data:image/png;base64,{base64.b64encode(PNG_BYTES).decode()}"
+
+    response = litellm.image_edit(
+        model="tencent/seedream-image-v5.0-lite",
+        image=[PNG_BYTES, io.BytesIO(PNG_BYTES), "https://ref.example/style.png"],
+        prompt="turn the square into a paper lantern",
+        size="2K",
+        api_key="sk-tokenhub",
+        client=_client(captured),
+    )
+
+    assert captured["url"] == "https://tokenhub-intl.tencentcloudmaas.com/v1/wand/si-image/generation"
+    assert captured["content_type"] == "application/json"
+    assert captured["authorization"] == "Bearer sk-tokenhub"
+    assert captured["body"] == {
+        "model": "seedream-image-v5.0-lite",
+        "prompt": "turn the square into a paper lantern",
+        "images": [png_data_url, png_data_url, "https://ref.example/style.png"],
+        "size": "2K",
+    }
+    assert response.data[0].url == "https://cos.example/edited.jpg"
+    assert response.usage.total_tokens == 21875
+
+
+def test_edit_rejects_a_mask_without_drop_params():
+    captured = {}
+
+    with pytest.raises(litellm.UnsupportedParamsError):
+        litellm.image_edit(
+            model="tencent/seedream-image-v5.0-pro",
+            image=PNG_BYTES,
+            mask=PNG_BYTES,
+            prompt="turn the square into a paper lantern",
+            api_key="sk-tokenhub",
+            client=_client(captured),
+        )
+
+    assert captured == {}
+
+
+def test_edit_drops_a_mask_under_drop_params():
+    captured = {}
+
+    litellm.image_edit(
+        model="tencent/seedream-image-v5.0-pro",
+        image=PNG_BYTES,
+        mask=PNG_BYTES,
+        prompt="turn the square into a paper lantern",
+        api_key="sk-tokenhub",
+        client=_client(captured),
+        drop_params=True,
+    )
+
+    assert captured["body"] == {
+        "model": "seedream-image-v5.0-pro",
+        "prompt": "turn the square into a paper lantern",
+        "images": [f"data:image/png;base64,{base64.b64encode(PNG_BYTES).decode()}"],
+    }

tests/unit/llms/tencent/image_generation/test_tencent_image_generation_transformation.py

@@ -0,0 +1,180 @@
+import json
+
+import httpx
+import pytest
+
+import litellm
+from litellm.llms.custom_httpx.http_handler import HTTPHandler
+
+TOKENHUB_RESPONSE = {
+    "created": 1792551825,
+    "data": [{"output_format": "jpeg", "size": "1248x832", "url": "https://cos.example/lantern.jpg"}],
+    "model": "doubao-seedream-5-0-pro-260628",
+    "tokenhub_usage": {"total_tokens": 28125},
+    "request_id": "c57b0374-b5fe-47db-b4b2-751315f920aa",
+}
+
+
+def _client(captured: dict, status_code: int = 200, body: dict | None = None) -> HTTPHandler:
+    def respond(request: httpx.Request) -> httpx.Response:
+        captured["url"] = str(request.url)
+        captured["authorization"] = request.headers.get("authorization")
+        captured["body"] = json.loads(request.content)
+        return httpx.Response(status_code, json=TOKENHUB_RESPONSE if body is None else body)
+
+    return HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(respond)))
+
+
+def test_generation_posts_openai_and_tokenhub_params_to_the_wand_route():
+    captured = {}
+
+    response = litellm.image_generation(
+        model="tencent/seedream-image-v5.0-pro",
+        prompt="a red paper lantern",
+        size="2K",
+        response_format="url",
+        api_key="sk-tokenhub",
+        client=_client(captured),
+        images=["https://ref.example/lantern.png"],
+        watermark=False,
+    )
+
+    assert captured["url"] == "https://tokenhub-intl.tencentcloudmaas.com/v1/wand/si-image/generation"
+    assert captured["authorization"] == "Bearer sk-tokenhub"
+    assert captured["body"] == {
+        "model": "seedream-image-v5.0-pro",
+        "prompt": "a red paper lantern",
+        "size": "2K",
+        "response_format": "url",
+        "images": ["https://ref.example/lantern.png"],
+        "watermark": False,
+    }
+    assert response.data[0].url == "https://cos.example/lantern.jpg"
+    assert response.created == 1792551825
+    assert response.usage.total_tokens == 28125
+    assert response.usage.output_tokens == 28125
+
+
+def test_generation_returns_base64_images():
+    captured = {}
+    body = {**TOKENHUB_RESPONSE, "data": [{"output_format": "png", "size": "1024x1024", "b64_json": "aW1n"}]}
+
+    response = litellm.image_generation(
+        model="tencent/seedream-image-v5.0-flash",
+        prompt="a red paper lantern",
+        response_format="b64_json",
+        api_key="sk-tokenhub",
+        client=_client(captured, body=body),
+    )
+
+    assert response.data[0].b64_json == "aW1n"
+    assert response.data[0].url is None
+
+
+@pytest.mark.parametrize(
+    ("api_base", "expected_url"),
+    [
+        ("https://tokenhub.example", "https://tokenhub.example/v1/wand/si-image/generation"),
+        ("https://tokenhub.example/v1", "https://tokenhub.example/v1/wand/si-image/generation"),
+        ("https://tokenhub.example/v1/", "https://tokenhub.example/v1/wand/si-image/generation"),
+    ],
+)
+def test_generation_resolves_a_configured_api_base_to_the_wand_route(api_base, expected_url):
+    captured = {}
+
+    litellm.image_generation(
+        model="tencent/seedream-image-v5.0-lite",
+        prompt="a red paper lantern",
+        api_key="sk-tokenhub",
+        api_base=api_base,
+        client=_client(captured),
+    )
+
+    assert captured["url"] == expected_url
+
+
+@pytest.mark.parametrize(
+    ("params", "expected_body_params"),
+    [
+        ({"n": 1}, {}),
+        ({"n": 3, "drop_params": True}, {}),
+        ({"background": "auto"}, {}),
+        ({"background": "transparent", "output_format": "png"}, {"background": "transparent", "output_format": "png"}),
+    ],
+)
+def test_generation_sends_only_values_tokenhub_accepts(params, expected_body_params):
+    captured = {}
+
+    litellm.image_generation(
+        model="tencent/seedream-image-v5.0-pro",
+        prompt="a red paper lantern",
+        api_key="sk-tokenhub",
+        client=_client(captured),
+        **params,
+    )
+
+    assert captured["body"] == {
+        "model": "seedream-image-v5.0-pro",
+        "prompt": "a red paper lantern",
+        **expected_body_params,
+    }
+
+
+def test_generation_rejects_more_than_one_image_without_drop_params():
+    captured = {}
+
+    with pytest.raises(litellm.UnsupportedParamsError) as exc_info:
+        litellm.image_generation(
+            model="tencent/seedream-image-v5.0-pro",
+            prompt="a red paper lantern",
+            n=2,
+            api_key="sk-tokenhub",
+            client=_client(captured),
+        )
+
+    assert exc_info.value.status_code == 400
+    assert captured == {}
+
+
+@pytest.mark.parametrize(
+    ("error", "expected_message"),
+    [
+        (
+            {
+                "error": {
+                    "message": "The model or service ID 'seedream-image-v5.0-pro' does not exist.",
+                    "code": "400004",
+                    "type": "gateway_error",
+                }
+            },
+            "does not exist",
+        ),
+        ({"Type": 2, "Code": 1001, "Msg": "Invalid param", "Desc": ""}, "Invalid param"),
+    ],
+)
+def test_generation_surfaces_the_tokenhub_error_message(error, expected_message):
+    with pytest.raises(litellm.BadRequestError) as exc_info:
+        litellm.image_generation(
+            model="tencent/seedream-image-v5.0-pro",
+            prompt="a red paper lantern",
+            api_key="sk-tokenhub",
+            client=_client({}, status_code=400, body=error),
+        )
+
+    assert expected_message in str(exc_info.value)
+
+
+def test_generation_without_a_tokenhub_key_is_an_authentication_error(monkeypatch):
+    monkeypatch.delenv("TENCENT_API_KEY", raising=False)
+    captured = {}
+
+    with pytest.raises(litellm.AuthenticationError) as exc_info:
+        litellm.image_generation(
+            model="tencent/seedream-image-v5.0-pro",
+            prompt="a red paper lantern",
+            client=_client(captured),
+        )
+
+    assert exc_info.value.status_code == 401
+    assert "TENCENT_API_KEY" in str(exc_info.value)
+    assert captured == {}

tests/unit/llms/tencent/test_tencent_cost_calculator.py

@@ -0,0 +1,52 @@
+import pytest
+
+import litellm
+from litellm.types.utils import ImageObject, ImageResponse, ImageUsage, ImageUsageInputTokensDetails
+
+SEEDREAM_MODELS = (
+    "tencent/seedream-image-v5.0-pro",
+    "tencent/seedream-image-v5.0-flash",
+    "tencent/seedream-image-v5.0-lite",
+)
+
+
+def _image_response(total_tokens: int, images: int = 1) -> ImageResponse:
+    return ImageResponse(
+        created=1792551825,
+        data=[ImageObject(url=f"https://cos.example/{index}.jpg") for index in range(images)],
+        usage=ImageUsage(
+            input_tokens=0,
+            input_tokens_details=ImageUsageInputTokensDetails(image_tokens=0, text_tokens=0),
+            output_tokens=total_tokens,
+            total_tokens=total_tokens,
+        ),
+    )
+
+
+@pytest.mark.parametrize("model", SEEDREAM_MODELS)
+@pytest.mark.parametrize("call_type", ["image_generation", "image_edit"])
+def test_seedream_bills_tokenhub_usage_tokens_for_a_tier_size(local_model_cost_map, model, call_type):
+    token_rate = litellm.model_cost[model]["output_cost_per_image_token"]
+
+    cost = litellm.completion_cost(
+        completion_response=_image_response(total_tokens=56250),
+        model=model,
+        call_type=call_type,
+        size="2K",
+    )
+
+    assert cost == pytest.approx(56250 * token_rate)
+
+
+@pytest.mark.parametrize("model", SEEDREAM_MODELS)
+def test_seedream_falls_back_to_the_per_image_price_without_usage(local_model_cost_map, model):
+    per_image_price = litellm.model_cost[model]["output_cost_per_image"]
+
+    cost = litellm.completion_cost(
+        completion_response=_image_response(total_tokens=0, images=2),
+        model=model,
+        call_type="image_generation",
+        size="2K",
+    )
+
+    assert cost == pytest.approx(2 * per_image_price)