scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-12 08:42:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 455

fix(proxy): bill an OpenAI pass-through background response once across repeated reads (#50841)

Browse source 
* fix(proxy): bill an OpenAI pass-through background response once across repeated reads

A finished background response now bills only on the first pass-through call
that sees it finished with usage. The claim is an insert on the new
LiteLLM_BackgroundResponseSettlement table, keyed by response id, so it holds
across workers and replicas. Every later retrieve, cancel, streamed create, or
resumed stream logs like an unbilled read. With no database or a failed insert
the call still bills.

* fix(proxy): keep pass-through background billing fail-open without prisma and skip foreground stream re-parse

* fix(proxy): claim background responses on failed streams and bound the claim

A resumed stream that fails after response.completed billed its usage through
the partial-usage path without claiming the response id, so a later read billed
it again. That path now claims the id and logs $0 when an earlier read paid.

The claim uses create_many(skip_duplicates=True) so a repeat read no longer
raises a unique violation in Postgres, and it waits at most 2 s before billing,
since the claim now runs before the client gets its response.

* fix(pass_through): keep the background response claim and its spend write together under cancellation and managed ids

* test(pass_through): cover a failed stream whose failure logging cannot be scheduled

* fix(proxy): start the background response settlement migration with a lock timeout

* fix(passthrough): leave a failed background stream's charge to its next read

A stream that failed after response.completed claimed the response id while its failure row billed $0, so no later read paid. The failure path is back to main's behavior, so the next finished read claims and bills once

---------

Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
tavren27 2026-11-10 07:33:25 -08:00 • committed by GitHub
parent 7bcc46246b
commit 1c68ec429a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
14 changed files with 714 additions and 28 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

9
litellm-proxy-extras/litellm_proxy_extras/migrations/20261029000000_add_background_response_settlement/migration.sql Normal file
View file

@ -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")
);

5
litellm-proxy-extras/litellm_proxy_extras/schema.prisma
View file

@ -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

1
litellm/constants.py
View file

@ -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))

1
litellm/integrations/otel/model/spans.py
View file

@ -292,6 +292,7 @@ _PRISMA_MODELS: Final[frozenset[str]] = frozenset(
"LiteLLM_SpendLogs",
"LiteLLM_BudgetWindowSpend",
"LiteLLM_BackgroundInteractionSettlement",
"LiteLLM_BackgroundResponseSettlement",
"LiteLLM_ErrorLogs",
"LiteLLM_UserNotifications",
"LiteLLM_UserProviderCredentials",

2
litellm/llms/openai/responses/transformation.py
View file

@ -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

52
litellm/proxy/pass_through_endpoints/background_response_settlement.py Normal file
View file

@ -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

54
litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py
View file

@ -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
)

34
litellm/proxy/pass_through_endpoints/streaming_handler.py
View file

@ -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"]

86
litellm/proxy/pass_through_endpoints/success_handler.py
View file

@ -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[

5
litellm/proxy/schema.prisma
View file

@ -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

6
litellm/repositories/table_repositories.py
View file

@ -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"

5
schema.prisma
View file

@ -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

60
tests/integration/spend/test_passthrough_stream_pricing_paths_spend.py
View file

@ -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}"}

422
tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py
View file

@ -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)