mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-07 07:58:09 +00:00
fix(spend): keep a NaN or infinite cost from wedging daily spend writes
A non-finite provider-reported cost now counts as not reported, so the token price bills the request. Any other non-finite cost is logged and recorded as 0, and spend log JSON writes NaN and Infinity as null so Postgres accepts the row.
This commit is contained in:
mateo-berri
parent
b97fcf1af0
commit
f82002f439
12 changed files with 347 additions and 17 deletions
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
164
tests/integration/spend/test_non_finite_cost_spend_writes.py
Normal file
164
tests/integration/spend/test_non_finite_cost_spend_writes.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
||||
|
|
|
|||
|
|
@ -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():
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue