diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index 00c6ce97e78..fb430543f69 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -2,6 +2,7 @@ ## Helper utilities import copy import logging +import math import re from collections.abc import Collection, Iterable, Iterator, Mapping from types import MappingProxyType @@ -937,8 +938,8 @@ def _usage_field(usage: object, name: str) -> object: return getattr(usage, name, None) -def _cost_number(value: object) -> float | None: - if isinstance(value, bool) or not isinstance(value, int | float): +def finite_cost(value: object) -> float | None: + if isinstance(value, bool) or not isinstance(value, int | float) or not math.isfinite(value): return None return float(value) @@ -947,15 +948,15 @@ def _reported_cost_value(value: object) -> float | None: """A reported cost is a number, or a breakdown whose total lives under ``total_cost``.""" raw: Final[object] = value if isinstance(value, Mapping): - return _cost_number(_usage_field(raw, "total_cost")) - return _cost_number(value) + return finite_cost(_usage_field(raw, "total_cost")) + return finite_cost(value) def response_cost_in_hidden_params(response: _CarriesHiddenParams) -> float | None: additional_headers: Final[object] = response.hidden_params.get("additional_headers") if not isinstance(additional_headers, Mapping): return None - return _cost_number(_USAGE_BLOCK_ADAPTER.validate_python(additional_headers).get(RESPONSE_COST_HEADER)) + return finite_cost(_USAGE_BLOCK_ADAPTER.validate_python(additional_headers).get(RESPONSE_COST_HEADER)) def override_stamped_response_cost(response: _CarriesHiddenParams, cost: float) -> None: diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 48e8a44b9a4..450010f6f28 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -75,6 +75,7 @@ from litellm.litellm_core_utils.classifier_logging import ( is_classifier_call, ) from litellm.litellm_core_utils.core_helpers import ( + finite_cost, get_provider_response_headers_from_hidden_params, is_expected_client_error, native_messages_reported_cost, @@ -7449,7 +7450,15 @@ def get_standard_logging_object_payload( ) custom_pricing: Final = False if spilled_over else logging_obj.custom_pricing_for(init_response_obj) raw_response_cost: Final = kwargs.get("response_cost") - llm_response_cost: Final[float] = raw_response_cost or 0.0 + numeric_response_cost: Final = raw_response_cost if isinstance(raw_response_cost, int | float) else None + priced_response_cost: Final = finite_cost(numeric_response_cost) + if numeric_response_cost is not None and priced_response_cost is None: + verbose_logger.warning( + "Logging a non-finite response_cost=%s as 0 for litellm_call_id=%s", + numeric_response_cost, + logging_obj.litellm_call_id, + ) + llm_response_cost: Final[float] = priced_response_cost or 0.0 guardrail_cost: Final = guardrail_information_cost(metadata.get("standard_logging_guardrail_information")) response_cost: Final[float] = llm_response_cost + guardrail_cost diff --git a/litellm/litellm_core_utils/safe_json_dumps.py b/litellm/litellm_core_utils/safe_json_dumps.py index 63242a580e7..57cb6574bad 100644 --- a/litellm/litellm_core_utils/safe_json_dumps.py +++ b/litellm/litellm_core_utils/safe_json_dumps.py @@ -1,4 +1,5 @@ import json +import math from collections.abc import Callable from typing import Any, Final @@ -14,6 +15,10 @@ def strip_null_bytes(value: str) -> str: return value.replace("\x00", "") +def _finite_or_none(value: float | None) -> float | None: + return None if isinstance(value, float) and not math.isfinite(value) else value + + def safe_json_structure( data: object, max_depth: int = DEFAULT_MAX_RECURSE_DEPTH, @@ -23,7 +28,8 @@ def safe_json_structure( """ Rebuild data out of JSON-native pieces while detecting circular references. If a circular reference is detected then a marker string is returned. - NUL bytes are stripped from strings to prevent PostgreSQL 22P05 errors. + NUL bytes are stripped from strings to prevent PostgreSQL 22P05 errors, and NaN and + infinite floats become null, since JSON has no token for them and jsonb rejects them. value_transform, when given, is applied to every string leaf (and to the str() fallback for non-serializable objects) with the mapping key the leaf @@ -43,7 +49,7 @@ def safe_json_structure( cleaned = obj.replace("\x00", "") if "\x00" in obj else obj return _transform(key, cleaned) if isinstance(obj, (int, float, bool, type(None))): - return obj + return _finite_or_none(obj) # Check for circular reference. if id(obj) in seen: return "CircularReference Detected" diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 8d243e572fa..4ac5693e493 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -12,7 +12,7 @@ from typing_extensions import ReadOnly, Required import litellm from litellm._logging import verbose_logger -from litellm.litellm_core_utils.core_helpers import chunk_usage, provider_reported_usage_cost +from litellm.litellm_core_utils.core_helpers import chunk_usage, finite_cost, provider_reported_usage_cost from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR, set_hidden_params from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import BEDROCK_TRACE_FIELDS, object_items, object_mapping from litellm.litellm_core_utils.prompt_templates.common_utils import ( @@ -1375,9 +1375,9 @@ def _stream_builder_response_cost( for chunk in reversed(chunks) ) return next((cost for cost in reported if cost is not None), None) - usage_cost: Final = getattr(getattr(response, "usage", None), "cost", None) - if isinstance(usage_cost, (int, float)) and not _reported_cost_is_priced_by_calculator(logging_obj): - return float(usage_cost) + usage_cost: Final = finite_cost(getattr(getattr(response, "usage", None), "cost", None)) + if usage_cost is not None and not _reported_cost_is_priced_by_calculator(logging_obj): + return usage_cost if logging_obj is not None: return None provider_hint: Final = response._hidden_params.get( # pyright: ignore[reportPrivateUsage] # no public accessor @@ -1424,7 +1424,7 @@ def _set_stream_builder_response_cost( def _stamp_streaming_usage_cost(usage: Usage, response: ModelResponse, logging_obj: Optional["Logging"]) -> None: if logging_obj is None: return - if isinstance(getattr(usage, "cost", None), (int, float)): + if finite_cost(getattr(usage, "cost", None)) is not None: return computed_cost: Final = logging_obj.response_cost_calculator(result=response) if not isinstance(computed_cost, (int, float)) or computed_cost <= 0: diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 8ac2597c99b..6221eaeb8bc 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -34,6 +34,7 @@ from litellm.constants import ( DB_SPEND_UPDATE_JOB_NAME, INTERNAL_CALL_ORIGIN_METADATA_KEY, ) +from litellm.litellm_core_utils.core_helpers import finite_cost from litellm.litellm_core_utils.litellm_logging import coerce_model_access_groups from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.proxy._types import ( @@ -647,7 +648,14 @@ class DBSpendUpdateWriter: end_time=end_time, llm_router=llm_router, ) - payload["spend"] = response_cost or 0.0 + billable_cost: Final = finite_cost(response_cost) + if response_cost is not None and billable_cost is None: + verbose_proxy_logger.warning( + "Spend tracking - recording a non-finite response_cost=%s as 0 for request_id=%s", + response_cost, + payload["request_id"], + ) + payload["spend"] = billable_cost or 0.0 if isinstance(payload["startTime"], datetime): payload["startTime"] = payload["startTime"].isoformat() if isinstance(payload["endTime"], datetime): @@ -689,7 +697,7 @@ class DBSpendUpdateWriter: spawn_logging_task( asyncio.get_running_loop(), self._batch_database_updates( - response_cost=response_cost, + response_cost=billable_cost, user_id=user_id, hashed_token=hashed_token, team_id=team_id, diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index c11b59a809a..350dab94706 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -16,6 +16,7 @@ from litellm.constants import BACKGROUND_INTERACTION_COST_POLLING_ENABLED, BILLE from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.core_helpers import ( budget_reservation_from_metadata, + finite_cost, get_litellm_metadata_from_kwargs, get_metadata_variable_name_from_kwargs, get_parent_otel_span_from_kwargs, @@ -627,7 +628,7 @@ class _ProxyDBLogger(CustomLogger): raw_cost: Final = _str_keyed_fields( sl_object if sl_object is not None else kwargs # pyright: ignore[reportUnknownArgumentType] # untyped success-callback payload ).get("response_cost") - priced_cost: Final = raw_cost if isinstance(raw_cost, (int, float)) else None + priced_cost: Final = finite_cost(raw_cost) response_cost: Final = ( await apply_search_free_tier( kwargs=kwargs, # pyright: ignore[reportUnknownArgumentType] # untyped success-callback payload diff --git a/tests/integration/spend/test_non_finite_cost_spend_writes.py b/tests/integration/spend/test_non_finite_cost_spend_writes.py new file mode 100644 index 00000000000..ec1e2f87b3e --- /dev/null +++ b/tests/integration/spend/test_non_finite_cost_spend_writes.py @@ -0,0 +1,164 @@ +import json +from hashlib import sha256 +from pathlib import Path +from typing import Final, Literal +from uuid import uuid4 + +import pytest +from integration._support.client import Gateway, eventually +from integration._support.database import read_rows +from integration._support.openai_wire import answering_model_discovery +from integration._support.process import owned_proxy +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue + +pytestmark: Final = pytest.mark.shares_worker_state + +PROMPT_TOKENS: Final = 5 +COMPLETION_TOKENS: Final = 3 +INPUT_RATE: Final = 0.001 +OUTPUT_RATE: Final = 0.002 +TOKEN_PRICE: Final = PROMPT_TOKENS * INPUT_RATE + COMPLETION_TOKENS * OUTPUT_RATE +OVERFLOWING_INPUT_RATE: Final = 1e308 +POISONED_UPSTREAM: Final = "poisoned-upstream" +CONTROL_UPSTREAM: Final = "control-upstream" +DAILY_TOTALS: Final = "SELECT model_group, SUM(spend)::float8 AS spend, SUM(api_requests)::int8 AS api_requests" +DAILY_USER_QUERY: Final = f'{DAILY_TOTALS} FROM "LiteLLM_DailyUserSpend" WHERE api_key=%s GROUP BY model_group' +DAILY_TEAM_QUERY: Final = f'{DAILY_TOTALS} FROM "LiteLLM_DailyTeamSpend" WHERE api_key=%s GROUP BY model_group' +DAILY_TAG_QUERY: Final = f'{DAILY_TOTALS} FROM "LiteLLM_DailyTagSpend" WHERE api_key=%s AND tag=%s GROUP BY model_group' + +Poison = Literal["reported_nan_cost", "overflowing_price"] + + +def _frame(payload: dict[str, JsonValue]) -> bytes: + return b"data: " + json.dumps(payload).encode() + b"\n\n" + + +def _stream(upstream_model: str, reported_cost: float | None) -> tuple[bytes, ...]: + chunk: Final = {"id": f"chatcmpl-{uuid4().hex}", "object": "chat.completion.chunk", "created": 1} + usage: Final[dict[str, JsonValue]] = { + "prompt_tokens": PROMPT_TOKENS, + "completion_tokens": COMPLETION_TOKENS, + "total_tokens": PROMPT_TOKENS + COMPLETION_TOKENS, + } + reported_usage: Final = usage if reported_cost is None else {**usage, "cost": reported_cost} + return ( + _frame({**chunk, "model": upstream_model, "choices": [{"index": 0, "delta": {"content": "ok"}}]}), + _frame( + { + **chunk, + "model": upstream_model, + "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], + "usage": reported_usage, + } + ), + b"data: [DONE]\n\n", + ) + + +def _deployment(alias: str, upstream_model: str, wire_url: str, input_rate: float) -> dict[str, JsonValue]: + return { + "model_name": alias, + "litellm_params": { + "model": f"openai/{upstream_model}", + "api_base": f"{wire_url}/v1", + "api_key": "integration-provider-key", + "input_cost_per_token": input_rate, + "output_cost_per_token": OUTPUT_RATE, + }, + } + + +def _proxy_config(directory: Path, deployments: tuple[dict[str, JsonValue], ...]) -> Path: + config: Final = directory / "non_finite_cost_config.yaml" + config.write_text( + json.dumps( + { + "model_list": list(deployments), + "general_settings": { + "master_key": "os.environ/LITELLM_MASTER_KEY", + "database_url": "os.environ/DATABASE_URL", + "disable_spend_logs": False, + "proxy_batch_write_at": 1, + "proxy_batch_polling_interval": 1, + }, + "router_settings": {"disable_cooldowns": True}, + } + ) + ) + return config + + +def _by_model_group(rows: list[dict[str, JsonValue]]) -> dict[str, float]: + return {str(row["model_group"]): float(str(row["spend"])) for row in rows} + + +def _both_requests_rows(query: str, parameters: tuple[str, ...]) -> list[dict[str, JsonValue]]: + return eventually(lambda: read_rows(query, parameters), lambda rows: len(rows) == 2, seconds=70) + + +@pytest.mark.parametrize("poison", ["reported_nan_cost", "overflowing_price"]) +@pytest.mark.timeout(240) +def test_a_non_finite_cost_never_blocks_the_next_requests_spend_rows( + gateway: Gateway, tmp_path: Path, poison: Poison +) -> None: + poisoned_alias: Final = f"non-finite-cost-{uuid4().hex}" + control_alias: Final = f"finite-cost-{uuid4().hex}" + tag: Final = f"non-finite-cost-{uuid4().hex}" + poisoned_spend: Final = TOKEN_PRICE if poison == "reported_nan_cost" else 0.0 + poisoned_input_rate: Final = INPUT_RATE if poison == "reported_nan_cost" else OVERFLOWING_INPUT_RATE + + def respond(request: Request) -> Reply: + assert request.target.endswith("/chat/completions"), request.target + upstream_model: Final = json.loads(request.body)["model"] + reported_cost: Final = ( + float("nan") if upstream_model == POISONED_UPSTREAM and poison == "reported_nan_cost" else None + ) + return Reply(content_type="text/event-stream", chunks=_stream(upstream_model, reported_cost)) + + with wire_server(answering_model_discovery(respond)) as wire: + config: Final = _proxy_config( + tmp_path, + ( + _deployment(poisoned_alias, POISONED_UPSTREAM, wire.url, poisoned_input_rate), + _deployment(control_alias, CONTROL_UPSTREAM, wire.url, INPUT_RATE), + ), + ) + with owned_proxy(gateway, tmp_path, {}, config=config) as candidate, candidate.scenario() as scenario: + team: Final = scenario.team() + key: Final = scenario.key(user_id=f"non-finite-cost-{uuid4().hex}", team_id=team) + digest: Final = sha256(key.encode()).hexdigest() + for alias in (poisoned_alias, control_alias): + response = candidate.request( + "POST", + "/v1/chat/completions", + { + "model": alias, + "messages": [{"role": "user", "content": f"non-finite cost {uuid4().hex}"}], + "stream": True, + "metadata": {"tags": [tag]}, + }, + key=key, + ) + assert response.status_code == 200, response.text + + expected: Final = {poisoned_alias: pytest.approx(poisoned_spend), control_alias: pytest.approx(TOKEN_PRICE)} + spend_logs: Final = _both_requests_rows( + 'SELECT model_group, spend FROM "LiteLLM_SpendLogs" WHERE api_key=%s', (digest,) + ) + assert _by_model_group(spend_logs) == expected, spend_logs + daily_rows: Final = ( + _both_requests_rows(DAILY_USER_QUERY, (digest,)), + _both_requests_rows(DAILY_TEAM_QUERY, (digest,)), + _both_requests_rows(DAILY_TAG_QUERY, (digest, tag)), + ) + assert [_by_model_group(rows) for rows in daily_rows] == [expected] * 3, daily_rows + assert [{row["api_requests"] for row in rows} for rows in daily_rows] == [{1}] * 3, daily_rows + key_total: Final = poisoned_spend + TOKEN_PRICE + key_spend: Final = eventually( + lambda: read_rows('SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s', (digest,)), + lambda rows: len(rows) == 1 and float(str(rows[0]["spend"])) == pytest.approx(key_total), + seconds=70, + return_last_on_timeout=True, + ) + assert float(str(key_spend[0]["spend"])) == pytest.approx(key_total), key_spend diff --git a/tests/unit/litellm_core_utils/test_core_helpers.py b/tests/unit/litellm_core_utils/test_core_helpers.py index d66c37ba8c2..350ba1c6941 100644 --- a/tests/unit/litellm_core_utils/test_core_helpers.py +++ b/tests/unit/litellm_core_utils/test_core_helpers.py @@ -26,13 +26,22 @@ from litellm.litellm_core_utils.core_helpers import ( normalize_drop_params, preserve_upstream_non_openai_attributes, process_response_headers, + provider_reported_usage_cost, reconstruct_model_name, redact_nested_match_and_regex_keys, + response_cost_in_hidden_params, set_provider_response_headers_in_hidden_params, unbind_budget_reservation_from_callbacks, ) from litellm.proxy._types import UserAPIKeyAuth -from litellm.types.utils import Delta, ImageResponse, ModelResponseStream, StreamingChoices, TranscriptionResponse +from litellm.types.utils import ( + Delta, + ImageResponse, + ModelResponse, + ModelResponseStream, + StreamingChoices, + TranscriptionResponse, +) def test_parent_otel_span_reads_dict_like_metadata_by_index(): @@ -719,3 +728,20 @@ class TestReceiveStampSurvivesASecondMetadataBucket: } assert get_litellm_metadata_from_kwargs(kwargs)["litellm_received_at"] is own + + +@pytest.mark.parametrize("reported", [float("nan"), float("inf"), {"total_cost": float("nan")}]) +def test_a_non_finite_provider_reported_cost_counts_as_unreported(reported: object) -> None: + assert provider_reported_usage_cost({"cost": reported}, "openrouter") is None + assert provider_reported_usage_cost({"cost": {"total_cost": 0.25}}, "openrouter") == 0.25 + + +@pytest.mark.parametrize("header_cost", [float("nan"), float("inf")]) +def test_a_non_finite_response_cost_header_counts_as_unreported(header_cost: float) -> None: + response: Final = ModelResponse() + response._hidden_params = {"additional_headers": {RESPONSE_COST_HEADER: header_cost}} + finite: Final = ModelResponse() + finite._hidden_params = {"additional_headers": {RESPONSE_COST_HEADER: 0.25}} + + assert response_cost_in_hidden_params(response) is None + assert response_cost_in_hidden_params(finite) == 0.25 diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index e786753e1a8..d1fc16f571b 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -6112,6 +6112,28 @@ def test_standard_logging_payload_resolves_kubernetes_pod_routing_from_metadata_ assert payload["metadata"][KUBERNETES_POD_ROUTING_KEY] == expected +@pytest.mark.parametrize("response_cost", [float("nan"), float("inf"), float("-inf")]) +def test_get_standard_logging_object_payload_logs_a_non_finite_response_cost_as_zero( + logging_obj: LitellmLogging, response_cost: float +) -> None: + from datetime import datetime + + from litellm.litellm_core_utils.litellm_logging import get_standard_logging_object_payload + + now: Final = datetime(2026, 1, 1) + payload: Final = get_standard_logging_object_payload( + kwargs={"model": "gpt-4o", "messages": [], "litellm_params": {}, "response_cost": response_cost}, + init_response_obj={}, + start_time=now, + end_time=now, + logging_obj=logging_obj, + status="success", + ) + + assert payload is not None + assert payload["response_cost"] == 0 + + def test_get_standard_logging_object_payload_carries_matched_access_groups(logging_obj): """Access groups stamped at auth time reach the logging payload, so integrations see what a request billed.""" from datetime import datetime diff --git a/tests/unit/litellm_core_utils/test_safe_json_dumps.py b/tests/unit/litellm_core_utils/test_safe_json_dumps.py index 1f2664f33cb..4b46ea3a35d 100644 --- a/tests/unit/litellm_core_utils/test_safe_json_dumps.py +++ b/tests/unit/litellm_core_utils/test_safe_json_dumps.py @@ -169,6 +169,16 @@ def test_null_byte_stripped_in_fallback_str(): assert json.loads(out)["obj"] == "objrepr" +@pytest.mark.parametrize("value", [float("nan"), float("inf"), float("-inf")]) +def test_non_finite_floats_are_written_as_json_null(value: float) -> None: + out = safe_dumps({"cost_breakdown": {"total_cost": value, "output_cost": 0.006}, "costs": [value]}) + + assert json.loads(out, parse_constant=lambda token: pytest.fail(f"non-JSON token {token}")) == { + "cost_breakdown": {"total_cost": None, "output_cost": 0.006}, + "costs": [None], + } + + def test_clean_strings_are_not_run_through_replace(): """Regression for LIT-3910. diff --git a/tests/unit/proxy/db/test_db_spend_update_writer.py b/tests/unit/proxy/db/test_db_spend_update_writer.py index a42b6404278..f273711d9dc 100644 --- a/tests/unit/proxy/db/test_db_spend_update_writer.py +++ b/tests/unit/proxy/db/test_db_spend_update_writer.py @@ -2225,6 +2225,67 @@ async def test_shutdown_drain_finishes_the_rollups_update_database_started(): assert (team_row["api_requests"], team_row["spend"]) == (1, 0.1) +@pytest.mark.asyncio +@pytest.mark.parametrize("response_cost", [float("nan"), float("inf"), float("-inf")]) +async def test_update_database_records_a_non_finite_response_cost_as_zero(response_cost: float) -> None: + db_writer: Final = DBSpendUpdateWriter() + db_writer._insert_spend_log_to_db = AsyncMock() + prisma: Final = MagicMock() + prisma.get_request_status = MagicMock(return_value="success") + fake_payload: Final = { + "request_id": "req-non-finite-cost", + "api_key": "hashed-key", + "user": "test-user", + "team_id": "test-team", + "agent_id": None, + "startTime": "2026-01-01T00:00:00", + "endTime": "2026-01-01T00:00:01", + "model": "gpt-4o-mini", + "custom_llm_provider": "openai", + "model_group": "gpt-4o-mini", + "prompt_tokens": 10, + "completion_tokens": 5, + "spend": response_cost, + "metadata": '{"usage_object": {}}', + } + with ( + patch("litellm.proxy.proxy_server.disable_spend_logs", False), + patch("litellm.proxy.proxy_server.prisma_client", prisma), + patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), + patch("litellm.proxy.proxy_server.litellm_proxy_budget_name", "test-budget"), + patch("litellm.proxy.spend_tracking.spend_tracking_utils.get_logging_payload", return_value=fake_payload), + ): + await db_writer.update_database( + token="test-token", + user_id="test-user", + end_user_id=None, + team_id="test-team", + org_id=None, + kwargs={"model": "gpt-4o-mini", "custom_llm_provider": "openai"}, + completion_response=MagicMock(), + start_time=datetime.now(), + end_time=datetime.now(), + response_cost=response_cost, + ) + await drain_logging_tasks() + + assert db_writer._insert_spend_log_to_db.call_args.kwargs["payload"]["spend"] == 0 + spend: Final = await db_writer.spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions() + assert spend["key_list_transactions"] == {"test-token": 0} + assert spend["team_list_transactions"] == {"test-team": 0} + daily_user: Final = ( + await db_writer.daily_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() + ) + daily_team: Final = ( + await db_writer.daily_team_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() + ) + assert daily_user is not None and daily_team is not None + (user_row,) = daily_user.values() + (team_row,) = daily_team.values() + assert (user_row["api_requests"], user_row["spend"]) == (1, 0) + assert (team_row["api_requests"], team_row["spend"]) == (1, 0) + + @pytest.mark.asyncio async def test_shutdown_drain_waits_for_an_in_flight_rollup_enqueue(): """ diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index 70b268bf2a2..91f33693a6e 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -5041,6 +5041,28 @@ def test_stream_chunk_builder_keeps_provider_reported_usage_cost(): assert response._hidden_params["response_cost"] == pytest.approx(0.5) +@pytest.mark.parametrize("reported_cost", [float("nan"), float("inf")]) +def test_stream_chunk_builder_bills_the_token_price_when_the_provider_cost_is_not_finite(reported_cost: float): + usage_chunk: Final = _stream_builder_text_chunk("gpt-4o", "") + usage_chunk.usage = Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15, cost=reported_cost) + chunks: Final = [ + _stream_builder_text_chunk("gpt-4o", "Hello "), + _stream_builder_text_chunk("gpt-4o", "world.", finish_reason="stop"), + usage_chunk, + ] + + response: Final = litellm.stream_chunk_builder( + chunks=chunks, messages=[{"role": "user", "content": "hi"}], logging_obj=_stream_builder_logging_obj() + ) + + assert response is not None + prompt_cost, completion_cost = litellm.cost_per_token(model="gpt-4o", usage_object=response.usage) + token_price: Final = prompt_cost + completion_cost + assert token_price > 0 + assert getattr(response.usage, "cost", None) == pytest.approx(token_price) + assert response._hidden_params["response_cost"] == pytest.approx(token_price) + + def test_stream_chunk_builder_prices_alias_from_openai_sdk_usage_chunk(): from openai.types.completion_usage import CompletionUsage