llm-connect/llm_connect/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

564 lines
19 KiB
Python

"""Text-to-image rendering with a spend gate and a provenance receipt.
This is a separate surface from :class:`~llm_connect.adapter.LLMAdapter`.
A render takes a prompt and returns image bytes plus a receipt that names
the provider, the model, a digest of the prompt, and the cost. Credentials
stay inside the client that talks to the provider.
The first client speaks OpenRouter's ``POST /api/v1/images`` API. The model
is always chosen by the caller. A per-call or batch cap requires a quote so
the gate can refuse before the provider is contacted.
"""
from __future__ import annotations
import base64
import hashlib
import re
import threading
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Protocol
from llm_connect._http import post_json
from llm_connect.config import find_project_root, resolve_api_key
from llm_connect.exceptions import (
LLMAPIError,
LLMBudgetExceededError,
LLMConfigurationError,
)
from llm_connect.fx import FxRate, resolve_fx_rate
from llm_connect.usage import UsageEvent
_SIZE_RE = re.compile(r"^(\d+)x(\d+)$")
_MIN_EDGE = 1024
_MAX_EDGE = 4096
_MAX_PROMPT_CHARS = 16_000
_MONEY_EPSILON = 1e-9
_MAGIC: dict[str, bytes] = {
"png": b"\x89PNG\r\n\x1a\n",
"jpeg": b"\xff\xd8\xff",
}
_MEDIA_TYPE = {"png": "image/png", "jpeg": "image/jpeg"}
_SUFFIX_FORMAT = {".png": "png", ".jpg": "jpeg", ".jpeg": "jpeg"}
_IMAGE_PROVIDERS: dict[str, dict[str, str]] = {
"openrouter": {
"api_base": "https://openrouter.ai/api/v1",
"env_var": "OPENROUTER_API_KEY",
},
}
class ImageClient(Protocol):
"""Provider call used by :func:`render_image`.
``generate`` returns image bytes and the provider-reported USD cost.
``None`` means the provider omitted a cost.
"""
provider: str
model: str
def generate(self, request: ImageRenderRequest) -> tuple[bytes, float | None]:
"""Call the provider once. Do not write files or enforce spend caps."""
@dataclass(frozen=True)
class ImageRenderRequest:
"""One square text-to-image render.
``size`` is ``WIDTHxHEIGHT`` with equal edges of at least 1024.
``output_format`` is ``png`` or ``jpeg``.
"""
prompt: str
size: str = "1024x1024"
output_format: str = "png"
output_path: Path | None = None
def edge(self) -> int:
"""Return the square edge length after validation."""
return parse_square_size(self.size)
def format(self) -> str:
"""Return ``png`` or ``jpeg`` after validation against ``output_path``."""
return resolve_output_format(self.output_format, self.output_path)
@dataclass(frozen=True)
class ImageReceipt:
"""Provenance for one render. Image bytes are not part of the receipt."""
provider: str
model: str
prompt_digest: str
cost_usd: float | None
cost_eur: float | None
cost_source: str
width: int
height: int
media_type: str
output_format: str
byte_length: int
output_path: str | None
within_cap: bool
fx_source: str | None = None
def to_dict(self) -> dict[str, Any]:
"""JSON-ready receipt. Suitable for the hall to store next to a portrait."""
return {
"provider": self.provider,
"model": self.model,
"prompt_digest": self.prompt_digest,
"cost_usd": self.cost_usd,
"cost_eur": self.cost_eur,
"cost_source": self.cost_source,
"width": self.width,
"height": self.height,
"media_type": self.media_type,
"output_format": self.output_format,
"byte_length": self.byte_length,
"output_path": self.output_path,
"within_cap": self.within_cap,
"fx_source": self.fx_source,
}
@dataclass(frozen=True)
class ImageRenderResult:
"""Image bytes plus the receipt that describes them."""
image: bytes
receipt: ImageReceipt
@dataclass(frozen=True)
class ImageSpendPolicy:
"""Caps applied before a provider call.
When either cap is set, ``quote_usd`` is required. The quote is the
most the gate is willing to reserve for the call. The provider-reported
cost, when present, is what the receipt and the running total record.
"""
max_cost_usd: float | None = None
max_batch_cost_usd: float | None = None
quote_usd: float | None = None
def __post_init__(self) -> None:
for name in ("max_cost_usd", "max_batch_cost_usd", "quote_usd"):
value = getattr(self, name)
if value is not None:
object.__setattr__(self, name, _non_negative_money(name, value))
capped = self.max_cost_usd is not None or self.max_batch_cost_usd is not None
if capped and self.quote_usd is None:
raise LLMConfigurationError(
"quote_usd is required when a per-call or batch cost cap is set",
context={"max_cost_usd": self.max_cost_usd, "max_batch_cost_usd": self.max_batch_cost_usd},
)
if (
self.quote_usd is not None
and self.max_cost_usd is not None
and self.quote_usd > self.max_cost_usd + _MONEY_EPSILON
):
raise LLMBudgetExceededError(
"Image quote exceeds the per-call cost cap",
context={"quote_usd": self.quote_usd, "max_cost_usd": self.max_cost_usd},
)
if (
self.quote_usd is not None
and self.max_batch_cost_usd is not None
and self.quote_usd > self.max_batch_cost_usd + _MONEY_EPSILON
):
raise LLMBudgetExceededError(
"Image quote exceeds the batch cost cap",
context={
"quote_usd": self.quote_usd,
"max_batch_cost_usd": self.max_batch_cost_usd,
},
)
class ImageSpendGate:
"""Reserve a quote before the provider call and commit the actual cost after.
Concurrent callers share one gate. A reservation counts against the batch
cap until :meth:`commit` or :meth:`release`.
"""
def __init__(self, policy: ImageSpendPolicy | None = None) -> None:
self.policy = policy or ImageSpendPolicy()
self.spent_usd = 0.0
self._reserved_usd = 0.0
self._lock = threading.Lock()
def reserve(self) -> None:
"""Hold this call's quote against the batch cap, or refuse first."""
quote = self.policy.quote_usd or 0.0
with self._lock:
projected = self.spent_usd + self._reserved_usd + quote
batch = self.policy.max_batch_cost_usd
if batch is not None and projected > batch + _MONEY_EPSILON:
raise LLMBudgetExceededError(
"Image batch cost cap would be exceeded",
context={
"spent_usd": self.spent_usd,
"reserved_usd": self._reserved_usd,
"quote_usd": self.policy.quote_usd,
"max_batch_cost_usd": batch,
},
)
per_call = self.policy.max_cost_usd
if (
per_call is not None
and self.policy.quote_usd is not None
and self.policy.quote_usd > per_call + _MONEY_EPSILON
):
raise LLMBudgetExceededError(
"Image quote exceeds the per-call cost cap",
context={"quote_usd": self.policy.quote_usd, "max_cost_usd": per_call},
)
self._reserved_usd += quote
def release(self) -> None:
"""Drop a reservation when the provider call fails."""
quote = self.policy.quote_usd or 0.0
with self._lock:
self._reserved_usd = max(0.0, self._reserved_usd - quote)
def commit(self, accounted_usd: float) -> bool:
"""Replace the reservation with the accounted cost.
Returns ``False`` when the accounted cost is over a configured cap.
The cost is still counted, because the provider already charged it.
"""
accounted = _non_negative_money("accounted_usd", accounted_usd)
quote = self.policy.quote_usd or 0.0
with self._lock:
self._reserved_usd = max(0.0, self._reserved_usd - quote)
self.spent_usd += accounted
within = True
per_call = self.policy.max_cost_usd
if per_call is not None and accounted > per_call + _MONEY_EPSILON:
within = False
batch = self.policy.max_batch_cost_usd
if batch is not None and self.spent_usd > batch + _MONEY_EPSILON:
within = False
return within
def prompt_digest(prompt: str) -> str:
"""SHA-256 hex digest of the prompt's UTF-8 bytes."""
return hashlib.sha256(prompt.encode("utf-8")).hexdigest()
def parse_square_size(size: str) -> int:
"""Accept ``1024x1024`` and larger square sizes up to 4096."""
match = _SIZE_RE.fullmatch(size.strip())
if match is None:
raise LLMConfigurationError(
"Image size must look like 1024x1024",
context={"size": size},
)
width = int(match.group(1))
height = int(match.group(2))
if width != height or width < _MIN_EDGE or width > _MAX_EDGE:
raise LLMConfigurationError(
f"Image size must be square and between {_MIN_EDGE} and {_MAX_EDGE} pixels",
context={"size": size},
)
return width
def resolve_output_format(output_format: str, output_path: Path | None) -> str:
"""Return ``png`` or ``jpeg``, and require the path suffix to agree."""
fmt = output_format.strip().lower()
if fmt == "jpg":
fmt = "jpeg"
if fmt not in _MAGIC:
raise LLMConfigurationError(
"Image output_format must be png or jpeg",
context={"output_format": output_format},
)
if output_path is not None:
suffix = output_path.suffix.lower()
expected = _SUFFIX_FORMAT.get(suffix)
if expected is None:
raise LLMConfigurationError(
"Image output path must end in .png, .jpg, or .jpeg",
context={"output_path": str(output_path)},
)
if expected != fmt:
raise LLMConfigurationError(
"Image output path suffix does not match output_format",
context={"output_path": str(output_path), "output_format": fmt},
)
return fmt
def render_image(
client: ImageClient,
request: ImageRenderRequest,
*,
gate: ImageSpendGate | None = None,
fx: FxRate | float | None = None,
force: bool = False,
) -> ImageRenderResult:
"""Render one image and, when requested, write it to ``output_path``.
The provider is not called when the spend gate refuses the quote.
A provider HTTP failure releases the reservation. A successful response
is counted even when the bytes are rejected or the file cannot be written.
"""
_validate_prompt(request.prompt)
edge = request.edge()
fmt = request.format()
if request.output_path is not None and request.output_path.exists() and not force:
raise LLMConfigurationError(
"Image output already exists",
context={"output_path": str(request.output_path)},
)
spend = gate or ImageSpendGate()
rate = resolve_fx_rate(fx)
spend.reserve()
try:
raw, provider_cost = client.generate(request)
except Exception:
spend.release()
raise
cost_usd, cost_source = _accounted_cost(provider_cost, spend.policy.quote_usd)
within_cap = spend.commit(0.0 if cost_usd is None else cost_usd)
if not raw.startswith(_MAGIC[fmt]):
raise LLMAPIError(
f"Provider returned bytes that are not {fmt}",
context={
"provider": client.provider,
"model": client.model,
"output_format": fmt,
"cost_usd": cost_usd,
"cost_source": cost_source,
"within_cap": within_cap,
},
)
output_path: str | None = None
if request.output_path is not None:
write_image_file(request.output_path, raw, force=force)
output_path = str(request.output_path)
cost_eur = None if cost_usd is None or rate is None else rate.usd_to_eur(cost_usd)
receipt = ImageReceipt(
provider=client.provider,
model=client.model,
prompt_digest=prompt_digest(request.prompt),
cost_usd=cost_usd,
cost_eur=cost_eur,
cost_source=cost_source,
width=edge,
height=edge,
media_type=_MEDIA_TYPE[fmt],
output_format=fmt,
byte_length=len(raw),
output_path=output_path,
within_cap=within_cap,
fx_source=None if rate is None else rate.source,
)
return ImageRenderResult(image=raw, receipt=receipt)
def write_image_file(path: Path, data: bytes, *, force: bool = False) -> None:
"""Write image bytes atomically. Refuse to replace an existing file unless forced."""
if path.exists() and not force:
raise LLMConfigurationError(
"Image output already exists",
context={"output_path": str(path)},
)
path.parent.mkdir(parents=True, exist_ok=True)
temporary = path.with_name(path.name + ".partial")
try:
temporary.write_bytes(data)
temporary.replace(path)
except Exception:
temporary.unlink(missing_ok=True)
raise
def usage_event_from_receipt(receipt: ImageReceipt, *, source: str) -> UsageEvent:
"""Ledger row for one image render. The prompt itself is not stored."""
return UsageEvent(
provider=receipt.provider,
model_id=receipt.model,
prompt_tokens=0,
completion_tokens=0,
total_tokens=0,
cost_usd=receipt.cost_usd,
cost_eur=receipt.cost_eur,
cost_source=receipt.cost_source,
source=source,
fx_source=receipt.fx_source,
tags={
"kind": "image",
"prompt_digest": receipt.prompt_digest,
"width": receipt.width,
"height": receipt.height,
"output_format": receipt.output_format,
"within_cap": receipt.within_cap,
},
)
def create_image_client(
provider: str = "openrouter",
model: str | None = None,
api_key: str | None = None,
**kwargs: Any,
) -> ImageClient:
"""Build an image client. ``model`` is required; there is no priced default."""
if provider not in _IMAGE_PROVIDERS:
known = ", ".join(sorted(_IMAGE_PROVIDERS))
raise LLMConfigurationError(
f"Unknown image provider {provider!r}. Choose from: {known}",
context={"provider": provider},
)
if model is None or not str(model).strip():
raise LLMConfigurationError(
"An image model is required",
context={"provider": provider},
)
if provider == "openrouter":
return OpenRouterImageClient(model=model, api_key=api_key, **kwargs)
raise LLMConfigurationError(
f"Image provider {provider!r} is not implemented",
context={"provider": provider},
)
class OpenRouterImageClient:
"""OpenRouter ``POST /images`` client. The API key never appears on a receipt."""
def __init__(
self,
model: str,
api_key: str | None = None,
api_base: str | None = None,
timeout: int = 300,
) -> None:
defaults = _IMAGE_PROVIDERS["openrouter"]
self.provider = "openrouter"
self.model = model.strip()
if not self.model:
raise LLMConfigurationError("An image model is required", context={"provider": "openrouter"})
self._api_base = (api_base or defaults["api_base"]).rstrip("/")
self._timeout = timeout
root = find_project_root()
key_files = [root / "apikey-openrouter.txt"] if root else []
self._api_key = resolve_api_key(
explicit=api_key,
env_var=defaults["env_var"],
key_file_paths=key_files,
)
def generate(self, request: ImageRenderRequest) -> tuple[bytes, float | None]:
"""Generate one image. Returns raw bytes and the USD cost, if reported."""
if not self._api_key:
raise LLMConfigurationError(
"No API key configured for image rendering",
context={"provider": self.provider, "env_var": "OPENROUTER_API_KEY"},
)
fmt = request.format()
edge = request.edge()
payload: dict[str, Any] = {
"model": self.model,
"prompt": request.prompt,
"n": 1,
"aspect_ratio": "1:1",
"size": f"{edge}x{edge}",
"output_format": fmt,
}
data = post_json(
f"{self._api_base}/images",
payload,
headers={"Authorization": f"Bearer {self._api_key}"},
timeout=self._timeout,
)
return _decode_openrouter_image(data), _provider_cost(data)
def _validate_prompt(prompt: str) -> None:
if not isinstance(prompt, str) or not prompt.strip():
raise LLMConfigurationError("Image prompt must be a non-empty string")
if len(prompt) > _MAX_PROMPT_CHARS:
raise LLMConfigurationError(
"Image prompt is too long",
context={"length": len(prompt), "max_chars": _MAX_PROMPT_CHARS},
)
def _accounted_cost(
provider_cost: float | None,
quote_usd: float | None,
) -> tuple[float | None, str]:
if provider_cost is not None:
return _non_negative_money("cost_usd", provider_cost), "provider"
if quote_usd is not None:
return quote_usd, "quote"
return None, "unpriced"
def _provider_cost(data: dict[str, Any]) -> float | None:
usage = data.get("usage")
if not isinstance(usage, dict) or "cost" not in usage or usage["cost"] is None:
return None
return _non_negative_money("cost_usd", usage["cost"])
def _decode_openrouter_image(data: dict[str, Any]) -> bytes:
items = data.get("data")
if not isinstance(items, list) or not items:
raise LLMAPIError("Image response did not include data", context={"provider": "openrouter"})
first = items[0]
if not isinstance(first, dict):
raise LLMAPIError("Image response data[0] was not an object", context={"provider": "openrouter"})
encoded = first.get("b64_json")
if not isinstance(encoded, str) or not encoded.strip():
raise LLMAPIError(
"Image response did not include b64_json",
context={"provider": "openrouter"},
)
payload = encoded.strip()
if payload.startswith("data:") and "," in payload:
payload = payload.split(",", 1)[1]
try:
raw = base64.b64decode(payload, validate=True)
except Exception as exc:
raise LLMAPIError(
"Image response b64_json could not be decoded",
cause=exc,
context={"provider": "openrouter"},
) from exc
if not raw:
raise LLMAPIError("Image response decoded to empty bytes", context={"provider": "openrouter"})
return raw
def _non_negative_money(name: str, value: Any) -> float:
try:
amount = float(value)
except (TypeError, ValueError) as exc:
raise LLMConfigurationError(
f"{name} must be a number",
cause=exc,
context={name: value},
) from exc
if amount < 0 or amount != amount: # NaN
raise LLMConfigurationError(
f"{name} must be a non-negative number",
context={name: value},
)
return amount