mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-06 07:48:06 +00:00
fix(mcp): bill every chat MCP stream round once on failure, hang-up, and sync drop
When the follow-up call fails, round 1's success row is now logged on a copy of the logging object that does not share the request's guardrail results, so a post-call guardrail that reports on the failed request counts once instead of on both rows A client hang-up during the follow-up now bills round 1 plus the follow-up's partial answer on one success row with round 1's id and the fixed margin counted once, the same way a plain stream's hang-up bills its partial answer The last usage chunk now carries round 1's usage and cost when only round 1 is billed, and a synchronous reader that drops the stream mid-answer bills round 1 on the loop it read with instead of discarding it The deferred release's logged_on argument is typed and every release closure accepts it, and the parked dispatch constants and protocols both MCP billing paths use now live in one module. The unit tests wait on a logging signal instead of a fixed delay and check the billed row through the proxy stream hook instead of private callback state
This commit is contained in:
mateo-berri
parent
e512269bbc
commit
c32692bd9b
10 changed files with 624 additions and 145 deletions
|
|
@ -564,6 +564,11 @@ def _resolve_mantle_region_for_cost(
|
|||
_PER_CALL_LITELLM_PARAM_DICTS: Final = frozenset({"metadata", "litellm_metadata"})
|
||||
_LITELLM_PARAMS: Final = TypeAdapter(dict[str, object])
|
||||
_REDACTED_RESULT: Final = TypeAdapter(object)
|
||||
_GUARDRAIL_INFORMATION_KEY: Final = "standard_logging_guardrail_information"
|
||||
|
||||
|
||||
def _without_guardrail_information(metadata: Mapping[str, object]) -> Mapping[str, object]:
|
||||
return {key: value for key, value in metadata.items() if key != _GUARDRAIL_INFORMATION_KEY}
|
||||
|
||||
|
||||
def _provider_response_id(source: object) -> str | None:
|
||||
|
|
@ -2494,6 +2499,23 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
attempt.model_call_details = dict(self.model_call_details)
|
||||
return attempt
|
||||
|
||||
def guardrail_scoped_copy(self) -> "Logging":
|
||||
request_params: Final = _LITELLM_PARAMS.validate_python(
|
||||
self.litellm_params or {} # pyright: ignore[reportUnknownMemberType] # Logging.litellm_params is untyped
|
||||
)
|
||||
litellm_params: Final = {
|
||||
**request_params,
|
||||
**{
|
||||
key: _without_guardrail_information(_METADATA_MAPPING.validate_python(value))
|
||||
for key, value in request_params.items()
|
||||
if key in _PER_CALL_LITELLM_PARAM_DICTS and isinstance(value, dict)
|
||||
},
|
||||
}
|
||||
scoped: Final = copy.copy(self)
|
||||
scoped.litellm_params = litellm_params
|
||||
scoped.model_call_details = {**self.model_call_details, "litellm_params": litellm_params}
|
||||
return scoped
|
||||
|
||||
def dispatch_attempt_failure_handlers(
|
||||
self, exception: Exception, traceback_exception: str
|
||||
) -> Coroutine[None, None, None]:
|
||||
|
|
|
|||
|
|
@ -2458,6 +2458,25 @@ class CustomStreamWrapper:
|
|||
def bill_failure_after_earlier_rounds(self, earlier: "BilledRoundsSpend") -> None:
|
||||
self._earlier_rounds_spend = earlier
|
||||
|
||||
def partial_response(self) -> "ModelResponse | TextCompletionResponse | None":
|
||||
if not self.chunks:
|
||||
return None
|
||||
partial: Final = litellm.stream_chunk_builder(
|
||||
chunks=self.chunks,
|
||||
messages=self.messages if isinstance(self.messages, list) else None,
|
||||
logging_obj=self.logging_obj,
|
||||
count_prompt_tokens=self.count_prompt_tokens,
|
||||
)
|
||||
if partial is None:
|
||||
return None
|
||||
usage: Final[object] = getattr(partial, "usage", None)
|
||||
if not isinstance(usage, Usage):
|
||||
return None
|
||||
if self.model:
|
||||
partial.model = self.model
|
||||
backfill_missing_cache_usage_fields(usage)
|
||||
return partial
|
||||
|
||||
def _record_partial_usage_for_failure(self) -> None:
|
||||
"""
|
||||
A stream that breaks mid-flight still billed the provider for the chunks
|
||||
|
|
@ -2467,23 +2486,15 @@ class CustomStreamWrapper:
|
|||
moves it into terminal failure usage before opening the next attempt,
|
||||
so a recovered attempt logs its own usage.
|
||||
"""
|
||||
if self.logging_obj is None or not self.chunks:
|
||||
if self.logging_obj is None:
|
||||
return
|
||||
try:
|
||||
partial_response: Final = litellm.stream_chunk_builder(
|
||||
chunks=self.chunks,
|
||||
messages=self.messages if isinstance(self.messages, list) else None,
|
||||
logging_obj=self.logging_obj,
|
||||
count_prompt_tokens=self.count_prompt_tokens,
|
||||
)
|
||||
partial_response: Final = self.partial_response()
|
||||
if partial_response is None:
|
||||
return
|
||||
usage: Final = cast(Usage | None, getattr(partial_response, "usage", None))
|
||||
if usage is None:
|
||||
usage: Final[object] = getattr(partial_response, "usage", None)
|
||||
if not isinstance(usage, Usage):
|
||||
return
|
||||
if self.model:
|
||||
partial_response.model = self.model
|
||||
backfill_missing_cache_usage_fields(usage)
|
||||
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
|
||||
|
|
|
|||
|
|
@ -137,6 +137,7 @@ from litellm.proxy.utils import ( # noqa: F401 # legacy module exports
|
|||
_check_and_merge_model_level_guardrails, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
|
||||
check_and_merge_model_level_guardrails,
|
||||
)
|
||||
from litellm.responses.mcp.parked_stream_dispatch import parked_success
|
||||
from litellm.responses.streaming_iterator import partial_responses_incomplete_event_on_disconnect, spawn_logging_task
|
||||
from litellm.router import Router
|
||||
from litellm.router_strategy.auto_router.auto_router_model_naming import AUTO_ROUTER_MODEL_PREFIX
|
||||
|
|
@ -527,6 +528,20 @@ async def _bill_partial_streamed_spend_on_disconnect(request_data: dict, respons
|
|||
return await _dispatch_disconnect_billing_event(logging_obj, partial_response)
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class _BillsInFlightRoundOnDisconnect(Protocol):
|
||||
def bill_in_flight_round_on_disconnect(self) -> None: ...
|
||||
|
||||
|
||||
def _bill_in_flight_round_on_disconnect(response: object) -> None:
|
||||
if not isinstance(response, _BillsInFlightRoundOnDisconnect):
|
||||
return
|
||||
try:
|
||||
response.bill_in_flight_round_on_disconnect()
|
||||
except Exception as e: # noqa: BLE001 # disconnect billing is best-effort; never break stream teardown
|
||||
verbose_proxy_logger.debug("Failed to bill the in-flight round after client disconnect: %s", e)
|
||||
|
||||
|
||||
async def _dispatch_disconnect_billing_event(logging_obj: LiteLLMLoggingObj, partial_response: object) -> bool:
|
||||
try:
|
||||
await logging_obj.dispatch_success_handlers(
|
||||
|
|
@ -2052,6 +2067,13 @@ class _ClientStreamFailureReport:
|
|||
)
|
||||
|
||||
|
||||
def _metadata_logged_on(logged_on: LiteLLMLoggingObj) -> Mapping[str, object]:
|
||||
logged_params: Final = object_mapping(logged_on.model_call_details.get("litellm_params")) or MappingProxyType({})
|
||||
return MappingProxyType(
|
||||
{key: logged_params[key] for key in ("metadata", "litellm_metadata") if key in logged_params}
|
||||
)
|
||||
|
||||
|
||||
class ProxyBaseLLMRequestProcessing:
|
||||
def __init__(self, data: dict, served_model: str | None = None):
|
||||
self.data = data
|
||||
|
|
@ -4065,7 +4087,15 @@ class ProxyBaseLLMRequestProcessing:
|
|||
assembled_response: object, cache_hit: object, logged_on: LiteLLMLoggingObj | None = None
|
||||
) -> None:
|
||||
await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails(
|
||||
captured_data=_captured_data,
|
||||
captured_data=_captured_data
|
||||
if logged_on is None
|
||||
else {
|
||||
**_captured_data,
|
||||
**{
|
||||
key: value for key, value in _metadata_logged_on(logged_on).items() if key in _captured_data
|
||||
},
|
||||
"litellm_logging_obj": logged_on,
|
||||
},
|
||||
captured_user_api_key_dict=_captured_user_api_key_dict,
|
||||
captured_logging_obj=_captured_logging_obj if logged_on is None else logged_on,
|
||||
assembled_response=assembled_response,
|
||||
|
|
@ -4088,14 +4118,8 @@ class ProxyBaseLLMRequestProcessing:
|
|||
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):
|
||||
await _as_success_dispatcher(_captured_native_logging_obj).dispatch_success_handlers(
|
||||
assembled_response,
|
||||
cache_hit=cache_hit,
|
||||
start_time=None,
|
||||
end_time=None,
|
||||
prefer_async_handlers=True,
|
||||
)
|
||||
case _ if (success := parked_success(args)) is not None:
|
||||
await success.dispatch(_captured_native_logging_obj)
|
||||
case _:
|
||||
verbose_proxy_logger.error(
|
||||
"Deferred stream logging dropped: unexpected stored args shape %s",
|
||||
|
|
@ -4423,6 +4447,7 @@ class ProxyBaseLLMRequestProcessing:
|
|||
client_disconnected,
|
||||
)
|
||||
if recorded_client_disconnect:
|
||||
_bill_in_flight_round_on_disconnect(response)
|
||||
deferred_stream_logging_armed: Final = _deferred_stream_logging_is_armed(request_data)
|
||||
ProxyLogging.fire_deferred_stream_logging(request_data)
|
||||
# A disconnect-time success event (the deferred-guardrail flush
|
||||
|
|
|
|||
|
|
@ -273,6 +273,18 @@ async def acompletion_with_mcp(
|
|||
if isinstance(round_billing, ChatStreamRoundBilling):
|
||||
round_billing.release_abandoned()
|
||||
|
||||
def bill_in_flight_round_on_disconnect(self) -> None:
|
||||
follow_up: Final[object] = self.follow_up_stream
|
||||
if (
|
||||
self._round_billing is None
|
||||
or self.follow_up_exhausted
|
||||
or not isinstance(follow_up, CustomStreamWrapper)
|
||||
):
|
||||
return
|
||||
partial: Final = follow_up.partial_response()
|
||||
if isinstance(partial, ModelResponse):
|
||||
self._round_billing.bill_partial_follow_up(partial)
|
||||
|
||||
def _add_mcp_list_tools_to_chunk(self, chunk: ModelResponseStream) -> ModelResponseStream:
|
||||
"""Add mcp_list_tools to the first chunk."""
|
||||
from litellm.types.utils import (
|
||||
|
|
@ -598,6 +610,9 @@ async def acompletion_with_mcp(
|
|||
self._custom_iterator.release_abandoned()
|
||||
await super().aclose()
|
||||
|
||||
def bill_in_flight_round_on_disconnect(self) -> None:
|
||||
self._custom_iterator.bill_in_flight_round_on_disconnect()
|
||||
|
||||
def __iter__(self):
|
||||
# For synchronous iteration, create a sync wrapper
|
||||
if self._sync_iterator is None:
|
||||
|
|
|
|||
|
|
@ -1,39 +1,26 @@
|
|||
import asyncio
|
||||
from collections.abc import Callable, Coroutine, Generator
|
||||
from collections.abc import Generator
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass, replace
|
||||
from types import MappingProxyType
|
||||
from typing import Final, Protocol, TypeAlias
|
||||
from typing import Final, TypeAlias
|
||||
|
||||
from litellm.litellm_core_utils.asyncify import is_event_loop_running
|
||||
from litellm.litellm_core_utils.hidden_params import get_hidden_params, served_from_cache, set_hidden_params
|
||||
from litellm.litellm_core_utils.litellm_logging import BilledRoundsSpend, UncachedPrice
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.responses.mcp.parked_stream_dispatch import (
|
||||
DISPATCH_RELEASE,
|
||||
PARKED_DISPATCH,
|
||||
SERVED_FROM_CACHE,
|
||||
DeferredStreamRelease,
|
||||
parked_success,
|
||||
)
|
||||
from litellm.responses.mcp.round_billing import billed_for_every_round, merged_chat_usage, summed_cost_breakdown
|
||||
from litellm.responses.streaming_iterator import spawn_logging_task
|
||||
from litellm.types.utils import ModelResponse, ModelResponseStream, Usage
|
||||
|
||||
Parked: TypeAlias = tuple[object, ...]
|
||||
_Release: TypeAlias = Callable[..., Coroutine[object, object, 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
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
|
|
@ -73,8 +60,13 @@ class _FirstRound:
|
|||
def client_usage_chunk(self, billed: ModelResponse | None) -> ModelResponseStream | None:
|
||||
if self.usage_chunk is None:
|
||||
return None
|
||||
merged: Final = merged_chat_usage(self.priced.usage, self.usage_chunk.usage)
|
||||
usage: Final = merged if billed is None else merged.model_copy(update=MappingProxyType({"cost": _cost(billed)}))
|
||||
usage: Final = (
|
||||
self.priced.usage.model_copy(update=MappingProxyType({"cost": self.priced.price.cost}))
|
||||
if billed is None
|
||||
else merged_chat_usage(self.priced.usage, self.usage_chunk.usage).model_copy(
|
||||
update=MappingProxyType({"cost": _cost(billed)})
|
||||
)
|
||||
)
|
||||
return self.usage_chunk.chunk.model_copy(update=MappingProxyType({"usage": usage}))
|
||||
|
||||
|
||||
|
|
@ -99,12 +91,14 @@ class ChatStreamRoundBilling:
|
|||
def __init__(self, logging_obj: LiteLLMLoggingObj) -> None:
|
||||
self._logging_obj: Final = logging_obj
|
||||
self._first: _FirstRound | None = None
|
||||
self._reading_loop: asyncio.AbstractEventLoop | None = None
|
||||
|
||||
@contextmanager
|
||||
def armed_while_reading(self) -> Generator[None]:
|
||||
armed: Final = getattr(self._logging_obj, _DISPATCH_RELEASE, None) is None
|
||||
armed: Final = getattr(self._logging_obj, DISPATCH_RELEASE, None) is None
|
||||
if armed:
|
||||
setattr(self._logging_obj, _DISPATCH_RELEASE, self._log_success)
|
||||
setattr(self._logging_obj, DISPATCH_RELEASE, self._log_success)
|
||||
self._reading_loop = asyncio.get_running_loop()
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
|
|
@ -135,20 +129,30 @@ class ChatStreamRoundBilling:
|
|||
def fail(self) -> None:
|
||||
self._bank_apart_from_the_follow_up()
|
||||
first: Final = self._first
|
||||
if first is None or first.banked is None or self._peek_parked() is not first.parked:
|
||||
if first is None or first.banked is None or not self._holds_round_one_alone(first):
|
||||
self.finish()
|
||||
return
|
||||
self._first = None
|
||||
setattr(self._logging_obj, _PARKED_DISPATCH, None)
|
||||
self._log_through_the_release((*first.parked, first.banked))
|
||||
setattr(self._logging_obj, PARKED_DISPATCH, None)
|
||||
self._log_through_the_release((*first.parked, first.banked.guardrail_scoped_copy()))
|
||||
|
||||
def bill_partial_follow_up(self, partial: ModelResponse) -> None:
|
||||
first: Final = self._first
|
||||
final_round: Final = _assembled_round((partial, False))
|
||||
if first is None or final_round is None or not self._holds_round_one_alone(first):
|
||||
return
|
||||
self._first = None
|
||||
setattr(
|
||||
self._logging_obj, PARKED_DISPATCH, self._billed_for_both_rounds(first.priced, self._priced(final_round))
|
||||
)
|
||||
|
||||
def release_abandoned(self) -> None:
|
||||
if not self._owns_release():
|
||||
return
|
||||
if is_event_loop_running():
|
||||
if is_event_loop_running() or self._idle_reading_loop() is not None:
|
||||
self.finish()
|
||||
return
|
||||
setattr(self._logging_obj, _PARKED_DISPATCH, None)
|
||||
setattr(self._logging_obj, PARKED_DISPATCH, None)
|
||||
self._first = None
|
||||
self._disarm()
|
||||
|
||||
|
|
@ -162,7 +166,7 @@ class ChatStreamRoundBilling:
|
|||
return
|
||||
logged_alone: Final = (*first.parked, banked)
|
||||
self._first = replace(first, logged_alone=logged_alone, banked=banked)
|
||||
setattr(self._logging_obj, _PARKED_DISPATCH, logged_alone)
|
||||
setattr(self._logging_obj, PARKED_DISPATCH, logged_alone)
|
||||
|
||||
def _release_rounds(self) -> ModelResponse | None:
|
||||
parked: Final = self._take_parked()
|
||||
|
|
@ -196,7 +200,7 @@ class ChatStreamRoundBilling:
|
|||
|
||||
def _priced(self, assembled: _AssembledRound) -> _PricedRound:
|
||||
served: Final = assembled.cache_hit or served_from_cache(assembled.response)
|
||||
price: Final = _SERVED_FROM_CACHE if served else self._logging_obj.uncached_price(assembled.response)
|
||||
price: Final = SERVED_FROM_CACHE if served else self._logging_obj.uncached_price(assembled.response)
|
||||
priced: Final = assembled.response.model_copy()
|
||||
set_hidden_params(
|
||||
priced,
|
||||
|
|
@ -205,36 +209,47 @@ class ChatStreamRoundBilling:
|
|||
return _PricedRound(priced, assembled.usage, price, served)
|
||||
|
||||
def _peek_parked(self) -> Parked | None:
|
||||
parked: Final[Parked | None] = getattr(self._logging_obj, _PARKED_DISPATCH, None)
|
||||
parked: Final[Parked | None] = getattr(self._logging_obj, PARKED_DISPATCH, None)
|
||||
return parked if isinstance(parked, tuple) else None
|
||||
|
||||
def _take_parked(self) -> Parked | None:
|
||||
parked: Final = self._peek_parked()
|
||||
setattr(self._logging_obj, _PARKED_DISPATCH, None)
|
||||
setattr(self._logging_obj, PARKED_DISPATCH, None)
|
||||
return parked
|
||||
|
||||
def _holds_round_one_alone(self, first: _FirstRound) -> bool:
|
||||
parked: Final = self._peek_parked()
|
||||
return parked is first.parked or parked is first.logged_alone
|
||||
|
||||
def _owns_release(self) -> bool:
|
||||
release: Final[object] = getattr(self._logging_obj, _DISPATCH_RELEASE, None)
|
||||
release: Final[object] = getattr(self._logging_obj, DISPATCH_RELEASE, None)
|
||||
return release == self._log_success
|
||||
|
||||
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)
|
||||
if getattr(self._logging_obj, DISPATCH_RELEASE, None) is not None and not self._owns_release():
|
||||
setattr(self._logging_obj, PARKED_DISPATCH, parked)
|
||||
return
|
||||
self._log_through_the_release(parked)
|
||||
|
||||
def _log_through_the_release(self, parked: Parked) -> None:
|
||||
armed: Final[_Release | None] = getattr(self._logging_obj, _DISPATCH_RELEASE, None)
|
||||
setattr(self._logging_obj, _DISPATCH_RELEASE, None)
|
||||
armed: Final[DeferredStreamRelease | None] = getattr(self._logging_obj, DISPATCH_RELEASE, None)
|
||||
setattr(self._logging_obj, DISPATCH_RELEASE, None)
|
||||
logged: Final = (self._log_success if armed is None else armed)(*parked)
|
||||
idle_reading_loop: Final = None if is_event_loop_running() else self._idle_reading_loop()
|
||||
if idle_reading_loop is not None:
|
||||
idle_reading_loop.run_until_complete(logged)
|
||||
return
|
||||
spawn_logging_task(asyncio.get_running_loop(), logged, task_name="Chat completions MCP stream success log")
|
||||
|
||||
def _idle_reading_loop(self) -> asyncio.AbstractEventLoop | None:
|
||||
loop: Final = self._reading_loop
|
||||
return loop if loop is not None and not loop.is_closed() and not loop.is_running() else None
|
||||
|
||||
def _disarm(self) -> None:
|
||||
if self._owns_release():
|
||||
setattr(self._logging_obj, _DISPATCH_RELEASE, None)
|
||||
setattr(self._logging_obj, DISPATCH_RELEASE, None)
|
||||
|
||||
async def _log_success(self, assembled: object, cache_hit: object, banked: object = None) -> None:
|
||||
logger: Final = banked if isinstance(banked, LiteLLMLoggingObj) else self._logging_obj
|
||||
await _success_dispatcher(logger).dispatch_success_handlers(
|
||||
assembled, cache_hit=cache_hit, start_time=None, end_time=None, prefer_async_handlers=True
|
||||
)
|
||||
async def _log_success(self, *parked: object) -> None:
|
||||
success: Final = parked_success(parked)
|
||||
if success is not None:
|
||||
await success.dispatch(self._logging_obj)
|
||||
|
|
|
|||
52
litellm/responses/mcp/parked_stream_dispatch.py
Normal file
52
litellm/responses/mcp/parked_stream_dispatch.py
Normal file
|
|
@ -0,0 +1,52 @@
|
|||
from collections.abc import Coroutine
|
||||
from dataclasses import dataclass
|
||||
from typing import Final, Protocol
|
||||
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.litellm_core_utils.litellm_logging import UncachedPrice
|
||||
|
||||
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: ...
|
||||
|
||||
|
||||
class DeferredStreamRelease(Protocol):
|
||||
def __call__(self, *parked: object) -> Coroutine[object, object, None]: ...
|
||||
|
||||
|
||||
def success_dispatcher(logger: LiteLLMLoggingObj) -> DispatchesSuccessHandlers:
|
||||
return logger
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ParkedSuccess:
|
||||
assembled: object
|
||||
cache_hit: object
|
||||
logged_on: LiteLLMLoggingObj | None = None
|
||||
|
||||
async def dispatch(self, request_logger: LiteLLMLoggingObj) -> None:
|
||||
logger: Final = request_logger if self.logged_on is None else self.logged_on
|
||||
await success_dispatcher(logger).dispatch_success_handlers(
|
||||
self.assembled, cache_hit=self.cache_hit, start_time=None, end_time=None, prefer_async_handlers=True
|
||||
)
|
||||
|
||||
|
||||
def parked_success(parked: tuple[object, ...]) -> ParkedSuccess | None:
|
||||
match parked:
|
||||
case (assembled, cache_hit):
|
||||
return ParkedSuccess(assembled, cache_hit)
|
||||
case (assembled, cache_hit, LiteLLMLoggingObj() as logged_on):
|
||||
return ParkedSuccess(assembled, cache_hit, logged_on)
|
||||
case _:
|
||||
return None
|
||||
|
|
@ -2,11 +2,18 @@ import asyncio
|
|||
from collections.abc import Callable, Coroutine, Sequence
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from typing import Final, Protocol, TypeAlias
|
||||
from typing import Final, 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.parked_stream_dispatch import (
|
||||
DISPATCH_RELEASE,
|
||||
PARKED_DISPATCH,
|
||||
SERVED_FROM_CACHE,
|
||||
parked_success,
|
||||
success_dispatcher,
|
||||
)
|
||||
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
|
||||
|
|
@ -16,25 +23,6 @@ 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:
|
||||
|
|
@ -198,8 +186,8 @@ class StreamedRoundBilling:
|
|||
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)
|
||||
if getattr(logging_obj, DISPATCH_RELEASE, None) is None:
|
||||
setattr(logging_obj, DISPATCH_RELEASE, self._run_parked)
|
||||
|
||||
def complete_round(
|
||||
self,
|
||||
|
|
@ -286,7 +274,7 @@ class StreamedRoundBilling:
|
|||
|
||||
def _price(self, response: ResponsesAPIResponse, round_parked: Parked | None) -> UncachedPrice:
|
||||
if served_from_cache(response):
|
||||
return _SERVED_FROM_CACHE
|
||||
return SERVED_FROM_CACHE
|
||||
match round_parked:
|
||||
case (ModelResponse() as assembled, bool()):
|
||||
return self._logging_obj.uncached_price(assembled)
|
||||
|
|
@ -297,7 +285,7 @@ class StreamedRoundBilling:
|
|||
self._banked = parked
|
||||
self._banked_row = row
|
||||
if parked is not None:
|
||||
setattr(self._logging_obj, _PARKED_DISPATCH, parked)
|
||||
setattr(self._logging_obj, PARKED_DISPATCH, parked)
|
||||
|
||||
def _bank_apart_from_later_rounds(self, billed: _BilledRounds) -> None:
|
||||
row: Final = self._banked_row
|
||||
|
|
@ -320,7 +308,7 @@ class StreamedRoundBilling:
|
|||
@staticmethod
|
||||
async def _log_success(success: _SuccessLog) -> None:
|
||||
success.carry_bridged_usage()
|
||||
await _success_dispatcher(success.logger).dispatch_success_handlers(
|
||||
await success_dispatcher(success.logger).dispatch_success_handlers(
|
||||
success.row.logged,
|
||||
start_time=success.logger.start_time,
|
||||
end_time=success.row.end_time,
|
||||
|
|
@ -329,10 +317,10 @@ class StreamedRoundBilling:
|
|||
)
|
||||
|
||||
def _take_parked(self) -> Parked | None:
|
||||
parked: Final[Parked | None] = getattr(self._logging_obj, _PARKED_DISPATCH, 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)
|
||||
setattr(self._logging_obj, PARKED_DISPATCH, None)
|
||||
return parked
|
||||
|
||||
async def _run_parked(self, *parked: object) -> None:
|
||||
|
|
@ -341,29 +329,27 @@ class StreamedRoundBilling:
|
|||
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 _ if (success := parked_success(parked)) is not None:
|
||||
await success.dispatch(self._logging_obj)
|
||||
case _:
|
||||
return
|
||||
|
||||
def _owns_release(self) -> bool:
|
||||
release: Final[object] = getattr(self._logging_obj, _DISPATCH_RELEASE, None)
|
||||
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)
|
||||
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)
|
||||
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)
|
||||
setattr(self._logging_obj, DISPATCH_RELEASE, None)
|
||||
|
||||
@staticmethod
|
||||
def _close(parked: Parked | None) -> None:
|
||||
|
|
|
|||
|
|
@ -104,11 +104,12 @@ def gateway(_llm_gateway: tuple[Gateway, Mapping[str, str]]) -> Iterator[Gateway
|
|||
|
||||
|
||||
ROUND_USAGE_GUARDRAIL: Final = "mcp-llm-round-usage"
|
||||
CONTENT_FILTER: Final = f"mcp-llm-content-filter-{uuid.uuid4().hex[:8]}"
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def _guarded_llm_gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[Gateway, Mapping[str, str]]]:
|
||||
guardrail: Final[Mapping[str, JsonValue]] = {
|
||||
round_usage: Final[Mapping[str, JsonValue]] = {
|
||||
"guardrail_name": ROUND_USAGE_GUARDRAIL,
|
||||
"litellm_params": {
|
||||
"guardrail": "integration._support.hook_count_probe.UsageRecordingPostCallGuardrail",
|
||||
|
|
@ -116,7 +117,17 @@ def _guarded_llm_gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[t
|
|||
"default_on": True,
|
||||
},
|
||||
}
|
||||
with _private_gateway(tmp_path_factory.mktemp("mcp-llm-guarded-gateway"), (guardrail,)) as served:
|
||||
content_filter: Final[Mapping[str, JsonValue]] = {
|
||||
"guardrail_name": CONTENT_FILTER,
|
||||
"litellm_params": {
|
||||
"guardrail": "litellm_content_filter",
|
||||
"mode": "post_call",
|
||||
"default_on": False,
|
||||
"blocked_words": [{"keyword": "forbidden-keeper", "action": "BLOCK"}],
|
||||
},
|
||||
}
|
||||
guardrails: Final = (round_usage, content_filter)
|
||||
with _private_gateway(tmp_path_factory.mktemp("mcp-llm-guarded-gateway"), guardrails) as served:
|
||||
yield served
|
||||
|
||||
|
||||
|
|
@ -1153,6 +1164,59 @@ def test_streaming_chat_auto_execute_whose_follow_up_fails_runs_the_post_call_gu
|
|||
]
|
||||
|
||||
|
||||
def _chat_rows(call_id: str, statuses: Sequence[str]) -> list[dict[str, JsonValue]]:
|
||||
return eventually(
|
||||
lambda: read_rows(
|
||||
'SELECT status, metadata FROM "LiteLLM_SpendLogs" WHERE litellm_call_id=%s AND call_type=%s',
|
||||
(call_id, "acompletion"),
|
||||
),
|
||||
lambda found: _statuses(found) == sorted(statuses),
|
||||
seconds=70,
|
||||
)
|
||||
|
||||
|
||||
def _entries_of(guardrail_name: str, row: Mapping[str, JsonValue]) -> tuple[dict[str, JsonValue], ...]:
|
||||
return tuple(
|
||||
entry
|
||||
for entry in _guardrail_information(row["metadata"])
|
||||
if isinstance(entry, dict) and entry.get("guardrail_name") == guardrail_name
|
||||
)
|
||||
|
||||
|
||||
def test_streaming_chat_auto_execute_whose_follow_up_fails_counts_its_post_call_content_filter_once(
|
||||
guarded_gateway: Gateway,
|
||||
) -> None:
|
||||
billing: Final = Billing(FIRST_ROUND, FOLLOW_UP_ROUND, follow_up_fails=True)
|
||||
with _rig(guarded_gateway, "chat", billing=billing, prices=PRICES) as rig:
|
||||
response: Final = rig.send(
|
||||
_granted_key(rig),
|
||||
[AUTO],
|
||||
stream=True,
|
||||
stream_options={"include_usage": True},
|
||||
guardrails=[CONTENT_FILTER],
|
||||
)
|
||||
assert FOLLOW_UP_OUTAGE in response.text, response.text
|
||||
rows: Final = _chat_rows(response.headers["x-litellm-call-id"], ("success", "failure"))
|
||||
filtered: Final = tuple(
|
||||
(row["status"], entry.get("guardrail_status")) for row in rows for entry in _entries_of(CONTENT_FILTER, row)
|
||||
)
|
||||
assert filtered == (("failure", "guardrail_failed_to_respond"),), rows
|
||||
(guardrail_id,) = frozenset(
|
||||
str(entry.get("guardrail_id") or CONTENT_FILTER)
|
||||
for row in rows
|
||||
for entry in _entries_of(CONTENT_FILTER, row)
|
||||
)
|
||||
counted: Final = eventually(
|
||||
lambda: read_rows(
|
||||
'SELECT requests_evaluated, flagged_count FROM "LiteLLM_DailyGuardrailMetrics" WHERE guardrail_id=%s',
|
||||
(guardrail_id,),
|
||||
),
|
||||
lambda found: bool(found),
|
||||
seconds=70,
|
||||
)
|
||||
assert counted == [{"requests_evaluated": 1, "flagged_count": 1}], counted
|
||||
|
||||
|
||||
@pytest.mark.parametrize("include_usage", (True, False), ids=("include-usage", "no-usage"))
|
||||
def test_streaming_chat_auto_execute_answered_in_one_round_ends_with_the_usage_it_billed(
|
||||
gateway: Gateway, include_usage: bool
|
||||
|
|
@ -1430,6 +1494,27 @@ def test_streaming_responses_follow_up_broken_mid_stream_bills_round_one_once_an
|
|||
assert float(str(failure["spend"])) == pytest.approx(_price(failed_round)), rows
|
||||
|
||||
|
||||
def test_streaming_chat_follow_up_broken_mid_answer_bills_round_one_and_its_partial_answer_with_one_fixed_margin(
|
||||
gateway: Gateway,
|
||||
) -> None:
|
||||
billing: Final = Billing(FIRST_ROUND, FOLLOW_UP_ROUND, follow_up_breaks=True)
|
||||
with _rig(gateway, "chat", upstream="messages_bridge", billing=billing, prices=PRICES) as rig:
|
||||
response: Final = rig.send(_granted_key(rig), [AUTO], stream=True, stream_options={"include_usage": True})
|
||||
assert "the sum" in response.text, response.text
|
||||
assert ANSWER not in response.text, response.text
|
||||
assert len(rig.upstream_tools()) == 2
|
||||
round_one: Final = _price(FIRST_ROUND) + FIXED_MARGIN
|
||||
rows: Final = _model_rows(response.headers["x-litellm-call-id"], _logs_a_failure_billing(round_one))
|
||||
assert _statuses(rows) == ["failure", "success"], rows
|
||||
(success,) = (row for row in rows if row["status"] == "success")
|
||||
(failure,) = (row for row in rows if row["status"] == "failure")
|
||||
assert (success["prompt_tokens"], success["completion_tokens"]) == FIRST_ROUND, rows
|
||||
assert float(str(success["spend"])) == pytest.approx(round_one), rows
|
||||
partial_answer: Final = (int(str(failure["prompt_tokens"])), int(str(failure["completion_tokens"])))
|
||||
assert min(partial_answer) > 0, rows
|
||||
assert float(str(failure["spend"])) == pytest.approx(_price(partial_answer)), rows
|
||||
|
||||
|
||||
@pytest.mark.parametrize("upstream", STREAMED_UPSTREAMS, ids=STREAMED_IDS)
|
||||
def test_streaming_responses_blocked_by_a_post_call_guardrail_bills_every_round_once(
|
||||
gateway: Gateway, upstream: StreamedUpstream
|
||||
|
|
|
|||
|
|
@ -9056,6 +9056,34 @@ class TestStreamingClientDisconnectBilling:
|
|||
|
||||
proxy_logging_obj.arelease_max_parallel_requests_on_disconnect.assert_awaited_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_disconnect_whose_in_flight_round_billing_raises_still_closes_the_stream_and_frees_the_slot(self):
|
||||
import types
|
||||
|
||||
response = await self._start_partial_stream()
|
||||
closed: Final = AsyncMock()
|
||||
|
||||
def cannot_price_the_in_flight_round() -> None:
|
||||
raise RuntimeError("the in-flight round could not be priced")
|
||||
|
||||
stream: Final = types.SimpleNamespace(
|
||||
chunks=[], messages=None, aclose=closed, bill_in_flight_round_on_disconnect=cannot_price_the_in_flight_round
|
||||
)
|
||||
proxy_logging_obj: Final = types.SimpleNamespace(arelease_max_parallel_requests_on_disconnect=AsyncMock())
|
||||
|
||||
await ProxyBaseLLMRequestProcessing._finalize_streaming_generator_cleanup(
|
||||
request=None,
|
||||
request_data={"litellm_logging_obj": response.logging_obj},
|
||||
response=stream,
|
||||
stream_completed=False,
|
||||
client_disconnected=True,
|
||||
user_api_key_dict=MagicMock(),
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
|
||||
closed.assert_awaited_once()
|
||||
proxy_logging_obj.arelease_max_parallel_requests_on_disconnect.assert_awaited_once()
|
||||
|
||||
async def _bill_and_collect_success_event(self, prepare=None, request_data=None):
|
||||
recorder = _RecordingSuccessLogger()
|
||||
original_callbacks = litellm.callbacks
|
||||
|
|
|
|||
|
|
@ -1562,6 +1562,7 @@ _CHAT_ROUND_TWO: Final = MappingProxyType(
|
|||
),
|
||||
}
|
||||
)
|
||||
_NEVER_LOGGED_SECONDS: Final = 60
|
||||
_CONTEXT_WINDOW_ERROR: Final = MappingProxyType(
|
||||
{
|
||||
"error": MappingProxyType(
|
||||
|
|
@ -1581,16 +1582,31 @@ class _ChatRoundRecorder(CustomLogger):
|
|||
super().__init__()
|
||||
self.logged: tuple[StandardLoggingPayload, ...] = ()
|
||||
self.failed: tuple[StandardLoggingPayload, ...] = ()
|
||||
self._recorded: Final = asyncio.Condition()
|
||||
|
||||
async def async_log_success_event(
|
||||
self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime
|
||||
) -> None:
|
||||
self.logged = (*self.logged, cast(StandardLoggingPayload, kwargs["standard_logging_object"]))
|
||||
await self._announce()
|
||||
|
||||
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"]))
|
||||
await self._announce()
|
||||
|
||||
async def _announce(self) -> None:
|
||||
async with self._recorded:
|
||||
self._recorded.notify_all()
|
||||
|
||||
async def settled(self, logged: int = 0, failed: int = 0) -> None:
|
||||
async with self._recorded:
|
||||
await asyncio.wait_for(
|
||||
self._recorded.wait_for(lambda: len(self.logged) >= logged and len(self.failed) >= failed),
|
||||
timeout=_NEVER_LOGGED_SECONDS,
|
||||
)
|
||||
await _settled()
|
||||
|
||||
|
||||
def _chat_round_cost(body: Mapping[str, object]) -> float:
|
||||
|
|
@ -1606,14 +1622,6 @@ async def _settled() -> None:
|
|||
await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10)
|
||||
|
||||
|
||||
async def _settled_once(logged: Callable[[], bool]) -> None:
|
||||
for _ in range(200):
|
||||
if logged():
|
||||
break
|
||||
await asyncio.sleep(0.01)
|
||||
await _settled()
|
||||
|
||||
|
||||
def _serve_chat_rounds(
|
||||
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, rounds: Sequence[Response]
|
||||
) -> None:
|
||||
|
|
@ -1785,7 +1793,9 @@ def _cut_once_the_answer_starts() -> Response:
|
|||
return Response(200, stream=_CutOnceTheAnswerStarts(), headers={"content-type": "text/event-stream"})
|
||||
|
||||
|
||||
async def _open_chat_echo_stream(logging_obj: Logging) -> CustomStreamWrapper:
|
||||
async def _open_chat_echo_stream(
|
||||
logging_obj: Logging, metadata: dict[str, object] | None = None
|
||||
) -> CustomStreamWrapper:
|
||||
stream: Final = await litellm.acompletion(
|
||||
model="openai/gpt-5.6",
|
||||
messages=[{"role": "user", "content": "Call echo with PING"}],
|
||||
|
|
@ -1795,6 +1805,7 @@ async def _open_chat_echo_stream(logging_obj: Logging) -> CustomStreamWrapper:
|
|||
stream=True,
|
||||
stream_options={"include_usage": True},
|
||||
litellm_logging_obj=logging_obj,
|
||||
metadata=metadata,
|
||||
)
|
||||
assert isinstance(stream, CustomStreamWrapper)
|
||||
return stream
|
||||
|
|
@ -1860,7 +1871,7 @@ async def test_stream_auto_execute_deferred_for_a_post_call_guardrail_logs_both_
|
|||
await _settled()
|
||||
logged_before_release: Final = recorder.logged
|
||||
ProxyLogging.fire_deferred_stream_logging({"litellm_logging_obj": logging_obj})
|
||||
await _settled_once(lambda: len(recorder.logged) > 0)
|
||||
await recorder.settled(logged=1)
|
||||
|
||||
assert logged_before_release == ()
|
||||
_assert_one_stream_row_bills_both_rounds_once(recorder.logged, chunks)
|
||||
|
|
@ -1888,7 +1899,7 @@ async def test_stream_auto_execute_blocked_by_a_post_call_guardrail_bills_both_r
|
|||
await _read_all(stream)
|
||||
assert ProxyLogging._discard_deferred_stream_logging_for_failure({"litellm_logging_obj": logging_obj}, blocked)
|
||||
await logging_obj.async_failure_handler(blocked, "blocked by output-policy")
|
||||
await _settled_once(lambda: len(recorder.failed) > 0)
|
||||
await recorder.settled(failed=1)
|
||||
|
||||
charged_once: Final = _chat_round_cost(_CHAT_ROUND_ONE) + _chat_round_cost(_CHAT_ROUND_TWO) - _FIXED_MARGIN
|
||||
assert recorder.logged == ()
|
||||
|
|
@ -1910,7 +1921,7 @@ async def test_stream_auto_execute_answered_without_a_tool_call_logs_its_one_rou
|
|||
async for chunk in stream:
|
||||
if chunk.choices and chunk.choices[0].finish_reason is not None:
|
||||
break
|
||||
await _settled_once(lambda: len(recorder.logged) > 0)
|
||||
await recorder.settled(logged=1)
|
||||
logged_once_the_answer_finished: Final = recorder.logged
|
||||
after_the_answer: Final = await _read_all(stream)
|
||||
await _settled()
|
||||
|
|
@ -1947,11 +1958,30 @@ async def test_stream_auto_execute_bills_the_finished_first_round_when_the_calle
|
|||
|
||||
with _echo_tool_served():
|
||||
await stop_reading_once_the_follow_up_answers()
|
||||
await _settled_once(lambda: len(recorder.logged) > 0)
|
||||
await recorder.settled(logged=1)
|
||||
|
||||
_assert_one_row_bills_round_one(recorder.logged)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_auto_execute_whose_first_round_breaks_mid_answer_bills_only_its_failure_row(
|
||||
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
|
||||
):
|
||||
_serve_chat_rounds(respx_mock, monkeypatch, (_cut_once_the_answer_starts(),))
|
||||
recorder: Final = _ChatRoundRecorder()
|
||||
logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-first-round-breaks", stream=True)
|
||||
|
||||
with _echo_tool_served():
|
||||
stream: Final = await _open_chat_echo_stream(logging_obj)
|
||||
with pytest.raises(litellm.ServiceUnavailableError):
|
||||
await _read_all(stream)
|
||||
await recorder.settled(failed=1)
|
||||
|
||||
assert recorder.logged == ()
|
||||
assert len(recorder.failed) == 1
|
||||
assert recorder.failed[0]["completion_tokens"] > 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("deferred_for_a_post_call_guardrail", [False, True])
|
||||
async def test_stream_auto_execute_bills_a_broken_follow_up_on_its_own_failure_row_with_no_second_fixed_margin(
|
||||
|
|
@ -1976,7 +2006,7 @@ async def test_stream_auto_execute_bills_a_broken_follow_up_on_its_own_failure_r
|
|||
await _read_all(stream)
|
||||
if not ProxyLogging._discard_deferred_stream_logging_for_failure(dict(request_data), broken.value):
|
||||
ProxyLogging.fire_deferred_stream_logging(dict(request_data))
|
||||
await _settled_once(lambda: len(recorder.logged) > 0 and len(recorder.failed) > 0)
|
||||
await recorder.settled(logged=1, failed=1)
|
||||
|
||||
_assert_one_row_bills_round_one(recorder.logged)
|
||||
assert len(recorder.failed) == 1
|
||||
|
|
@ -2004,23 +2034,67 @@ async def test_stream_auto_execute_bills_a_broken_follow_up_on_its_own_failure_r
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_auto_execute_leaves_no_deferred_log_armed_for_the_proxy_to_orphan_after_its_first_chunk(
|
||||
@pytest.mark.parametrize("deferred_for_a_post_call_guardrail", [False, True])
|
||||
async def test_stream_auto_execute_whose_follow_up_breaks_reports_the_request_guardrail_result_on_the_failure_row_only(
|
||||
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, deferred_for_a_post_call_guardrail: bool
|
||||
):
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
|
||||
_serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _cut_once_the_answer_starts()))
|
||||
recorder: Final = _ChatRoundRecorder()
|
||||
logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-follow-up-breaks-guarded", stream=True)
|
||||
metadata: Final[dict[str, object]] = {}
|
||||
request_data: Final = {"metadata": metadata, "litellm_logging_obj": logging_obj}
|
||||
content_filter: Final = CustomGuardrail(guardrail_name="request-content-filter", event_hook="post_call")
|
||||
|
||||
with _echo_tool_served():
|
||||
stream: Final = await _open_chat_echo_stream(logging_obj, metadata)
|
||||
if deferred_for_a_post_call_guardrail:
|
||||
ProxyBaseLLMRequestProcessing(data={})._arm_deferred_stream_dispatch(
|
||||
response=stream, route_type="acompletion", user_api_key_dict=MagicMock(), logging_obj=logging_obj
|
||||
)
|
||||
with pytest.raises(litellm.ServiceUnavailableError) as broken:
|
||||
await _read_all(stream)
|
||||
content_filter.add_standard_logging_guardrail_information_to_request_data(
|
||||
guardrail_json_response=str(broken.value),
|
||||
request_data=request_data,
|
||||
guardrail_status="guardrail_failed_to_respond",
|
||||
)
|
||||
if not ProxyLogging._discard_deferred_stream_logging_for_failure(request_data, broken.value):
|
||||
ProxyLogging.fire_deferred_stream_logging(request_data)
|
||||
await recorder.settled(logged=1, failed=1)
|
||||
|
||||
_assert_one_row_bills_round_one(recorder.logged)
|
||||
assert recorder.logged[0]["guardrail_information"] is None
|
||||
assert [entry["guardrail_name"] for entry in recorder.failed[0]["guardrail_information"] or ()] == [
|
||||
"request-content-filter"
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_auto_execute_served_through_the_proxy_stream_hook_logs_both_rounds_once_when_it_ends(
|
||||
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
|
||||
):
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
|
||||
monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}})
|
||||
_serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _streamed(_CHAT_ROUND_TWO)))
|
||||
recorder: Final = _ChatRoundRecorder()
|
||||
logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-first-chunk-read-ahead", stream=True)
|
||||
logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-proxy-stream-hook", stream=True)
|
||||
|
||||
with _echo_tool_served():
|
||||
stream: Final = await _open_chat_echo_stream(logging_obj)
|
||||
rest: Final = stream.__aiter__()
|
||||
first_chunk: Final = await rest.__anext__()
|
||||
armed_after_the_first_chunk: Final = getattr(logging_obj, "_on_deferred_stream_complete", None)
|
||||
chunks: Final = (first_chunk, *[chunk async for chunk in rest])
|
||||
await _settled_once(lambda: len(recorder.logged) > 0)
|
||||
served: Final = ProxyLogging(user_api_key_cache=DualCache()).async_post_call_streaming_iterator_hook(
|
||||
response=await _open_chat_echo_stream(logging_obj),
|
||||
user_api_key_dict=UserAPIKeyAuth(),
|
||||
request_data={"litellm_logging_obj": logging_obj},
|
||||
)
|
||||
chunks: Final = tuple([chunk async for chunk in served])
|
||||
await recorder.settled(logged=1)
|
||||
|
||||
assert armed_after_the_first_chunk is None
|
||||
_assert_one_stream_row_bills_both_rounds_once(recorder.logged, chunks)
|
||||
|
||||
|
||||
|
|
@ -2035,17 +2109,152 @@ async def test_stream_auto_execute_whose_follow_up_streams_nothing_still_bills_t
|
|||
|
||||
with _echo_tool_served():
|
||||
await _read_all(await _open_chat_echo_stream(logging_obj))
|
||||
await _settled_once(lambda: len(recorder.logged) > 0)
|
||||
await recorder.settled(logged=1)
|
||||
|
||||
_assert_one_row_bills_round_one(recorder.logged)
|
||||
|
||||
|
||||
async def _stream_read_until_the_follow_up_answers(
|
||||
logging_obj: Logging, opened: "queue.SimpleQueue[CustomStreamWrapper]"
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_auto_execute_whose_follow_up_cannot_be_assembled_ends_with_the_usage_of_the_round_it_billed(
|
||||
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
|
||||
):
|
||||
_serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _streamed(_CHAT_ROUND_TWO)))
|
||||
recorder: Final = _ChatRoundRecorder()
|
||||
logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-follow-up-unassembled", stream=True)
|
||||
|
||||
def cannot_assemble(*args: object, **kwargs: object) -> ModelResponse:
|
||||
raise ValueError("the follow-up's chunks cannot be assembled")
|
||||
|
||||
with _echo_tool_served():
|
||||
stream: Final = await _open_chat_echo_stream(logging_obj)
|
||||
await _read_until(stream, "PING")
|
||||
monkeypatch.setattr(litellm, "stream_chunk_builder", cannot_assemble)
|
||||
rest: Final = await _read_all(stream)
|
||||
await recorder.settled(logged=1)
|
||||
|
||||
_assert_one_row_bills_round_one(recorder.logged)
|
||||
usage: Final = _streamed_usage(rest)
|
||||
assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (71, 19, 90)
|
||||
assert usage.cost == pytest.approx(recorder.logged[0]["response_cost"])
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("deferred_for_a_post_call_guardrail", [False, True])
|
||||
async def test_stream_auto_execute_hung_up_during_the_follow_up_bills_round_one_and_the_partial_answer_on_one_row(
|
||||
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, deferred_for_a_post_call_guardrail: bool
|
||||
):
|
||||
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
|
||||
|
||||
monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}})
|
||||
_serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _streamed(_CHAT_ROUND_TWO)))
|
||||
recorder: Final = _ChatRoundRecorder()
|
||||
logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-hung-up-mid-answer", stream=True)
|
||||
|
||||
with _echo_tool_served():
|
||||
stream: Final = await _open_chat_echo_stream(logging_obj)
|
||||
if deferred_for_a_post_call_guardrail:
|
||||
ProxyBaseLLMRequestProcessing(data={})._arm_deferred_stream_dispatch(
|
||||
response=stream, route_type="acompletion", user_api_key_dict=MagicMock(), logging_obj=logging_obj
|
||||
)
|
||||
await _read_until(stream, "PING")
|
||||
await ProxyBaseLLMRequestProcessing.finalize_streaming_generator_cleanup(
|
||||
request=None, request_data={"litellm_logging_obj": logging_obj}, response=stream, client_disconnected=True
|
||||
)
|
||||
await recorder.settled(logged=1)
|
||||
|
||||
assert len(recorder.logged) == 1
|
||||
row: Final = recorder.logged[0]
|
||||
assert row["id"] == _CHAT_ROUND_ONE["id"]
|
||||
partial_prompt, partial_completion = (row["prompt_tokens"] - 71, row["completion_tokens"] - 19)
|
||||
assert partial_prompt > 0 and partial_completion > 0
|
||||
partial_answer: Final = litellm.completion_cost(
|
||||
completion_response=ModelResponse(
|
||||
model="gpt-5.6",
|
||||
usage=Usage(
|
||||
prompt_tokens=partial_prompt,
|
||||
completion_tokens=partial_completion,
|
||||
total_tokens=partial_prompt + partial_completion,
|
||||
),
|
||||
),
|
||||
model="gpt-5.6",
|
||||
custom_llm_provider="openai",
|
||||
)
|
||||
assert row["response_cost"] == pytest.approx(_chat_round_cost(_CHAT_ROUND_ONE) + partial_answer - _FIXED_MARGIN)
|
||||
assert recorder.failed == ()
|
||||
|
||||
|
||||
async def _read_through_the_usage_chunk(
|
||||
chunks: AsyncIterator[ModelResponseStream], read: tuple[ModelResponseStream, ...] = ()
|
||||
) -> tuple[ModelResponseStream, ...]:
|
||||
chunk: Final = await anext(chunks)
|
||||
if isinstance(getattr(chunk, "usage", None), Usage):
|
||||
return (*read, chunk)
|
||||
return await _read_through_the_usage_chunk(chunks, (*read, chunk))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("deferred_for_a_post_call_guardrail", [False, True])
|
||||
async def test_stream_auto_execute_hung_up_after_its_whole_answer_bills_both_rounds_once(
|
||||
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, deferred_for_a_post_call_guardrail: bool
|
||||
):
|
||||
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
|
||||
|
||||
monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}})
|
||||
_serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _streamed(_CHAT_ROUND_TWO)))
|
||||
recorder: Final = _ChatRoundRecorder()
|
||||
logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-hung-up-after-answer", stream=True)
|
||||
|
||||
with _echo_tool_served():
|
||||
stream: Final = await _open_chat_echo_stream(logging_obj)
|
||||
if deferred_for_a_post_call_guardrail:
|
||||
ProxyBaseLLMRequestProcessing(data={})._arm_deferred_stream_dispatch(
|
||||
response=stream, route_type="acompletion", user_api_key_dict=MagicMock(), logging_obj=logging_obj
|
||||
)
|
||||
chunks: Final = await _read_through_the_usage_chunk(aiter(stream))
|
||||
await ProxyBaseLLMRequestProcessing.finalize_streaming_generator_cleanup(
|
||||
request=None, request_data={"litellm_logging_obj": logging_obj}, response=stream, client_disconnected=True
|
||||
)
|
||||
await recorder.settled(logged=1)
|
||||
|
||||
_assert_one_stream_row_bills_both_rounds_once(recorder.logged, chunks)
|
||||
assert recorder.failed == ()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_auto_execute_bridged_to_messages_bills_round_one_when_its_follow_up_breaks(
|
||||
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
|
||||
):
|
||||
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
|
||||
|
||||
_serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _cut_once_the_answer_starts()))
|
||||
recorder: Final = _ChatRoundRecorder()
|
||||
logging_obj: Final = _chat_rounds_logging_obj(recorder, "mcp-chat-stream-bridged-to-messages", stream=True)
|
||||
|
||||
async def bridged(stream: CustomStreamWrapper) -> AsyncIterator[ModelResponseStream]:
|
||||
async for chunk in stream:
|
||||
yield chunk
|
||||
|
||||
with _echo_tool_served():
|
||||
relayed: Final = bridged(await _open_chat_echo_stream(logging_obj))
|
||||
ProxyBaseLLMRequestProcessing(data={})._arm_deferred_stream_dispatch(
|
||||
response=relayed, route_type="anthropic_messages", user_api_key_dict=MagicMock(), logging_obj=logging_obj
|
||||
)
|
||||
with pytest.raises(litellm.ServiceUnavailableError):
|
||||
_ = [chunk async for chunk in relayed]
|
||||
await recorder.settled(logged=1, failed=1)
|
||||
|
||||
_assert_one_row_bills_round_one(recorder.logged)
|
||||
|
||||
|
||||
def _read_synchronously_until_the_follow_up_answers(
|
||||
logging_obj: Logging, loop: asyncio.AbstractEventLoop, before_dropping: Callable[[], None] = lambda: None
|
||||
) -> None:
|
||||
stream: Final = await _open_chat_echo_stream(logging_obj)
|
||||
await _read_until(stream, "PING")
|
||||
opened.put(stream)
|
||||
stream: Final = loop.run_until_complete(_open_chat_echo_stream(logging_obj))
|
||||
for chunk in stream:
|
||||
if chunk.choices and chunk.choices[0].delta.content == "PING":
|
||||
before_dropping()
|
||||
return
|
||||
pytest.fail("the stream never sent 'PING'")
|
||||
|
||||
|
||||
_CHAT_STREAM_BILLING_NAMES: Final = ("MCPStreamingIterator", "ChatStreamRoundBilling", "_log_success")
|
||||
|
|
@ -2063,26 +2272,57 @@ class _UnraisableRecorder:
|
|||
self.raised = (*self.raised, f"{unraisable.object!r}: {unraisable.exc_value!r}")
|
||||
|
||||
|
||||
def test_stream_auto_execute_dropped_after_its_event_loop_closed_discards_its_parked_log_cleanly(
|
||||
def test_stream_auto_execute_read_synchronously_and_dropped_mid_answer_bills_the_finished_first_round(
|
||||
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
|
||||
):
|
||||
_serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _streamed(_CHAT_ROUND_TWO)))
|
||||
recorder: Final = _ChatRoundRecorder()
|
||||
unraisable: Final = _UnraisableRecorder()
|
||||
opened: Final = queue.SimpleQueue[CustomStreamWrapper]()
|
||||
loop: Final = asyncio.new_event_loop()
|
||||
gc.collect()
|
||||
monkeypatch.setattr(sys, "unraisablehook", unraisable)
|
||||
asyncio.set_event_loop(loop)
|
||||
|
||||
with _echo_tool_served(), warnings.catch_warnings(record=True) as caught:
|
||||
warnings.simplefilter("always")
|
||||
asyncio.run(
|
||||
_stream_read_until_the_follow_up_answers(
|
||||
_chat_rounds_logging_obj(recorder, "mcp-chat-stream-loop-closed", stream=True), opened
|
||||
try:
|
||||
with _echo_tool_served(), warnings.catch_warnings(record=True) as caught:
|
||||
warnings.simplefilter("always")
|
||||
_read_synchronously_until_the_follow_up_answers(
|
||||
_chat_rounds_logging_obj(recorder, "mcp-chat-stream-sync-dropped", stream=True), loop
|
||||
)
|
||||
)
|
||||
opened.get_nowait()
|
||||
gc.collect()
|
||||
gc.collect()
|
||||
loop.run_until_complete(recorder.settled(logged=1))
|
||||
finally:
|
||||
asyncio.set_event_loop(None)
|
||||
loop.close()
|
||||
|
||||
_assert_one_row_bills_round_one(recorder.logged)
|
||||
assert [raised for raised in unraisable.raised if _names_chat_stream_billing(raised)] == []
|
||||
assert [str(warning.message) for warning in caught if _names_chat_stream_billing(str(warning.message))] == []
|
||||
|
||||
|
||||
def test_stream_auto_execute_read_synchronously_and_dropped_after_its_loop_closed_drops_without_errors(
|
||||
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
|
||||
):
|
||||
_serve_chat_rounds(respx_mock, monkeypatch, (_streamed(_CHAT_ROUND_ONE), _streamed(_CHAT_ROUND_TWO)))
|
||||
recorder: Final = _ChatRoundRecorder()
|
||||
unraisable: Final = _UnraisableRecorder()
|
||||
loop: Final = asyncio.new_event_loop()
|
||||
gc.collect()
|
||||
monkeypatch.setattr(sys, "unraisablehook", unraisable)
|
||||
asyncio.set_event_loop(loop)
|
||||
|
||||
try:
|
||||
with _echo_tool_served(), warnings.catch_warnings(record=True) as caught:
|
||||
warnings.simplefilter("always")
|
||||
_read_synchronously_until_the_follow_up_answers(
|
||||
_chat_rounds_logging_obj(recorder, "mcp-chat-stream-sync-dropped-after-close", stream=True),
|
||||
loop,
|
||||
before_dropping=loop.close,
|
||||
)
|
||||
gc.collect()
|
||||
finally:
|
||||
asyncio.set_event_loop(None)
|
||||
loop.close()
|
||||
|
||||
assert recorder.logged == ()
|
||||
assert [raised for raised in unraisable.raised if _names_chat_stream_billing(raised)] == []
|
||||
assert [str(warning.message) for warning in caught if _names_chat_stream_billing(str(warning.message))] == []
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue