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
564 lines
19 KiB
Python
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
|
|
|
|
|