mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-27 06:13:07 +00:00
fix(responses): bill every round of a streamed MCP auto-execute on one spend row (#47445)
* test(responses): cover billing every round of a streamed MCP auto-execute * fix(responses): bill every round of a streamed MCP auto-execute on one spend row * fix(responses): bill every round of a streamed MCP auto-execute on bridged providers A streamed /v1/responses MCP auto-execute on a provider bridged through chat completions (for example Anthropic) billed only round 1, because each round's chat stream dispatched on the shared logging object and the first dispatch won. Bridged rounds now go through the same one-row billing as native rounds. A round is billed as soon as its chat stream ends, before its closing events reach the client. The row carries the final response, usage and cost summed over every round, and the fixed margin once. A round served from the response cache is billed at zero. A follow-up chat stream that breaks bills the paid rounds on its failure row. A single-round stream logs exactly as before. * fix(responses): charge one fixed margin and keep bridged usage fields on streamed MCP rounds A follow-up round that fails after earlier rounds were billed now charges the per-request fixed margin once, and its failure row carries no cost breakdown that would disagree with the summed cost. A bridged multi-round stream logs a plain cache_hit flag like a single-round stream does, and its usage keeps the summed provider usage fields of every round. A new test covers a stream dropped after its event loop closed, which must release its parked log without an unawaited coroutine warning. * fix(responses): price each streamed MCP round on its own and bill every round when a guardrail blocks Each round of a streamed MCP auto-execute is now priced from its own uncached cost, so a provider that states its cost (xAI) charges the stated cost of every round plus the fixed margin once, and the cost breakdown always agrees with the billed cost. Summed usage goes through a JSON-safe dump, so a nested provider field such as xAI's context_details no longer breaks a failed stream's response.failed event. A post-call guardrail that blocks a multi-round bridged stream now drops the parked success log and bills every round on the failure row, and a failed follow-up round's failure row carries the summed breakdown of both rounds. * fix(responses): keep earlier MCP rounds on their success row when a later round fails When a later round of a streamed MCP auto-execute fails, the rounds before it are now logged on their success row, and the failure row bills only the failed round with no second fixed margin. That failure log runs on its own copy of the request's logging state, so its usage and cost never leak into the earlier rounds' row. A post-call guardrail that blocks a native multi-round stream keeps the summed success row beside the $0 failure row, as at the merge base, while a bridged stream still bills every round on its failure row. * fix(responses): keep a failed later MCP round on its own failure row A streamed MCP auto-execute whose later round failed wrote only the earlier rounds' success row. The failure was logged on a copy of the logging state, while the proxy builds its failure row from the request's own state, which by then held the success payload, so the failure row took the success row's id and was dropped. Now the failure is logged on the request's own state as at the merge base, and once a later round starts, the success row for the rounds before it is logged from a copy of the state taken at that moment, so the two rows never mix. * fix(proxy): let an MCP tool call reuse its request's max_parallel_requests slot * test(proxy): create the other-request task in a way Python 3.10 supports and mark its bindings Final * fix(proxy): keep a reused max_parallel slot until the request's last round An MCP tool call that reuses its request's max_parallel_requests slot no longer loses it to the first round's late success callback, so the key stays counted for the rounds still to run. The reused call also stops replacing the request's rate limit headers. * fix(responses): keep an unpriced streamed MCP round unpriced instead of billing it as free * fix(responses): log a bridged MCP stream as a chat completion with every round's usage
This commit is contained in:
Mateo Wang
GitHub
parent
647ebd63e0
commit
0e24bbb8a2
14 changed files with 1848 additions and 64 deletions
|
|
@ -12,7 +12,7 @@ import subprocess
|
|||
import sys
|
||||
import time
|
||||
import traceback
|
||||
from collections.abc import Awaitable, Callable, Iterable, Iterator, Mapping, Sequence
|
||||
from collections.abc import Awaitable, Callable, Coroutine, Iterable, Iterator, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime as dt_object
|
||||
from functools import lru_cache
|
||||
|
|
@ -575,6 +575,27 @@ class UncachedPrice:
|
|||
cost_breakdown: CostBreakdown | None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class BilledRoundsSpend:
|
||||
usage: Usage
|
||||
cost: float
|
||||
cost_breakdown: CostBreakdown | None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class BilledRoundsDispatch:
|
||||
"""A parked success log for a stream bridged through chat completions that billed several rounds, and the spend
|
||||
its failure log bills when a post-call hook rejects the response, as it does for a rejected one-round bridged
|
||||
stream."""
|
||||
|
||||
dispatch: Coroutine[object, object, None]
|
||||
spend: BilledRoundsSpend
|
||||
|
||||
|
||||
def _fixed_margin(cost_breakdown: CostBreakdown | None) -> float:
|
||||
return 0.0 if cost_breakdown is None else cost_breakdown.get("margin_fixed_amount", 0.0)
|
||||
|
||||
|
||||
def _timestamp_seconds(moment: object) -> float | None:
|
||||
if isinstance(moment, datetime.datetime):
|
||||
return moment.timestamp()
|
||||
|
|
@ -2421,6 +2442,31 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
self.model_call_details["combined_usage_object"] = usage
|
||||
self.model_call_details["response_cost"] = response_cost
|
||||
|
||||
def record_failure_after_earlier_rounds(
|
||||
self, earlier: BilledRoundsSpend, failed_usage: Usage, failed_cost: float
|
||||
) -> None:
|
||||
"""Bill a failed round on the failure log after the rounds of its request before it were billed on the success
|
||||
log.
|
||||
|
||||
The fixed margin is a per-request fee, so it is charged once: the failed round's margin is dropped when the
|
||||
earlier rounds already paid one."""
|
||||
from litellm.responses.mcp.round_billing import breakdown_less_repeated_fixed_margin
|
||||
|
||||
failed_breakdown: Final = None if self.cost_breakdown is earlier.cost_breakdown else self.cost_breakdown
|
||||
repeated_fixed_margin: Final = (
|
||||
0.0
|
||||
if failed_breakdown is None
|
||||
else min(_fixed_margin(earlier.cost_breakdown), _fixed_margin(failed_breakdown), earlier.cost, failed_cost)
|
||||
)
|
||||
self.model_call_details["combined_usage_object"] = failed_usage
|
||||
self.model_call_details["response_cost"] = failed_cost - repeated_fixed_margin
|
||||
self.cost_breakdown = breakdown_less_repeated_fixed_margin(failed_breakdown, repeated_fixed_margin)
|
||||
|
||||
def record_billed_rounds_for_failure(self, spend: BilledRoundsSpend) -> None:
|
||||
"""Bill every round of a fully streamed request on the failure log when a post-call hook rejects it."""
|
||||
self.record_partial_usage_for_failure(spend.usage, spend.cost)
|
||||
self.cost_breakdown = spend.cost_breakdown
|
||||
|
||||
def attempt_scoped_copy(self) -> "Logging":
|
||||
"""A copy with its own model_call_details, so logging one attempt of a request through it leaves
|
||||
the state the request's other attempts share untouched."""
|
||||
|
|
|
|||
|
|
@ -10,7 +10,7 @@ import traceback
|
|||
from collections.abc import AsyncIterator, Callable, Iterable, Iterator, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from types import MappingProxyType
|
||||
from typing import Any, Final, NoReturn, Protocol, TypeVar, cast
|
||||
from typing import TYPE_CHECKING, Any, Final, NoReturn, Protocol, TypeVar, cast
|
||||
|
||||
import anyio
|
||||
import httpx
|
||||
|
|
@ -60,6 +60,9 @@ from .exception_mapping_utils import exception_type, with_deployment_api_base
|
|||
from .llm_response_utils.get_api_base import get_api_base
|
||||
from .rules import Rules
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .litellm_logging import BilledRoundsSpend
|
||||
|
||||
# Constants for special delta attribute names
|
||||
AUDIO_ATTRIBUTE: Final = "audio"
|
||||
IMAGE_ATTRIBUTE: Final = "images"
|
||||
|
|
@ -259,6 +262,8 @@ def _provider_hidden_params(
|
|||
|
||||
|
||||
class CustomStreamWrapper:
|
||||
_earlier_rounds_spend: "BilledRoundsSpend | None" = None
|
||||
|
||||
@property
|
||||
def hidden_params(self) -> dict[str, object]: # mutable-ok: API requires mutation
|
||||
return self._hidden_params
|
||||
|
|
@ -2399,6 +2404,9 @@ class CustomStreamWrapper:
|
|||
return
|
||||
await asyncio.wait((self._failure_logging,), timeout=timeout)
|
||||
|
||||
def bill_failure_after_earlier_rounds(self, earlier: "BilledRoundsSpend") -> None:
|
||||
self._earlier_rounds_spend = earlier
|
||||
|
||||
def _record_partial_usage_for_failure(self) -> None:
|
||||
"""
|
||||
A stream that breaks mid-flight still billed the provider for the chunks
|
||||
|
|
@ -2425,10 +2433,12 @@ class CustomStreamWrapper:
|
|||
if self.model:
|
||||
partial_response.model = self.model
|
||||
backfill_missing_cache_usage_fields(usage)
|
||||
self.logging_obj.model_call_details["combined_usage_object"] = usage
|
||||
self.logging_obj.model_call_details["response_cost"] = (
|
||||
self.logging_obj.response_cost_calculator(result=partial_response) or 0.0
|
||||
)
|
||||
partial_cost: Final = self.logging_obj.response_cost_calculator(result=partial_response) or 0.0
|
||||
if self._earlier_rounds_spend is None:
|
||||
self.logging_obj.model_call_details["combined_usage_object"] = usage
|
||||
self.logging_obj.model_call_details["response_cost"] = partial_cost
|
||||
return
|
||||
self.logging_obj.record_failure_after_earlier_rounds(self._earlier_rounds_spend, usage, partial_cost)
|
||||
except Exception as recover_error:
|
||||
verbose_logger.debug(
|
||||
"could not recover partial usage for interrupted stream: %s",
|
||||
|
|
|
|||
|
|
@ -67,6 +67,7 @@ from litellm.litellm_core_utils.get_supported_openai_params import (
|
|||
get_supported_openai_params,
|
||||
)
|
||||
from litellm.litellm_core_utils.internal_call_metadata import is_unbilled_non_inference_call_from_params
|
||||
from litellm.litellm_core_utils.litellm_logging import BilledRoundsDispatch
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import guardrail_information_cost
|
||||
from litellm.litellm_core_utils.llm_response_utils.get_headers import (
|
||||
|
|
@ -3996,6 +3997,8 @@ class ProxyBaseLLMRequestProcessing:
|
|||
|
||||
async def _on_deferred_native_stream_complete(*args: object) -> None:
|
||||
match args:
|
||||
case (BilledRoundsDispatch(dispatch=logging_coroutine),):
|
||||
GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=logging_coroutine)
|
||||
case (logging_coroutine,) if asyncio.iscoroutine(logging_coroutine):
|
||||
GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=logging_coroutine)
|
||||
case (assembled_response, cache_hit):
|
||||
|
|
|
|||
|
|
@ -695,6 +695,7 @@ class RequestRateLimiterStash:
|
|||
owner_litellm_call_id: str | None = None
|
||||
rate_limit_response: RateLimitResponse | None = None
|
||||
parallel_slot: ParallelSlotAcquisition | None = None
|
||||
parallel_slot_reused_at: datetime | None = None
|
||||
parallel_slot_release_lock: asyncio.Lock = field(default_factory=asyncio.Lock, repr=False, compare=False)
|
||||
reserved_tokens: int = 0
|
||||
reserved_model: RateLimitedModel | None = None
|
||||
|
|
@ -791,6 +792,19 @@ def claim_request_stash_for_data(data: dict) -> RequestRateLimiterStash:
|
|||
return stash
|
||||
|
||||
|
||||
def _ended_before_slot_reuse(ended_at: datetime | None, reused_at: datetime | None) -> bool:
|
||||
if ended_at is None or reused_at is None or (ended_at.tzinfo is None) != (reused_at.tzinfo is None):
|
||||
return False
|
||||
return ended_at <= reused_at
|
||||
|
||||
|
||||
def _callback_kwargs_mapping(kwargs: object) -> Mapping[str, object]:
|
||||
try:
|
||||
return _REQUEST_RATE_LIMIT_DATA.validate_python(kwargs)
|
||||
except ValidationError:
|
||||
return MappingProxyType({})
|
||||
|
||||
|
||||
def get_request_stash_for_call(litellm_call_id: str | None) -> RequestRateLimiterStash | None:
|
||||
stash: Final = _request_stash.get()
|
||||
if stash is None:
|
||||
|
|
@ -1652,6 +1666,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
read_only: bool = False,
|
||||
skip_tpm_check: bool = False,
|
||||
parallel_slot_id: str | None = None,
|
||||
skip_parallel_check: bool = False,
|
||||
) -> RateLimitResponse:
|
||||
"""
|
||||
Check if any of the rate limit descriptors should be rate limited.
|
||||
|
|
@ -1668,6 +1683,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
``reserve_tpm_tokens`` reservation path should set this to
|
||||
avoid the +1-per-key Lua / in-memory increment double-charging
|
||||
the tokens counter.
|
||||
skip_parallel_check: If True, ``max_parallel_requests`` gauges are
|
||||
neither checked nor acquired, for a call that runs inside a
|
||||
request already holding its slot.
|
||||
|
||||
``max_parallel_requests`` descriptors are enforced by the dedicated
|
||||
concurrency-gauge path (``_check_parallel_request_gauges``), never by
|
||||
|
|
@ -1760,7 +1778,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
if windowed_response["overall_code"] == "OVER_LIMIT":
|
||||
return windowed_response
|
||||
|
||||
if not gauges:
|
||||
if not gauges or skip_parallel_check:
|
||||
return windowed_response
|
||||
|
||||
gauge_response: Final = await self._check_parallel_request_gauges(
|
||||
|
|
@ -2023,12 +2041,13 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
parent_otel_span: Span | None,
|
||||
*,
|
||||
in_logging_callback: bool = False,
|
||||
round_ended_at: datetime | None = None,
|
||||
) -> None:
|
||||
if stash is None:
|
||||
return
|
||||
async with stash.parallel_slot_release_lock:
|
||||
acquisition: Final = stash.parallel_slot
|
||||
if acquisition is None:
|
||||
if acquisition is None or _ended_before_slot_reuse(round_ended_at, stash.parallel_slot_reused_at):
|
||||
return
|
||||
deferred: Final = in_logging_callback and await self._defer_parallel_slot_release(
|
||||
acquisition, parent_otel_span
|
||||
|
|
@ -3947,6 +3966,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
tpm_reservation_amount: int,
|
||||
call_type: str | None = None,
|
||||
endpoint_type: EndpointType = EndpointType.GENERIC,
|
||||
reuses_request_slot: bool = False,
|
||||
) -> None:
|
||||
"""
|
||||
Reserve project-scoped ITPM/OTPM tokens (Bedrock Mantle-style
|
||||
|
|
@ -4026,7 +4046,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
parent_otel_span=user_api_key_dict.parent_otel_span,
|
||||
)
|
||||
stash.reservation_released = True
|
||||
await self._release_stashed_parallel_slot(stash, user_api_key_dict.parent_otel_span)
|
||||
if not reuses_request_slot:
|
||||
await self._release_stashed_parallel_slot(stash, user_api_key_dict.parent_otel_span)
|
||||
self._handle_rate_limit_error(
|
||||
response=io_response,
|
||||
descriptors=descriptors,
|
||||
|
|
@ -4187,6 +4208,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
verbose_proxy_logger.debug("Inside Rate Limit Pre-Call Hook")
|
||||
|
||||
stash: Final = claim_request_stash_for_data(data)
|
||||
reuses_request_slot: Final = call_type == CallTypes.call_mcp_tool.value and stash.parallel_slot is not None
|
||||
if reuses_request_slot:
|
||||
stash.parallel_slot_reused_at = self._get_current_time()
|
||||
|
||||
#########################################################
|
||||
# Check if the call type has a specific rate limiter
|
||||
|
|
@ -4231,11 +4255,15 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
# then because _reserve_project_io_tokens_or_raise below charges
|
||||
# them unconditionally and counting them here too would
|
||||
# double-charge every request.
|
||||
parallel_counter_keys: Final = [
|
||||
self.create_rate_limit_keys(d["key"], d["value"], "max_parallel_requests")
|
||||
for d in descriptors
|
||||
if (d.get("rate_limit") or {}).get("max_parallel_requests") is not None
|
||||
]
|
||||
parallel_counter_keys: Final = (
|
||||
[]
|
||||
if reuses_request_slot
|
||||
else [
|
||||
self.create_rate_limit_keys(d["key"], d["value"], "max_parallel_requests")
|
||||
for d in descriptors
|
||||
if (d.get("rate_limit") or {}).get("max_parallel_requests") is not None
|
||||
]
|
||||
)
|
||||
parallel_slot_id: Final = uuid.uuid4().hex if parallel_counter_keys else None
|
||||
|
||||
first_pass_descriptors: Final = (
|
||||
|
|
@ -4250,6 +4278,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
parent_otel_span=user_api_key_dict.parent_otel_span,
|
||||
skip_tpm_check=self.tpm_reservation_enabled,
|
||||
parallel_slot_id=parallel_slot_id,
|
||||
skip_parallel_check=reuses_request_slot,
|
||||
)
|
||||
|
||||
if response["overall_code"] == "OVER_LIMIT":
|
||||
|
|
@ -4259,7 +4288,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
requested_model=requested_model,
|
||||
)
|
||||
else:
|
||||
stash.rate_limit_response = response
|
||||
if not reuses_request_slot:
|
||||
stash.rate_limit_response = response
|
||||
if parallel_slot_id is not None:
|
||||
stash.parallel_slot = ParallelSlotAcquisition(
|
||||
slot_id=parallel_slot_id,
|
||||
|
|
@ -4357,7 +4387,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
raw_estimate=estimated_tokens,
|
||||
)
|
||||
if stash.ptu_ceiling is not None and ptu_estimated_tokens > stash.ptu_ceiling.tpm_limit:
|
||||
await self._release_stashed_parallel_slot(stash, user_api_key_dict.parent_otel_span)
|
||||
if not reuses_request_slot:
|
||||
await self._release_stashed_parallel_slot(stash, user_api_key_dict.parent_otel_span)
|
||||
raise PTUShareExceededError(stash.ptu_ceiling, ptu_estimated_tokens)
|
||||
tpm_response: Final = await self.reserve_tpm_tokens(
|
||||
descriptors=descriptors,
|
||||
|
|
@ -4367,7 +4398,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
)
|
||||
|
||||
if tpm_response["overall_code"] == "OVER_LIMIT":
|
||||
await self._release_stashed_parallel_slot(stash, user_api_key_dict.parent_otel_span)
|
||||
if not reuses_request_slot:
|
||||
await self._release_stashed_parallel_slot(stash, user_api_key_dict.parent_otel_span)
|
||||
self._handle_rate_limit_error(
|
||||
response=tpm_response,
|
||||
descriptors=descriptors,
|
||||
|
|
@ -4418,6 +4450,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
tpm_reservation_amount=tpm_reservation_amount,
|
||||
call_type=call_type,
|
||||
endpoint_type=endpoint_type,
|
||||
reuses_request_slot=reuses_request_slot,
|
||||
)
|
||||
|
||||
def _create_pipeline_operations(
|
||||
|
|
@ -5332,8 +5365,16 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
try:
|
||||
verbose_proxy_logger.debug("INSIDE parallel request limiter ASYNC SUCCESS LOGGING")
|
||||
|
||||
stash: Final = get_request_stash_for_call(_call_id_from_callback_kwargs(kwargs))
|
||||
await self._release_stashed_parallel_slot(stash, litellm_parent_otel_span, in_logging_callback=True)
|
||||
callback_kwargs: Final = _callback_kwargs_mapping(kwargs)
|
||||
stash: Final = get_request_stash_for_call(_call_id_from_callback_kwargs(callback_kwargs))
|
||||
await self._release_stashed_parallel_slot(
|
||||
stash,
|
||||
litellm_parent_otel_span,
|
||||
in_logging_callback=True,
|
||||
round_ended_at=end_time
|
||||
if isinstance(end_time, datetime) and not callback_kwargs.get("stream")
|
||||
else None,
|
||||
)
|
||||
|
||||
pipeline_operations: Final = self._build_success_event_pipeline_operations(
|
||||
kwargs=kwargs,
|
||||
|
|
|
|||
|
|
@ -158,7 +158,7 @@ from litellm.litellm_core_utils.core_helpers import (
|
|||
)
|
||||
from litellm.litellm_core_utils.duration_parser import get_budget_window_start
|
||||
from litellm.litellm_core_utils.env_utils import get_env_float, get_env_int
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging, deployment_pricing_model_info
|
||||
from litellm.litellm_core_utils.litellm_logging import BilledRoundsDispatch, Logging, deployment_pricing_model_info
|
||||
from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token
|
||||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
||||
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
|
||||
|
|
@ -4732,19 +4732,25 @@ class ProxyLogging:
|
|||
|
||||
@staticmethod
|
||||
def _discard_deferred_stream_logging_for_failure(request_data: Mapping[str, object], error: Exception) -> bool:
|
||||
"""Drop the parked success dispatch for an assembled chat stream that ends in an error
|
||||
``post_call_failure_hook`` logs as a failure, billing its usage on the failure row instead.
|
||||
"""Drop the parked success dispatch for an assembled chat stream or a stream that billed several rounds when it
|
||||
ends in an error ``post_call_failure_hook`` logs as a failure, billing its usage on the failure row instead.
|
||||
Returns False when the parked dispatch should still be flushed by the caller."""
|
||||
logging_obj: Final = request_data.get("litellm_logging_obj")
|
||||
if not isinstance(logging_obj, Logging):
|
||||
return False
|
||||
_args: Final[tuple[object, ...] | None] = getattr(logging_obj, "_deferred_stream_complete_args", None)
|
||||
assembled: Final = _args[0] if _args else None
|
||||
if not isinstance(error, _PROXY_ONLY_LLM_API_ERRORS) or not isinstance(assembled, ModelResponse):
|
||||
if not isinstance(error, _PROXY_ONLY_LLM_API_ERRORS):
|
||||
return False
|
||||
match _args:
|
||||
case (ModelResponse() as assembled, *_):
|
||||
logging_obj.record_assembled_response_for_failure(assembled)
|
||||
case (BilledRoundsDispatch() as billed,):
|
||||
billed.dispatch.close()
|
||||
logging_obj.record_billed_rounds_for_failure(billed.spend)
|
||||
case _:
|
||||
return False
|
||||
logging_obj._on_deferred_stream_complete = None
|
||||
logging_obj._deferred_stream_complete_args = None
|
||||
logging_obj.record_assembled_response_for_failure(assembled)
|
||||
return True
|
||||
|
||||
async def arelease_max_parallel_requests_on_disconnect(
|
||||
|
|
|
|||
|
|
@ -5,7 +5,8 @@ from itertools import filterfalse
|
|||
from typing import Any, Final, cast
|
||||
|
||||
import litellm
|
||||
from litellm.litellm_core_utils.hidden_params import get_or_create_hidden_params
|
||||
from litellm.litellm_core_utils.hidden_params import get_or_create_hidden_params, served_from_cache
|
||||
from litellm.litellm_core_utils.litellm_logging import BilledRoundsSpend
|
||||
from litellm.main import stream_chunk_builder
|
||||
from litellm.responses.litellm_completion_transformation.custom_tools import (
|
||||
build_tool_call_item_kwargs,
|
||||
|
|
@ -1363,6 +1364,17 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
|
|||
return False
|
||||
return not any(getattr(part, "text", None) for part in getattr(item, "content", None) or ())
|
||||
|
||||
def finished_round_event(self) -> ResponseCompletedEvent | ResponseIncompleteEvent | None:
|
||||
if not isinstance(self.litellm_model_response, ModelResponse):
|
||||
return None
|
||||
finished: Final = self._emit_terminal_response_event(self.litellm_model_response)
|
||||
if finished is not None and served_from_cache(self.litellm_custom_stream_wrapper):
|
||||
finished.response.hidden_params = {**finished.response.hidden_params, "cache_hit": True}
|
||||
return finished
|
||||
|
||||
def bill_failure_after_earlier_rounds(self, earlier: BilledRoundsSpend) -> None:
|
||||
self.litellm_custom_stream_wrapper.bill_failure_after_earlier_rounds(earlier)
|
||||
|
||||
def partial_response_on_disconnect(self) -> ResponseIncompleteEvent | None:
|
||||
"""
|
||||
The response.incomplete event to bill when the client went away mid-stream.
|
||||
|
|
|
|||
|
|
@ -3,7 +3,9 @@ from typing import TYPE_CHECKING, Any, Final, cast
|
|||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm._uuid import uuid
|
||||
from litellm.responses.litellm_completion_transformation.streaming_iterator import LiteLLMCompletionStreamingIterator
|
||||
from litellm.responses.mcp.request_context import MCPRequestContext
|
||||
from litellm.responses.mcp.stream_round_billing import StreamedRoundBilling
|
||||
from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator
|
||||
from litellm.types.llms.openai import (
|
||||
BaseLiteLLMOpenAIResponseObject,
|
||||
|
|
@ -352,6 +354,9 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator):
|
|||
self._round_max_output_index = -1
|
||||
self._composed_output: list[object] = [] # mutable-ok: grows as each round finishes
|
||||
self._pending_mcp_call_items: list[dict[str, object]] = [] # mutable-ok: grows per executed tool
|
||||
self._iteration_started = False
|
||||
self._round_billing: StreamedRoundBilling | None = None
|
||||
self._bridged_round_billed: LiteLLMCompletionStreamingIterator | None = None
|
||||
|
||||
def _extract_mcp_headers_from_params(self) -> None:
|
||||
"""Extract MCP headers from original request params to pass to tool calls"""
|
||||
|
|
@ -436,7 +441,15 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator):
|
|||
return self
|
||||
|
||||
async def __anext__(self) -> ResponsesAPIStreamingResponse:
|
||||
chunk: Final = await self._anext_impl()
|
||||
if not self._iteration_started:
|
||||
self._iteration_started = True
|
||||
self._round_billing = self._bill_rounds_on_one_row()
|
||||
try:
|
||||
chunk: Final = await self._anext_impl()
|
||||
except Exception:
|
||||
if self._round_billing is not None:
|
||||
self._round_billing.finish()
|
||||
raise
|
||||
sequence_number: Final = getattr(chunk, "sequence_number", None)
|
||||
if isinstance(sequence_number, int):
|
||||
if sequence_number <= self._last_sequence_number and self._last_sequence_number > 0:
|
||||
|
|
@ -446,6 +459,49 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator):
|
|||
self._last_sequence_number = max(self._last_sequence_number, sequence_number)
|
||||
return chunk
|
||||
|
||||
def __del__(self) -> None:
|
||||
round_billing: Final[object] = getattr(self, "_round_billing", None)
|
||||
if isinstance(round_billing, StreamedRoundBilling):
|
||||
round_billing.release_abandoned()
|
||||
|
||||
def _bill_rounds_on_one_row(self) -> StreamedRoundBilling | None:
|
||||
if not self.should_auto_execute:
|
||||
return None
|
||||
if isinstance(self.base_iterator, LiteLLMCompletionStreamingIterator):
|
||||
return StreamedRoundBilling(self.base_iterator.litellm_logging_obj)
|
||||
if isinstance(self.base_iterator, BaseResponsesAPIStreamingIterator):
|
||||
return StreamedRoundBilling(self.base_iterator.logging_obj)
|
||||
return None
|
||||
|
||||
def _restore_round_headers(self, logged_event: ResponsesAPIStreamingResponse) -> None:
|
||||
if isinstance(self.base_iterator, LiteLLMCompletionStreamingIterator):
|
||||
return
|
||||
if isinstance(self.base_iterator, BaseResponsesAPIStreamingIterator):
|
||||
self.base_iterator._restore_provider_response_headers(logged_event)
|
||||
|
||||
def _round_billed_early(self) -> bool:
|
||||
return self._bridged_round_billed is not None and self._bridged_round_billed is self.base_iterator
|
||||
|
||||
def _bill_finished_bridged_round(self) -> None:
|
||||
base: Final = self.base_iterator
|
||||
if self._round_billing is None or not isinstance(base, LiteLLMCompletionStreamingIterator):
|
||||
return
|
||||
if self._round_billed_early():
|
||||
return
|
||||
finished: Final = base.finished_round_event()
|
||||
if finished is None:
|
||||
return
|
||||
self._bridged_round_billed = base
|
||||
round_output: Final = _output_items(finished.response)
|
||||
self._round_billing.complete_round(
|
||||
finished,
|
||||
finished.response,
|
||||
round_output
|
||||
if finished.type == ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE
|
||||
else (*self._composed_output, *round_output),
|
||||
self._restore_round_headers,
|
||||
)
|
||||
|
||||
async def _anext_impl(self) -> ResponsesAPIStreamingResponse:
|
||||
"""
|
||||
Phase-based streaming:
|
||||
|
|
@ -617,6 +673,7 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator):
|
|||
follow-up rounds past the items already emitted, and lists every
|
||||
round's items on the final response.completed.
|
||||
"""
|
||||
self._bill_finished_bridged_round()
|
||||
chunk_type: Final[object] = getattr(chunk, "type", None)
|
||||
if self._round_index > 0 and chunk_type in (
|
||||
ResponsesAPIStreamEvents.RESPONSE_CREATED,
|
||||
|
|
@ -630,12 +687,21 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator):
|
|||
if self._output_index_offset:
|
||||
_set_event_field(chunk, "output_index", output_index + self._output_index_offset)
|
||||
|
||||
if self.should_auto_execute and chunk_type == ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE:
|
||||
return self._charge_incomplete_round(chunk)
|
||||
if not (self.should_auto_execute and self._is_response_completed(chunk)):
|
||||
return chunk
|
||||
|
||||
response_obj: Final[object] = getattr(chunk, "response", None)
|
||||
if isinstance(response_obj, ResponsesAPIResponse):
|
||||
self.collected_response = response_obj
|
||||
if self._round_billing is not None and not self._round_billed_early():
|
||||
self._round_billing.complete_round(
|
||||
chunk,
|
||||
response_obj,
|
||||
(*self._composed_output, *_output_items(response_obj)),
|
||||
self._restore_round_headers,
|
||||
)
|
||||
# Move to tool execution phase after this chunk
|
||||
self.phase = "tool_execution"
|
||||
await self._generate_tool_execution_events()
|
||||
|
|
@ -645,13 +711,33 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator):
|
|||
if self._follow_up_pending():
|
||||
self._absorb_round(response_obj)
|
||||
return None
|
||||
if self._composed_output:
|
||||
merged_output: Final[list[object]] = [ # mutable-ok: the response model declares output as a list
|
||||
*self._composed_output,
|
||||
*_output_items(response_obj),
|
||||
]
|
||||
merged_response: Final = response_obj.model_copy(update={"output": merged_output})
|
||||
_set_event_field(chunk, "response", merged_response)
|
||||
merged_response: Final = (
|
||||
response_obj.model_copy(update={"output": [*self._composed_output, *_output_items(response_obj)]})
|
||||
if self._composed_output
|
||||
else response_obj
|
||||
)
|
||||
client_response: Final = (
|
||||
merged_response
|
||||
if self._round_billing is None
|
||||
else self._round_billing.charge_final_round(chunk, merged_response, self._restore_round_headers)
|
||||
)
|
||||
if client_response is not response_obj:
|
||||
_set_event_field(chunk, "response", client_response)
|
||||
return chunk
|
||||
|
||||
def _charge_incomplete_round(self, chunk: ResponsesAPIStreamingResponse) -> ResponsesAPIStreamingResponse:
|
||||
response_obj: Final[object] = getattr(chunk, "response", None)
|
||||
if self._round_billing is None or not isinstance(response_obj, ResponsesAPIResponse):
|
||||
return chunk
|
||||
if not self._round_billed_early():
|
||||
self._round_billing.complete_round(
|
||||
chunk, response_obj, _output_items(response_obj), self._restore_round_headers
|
||||
)
|
||||
client_response: Final = self._round_billing.charge_final_round(
|
||||
chunk, response_obj, self._restore_round_headers
|
||||
)
|
||||
if client_response is not response_obj:
|
||||
_set_event_field(chunk, "response", client_response)
|
||||
return chunk
|
||||
|
||||
async def _process_base_iterator_chunk(self) -> ResponsesAPIStreamingResponse:
|
||||
|
|
@ -949,6 +1035,8 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator):
|
|||
self.collected_response = None
|
||||
self._cached_response_id = None
|
||||
self._round_index += 1
|
||||
if self._round_billing is not None:
|
||||
self._round_billing.record_billed_rounds(follow_up_response)
|
||||
|
||||
except Exception as e:
|
||||
verbose_logger.error("Error creating follow-up iterator: %s", e)
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ from contextlib import contextmanager
|
|||
from types import MappingProxyType
|
||||
from typing import Final, TypeVar
|
||||
|
||||
from pydantic import TypeAdapter
|
||||
from pydantic import BaseModel, TypeAdapter
|
||||
from pydantic_core import to_jsonable_python
|
||||
|
||||
from litellm.litellm_core_utils.hidden_params import get_hidden_params, served_from_cache, set_hidden_params
|
||||
|
|
@ -46,15 +46,18 @@ def _summed_mapping(first: Mapping[str, object], final: Mapping[str, object]) ->
|
|||
return MappingProxyType({key: _summed_entry(key, first, final) for key in keys})
|
||||
|
||||
|
||||
def _summed_usage(first: BaseModel, final: BaseModel) -> object:
|
||||
return to_jsonable_python(_summed_mapping(first.model_dump(), final.model_dump()), fallback=dict)
|
||||
|
||||
|
||||
def merged_round_usage(first: ResponseAPIUsage | None, final: ResponseAPIUsage | None) -> ResponseAPIUsage | None:
|
||||
if first is None or final is None:
|
||||
return final if first is None else first
|
||||
return ResponseAPIUsage.model_validate(_summed_mapping(first.model_dump(), final.model_dump()))
|
||||
return ResponseAPIUsage.model_validate(_summed_usage(first, final))
|
||||
|
||||
|
||||
def _merged_chat_usage(first: Usage, final: Usage) -> Usage:
|
||||
summed: Final = _summed_mapping(first.model_dump(), final.model_dump())
|
||||
return Usage.model_validate(to_jsonable_python(summed, fallback=dict))
|
||||
def merged_chat_usage(first: Usage, final: Usage) -> Usage:
|
||||
return Usage.model_validate(_summed_usage(first, final))
|
||||
|
||||
|
||||
def _chat_usage(round_response: ResponsesAPIResponse | ModelResponse) -> Usage | None:
|
||||
|
|
@ -82,7 +85,7 @@ def _usage_of_both_rounds(
|
|||
final_usage: Final = _chat_usage(final)
|
||||
if first_usage is None or final_usage is None:
|
||||
return final_usage if first_usage is None else first_usage
|
||||
return _merged_chat_usage(first_usage, final_usage)
|
||||
return merged_chat_usage(first_usage, final_usage)
|
||||
|
||||
|
||||
def _repeated_fixed_margin(
|
||||
|
|
@ -105,6 +108,14 @@ def _less_repeated_fixed_margin(summed: Mapping[str, object], repeated: float) -
|
|||
return MappingProxyType({**summed, **charged_once})
|
||||
|
||||
|
||||
def breakdown_less_repeated_fixed_margin(
|
||||
breakdown: CostBreakdown | None, repeated_fixed_margin: float
|
||||
) -> CostBreakdown | None:
|
||||
if breakdown is None:
|
||||
return None
|
||||
return _COST_BREAKDOWN.validate_python(_less_repeated_fixed_margin(breakdown, repeated_fixed_margin))
|
||||
|
||||
|
||||
def summed_cost_breakdown(
|
||||
first: ResponsesAPIResponse | ModelResponse,
|
||||
first_breakdown: CostBreakdown | None,
|
||||
|
|
|
|||
376
litellm/responses/mcp/stream_round_billing.py
Normal file
376
litellm/responses/mcp/stream_round_billing.py
Normal file
|
|
@ -0,0 +1,376 @@
|
|||
import asyncio
|
||||
from collections.abc import Callable, Coroutine, Sequence
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from typing import Final, Protocol, TypeAlias
|
||||
|
||||
from litellm.litellm_core_utils.hidden_params import get_hidden_params, served_from_cache
|
||||
from litellm.litellm_core_utils.litellm_logging import BilledRoundsDispatch, BilledRoundsSpend, UncachedPrice
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.responses.mcp.round_billing import billed_for_every_round, merged_chat_usage, summed_cost_breakdown
|
||||
from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator, spawn_logging_task
|
||||
from litellm.responses.utils import ResponseAPILoggingUtils
|
||||
from litellm.types.llms.openai import ResponsesAPIResponse, ResponsesAPIStreamingResponse
|
||||
from litellm.types.utils import CostBreakdown, ModelResponse, Usage
|
||||
|
||||
Parked: TypeAlias = tuple[object, ...]
|
||||
RestoresHeaders: TypeAlias = Callable[[ResponsesAPIStreamingResponse], None]
|
||||
|
||||
_PARKED_DISPATCH: Final = "_deferred_stream_complete_args"
|
||||
_DISPATCH_RELEASE: Final = "_on_deferred_stream_complete"
|
||||
_SERVED_FROM_CACHE: Final = UncachedPrice(cost=0.0, cost_breakdown=None)
|
||||
|
||||
|
||||
class _DispatchesSuccessHandlers(Protocol):
|
||||
async def dispatch_success_handlers(
|
||||
self,
|
||||
result: object = None,
|
||||
start_time: object = None,
|
||||
end_time: object = None,
|
||||
cache_hit: object = None,
|
||||
prefer_async_handlers: bool = False,
|
||||
) -> None: ...
|
||||
|
||||
|
||||
def _success_dispatcher(logger: LiteLLMLoggingObj) -> _DispatchesSuccessHandlers:
|
||||
return logger
|
||||
|
||||
|
||||
def _event_loop_running() -> bool:
|
||||
try:
|
||||
asyncio.get_running_loop()
|
||||
except RuntimeError:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def _priced(response: ResponsesAPIResponse, price: UncachedPrice) -> ResponsesAPIResponse:
|
||||
priced: Final = response.model_copy()
|
||||
priced.hidden_params = {**response.hidden_params, "response_cost": price.cost}
|
||||
return priced
|
||||
|
||||
|
||||
def _assembled(parked: Parked | None) -> ModelResponse | None:
|
||||
match parked:
|
||||
case (ModelResponse() as assembled, bool()):
|
||||
return assembled
|
||||
case _:
|
||||
return None
|
||||
|
||||
|
||||
def _bridged_round_usage(parked: Parked | None) -> Usage | None:
|
||||
usage: Final[object] = getattr(_assembled(parked), "usage", None)
|
||||
return usage if isinstance(usage, Usage) else None
|
||||
|
||||
|
||||
def _usage_of_bridged_rounds(billed: Usage | None, round_usage: Usage | None, cost: object) -> Usage | None:
|
||||
if billed is None or round_usage is None:
|
||||
return None
|
||||
return merged_chat_usage(billed, round_usage).model_copy(update={"cost": cost})
|
||||
|
||||
|
||||
def _logged_cache_hit(response: ResponsesAPIResponse, chat_usage: Usage | None) -> bool | None:
|
||||
if chat_usage is not None:
|
||||
return served_from_cache(response)
|
||||
return served_from_cache(response) or None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _RoundRow:
|
||||
logged: ResponsesAPIStreamingResponse | ModelResponse
|
||||
cache_hit: bool | None
|
||||
end_time: datetime
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _SuccessLog:
|
||||
logger: LiteLLMLoggingObj
|
||||
row: _RoundRow
|
||||
chat_usage: Usage | None
|
||||
|
||||
def carry_bridged_usage(self) -> None:
|
||||
if self.chat_usage is not None:
|
||||
self.logger.model_call_details["combined_usage_object"] = self.chat_usage
|
||||
|
||||
|
||||
def _round_row(
|
||||
completed: ResponsesAPIStreamingResponse,
|
||||
response: ResponsesAPIResponse,
|
||||
chat_usage: Usage | None,
|
||||
restore_headers: RestoresHeaders,
|
||||
) -> _RoundRow:
|
||||
logged_response: Final = response.model_copy()
|
||||
logged_response.hidden_params = {**response.hidden_params}
|
||||
logged: Final = completed.model_copy(update={"response": logged_response})
|
||||
restore_headers(logged)
|
||||
return _RoundRow(
|
||||
logged,
|
||||
_logged_cache_hit(response, chat_usage),
|
||||
datetime.now(), # noqa: DTZ005 # naive to match the logging pipeline start_time
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _BilledRounds:
|
||||
response: ResponsesAPIResponse
|
||||
breakdown: CostBreakdown | None
|
||||
chat_usage: Usage | None
|
||||
rounds: int = 1
|
||||
|
||||
@property
|
||||
def stamped_cost(self) -> float | None:
|
||||
cost: Final = self.response.hidden_params.get("response_cost")
|
||||
return float(cost) if isinstance(cost, (int, float)) else None
|
||||
|
||||
@property
|
||||
def cost(self) -> float:
|
||||
return self.stamped_cost or 0.0
|
||||
|
||||
@property
|
||||
def spend(self) -> BilledRoundsSpend:
|
||||
usage: Final = (
|
||||
ResponseAPILoggingUtils.transform_response_api_usage_to_chat_usage(self.response.usage)
|
||||
if self.chat_usage is None
|
||||
else self.chat_usage
|
||||
)
|
||||
return BilledRoundsSpend(usage, self.cost, self.breakdown)
|
||||
|
||||
def with_round(
|
||||
self, priced: ResponsesAPIResponse, breakdown: CostBreakdown | None, chat_usage: Usage | None
|
||||
) -> "_BilledRounds":
|
||||
merged: Final = billed_for_every_round(self.response, self.breakdown, priced, breakdown)
|
||||
total: Final = merged.hidden_params.get("response_cost")
|
||||
charged: Final = merged.model_copy(
|
||||
update={"usage": None if merged.usage is None else merged.usage.model_copy(update={"cost": total})}
|
||||
)
|
||||
charged.hidden_params = {
|
||||
**merged.hidden_params,
|
||||
"cache_hit": served_from_cache(self.response) and served_from_cache(priced),
|
||||
}
|
||||
return _BilledRounds(
|
||||
charged,
|
||||
summed_cost_breakdown(self.response, self.breakdown, priced, breakdown),
|
||||
_usage_of_bridged_rounds(self.chat_usage, chat_usage, total),
|
||||
self.rounds + 1,
|
||||
)
|
||||
|
||||
def row(
|
||||
self,
|
||||
completed: ResponsesAPIStreamingResponse,
|
||||
client_output: Sequence[object],
|
||||
restore_headers: RestoresHeaders,
|
||||
) -> _RoundRow:
|
||||
return _round_row(
|
||||
completed,
|
||||
self.response.model_copy(update={"output": list(client_output)}),
|
||||
self.chat_usage,
|
||||
restore_headers,
|
||||
)
|
||||
|
||||
def bridged_row(self, assembled: ModelResponse | None, logged_id: str) -> _RoundRow | None:
|
||||
if assembled is None or self.chat_usage is None:
|
||||
return None
|
||||
logged: Final = assembled.model_copy(update={"id": logged_id, "usage": self.chat_usage})
|
||||
logged.hidden_params = {
|
||||
**(get_hidden_params(assembled) or {}),
|
||||
"response_cost": self.stamped_cost,
|
||||
}
|
||||
return _RoundRow(
|
||||
logged,
|
||||
served_from_cache(self.response),
|
||||
datetime.now(), # noqa: DTZ005 # naive to match the logging pipeline start_time
|
||||
)
|
||||
|
||||
|
||||
class StreamedRoundBilling:
|
||||
"""One success row per streamed auto-execute request, billed for every round it ran.
|
||||
|
||||
Each round's iterator parks its success dispatch on the shared logging object rather than sending it, either as a
|
||||
ready dispatch or, for a stream bridged through chat completions, as the assembled response and its cache hit.
|
||||
When a round completes, before its tools run, it is folded into the rounds before it and a dispatch for every
|
||||
round so far stays parked, so a client disconnect or a later round that fails still bills them. Once a later round
|
||||
starts, that dispatch logs from a copy of the logging state taken then, so the later round's failure, logged on the
|
||||
request's own state where the proxy reads its failure row, never reaches their success row. Only the latest
|
||||
dispatch is sent, through the proxy's deferred release when the proxy armed one."""
|
||||
|
||||
def __init__(self, logging_obj: LiteLLMLoggingObj) -> None:
|
||||
self._logging_obj: Final = logging_obj
|
||||
self._billed: _BilledRounds | None = None
|
||||
self._banked: Parked | None = None
|
||||
self._banked_row: _RoundRow | None = None
|
||||
if getattr(logging_obj, _DISPATCH_RELEASE, None) is None:
|
||||
setattr(logging_obj, _DISPATCH_RELEASE, self._run_parked)
|
||||
|
||||
def complete_round(
|
||||
self,
|
||||
completed: ResponsesAPIStreamingResponse,
|
||||
response: ResponsesAPIResponse,
|
||||
client_output: Sequence[object],
|
||||
restore_headers: RestoresHeaders,
|
||||
) -> None:
|
||||
round_parked: Final = self._take_parked()
|
||||
price: Final = self._price(response, round_parked)
|
||||
priced: Final = _priced(response, price)
|
||||
chat_usage: Final = _bridged_round_usage(round_parked)
|
||||
if self._billed is None:
|
||||
first_round: Final = _BilledRounds(priced, price.cost_breakdown, chat_usage)
|
||||
self._billed = first_round
|
||||
self._bank(
|
||||
round_parked,
|
||||
first_round.bridged_row(_assembled(round_parked), priced.id)
|
||||
or first_round.row(completed, client_output, restore_headers),
|
||||
)
|
||||
return
|
||||
billed: Final = self._billed.with_round(priced, price.cost_breakdown, chat_usage)
|
||||
self._close(round_parked)
|
||||
self._close(self._banked)
|
||||
self._logging_obj.model_call_details.pop("combined_usage_object", None)
|
||||
self._logging_obj.model_call_details.pop("response_cost", None)
|
||||
self._billed = billed
|
||||
self._logging_obj.cost_breakdown = billed.breakdown
|
||||
row: Final = billed.bridged_row(_assembled(round_parked), billed.response.id) or billed.row(
|
||||
completed, client_output, restore_headers
|
||||
)
|
||||
self._bank(None if round_parked is None else (self._dispatch(row, billed, self._logging_obj),), row)
|
||||
|
||||
def record_billed_rounds(self, follow_up: object) -> None:
|
||||
if self._billed is None or not isinstance(follow_up, BaseResponsesAPIStreamingIterator):
|
||||
return
|
||||
follow_up.bill_failure_after_earlier_rounds(self._billed.spend)
|
||||
self._bank_apart_from_later_rounds(self._billed)
|
||||
|
||||
def charge_final_round(
|
||||
self, completed: ResponsesAPIStreamingResponse, response: ResponsesAPIResponse, restore_headers: RestoresHeaders
|
||||
) -> ResponsesAPIResponse:
|
||||
parked: Final = self._take_parked()
|
||||
self._banked = None
|
||||
if self._billed is None or self._billed.rounds == 1:
|
||||
if parked is not None:
|
||||
self._release(parked)
|
||||
return response
|
||||
self._close(parked)
|
||||
charged: Final = response.model_copy(update={"usage": self._billed.response.usage})
|
||||
charged.hidden_params = {**self._billed.response.hidden_params}
|
||||
if parked is not None:
|
||||
final_row: Final = self._billed.bridged_row(self._latest_assembled(), charged.id) or _round_row(
|
||||
completed, charged, self._billed.chat_usage, restore_headers
|
||||
)
|
||||
self._release((self._dispatch(final_row, self._billed, self._logging_obj),))
|
||||
return charged
|
||||
|
||||
def finish(self) -> None:
|
||||
parked: Final = self._take_parked()
|
||||
self._banked = None
|
||||
if parked is not None:
|
||||
self._release(parked)
|
||||
return
|
||||
self._disarm()
|
||||
|
||||
def release_abandoned(self) -> None:
|
||||
if not self._owns_release():
|
||||
return
|
||||
if _event_loop_running():
|
||||
self.finish()
|
||||
return
|
||||
self.abandon()
|
||||
|
||||
def abandon(self) -> None:
|
||||
self._close(self._take_parked())
|
||||
self._close(self._banked)
|
||||
self._banked = None
|
||||
self._disarm()
|
||||
|
||||
def _latest_assembled(self) -> ModelResponse | None:
|
||||
row: Final = self._banked_row
|
||||
return row.logged if row is not None and isinstance(row.logged, ModelResponse) else None
|
||||
|
||||
def _price(self, response: ResponsesAPIResponse, round_parked: Parked | None) -> UncachedPrice:
|
||||
if served_from_cache(response):
|
||||
return _SERVED_FROM_CACHE
|
||||
match round_parked:
|
||||
case (ModelResponse() as assembled, bool()):
|
||||
return self._logging_obj.uncached_price(assembled)
|
||||
case _:
|
||||
return self._logging_obj.uncached_price(response)
|
||||
|
||||
def _bank(self, parked: Parked | None, row: _RoundRow) -> None:
|
||||
self._banked = parked
|
||||
self._banked_row = row
|
||||
if parked is not None:
|
||||
setattr(self._logging_obj, _PARKED_DISPATCH, parked)
|
||||
|
||||
def _bank_apart_from_later_rounds(self, billed: _BilledRounds) -> None:
|
||||
row: Final = self._banked_row
|
||||
if row is None:
|
||||
return
|
||||
parked: Final = self._take_parked()
|
||||
if parked is None:
|
||||
return
|
||||
self._close(parked)
|
||||
self._bank((self._dispatch(row, billed, self._logging_obj.attempt_scoped_copy()),), row)
|
||||
|
||||
def _dispatch(
|
||||
self, row: _RoundRow, billed: _BilledRounds, logger: LiteLLMLoggingObj
|
||||
) -> BilledRoundsDispatch | Coroutine[object, object, None]:
|
||||
dispatch: Final = self._log_success(_SuccessLog(logger, row, billed.chat_usage))
|
||||
if billed.chat_usage is None:
|
||||
return dispatch
|
||||
return BilledRoundsDispatch(dispatch, billed.spend)
|
||||
|
||||
@staticmethod
|
||||
async def _log_success(success: _SuccessLog) -> None:
|
||||
success.carry_bridged_usage()
|
||||
await _success_dispatcher(success.logger).dispatch_success_handlers(
|
||||
success.row.logged,
|
||||
start_time=success.logger.start_time,
|
||||
end_time=success.row.end_time,
|
||||
cache_hit=success.row.cache_hit,
|
||||
prefer_async_handlers=True,
|
||||
)
|
||||
|
||||
def _take_parked(self) -> Parked | None:
|
||||
parked: Final[Parked | None] = getattr(self._logging_obj, _PARKED_DISPATCH, None)
|
||||
if not isinstance(parked, tuple):
|
||||
return None
|
||||
setattr(self._logging_obj, _PARKED_DISPATCH, None)
|
||||
return parked
|
||||
|
||||
async def _run_parked(self, *parked: object) -> None:
|
||||
match parked:
|
||||
case (BilledRoundsDispatch(dispatch=dispatch),):
|
||||
await dispatch
|
||||
case (dispatch,) if asyncio.iscoroutine(dispatch):
|
||||
await dispatch
|
||||
case (assembled, cache_hit):
|
||||
await _success_dispatcher(self._logging_obj).dispatch_success_handlers(
|
||||
assembled, cache_hit=cache_hit, start_time=None, end_time=None, prefer_async_handlers=True
|
||||
)
|
||||
case _:
|
||||
return
|
||||
|
||||
def _owns_release(self) -> bool:
|
||||
release: Final[object] = getattr(self._logging_obj, _DISPATCH_RELEASE, None)
|
||||
return release == self._run_parked
|
||||
|
||||
def _release(self, parked: Parked) -> None:
|
||||
if getattr(self._logging_obj, _DISPATCH_RELEASE, None) is not None and not self._owns_release():
|
||||
setattr(self._logging_obj, _PARKED_DISPATCH, parked)
|
||||
return
|
||||
setattr(self._logging_obj, _DISPATCH_RELEASE, None)
|
||||
spawn_logging_task(
|
||||
asyncio.get_running_loop(), self._run_parked(*parked), task_name="Responses MCP stream success logging"
|
||||
)
|
||||
|
||||
def _disarm(self) -> None:
|
||||
if self._owns_release():
|
||||
setattr(self._logging_obj, _DISPATCH_RELEASE, None)
|
||||
|
||||
@staticmethod
|
||||
def _close(parked: Parked | None) -> None:
|
||||
match parked:
|
||||
case (BilledRoundsDispatch(dispatch=dispatch),):
|
||||
dispatch.close()
|
||||
case (dispatch,) if asyncio.iscoroutine(dispatch):
|
||||
dispatch.close()
|
||||
case _:
|
||||
return
|
||||
|
|
@ -42,6 +42,7 @@ from litellm.litellm_core_utils.asyncify import run_async_function
|
|||
from litellm.litellm_core_utils.core_helpers import process_response_headers
|
||||
from litellm.litellm_core_utils.exception_mapping_utils import with_deployment_api_base
|
||||
from litellm.litellm_core_utils.internal_call_metadata import without_budget_reservation
|
||||
from litellm.litellm_core_utils.litellm_logging import BilledRoundsSpend
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.litellm_core_utils.llm_response_utils.get_api_base import get_api_base
|
||||
from litellm.litellm_core_utils.llm_response_utils.get_headers import provider_request_id_headers
|
||||
|
|
@ -363,6 +364,7 @@ class BaseResponsesAPIStreamingIterator:
|
|||
"""
|
||||
|
||||
_pending_logging_tasks: tuple[asyncio.Task[object], ...] = ()
|
||||
_earlier_rounds_spend: BilledRoundsSpend | None = None
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
|
|
@ -705,18 +707,22 @@ class BaseResponsesAPIStreamingIterator:
|
|||
if usage_obj is None:
|
||||
return
|
||||
try:
|
||||
self.logging_obj.model_call_details["combined_usage_object"] = (
|
||||
ResponseAPILoggingUtils.transform_response_api_usage_to_chat_usage(usage_obj)
|
||||
)
|
||||
failed_usage: Final = ResponseAPILoggingUtils.transform_response_api_usage_to_chat_usage(usage_obj)
|
||||
except (TypeError, ValueError) as usage_error:
|
||||
verbose_logger.debug(
|
||||
"could not record usage for failed responses stream: %s",
|
||||
usage_error,
|
||||
)
|
||||
return
|
||||
self.logging_obj.model_call_details["response_cost"] = (
|
||||
self.logging_obj.response_cost_calculator(result=response_obj) or 0.0
|
||||
)
|
||||
failed_cost: Final = self.logging_obj.response_cost_calculator(result=response_obj) or 0.0
|
||||
if self._earlier_rounds_spend is None:
|
||||
self.logging_obj.model_call_details["combined_usage_object"] = failed_usage
|
||||
self.logging_obj.model_call_details["response_cost"] = failed_cost
|
||||
return
|
||||
self.logging_obj.record_failure_after_earlier_rounds(self._earlier_rounds_spend, failed_usage, failed_cost)
|
||||
|
||||
def bill_failure_after_earlier_rounds(self, earlier: BilledRoundsSpend) -> None:
|
||||
self._earlier_rounds_spend = earlier
|
||||
|
||||
def _map_error_event_exception(self, error_obj: object) -> Exception:
|
||||
return _map_stream_error_to_exception(
|
||||
|
|
|
|||
|
|
@ -42,6 +42,7 @@ from litellm.integrations.custom_logger import CustomLogger
|
|||
from litellm.litellm_core_utils import litellm_logging
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging
|
||||
from litellm.litellm_core_utils.litellm_logging import (
|
||||
BilledRoundsSpend,
|
||||
Logging,
|
||||
StandardLoggingPayloadSetup,
|
||||
_extract_response_obj_and_hidden_params,
|
||||
|
|
@ -6837,6 +6838,26 @@ def test_failure_handler_zeroes_spend_without_recovered_usage(logging_obj):
|
|||
assert payload["total_tokens"] == 0
|
||||
|
||||
|
||||
def test_failure_after_earlier_rounds_bills_only_the_failed_round_when_it_priced_no_breakdown(
|
||||
logging_obj: LitellmLogging,
|
||||
) -> None:
|
||||
earlier_breakdown: Final = CostBreakdown(total_cost=0.3, margin_fixed_amount=0.1)
|
||||
logging_obj.cost_breakdown = earlier_breakdown
|
||||
logging_obj.model_call_details["litellm_params"] = {}
|
||||
|
||||
logging_obj.record_failure_after_earlier_rounds(
|
||||
BilledRoundsSpend(Usage(prompt_tokens=5, completion_tokens=2, total_tokens=7), 0.3, earlier_breakdown),
|
||||
Usage(prompt_tokens=7, completion_tokens=0, total_tokens=7),
|
||||
0.0,
|
||||
)
|
||||
logging_obj.failure_handler(litellm.ServiceUnavailableError("503", "anthropic", "claude"), "tb")
|
||||
|
||||
row: Final = logging_obj.model_call_details["standard_logging_object"]
|
||||
assert (row["prompt_tokens"], row["completion_tokens"]) == (7, 0)
|
||||
assert row["response_cost"] == pytest.approx(0.0)
|
||||
assert row.get("cost_breakdown") is None or row["cost_breakdown"]["total_cost"] == pytest.approx(0.0)
|
||||
|
||||
|
||||
def test_set_cost_breakdown_stores_reasoning_cost():
|
||||
"""reasoning_cost is stored only when positive, mirroring the cache-cost fields."""
|
||||
from datetime import datetime
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@ Unit Tests for the max parallel request limiter v3 for the proxy
|
|||
"""
|
||||
|
||||
import asyncio
|
||||
import contextvars
|
||||
import logging
|
||||
import sys
|
||||
import time
|
||||
|
|
@ -54,6 +55,7 @@ from litellm.types.mcp import MCPPreCallRequestObject, MCPTransport
|
|||
from litellm.types.mcp_server.mcp_server_manager import MCPServer
|
||||
from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem
|
||||
from litellm.types.utils import (
|
||||
CallTypes,
|
||||
EmbeddingResponse,
|
||||
ModelResponse,
|
||||
PromptTokensDetailsWrapper,
|
||||
|
|
@ -4206,6 +4208,175 @@ async def test_completed_responses_post_call_releases_parallel_slot() -> None:
|
|||
) == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mcp_tool_call_inside_admitted_request_reuses_its_parallel_slot() -> None:
|
||||
api_key: Final = hash_token("sk-mcp-auto-execute")
|
||||
local_cache: Final = DualCache()
|
||||
handler: Final = _PROXY_MaxParallelRequestsHandler(
|
||||
internal_usage_cache=InternalUsageCache(local_cache)
|
||||
)
|
||||
user_api_key_dict: Final = UserAPIKeyAuth(api_key=api_key, max_parallel_requests=1)
|
||||
parallel_key: Final = f"{{api_key:{api_key}}}:max_parallel_requests"
|
||||
owner_call_id: Final = "mcp-auto-execute-owner"
|
||||
|
||||
async def in_flight() -> int:
|
||||
return handler._gauge_in_flight_from_cache_value(
|
||||
await local_cache.async_get_cache(key=parallel_key)
|
||||
)
|
||||
|
||||
def mcp_tool_call() -> dict[str, object]:
|
||||
return {
|
||||
"model": "gpt-4o-mini",
|
||||
"messages": [{"role": "user", "content": "Tool: add"}],
|
||||
"mcp_server_name": "math",
|
||||
}
|
||||
|
||||
async def pre_call(data: dict[str, object], call_type: str) -> None:
|
||||
await handler.async_pre_call_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
cache=local_cache,
|
||||
data=data,
|
||||
call_type=call_type,
|
||||
)
|
||||
|
||||
async def pre_call_from_another_request(
|
||||
data: dict[str, object], call_type: str
|
||||
) -> None:
|
||||
await contextvars.Context().run(asyncio.create_task, pre_call(data, call_type))
|
||||
|
||||
await pre_call(
|
||||
{"model": "gpt-4o-mini", "input": "add 2 and 3", "litellm_call_id": owner_call_id},
|
||||
CallTypes.aresponses.value,
|
||||
)
|
||||
await pre_call(mcp_tool_call(), CallTypes.call_mcp_tool.value)
|
||||
await pre_call(mcp_tool_call(), CallTypes.call_mcp_tool.value)
|
||||
assert await in_flight() == 1
|
||||
|
||||
with pytest.raises(HTTPException):
|
||||
await pre_call_from_another_request(
|
||||
mcp_tool_call(), CallTypes.call_mcp_tool.value
|
||||
)
|
||||
|
||||
await handler.async_log_success_event(
|
||||
kwargs={"litellm_call_id": owner_call_id},
|
||||
response_obj=None,
|
||||
start_time=None,
|
||||
end_time=None,
|
||||
)
|
||||
assert await in_flight() == 0
|
||||
await pre_call_from_another_request(
|
||||
{"model": "gpt-4o-mini", "input": "next request"},
|
||||
CallTypes.aresponses.value,
|
||||
)
|
||||
assert await in_flight() == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("stream", "in_flight_after_first_round_logs"),
|
||||
[(False, 1), (True, 0)],
|
||||
ids=["non_stream_round_logged_late_keeps_slot", "stream_logged_at_its_end_releases_slot"],
|
||||
)
|
||||
async def test_first_round_logged_after_tool_call_reuse_releases_only_at_request_end(
|
||||
stream: bool, in_flight_after_first_round_logs: int
|
||||
) -> None:
|
||||
reused_at: Final = datetime(2026, 1, 5, 12, 0, 0)
|
||||
api_key: Final = hash_token(f"sk-mcp-late-round-log-{stream}")
|
||||
local_cache: Final = DualCache()
|
||||
handler: Final = _PROXY_MaxParallelRequestsHandler(
|
||||
internal_usage_cache=InternalUsageCache(local_cache),
|
||||
time_provider=lambda: reused_at,
|
||||
)
|
||||
user_api_key_dict: Final = UserAPIKeyAuth(api_key=api_key, max_parallel_requests=1)
|
||||
parallel_key: Final = f"{{api_key:{api_key}}}:max_parallel_requests"
|
||||
owner_call_id: Final = f"mcp-late-round-log-{stream}"
|
||||
|
||||
async def in_flight() -> int:
|
||||
return handler._gauge_in_flight_from_cache_value(
|
||||
await local_cache.async_get_cache(key=parallel_key)
|
||||
)
|
||||
|
||||
async def pre_call(data: dict[str, object], call_type: str) -> None:
|
||||
await handler.async_pre_call_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
cache=local_cache,
|
||||
data=data,
|
||||
call_type=call_type,
|
||||
)
|
||||
|
||||
async def round_logged(ended_at: datetime) -> None:
|
||||
await handler.async_log_success_event(
|
||||
kwargs={"litellm_call_id": owner_call_id, "stream": stream},
|
||||
response_obj=None,
|
||||
start_time=None,
|
||||
end_time=ended_at,
|
||||
)
|
||||
|
||||
await pre_call(
|
||||
{
|
||||
"model": "claude-haiku-5-5",
|
||||
"messages": [{"role": "user", "content": "add 2 and 3"}],
|
||||
"litellm_call_id": owner_call_id,
|
||||
},
|
||||
CallTypes.anthropic_messages.value,
|
||||
)
|
||||
await pre_call(
|
||||
{"model": "MCP: add", "messages": [{"role": "user", "content": "Tool: add"}], "mcp_server_name": "math"},
|
||||
CallTypes.call_mcp_tool.value,
|
||||
)
|
||||
|
||||
await round_logged(reused_at - timedelta(seconds=1))
|
||||
assert await in_flight() == in_flight_after_first_round_logs
|
||||
|
||||
await round_logged(reused_at + timedelta(seconds=1))
|
||||
assert await in_flight() == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mcp_tool_call_reusing_request_slot_keeps_request_max_parallel_headers() -> None:
|
||||
api_key: Final = hash_token("sk-mcp-reuse-headers")
|
||||
local_cache: Final = DualCache()
|
||||
handler: Final = _PROXY_MaxParallelRequestsHandler(
|
||||
internal_usage_cache=InternalUsageCache(local_cache)
|
||||
)
|
||||
user_api_key_dict: Final = UserAPIKeyAuth(api_key=api_key, max_parallel_requests=2)
|
||||
data: Final[dict[str, object]] = {
|
||||
"model": "gpt-4o-mini",
|
||||
"input": "add 2 and 3",
|
||||
"litellm_call_id": "mcp-reuse-headers-owner",
|
||||
}
|
||||
for request_data, call_type in (
|
||||
(data, CallTypes.aresponses.value),
|
||||
(
|
||||
{"model": "MCP: add", "messages": [{"role": "user", "content": "Tool: add"}], "mcp_server_name": "math"},
|
||||
CallTypes.call_mcp_tool.value,
|
||||
),
|
||||
):
|
||||
await handler.async_pre_call_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
cache=local_cache,
|
||||
data=request_data,
|
||||
call_type=call_type,
|
||||
)
|
||||
response: Final = ModelResponse(
|
||||
id="mcp-reuse-headers",
|
||||
object="chat.completion",
|
||||
created=0,
|
||||
model="gpt-4o-mini",
|
||||
usage=Usage(prompt_tokens=1, completion_tokens=1, total_tokens=2),
|
||||
choices=[],
|
||||
)
|
||||
response._hidden_params = {}
|
||||
|
||||
await handler.async_post_call_success_hook(
|
||||
data=data, user_api_key_dict=user_api_key_dict, response=response
|
||||
)
|
||||
|
||||
headers: Final = response._hidden_params["additional_headers"]
|
||||
assert headers["x-ratelimit-api_key-limit-max_parallel_requests"] == 2
|
||||
assert headers["x-ratelimit-api_key-remaining-max_parallel_requests"] == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_concurrent_success_callbacks_release_parallel_slot_once_when_redis_fails() -> None:
|
||||
from unittest.mock import AsyncMock
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -23,7 +23,8 @@ import asyncio
|
|||
import json
|
||||
from datetime import datetime
|
||||
from importlib import import_module
|
||||
from typing import Final
|
||||
from collections.abc import Mapping
|
||||
from typing import Final, cast
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
import httpx
|
||||
|
|
@ -31,6 +32,7 @@ import pytest
|
|||
|
||||
import litellm
|
||||
from litellm.exceptions import MidStreamFallbackError
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.litellm_core_utils.get_litellm_params import get_litellm_params
|
||||
from litellm.litellm_core_utils.litellm_logging import _get_provider_request_id
|
||||
|
|
@ -44,6 +46,7 @@ from litellm.responses.streaming_iterator import (
|
|||
_status_code_for_error_fields,
|
||||
)
|
||||
from litellm.utils import ProviderConfigManager
|
||||
from litellm.types.utils import StandardLoggingPayload
|
||||
from litellm.types.llms.openai import (
|
||||
ErrorEvent,
|
||||
ErrorEventError,
|
||||
|
|
@ -523,6 +526,82 @@ def test_handle_logging_failed_response_records_usage_and_cost():
|
|||
iterator.logging_obj.response_cost_calculator.assert_called_once_with(result=chunk.response)
|
||||
|
||||
|
||||
class _FailureRecorder(CustomLogger):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.failed: tuple[StandardLoggingPayload, ...] = ()
|
||||
|
||||
async def async_log_failure_event(
|
||||
self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime
|
||||
) -> None:
|
||||
self.failed = (*self.failed, cast(StandardLoggingPayload, kwargs["standard_logging_object"]))
|
||||
|
||||
|
||||
async def _first_failure_row(recorder: _FailureRecorder) -> StandardLoggingPayload:
|
||||
for _ in range(200):
|
||||
if recorder.failed:
|
||||
return recorder.failed[0]
|
||||
await asyncio.sleep(0.01)
|
||||
pytest.fail("the failed stream logged no failure row")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_failed_stream_bills_only_its_own_usage_over_an_earlier_attempts_stash():
|
||||
recorder: Final = _FailureRecorder()
|
||||
logging_obj: Final = LiteLLMLoggingObj(
|
||||
model="gpt-5.6",
|
||||
messages=[],
|
||||
stream=True,
|
||||
call_type="aresponses",
|
||||
start_time=datetime(2026, 1, 1),
|
||||
litellm_call_id="responses-failed-over-stash",
|
||||
function_id="responses-failed-over-stash",
|
||||
dynamic_async_failure_callbacks=[recorder],
|
||||
)
|
||||
logging_obj.update_environment_variables(
|
||||
model="gpt-5.6", optional_params={}, litellm_params=get_litellm_params(custom_llm_provider="openai")
|
||||
)
|
||||
logging_obj.model_call_details["combined_usage_object"] = litellm.Usage(
|
||||
prompt_tokens=100, completion_tokens=50, total_tokens=150
|
||||
)
|
||||
logging_obj.model_call_details["response_cost"] = 0.5
|
||||
failed: Final = {
|
||||
"type": "response.failed",
|
||||
"sequence_number": 1,
|
||||
"response": {
|
||||
"id": "resp_failed_over_stash",
|
||||
"object": "response",
|
||||
"created_at": 1790000000,
|
||||
"status": "failed",
|
||||
"model": "gpt-5.6",
|
||||
"output": [],
|
||||
"error": {"code": "server_error", "message": "boom"},
|
||||
"usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15},
|
||||
},
|
||||
}
|
||||
iterator: Final = ResponsesAPIStreamingIterator(
|
||||
response=httpx.Response(
|
||||
200, headers={"content-type": "text/event-stream"}, content=f"data: {json.dumps(failed)}\n\n".encode()
|
||||
),
|
||||
model="gpt-5.6",
|
||||
responses_api_provider_config=ProviderConfigManager.get_provider_responses_api_config(
|
||||
provider="openai", model="gpt-5.6"
|
||||
),
|
||||
logging_obj=logging_obj,
|
||||
custom_llm_provider="openai",
|
||||
)
|
||||
|
||||
with pytest.raises(MidStreamFallbackError):
|
||||
await anext(aiter(iterator))
|
||||
|
||||
row: Final = await _first_failure_row(recorder)
|
||||
prompt_cost, completion_cost = litellm.cost_per_token(
|
||||
model="gpt-5.6", prompt_tokens=10, completion_tokens=5, custom_llm_provider="openai"
|
||||
)
|
||||
assert (row["prompt_tokens"], row["completion_tokens"]) == (10, 5)
|
||||
assert row["response_cost"] == pytest.approx(prompt_cost + completion_cost)
|
||||
|
||||
|
||||
def test_handle_logging_failed_response_without_usage_skips_recording():
|
||||
iterator = _make_iterator()
|
||||
iterator.completed_response = _make_failed_chunk(
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue