llm-connect/tests/test_images.py
tegwick e3ab268d6e
Some checks failed
CI Smoke / host-smoke (push) Successful in 0s
CI Smoke / container-smoke (push) Has been cancelled
Add square text-to-image rendering with a spend gate
Hall of Helix needs a credential-safe render that returns a provenance
receipt and refuses a batch before the quote exceeds its cap. OpenRouter
is the only client; the model and the live budget stay with the founder.

Assistant: grok
Assistant-Session: 01a11ab6-1534-7d63-8606-662b92a80ed6
2026-10-08 11:01:49 +02:00

250 lines
8.8 KiB
Python

"""Text-to-image rendering, spend gate, and OpenRouter payload parsing."""
import base64
import json
import pytest
from llm_connect.cli import main
from llm_connect.exceptions import LLMAPIError, LLMBudgetExceededError, LLMConfigurationError
from llm_connect.images import (
ImageRenderRequest,
ImageSpendGate,
ImageSpendPolicy,
OpenRouterImageClient,
create_image_client,
prompt_digest,
render_image,
)
from llm_connect.usage import UsageLedger
_PNG = b"\x89PNG\r\n\x1a\nrest"
_JPEG = b"\xff\xd8\xff\xd9"
class FakeImageClient:
def __init__(self, payload: bytes = _PNG, cost: float | None = 0.02) -> None:
self.provider = "openrouter"
self.model = "test/image"
self.payload = payload
self.cost = cost
self.calls: list[ImageRenderRequest] = []
def generate(self, request: ImageRenderRequest) -> tuple[bytes, float | None]:
self.calls.append(request)
return self.payload, self.cost
def test_square_size_rejects_small_and_non_square() -> None:
client = FakeImageClient()
with pytest.raises(LLMConfigurationError):
render_image(client, ImageRenderRequest(prompt="a gear", size="512x512"))
with pytest.raises(LLMConfigurationError):
render_image(client, ImageRenderRequest(prompt="a gear", size="1024x768"))
assert client.calls == []
def test_render_writes_jpeg_and_receipt_omits_prompt_and_credentials(tmp_path) -> None:
client = FakeImageClient(payload=_JPEG, cost=0.04)
output = tmp_path / "visuals" / "ada.jpg"
result = render_image(
client,
ImageRenderRequest(
prompt="precise technical portrait of a gear",
output_format="jpeg",
output_path=output,
),
gate=ImageSpendGate(ImageSpendPolicy(max_cost_usd=0.10, quote_usd=0.05)),
fx=0.5,
)
assert output.read_bytes() == _JPEG
assert result.receipt.provider == "openrouter"
assert result.receipt.model == "test/image"
assert result.receipt.prompt_digest == prompt_digest("precise technical portrait of a gear")
assert result.receipt.cost_usd == 0.04
assert result.receipt.cost_eur == pytest.approx(0.02)
assert result.receipt.cost_source == "provider"
assert result.receipt.width == 1024
assert result.receipt.height == 1024
assert result.receipt.media_type == "image/jpeg"
assert result.receipt.within_cap is True
encoded = json.dumps(result.receipt.to_dict())
assert "precise technical portrait" not in encoded
assert "api_key" not in encoded
assert "Bearer" not in encoded
def test_batch_cap_stops_before_the_provider_call() -> None:
client = FakeImageClient(cost=0.05)
gate = ImageSpendGate(
ImageSpendPolicy(max_cost_usd=0.05, max_batch_cost_usd=0.10, quote_usd=0.05)
)
request = ImageRenderRequest(prompt="one")
render_image(client, request, gate=gate)
render_image(client, ImageRenderRequest(prompt="two"), gate=gate)
with pytest.raises(LLMBudgetExceededError):
render_image(client, ImageRenderRequest(prompt="three"), gate=gate)
assert len(client.calls) == 2
assert gate.spent_usd == pytest.approx(0.10)
def test_cap_requires_a_quote_and_quote_cannot_exceed_it() -> None:
with pytest.raises(LLMConfigurationError):
ImageSpendPolicy(max_cost_usd=0.10)
with pytest.raises(LLMBudgetExceededError):
ImageSpendPolicy(max_cost_usd=0.10, quote_usd=0.20)
def test_existing_output_is_left_untouched(tmp_path) -> None:
output = tmp_path / "ada.png"
output.write_bytes(b"keep")
client = FakeImageClient()
with pytest.raises(LLMConfigurationError):
render_image(
client,
ImageRenderRequest(prompt="a gear", output_path=output),
)
assert output.read_bytes() == b"keep"
assert client.calls == []
def test_quote_is_used_when_the_provider_omits_cost() -> None:
client = FakeImageClient(cost=None)
result = render_image(
client,
ImageRenderRequest(prompt="a gear"),
gate=ImageSpendGate(ImageSpendPolicy(max_batch_cost_usd=1.0, quote_usd=0.03)),
)
assert result.receipt.cost_usd == 0.03
assert result.receipt.cost_source == "quote"
def test_openrouter_client_posts_square_size_and_decodes_b64(monkeypatch) -> None:
seen: dict = {}
def fake_post(url, payload, headers, timeout=300):
seen["url"] = url
seen["payload"] = payload
seen["headers"] = headers
seen["timeout"] = timeout
encoded = base64.b64encode(_PNG).decode("ascii")
return {"data": [{"b64_json": encoded}], "usage": {"cost": 0.07}}
monkeypatch.setattr("llm_connect.images.post_json", fake_post)
client = OpenRouterImageClient(model="black-forest-labs/flux-3-image", api_key="test-key")
raw, cost = client.generate(
ImageRenderRequest(prompt="a gear", size="1024x1024", output_format="png")
)
assert raw == _PNG
assert cost == 0.07
assert seen["url"] == "https://openrouter.ai/api/v1/images"
assert seen["payload"]["model"] == "black-forest-labs/flux-3-image"
assert seen["payload"]["aspect_ratio"] == "1:1"
assert seen["payload"]["size"] == "1024x1024"
assert seen["payload"]["n"] == 1
assert seen["payload"]["output_format"] == "png"
assert seen["headers"]["Authorization"] == "Bearer test-key"
def test_openrouter_client_rejects_a_response_without_bytes(monkeypatch) -> None:
monkeypatch.setattr(
"llm_connect.images.post_json",
lambda *args, **kwargs: {"data": [{"url": "https://example.invalid/secret.png"}]},
)
client = OpenRouterImageClient(model="test/image", api_key="test-key")
with pytest.raises(LLMAPIError, match="b64_json"):
client.generate(ImageRenderRequest(prompt="a gear"))
def test_unknown_provider_and_missing_model_fail_closed(monkeypatch) -> None:
with pytest.raises(LLMConfigurationError):
create_image_client("openai", model="gpt-image-1")
with pytest.raises(LLMConfigurationError):
create_image_client("openrouter", model=None)
monkeypatch.setattr("llm_connect.images.resolve_api_key", lambda **kwargs: None)
with pytest.raises(LLMConfigurationError, match="No API key"):
OpenRouterImageClient(model="test/image").generate(ImageRenderRequest(prompt="a gear"))
def test_cli_render_writes_receipt_and_ledger(tmp_path, monkeypatch, capsys) -> None:
client = FakeImageClient(payload=_JPEG, cost=0.04)
monkeypatch.setattr("llm_connect.cli.create_image_client", lambda *args, **kwargs: client)
output = tmp_path / "ada.jpg"
ledger = tmp_path / "usage.jsonl"
code = main(
[
"image",
"render",
"--prompt",
"a gear",
"--model",
"test/image",
"--output",
str(output),
"--max-cost-usd",
"0.10",
"--quote-usd",
"0.05",
"--ledger",
str(ledger),
"--json",
"--eur-per-usd",
"0.5",
]
)
assert code == 0
assert output.read_bytes() == _JPEG
receipt = json.loads(capsys.readouterr().out)
assert receipt["model"] == "test/image"
assert receipt["prompt_digest"] == prompt_digest("a gear")
assert receipt["cost_usd"] == 0.04
assert receipt["output_path"] == str(output)
events = UsageLedger(ledger).read_all()
assert len(events) == 1
assert events[0].tags["kind"] == "image"
assert events[0].tags["prompt_digest"] == receipt["prompt_digest"]
assert "a gear" not in json.dumps(events[0].to_dict())
def test_cli_batch_stops_when_the_next_quote_does_not_fit(tmp_path, monkeypatch, capsys) -> None:
client = FakeImageClient(payload=_JPEG, cost=0.05)
monkeypatch.setattr("llm_connect.cli.create_image_client", lambda *args, **kwargs: client)
batch = tmp_path / "batch.jsonl"
first = tmp_path / "one.jpg"
second = tmp_path / "two.jpg"
third = tmp_path / "three.jpg"
lines = [
{"prompt": "one", "output_path": str(first)},
{"prompt": "two", "output_path": str(second)},
{"prompt": "three", "output_path": str(third)},
]
batch.write_text("\n".join(json.dumps(line) for line in lines), encoding="utf-8")
code = main(
[
"image",
"render",
"--batch",
str(batch),
"--model",
"test/image",
"--max-cost-usd",
"0.05",
"--max-batch-cost-usd",
"0.10",
"--quote-usd",
"0.05",
"--no-ledger",
]
)
captured = capsys.readouterr()
assert code == 2
assert first.is_file()
assert second.is_file()
assert not third.exists()
receipts = json.loads(captured.out)
assert [item["output_path"] for item in receipts] == [str(first), str(second)]
assert "batch cost cap" in captured.err