mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-11 08:38:24 +00:00
feat(xai): bill image calls from cost_in_usd_ticks and add grok-imagine-image-2.0 resolution prices (#50481)
* feat(xai): use billed image cost and resolution prices * test(xai): mark image generation tests as a package --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
6b745635c8
commit
94861ecdf7
9 changed files with 322 additions and 2 deletions
|
|
@ -9,6 +9,7 @@ from types import MappingProxyType
|
|||
from typing import Final, Literal, Protocol, SupportsFloat, SupportsIndex, TypedDict, cast
|
||||
from zoneinfo import ZoneInfo, ZoneInfoNotFoundError
|
||||
|
||||
from pydantic import TypeAdapter
|
||||
from typing_extensions import Buffer, ReadOnly
|
||||
|
||||
import litellm
|
||||
|
|
@ -51,6 +52,7 @@ _IMAGE_RESPONSE_CALL_TYPES: Final = frozenset(
|
|||
_VALID_DATA_RESIDENCIES: Final = frozenset(r.value for r in DataResidency)
|
||||
|
||||
_DEPLOYMENT_PRICING_KEYS: Final[frozenset[str]] = frozenset(CustomPricingLiteLLMParams.model_fields)
|
||||
_IMAGE_REQUEST_PARAMS: Final[TypeAdapter[Mapping[str, object] | None]] = TypeAdapter(Mapping[str, object] | None)
|
||||
|
||||
_BASE_CACHE_WRITE_RATE_KEY: Final = "cache_creation_input_token_cost"
|
||||
|
||||
|
|
@ -2324,6 +2326,20 @@ class CostCalculatorUtils:
|
|||
)
|
||||
resolved_n: Final = n if n is not None else (len(completion_response.data) if completion_response.data else 0)
|
||||
pricing: Final = deployment_pricing(model_info)
|
||||
if custom_llm_provider == litellm.LlmProviders.XAI.value:
|
||||
from litellm.llms.xai.image_generation.cost_calculator import cost_calculator as xai_image_cost_calculator
|
||||
|
||||
xai_cost: Final = xai_image_cost_calculator(
|
||||
model=model,
|
||||
image_response=completion_response,
|
||||
quality=quality,
|
||||
n=resolved_n,
|
||||
optional_params=_IMAGE_REQUEST_PARAMS.validate_python(optional_params),
|
||||
call_type=call_type,
|
||||
model_info=pricing,
|
||||
)
|
||||
if xai_cost is not None:
|
||||
return xai_cost
|
||||
model_priced_calculator: Final = _model_priced_image_cost_calculator(custom_llm_provider, model, pricing)
|
||||
|
||||
if model_priced_calculator is not None:
|
||||
|
|
|
|||
67
litellm/llms/xai/image_generation/cost_calculator.py
Normal file
67
litellm/llms/xai/image_generation/cost_calculator.py
Normal file
|
|
@ -0,0 +1,67 @@
|
|||
from collections.abc import Mapping
|
||||
from typing import Final
|
||||
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
from litellm.litellm_core_utils.llm_cost_calc.utils import get_cost_per_unit, resolve_image_model_info
|
||||
from litellm.llms.xai.common_utils import xai_reported_cost_in_usd
|
||||
from litellm.types.utils import ImageResponse, ModelInfo
|
||||
|
||||
_PARAMS: Final = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
||||
def _request_param(optional_params: Mapping[str, object] | None, key: str) -> object:
|
||||
if optional_params is None:
|
||||
return None
|
||||
if key in optional_params:
|
||||
return optional_params[key]
|
||||
extra_body: Final = optional_params.get("extra_body")
|
||||
return _PARAMS.validate_python(extra_body).get(key) if isinstance(extra_body, Mapping) else None
|
||||
|
||||
|
||||
def cost_calculator(
|
||||
model: str,
|
||||
image_response: ImageResponse,
|
||||
n: int,
|
||||
quality: str | None = None,
|
||||
optional_params: Mapping[str, object] | None = None,
|
||||
call_type: str | None = None,
|
||||
model_info: ModelInfo | None = None,
|
||||
) -> float | None:
|
||||
caller_set_prices: Final = model_info is not None and any(
|
||||
model_info.get(key) is not None
|
||||
for key in (
|
||||
"input_cost_per_image",
|
||||
"output_cost_per_image",
|
||||
"output_cost_per_image_1K",
|
||||
"output_cost_per_image_2K",
|
||||
)
|
||||
)
|
||||
billed_cost: Final = xai_reported_cost_in_usd(getattr(image_response.usage, "cost_in_usd_ticks", None))
|
||||
if billed_cost is not None and not caller_set_prices:
|
||||
return billed_cost
|
||||
|
||||
requested_quality: Final = quality or image_response.quality or _request_param(optional_params, "quality") or "auto"
|
||||
resolution: Final = _request_param(optional_params, "resolution") or "1k"
|
||||
is_edit: Final = call_type in ("image_edit", "aimage_edit")
|
||||
low_quality: Final = requested_quality == "low" or (requested_quality == "auto" and not is_edit)
|
||||
if low_quality and resolution in ("1k", "2k"):
|
||||
try:
|
||||
pricing: Final = resolve_image_model_info(model, "xai", model_info)
|
||||
except Exception: # noqa: BLE001 # model lookup raises bare Exception for unmapped models
|
||||
return None
|
||||
resolution_price: Final = get_cost_per_unit(pricing, f"output_cost_per_image_{str(resolution).upper()}", None)
|
||||
deployment_flat_price: Final = next(
|
||||
(
|
||||
get_cost_per_unit(model_info, key, None)
|
||||
for key in ("input_cost_per_image", "output_cost_per_image")
|
||||
if model_info is not None and model_info.get(key) is not None
|
||||
),
|
||||
None,
|
||||
)
|
||||
if deployment_flat_price is not None:
|
||||
return deployment_flat_price * n
|
||||
if resolution_price is not None:
|
||||
return resolution_price * n
|
||||
|
||||
return None
|
||||
|
|
@ -68030,7 +68030,9 @@
|
|||
"input_cost_per_image": 0.06,
|
||||
"litellm_provider": "xai",
|
||||
"mode": "image_generation",
|
||||
"source": "https://docs.x.ai/docs/models",
|
||||
"output_cost_per_image_1K": 0.04,
|
||||
"output_cost_per_image_2K": 0.06,
|
||||
"source": "https://docs.x.ai/developers/migration/imagine-image-quality-nov-2",
|
||||
"supported_endpoints": [
|
||||
"/v1/images/generations"
|
||||
],
|
||||
|
|
|
|||
|
|
@ -476,6 +476,8 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False):
|
|||
output_cost_per_image_512: ReadOnly[float | None]
|
||||
output_cost_per_image_1024: ReadOnly[float | None]
|
||||
output_cost_per_image_1536: ReadOnly[float | None]
|
||||
output_cost_per_image_1K: ReadOnly[float | None]
|
||||
output_cost_per_image_2K: ReadOnly[float | None]
|
||||
ocr_cost_per_page: float | None # for OCR models
|
||||
ocr_cost_per_page_batches: ReadOnly[float | None]
|
||||
ocr_cost_per_credit: float | None # for OCR models priced by credit
|
||||
|
|
@ -4151,6 +4153,8 @@ class CustomPricingLiteLLMParams(MirroredPricingParams):
|
|||
output_cost_per_image_512: float | None = None
|
||||
output_cost_per_image_1024: float | None = None
|
||||
output_cost_per_image_1536: float | None = None
|
||||
output_cost_per_image_1K: float | None = None
|
||||
output_cost_per_image_2K: float | None = None
|
||||
input_cost_per_pixel: float | None = None
|
||||
output_cost_per_pixel: float | None = None
|
||||
|
||||
|
|
|
|||
|
|
@ -7584,6 +7584,8 @@ def get_model_info_helper(
|
|||
output_cost_per_second_4k=_model_info.get("output_cost_per_second_4k", None),
|
||||
output_cost_per_video_per_second=_model_info.get("output_cost_per_video_per_second", None),
|
||||
output_cost_per_image=_model_info.get("output_cost_per_image", None),
|
||||
output_cost_per_image_1K=_model_info.get("output_cost_per_image_1K", None),
|
||||
output_cost_per_image_2K=_model_info.get("output_cost_per_image_2K", None),
|
||||
output_cost_per_pixel=_model_info.get("output_cost_per_pixel", None),
|
||||
output_cost_per_image_token=_model_info.get("output_cost_per_image_token", None),
|
||||
output_cost_per_image_token_batches=_model_info.get("output_cost_per_image_token_batches", None),
|
||||
|
|
|
|||
|
|
@ -68030,7 +68030,9 @@
|
|||
"input_cost_per_image": 0.06,
|
||||
"litellm_provider": "xai",
|
||||
"mode": "image_generation",
|
||||
"source": "https://docs.x.ai/docs/models",
|
||||
"output_cost_per_image_1K": 0.04,
|
||||
"output_cost_per_image_2K": 0.06,
|
||||
"source": "https://docs.x.ai/developers/migration/imagine-image-quality-nov-2",
|
||||
"supported_endpoints": [
|
||||
"/v1/images/generations"
|
||||
],
|
||||
|
|
|
|||
0
tests/unit/llms/xai/image_generation/__init__.py
Normal file
0
tests/unit/llms/xai/image_generation/__init__.py
Normal file
219
tests/unit/llms/xai/image_generation/test_cost_calculator.py
Normal file
219
tests/unit/llms/xai/image_generation/test_cost_calculator.py
Normal file
|
|
@ -0,0 +1,219 @@
|
|||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from openai import AsyncOpenAI, OpenAI
|
||||
|
||||
import litellm
|
||||
from litellm.litellm_core_utils.llm_cost_calc.utils import CostCalculatorUtils
|
||||
from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import convert_to_model_response_object
|
||||
from litellm.types.utils import ImageResponse
|
||||
|
||||
|
||||
def _response(ticks: object = None) -> ImageResponse:
|
||||
return convert_to_model_response_object(
|
||||
response_object={
|
||||
"created": 123,
|
||||
"data": [{"url": "https://image.test/image.png"}],
|
||||
"usage": {"cost_in_usd_ticks": ticks},
|
||||
},
|
||||
response_type="image_generation",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("ticks", (0, 1234567890))
|
||||
@pytest.mark.parametrize("call_type", ("image_generation", "aimage_generation", "image_edit", "aimage_edit"))
|
||||
def test_billed_cost_is_request_total_even_for_multiple_images(ticks: int, call_type: str) -> None:
|
||||
assert (
|
||||
CostCalculatorUtils.route_image_generation_cost_calculator(
|
||||
model="grok-imagine-image-2.0",
|
||||
custom_llm_provider="xai",
|
||||
completion_response=_response(ticks),
|
||||
n=4,
|
||||
call_type=call_type,
|
||||
optional_params={"resolution": "2k", "quality": "medium"},
|
||||
)
|
||||
== ticks / 10_000_000_000
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("ticks", (None, -1, True, "400000000", 1.5))
|
||||
@pytest.mark.parametrize("resolution,unit_price", ((None, 0.17), ("1k", 0.17), ("2k", 0.29)))
|
||||
def test_missing_or_invalid_ticks_use_requested_resolution(
|
||||
ticks: object, resolution: str | None, unit_price: float
|
||||
) -> None:
|
||||
params: Final = {"resolution": resolution} if resolution else {}
|
||||
assert CostCalculatorUtils.route_image_generation_cost_calculator(
|
||||
model="unmapped-image-fixture",
|
||||
custom_llm_provider="xai",
|
||||
completion_response=_response(ticks),
|
||||
n=3,
|
||||
optional_params={"extra_body": params},
|
||||
model_info={"output_cost_per_image_1K": 0.17, "output_cost_per_image_2K": 0.29},
|
||||
) == pytest.approx(unit_price * 3)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"quality,call_type,resolution",
|
||||
(
|
||||
("medium", "image_generation", "1k"),
|
||||
("high", "image_generation", "2k"),
|
||||
(None, "image_edit", "1k"),
|
||||
("auto", "aimage_edit", "2k"),
|
||||
("low", "image_generation", "1.5k"),
|
||||
),
|
||||
)
|
||||
def test_unpublished_rates_keep_legacy_fallback(quality: str | None, call_type: str, resolution: str) -> None:
|
||||
assert CostCalculatorUtils.route_image_generation_cost_calculator(
|
||||
model="unmapped-image-fixture",
|
||||
custom_llm_provider="xai",
|
||||
completion_response=_response(),
|
||||
n=2,
|
||||
quality=quality,
|
||||
call_type=call_type,
|
||||
optional_params={"resolution": resolution},
|
||||
model_info={"input_cost_per_image": 0.71, "output_cost_per_image_1K": 0.17, "output_cost_per_image_2K": 0.29},
|
||||
) == pytest.approx(0.71 * 2)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("provider", ("xai", "openai"))
|
||||
def test_flat_price_and_empty_response_remain_supported(provider: str) -> None:
|
||||
response: Final = ImageResponse(data=[])
|
||||
assert (
|
||||
CostCalculatorUtils.route_image_generation_cost_calculator(
|
||||
model="unmapped-image-fixture",
|
||||
custom_llm_provider=provider,
|
||||
completion_response=response,
|
||||
model_info={"output_cost_per_image": 0.37},
|
||||
)
|
||||
== 0.0
|
||||
)
|
||||
assert CostCalculatorUtils.route_image_generation_cost_calculator(
|
||||
model="unmapped-image-fixture",
|
||||
custom_llm_provider=provider,
|
||||
completion_response=response,
|
||||
n=2,
|
||||
model_info={"output_cost_per_image": 0.37},
|
||||
) == pytest.approx(0.37 * 2)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("quality,call_type", (("low", "image_edit"), ("auto", "image_generation"), (None, None)))
|
||||
def test_low_resolution_price_applies_without_flat_override(quality: str | None, call_type: str | None) -> None:
|
||||
assert CostCalculatorUtils.route_image_generation_cost_calculator(
|
||||
model="unmapped-image-fixture",
|
||||
custom_llm_provider="xai",
|
||||
completion_response=_response(),
|
||||
n=2,
|
||||
quality=quality,
|
||||
call_type=call_type,
|
||||
optional_params={"resolution": "2k"},
|
||||
model_info={"output_cost_per_image_2K": 0.29},
|
||||
) == pytest.approx(0.29 * 2)
|
||||
|
||||
|
||||
def test_deployment_flat_price_overrides_shared_resolution_price() -> None:
|
||||
assert CostCalculatorUtils.route_image_generation_cost_calculator(
|
||||
model="grok-imagine-image-2.0",
|
||||
custom_llm_provider="xai",
|
||||
completion_response=_response(),
|
||||
n=2,
|
||||
model_info={"output_cost_per_image": 0.83},
|
||||
) == pytest.approx(0.83 * 2)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"price_key,resolution",
|
||||
(
|
||||
("input_cost_per_image", "1k"),
|
||||
("output_cost_per_image", "1k"),
|
||||
("output_cost_per_image_1K", "1k"),
|
||||
("output_cost_per_image_2K", "2k"),
|
||||
),
|
||||
)
|
||||
def test_deployment_image_price_overrides_billed_ticks(price_key: str, resolution: str) -> None:
|
||||
assert CostCalculatorUtils.route_image_generation_cost_calculator(
|
||||
model="grok-imagine-image-2.0",
|
||||
custom_llm_provider="xai",
|
||||
completion_response=_response(1234567890),
|
||||
n=3,
|
||||
optional_params={"resolution": resolution},
|
||||
model_info={price_key: 0.83},
|
||||
) == pytest.approx(0.83 * 3)
|
||||
|
||||
|
||||
def test_shared_resolution_prices_are_used_when_ticks_are_absent() -> None:
|
||||
pricing: Final = litellm.get_model_info("xai/grok-imagine-image-2.0")
|
||||
assert CostCalculatorUtils.route_image_generation_cost_calculator(
|
||||
model="grok-imagine-image-2.0",
|
||||
custom_llm_provider="xai",
|
||||
completion_response=_response(),
|
||||
n=2,
|
||||
optional_params={"resolution": "2k"},
|
||||
) == pytest.approx(pricing["output_cost_per_image_2K"] * 2)
|
||||
|
||||
|
||||
def test_missing_all_prices_keeps_missing_cost_error() -> None:
|
||||
with pytest.raises(Exception, match="Model not found in cost map"):
|
||||
CostCalculatorUtils.route_image_generation_cost_calculator(
|
||||
model="unmapped-image-fixture",
|
||||
custom_llm_provider="xai",
|
||||
completion_response=_response(),
|
||||
)
|
||||
|
||||
|
||||
def test_non_xai_usage_ticks_do_not_override_flat_prices() -> None:
|
||||
assert CostCalculatorUtils.route_image_generation_cost_calculator(
|
||||
model="unmapped-image-fixture",
|
||||
custom_llm_provider="openai",
|
||||
completion_response=_response(1234567890),
|
||||
n=2,
|
||||
model_info={"output_cost_per_image": 0.37},
|
||||
) == pytest.approx(0.37 * 2)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("is_async", (False, True))
|
||||
async def test_generation_sdk_preserves_billed_ticks_through_normalization(is_async: bool) -> None:
|
||||
ticks: Final = 1234567890
|
||||
transport: Final = httpx.MockTransport(
|
||||
lambda request: httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"created": 123,
|
||||
"data": [{"url": "https://image.test/output.png"}],
|
||||
"usage": {"cost_in_usd_ticks": ticks},
|
||||
},
|
||||
)
|
||||
)
|
||||
if is_async:
|
||||
async with AsyncOpenAI(
|
||||
api_key="test", base_url="https://image.test/v1", http_client=httpx.AsyncClient(transport=transport)
|
||||
) as client:
|
||||
response: Final = await litellm.aimage_generation(
|
||||
model="xai/grok-imagine-image-2.0",
|
||||
prompt="circle",
|
||||
api_key="test",
|
||||
api_base="https://image.test/v1",
|
||||
client=client,
|
||||
)
|
||||
assert getattr(response.usage, "cost_in_usd_ticks") == ticks
|
||||
assert (
|
||||
litellm.completion_cost(completion_response=response, model="xai/grok-imagine-image-2.0", n=4)
|
||||
== ticks / 10_000_000_000
|
||||
)
|
||||
else:
|
||||
with OpenAI(
|
||||
api_key="test", base_url="https://image.test/v1", http_client=httpx.Client(transport=transport)
|
||||
) as client:
|
||||
sync_response: Final = litellm.image_generation(
|
||||
model="xai/grok-imagine-image-2.0",
|
||||
prompt="circle",
|
||||
api_key="test",
|
||||
api_base="https://image.test/v1",
|
||||
client=client,
|
||||
)
|
||||
assert getattr(sync_response.usage, "cost_in_usd_ticks") == ticks
|
||||
assert (
|
||||
litellm.completion_cost(completion_response=sync_response, model="xai/grok-imagine-image-2.0", n=4)
|
||||
== ticks / 10_000_000_000
|
||||
)
|
||||
8
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
8
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -36733,6 +36733,10 @@ export interface components {
|
|||
output_cost_per_image_1024?: number | null;
|
||||
/** Output Cost Per Image 1536 */
|
||||
output_cost_per_image_1536?: number | null;
|
||||
/** Output Cost Per Image 1K */
|
||||
output_cost_per_image_1K?: number | null;
|
||||
/** Output Cost Per Image 2K */
|
||||
output_cost_per_image_2K?: number | null;
|
||||
/** Output Cost Per Image 512 */
|
||||
output_cost_per_image_512?: number | null;
|
||||
/** Output Cost Per Image Token */
|
||||
|
|
@ -52864,6 +52868,10 @@ export interface components {
|
|||
output_cost_per_image_1024?: number | null;
|
||||
/** Output Cost Per Image 1536 */
|
||||
output_cost_per_image_1536?: number | null;
|
||||
/** Output Cost Per Image 1K */
|
||||
output_cost_per_image_1K?: number | null;
|
||||
/** Output Cost Per Image 2K */
|
||||
output_cost_per_image_2K?: number | null;
|
||||
/** Output Cost Per Image 512 */
|
||||
output_cost_per_image_512?: number | null;
|
||||
/** Output Cost Per Image Token */
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue