mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-02 07:13:02 +00:00
fix(accounting): preserve attempt evidence through retries and failures
This commit is contained in:
Joshua Valluru
parent
264fbde4f7
commit
f95f20b54a
7 changed files with 391 additions and 46 deletions
|
|
@ -271,6 +271,7 @@ if TYPE_CHECKING:
|
|||
from litellm.proxy.hooks.autorouter_baseline_cache import BaselineCacheContext, CapturedBaselineObservation
|
||||
from litellm.proxy.hooks.search_free_tier import SearchAllowanceAllocation
|
||||
from litellm.proxy.spend_tracking.usage_capture import CapturedUsage
|
||||
from litellm.proxy.spend_tracking.usage_meter import UsageCaptureError
|
||||
try:
|
||||
from litellm_enterprise.enterprise_callbacks.callback_controls import (
|
||||
EnterpriseCallbackControls,
|
||||
|
|
@ -642,9 +643,34 @@ class UsageObservation(Enum):
|
|||
ESTIMATED = "estimated"
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class UsageCaptureObservation:
|
||||
usage_json: str | None
|
||||
provenance: UsageObservation
|
||||
response_cost: float | None = None
|
||||
|
||||
@classmethod
|
||||
def from_usage(
|
||||
cls, usage: Usage | None, provenance: UsageObservation, response_cost: float | None = None
|
||||
) -> "UsageCaptureObservation | UsageCaptureError":
|
||||
from litellm.proxy.spend_tracking.usage_meter import UsageCaptureError
|
||||
|
||||
try:
|
||||
serialized: Final = (
|
||||
None if usage is None or provenance is UsageObservation.UNKNOWN else usage.model_dump_json()
|
||||
)
|
||||
except Exception as error: # noqa: BLE001 # observation failures must not interrupt legacy retry handling
|
||||
return UsageCaptureError(str(error))
|
||||
else:
|
||||
return cls(serialized, provenance, response_cost)
|
||||
|
||||
|
||||
class Logging(LiteLLMLoggingBaseClass):
|
||||
usage_capture_callback: ClassVar[
|
||||
Callable[["Logging", object, datetime.datetime, datetime.datetime, bool], None] | None
|
||||
Callable[
|
||||
["Logging", object, datetime.datetime, datetime.datetime, bool], "CapturedUsage | UsageCaptureError | None"
|
||||
]
|
||||
| None
|
||||
] = None
|
||||
|
||||
def capture_usage_before_logging(
|
||||
|
|
@ -2592,7 +2618,9 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
},
|
||||
usage,
|
||||
)
|
||||
self.record_partial_usage_for_failure(usage, raw_cost)
|
||||
self.record_partial_usage_for_failure(
|
||||
usage, raw_cost, observation=self.model_call_details["_accounting_usage_observed"]
|
||||
)
|
||||
return self.prepared_cost
|
||||
|
||||
def record_finalized_cost(self, finalized: FinalizedCost) -> None:
|
||||
|
|
@ -2605,10 +2633,14 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
def record_usage_pricing(self, pricing: object) -> None:
|
||||
self.model_call_details["_accounting_pricing_snapshot"] = pricing
|
||||
|
||||
def record_partial_usage_for_failure(self, usage: Usage, response_cost: float) -> None:
|
||||
def record_partial_usage_for_failure(
|
||||
self, usage: Usage, response_cost: float, *, observation: UsageObservation = UsageObservation.KNOWN
|
||||
) -> None:
|
||||
"""Stash what an interrupted stream already consumed so the failure log bills it instead of zero."""
|
||||
self.model_call_details["combined_usage_object"] = usage
|
||||
self.model_call_details["response_cost"] = response_cost
|
||||
if type(self).usage_capture_callback is not None:
|
||||
self.model_call_details["_accounting_usage_observed"] = observation
|
||||
|
||||
def record_failure_after_earlier_rounds(
|
||||
self, earlier: BilledRoundsSpend, failed_usage: Usage, failed_cost: float
|
||||
|
|
@ -2704,38 +2736,50 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
|
||||
def capture_failed_attempt_usage(self, usage: Usage) -> None:
|
||||
from litellm.constants import BILLED_RETRY_ATTEMPT_METADATA_KEY
|
||||
from litellm.proxy.spend_tracking.usage_capture import CapturedUsage
|
||||
|
||||
capture: Final = type(self).usage_capture_callback
|
||||
if capture is None:
|
||||
return
|
||||
attempt: Final = self.attempt_scoped_copy()
|
||||
attempt.prepared_cost = None
|
||||
attempt.finalized_cost = None
|
||||
identity: Final = (
|
||||
f"{self.litellm_call_id}-attempt-{len(self.model_call_details.get('failed_attempts_usage', ()))}"
|
||||
)
|
||||
params: Final = _LITELLM_PARAMS.validate_python(self.model_call_details.get("litellm_params") or {})
|
||||
metadata: Final = StandardLoggingPayloadSetup.merge_litellm_metadata(params)
|
||||
observed: Final[object] = self.model_call_details.get("_accounting_usage_observed")
|
||||
cost: Final[object] = self.model_call_details.get("response_cost")
|
||||
attempt.model_call_details = {
|
||||
**attempt.model_call_details,
|
||||
"litellm_call_id": identity,
|
||||
"parent_request_id": self.litellm_call_id,
|
||||
"_accounting_capture": None,
|
||||
"_accounting_prior_captures": (),
|
||||
"_accounting_usage_snapshot": UsageCaptureObservation.from_usage(
|
||||
usage,
|
||||
observed if isinstance(observed, UsageObservation) else UsageObservation.KNOWN,
|
||||
float(cost) if isinstance(cost, (int, float)) else None,
|
||||
),
|
||||
"litellm_params": {
|
||||
**params,
|
||||
"metadata": {**metadata, BILLED_RETRY_ATTEMPT_METADATA_KEY: True, "status": "failure"},
|
||||
"litellm_metadata": {**metadata, BILLED_RETRY_ATTEMPT_METADATA_KEY: True, "status": "failure"},
|
||||
},
|
||||
}
|
||||
capture(
|
||||
captured: Final = capture(
|
||||
attempt,
|
||||
{"id": identity, "usage": usage.model_dump(), "status": "failed"},
|
||||
self.start_time,
|
||||
datetime.datetime.now(datetime.timezone.utc),
|
||||
False,
|
||||
)
|
||||
self.model_call_details["_accounting_prior_attempt_ids"] = (
|
||||
*self.model_call_details.get("_accounting_prior_attempt_ids", ()),
|
||||
identity,
|
||||
)
|
||||
if isinstance(captured, CapturedUsage):
|
||||
self.model_call_details["_accounting_prior_captures"] = (
|
||||
*self.model_call_details.get("_accounting_prior_captures", ()),
|
||||
captured,
|
||||
)
|
||||
|
||||
def record_continued_partial(self, usage: Usage, response_cost: float, response_id: str) -> None:
|
||||
"""Bill what a cut stream already consumed on this continuation's success log, under the id the caller saw."""
|
||||
|
|
|
|||
|
|
@ -22,7 +22,11 @@ from litellm.litellm_core_utils.core_helpers import (
|
|||
get_metadata_variable_name_from_kwargs,
|
||||
get_parent_otel_span_from_kwargs,
|
||||
)
|
||||
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup, UsageObservation
|
||||
from litellm.litellm_core_utils.litellm_logging import (
|
||||
StandardLoggingPayloadSetup,
|
||||
UsageCaptureObservation,
|
||||
UsageObservation,
|
||||
)
|
||||
from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import guardrail_information_cost
|
||||
from litellm.litellm_core_utils.redact_messages import MESSAGE_REDACTION_ENABLE_HEADERS
|
||||
from litellm.llms.azure.ptu_capacity import is_str_keyed_mapping
|
||||
|
|
@ -346,6 +350,10 @@ def _attempt_spend_row_kwargs(
|
|||
"combined_usage_object": attempt.usage,
|
||||
"response_cost": attempt.response_cost,
|
||||
"_accounting_usage_observed": UsageObservation.ESTIMATED,
|
||||
"_accounting_usage_snapshot": UsageCaptureObservation.from_usage(
|
||||
attempt.usage, UsageObservation.ESTIMATED, attempt.response_cost
|
||||
),
|
||||
"_accounting_prior_captures": (),
|
||||
}
|
||||
row.pop("billed_failed_attempts", None)
|
||||
row.pop("_accounting_capture", None)
|
||||
|
|
|
|||
|
|
@ -472,7 +472,6 @@ from litellm.proxy.common_request_processing import ( # noqa: F401, RUF100 # l
|
|||
log_llm_api_exception,
|
||||
open_sse_before_first_byte,
|
||||
proxy_exception_from_route_error,
|
||||
request_litellm_call_id,
|
||||
resolve_litellm_call_id,
|
||||
should_return_raw_model_name,
|
||||
sse_error_payload,
|
||||
|
|
@ -13034,8 +13033,10 @@ async def completion(
|
|||
"""
|
||||
global user_temperature, user_request_timeout, user_max_tokens, user_api_base
|
||||
data = {}
|
||||
base_llm_response_processor: Final = ProxyBaseLLMRequestProcessing(data=data)
|
||||
try:
|
||||
data = await read_request_body(request=request) # rebind-ok: pre-existing rebinding on a rename-only line
|
||||
base_llm_response_processor.data = data
|
||||
if user_api_key_dict is not None:
|
||||
if data.get("metadata") is None:
|
||||
data["metadata"] = {}
|
||||
|
|
@ -13049,7 +13050,6 @@ async def completion(
|
|||
data["metadata"]["user_api_key_org_alias"] = user_api_key_dict.organization_alias
|
||||
if hasattr(user_api_key_dict, "agent_id") and user_api_key_dict.agent_id is not None:
|
||||
data["metadata"]["agent_id"] = user_api_key_dict.agent_id
|
||||
base_llm_response_processor: Final = ProxyBaseLLMRequestProcessing(data=data)
|
||||
return await base_llm_response_processor.base_process_llm_request(
|
||||
request=request,
|
||||
fastapi_response=fastapi_response,
|
||||
|
|
@ -13151,9 +13151,12 @@ async def completion(
|
|||
return _response
|
||||
except Exception as e:
|
||||
await proxy_logging_obj.post_call_failure_hook(
|
||||
user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
original_exception=e,
|
||||
request_data=base_llm_response_processor.data,
|
||||
dispatched_call_type=base_llm_response_processor.route_type,
|
||||
)
|
||||
litellm_call_id: Final = request_litellm_call_id(data)
|
||||
litellm_call_id: Final = base_llm_response_processor.litellm_call_id
|
||||
log_llm_api_exception(e, litellm_call_id)
|
||||
if isinstance(e, ProxyException):
|
||||
raise with_litellm_call_id(e, litellm_call_id)
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ from collections.abc import Awaitable, Callable, Mapping
|
|||
from dataclasses import dataclass, replace
|
||||
from datetime import datetime, timezone
|
||||
from decimal import Decimal
|
||||
from itertools import chain
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
from pydantic import BaseModel, JsonValue, TypeAdapter
|
||||
|
|
@ -129,6 +130,7 @@ class UsageCaptureContext:
|
|||
provider_timestamps: str = "{}"
|
||||
usage_source: str = "measured"
|
||||
included_attempt_ids: tuple[str, ...] = ()
|
||||
related_attempt_ids: tuple[str, ...] = ()
|
||||
|
||||
|
||||
def _database_usage_write(event: InboundUsageEvent | OutboundUsageEvent, envelope: Mapping[str, object]) -> UsageWrite:
|
||||
|
|
@ -175,6 +177,7 @@ def usage_writes(
|
|||
"usage": usage,
|
||||
"usage_source": "unknown" if usage is None else context.usage_source,
|
||||
"included_attempt_ids": context.included_attempt_ids,
|
||||
"related_attempt_ids": context.related_attempt_ids,
|
||||
"pricing": _JSON_OBJECT.validate_json(context.pricing),
|
||||
"charge_eligible": context.charge_eligible,
|
||||
"request_count_eligible": context.request_count_eligible,
|
||||
|
|
@ -267,7 +270,12 @@ def capture_context(payload: SpendLogsPayload, kwargs: Mapping[str, object], res
|
|||
|
||||
from litellm.litellm_core_utils.hidden_params import get_hidden_params
|
||||
from litellm.litellm_core_utils.internal_call_metadata import is_background_response, is_unbilled_non_inference_call
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging, UsageObservation, mask_api_base_credentials
|
||||
from litellm.litellm_core_utils.litellm_logging import (
|
||||
Logging,
|
||||
UsageCaptureObservation,
|
||||
UsageObservation,
|
||||
mask_api_base_credentials,
|
||||
)
|
||||
from litellm.litellm_core_utils.llm_cost_calc.finalization import FinalizedCost
|
||||
from litellm.proxy.hooks.search_free_tier import SearchAllowanceAllocation
|
||||
from litellm.proxy.spend_tracking.spend_tracking_utils import get_request_model_access_groups
|
||||
|
|
@ -324,6 +332,7 @@ def capture_context(payload: SpendLogsPayload, kwargs: Mapping[str, object], res
|
|||
)
|
||||
allocation: Final = allocation_value if isinstance(allocation_value, SearchAllowanceAllocation) else None
|
||||
pricing_snapshot: Final = kwargs.get("_accounting_pricing_snapshot")
|
||||
observation: Final = kwargs.get("_accounting_usage_snapshot")
|
||||
captured_pricing: Final = {
|
||||
"model_map_information": (
|
||||
_JSON_OBJECT.validate_json(pricing_snapshot.model_map_information)
|
||||
|
|
@ -353,7 +362,9 @@ def capture_context(payload: SpendLogsPayload, kwargs: Mapping[str, object], res
|
|||
"service_tier": raw_response.get("service_tier"),
|
||||
"region_name": hidden.get("region_name"),
|
||||
},
|
||||
"reported_cost": standard.get(
|
||||
"reported_cost": observation.response_cost
|
||||
if isinstance(observation, UsageCaptureObservation)
|
||||
else standard.get(
|
||||
"response_cost", kwargs.get("response_cost", _fields(get_hidden_params(response)).get("response_cost"))
|
||||
),
|
||||
"adjustments": (
|
||||
|
|
@ -413,10 +424,9 @@ def capture_context(payload: SpendLogsPayload, kwargs: Mapping[str, object], res
|
|||
if kwargs.get("_accounting_usage_observed") is UsageObservation.ESTIMATED
|
||||
else "measured",
|
||||
included_attempt_ids=TypeAdapter(tuple[str, ...]).validate_python(
|
||||
kwargs.get(
|
||||
"_accounting_included_attempt_ids", () if retry else kwargs.get("_accounting_prior_attempt_ids", ())
|
||||
)
|
||||
kwargs.get("_accounting_included_attempt_ids", ())
|
||||
),
|
||||
related_attempt_ids=tuple(write.event.event_id for write in _prior_writes(kwargs)),
|
||||
provider_timestamps=_json(
|
||||
{
|
||||
key: raw_response.get(key)
|
||||
|
|
@ -463,7 +473,13 @@ class _CaptureBuilder:
|
|||
self.writes: tuple[UsageWrite, ...] = ()
|
||||
|
||||
def observe(self, payload: SpendLogsPayload, usage: str | None) -> None:
|
||||
parsed: Final = _JSON_OBJECT.validate_json(usage) if usage is not None else {}
|
||||
from litellm.litellm_core_utils.litellm_logging import UsageCaptureObservation, UsageObservation
|
||||
|
||||
snapshot: Final = self.kwargs.get("_accounting_usage_snapshot")
|
||||
if isinstance(snapshot, UsageCaptureError):
|
||||
raise snapshot
|
||||
effective: Final = snapshot.usage_json if isinstance(snapshot, UsageCaptureObservation) else usage
|
||||
parsed: Final = _JSON_OBJECT.validate_json(effective) if effective is not None else {}
|
||||
batch: Final = _fields(parsed.get("batch_usage_by_model"))
|
||||
unknown: Final = (bool(batch) or bool(parsed.get("batch_unknown_usage_requests"))) and all(
|
||||
value is None for value in batch.values()
|
||||
|
|
@ -476,10 +492,32 @@ class _CaptureBuilder:
|
|||
}
|
||||
)
|
||||
if unknown
|
||||
else usage
|
||||
else effective
|
||||
)
|
||||
context: Final = capture_context(payload, self.kwargs, self.response)
|
||||
self.writes = usage_writes(payload, observed, replace(context, usage_source="unknown") if unknown else context)
|
||||
provenance: Final = (
|
||||
"unknown"
|
||||
if unknown
|
||||
else "estimated"
|
||||
if isinstance(snapshot, UsageCaptureObservation) and snapshot.provenance is UsageObservation.ESTIMATED
|
||||
else context.usage_source
|
||||
)
|
||||
self.writes = usage_writes(payload, observed, replace(context, usage_source=provenance))
|
||||
|
||||
|
||||
def _prior_writes(kwargs: Mapping[str, object]) -> tuple[UsageWrite, ...]:
|
||||
prior: Final = kwargs.get("_accounting_prior_captures", ())
|
||||
return (
|
||||
tuple(
|
||||
chain.from_iterable(
|
||||
item.writes
|
||||
for item in TypeAdapter(tuple[object, ...]).validate_python(prior)
|
||||
if isinstance(item, CapturedUsage)
|
||||
)
|
||||
)
|
||||
if isinstance(prior, tuple)
|
||||
else ()
|
||||
)
|
||||
|
||||
|
||||
def capture_callback(
|
||||
|
|
@ -499,7 +537,8 @@ def capture_callback(
|
|||
payload: Final = get_logging_payload(
|
||||
dict(kwargs), response, start_time, end_time, llm_router=get_llm_router(), capture_usage=builder.observe
|
||||
)
|
||||
return CapturedUsage(payload, builder.writes, meter.submit(builder.writes))
|
||||
dependencies: Final = _prior_writes(kwargs)
|
||||
return CapturedUsage(payload, builder.writes, meter.submit((*dependencies, *builder.writes)))
|
||||
|
||||
|
||||
def captured_callback_kwargs(
|
||||
|
|
@ -559,7 +598,7 @@ def capture_logging_usage(
|
|||
start_time: datetime,
|
||||
end_time: datetime,
|
||||
freeze_only: bool = False,
|
||||
) -> None:
|
||||
) -> CapturedUsage | UsageCaptureError | None:
|
||||
from litellm._internal_context import internal_parent_request_id, is_internal_call
|
||||
from litellm.batches.batch_utils import batch_cost_is_final
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
|
|
@ -572,10 +611,7 @@ def capture_logging_usage(
|
|||
return
|
||||
try:
|
||||
if not internal and isinstance(logging_obj.model_call_details.get("_accounting_capture"), CapturedUsage):
|
||||
logging_obj.record_usage_capture(
|
||||
capture_callback(meter, logging_obj.model_call_details, response, start_time, end_time)
|
||||
)
|
||||
return
|
||||
return capture_callback(meter, logging_obj.model_call_details, response, start_time, end_time)
|
||||
normalized_result: Final = logging_obj.usage_capture_response(response)
|
||||
prepared: Final = logging_obj.prepared_cost
|
||||
normalized: Final = (
|
||||
|
|
@ -616,8 +652,11 @@ def capture_logging_usage(
|
|||
)
|
||||
if not internal:
|
||||
logging_obj.record_usage_capture(captured)
|
||||
except Exception: # noqa: BLE001 # failed capture must not prevent normal logging and legacy charging
|
||||
except Exception as error: # noqa: BLE001 # failed capture must not prevent normal logging and legacy charging
|
||||
verbose_proxy_logger.exception("Usage capture failed at the execution boundary")
|
||||
return UsageCaptureError(str(error))
|
||||
else:
|
||||
return captured
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
|
|
|
|||
|
|
@ -1744,7 +1744,39 @@ def _failure_fields_to_lift(
|
|||
if _prompt_estimate is not None
|
||||
else ()
|
||||
)
|
||||
from litellm.litellm_core_utils.litellm_logging import UsageCaptureObservation
|
||||
|
||||
_recorded_observation: Final[object] = _model_call_details.get("_accounting_usage_observed")
|
||||
_current_usage: Final = (
|
||||
_current_failure_usage_to_lift(
|
||||
_model_call_details,
|
||||
request_data,
|
||||
_model_call_details.get("current_attempt_api_call_start_time", _first_handoff) is not None,
|
||||
_billed,
|
||||
_deployment_id,
|
||||
_base_model,
|
||||
dispatched_call_type,
|
||||
)
|
||||
if Logging.usage_capture_callback is not None
|
||||
and (_billed or _model_call_details.get("combined_usage_object") is not None)
|
||||
else None
|
||||
)
|
||||
_current_observation: Final = (
|
||||
_recorded_observation
|
||||
if isinstance(_recorded_observation, UsageObservation)
|
||||
else UsageObservation.UNKNOWN
|
||||
if _current_usage is None
|
||||
else UsageObservation.KNOWN
|
||||
if _model_call_details.get("combined_usage_object") is not None
|
||||
else UsageObservation.ESTIMATED
|
||||
)
|
||||
_capture_observation: Final = UsageCaptureObservation.from_usage(
|
||||
_current_usage[0] if _current_usage is not None and isinstance(_current_usage[0], Usage) else None,
|
||||
_current_observation,
|
||||
float(_current_usage[1])
|
||||
if _current_usage is not None and isinstance(_current_usage[1], (int, float))
|
||||
else None,
|
||||
)
|
||||
_entries: Final = (
|
||||
("start_time", _model_call_details.get("start_time")),
|
||||
("first_api_call_start_time", _first_handoff),
|
||||
|
|
@ -1760,22 +1792,9 @@ def _failure_fields_to_lift(
|
|||
"_accounting_allocation",
|
||||
_logging_obj.cost_allowance_allocation if isinstance(_logging_obj, Logging) else None,
|
||||
),
|
||||
(
|
||||
"_accounting_usage_observed",
|
||||
(
|
||||
None
|
||||
if Logging.usage_capture_callback is None
|
||||
else _recorded_observation
|
||||
if isinstance(_recorded_observation, UsageObservation)
|
||||
else UsageObservation.UNKNOWN
|
||||
if _usage_to_lift is None
|
||||
else UsageObservation.KNOWN
|
||||
if _model_call_details.get("combined_usage_object") is not None
|
||||
or _model_call_details.get("failed_attempts_usage")
|
||||
else UsageObservation.ESTIMATED
|
||||
),
|
||||
),
|
||||
("_accounting_included_attempt_ids", _model_call_details.get("_accounting_prior_attempt_ids")),
|
||||
("_accounting_usage_observed", _current_observation if Logging.usage_capture_callback is not None else None),
|
||||
("_accounting_usage_snapshot", _capture_observation if Logging.usage_capture_callback is not None else None),
|
||||
("_accounting_prior_captures", _model_call_details.get("_accounting_prior_captures")),
|
||||
)
|
||||
return MappingProxyType({key: value for key, value in _entries if value is not None})
|
||||
|
||||
|
|
|
|||
|
|
@ -133,7 +133,10 @@ class StoredUsage:
|
|||
existing: Final = {row.event.event_id: row for row in self.rows}
|
||||
if any(existing.get(row.event.event_id, row) != row for row in writes):
|
||||
raise AccountingConflictError("Same identity has different content")
|
||||
self.rows += tuple(row for row in writes if row.event.event_id not in existing)
|
||||
unique: Final = {row.event.event_id: row for row in writes}
|
||||
if any(unique[row.event.event_id] != row for row in writes):
|
||||
raise AccountingConflictError("Submission contains conflicting content")
|
||||
self.rows += tuple(row for identity, row in unique.items() if identity not in existing)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -537,7 +540,8 @@ async def test_failed_attempt_keeps_its_usage_without_claiming_the_callers_succe
|
|||
assert failed["pricing"]["reported_cost"] == 0.4
|
||||
assert all(json.loads(row.envelope)["included_attempt_ids"] == [] for row in database.rows[:2])
|
||||
successful: Final = json.loads(database.rows[-1].envelope)
|
||||
assert successful["included_attempt_ids"] == ["caller-attempt-0", "caller-attempt-1"]
|
||||
assert successful["related_attempt_ids"] == ["caller-attempt-0", "caller-attempt-1"]
|
||||
assert successful["included_attempt_ids"] == []
|
||||
assert successful["usage"]["total_tokens"] == 3
|
||||
assert successful["charge_eligible"] is True
|
||||
assert database.rows[2].event.outcome is Outcome.SUCCEEDED
|
||||
|
|
@ -1146,3 +1150,155 @@ async def test_prepared_native_google_capture_preserves_reported_usage(
|
|||
) == tokens
|
||||
assert envelope["usage_source"] == ("unknown" if tokens is None else "measured")
|
||||
assert response == native_response.json()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("rejected", [False, True])
|
||||
async def test_retry_links_only_admitted_usage_and_preserves_later_measurement(rejected: bool) -> None:
|
||||
from litellm.litellm_core_utils.litellm_logging import UsageObservation
|
||||
from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime
|
||||
from litellm.types.utils import ModelResponse, Usage
|
||||
|
||||
database: Final = StoredUsage()
|
||||
runtime: Final = UsageCaptureRuntime()
|
||||
runtime.start(database.insert, UsageCaptureSettings(max_pending_records=4))
|
||||
logger: Final = _logger()
|
||||
assert runtime.meter is not None
|
||||
try:
|
||||
if rejected:
|
||||
runtime.meter.submit(tuple(_write(f"occupied-{index}") for index in range(4)))
|
||||
logger.model_call_details["_accounting_usage_observed"] = UsageObservation.UNKNOWN
|
||||
logger.record_partial_usage_for_failure(Usage(prompt_tokens=5, total_tokens=5), 0.4)
|
||||
logger.carry_failed_attempt_usage()
|
||||
await runtime.meter.flush()
|
||||
logger.capture_usage_before_logging(ModelResponse(id="success", usage=Usage(prompt_tokens=3)), NOW, NOW)
|
||||
finally:
|
||||
await runtime.stop(1)
|
||||
successful: Final = json.loads(next(row.envelope for row in database.rows if row.event.event_id == "success"))
|
||||
assert successful.get("related_attempt_ids", successful["included_attempt_ids"]) == (
|
||||
[] if rejected else ["caller-attempt-0"]
|
||||
)
|
||||
assert successful["usage"]["prompt_tokens"] == 3
|
||||
if not rejected:
|
||||
failed: Final = json.loads(
|
||||
next(row.envelope for row in database.rows if row.event.event_id == "caller-attempt-0")
|
||||
)
|
||||
assert failed["usage"]["prompt_tokens"] == 5
|
||||
assert failed["usage_source"] == "measured"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("current_tokens", [None, 0, 3])
|
||||
async def test_final_failure_captures_current_attempt_without_legacy_retry_aggregate(
|
||||
current_tokens: int | None,
|
||||
) -> None:
|
||||
from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime, capture_callback
|
||||
from litellm.proxy.utils import _failure_fields_to_lift
|
||||
from litellm.types.utils import Usage
|
||||
|
||||
database: Final = StoredUsage()
|
||||
runtime: Final = UsageCaptureRuntime()
|
||||
runtime.start(database.insert, UsageCaptureSettings())
|
||||
logger: Final = _logger()
|
||||
try:
|
||||
logger.record_partial_usage_for_failure(Usage(prompt_tokens=5, total_tokens=5), 0.4)
|
||||
logger.carry_failed_attempt_usage()
|
||||
if current_tokens is not None:
|
||||
logger.record_partial_usage_for_failure(
|
||||
Usage(prompt_tokens=current_tokens, total_tokens=current_tokens), 0.2
|
||||
)
|
||||
lifted: Final = _failure_fields_to_lift({"litellm_logging_obj": logger}, provider_billed=True)
|
||||
assert lifted["combined_usage_object"].total_tokens == 5 + (current_tokens or 0)
|
||||
assert runtime.meter is not None
|
||||
capture_callback(
|
||||
runtime.meter,
|
||||
{**logger.model_call_details, **lifted, "log_event_type": "failed_api_call"},
|
||||
None,
|
||||
NOW,
|
||||
NOW,
|
||||
)
|
||||
finally:
|
||||
await runtime.stop(1)
|
||||
final: Final = json.loads(next(row.envelope for row in database.rows if row.event.event_id == "caller"))
|
||||
assert (None if final["usage"] is None else final["usage"]["total_tokens"]) == current_tokens
|
||||
assert final["usage_source"] == ("unknown" if current_tokens is None else "measured")
|
||||
assert final["pricing"]["reported_cost"] == (None if current_tokens is None else 0.2)
|
||||
assert final["included_attempt_ids"] == []
|
||||
assert final["related_attempt_ids"] == ["caller-attempt-0"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_estimated_final_attempt_keeps_its_provenance_after_measured_retry() -> None:
|
||||
from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime, capture_callback
|
||||
from litellm.proxy.utils import _failure_fields_to_lift
|
||||
from litellm.types.utils import Usage
|
||||
|
||||
database: Final = StoredUsage()
|
||||
runtime: Final = UsageCaptureRuntime()
|
||||
runtime.start(database.insert, UsageCaptureSettings())
|
||||
logger: Final = _logger()
|
||||
logger.messages = [{"role": "user", "content": "hello"}]
|
||||
logger.model_call_details["messages"] = logger.messages
|
||||
try:
|
||||
logger.record_partial_usage_for_failure(Usage(prompt_tokens=5, total_tokens=5), 0.4)
|
||||
logger.carry_failed_attempt_usage()
|
||||
logger.model_call_details["current_attempt_api_call_start_time"] = NOW
|
||||
lifted: Final = _failure_fields_to_lift({"litellm_logging_obj": logger}, provider_billed=True)
|
||||
assert runtime.meter is not None
|
||||
capture_callback(runtime.meter, {**logger.model_call_details, **lifted}, None, NOW, NOW)
|
||||
finally:
|
||||
await runtime.stop(1)
|
||||
final: Final = json.loads(next(row.envelope for row in database.rows if row.event.event_id == "caller"))
|
||||
assert final["usage_source"] == "estimated"
|
||||
assert final["usage"]["total_tokens"] > 0
|
||||
assert lifted["combined_usage_object"].total_tokens == final["usage"]["total_tokens"] + 5
|
||||
assert final["pricing"]["reported_cost"] == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_parent_receipt_cannot_acknowledge_when_a_linked_attempt_conflicts() -> None:
|
||||
from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime
|
||||
from litellm.types.utils import ModelResponse, Usage
|
||||
|
||||
database: Final = StoredUsage()
|
||||
runtime: Final = UsageCaptureRuntime()
|
||||
runtime.start(database.insert, UsageCaptureSettings())
|
||||
logger: Final = _logger()
|
||||
try:
|
||||
logger.record_partial_usage_for_failure(Usage(prompt_tokens=5), 0.4)
|
||||
logger.carry_failed_attempt_usage()
|
||||
assert runtime.meter is not None
|
||||
await runtime.meter.flush()
|
||||
other: Final = _logger()
|
||||
other.record_partial_usage_for_failure(Usage(prompt_tokens=8), 0.8)
|
||||
other.carry_failed_attempt_usage()
|
||||
other.capture_usage_before_logging(ModelResponse(id="success", usage=Usage(prompt_tokens=3)), NOW, NOW)
|
||||
captured: Final = other.model_call_details["_accounting_capture"]
|
||||
await runtime.meter.flush()
|
||||
with pytest.raises(UsageCaptureError, match="identity"):
|
||||
await captured.receipt
|
||||
finally:
|
||||
await runtime.stop(1)
|
||||
assert tuple(row.event.event_id for row in database.rows) == ("caller-attempt-0",)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unserializable_observation_does_not_interrupt_retry_or_legacy_usage() -> None:
|
||||
from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime
|
||||
from litellm.types.utils import ModelResponse, Usage
|
||||
|
||||
database: Final = StoredUsage()
|
||||
runtime: Final = UsageCaptureRuntime()
|
||||
runtime.start(database.insert, UsageCaptureSettings())
|
||||
logger: Final = _logger()
|
||||
usage: Final = Usage(prompt_tokens=5, total_tokens=5, opaque_usage=object())
|
||||
try:
|
||||
logger.record_partial_usage_for_failure(usage, 0.4)
|
||||
assert logger.model_call_details["combined_usage_object"] is usage
|
||||
logger.carry_failed_attempt_usage()
|
||||
assert logger.model_call_details["failed_attempts_usage"] == ((usage, 0.4),)
|
||||
logger.capture_usage_before_logging(ModelResponse(id="success", usage=Usage(prompt_tokens=3)), NOW, NOW)
|
||||
finally:
|
||||
await runtime.stop(1)
|
||||
assert tuple(row.event.event_id for row in database.rows) == ("caller:inbound", "success")
|
||||
assert json.loads(database.rows[-1].envelope)["related_attempt_ids"] == []
|
||||
|
|
|
|||
|
|
@ -4465,3 +4465,79 @@ async def test_sync_ack_preserves_async_request_support(
|
|||
else:
|
||||
assert "hi" in response.text
|
||||
assert not persisted.empty(), "Terminal usage must still be captured"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("sync_ack", [False, True])
|
||||
async def test_text_completion_provider_rejection_captures_unknown_usage(
|
||||
sync_ack: bool, monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter
|
||||
) -> None:
|
||||
from fastapi import Request
|
||||
|
||||
from litellm.accounting import Outcome
|
||||
from litellm.proxy import proxy_server
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.db.accounting_ledger import UsageWrite
|
||||
from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger
|
||||
from litellm.proxy.spend_tracking.usage_capture import UsageCaptureRuntime
|
||||
from litellm.proxy.spend_tracking.usage_meter import UsageCaptureSettings
|
||||
|
||||
persisted: Final[asyncio.Queue[UsageWrite]] = asyncio.Queue()
|
||||
|
||||
async def persist(writes: tuple[UsageWrite, ...]) -> None:
|
||||
for write in writes:
|
||||
persisted.put_nowait(write)
|
||||
|
||||
async def authorize(request: Request) -> UserAPIKeyAuth:
|
||||
return UserAPIKeyAuth(api_key="test-key", request_route=request.url.path)
|
||||
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
monkeypatch.setenv("OPENAI_API_KEY", "test-provider-key")
|
||||
respx_mock.post().respond(
|
||||
404,
|
||||
json={
|
||||
"error": {"message": "The model is unavailable", "type": "invalid_request_error", "code": "model_not_found"}
|
||||
},
|
||||
)
|
||||
runtime: Final = UsageCaptureRuntime()
|
||||
monkeypatch.setattr(proxy_server, "usage_capture", runtime)
|
||||
monkeypatch.setattr(proxy_server, "general_settings", {"accounting_sync_ack": sync_ack})
|
||||
monkeypatch.setattr(proxy_server, "prisma_client", None)
|
||||
monkeypatch.setattr(
|
||||
proxy_server,
|
||||
"llm_router",
|
||||
litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "instruct",
|
||||
"litellm_params": {"model": "openai/gpt-3.5-turbo-instruct", "api_key": "test-provider-key"},
|
||||
}
|
||||
],
|
||||
num_retries=0,
|
||||
),
|
||||
)
|
||||
monkeypatch.setattr(proxy_server, "proxy_logging_obj", ProxyLogging(UserApiKeyCache()))
|
||||
monkeypatch.setattr(litellm, "callbacks", [ProxyDBLogger()])
|
||||
monkeypatch.setitem(proxy_server.app.dependency_overrides, proxy_server.user_api_key_auth, authorize)
|
||||
runtime.start(persist, UsageCaptureSettings())
|
||||
try:
|
||||
async with httpx.AsyncClient(
|
||||
transport=httpx.ASGITransport(app=proxy_server.app), base_url="http://test"
|
||||
) as client:
|
||||
response: Final = await client.post(
|
||||
"/v1/completions", json={"model": "instruct", "prompt": "hi", "max_tokens": 1}
|
||||
)
|
||||
finally:
|
||||
await runtime.stop(1)
|
||||
assert response.status_code == 404, response.text
|
||||
assert len(respx_mock.calls) == 1
|
||||
assert persisted.qsize() == 2
|
||||
inbound: Final = persisted.get_nowait()
|
||||
outbound: Final = persisted.get_nowait()
|
||||
assert outbound.event.outcome is Outcome.FAILED
|
||||
assert outbound.event.call_type.value == "atext_completion"
|
||||
assert outbound.event.metrics is None
|
||||
assert inbound.event.request_id == outbound.event.request_id
|
||||
envelope: Final = json.loads(outbound.envelope)
|
||||
assert envelope["usage"] is None
|
||||
assert envelope["usage_source"] == "unknown"
|
||||
assert envelope["pricing"]["reported_cost"] is None
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue