diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261029000000_add_background_response_settlement/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261029000000_add_background_response_settlement/migration.sql new file mode 100644 index 00000000000..4965ed3cfc4 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261029000000_add_background_response_settlement/migration.sql @@ -0,0 +1,9 @@ +SET LOCAL lock_timeout = '10s'; + +-- CreateTable +CREATE TABLE IF NOT EXISTS "LiteLLM_BackgroundResponseSettlement" ( + "response_id" TEXT NOT NULL, + "settled_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + + CONSTRAINT "LiteLLM_BackgroundResponseSettlement_pkey" PRIMARY KEY ("response_id") +); diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index e63580da721..37783275d9b 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -2055,6 +2055,11 @@ model LiteLLM_BackgroundInteractionSettlement { @@index([settled_at], map: "idx_background_interaction_settlement_settled_at") } +model LiteLLM_BackgroundResponseSettlement { + response_id String @id + settled_at DateTime @default(now()) +} + model LiteLLM_Lens { reviews LiteLLM_LensReview[] id String @id diff --git a/litellm/constants.py b/litellm/constants.py index 8df92c46630..b29c22810e7 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1969,6 +1969,7 @@ BATCH_LINE_ITEMS_MAX_PART_BYTES: Final = get_env_int_in_range( MANAGED_OBJECT_STALENESS_CUTOFF_DAYS: Final = max(1, get_env_int("MANAGED_OBJECT_STALENESS_CUTOFF_DAYS", 7)) PENDING_RESPONSE_STATUSES: Final = ("queued", "in_progress") TERMINAL_RESPONSE_STATUSES: Final = frozenset({"completed", "failed", "cancelled", "incomplete"}) +BACKGROUND_RESPONSE_SETTLEMENT_TIMEOUT_SECONDS: Final = 2.0 BATCH_OUTPUT_FILE_LOOKUP_TIMEOUT_SECONDS: Final = 10.0 BATCH_OUTPUT_FILE_FALLBACK_RETRY_AFTER_SECONDS: Final = 60 STALE_OBJECT_CLEANUP_BATCH_SIZE: Final = max(1, get_env_int("STALE_OBJECT_CLEANUP_BATCH_SIZE", 1000)) diff --git a/litellm/integrations/otel/model/spans.py b/litellm/integrations/otel/model/spans.py index 0f98a8fce0e..8f9e2c8c6db 100644 --- a/litellm/integrations/otel/model/spans.py +++ b/litellm/integrations/otel/model/spans.py @@ -292,6 +292,7 @@ _PRISMA_MODELS: Final[frozenset[str]] = frozenset( "LiteLLM_SpendLogs", "LiteLLM_BudgetWindowSpend", "LiteLLM_BackgroundInteractionSettlement", + "LiteLLM_BackgroundResponseSettlement", "LiteLLM_ErrorLogs", "LiteLLM_UserNotifications", "LiteLLM_UserProviderCredentials", diff --git a/litellm/llms/openai/responses/transformation.py b/litellm/llms/openai/responses/transformation.py index 314a1d22508..199f9525484 100644 --- a/litellm/llms/openai/responses/transformation.py +++ b/litellm/llms/openai/responses/transformation.py @@ -1136,7 +1136,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): return None @staticmethod - def parse_terminal_response_from_stream_chunks(all_chunks: list[str]) -> ResponsesAPIResponse | None: + def parse_terminal_response_from_stream_chunks(all_chunks: Sequence[str]) -> ResponsesAPIResponse | None: terminal_event: Final = OpenAIResponsesAPIConfig.parse_terminal_event_from_stream_chunks(all_chunks) return None if terminal_event is None else terminal_event.response diff --git a/litellm/proxy/pass_through_endpoints/background_response_settlement.py b/litellm/proxy/pass_through_endpoints/background_response_settlement.py new file mode 100644 index 00000000000..c5497e51fa4 --- /dev/null +++ b/litellm/proxy/pass_through_endpoints/background_response_settlement.py @@ -0,0 +1,52 @@ +import asyncio +from collections.abc import Awaitable, Callable, Sequence +from dataclasses import dataclass +from typing import Final, Protocol + +from typing_extensions import ReadOnly, TypedDict + +from litellm._logging import verbose_proxy_logger +from litellm.constants import BACKGROUND_RESPONSE_SETTLEMENT_TIMEOUT_SECONDS +from litellm.repositories.table_repositories import BackgroundResponseSettlementRepository + + +class _NewSettlementRow(TypedDict): + response_id: ReadOnly[str] + + +class BackgroundResponseSettlementTable(Protocol): + def create_many(self, data: Sequence[_NewSettlementRow], *, skip_duplicates: bool) -> Awaitable[int]: ... + + +class BackgroundResponseSettlementStore(Protocol): + async def claim(self, response_id: str) -> bool: ... + + +def _proxy_settlement_table() -> BackgroundResponseSettlementTable | None: + from litellm.proxy.proxy_server import prisma_client # noqa: PLC0415 # runtime global, set after load + + if prisma_client is None: + return None + return BackgroundResponseSettlementRepository(prisma_client).table + + +@dataclass(frozen=True, slots=True) +class PrismaBackgroundResponseSettlementStore: + table: Callable[[], BackgroundResponseSettlementTable | None] = _proxy_settlement_table + timeout_seconds: float = BACKGROUND_RESPONSE_SETTLEMENT_TIMEOUT_SECONDS + + async def claim(self, response_id: str) -> bool: + try: + table: Final = self.table() + if table is None: + return True + inserted: Final = await asyncio.wait_for( + table.create_many([_NewSettlementRow(response_id=response_id)], skip_duplicates=True), + timeout=self.timeout_seconds, + ) + except Exception as e: # noqa: BLE001 # any DB failure must bill as before rather than drop the charge + verbose_proxy_logger.warning( + "Could not record the settlement of background response %s, billing it anyway: %s", response_id, e + ) + return True + return inserted > 0 diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py index e4e3fec855d..0b97820fb0d 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py @@ -124,6 +124,10 @@ def _is_finished_background_response(response_body: Mapping[str, object]) -> boo return is_background_response(response_body) and response_body.get("usage") is not None +def _settles_background_response(terminal_response: ResponsesAPIResponse) -> bool: + return is_background_response(terminal_response) and terminal_response.usage is not None + + def _responses_sub_route(path: str) -> tuple[str, ...] | None: segments: Final = path.rstrip("/").split("/") anchor: Final = max((index for index, segment in enumerate(segments) if segment == "responses"), default=None) @@ -263,7 +267,7 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): @staticmethod def openai_responses_call_billing( - method: str, url_route: str, response_body: Mapping[str, object] + method: str, url_route: str, response_body: Mapping[str, object], settled_before: bool = False ) -> ResponsesCallBilling: if not OpenAIPassthroughLoggingHandler.is_openai_responses_route(url_route): return "unbilled" @@ -271,10 +275,35 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): case ("POST", ("input_tokens",)) | ("GET", (_, "input_items")) | ("DELETE", (_,)): return "unbilled" case ("GET", (_,)) | ("POST", (_, "cancel")): - return "background_settlement" if _is_finished_background_response(response_body) else "unbilled" + settles: Final = not settled_before and _is_finished_background_response(response_body) + return "background_settlement" if settles else "unbilled" case _: return "billed" + @staticmethod + def background_response_settled_by_call( + method: str, url_route: str, response_body: Mapping[str, object] + ) -> str | None: + response_id: Final = response_body.get("id") + billing: Final = OpenAIPassthroughLoggingHandler.openai_responses_call_billing(method, url_route, response_body) + return response_id if billing == "background_settlement" and isinstance(response_id, str) else None + + @staticmethod + def stream_may_settle_background_response(url_route: str, request_body: Mapping[str, object]) -> bool: + sub_route: Final = _responses_sub_route(urlparse(url_route).path) + creates_in_background: Final = sub_route == () and request_body.get("background") is True + resumes_a_response: Final = sub_route is not None and len(sub_route) == 1 + return OpenAIPassthroughLoggingHandler.is_openai_responses_route(url_route) and ( + creates_in_background or resumes_a_response + ) + + @staticmethod + def background_response_settled_by_stream(all_chunks: Sequence[str]) -> str | None: + terminal_response: Final = OpenAIResponsesAPIConfig.parse_terminal_response_from_stream_chunks(all_chunks) + if terminal_response is None or not _settles_background_response(terminal_response): + return None + return terminal_response.id + @staticmethod def is_openai_agents_route(url_route: str) -> bool: return is_openai_compatible_url(url_route) and is_openai_agents_api_path(urlparse(url_route).path) @@ -483,6 +512,7 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): end_time: datetime, cache_hit: bool, request_body: dict, + background_response_settled_before: bool = False, **kwargs, ) -> PassThroughEndpointLoggingTypedDict: """ @@ -494,7 +524,7 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): is_image_generation: Final = OpenAIPassthroughLoggingHandler.is_openai_image_generation_route(url_route) is_image_editing: Final = OpenAIPassthroughLoggingHandler.is_openai_image_editing_route(url_route) responses_billing: Final = OpenAIPassthroughLoggingHandler.openai_responses_call_billing( - httpx_response.request.method, url_route, response_body + httpx_response.request.method, url_route, response_body, background_response_settled_before ) is_responses: Final = responses_billing != "unbilled" is_decisions: Final = ( @@ -801,6 +831,7 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): start_time: datetime, all_chunks: list[str], end_time: datetime, + background_response_settled_before: bool = False, ) -> PassThroughEndpointLoggingTypedDict: """ Handle logging for collected OpenAI streaming chunks with cost tracking. @@ -811,11 +842,16 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): ) try: is_responses: Final = OpenAIPassthroughLoggingHandler.is_openai_responses_route(url_route) - terminal_response: Final = ( + parsed_terminal_response: Final = ( OpenAIResponsesAPIConfig.parse_terminal_response_from_stream_chunks(all_chunks=all_chunks) if is_responses else None ) + terminal_response: Final = ( + parsed_terminal_response.model_copy(update={"usage": None}) + if background_response_settled_before and parsed_terminal_response is not None + else parsed_terminal_response + ) requested_model: Final[object] = request_body.get("model") model: Final = ( requested_model @@ -846,7 +882,9 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): custom_llm_provider: Final = litellm_logging_obj.model_call_details.get("custom_llm_provider", "openai") # Calculate cost using LiteLLM's cost calculator response_cost: Final = ( - litellm.completion_cost( + 0.0 + if background_response_settled_before + else litellm.completion_cost( completion_response=complete_response, model=model, custom_llm_provider=custom_llm_provider, @@ -862,9 +900,9 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): settles_background_response: Final = ( BackgroundResponseSettlement.SETTLES - if terminal_response is not None - and is_background_response(terminal_response) - and terminal_response.usage is not None + if not background_response_settled_before + and terminal_response is not None + and _settles_background_response(terminal_response) else None ) diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py index ef8ffdf640e..345e17c435d 100644 --- a/litellm/proxy/pass_through_endpoints/streaming_handler.py +++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py @@ -5,6 +5,7 @@ from datetime import datetime, timezone from typing import Final, Protocol import httpx +from pydantic import TypeAdapter from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.asyncify import asyncify @@ -42,6 +43,8 @@ from .llm_provider_handlers.vertex_passthrough_logging_handler import ( ) from .success_handler import PassThroughEndpointLogging +_REQUEST_FIELDS: Final = TypeAdapter(dict[str, object]) + class RouteStreamingLogging(Protocol): def __call__( @@ -395,6 +398,15 @@ class PassThroughStreamingHandler: **tinyfish_payload["kwargs"], ) return + background_response_settled_before: Final = ( + endpoint_type == EndpointType.OPENAI + and await PassThroughStreamingHandler._stream_background_response_settled_before( + passthrough_success_handler_obj=passthrough_success_handler_obj, + url_route=url_route, + request_body=_REQUEST_FIELDS.validate_python(request_body), + raw_bytes=raw_bytes, + ) + ) ( standard_logging_response_object, kwargs, @@ -408,6 +420,7 @@ class PassThroughStreamingHandler: raw_bytes=raw_bytes, end_time=end_time, model=model, + background_response_settled_before=background_response_settled_before, ) # Always reached from an async context (anthropic_messages, # google_genai, and proxy pass-through stream tasks). prefer_async_handlers @@ -426,6 +439,25 @@ class PassThroughStreamingHandler: _route_streaming_logging_to_handler = route_streaming_logging_to_handler + @staticmethod + async def _stream_background_response_settled_before( + passthrough_success_handler_obj: PassThroughEndpointLogging, + url_route: str, + request_body: Mapping[str, object], + raw_bytes: Sequence[bytes], + ) -> bool: + return OpenAIPassthroughLoggingHandler.stream_may_settle_background_response( + url_route, request_body + ) and await passthrough_success_handler_obj.background_response_settled_before( + await asyncify(PassThroughStreamingHandler._background_response_settled_by_stream)(raw_bytes) + ) + + @staticmethod + def _background_response_settled_by_stream(raw_bytes: Sequence[bytes]) -> str | None: + return OpenAIPassthroughLoggingHandler.background_response_settled_by_stream( + PassThroughStreamingHandler._convert_raw_bytes_to_str_lines(raw_bytes) + ) + @staticmethod def _build_passthrough_logging_result( litellm_logging_obj: LiteLLMLoggingObj, @@ -437,6 +469,7 @@ class PassThroughStreamingHandler: raw_bytes: Sequence[bytes], end_time: datetime, model: str | None, + background_response_settled_before: bool = False, ) -> tuple[PassThroughEndpointLoggingResultValues, dict]: """ Synchronous, CPU-bound reconstruction of the standard logging payload @@ -510,6 +543,7 @@ class PassThroughStreamingHandler: start_time=start_time, all_chunks=all_chunks, end_time=end_time, + background_response_settled_before=background_response_settled_before, ) ) standard_logging_response_object = openai_passthrough_logging_handler_result["result"] diff --git a/litellm/proxy/pass_through_endpoints/success_handler.py b/litellm/proxy/pass_through_endpoints/success_handler.py index 8f55906498b..4552ac98624 100644 --- a/litellm/proxy/pass_through_endpoints/success_handler.py +++ b/litellm/proxy/pass_through_endpoints/success_handler.py @@ -1,14 +1,16 @@ import asyncio +import functools import json -from collections.abc import Mapping, Sequence +from collections.abc import Callable, Coroutine, Mapping, Sequence from contextlib import suppress from datetime import datetime from types import MappingProxyType -from typing import Any, Final +from typing import Any, Final, ParamSpec from urllib.parse import urlparse import anyio import httpx +from pydantic import TypeAdapter from litellm._logging import verbose_proxy_logger from litellm.constants import AZURE_SPEECH_CUSTOM_LLM_PROVIDER @@ -20,6 +22,7 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import ( ) from litellm.types.utils import StandardPassThroughResponseObject +from .background_response_settlement import BackgroundResponseSettlementStore, PrismaBackgroundResponseSettlementStore from .llm_provider_handlers.anthropic_passthrough_logging_handler import ( AnthropicPassthroughLoggingHandler, ) @@ -54,9 +57,43 @@ from .llm_provider_handlers.transcribe_passthrough_logging_handler import ( from .llm_provider_handlers.vertex_passthrough_logging_handler import ( VertexPassthroughLoggingHandler, ) +from .managed_id_codec import decode as decode_managed_id from .upstream_usage_headers import has_upstream_reported_usage cohere_passthrough_logging_handler: Final = CoherePassthroughLoggingHandler() +_RESPONSE_BODY: Final = TypeAdapter(dict[str, object]) + + +async def complete_despite_cancellation(work: Coroutine[object, object, None]) -> None: + pending: Final = asyncio.create_task(work) + try: + await asyncio.shield(pending) + except asyncio.CancelledError: + with anyio.CancelScope(shield=True): + while not pending.done(): + with suppress(asyncio.CancelledError): + await asyncio.shield(pending) + pending.result() + raise + + +_HandlerParams: Final = ParamSpec("_HandlerParams") + + +def _finishes_despite_cancellation_when_finalizing( + handler: Callable[_HandlerParams, Coroutine[object, object, None]], +) -> Callable[_HandlerParams, Coroutine[object, object, None]]: + @functools.wraps(handler) + async def handle( + *args: _HandlerParams.args, + **kwargs: _HandlerParams.kwargs, # kwargs-ok: ParamSpec forwards the wrapped handler's exact signature + ) -> None: + if kwargs.get("finalize_before_dispatch") is True: + await complete_despite_cancellation(handler(*args, **kwargs)) + return + await handler(*args, **kwargs) + + return handle def _safe_response_text(httpx_response: httpx.Response) -> str: @@ -84,11 +121,17 @@ class PassThroughEndpointLogging: self, transcribe_handler: TranscribePassthroughLoggingHandler | None = None, log_dispatch: PassThroughLogDispatch | None = None, + background_response_settlement: BackgroundResponseSettlementStore | None = None, ): self.transcribe_passthrough_logging_handler: Final = ( transcribe_handler if transcribe_handler is not None else TranscribePassthroughLoggingHandler() ) self._injected_log_dispatch: Final = log_dispatch + self.background_response_settlement: Final[BackgroundResponseSettlementStore] = ( + background_response_settlement + if background_response_settlement is not None + else PrismaBackgroundResponseSettlementStore() + ) self.TRACKED_VERTEX_METHOD_ROUTES = ( "generateContent", "streamGenerateContent", @@ -174,16 +217,7 @@ class PassThroughEndpointLogging: ) bind_budget_reservation_to_callbacks(logging_obj.litellm_params) - pending: Final = asyncio.create_task(finalize_and_enqueue()) - try: - await asyncio.shield(pending) - except asyncio.CancelledError: - with anyio.CancelScope(shield=True): - while not pending.done(): - with suppress(asyncio.CancelledError): - await asyncio.shield(pending) - pending.result() - raise + await complete_despite_cancellation(finalize_and_enqueue()) return await logging_obj.dispatch_success_handlers( @@ -197,6 +231,13 @@ class PassThroughEndpointLogging: _handle_logging = handle_logging + async def background_response_settled_before(self, settled_response_id: str | None) -> bool: + if settled_response_id is None: + return False + managed_id: Final = decode_managed_id(settled_response_id) + upstream_response_id: Final = settled_response_id if managed_id is None else managed_id.raw_provider_id + return not await self.background_response_settlement.claim(upstream_response_id) + def normalize_llm_passthrough_logging_payload( self, httpx_response: httpx.Response, @@ -209,6 +250,7 @@ class PassThroughEndpointLogging: end_time: datetime, cache_hit: bool, custom_llm_provider: str | None = None, + background_response_settled_before: bool = False, **kwargs, ): return_dict: Final = { @@ -294,6 +336,7 @@ class PassThroughEndpointLogging: end_time=end_time, cache_hit=cache_hit, request_body=request_body, + background_response_settled_before=background_response_settled_before, **kwargs, ) standard_logging_response_object = openai_passthrough_logging_handler_result["result"] @@ -489,6 +532,7 @@ class PassThroughEndpointLogging: ) return None + @_finishes_despite_cancellation_when_finalizing async def pass_through_async_success_handler( self, httpx_response: httpx.Response, @@ -504,7 +548,7 @@ class PassThroughEndpointLogging: custom_llm_provider: str | None = None, finalize_before_dispatch: bool = False, **kwargs, - ): + ) -> None: standard_logging_response_object: PassThroughEndpointLoggingResultValues | None = None logging_obj.model_call_details["passthrough_logging_payload"] = passthrough_logging_payload if self.is_tinyfish_route(url_route, custom_llm_provider): @@ -558,7 +602,22 @@ class PassThroughEndpointLogging: ) return else: + from .llm_provider_handlers.openai_passthrough_logging_handler import ( + OpenAIPassthroughLoggingHandler, + ) + try: + settled_response_id: Final = ( + OpenAIPassthroughLoggingHandler.background_response_settled_by_call( + httpx_response.request.method, url_route, _RESPONSE_BODY.validate_python(response_body) + ) + if isinstance(response_body, dict) + and OpenAIPassthroughLoggingHandler.is_openai_responses_route(url_route) + else None + ) + background_response_settled_before: Final = await self.background_response_settled_before( + settled_response_id + ) normalized_llm_passthrough_logging_payload: Final = self.normalize_llm_passthrough_logging_payload( httpx_response=httpx_response, response_body=response_body, @@ -570,6 +629,7 @@ class PassThroughEndpointLogging: end_time=end_time, cache_hit=cache_hit, custom_llm_provider=custom_llm_provider, + background_response_settled_before=background_response_settled_before, **kwargs, ) standard_logging_response_object = normalized_llm_passthrough_logging_payload[ diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index e63580da721..37783275d9b 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -2055,6 +2055,11 @@ model LiteLLM_BackgroundInteractionSettlement { @@index([settled_at], map: "idx_background_interaction_settlement_settled_at") } +model LiteLLM_BackgroundResponseSettlement { + response_id String @id + settled_at DateTime @default(now()) +} + model LiteLLM_Lens { reviews LiteLLM_LensReview[] id String @id diff --git a/litellm/repositories/table_repositories.py b/litellm/repositories/table_repositories.py index fbf9e5d07f0..fbef73a6114 100644 --- a/litellm/repositories/table_repositories.py +++ b/litellm/repositories/table_repositories.py @@ -408,6 +408,12 @@ class BackgroundInteractionSettlementRepository( table_name = "litellm_backgroundinteractionsettlement" +class BackgroundResponseSettlementRepository( + PrismaTableRepository["prisma_models.LiteLLM_BackgroundResponseSettlement"] +): + table_name = "litellm_backgroundresponsesettlement" + + class RetiredAgentRepository(PrismaTableRepository["prisma_models.LiteLLM_RetiredAgent"]): table_name = "litellm_retiredagent" diff --git a/schema.prisma b/schema.prisma index e63580da721..37783275d9b 100644 --- a/schema.prisma +++ b/schema.prisma @@ -2055,6 +2055,11 @@ model LiteLLM_BackgroundInteractionSettlement { @@index([settled_at], map: "idx_background_interaction_settlement_settled_at") } +model LiteLLM_BackgroundResponseSettlement { + response_id String @id + settled_at DateTime @default(now()) +} + model LiteLLM_Lens { reviews LiteLLM_LensReview[] id String @id diff --git a/tests/integration/spend/test_passthrough_stream_pricing_paths_spend.py b/tests/integration/spend/test_passthrough_stream_pricing_paths_spend.py index 0ead544a2ee..508ceab38f9 100644 --- a/tests/integration/spend/test_passthrough_stream_pricing_paths_spend.py +++ b/tests/integration/spend/test_passthrough_stream_pricing_paths_spend.py @@ -11,7 +11,7 @@ from urllib.parse import urlsplit import httpx import pytest import yaml -from integration._support.client import Gateway, eventually, gateway_from_environment +from integration._support.client import Gateway, eventually, gateway_from_environment, held from integration._support.forward_proxy import tunnelling_forward_proxy from integration._support.process import graceful_stop_seconds, owned_proxy from integration._support.tls import server_context, write_self_signed_cert @@ -50,6 +50,7 @@ _GEMINI_USAGE: Final = { _DEPLOYMENT_LOCATIONS: Final = MappingProxyType({"audit-vertex-global": "global", "audit-vertex-us": "us"}) _GEMINI_ACTIONS: Final = ("generateContent", "streamGenerateContent?alt=sse") _CONFIGURED_GATEWAY_PATH: Final = "/google-egress" +_SPEND_HOLD_SECONDS: Final = 10 @dataclass(frozen=True, slots=True) @@ -157,9 +158,40 @@ def _openai_reply(request: Request) -> Reply: return Reply(chunks=(*encoded, b"data: [DONE]\n\n"), content_type="text/event-stream") +def _finished_background_response(request: Request) -> Reply: + return Reply( + body=json.dumps( + { + "id": urlsplit(request.target).path.rsplit("/", 1)[-1], + "object": "response", + "created_at": 1, + "status": "completed", + "background": True, + "model": _OPENAI_MODEL, + "output": [ + { + "id": f"msg_{uuid.uuid4().hex}", + "type": "message", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "scripted", "annotations": []}], + } + ], + "usage": { + "input_tokens": _PROMPT_TOKENS, + "output_tokens": _COMPLETION_TOKENS, + "total_tokens": _PROMPT_TOKENS + _COMPLETION_TOKENS, + }, + } + ).encode() + ) + + def _upstream(request: Request) -> Reply: if request.target.startswith("/v1/chat/completions"): return _openai_reply(request) + if request.target.startswith("/v1/responses/"): + return _finished_background_response(request) if "/publishers/anthropic/models/" in request.target: assert f"/v1/projects/{_PROJECT}/locations/" in request.target, request.target return _claude_reply(request) @@ -300,6 +332,15 @@ def _spend_row(gateway: Gateway, key: str) -> dict[str, JsonValue]: return rows[0] +def _key_spend(gateway: Gateway, key: str) -> float: + rows: Final = gateway.rows( + 'SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s', (sha256(key.encode()).hexdigest(),) + ) + spend: Final = rows[0]["spend"] + assert isinstance(spend, float), rows + return spend + + def _spend_for_key(gateway: Gateway, key: str) -> float: spend: Final = _spend_row(gateway, key)["spend"] assert isinstance(spend, float), key @@ -448,6 +489,23 @@ def test_openai_passthrough_stream_cut_off_after_its_usage_bills_what_the_finish assert spend == pytest.approx({"finished": list_price, "cut off": list_price}) +def test_an_openai_passthrough_background_response_read_three_times_bills_one_read(rig: _Rig) -> None: + response_path: Final = f"/openai_passthrough/v1/responses/resp_{uuid.uuid4().hex}" + with rig.proxy.scenario() as scenario: + key: Final = scenario.key() + statuses: Final = [ + rig.proxy.client.get(response_path, headers={"Authorization": f"Bearer {key}"}).status_code + for _ in range(3) + ] + assert statuses == [200, 200, 200] + held( + lambda: _key_spend(rig.proxy, key), + lambda spend: spend == pytest.approx(_list_price(_OPENAI_MODEL)), + holding=_SPEND_HOLD_SECONDS, + seconds=70, + ) + + def _caller_headers(route: str, key: str) -> Mapping[str, str]: if route == "vertex_ai": return {"x-litellm-api-key": f"Bearer {key}", "Authorization": f"Bearer {_CALLER_GOOGLE_TOKEN}"} diff --git a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py index bb4b41fd09d..2ed75548702 100644 --- a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py +++ b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py @@ -2,17 +2,20 @@ from dataclasses import dataclass import asyncio import json import logging -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from datetime import datetime +from types import SimpleNamespace from typing import Any, Dict, Final, List from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest +from prisma.errors import UniqueViolationError from pydantic import BaseModel, ConfigDict import litellm +from litellm.proxy import proxy_server from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.token_counter import high_detail_image_token_upper_bound @@ -21,10 +24,15 @@ from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_passthrou count_relayed_prompt_tokens, ) from litellm.proxy._types import SpendLogsPayload +from litellm.proxy.pass_through_endpoints.background_response_settlement import ( + PrismaBackgroundResponseSettlementStore, +) +from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER +from litellm.proxy.pass_through_endpoints.managed_id_codec import new_managed_id +from litellm.proxy.pass_through_endpoints.streaming_handler import PassThroughStreamContext, PassThroughStreamingHandler from litellm.proxy.pass_through_endpoints.success_handler import ( PassThroughEndpointLogging, ) -from litellm.proxy.pass_through_endpoints.streaming_handler import PassThroughStreamingHandler from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload from litellm.types.passthrough_endpoints.pass_through_endpoints import ( EndpointType, @@ -2174,6 +2182,60 @@ class _SpendRowRecorder(CustomLogger): ) +class _SettlementTable: + def __init__(self) -> None: + self.settled_response_ids: Final[set[str]] = set() + + async def create_many(self, data: Sequence[Mapping[str, str]], *, skip_duplicates: bool) -> int: + new_ids: Final = {row["response_id"] for row in data} - self.settled_response_ids + if not skip_duplicates and len(new_ids) < len(data): + raise UniqueViolationError({"user_facing_error": {"error_code": "P2002"}}) + self.settled_response_ids.update(new_ids) + return len(new_ids) + + +class _UnreachableSettlementTable: + async def create_many(self, data: Sequence[Mapping[str, str]], *, skip_duplicates: bool) -> int: + raise ConnectionResetError("server closed the connection unexpectedly") + + +@dataclass(frozen=True, slots=True) +class _StalledSettlementTable: + table: _SettlementTable + + async def create_many(self, data: Sequence[Mapping[str, str]], *, skip_duplicates: bool) -> int: + await asyncio.sleep(3) + return await self.table.create_many(data, skip_duplicates=skip_duplicates) + + +@dataclass(frozen=True, slots=True) +class _SlowAckSettlementTable: + table: _SettlementTable + inserted: asyncio.Event + + async def create_many(self, data: Sequence[Mapping[str, str]], *, skip_duplicates: bool) -> int: + inserted_count: Final = await self.table.create_many(data, skip_duplicates=skip_duplicates) + self.inserted.set() + await asyncio.sleep(0.3) + return inserted_count + + +def _settling_passthrough_logging( + table: _SettlementTable | _UnreachableSettlementTable | _StalledSettlementTable | _SlowAckSettlementTable | None, + timeout_seconds: float = 0.05, +) -> PassThroughEndpointLogging: + return PassThroughEndpointLogging( + background_response_settlement=PrismaBackgroundResponseSettlementStore( + table=lambda: table, timeout_seconds=timeout_seconds + ) + ) + + +def _completed_stream_bytes(response: Mapping[str, object]) -> list[bytes]: + completed: Final = {"type": "response.completed", "sequence_number": 9, "response": response} + return [f"event: response.completed\ndata: {json.dumps(completed)}\n\n".encode()] + + class TestOpenAIPassthroughResponsesSubRouteBilling: def _dispatch( self, @@ -2362,12 +2424,31 @@ class TestOpenAIPassthroughResponsesSubRouteBilling: response_body: dict[str, object], litellm_call_id: str, base_url: str = RESPONSES_URL, + *, + passthrough_logging: PassThroughEndpointLogging | None = None, ) -> _LoggedCall: - url_route: Final = f"{base_url}{sub_route}" recorder: Final = _SpendRowRecorder() - logging_obj: Final = self._recorded_logging_obj(recorder, litellm_call_id, stream=False) + await self._handled_call( + method, + f"{base_url}{sub_route}", + request_body, + response_body, + self._recorded_logging_obj(recorder, litellm_call_id, stream=False), + passthrough_logging or PassThroughEndpointLogging(), + ) + return await asyncio.wait_for(recorder.call, timeout=5) - await PassThroughEndpointLogging().pass_through_async_success_handler( + @staticmethod + async def _handled_call( + method: str, + url_route: str, + request_body: dict[str, object], + response_body: dict[str, object], + logging_obj: LiteLLMLoggingObj, + passthrough_logging: PassThroughEndpointLogging, + finalize_before_dispatch: bool = False, + ) -> None: + await passthrough_logging.pass_through_async_success_handler( httpx_response=httpx.Response(200, json=response_body, request=httpx.Request(method, url_route)), response_body=response_body, logging_obj=logging_obj, @@ -2381,9 +2462,340 @@ class TestOpenAIPassthroughResponsesSubRouteBilling: url=url_route, request_body=request_body, request_method=method ), custom_llm_provider="openai", + finalize_before_dispatch=finalize_before_dispatch, + ) + + async def _streamed_call( + self, + sub_route: str, + request_body: dict[str, object], + raw_bytes: list[bytes], + litellm_call_id: str, + passthrough_logging: PassThroughEndpointLogging, + ) -> _LoggedCall: + recorder: Final = _SpendRowRecorder() + logging_obj: Final = self._recorded_logging_obj(recorder, litellm_call_id, stream=True) + logging_obj.model_call_details["custom_llm_provider"] = "openai" + + await PassThroughStreamingHandler.route_streaming_logging_to_handler( + litellm_logging_obj=logging_obj, + passthrough_success_handler_obj=passthrough_logging, + url_route=f"{RESPONSES_URL}{sub_route}", + request_body=request_body, + endpoint_type=EndpointType.OPENAI, + start_time=DISPATCH_TIME, + raw_bytes=raw_bytes, + end_time=DISPATCH_TIME, ) return await asyncio.wait_for(recorder.call, timeout=5) + def _failed_stream_logging_obj(self, litellm_call_id: str) -> LiteLLMLoggingObj: + logging_obj: Final = self._recorded_logging_obj(_SpendRowRecorder(), litellm_call_id, stream=True) + logging_obj.model_call_details["custom_llm_provider"] = "openai" + return logging_obj + + @staticmethod + async def _failed_stream( + logging_obj: LiteLLMLoggingObj, + sub_route: str, + raw_bytes: list[bytes], + passthrough_logging: PassThroughEndpointLogging, + ) -> None: + await PassThroughStreamingHandler.schedule_stream_failure_logging( + litellm_logging_obj=logging_obj, + endpoint_type=EndpointType.OPENAI, + request_body={}, + raw_bytes=raw_bytes, + exception=httpx.RemoteProtocolError("peer closed connection without sending complete message body"), + stream_context=PassThroughStreamContext(passthrough_logging, f"{RESPONSES_URL}{sub_route}", DISPATCH_TIME), + ) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("method", "sub_route", "response_body"), + [ + pytest.param( + "GET", + f"/{STORED_RESPONSE_ID}", + _stored_response("completed", True, RESPONSES_USAGE), + id="retrieve_then_retrieve", + ), + pytest.param( + "POST", + f"/{STORED_RESPONSE_ID}/cancel", + _stored_response("cancelled", True, RESPONSES_USAGE), + id="retrieve_then_cancel", + ), + ], + ) + async def test_a_background_response_read_twice_bills_only_the_first_read( + self, local_model_cost_map, method: str, sub_route: str, response_body: dict[str, object] + ): + passthrough_logging: Final = _settling_passthrough_logging(_SettlementTable()) + later_call_id: Final = "e4a1c7b2-3d58-4f9e-8b06-91c2d7f3a5e8" + + first: Final = await self._logged_call( + "GET", + f"/{STORED_RESPONSE_ID}", + {}, + _stored_response("completed", True, RESPONSES_USAGE), + "2f9d6b13-8c4e-4a71-b5d0-3e7a1c9f2b64", + passthrough_logging=passthrough_logging, + ) + later: Final = await self._logged_call( + method, sub_route, {}, response_body, later_call_id, passthrough_logging=passthrough_logging + ) + + assert first.spend_row["spend"] > 0 + assert first.spend_row["total_tokens"] > 0 + assert first.spend_row["request_id"].startswith(STORED_RESPONSE_ID) + assert not later.spend_row["spend"] + assert not later.spend_row["total_tokens"] + assert later.spend_row["request_id"] == later_call_id + + @pytest.mark.asyncio + async def test_the_proxy_database_keeps_a_background_response_billed_once(self, local_model_cost_map, monkeypatch): + settlement_table: Final = _SettlementTable() + monkeypatch.setattr( + proxy_server, + "prisma_client", + SimpleNamespace(db=SimpleNamespace(litellm_backgroundresponsesettlement=settlement_table)), + ) + passthrough_logging: Final = PassThroughEndpointLogging() + finished: Final = _stored_response("completed", True, RESPONSES_USAGE) + + first: Final = await self._logged_call( + "GET", + f"/{STORED_RESPONSE_ID}", + {}, + finished, + "5a9c2e71-3f4b-4d86-a1e7-c8b6d3f05e92", + passthrough_logging=passthrough_logging, + ) + second: Final = await self._logged_call( + "GET", + f"/{STORED_RESPONSE_ID}", + {}, + finished, + "b2d7f4a8-6e1c-4a39-8f5b-0c9e3a7d1b64", + passthrough_logging=passthrough_logging, + ) + + assert first.spend_row["spend"] > 0 + assert not second.spend_row["spend"] + assert settlement_table.settled_response_ids == {STORED_RESPONSE_ID} + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "retrieved_id", + [ + pytest.param(STORED_RESPONSE_ID, id="raw_id"), + pytest.param(new_managed_id("openai", STORED_RESPONSE_ID), id="managed_id"), + ], + ) + async def test_a_streamed_background_create_bills_once_and_its_later_retrieve_does_not( + self, local_model_cost_map, retrieved_id: str + ): + passthrough_logging: Final = _settling_passthrough_logging(_SettlementTable()) + finished: Final = _stored_response("completed", True, RESPONSES_USAGE) + + streamed_create: Final = await self._streamed_call( + "", + {**CREATE_REQUEST, "background": True, "stream": True}, + _completed_stream_bytes(finished), + "8d2f4a61-7b3c-4e95-a0d8-c5e1f9b7234a", + passthrough_logging, + ) + retrieve: Final = await self._logged_call( + "GET", + f"/{STORED_RESPONSE_ID}", + {}, + {**finished, "id": retrieved_id}, + "4b7e9c25-1a6d-4f38-9e02-d8c3a5f1b760", + passthrough_logging=passthrough_logging, + ) + + assert streamed_create.spend_row["spend"] > 0 + assert not retrieve.spend_row["spend"] + + @pytest.mark.asyncio + async def test_a_finished_read_cancelled_while_it_settles_still_bills_and_its_later_read_does_not( + self, local_model_cost_map + ): + settled: Final = asyncio.Event() + passthrough_logging: Final = _settling_passthrough_logging( + _SlowAckSettlementTable(_SettlementTable(), settled), timeout_seconds=2 + ) + finished: Final = _stored_response("completed", True, RESPONSES_USAGE) + recorder: Final = _SpendRowRecorder() + cancelled_read: Final = asyncio.create_task( + self._handled_call( + "GET", + f"{RESPONSES_URL}/{STORED_RESPONSE_ID}", + {}, + finished, + self._recorded_logging_obj(recorder, "6c3e8a14-2f7b-4d95-b1e0-a8d4c6f2e735", stream=False), + passthrough_logging, + finalize_before_dispatch=True, + ) + ) + await settled.wait() + cancelled_read.cancel() + with pytest.raises(asyncio.CancelledError): + await cancelled_read + first: Final = await asyncio.wait_for(recorder.call, timeout=5) + later: Final = await self._logged_call( + "GET", + f"/{STORED_RESPONSE_ID}", + {}, + finished, + "d1a7f3c9-5e2b-4806-9c4d-3b8e6f0a2d17", + passthrough_logging=passthrough_logging, + ) + + assert first.spend_row["spend"] > 0 + assert not later.spend_row["spend"] + + @pytest.mark.asyncio + async def test_a_background_create_stream_cut_before_it_finishes_leaves_the_bill_to_the_finished_retrieve( + self, local_model_cost_map + ): + passthrough_logging: Final = _settling_passthrough_logging(_SettlementTable()) + created: Final = { + "type": "response.created", + "sequence_number": 0, + "response": _stored_response("in_progress", True, None), + } + + await self._streamed_call( + "", + {**CREATE_REQUEST, "background": True, "stream": True}, + [f"event: response.created\ndata: {json.dumps(created)}\n\n".encode()], + "d4b8e2f6-9a1c-4573-8e6d-1f7c3a5b9e20", + passthrough_logging, + ) + retrieve: Final = await self._logged_call( + "GET", + f"/{STORED_RESPONSE_ID}", + {}, + _stored_response("completed", True, RESPONSES_USAGE), + "f1a6c3e8-4d2b-4b97-a5e0-7c9d2f4b6a13", + passthrough_logging=passthrough_logging, + ) + + assert retrieve.spend_row["spend"] > 0 + + @pytest.mark.asyncio + async def test_a_streamed_foreground_create_bills_without_recording_a_settlement(self, local_model_cost_map): + settlement_table: Final = _SettlementTable() + + streamed_create: Final = await self._streamed_call( + "", + {**CREATE_REQUEST, "stream": True}, + _completed_stream_bytes(_stored_response("completed", False, RESPONSES_USAGE)), + "3e8b1f57-6c2a-4d94-b7e0-9a5c4f2d816b", + _settling_passthrough_logging(settlement_table), + ) + + assert streamed_create.spend_row["spend"] > 0 + assert not settlement_table.settled_response_ids + + @pytest.mark.asyncio + async def test_a_resumed_stream_after_a_billed_retrieve_does_not_bill_again(self, local_model_cost_map): + passthrough_logging: Final = _settling_passthrough_logging(_SettlementTable()) + finished: Final = _stored_response("completed", True, RESPONSES_USAGE) + + retrieve: Final = await self._logged_call( + "GET", + f"/{STORED_RESPONSE_ID}", + {}, + finished, + "6c3a8e14-5f2b-4d97-b1e6-a7d4c2f9e853", + passthrough_logging=passthrough_logging, + ) + resumed_stream: Final = await self._streamed_call( + f"/{STORED_RESPONSE_ID}?stream=true", + {}, + _completed_stream_bytes(finished), + "a9e5d3c7-2b8f-4146-8d0a-f3b6e1c4d295", + passthrough_logging, + ) + + assert retrieve.spend_row["spend"] > 0 + assert retrieve.spend_row["total_tokens"] > 0 + assert not resumed_stream.spend_row["spend"] + assert not resumed_stream.spend_row["total_tokens"] + assert resumed_stream.spend_row["request_id"] != retrieve.spend_row["request_id"] + + @pytest.mark.asyncio + async def test_a_resumed_stream_that_fails_after_it_finished_leaves_the_charge_to_its_next_retrieve( + self, local_model_cost_map + ): + passthrough_logging: Final = _settling_passthrough_logging(_SettlementTable()) + finished: Final = _stored_response("completed", True, RESPONSES_USAGE) + + await self._failed_stream( + self._failed_stream_logging_obj("5d9a3f72-1e8b-4c46-a7d0-b3f6e2c9a581"), + f"/{STORED_RESPONSE_ID}?stream=true", + _completed_stream_bytes(finished), + passthrough_logging, + ) + await GLOBAL_LOGGING_WORKER.flush() + first_retrieve: Final = await self._logged_call( + "GET", + f"/{STORED_RESPONSE_ID}", + {}, + finished, + "c1e7b4a9-6f2d-4835-9b0e-a8d5f3c2e176", + passthrough_logging=passthrough_logging, + ) + later_retrieve: Final = await self._logged_call( + "GET", + f"/{STORED_RESPONSE_ID}", + {}, + finished, + "7a3e9c15-4b8d-4f62-a0e7-d5c2b9f1e384", + passthrough_logging=passthrough_logging, + ) + + assert first_retrieve.spend_row["spend"] > 0 + assert not later_retrieve.spend_row["spend"] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "table", + [ + pytest.param(None, id="no_database"), + pytest.param(_UnreachableSettlementTable(), id="database_error"), + pytest.param(_StalledSettlementTable(_SettlementTable()), id="database_stalled"), + ], + ) + async def test_finished_reads_bill_as_before_when_their_settlement_cannot_be_recorded( + self, local_model_cost_map, table: _UnreachableSettlementTable | _StalledSettlementTable | None + ): + passthrough_logging: Final = _settling_passthrough_logging(table) + finished: Final = _stored_response("completed", True, RESPONSES_USAGE) + + first: Final = await self._logged_call( + "GET", + f"/{STORED_RESPONSE_ID}", + {}, + finished, + "1d6b4f82-9e3a-4c75-a2d8-b5f7e0c3a916", + passthrough_logging=passthrough_logging, + ) + second: Final = await self._logged_call( + "GET", + f"/{STORED_RESPONSE_ID}", + {}, + finished, + "7f2c9a35-4e1d-4b68-9c07-e6a3d8b5f142", + passthrough_logging=passthrough_logging, + ) + + assert first.spend_row["spend"] > 0 + assert second.spend_row["spend"] == pytest.approx(first.spend_row["spend"]) + @pytest.mark.asyncio async def test_unbilled_retrieve_spend_log_keeps_the_response_model(self, local_model_cost_map): response_body: Final = _stored_response("completed", False, RESPONSES_USAGE)