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
250 lines
8.8 KiB
Python
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
|