diff --git a/migrations/versions/c8a1d2e3f4b5_add_image_pricing_to_models.py b/migrations/versions/c8a1d2e3f4b5_add_image_pricing_to_models.py new file mode 100644 index 00000000..193e4968 --- /dev/null +++ b/migrations/versions/c8a1d2e3f4b5_add_image_pricing_to_models.py @@ -0,0 +1,31 @@ +"""add image_pricing to models + +Revision ID: c8a1d2e3f4b5 +Revises: e4c7a1b9d520 +Create Date: 2026-09-24 00:00:00.000000 +""" + +import sqlalchemy as sa +import sqlmodel +from alembic import op + +# revision identifiers, used by Alembic. +revision = "c8a1d2e3f4b5" +down_revision = "e4c7a1b9d520" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + op.add_column( + "models", + sa.Column( + "image_pricing", + sqlmodel.sql.sqltypes.AutoString(), + nullable=True, + ), + ) + + +def downgrade() -> None: + op.drop_column("models", "image_pricing") diff --git a/routstr/auth.py b/routstr/auth.py index bf29e86e..46282106 100644 --- a/routstr/auth.py +++ b/routstr/auth.py @@ -1112,6 +1112,7 @@ async def adjust_payment_for_tokens( model_obj: "Model | None" = None, provider_fee: float | None = None, reservation_snapshot: ReservationSnapshot | None = None, + precomputed_cost: CostData | None = None, ) -> dict: """ Adjusts the payment based on token usage in the response. @@ -1188,8 +1189,14 @@ async def adjust_payment_for_tokens( extra={"error": str(e), "fee_msats": fee_msats}, ) - calculated_cost = await calculate_cost( - response_data, deducted_max_cost, model_obj, provider_fee + # A flat per-unit charge is priced by the caller; its response carries no + # usage object for ``calculate_cost`` to read. + calculated_cost: CostData | CostDataError = ( + precomputed_cost + if precomputed_cost is not None + else await calculate_cost( + response_data, deducted_max_cost, model_obj, provider_fee + ) ) if isinstance(calculated_cost, CostDataError): # Content was already served, so release instead of raising a 400. diff --git a/routstr/core/admin.py b/routstr/core/admin.py index 365e2f68..e37b2e58 100644 --- a/routstr/core/admin.py +++ b/routstr/core/admin.py @@ -474,6 +474,7 @@ class ModelCreate(BaseModel): context_length: int architecture: dict[str, object] pricing: dict[str, object] + image_pricing: dict[str, object] | None = None per_request_limits: dict[str, object] | None = None top_provider: dict[str, object] | None = None upstream_provider_id: int | None = None @@ -556,6 +557,9 @@ async def upsert_provider_model( existing_row.top_provider = ( json.dumps(payload.top_provider) if payload.top_provider else None ) + existing_row.image_pricing = ( + json.dumps(payload.image_pricing) if payload.image_pricing else None + ) existing_row.canonical_slug = payload.canonical_slug existing_row.alias_ids = ( json.dumps(payload.alias_ids) if payload.alias_ids else None @@ -583,6 +587,9 @@ async def upsert_provider_model( architecture=json.dumps(payload.architecture), pricing=json.dumps(payload.pricing), sats_pricing=None, + image_pricing=( + json.dumps(payload.image_pricing) if payload.image_pricing else None + ), per_request_limits=( json.dumps(payload.per_request_limits) if payload.per_request_limits is not None @@ -724,6 +731,11 @@ async def batch_override_provider_models( existing_row.architecture = json.dumps(model_data.architecture) existing_row.pricing = json.dumps(model_data.pricing) existing_row.sats_pricing = None + existing_row.image_pricing = ( + json.dumps(model_data.image_pricing) + if model_data.image_pricing + else None + ) existing_row.per_request_limits = ( json.dumps(model_data.per_request_limits) if model_data.per_request_limits is not None @@ -755,6 +767,11 @@ async def batch_override_provider_models( architecture=json.dumps(model_data.architecture), pricing=json.dumps(model_data.pricing), sats_pricing=None, + image_pricing=( + json.dumps(model_data.image_pricing) + if model_data.image_pricing + else None + ), per_request_limits=( json.dumps(model_data.per_request_limits) if model_data.per_request_limits is not None diff --git a/routstr/core/db.py b/routstr/core/db.py index 78019c32..062f4168 100644 --- a/routstr/core/db.py +++ b/routstr/core/db.py @@ -420,6 +420,9 @@ class ModelRow(SQLModel, table=True): # type: ignore architecture: str = Field() pricing: str = Field() sats_pricing: str | None = Field(default=None) + image_pricing: str | None = Field( + default=None, description="JSON per-image price book, for image models" + ) per_request_limits: str | None = Field(default=None) top_provider: str | None = Field(default=None) canonical_slug: str | None = Field(default=None, description="Canonical model slug") diff --git a/routstr/payment/cost_calculation.py b/routstr/payment/cost_calculation.py index b4b0ff6b..a6f0ed89 100644 --- a/routstr/payment/cost_calculation.py +++ b/routstr/payment/cost_calculation.py @@ -17,6 +17,7 @@ __all__ = [ "CostDataError", "MaxCostData", "calculate_cost", + "calculate_flat_cost", "parse_token_count", ] @@ -71,6 +72,30 @@ def _empty_cost(cls: type[CostData] = CostData) -> CostData: ) +def calculate_flat_cost(units: int, sats_per_unit: float) -> CostData: + """Price a request billed per produced unit instead of per token. + + Image generation is the case this exists for: no usage object, so the + token path would release the reservation and serve the request free. + """ + if units <= 0 or not is_usable_rate(sats_per_unit) or sats_per_unit <= 0: + return _empty_cost() + + total_sats = units * sats_per_unit + total_msats = math.ceil(total_sats * 1000) + try: + total_usd = total_sats / sats_usd_price() + except Exception: + total_usd = 0.0 + return CostData( + base_msats=total_msats, + input_msats=0, + output_msats=total_msats, + total_msats=total_msats, + total_usd=total_usd, + ) + + async def calculate_cost( response_data: dict, max_cost: int, diff --git a/routstr/payment/helpers.py b/routstr/payment/helpers.py index d088151f..2b6b2d50 100644 --- a/routstr/payment/helpers.py +++ b/routstr/payment/helpers.py @@ -24,6 +24,7 @@ from ..wallet import ( deserialize_token_from_string, is_trusted_source_mint, ) +from .image_pricing import image_reservation_msats from .responses_input import ( FILE_ID_URL_PREFIX, count_input_images, @@ -200,6 +201,10 @@ async def calculate_discounted_max_cost( if not model_pricing: return max_cost_for_model + image_reservation = image_reservation_msats(body, model_obj) + if image_reservation is not None: + return max(settings.min_request_msat, image_reservation) + tol = settings.tolerance_percentage tol_factor = max(0.0, 1 - float(tol) / 100.0) diff --git a/routstr/payment/image_pricing.py b/routstr/payment/image_pricing.py new file mode 100644 index 00000000..a4737d1d --- /dev/null +++ b/routstr/payment/image_pricing.py @@ -0,0 +1,208 @@ +"""Per-image pricing for models that return images instead of tokens. + +The same model costs a different amount per resolution, and the newer ones per +quality step within a resolution, so the price is only known once the request +is read. + +``Pricing.image`` carries only the ceiling so it stays a flat numeric record +the provider-fee and sats conversions can map over. Tier prices here are raw +upstream USD and are applied as a ratio against that ceiling, which is why +neither conversion is repeated below. +""" + +from __future__ import annotations + +import math +from typing import TYPE_CHECKING + +from pydantic.v1 import BaseModel + +if TYPE_CHECKING: + from .models import Model + +__all__ = [ + "MAX_RESERVED_IMAGES", + "ImagePriceTier", + "ImagePricing", + "image_reservation_msats", + "per_image_sats", + "produces_images", + "select_image_price_usd", +] + +# 1K is nominally 1024x1024, 2K is 2048x2048. A pixel size maps onto a class by +# area, bounded at twice the nominal area so a class's portrait and landscape +# variants (1024x1536) stay in it. A size exactly on a boundary takes the +# dearer class rather than under-billing. +_RESOLUTION_AREAS: tuple[tuple[int, str], ...] = ( + (2 * 1024 * 1024, "1K"), + (2 * 2048 * 2048, "2K"), +) +_LARGEST_RESOLUTION_LABEL = "4K" + +# Bounds what a single reservation can hold. A larger batch is still billed per +# image returned. +MAX_RESERVED_IMAGES = 10 + + +class ImagePriceTier(BaseModel): + """USD for one image at a resolution and quality the upstream names itself. + + Either axis may be absent when the upstream does not price along it. + """ + + resolution: str | None = None + quality: str | None = None + usd: float + + class Config: + extra = "ignore" + + +class ImagePricing(BaseModel): + """A model's per-image price book, in raw upstream USD.""" + + max_usd: float + tiers: list[ImagePriceTier] = [] + default_resolution: str | None = None + default_quality: str | None = None + resolutions: list[str] = [] + qualities: list[str] = [] + # Upscale factor to USD, for upstreams billing an upscale as its own call. + upscale: dict[str, float] = {} + + class Config: + extra = "ignore" + + def price_usd( + self, resolution: str | None = None, quality: str | None = None + ) -> float: + """Price for one image at the requested tier. + + A request naming neither axis is priced at the upstream's defaults, + which is what it will be charged; falling back to ``max_usd`` would + bill a default 1K request at the 4K rate. An unmatched tier does fall + back to ``max_usd``, under-billing being the worse error. + """ + wanted_resolution = resolution or self.default_resolution + wanted_quality = quality or self.default_quality + + for candidate in ( + (wanted_resolution, wanted_quality), + (wanted_resolution, None), + ): + for tier in self.tiers: + if tier.resolution == candidate[0] and tier.quality == candidate[1]: + return tier.usd + + return self.max_usd + + def upscale_usd(self, factor: str | None) -> float: + """Price for one upscale at ``factor``, or the dearest one offered.""" + if factor and factor in self.upscale: + return self.upscale[factor] + return max(self.upscale.values()) if self.upscale else self.max_usd + + +def _resolution_from_size(body: dict) -> str | None: + """Resolution label for an OpenAI-style ``size`` or explicit dimensions.""" + edges: list[int] = [] + size = body.get("size") + if isinstance(size, str) and "x" in size.lower(): + for part in size.lower().split("x"): + try: + edges.append(int(part.strip())) + except ValueError: + return None + else: + for dimension in ("width", "height"): + value = body.get(dimension) + if isinstance(value, int) and not isinstance(value, bool): + edges.append(value) + if not edges: + return None + + area = 1 + for edge in edges: + area *= edge + for limit, label in _RESOLUTION_AREAS: + if area < limit: + return label + return _LARGEST_RESOLUTION_LABEL + + +def select_image_price_usd(image_pricing: ImagePricing, body: dict) -> float: + """USD for one image, priced at the tier this request asks for. + + Reads the upstream's own ``resolution``/``quality`` first, then falls back + to the OpenAI-compatible ``size`` (or explicit ``width``/``height``). + """ + resolution = body.get("resolution") + label = ( + resolution.upper() + if isinstance(resolution, str) and resolution + else _resolution_from_size(body) + ) + quality = body.get("quality") + quality_label = quality.lower() if isinstance(quality, str) and quality else None + + # A tier the model never declared is priced at the ceiling: reading it as + # unspecified would bill it at the default tier, the cheaper one on every + # model that has tiers at all. + declared_resolutions = {r.upper() for r in image_pricing.resolutions} + if label is not None and declared_resolutions and label not in declared_resolutions: + return image_pricing.max_usd + declared_qualities = {q.lower() for q in image_pricing.qualities} + if ( + quality_label is not None + and declared_qualities + and quality_label not in declared_qualities + ): + return image_pricing.max_usd + + return image_pricing.price_usd(label, quality_label) + + +def per_image_sats(model: "Model | None", body: dict) -> float: + """Sats for one image from this model, at the tier ``body`` asks for. + + ``sats_pricing.image`` is the ceiling, already carrying the provider fee + and sats conversion, so a price book only scales it by the ratio of the + selected tier to that ceiling. + """ + if model is None or model.sats_pricing is None: + return 0.0 + + ceiling_sats = model.sats_pricing.image + book = model.image_pricing + if book is None or book.max_usd <= 0 or ceiling_sats <= 0: + return ceiling_sats + + return ceiling_sats * select_image_price_usd(book, body) / book.max_usd + + +def produces_images(model: "Model | None") -> bool: + architecture = getattr(model, "architecture", None) + return getattr(architecture, "output_modalities", None) == ["image"] + + +def image_reservation_msats(body: dict, model: "Model | None") -> int | None: + """Msats to hold for an image request, or ``None`` if not an image model. + + Token-window math means nothing for a model that returns images, so the + hold is the requested tier times the requested batch size. + """ + if not produces_images(model): + return None + + sats_per_image = per_image_sats(model, body) + if sats_per_image <= 0: + return None + + try: + count = int(body.get("n", 1)) + except (TypeError, ValueError): + count = 1 + count = min(max(count, 1), MAX_RESERVED_IMAGES) + + return math.ceil(count * sats_per_image * 1000) diff --git a/routstr/payment/models.py b/routstr/payment/models.py index 8d7788a1..591d224f 100644 --- a/routstr/payment/models.py +++ b/routstr/payment/models.py @@ -11,6 +11,7 @@ from sqlmodel.ext.asyncio.session import AsyncSession from ..core.db import ModelRow, UpstreamProviderRow, get_session from ..core.logging import get_logger from ..core.settings import settings +from .image_pricing import ImagePricing from .price import sats_usd_price from .rates import BILLABLE_PRICING_FIELDS, coerce_rate, is_usable_rate @@ -126,6 +127,7 @@ class Model(BaseModel): alias_ids: list[str] | None = None forwarded_model_id: str | None = None reasoning: Reasoning | None = None + image_pricing: ImagePricing | None = None class Config: extra = "ignore" @@ -345,6 +347,8 @@ def _build_model_from_row( json.loads(row.per_request_limits) if row.per_request_limits else None ) top_provider_dict = json.loads(row.top_provider) if row.top_provider else None + raw_image_pricing = getattr(row, "image_pricing", None) + image_pricing_dict = json.loads(raw_image_pricing) if raw_image_pricing else None if isinstance(pricing, dict) and float(pricing.get("request", 0.0)) <= 0.0: pricing["request"] = max(pricing.get("request", 0.0), 0.0) @@ -387,6 +391,9 @@ def _build_model_from_row( canonical_slug=getattr(row, "canonical_slug", None), alias_ids=json.loads(row.alias_ids) if row.alias_ids else None, forwarded_model_id=getattr(row, "forwarded_model_id", None), + image_pricing=ImagePricing.parse_obj(image_pricing_dict) + if image_pricing_dict + else None, ) if apply_provider_fee: diff --git a/routstr/upstream/__init__.py b/routstr/upstream/__init__.py index edac0020..85d094e8 100644 --- a/routstr/upstream/__init__.py +++ b/routstr/upstream/__init__.py @@ -13,6 +13,7 @@ from .ppqai import PPQAIUpstreamProvider from .routstr import RoutstrUpstreamProvider from .tinfoil import TinfoilUpstreamProvider from .typesafe import TypeSafeUpstreamProvider +from .venice import VeniceUpstreamProvider from .xai import XAIUpstreamProvider upstream_provider_classes: list[type[BaseUpstreamProvider]] = [ @@ -30,6 +31,7 @@ upstream_provider_classes: list[type[BaseUpstreamProvider]] = [ RoutstrUpstreamProvider, TinfoilUpstreamProvider, TypeSafeUpstreamProvider, + VeniceUpstreamProvider, XAIUpstreamProvider, ] """List of all upstream classes""" diff --git a/routstr/upstream/base.py b/routstr/upstream/base.py index f270442d..adffff28 100644 --- a/routstr/upstream/base.py +++ b/routstr/upstream/base.py @@ -38,8 +38,10 @@ from ..payment.cost_calculation import ( CostDataError, MaxCostData, calculate_cost, + calculate_flat_cost, ) from ..payment.helpers import create_error_response +from ..payment.image_pricing import per_image_sats from ..payment.models import ( Model, Pricing, @@ -62,6 +64,11 @@ from .cache_breakpoints import ( is_explicit_cache_model, ) from .count_tokens import MissingUsageEstimator, count_tokens_locally +from .image_generation import ( + count_generated_images, + is_image_generation_path, + parse_json_body, +) from .litellm_routing import detect_litellm_prefix from .model_paths import public_provider_url from .rate_limit import UPSTREAM_RATE_LIMIT, classify_rate_limit @@ -2051,6 +2058,74 @@ class BaseUpstreamProvider: ) raise + async def handle_image_generation( + self, + response: httpx.Response, + key: ApiKey, + session: AsyncSession, + max_cost_for_model: int, + model_obj: Model | None, + reservation_snapshot: ReservationSnapshot | None = None, + request_body: bytes | None = None, + ) -> Response: + """Settle an image response: one flat charge per image returned.""" + content = await response.aread() + content_type = response.headers.get("content-type") + image_count = count_generated_images( + content, _is_json_content_type(content_type) if content_type else True + ) + sats_per_image = per_image_sats(model_obj, parse_json_body(request_body)) + model_id = model_obj.id if model_obj else "unknown" + + if image_count > 0 and sats_per_image <= 0: + logger.warning( + "Image response carries no per-image price; releasing the " + "reservation instead of billing a rate we do not have", + extra={ + "model": model_id, + "image_count": image_count, + "key_hash": key.hashed_key[:8] + "...", + }, + ) + + cost_data = await adjust_payment_for_tokens( + key, + {"model": model_id, "usage": None}, + session, + max_cost_for_model, + model_obj, + self.provider_fee, + reservation_snapshot, + precomputed_cost=calculate_flat_cost(image_count, sats_per_image), + ) + + logger.info( + "Settled image generation request", + extra={ + "model": model_id, + "image_count": image_count, + "sats_per_image": sats_per_image, + "key_hash": key.hashed_key[:8] + "...", + }, + ) + + # httpx already decoded the body, so the upstream's framing headers no + # longer describe it. + headers = { + name: value + for name, value in response.headers.items() + if name.lower() + not in {"content-length", "content-encoding", "transfer-encoding"} + } + _inject_cost_response_headers(headers, cost_data) + + return Response( + content=content, + status_code=response.status_code, + headers=headers, + media_type=content_type, + ) + async def _finalize_generic_streaming_payment( self, key_hash: str, @@ -3304,6 +3379,21 @@ class BaseUpstreamProvider: if reservation_snapshot is None: reservation_snapshot = await get_reservation_snapshot(key, session) + if is_image_generation_path(path) and response.status_code == 200: + try: + return await self.handle_image_generation( + response, + key, + session, + max_cost_for_model, + model_obj, + reservation_snapshot=reservation_snapshot, + request_body=request_body, + ) + finally: + await response.aclose() + await client.aclose() + background_tasks = BackgroundTasks() background_tasks.add_task(response.aclose) background_tasks.add_task(client.aclose) diff --git a/routstr/upstream/image_generation.py b/routstr/upstream/image_generation.py new file mode 100644 index 00000000..1b444a74 --- /dev/null +++ b/routstr/upstream/image_generation.py @@ -0,0 +1,67 @@ +"""Recognising and reading image-generation responses. + +These endpoints answer with images and no usage object, so the settlement path +that reads tokens has nothing to work from. What it needs instead is here: is +this an image endpoint, and how many images came back. Pricing those images is +``payment.image_pricing``. + +Nothing here imports the provider, so the provider can import this. +""" + +from __future__ import annotations + +import json + +__all__ = [ + "count_generated_images", + "is_image_generation_path", + "parse_json_body", +] + +# OpenAI's images API plus the native image routes providers expose beside it. +_IMAGE_GENERATION_SUFFIXES = ( + "/images/generations", + "/images/edits", + "/images/variations", + "/image/generate", + "/image/edit", + "/image/inpaint", + "/image/upscale", +) + + +def is_image_generation_path(path: str) -> bool: + return ("/" + path.strip("/")).endswith(_IMAGE_GENERATION_SUFFIXES) + + +def parse_json_body(request_body: bytes | None) -> dict: + """The request body as a dict, or empty when absent or not JSON.""" + if not request_body: + return {} + try: + parsed = json.loads(request_body) + except (json.JSONDecodeError, UnicodeDecodeError): + return {} + return parsed if isinstance(parsed, dict) else {} + + +def count_generated_images(content: bytes, is_json: bool) -> int: + """How many images the response carried. + + A response that produced nothing counts zero and is never billed; the + reservation is released instead. + """ + if not is_json: + # Raw image bytes (Venice /image/edit, /image/upscale). + return 1 if content else 0 + try: + payload = json.loads(content) + except (json.JSONDecodeError, UnicodeDecodeError): + return 0 + if not isinstance(payload, dict): + return 0 + for field in ("data", "images"): + value = payload.get(field) + if isinstance(value, list): + return len(value) + return 0 diff --git a/routstr/upstream/venice.py b/routstr/upstream/venice.py new file mode 100644 index 00000000..7c88099f --- /dev/null +++ b/routstr/upstream/venice.py @@ -0,0 +1,325 @@ +from __future__ import annotations + +from collections.abc import Callable +from typing import TYPE_CHECKING, Any + +import httpx + +from ..core.logging import get_logger +from ..payment.image_pricing import ImagePriceTier, ImagePricing, produces_images +from ..payment.models import Architecture, Model, Pricing, TopProvider +from .base import BaseUpstreamProvider + +if TYPE_CHECKING: + from ..core.db import UpstreamProviderRow + +logger = get_logger(__name__) + +# ``GET /models`` defaults to ``type=text``, which is why a Venice account +# configured as a generic upstream never sees its image catalog. +_MODELS_TYPE_PARAM = "all" + +# Families this proxy can both route and price. Audio, music and video are +# billed per second or per clip and return no usage object to settle against, +# so exposing them would hand out unpriced inference. +_SUPPORTED_TYPES = frozenset({"text", "image", "inpaint", "upscale", "embedding"}) + +_IMAGE_TYPES = frozenset({"image", "inpaint", "upscale"}) + +# Venice prices text in USD per million tokens; Routstr prices per token. +_USD_PER_MILLION = 1_000_000.0 + +_ARCHITECTURES: dict[str, tuple[str, list[str], list[str]]] = { + "text": ("text->text", ["text"], ["text"]), + "image": ("text->image", ["text"], ["image"]), + "inpaint": ("text+image->image", ["text", "image"], ["image"]), + "upscale": ("image->image", ["image"], ["image"]), + "embedding": ("text->embedding", ["text"], ["embedding"]), +} + +# One call may ask for several images, so the ceiling covers a small batch. +_IMAGES_PER_RESERVATION = 4 + + +def _usd(entry: Any) -> float | None: + """Read the USD leg of a Venice ``{usd, diem}`` price pair.""" + if isinstance(entry, dict): + value = entry.get("usd") + if isinstance(value, (int, float)) and not isinstance(value, bool): + return float(value) + return None + + +def _max_usd(entry: Any) -> float | None: + """Worst-case USD price across a nested Venice price table.""" + direct = _usd(entry) + if direct is not None: + return direct + if not isinstance(entry, dict): + return None + prices = [p for p in (_max_usd(value) for value in entry.values()) if p is not None] + return max(prices) if prices else None + + +def _prices(table: Any, case: Callable[[str], str]) -> dict[str, float]: + """A Venice ``{label: {usd, diem}}`` table as ``{label: usd}``.""" + if not isinstance(table, dict): + return {} + priced = ((case(str(label)), _usd(entry)) for label, entry in table.items()) + return {label: usd for label, usd in priced if usd is not None} + + +def _tables(table: Any) -> dict[str, dict]: + """The nested sub-tables of a Venice price table, keyed by their label.""" + if not isinstance(table, dict): + return {} + return { + str(label): entry for label, entry in table.items() if isinstance(entry, dict) + } + + +def _label(value: Any, case: Callable[[str], str]) -> str | None: + return case(value) if isinstance(value, str) else None + + +def _labels(values: Any, case: Callable[[str], str]) -> list[str]: + if not isinstance(values, list): + return [] + return [case(str(value)) for value in values] + + +class VeniceUpstreamProvider(BaseUpstreamProvider): + """Upstream provider for the Venice.ai API. + + Venice publishes a complete price book on its own catalog, so models are + built from that rather than matched against OpenRouter, which has never + heard of most of Venice's image catalog. + """ + + provider_type = "venice" + default_base_url = "https://api.venice.ai/api/v1" + platform_url = "https://venice.ai/settings/api" + + def __init__(self, api_key: str, provider_fee: float = 1.01): + super().__init__( + base_url=self.default_base_url, api_key=api_key, provider_fee=provider_fee + ) + + @classmethod + def _build_from_row( + cls, provider_row: "UpstreamProviderRow" + ) -> "VeniceUpstreamProvider": + return cls( + api_key=provider_row.api_key, + provider_fee=provider_row.provider_fee, + ) + + @classmethod + def get_provider_metadata(cls) -> dict[str, object]: + return { + "id": cls.provider_type, + "name": "Venice AI", + "default_base_url": cls.default_base_url, + "fixed_base_url": True, + "platform_url": cls.platform_url, + } + + def transform_model_name(self, model_id: str) -> str: + return model_id.removeprefix("venice/") + + async def _fetch_provider_models(self) -> dict: + url = f"{self.base_url.rstrip('/')}/models" + headers = {"Authorization": f"Bearer {self.api_key}"} if self.api_key else None + async with httpx.AsyncClient(timeout=30.0) as client: + response = await client.get( + url, params={"type": _MODELS_TYPE_PARAM}, headers=headers + ) + response.raise_for_status() + return response.json() + + async def fetch_models(self) -> list[Model]: + try: + payload = await self._fetch_provider_models() + except Exception as e: + logger.error( + "Error fetching Venice models", + extra={"error": str(e), "error_type": type(e).__name__}, + ) + return [] + + models: list[Model] = [] + skipped: list[str] = [] + for entry in payload.get("data", []): + if not isinstance(entry, dict): + continue + try: + model = self._parse_model(entry) + except Exception as e: + logger.warning( + "Failed to parse Venice model", + extra={ + "model_id": entry.get("id", "unknown"), + "error": str(e), + "error_type": type(e).__name__, + }, + ) + continue + if model is None: + skipped.append(str(entry.get("id", "unknown"))) + continue + models.append(model) + + if skipped: + logger.debug( + f"({len(skipped)}) Venice models skipped as unsupported or unpriced", + extra={"skipped_models": skipped}, + ) + return models + + def _parse_model(self, entry: dict[str, Any]) -> Model | None: + model_type = entry.get("type") + model_id = entry.get("id") + spec = entry.get("model_spec") + if not model_id or model_type not in _SUPPORTED_TYPES: + return None + if not isinstance(spec, dict) or spec.get("offline"): + return None + + pricing = self._parse_pricing(str(model_type), spec.get("pricing")) + if pricing is None: + return None + image_pricing = ( + self._build_image_pricing(spec.get("pricing"), spec) + if model_type in _IMAGE_TYPES + else None + ) + + modality, input_modalities, output_modalities = _ARCHITECTURES[str(model_type)] + capabilities = spec.get("capabilities") + if ( + model_type == "text" + and isinstance(capabilities, dict) + and capabilities.get("supportsVision") + ): + input_modalities = [*input_modalities, "image"] + modality = "text+image->text" + + context_length = spec.get("availableContextTokens") + max_completion_tokens = spec.get("maxCompletionTokens") + name = spec.get("name") or str(model_id) + + return Model( + id=str(model_id), + name=str(name), + created=int(entry.get("created") or 0), + description=str(spec.get("description") or f"Venice {model_type} model"), + context_length=int(context_length) if context_length else 0, + architecture=Architecture( + modality=modality, + input_modalities=input_modalities, + output_modalities=output_modalities, + tokenizer="Unknown", + instruct_type=None, + ), + pricing=pricing, + image_pricing=image_pricing, + top_provider=TopProvider( + context_length=int(context_length) if context_length else None, + max_completion_tokens=int(max_completion_tokens) + if max_completion_tokens + else None, + ), + ) + + def _parse_pricing(self, model_type: str, raw: Any) -> Pricing | None: + if not isinstance(raw, dict): + return None + + if model_type in _IMAGE_TYPES: + per_image = self._per_image_usd(raw) + if per_image is None: + return None + return Pricing(prompt=0.0, completion=0.0, image=per_image) + + # The ``extended`` tier some models charge past a context threshold is + # ignored: billing it would overcharge every request staying under it. + input_usd = _usd(raw.get("input")) + if input_usd is None: + return None + return Pricing( + prompt=input_usd / _USD_PER_MILLION, + completion=(_usd(raw.get("output")) or 0.0) / _USD_PER_MILLION, + input_cache_read=(_usd(raw.get("cache_input")) or 0.0) / _USD_PER_MILLION, + input_cache_write=(_usd(raw.get("cache_write")) or 0.0) / _USD_PER_MILLION, + ) + + def _build_image_pricing( + self, raw: Any, spec: dict[str, Any] + ) -> ImagePricing | None: + """Venice's per-tier image prices as the model's own price book. + + ``constraints`` carries the resolution and quality applied when the + request names neither, so a default request is priced at the default + tier rather than the ceiling. + """ + if not isinstance(raw, dict): + return None + max_usd = self._per_image_usd(raw) + if max_usd is None: + return None + + tiers = [ + ImagePriceTier(resolution=label, usd=price) + for label, price in _prices(raw.get("resolutions"), str.upper).items() + ] + for label, steps in _tables(raw.get("quality")).items(): + tiers += [ + ImagePriceTier(resolution=label.upper(), quality=step, usd=price) + for step, price in _prices(steps, str.lower).items() + ] + + constraints = spec.get("constraints") + constraints = constraints if isinstance(constraints, dict) else {} + + return ImagePricing( + max_usd=max_usd, + tiers=tiers, + default_resolution=_label(constraints.get("defaultResolution"), str.upper), + default_quality=_label(constraints.get("defaultQuality"), str.lower), + resolutions=_labels(constraints.get("resolutions"), str.upper), + qualities=_labels(constraints.get("qualities"), str.lower), + upscale=_prices(raw.get("upscale"), str.lower), + ) + + @staticmethod + def _per_image_usd(raw: dict[str, Any]) -> float | None: + """Worst-case USD for one generation. + + ``upscale`` and ``inputImages`` price a separate call and a per-extra- + image surcharge, so folding them in would inflate every reservation. + """ + candidates = [ + _usd(raw.get("generation")), + _usd(raw.get("inpaint")), + _max_usd(raw.get("resolutions")), + _max_usd(raw.get("quality")), + ] + priced = [c for c in candidates if c is not None] + return max(priced) if priced else None + + def _apply_provider_fee_to_model(self, model: Model) -> Model: + """Reserve a batch of images for image models, tokens for the rest. + + The inherited max-cost formula reads a per-image rate as a per-input- + image surcharge, reserving a hundred generations for one image. + """ + if not produces_images(model): + return super()._apply_provider_fee_to_model(model) + + adjusted = Pricing.parse_obj( + {k: v * self.provider_fee for k, v in model.pricing.dict().items()} + ) + adjusted.max_prompt_cost = 0.0 + adjusted.max_completion_cost = adjusted.image * _IMAGES_PER_RESERVATION + adjusted.max_cost = adjusted.max_completion_cost + return model.copy(update={"pricing": adjusted}) diff --git a/tests/integration/test_model_serialization.py b/tests/integration/test_model_serialization.py index 791669d7..b1477472 100644 --- a/tests/integration/test_model_serialization.py +++ b/tests/integration/test_model_serialization.py @@ -349,6 +349,8 @@ def _expected_serialised_model(provider_id: int) -> dict: "enabled": True, "forwarded_model_id": _SEEDED_MODEL_ID, "id": _SEEDED_MODEL_ID, + # Only models that return images carry a per-image price book. + "image_pricing": None, "name": f"SerTest {_SEEDED_MODEL_ID}", "per_request_limits": None, "pricing": { diff --git a/tests/unit/test_image_generation_billing.py b/tests/unit/test_image_generation_billing.py new file mode 100644 index 00000000..0216c8d3 --- /dev/null +++ b/tests/unit/test_image_generation_billing.py @@ -0,0 +1,324 @@ +"""Image-generation endpoints are billed per image returned. + +Image responses carry no usage object, so the generic non-chat path released +the reservation and served them free. These tests pin the flat per-image +settlement that replaces it, and that a response producing none costs nothing. +""" + +import json +import math +from typing import Any +from unittest.mock import AsyncMock, MagicMock, patch + +import httpx +import pytest +from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine +from sqlmodel import SQLModel +from sqlmodel.ext.asyncio.session import AsyncSession + +import routstr.auth as auth_module +from routstr.auth import ReservationSnapshot, get_reservation_snapshot, pay_for_request +from routstr.core.db import ApiKey, ModelRow, ReservationRelease +from routstr.payment.helpers import calculate_discounted_max_cost +from routstr.payment.image_pricing import ImagePriceTier, ImagePricing +from routstr.payment.models import Architecture, Model, Pricing, _row_to_model +from routstr.upstream.base import BaseUpstreamProvider + +BALANCE = 100_000 +RESERVED = 5_000 +SATS_PER_IMAGE = 1.0 + +IMAGE_MODEL = Model( + id="venice-sd35", + name="Venice SD35", + created=0, + description="", + context_length=0, + architecture=Architecture( + modality="text->image", + input_modalities=["text"], + output_modalities=["image"], + tokenizer="Unknown", + instruct_type=None, + ), + pricing=Pricing(prompt=0.0, completion=0.0, image=0.0005), + sats_pricing=Pricing(prompt=0.0, completion=0.0, image=SATS_PER_IMAGE), +) +UNPRICED_MODEL = IMAGE_MODEL.copy( + update={"sats_pricing": Pricing(prompt=0.0, completion=0.0, image=0.0)} +) + +# Ceiling $0.40, default tier 1K/medium at $0.10, cheapest 1K/low at $0.04. +PRICE_BOOK = ImagePricing( + max_usd=0.4, + tiers=[ + ImagePriceTier(resolution="1K", quality="low", usd=0.04), + ImagePriceTier(resolution="1K", quality="medium", usd=0.1), + ImagePriceTier(resolution="2K", quality="high", usd=0.4), + ], + default_resolution="1K", + default_quality="medium", + resolutions=["1K", "2K"], + qualities=["low", "medium", "high"], +) +TIERED_MODEL = IMAGE_MODEL.copy(update={"image_pricing": PRICE_BOOK}) + +BODY = {"model": IMAGE_MODEL.id, "prompt": "a cat"} + + +@pytest.fixture(autouse=True) +def patch_sats_usd_price() -> Any: + with patch("routstr.payment.cost_calculation.sats_usd_price", return_value=5.0e-4): + yield + + +async def _engine() -> AsyncEngine: + engine = create_async_engine("sqlite+aiosqlite://") + async with engine.begin() as connection: + await connection.run_sync(SQLModel.metadata.create_all) + return engine + + +def _upstream(content: bytes, content_type: str) -> httpx.Response: + return httpx.Response( + 200, + content=content, + headers={"content-type": content_type}, + request=httpx.Request("POST", "http://upstream"), + ) + + +async def _drain(response: Any) -> bytes: + body = b"" + if hasattr(response, "body_iterator"): + async for chunk in response.body_iterator: + body += chunk if isinstance(chunk, bytes) else chunk.encode() + else: + body = response.body + return body + + +async def _forward( + engine: AsyncEngine, + path: str, + upstream: httpx.Response, + model: Model = IMAGE_MODEL, + body: dict | None = None, +) -> tuple[bytes, ReservationSnapshot]: + provider = BaseUpstreamProvider( + base_url="http://upstream", api_key="k", provider_fee=1.0 + ) + request = MagicMock() + request.method = "POST" + request.query_params = {} + send = AsyncMock(return_value=upstream) + + async with AsyncSession(engine, expire_on_commit=False) as session: + key = ApiKey(hashed_key="key", balance=BALANCE) + session.add(key) + await session.commit() + await pay_for_request(key, RESERVED, session) + snapshot = await get_reservation_snapshot(key, session) + + with ( + patch("httpx.AsyncClient.send", send), + patch( + "routstr.upstream.base.create_session", + side_effect=lambda: AsyncSession(engine, expire_on_commit=False), + ), + patch( + "routstr.upstream.base.adjust_payment_for_tokens", + auth_module.adjust_payment_for_tokens, + ), + ): + response = await provider.forward_request( + request, + path, + {}, + json.dumps(body if body is not None else BODY).encode(), + key, + RESERVED, + session, + model, + snapshot, + ) + out = await _drain(response) + return out, snapshot + + +async def _ledger( + engine: AsyncEngine, snapshot: ReservationSnapshot +) -> tuple[int, int, int, str | None]: + async with AsyncSession(engine, expire_on_commit=False) as session: + key = await session.get(ApiKey, snapshot.key_hash) + record = await session.get(ReservationRelease, snapshot.release_id) + assert key is not None + return ( + key.balance, + key.total_spent, + key.reserved_balance, + record.status if record else None, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("path", "payload", "expected_images"), + [ + ("v1/images/generations", {"data": [{"b64_json": "a"}, {"b64_json": "b"}]}, 2), + ("v1/image/generate", {"id": "gen-1", "images": ["a", "b", "c"]}, 3), + ("openai/v1/images/edits", {"data": [{"url": "http://x"}]}, 1), + ], +) +async def test_images_are_charged_per_returned_image( + path: str, payload: dict, expected_images: int +) -> None: + engine = await _engine() + out, snapshot = await _forward( + engine, path, _upstream(json.dumps(payload).encode(), "application/json") + ) + + assert json.loads(out) == payload + balance, spent, reserved, status = await _ledger(engine, snapshot) + expected_msats = int(expected_images * SATS_PER_IMAGE * 1000) + assert spent == expected_msats + assert balance == BALANCE - expected_msats + assert reserved == 0 + assert status == "charged" + + +@pytest.mark.asyncio +async def test_binary_image_response_counts_as_one_image() -> None: + engine = await _engine() + out, snapshot = await _forward( + engine, "v1/image/upscale", _upstream(b"\x89PNG\r\n", "image/png") + ) + + assert out == b"\x89PNG\r\n" + _, spent, reserved, _ = await _ledger(engine, snapshot) + assert spent == int(SATS_PER_IMAGE * 1000) + assert reserved == 0 + + +@pytest.mark.asyncio +async def test_empty_image_response_is_not_charged() -> None: + engine = await _engine() + _, snapshot = await _forward( + engine, + "v1/images/generations", + _upstream(json.dumps({"data": []}).encode(), "application/json"), + ) + + balance, spent, reserved, _ = await _ledger(engine, snapshot) + assert spent == 0 + assert balance == BALANCE + assert reserved == 0 + + +@pytest.mark.asyncio +async def test_model_without_a_per_image_rate_is_not_charged() -> None: + engine = await _engine() + _, snapshot = await _forward( + engine, + "v1/images/generations", + _upstream( + json.dumps({"data": [{"b64_json": "a"}]}).encode(), "application/json" + ), + model=UNPRICED_MODEL, + ) + + balance, spent, reserved, _ = await _ledger(engine, snapshot) + assert spent == 0 + assert balance == BALANCE + assert reserved == 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("body", "expected_usd"), + [ + # No tier named: the model's own default, not its ceiling. + ({"model": IMAGE_MODEL.id, "prompt": "a cat"}, 0.1), + ({"model": IMAGE_MODEL.id, "resolution": "1K", "quality": "low"}, 0.04), + ({"model": IMAGE_MODEL.id, "resolution": "2K", "quality": "high"}, 0.4), + # A tier the model never declared is billed at the ceiling. + ({"model": IMAGE_MODEL.id, "resolution": "4K"}, 0.4), + ], +) +async def test_charge_uses_the_tier_the_request_asked_for( + body: dict, expected_usd: float +) -> None: + engine = await _engine() + _, snapshot = await _forward( + engine, + "v1/images/generations", + _upstream( + json.dumps({"data": [{"b64_json": "a"}]}).encode(), "application/json" + ), + model=TIERED_MODEL, + body=body, + ) + + # sats_pricing.image is the ceiling in sats; a tier scales against max_usd. + expected_msats = math.ceil( + SATS_PER_IMAGE * (expected_usd / PRICE_BOOK.max_usd) * 1000 + ) + _, spent, reserved, _ = await _ledger(engine, snapshot) + assert spent == expected_msats + assert reserved == 0 + + +@pytest.mark.asyncio +async def test_reservation_holds_the_requested_tier_not_the_ceiling() -> None: + ceiling = await calculate_discounted_max_cost( + 999_999, + {"model": TIERED_MODEL.id, "resolution": "2K", "quality": "high"}, + model_obj=TIERED_MODEL, + ) + cheapest = await calculate_discounted_max_cost( + 999_999, + {"model": TIERED_MODEL.id, "resolution": "1K", "quality": "low"}, + model_obj=TIERED_MODEL, + ) + + assert ceiling == math.ceil(SATS_PER_IMAGE * 1000) + assert cheapest == math.ceil(SATS_PER_IMAGE * (0.04 / 0.4) * 1000) + + +def test_price_book_survives_a_database_round_trip() -> None: + row = ModelRow( + id=TIERED_MODEL.id, + upstream_provider_id=1, + name=TIERED_MODEL.name, + created=0, + description="", + context_length=0, + architecture=TIERED_MODEL.architecture.json(), + pricing=TIERED_MODEL.pricing.json(), + image_pricing=PRICE_BOOK.json(), + ) + + restored = _row_to_model(row).image_pricing + + assert restored is not None + assert restored.max_usd == pytest.approx(0.4) + assert restored.default_resolution == "1K" + assert restored.default_quality == "medium" + assert [(t.resolution, t.quality, t.usd) for t in restored.tiers] == [ + ("1K", "low", 0.04), + ("1K", "medium", 0.1), + ("2K", "high", 0.4), + ] + + +@pytest.mark.asyncio +async def test_reservation_is_sized_from_the_requested_batch() -> None: + single = await calculate_discounted_max_cost( + 999_999, {"model": IMAGE_MODEL.id}, model_obj=IMAGE_MODEL + ) + batch = await calculate_discounted_max_cost( + 999_999, {"model": IMAGE_MODEL.id, "n": 3}, model_obj=IMAGE_MODEL + ) + + assert single == int(SATS_PER_IMAGE * 1000) + assert batch == int(3 * SATS_PER_IMAGE * 1000) diff --git a/tests/unit/test_upstream_venice.py b/tests/unit/test_upstream_venice.py new file mode 100644 index 00000000..28606cec --- /dev/null +++ b/tests/unit/test_upstream_venice.py @@ -0,0 +1,329 @@ +"""Unit tests for ``VeniceUpstreamProvider.fetch_models``. + +Venice answers ``/models`` with only its text catalog unless ``type`` is +passed, which is why the same account configured as a generic upstream shows +no image models. These tests pin that query parameter, the per-family pricing +shapes, and the tier lookup they feed. +""" + +from __future__ import annotations + +from typing import Any +from unittest.mock import patch + +import pytest + +from routstr.payment.image_pricing import per_image_sats +from routstr.upstream.image_generation import parse_json_body +from routstr.upstream.venice import VeniceUpstreamProvider + + +class _FakeResponse: + def __init__(self, payload: dict[str, Any]) -> None: + self._payload = payload + + def raise_for_status(self) -> None: + return None + + def json(self) -> dict[str, Any]: + return self._payload + + +class _FakeAsyncClient: + def __init__(self, payload: dict[str, Any], calls: list[dict[str, Any]]) -> None: + self._payload = payload + self._calls = calls + + async def __aenter__(self) -> "_FakeAsyncClient": + return self + + async def __aexit__(self, *_: object) -> None: + return None + + async def get( + self, + url: str, + params: dict[str, Any] | None = None, + headers: dict[str, str] | None = None, + ) -> _FakeResponse: + self._calls.append({"url": url, "params": params, "headers": headers}) + return _FakeResponse(self._payload) + + +CATALOG: dict[str, Any] = { + "object": "list", + "data": [ + { + "id": "venice-uncensored-1-2", + "type": "text", + "created": 1727966436, + "model_spec": { + "name": "Venice Uncensored 1.2", + "availableContextTokens": 128000, + "maxCompletionTokens": 8192, + "capabilities": {"supportsVision": True}, + "pricing": { + "input": {"usd": 0.2, "diem": 0.2}, + "output": {"usd": 0.9, "diem": 0.9}, + "cache_input": {"usd": 0.02, "diem": 0.02}, + }, + }, + }, + { + "id": "venice-sd35", + "type": "image", + "created": 1727966436, + "model_spec": { + "name": "Venice SD35", + "pricing": { + "generation": {"usd": 0.01, "diem": 0.01}, + "upscale": {"4x": {"usd": 0.08, "diem": 0.08}}, + }, + }, + }, + { + "id": "grok-imagine-image-quality", + "type": "image", + "created": 1727966436, + "model_spec": { + "name": "Grok Imagine High Quality", + "pricing": { + "resolutions": { + "1K": {"usd": 0.06, "diem": 0.06}, + "2K": {"usd": 0.09, "diem": 0.09}, + }, + "upscale": {"4x": {"usd": 0.08, "diem": 0.08}}, + }, + }, + }, + { + "id": "gpt-image-2", + "type": "image", + "created": 1727966436, + "model_spec": { + "name": "GPT Image 2", + "constraints": { + "defaultResolution": "1K", + "resolutions": ["1K", "2K"], + "defaultQuality": "medium", + "qualities": ["low", "medium", "high"], + }, + "pricing": { + "resolutions": { + "1K": {"usd": 0.07, "diem": 0.07}, + "2K": {"usd": 0.1, "diem": 0.1}, + }, + "quality": { + "1K": { + "low": {"usd": 0.02, "diem": 0.02}, + "medium": {"usd": 0.07, "diem": 0.07}, + "high": {"usd": 0.27, "diem": 0.27}, + }, + "2K": { + "low": {"usd": 0.03, "diem": 0.03}, + "high": {"usd": 0.5, "diem": 0.5}, + }, + }, + }, + }, + }, + { + "id": "flux-2-max-edit", + "type": "inpaint", + "created": 1727966436, + "model_spec": { + "name": "FLUX.2 Max Edit", + "pricing": { + "inpaint": {"usd": 0.12, "diem": 0.12}, + "inputImages": { + "included": 1, + "additional": {"usd": 0.0345, "diem": 0.0345}, + }, + }, + }, + }, + { + "id": "tts-kokoro", + "type": "tts", + "created": 1727966436, + "model_spec": { + "name": "Kokoro", + "pricing": {"input": {"usd": 3.5, "diem": 3.5}}, + }, + }, + { + "id": "offline-model", + "type": "image", + "created": 1727966436, + "model_spec": { + "name": "Offline", + "offline": True, + "pricing": {"generation": {"usd": 0.01, "diem": 0.01}}, + }, + }, + { + "id": "unpriced-video", + "type": "video", + "created": 1727966436, + "model_spec": {"name": "Video"}, + }, + ], +} + + +def _fetch(payload: dict[str, Any] = CATALOG) -> tuple[list[Any], list[dict[str, Any]]]: + import asyncio + + calls: list[dict[str, Any]] = [] + provider = VeniceUpstreamProvider(api_key="sk-test") + with patch( + "routstr.upstream.venice.httpx.AsyncClient", + lambda *a, **kw: _FakeAsyncClient(payload, calls), + ): + models = asyncio.run(provider.fetch_models()) + return models, calls + + +def test_requests_every_model_family() -> None: + _, calls = _fetch() + assert calls[0]["params"] == {"type": "all"} + assert calls[0]["url"] == "https://api.venice.ai/api/v1/models" + assert calls[0]["headers"] == {"Authorization": "Bearer sk-test"} + + +def test_image_models_are_listed() -> None: + models, _ = _fetch() + by_id = {m.id: m for m in models} + assert "venice-sd35" in by_id + assert by_id["venice-sd35"].architecture.output_modalities == ["image"] + assert by_id["venice-sd35"].pricing.image == pytest.approx(0.01) + + +def test_image_pricing_uses_worst_case_resolution_not_upscale() -> None: + models, _ = _fetch() + by_id = {m.id: m for m in models} + # 0.09 is the 2K generation price; 0.08 is a separate /image/upscale call. + assert by_id["grok-imagine-image-quality"].pricing.image == pytest.approx(0.09) + # inputImages is a per-extra-image surcharge, not the generation price. + assert by_id["flux-2-max-edit"].pricing.image == pytest.approx(0.12) + + +def test_text_pricing_is_per_token() -> None: + models, _ = _fetch() + model = next(m for m in models if m.id == "venice-uncensored-1-2") + assert model.pricing.prompt == pytest.approx(0.2 / 1_000_000) + assert model.pricing.completion == pytest.approx(0.9 / 1_000_000) + assert model.pricing.input_cache_read == pytest.approx(0.02 / 1_000_000) + assert model.context_length == 128000 + assert model.architecture.input_modalities == ["text", "image"] + + +def test_unsupported_offline_and_unpriced_models_are_dropped() -> None: + models, _ = _fetch() + ids = {m.id for m in models} + assert "tts-kokoro" not in ids + assert "unpriced-video" not in ids + assert "offline-model" not in ids + + +def test_worst_case_rate_covers_the_quality_table() -> None: + models, _ = _fetch() + by_id = {m.id: m for m in models} + # 0.5 is the 2K/high quality tier, above every resolutions entry. + assert by_id["gpt-image-2"].pricing.image == pytest.approx(0.5) + + +@pytest.mark.parametrize( + ("model_id", "body", "expected_usd"), + [ + # Venice's own resolution label, and the OpenAI size that maps to it. + ("grok-imagine-image-quality", {"resolution": "1K"}, 0.06), + ("grok-imagine-image-quality", {"size": "1024x1024"}, 0.06), + ("grok-imagine-image-quality", {"width": 2048, "height": 1024}, 0.09), + # Quality wins over the bare resolution price when both are known. + ("gpt-image-2", {"resolution": "1K", "quality": "low"}, 0.02), + ("gpt-image-2", {"size": "1024x1536", "quality": "HIGH"}, 0.27), + ("gpt-image-2", {"resolution": "1K"}, 0.07), + # No tier named: the upstream's own defaults, not the ceiling. + ("gpt-image-2", {}, 0.07), + # Declared but unpriced quality step still resolves through resolution. + ("gpt-image-2", {"resolution": "2K", "quality": "medium"}, 0.1), + # A tier the model never declared is billed at the ceiling. + ("gpt-image-2", {"resolution": "4K"}, 0.5), + ("gpt-image-2", {"quality": "ultra"}, 0.5), + # A model with no tiers at all prices everything the same. + ("venice-sd35", {"resolution": "2K", "quality": "low"}, 0.01), + ("venice-sd35", {}, 0.01), + ], +) +def test_per_image_price_narrows_to_the_requested_tier( + model_id: str, body: dict[str, Any], expected_usd: float +) -> None: + models, _ = _fetch() + model = next(m for m in models if m.id == model_id) + # sats_pricing is 1 sat per USD here, so the tier reads back in USD. + model = model.copy(update={"sats_pricing": model.pricing}) + + assert per_image_sats(model, body) == pytest.approx(expected_usd) + + +def test_per_image_price_survives_an_unreadable_body() -> None: + models, _ = _fetch() + model = next(m for m in models if m.id == "gpt-image-2") + model = model.copy(update={"sats_pricing": model.pricing}) + + # A body naming no tier prices at the model's own defaults. + assert per_image_sats(model, parse_json_body(b"not json")) == pytest.approx(0.07) + assert per_image_sats(model, parse_json_body(None)) == pytest.approx(0.07) + + +def test_image_price_book_is_carried_on_the_model() -> None: + models, _ = _fetch() + book = next(m for m in models if m.id == "gpt-image-2").image_pricing + + assert book is not None + assert book.max_usd == pytest.approx(0.5) + assert book.default_resolution == "1K" + assert book.default_quality == "medium" + assert book.resolutions == ["1K", "2K"] + assert book.qualities == ["low", "medium", "high"] + assert {(t.resolution, t.quality, t.usd) for t in book.tiers} == { + ("1K", None, 0.07), + ("2K", None, 0.1), + ("1K", "low", 0.02), + ("1K", "medium", 0.07), + ("1K", "high", 0.27), + ("2K", "low", 0.03), + ("2K", "high", 0.5), + } + + +def test_upscale_factors_are_priced_separately() -> None: + models, _ = _fetch() + book = next(m for m in models if m.id == "venice-sd35").image_pricing + + assert book is not None + assert book.upscale == {"4x": 0.08} + assert book.upscale_usd("4x") == pytest.approx(0.08) + # An unnamed factor takes the dearest one rather than under-billing. + assert book.upscale_usd(None) == pytest.approx(0.08) + + +def test_models_without_images_carry_no_price_book() -> None: + models, _ = _fetch() + by_id = {m.id: m for m in models} + + assert by_id["venice-uncensored-1-2"].image_pricing is None + + +def test_image_reservation_covers_a_batch_not_a_token_window() -> None: + models, _ = _fetch() + provider = VeniceUpstreamProvider(api_key="sk-test", provider_fee=1.0) + image_model = next(m for m in models if m.id == "venice-sd35") + priced = provider._apply_provider_fee_to_model(image_model) + assert priced.pricing.max_cost == pytest.approx(0.04) + assert priced.pricing.max_prompt_cost == 0.0 + + text_model = next(m for m in models if m.id == "venice-uncensored-1-2") + text_priced = provider._apply_provider_fee_to_model(text_model) + assert text_priced.pricing.max_cost > 0