mirror of
https://github.com/Routstr/routstr-core.git
synced 2026-10-05 20:28:23 +00:00
feat: add venice provider with per-image tier pricing and billing
This commit is contained in:
9qeklajc
@@ -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")
|
||||
+9
-2
@@ -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.
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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)
|
||||
@@ -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:
|
||||
|
||||
@@ -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"""
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
@@ -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})
|
||||
@@ -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": {
|
||||
|
||||
@@ -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)
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user