mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-29 06:33:28 +00:00
4e2dd65549
* fix(interactions): keep background interaction bills through crashes, late restarts, and deployment keys A settlement claim is now a lease, so a row whose settler died between claim and outcome is claimable again after the lease and boot resumes it. A resumed poll gets the full poll window and fetches with the serving deployment's credentials from the live router. The stored create context keeps the request tags and the request metadata that the creating worker's own spend row used * fix(interactions): bill a lapsed claim's takeover only when no spend row exists A settler that billed and then died, or failed to record the outcome, left the row unsettled, so the replica that took over its lapsed claim billed the interaction a second time. A takeover now looks up the spend row keyed by the interaction id and records the row as billed when it finds one. * fix(interactions): settle Vertex AI and edited deployments, never double bill with spend logs off A takeover with spend logs off records the interaction unsettled instead of billing it again, and a release records its outcome before freeing the budget reservation. A resumed or live fetch carries the deployment's Vertex AI project, location, and credentials, and a stale deployment id falls back to the model group only when one same-provider deployment serves it. A poll stops at the first 400, 401, 403, or 404, string request tags no longer fail the create, and settled_at gets an index built concurrently. * fix(interactions): leave a live replica's interaction to it when a resumed poll cannot read it A replica that boots with credentials unable to read an interaction got a 404 on its resumed poll and settled the row as unsettled at once, so the replica that created it, still polling, found it settled and never billed it. A resumed poll that gets a 400, 401, 403, or 404 now stops fetching and keeps checking the row until its poll window ends, and settles it as unsettled only when no other replica did. * fix(interactions): keep the creator's metadata on a resumed bill and retry outcome writes inside the claim lease A bill another replica writes now carries the metadata the creating worker's own spend row is built from, so a client's metadata.model_group no longer relabels it. A failed outcome write is retried for half the claim lease before the budget reservation is left to a takeover, and the creating replica retries every fetch error for its whole window again, as main does. The resume integration case now checks that a resumed fetch uses the deployment's own key. * fix(interactions): resume background interaction polls only after database models load Boot resume ran before the proxy loaded the models stored in the database, so a resumed poll for a model added through /model/new or the Admin UI fetched with env credentials. A 4xx there stops a resumed poll, so the bill was lost. The settlement now installs after the database models load. The integration test resumes a database model with an eager first fetch and fails with the old startup order. * fix(interactions): renew a settling claim so a slow settle is never taken over and billed twice A settle runs every success callback in turn, and one slow callback could hold it past the 5 minute claim lease. A booting replica or a DELETE then took the live claim over, found no spend row yet, and billed the interaction a second time. The settler now renews its claim every third of the lease until it records the outcome, so only a settler that stopped renewing is taken over. A takeover also read the spend log after writing its claim, so a failed read left the row under the taker's own fresh claim, which its poll then took for another settler's and left unbilled until a reboot. The spend row is now read before the claim is written, so a failed read takes nothing and the claim is retried on the next poll. * test(spend): cover lapsed claim takeovers, rotated deployment keys, late restarts, and resumed bill tags --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
1424 lines
53 KiB
Python
1424 lines
53 KiB
Python
import asyncio
|
|
import json
|
|
import time
|
|
from datetime import datetime, timedelta, timezone
|
|
from itertools import islice
|
|
from typing import Optional
|
|
|
|
import httpx
|
|
import pytest
|
|
|
|
import litellm
|
|
from litellm.interactions.background_cost_polling import (
|
|
_create_context,
|
|
_poll_intervals,
|
|
_rebuild_logging_obj,
|
|
BackgroundInteractionPollContext,
|
|
InMemoryBackgroundSettlementStore,
|
|
maybe_schedule_background_interaction_cost_polling,
|
|
maybe_settle_background_interaction_before_delete,
|
|
PendingBackgroundInteraction,
|
|
poll_and_log_background_interaction_cost,
|
|
PollSchedule,
|
|
resume_unsettled_background_interactions,
|
|
)
|
|
from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kwargs
|
|
from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging
|
|
from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload
|
|
from litellm.types.interactions import InteractionsAPIResponse
|
|
|
|
USAGE_BLOCK = {
|
|
"total_tokens": 175,
|
|
"total_input_tokens": 100,
|
|
"input_tokens_by_modality": [{"modality": "text", "tokens": 100}],
|
|
"total_cached_tokens": 0,
|
|
"total_output_tokens": 50,
|
|
"output_tokens_by_modality": [{"modality": "text", "tokens": 50}],
|
|
"total_tool_use_tokens": 0,
|
|
"total_thought_tokens": 25,
|
|
}
|
|
|
|
|
|
def _logging_obj(
|
|
call_type: str = "acreate_interaction",
|
|
litellm_params: Optional[dict] = None,
|
|
) -> LitellmLogging:
|
|
logging_obj = LitellmLogging(
|
|
model="gemini-2.5-flash",
|
|
messages=[],
|
|
stream=False,
|
|
call_type=call_type,
|
|
start_time=time.time(),
|
|
litellm_call_id="bg-interactions-call-id",
|
|
function_id="bg-interactions-fn-id",
|
|
)
|
|
logging_obj.update_environment_variables(
|
|
litellm_params=litellm_params or {},
|
|
optional_params={},
|
|
model="gemini-2.5-flash",
|
|
custom_llm_provider="gemini",
|
|
input="hi",
|
|
)
|
|
return logging_obj
|
|
|
|
|
|
def _reservation() -> dict:
|
|
return {"reserved_cost": 0.05, "entries": [], "finalized": False, "input_cost": 0.001}
|
|
|
|
|
|
def _logging_obj_with_reservation(reservation: dict) -> LitellmLogging:
|
|
return _logging_obj(litellm_params={"metadata": {"user_api_key_budget_reservation": reservation}})
|
|
|
|
|
|
async def _raise_on_billing(result: InteractionsAPIResponse) -> None:
|
|
raise RuntimeError("cost calculation failed for a settled background interaction")
|
|
|
|
|
|
def _context(
|
|
logging_obj: LitellmLogging,
|
|
timeout_seconds: float = 1.0,
|
|
store: Optional[InMemoryBackgroundSettlementStore] = None,
|
|
resumed: bool = False,
|
|
) -> BackgroundInteractionPollContext:
|
|
return BackgroundInteractionPollContext(
|
|
interaction_id="interactions/bg-abc",
|
|
custom_llm_provider="gemini",
|
|
logging_obj=logging_obj,
|
|
initial_interval_seconds=0.001,
|
|
max_interval_seconds=0.002,
|
|
timeout_seconds=timeout_seconds,
|
|
outcome_write_timeout_seconds=0.01,
|
|
store=store if store is not None else InMemoryBackgroundSettlementStore(),
|
|
resumed=resumed,
|
|
)
|
|
|
|
|
|
def _response(status: str, with_usage: bool) -> InteractionsAPIResponse:
|
|
return InteractionsAPIResponse(
|
|
id="interactions/bg-abc",
|
|
model="gemini-2.5-flash",
|
|
status=status,
|
|
steps=[],
|
|
usage=dict(USAGE_BLOCK) if with_usage else None,
|
|
)
|
|
|
|
|
|
def _fetch_sequence(*responses):
|
|
remaining = list(responses)
|
|
calls = []
|
|
|
|
async def fetch(context):
|
|
calls.append(context.interaction_id)
|
|
item = remaining.pop(0) if len(remaining) > 1 else remaining[0]
|
|
if isinstance(item, Exception):
|
|
raise item
|
|
return item
|
|
|
|
return fetch, calls
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"initial, maximum",
|
|
[(0.0, 0.002), (0.001, 0.0), (-1.0, 0.002), (0.0, 0.0)],
|
|
)
|
|
def test_poll_intervals_stops_instead_of_looping_on_a_non_positive_interval(initial, maximum):
|
|
intervals = list(islice(_poll_intervals(initial=initial, maximum=maximum, timeout=3600.0), 10))
|
|
|
|
assert len(intervals) < 10
|
|
assert all(interval > 0 for interval in intervals)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_poller_bills_once_when_interaction_completes():
|
|
logging_obj = _logging_obj()
|
|
fetch, calls = _fetch_sequence(
|
|
_response("in_progress", with_usage=False),
|
|
_response("completed", with_usage=True),
|
|
)
|
|
|
|
await poll_and_log_background_interaction_cost(_context(logging_obj), fetch_interaction=fetch)
|
|
|
|
assert len(calls) == 2
|
|
assert logging_obj.model_call_details["response_cost"] > 0
|
|
assert logging_obj.model_call_details["standard_logging_object"]["total_tokens"] == 175
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_poller_bills_an_interaction_paused_for_a_tool_result():
|
|
logging_obj = _logging_obj()
|
|
fetch, calls = _fetch_sequence(
|
|
_response("in_progress", with_usage=False),
|
|
_response("requires_action", with_usage=True),
|
|
)
|
|
|
|
await poll_and_log_background_interaction_cost(_context(logging_obj), fetch_interaction=fetch)
|
|
|
|
assert len(calls) == 2
|
|
assert logging_obj.model_call_details["response_cost"] > 0
|
|
assert logging_obj.model_call_details["standard_logging_object"]["total_tokens"] == 175
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_poller_does_not_pin_the_budget_for_an_interaction_paused_for_a_tool_result():
|
|
reservation = _reservation()
|
|
logging_obj = _logging_obj_with_reservation(reservation)
|
|
fetch, _ = _fetch_sequence(_response("requires_action", with_usage=True))
|
|
|
|
await poll_and_log_background_interaction_cost(_context(logging_obj), fetch_interaction=fetch)
|
|
|
|
assert logging_obj.model_call_details["response_cost"] > 0
|
|
assert reservation["finalized"] is False
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_poller_stops_without_billing_on_terminal_status_without_usage():
|
|
logging_obj = _logging_obj()
|
|
fetch, calls = _fetch_sequence(_response("failed", with_usage=False))
|
|
|
|
await poll_and_log_background_interaction_cost(_context(logging_obj), fetch_interaction=fetch)
|
|
|
|
assert len(calls) == 1
|
|
assert logging_obj.model_call_details.get("response_cost") is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_poller_gives_up_after_timeout_without_billing():
|
|
logging_obj = _logging_obj()
|
|
fetch, calls = _fetch_sequence(_response("in_progress", with_usage=False))
|
|
|
|
await poll_and_log_background_interaction_cost(
|
|
_context(logging_obj, timeout_seconds=0.01),
|
|
fetch_interaction=fetch,
|
|
)
|
|
|
|
assert len(calls) >= 2
|
|
assert logging_obj.model_call_details.get("response_cost") is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_poller_releases_budget_reservation_when_interaction_ends_without_usage():
|
|
reservation = _reservation()
|
|
logging_obj = _logging_obj_with_reservation(reservation)
|
|
fetch, _ = _fetch_sequence(_response("failed", with_usage=False))
|
|
|
|
await poll_and_log_background_interaction_cost(_context(logging_obj), fetch_interaction=fetch)
|
|
|
|
assert reservation["finalized"] is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_poller_releases_budget_reservation_on_timeout_give_up():
|
|
reservation = _reservation()
|
|
logging_obj = _logging_obj_with_reservation(reservation)
|
|
fetch, _ = _fetch_sequence(_response("in_progress", with_usage=False))
|
|
|
|
await poll_and_log_background_interaction_cost(
|
|
_context(logging_obj, timeout_seconds=0.01),
|
|
fetch_interaction=fetch,
|
|
)
|
|
|
|
assert reservation["finalized"] is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_poller_releases_budget_reservation_when_billing_raises():
|
|
reservation = _reservation()
|
|
logging_obj = _logging_obj_with_reservation(reservation)
|
|
fetch, _ = _fetch_sequence(_response("completed", with_usage=True))
|
|
logging_obj.async_log_background_interaction_completion = _raise_on_billing
|
|
|
|
with pytest.raises(RuntimeError):
|
|
await poll_and_log_background_interaction_cost(_context(logging_obj), fetch_interaction=fetch)
|
|
|
|
assert reservation["finalized"] is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_poller_leaves_reservation_reconciliation_to_the_completion_event():
|
|
reservation = _reservation()
|
|
logging_obj = _logging_obj_with_reservation(reservation)
|
|
fetch, _ = _fetch_sequence(
|
|
_response("in_progress", with_usage=False),
|
|
_response("completed", with_usage=True),
|
|
)
|
|
|
|
await poll_and_log_background_interaction_cost(_context(logging_obj), fetch_interaction=fetch)
|
|
|
|
assert logging_obj.model_call_details["response_cost"] > 0
|
|
assert reservation["finalized"] is False
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_poller_retries_after_fetch_error_and_still_bills():
|
|
logging_obj = _logging_obj()
|
|
fetch, calls = _fetch_sequence(
|
|
RuntimeError("transient network error"),
|
|
_response("completed", with_usage=True),
|
|
)
|
|
|
|
await poll_and_log_background_interaction_cost(_context(logging_obj), fetch_interaction=fetch)
|
|
|
|
assert len(calls) == 2
|
|
assert logging_obj.model_call_details["response_cost"] > 0
|
|
|
|
|
|
def _provider_error_response(status_code: int) -> httpx.Response:
|
|
return httpx.Response(status_code, request=httpx.Request("GET", "https://generativelanguage.googleapis.com"))
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"error",
|
|
[
|
|
litellm.BadRequestError(message="invalid interaction id", model="gemini-2.5-flash", llm_provider="gemini"),
|
|
litellm.AuthenticationError(message="API key not valid", model="gemini-2.5-flash", llm_provider="gemini"),
|
|
litellm.PermissionDeniedError(
|
|
message="permission denied",
|
|
model="gemini-2.5-flash",
|
|
llm_provider="gemini",
|
|
response=_provider_error_response(403),
|
|
),
|
|
litellm.NotFoundError(message="interaction not found", model="gemini-2.5-flash", llm_provider="gemini"),
|
|
],
|
|
)
|
|
@pytest.mark.asyncio
|
|
async def test_a_resumed_poll_stops_fetching_at_the_first_error_no_retry_changes_and_releases_the_reservation(error):
|
|
"""
|
|
The regression: a resumed fetch the provider refuses for good, like an
|
|
unknown id or a key rotated away, was retried every interval until the
|
|
poll timed out.
|
|
"""
|
|
reservation = _reservation()
|
|
logging_obj = _logging_obj_with_reservation(reservation)
|
|
fetch, calls = _fetch_sequence(error)
|
|
|
|
outcome = await poll_and_log_background_interaction_cost(
|
|
_context(logging_obj, timeout_seconds=0.05, resumed=True), fetch_interaction=fetch
|
|
)
|
|
|
|
assert outcome == "unsettled"
|
|
assert len(calls) == 1
|
|
assert reservation["finalized"] is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_a_resumed_poll_the_provider_refuses_leaves_the_bill_to_the_replica_still_polling_it():
|
|
"""
|
|
The regression: a replica that booted with credentials unable to read the
|
|
interaction settled it as unsettled on its first 404, and the replica that
|
|
created it, still polling, then found it settled and never billed it.
|
|
"""
|
|
store = InMemoryBackgroundSettlementStore()
|
|
creator_logging_obj = _logging_obj()
|
|
resumed_fetched = asyncio.Event()
|
|
|
|
async def creator_fetch(context):
|
|
if resumed_fetched.is_set():
|
|
return _response("completed", with_usage=True)
|
|
return _response("in_progress", with_usage=False)
|
|
|
|
async def resumed_fetch(context):
|
|
resumed_fetched.set()
|
|
raise litellm.NotFoundError(message="interaction not found", model="gemini-2.5-flash", llm_provider="gemini")
|
|
|
|
creator_outcome, resumed_outcome = await asyncio.gather(
|
|
poll_and_log_background_interaction_cost(
|
|
_context(creator_logging_obj, store=store), fetch_interaction=creator_fetch
|
|
),
|
|
poll_and_log_background_interaction_cost(
|
|
_context(_logging_obj(), store=store, resumed=True), fetch_interaction=resumed_fetch
|
|
),
|
|
)
|
|
|
|
assert (creator_outcome, resumed_outcome) == ("billed", None)
|
|
assert creator_logging_obj.model_call_details["response_cost"] > 0
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"error",
|
|
[
|
|
litellm.RateLimitError(message="quota exceeded", model="gemini-2.5-flash", llm_provider="gemini"),
|
|
litellm.ServiceUnavailableError(message="overloaded", model="gemini-2.5-flash", llm_provider="gemini"),
|
|
litellm.InternalServerError(message="internal error", model="gemini-2.5-flash", llm_provider="gemini"),
|
|
litellm.Timeout(message="timed out", model="gemini-2.5-flash", llm_provider="gemini"),
|
|
litellm.BadRequestError(message="invalid interaction id", model="gemini-2.5-flash", llm_provider="gemini"),
|
|
litellm.AuthenticationError(message="API key not valid", model="gemini-2.5-flash", llm_provider="gemini"),
|
|
litellm.NotFoundError(message="interaction not found", model="gemini-2.5-flash", llm_provider="gemini"),
|
|
],
|
|
)
|
|
@pytest.mark.asyncio
|
|
async def test_the_creating_poll_retries_through_any_provider_error(error):
|
|
logging_obj = _logging_obj()
|
|
fetch, calls = _fetch_sequence(error, _response("completed", with_usage=True))
|
|
|
|
outcome = await poll_and_log_background_interaction_cost(_context(logging_obj), fetch_interaction=fetch)
|
|
|
|
assert outcome == "billed"
|
|
assert len(calls) == 2
|
|
|
|
|
|
class _OutcomeWritesFail:
|
|
def __init__(self):
|
|
self.store = InMemoryBackgroundSettlementStore()
|
|
|
|
async def register(self, pending):
|
|
await self.store.register(pending)
|
|
|
|
async def pending(self, interaction_id):
|
|
return await self.store.pending(interaction_id)
|
|
|
|
async def is_claimed(self, interaction_id):
|
|
return await self.store.is_claimed(interaction_id)
|
|
|
|
async def claim(self, interaction_id):
|
|
return await self.store.claim(interaction_id)
|
|
|
|
async def record_outcome(self, interaction_id, outcome):
|
|
raise RuntimeError("database unavailable")
|
|
|
|
async def unsettled(self):
|
|
return await self.store.unsettled()
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"last_status, timeout_seconds",
|
|
[("failed", 1.0), ("in_progress", 0.01)],
|
|
)
|
|
@pytest.mark.asyncio
|
|
async def test_a_release_it_could_not_record_leaves_the_reservation_to_the_settler_that_takes_over(
|
|
last_status, timeout_seconds
|
|
):
|
|
"""
|
|
The regression: a settler released the reservation and then failed to
|
|
record it, so the replica that took over the lapsed claim released the
|
|
same reservation again and the key's budget lost that amount twice.
|
|
"""
|
|
reservation = _reservation()
|
|
logging_obj = _logging_obj_with_reservation(reservation)
|
|
fetch, _ = _fetch_sequence(_response(last_status, with_usage=False))
|
|
|
|
await poll_and_log_background_interaction_cost(
|
|
_context(logging_obj, timeout_seconds=timeout_seconds, store=_OutcomeWritesFail()),
|
|
fetch_interaction=fetch,
|
|
)
|
|
|
|
assert reservation["finalized"] is False
|
|
|
|
|
|
class _FirstOutcomeWriteFails(_OutcomeWritesFail):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.failed_once = False
|
|
|
|
async def record_outcome(self, interaction_id, outcome):
|
|
if not self.failed_once:
|
|
self.failed_once = True
|
|
raise RuntimeError("connection reset")
|
|
await self.store.record_outcome(interaction_id, outcome)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_a_release_whose_first_outcome_write_failed_still_frees_the_reservation():
|
|
"""
|
|
The regression: one failed outcome write left the key's budget
|
|
reservation held until a later boot or delete took the row over, and an
|
|
active key's spend counter never expired it in the meantime.
|
|
"""
|
|
reservation = _reservation()
|
|
store = _FirstOutcomeWriteFails()
|
|
fetch, _ = _fetch_sequence(_response("failed", with_usage=False))
|
|
|
|
outcome = await poll_and_log_background_interaction_cost(
|
|
_context(_logging_obj_with_reservation(reservation), store=store), fetch_interaction=fetch
|
|
)
|
|
|
|
assert outcome == "released"
|
|
assert reservation["finalized"] is True
|
|
assert await store.store.unsettled() == ()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_schedule_creates_poll_task_for_in_progress_create():
|
|
logging_obj = _logging_obj()
|
|
task = await maybe_schedule_background_interaction_cost_polling(
|
|
response=_response("in_progress", with_usage=False),
|
|
create_kwargs={"litellm_logging_obj": logging_obj},
|
|
custom_llm_provider="gemini",
|
|
store=InMemoryBackgroundSettlementStore(),
|
|
)
|
|
|
|
assert isinstance(task, asyncio.Task)
|
|
task.cancel()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await task
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_schedule_registers_an_agent_only_create_that_names_no_model():
|
|
logging_obj = LitellmLogging(
|
|
model=None,
|
|
messages=None,
|
|
stream=False,
|
|
call_type="acreate_interaction",
|
|
start_time=time.time(),
|
|
litellm_call_id="bg-agent-call-id",
|
|
function_id="bg-agent-fn-id",
|
|
)
|
|
logging_obj.update_environment_variables(litellm_params={}, optional_params={}, custom_llm_provider="gemini")
|
|
store = InMemoryBackgroundSettlementStore()
|
|
|
|
task = await maybe_schedule_background_interaction_cost_polling(
|
|
response=_response("in_progress", with_usage=False),
|
|
create_kwargs={"litellm_logging_obj": logging_obj},
|
|
custom_llm_provider="gemini",
|
|
store=store,
|
|
)
|
|
|
|
assert isinstance(task, asyncio.Task)
|
|
task.cancel()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await task
|
|
pending = await store.pending("interactions/bg-abc")
|
|
assert pending is not None
|
|
assert pending.create_context.model is None
|
|
assert _rebuild_logging_obj(pending.create_context).model is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_schedule_registers_a_create_whose_tags_are_a_string_and_keeps_them_as_sent():
|
|
"""
|
|
The regression: storing the create's context read string tags as a list
|
|
and raised, so a create sent with a tags string answered with an error
|
|
after the provider had already started the interaction.
|
|
"""
|
|
logging_obj = _logging_obj(litellm_params={"metadata": _create_metadata(tags="team-a,nightly-batch")})
|
|
store = InMemoryBackgroundSettlementStore()
|
|
|
|
task = await maybe_schedule_background_interaction_cost_polling(
|
|
response=_response("in_progress", with_usage=False),
|
|
create_kwargs={"litellm_logging_obj": logging_obj},
|
|
custom_llm_provider="gemini",
|
|
store=store,
|
|
)
|
|
|
|
assert isinstance(task, asyncio.Task)
|
|
task.cancel()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await task
|
|
pending = await store.pending("interactions/bg-abc")
|
|
assert pending is not None
|
|
rebuilt = _rebuild_logging_obj(pending.create_context)
|
|
assert get_litellm_metadata_from_kwargs(kwargs=rebuilt.model_call_details)["tags"] == "team-a,nightly-batch"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_delete_fetches_with_the_vertex_project_location_and_credentials_the_request_used():
|
|
store = InMemoryBackgroundSettlementStore()
|
|
await _create_on_a_replica_that_then_dies(_logging_obj(litellm_params={"metadata": _create_metadata()}), store)
|
|
fetch, captured = _capturing_fetch(_response("completed", with_usage=True))
|
|
|
|
outcome = await maybe_settle_background_interaction_before_delete(
|
|
interaction_id="interactions/bg-abc",
|
|
delete_kwargs={
|
|
"vertex_project": "request-project",
|
|
"vertex_location": "europe-west4",
|
|
"vertex_credentials": {"type": "service_account", "project_id": "request-project"},
|
|
},
|
|
fetch_interaction=fetch,
|
|
store=store,
|
|
)
|
|
|
|
assert outcome == "billed"
|
|
credentials = captured[0].credentials
|
|
assert (credentials.vertex_project, credentials.vertex_location, credentials.vertex_credentials) == (
|
|
"request-project",
|
|
"europe-west4",
|
|
{"type": "service_account", "project_id": "request-project"},
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
"response,create_kwargs",
|
|
[
|
|
(_response("completed", with_usage=True), {"litellm_logging_obj": "placeholder"}),
|
|
(_response("in_progress", with_usage=False), {}),
|
|
("not a response", {"litellm_logging_obj": "placeholder"}),
|
|
],
|
|
)
|
|
async def test_schedule_skips_non_pollable_results(response, create_kwargs):
|
|
if create_kwargs.get("litellm_logging_obj") == "placeholder":
|
|
create_kwargs = {"litellm_logging_obj": _logging_obj()}
|
|
|
|
task = await maybe_schedule_background_interaction_cost_polling(
|
|
response=response,
|
|
create_kwargs=create_kwargs,
|
|
custom_llm_provider="gemini",
|
|
store=InMemoryBackgroundSettlementStore(),
|
|
)
|
|
|
|
assert task is None
|
|
|
|
|
|
def _register_poll(logging_obj: LitellmLogging, poll_fetch=None, store=None) -> asyncio.Task:
|
|
import litellm.interactions.background_cost_polling as bg
|
|
|
|
if poll_fetch is None:
|
|
poll_fetch, _ = _fetch_sequence(_response("in_progress", with_usage=False))
|
|
context = _context(logging_obj, store=store)
|
|
task = asyncio.create_task(poll_and_log_background_interaction_cost(context, fetch_interaction=poll_fetch))
|
|
bg._ACTIVE_POLLS[context.interaction_id] = bg._ActiveBackgroundPoll(task=task, context=context)
|
|
task.add_done_callback(lambda finished: bg._discard_poll(context.interaction_id, finished))
|
|
return task
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_delete_settlement_bills_an_interaction_paused_for_a_tool_result():
|
|
logging_obj = _logging_obj()
|
|
task = _register_poll(logging_obj)
|
|
fetch, calls = _fetch_sequence(_response("requires_action", with_usage=True))
|
|
|
|
await maybe_settle_background_interaction_before_delete(
|
|
interaction_id="interactions/bg-abc",
|
|
delete_kwargs={},
|
|
fetch_interaction=fetch,
|
|
)
|
|
|
|
assert len(calls) == 1
|
|
assert logging_obj.model_call_details["response_cost"] > 0
|
|
await asyncio.wait_for(task, timeout=5)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_delete_settlement_bills_pending_background_interaction():
|
|
logging_obj = _logging_obj()
|
|
task = _register_poll(logging_obj)
|
|
fetch, calls = _fetch_sequence(_response("completed", with_usage=True))
|
|
|
|
await maybe_settle_background_interaction_before_delete(
|
|
interaction_id="interactions/bg-abc",
|
|
delete_kwargs={},
|
|
fetch_interaction=fetch,
|
|
)
|
|
|
|
assert len(calls) == 1
|
|
assert logging_obj.model_call_details["response_cost"] > 0
|
|
assert logging_obj.model_call_details["standard_logging_object"]["total_tokens"] == 175
|
|
await asyncio.wait_for(task, timeout=5)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_delete_settlement_releases_reservation_when_still_in_progress():
|
|
reservation = _reservation()
|
|
logging_obj = _logging_obj_with_reservation(reservation)
|
|
task = _register_poll(logging_obj)
|
|
fetch, _ = _fetch_sequence(_response("in_progress", with_usage=False))
|
|
|
|
await maybe_settle_background_interaction_before_delete(
|
|
interaction_id="interactions/bg-abc",
|
|
delete_kwargs={},
|
|
fetch_interaction=fetch,
|
|
)
|
|
|
|
assert reservation["finalized"] is True
|
|
assert logging_obj.model_call_details.get("response_cost") is None
|
|
await asyncio.wait_for(task, timeout=5)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_delete_settlement_releases_reservation_when_prefetch_fails():
|
|
reservation = _reservation()
|
|
logging_obj = _logging_obj_with_reservation(reservation)
|
|
task = _register_poll(logging_obj)
|
|
fetch, _ = _fetch_sequence(RuntimeError("interaction already deleted"))
|
|
|
|
await maybe_settle_background_interaction_before_delete(
|
|
interaction_id="interactions/bg-abc",
|
|
delete_kwargs={},
|
|
fetch_interaction=fetch,
|
|
)
|
|
|
|
assert reservation["finalized"] is True
|
|
assert logging_obj.model_call_details.get("response_cost") is None
|
|
await asyncio.wait_for(task, timeout=5)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_delete_settlement_releases_reservation_when_billing_raises():
|
|
reservation = _reservation()
|
|
logging_obj = _logging_obj_with_reservation(reservation)
|
|
task = _register_poll(logging_obj)
|
|
fetch, _ = _fetch_sequence(_response("completed", with_usage=True))
|
|
logging_obj.async_log_background_interaction_completion = _raise_on_billing
|
|
|
|
with pytest.raises(RuntimeError):
|
|
await maybe_settle_background_interaction_before_delete(
|
|
interaction_id="interactions/bg-abc",
|
|
delete_kwargs={},
|
|
fetch_interaction=fetch,
|
|
)
|
|
|
|
assert reservation["finalized"] is True
|
|
await asyncio.wait_for(task, timeout=5)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_delete_settlement_ignores_interactions_without_pending_poll():
|
|
fetch, calls = _fetch_sequence(_response("completed", with_usage=True))
|
|
|
|
await maybe_settle_background_interaction_before_delete(
|
|
interaction_id="interactions/never-polled",
|
|
delete_kwargs={},
|
|
fetch_interaction=fetch,
|
|
)
|
|
|
|
assert calls == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_delete_settlement_noop_after_poll_task_finished():
|
|
logging_obj = _logging_obj()
|
|
poll_fetch, _ = _fetch_sequence(_response("completed", with_usage=True))
|
|
task = _register_poll(logging_obj, poll_fetch=poll_fetch)
|
|
await asyncio.wait_for(task, timeout=5)
|
|
assert logging_obj.model_call_details["response_cost"] > 0
|
|
|
|
settle_fetch, settle_calls = _fetch_sequence(_response("completed", with_usage=True))
|
|
await maybe_settle_background_interaction_before_delete(
|
|
interaction_id="interactions/bg-abc",
|
|
delete_kwargs={},
|
|
fetch_interaction=settle_fetch,
|
|
)
|
|
|
|
assert settle_calls == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_delete_settlement_does_not_rebill_when_gate_already_claimed():
|
|
logging_obj = _logging_obj()
|
|
store = InMemoryBackgroundSettlementStore()
|
|
assert await store.claim("interactions/bg-abc")
|
|
task = _register_poll(logging_obj, store=store)
|
|
fetch, calls = _fetch_sequence(_response("completed", with_usage=True))
|
|
|
|
await maybe_settle_background_interaction_before_delete(
|
|
interaction_id="interactions/bg-abc",
|
|
delete_kwargs={},
|
|
fetch_interaction=fetch,
|
|
)
|
|
|
|
assert len(calls) == 1
|
|
assert logging_obj.model_call_details.get("response_cost") is None
|
|
await asyncio.wait_for(task, timeout=5)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_poller_exits_without_billing_once_settled_elsewhere():
|
|
logging_obj = _logging_obj()
|
|
store = InMemoryBackgroundSettlementStore()
|
|
assert await store.claim("interactions/bg-abc")
|
|
fetch, calls = _fetch_sequence(_response("completed", with_usage=True))
|
|
|
|
await poll_and_log_background_interaction_cost(_context(logging_obj, store=store), fetch_interaction=fetch)
|
|
|
|
assert calls == []
|
|
assert logging_obj.model_call_details.get("response_cost") is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_schedule_respects_kill_switch(monkeypatch):
|
|
import litellm.interactions.background_cost_polling as module
|
|
|
|
monkeypatch.setattr(module, "BACKGROUND_INTERACTION_COST_POLLING_ENABLED", False)
|
|
|
|
task = await maybe_schedule_background_interaction_cost_polling(
|
|
response=_response("in_progress", with_usage=False),
|
|
create_kwargs={"litellm_logging_obj": _logging_obj()},
|
|
custom_llm_provider="gemini",
|
|
store=InMemoryBackgroundSettlementStore(),
|
|
)
|
|
|
|
assert task is None
|
|
|
|
|
|
def test_every_status_the_api_can_return_is_either_pollable_or_terminal():
|
|
"""
|
|
The proxy bills a usage-less create in exactly two ways: it polls the
|
|
interaction until it settles, or it recognises the status as terminal and
|
|
settles immediately. A status in neither set is billed by nobody, alerts
|
|
nobody, and releases its budget reservation, which is the zero-spend bug
|
|
this whole module exists to fix.
|
|
|
|
Pinned against the generated spec enum rather than a hand-written list, so
|
|
a status Google adds later breaks this test instead of silently shipping
|
|
another unbilled path.
|
|
"""
|
|
from litellm.interactions.background_cost_polling import _POLLABLE_STATUSES, _TERMINAL_STATUSES
|
|
from litellm.types.interactions.generated import Status1
|
|
|
|
spec_statuses = {member.value for member in Status1}
|
|
handled = _POLLABLE_STATUSES | _TERMINAL_STATUSES
|
|
|
|
assert spec_statuses - handled == set()
|
|
assert handled - spec_statuses == set()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_schedule_creates_poll_task_for_queued_create():
|
|
"""
|
|
``queued`` is the API's not-started-yet state. It carries no usage, so the
|
|
create cannot bill it, and it is not terminal, so nothing settles it:
|
|
without a poll task it is never charged at all.
|
|
"""
|
|
logging_obj = _logging_obj()
|
|
task = await maybe_schedule_background_interaction_cost_polling(
|
|
response=_response("queued", with_usage=False),
|
|
create_kwargs={"litellm_logging_obj": logging_obj},
|
|
custom_llm_provider="gemini",
|
|
store=InMemoryBackgroundSettlementStore(),
|
|
)
|
|
|
|
assert isinstance(task, asyncio.Task)
|
|
task.cancel()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await task
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_poller_bills_an_interaction_that_started_out_queued():
|
|
logging_obj = _logging_obj()
|
|
fetch, calls = _fetch_sequence(
|
|
_response("queued", with_usage=False),
|
|
_response("in_progress", with_usage=False),
|
|
_response("completed", with_usage=True),
|
|
)
|
|
|
|
await poll_and_log_background_interaction_cost(_context(logging_obj), fetch_interaction=fetch)
|
|
|
|
assert len(calls) == 3
|
|
assert logging_obj.model_call_details["response_cost"] > 0
|
|
assert logging_obj.model_call_details["standard_logging_object"]["total_tokens"] == 175
|
|
|
|
|
|
def test_poll_intervals_double_up_to_the_cap_and_stay_inside_the_timeout():
|
|
"""
|
|
The degenerate cases are covered above; this pins the shape the proxy
|
|
actually ships, so an off-by-one in the doubling or in the remaining-budget
|
|
check cannot pass green.
|
|
"""
|
|
intervals = list(_poll_intervals(initial=5.0, maximum=60.0, timeout=3600.0))
|
|
|
|
assert intervals[:6] == [5.0, 10.0, 20.0, 40.0, 60.0, 60.0]
|
|
assert max(intervals) == 60.0
|
|
assert sum(intervals) <= 3600.0
|
|
assert sum(intervals) + 60.0 > 3600.0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_giving_up_on_an_unrecognized_status_says_which_status_it_was(monkeypatch):
|
|
"""
|
|
A status outside both sets polls for the full timeout and then gives up.
|
|
The give-up line is the only trace it leaves, so it has to name the status
|
|
rather than reporting it as an interaction that was merely still running.
|
|
"""
|
|
import litellm.interactions.background_cost_polling as bg
|
|
|
|
errors = []
|
|
monkeypatch.setattr(bg.verbose_logger, "error", lambda *args, **kwargs: errors.append(args))
|
|
|
|
logging_obj = _logging_obj()
|
|
fetch, _ = _fetch_sequence(_response("halted_for_review", with_usage=False))
|
|
|
|
await poll_and_log_background_interaction_cost(
|
|
_context(logging_obj, timeout_seconds=0.01), fetch_interaction=fetch
|
|
)
|
|
|
|
assert len(errors) == 1
|
|
assert "halted_for_review" in errors[0]
|
|
|
|
|
|
KEY_HASH = "0123456789abcdef" * 4
|
|
|
|
FAST_SCHEDULE = PollSchedule(initial_interval_seconds=0.001, max_interval_seconds=0.002, timeout_seconds=1.0)
|
|
|
|
|
|
def _capturing_fetch(response: InteractionsAPIResponse):
|
|
captured = []
|
|
|
|
async def fetch(context):
|
|
captured.append(context)
|
|
return response
|
|
|
|
return fetch, captured
|
|
|
|
|
|
def _create_metadata(**extra) -> dict:
|
|
return {
|
|
"user_api_key": KEY_HASH,
|
|
"user_api_key_team_id": "team-1",
|
|
"user_api_key_auth": object(),
|
|
**extra,
|
|
}
|
|
|
|
|
|
async def _create_on_a_replica_that_then_dies(logging_obj: LitellmLogging, store) -> None:
|
|
import litellm.interactions.background_cost_polling as bg
|
|
|
|
poll_fetch, _ = _fetch_sequence(_response("in_progress", with_usage=False))
|
|
task = await maybe_schedule_background_interaction_cost_polling(
|
|
response=_response("in_progress", with_usage=False),
|
|
create_kwargs={"litellm_logging_obj": logging_obj},
|
|
custom_llm_provider="gemini",
|
|
store=store,
|
|
fetch_interaction=poll_fetch,
|
|
)
|
|
task.cancel()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await task
|
|
await asyncio.sleep(0)
|
|
assert "interactions/bg-abc" not in bg._ACTIVE_POLLS
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_delete_on_another_replica_bills_the_create_from_the_store():
|
|
"""
|
|
The regression: the replica that served the create owns the poll task, so
|
|
a delete served by any other replica used to find nothing to settle and
|
|
the work went unbilled. The store carries the create's attribution, never
|
|
its auth object, to whichever replica settles.
|
|
"""
|
|
store = InMemoryBackgroundSettlementStore()
|
|
logging_obj = _logging_obj(litellm_params={"metadata": _create_metadata()})
|
|
await _create_on_a_replica_that_then_dies(logging_obj, store)
|
|
fetch, captured = _capturing_fetch(_response("completed", with_usage=True))
|
|
|
|
outcome = await maybe_settle_background_interaction_before_delete(
|
|
interaction_id="interactions/bg-abc",
|
|
delete_kwargs={},
|
|
fetch_interaction=fetch,
|
|
store=store,
|
|
)
|
|
|
|
assert outcome == "billed"
|
|
settled = captured[0].logging_obj
|
|
assert settled is not logging_obj
|
|
assert settled.model_call_details["response_cost"] > 0
|
|
payload_metadata = settled.model_call_details["standard_logging_object"]["metadata"]
|
|
assert payload_metadata["user_api_key_hash"] == KEY_HASH
|
|
assert payload_metadata["user_api_key_team_id"] == "team-1"
|
|
assert "user_api_key_auth" not in get_litellm_metadata_from_kwargs(kwargs=settled.model_call_details)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_delete_on_another_replica_releases_the_create_reservation():
|
|
store = InMemoryBackgroundSettlementStore()
|
|
logging_obj = _logging_obj(
|
|
litellm_params={"metadata": _create_metadata(user_api_key_budget_reservation=_reservation())}
|
|
)
|
|
await _create_on_a_replica_that_then_dies(logging_obj, store)
|
|
fetch, captured = _capturing_fetch(_response("in_progress", with_usage=False))
|
|
|
|
outcome = await maybe_settle_background_interaction_before_delete(
|
|
interaction_id="interactions/bg-abc",
|
|
delete_kwargs={},
|
|
fetch_interaction=fetch,
|
|
store=store,
|
|
)
|
|
|
|
assert outcome == "released"
|
|
settled_metadata = get_litellm_metadata_from_kwargs(kwargs=captured[0].logging_obj.model_call_details)
|
|
assert settled_metadata["user_api_key_budget_reservation"]["finalized"] is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_delete_on_another_replica_fails_when_it_cannot_fetch_and_leaves_the_bill_to_the_creating_poll():
|
|
"""
|
|
The settling replica fetches with the delete's credentials, never the
|
|
create's, so a fetch it cannot make (a key only the deployment carries)
|
|
says nothing about the interaction. Deleting anyway would strand the bill
|
|
behind a deleted interaction, so the delete fails with the fetch's error
|
|
and the poll on the creating replica still owns the bill.
|
|
"""
|
|
store = InMemoryBackgroundSettlementStore()
|
|
logging_obj = _logging_obj(litellm_params={"metadata": _create_metadata()})
|
|
await _create_on_a_replica_that_then_dies(logging_obj, store)
|
|
fetch, _ = _fetch_sequence(RuntimeError("Google API key is required"))
|
|
|
|
with pytest.raises(RuntimeError, match="Google API key is required"):
|
|
await maybe_settle_background_interaction_before_delete(
|
|
interaction_id="interactions/bg-abc", delete_kwargs={}, fetch_interaction=fetch, store=store
|
|
)
|
|
|
|
assert await store.is_claimed("interactions/bg-abc") is False
|
|
poll_fetch, _ = _fetch_sequence(_response("completed", with_usage=True))
|
|
await asyncio.wait_for(_register_poll(logging_obj, poll_fetch=poll_fetch, store=store), timeout=5)
|
|
assert logging_obj.model_call_details["response_cost"] > 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_delete_settles_once_however_many_replicas_try():
|
|
store = InMemoryBackgroundSettlementStore()
|
|
await _create_on_a_replica_that_then_dies(_logging_obj(litellm_params={"metadata": _create_metadata()}), store)
|
|
first_fetch, first_calls = _capturing_fetch(_response("completed", with_usage=True))
|
|
second_fetch, second_calls = _capturing_fetch(_response("completed", with_usage=True))
|
|
|
|
first = await maybe_settle_background_interaction_before_delete(
|
|
interaction_id="interactions/bg-abc", delete_kwargs={}, fetch_interaction=first_fetch, store=store
|
|
)
|
|
second = await maybe_settle_background_interaction_before_delete(
|
|
interaction_id="interactions/bg-abc", delete_kwargs={}, fetch_interaction=second_fetch, store=store
|
|
)
|
|
|
|
assert (first, second) == ("billed", None)
|
|
assert len(first_calls) == 1
|
|
assert second_calls == []
|
|
|
|
|
|
def test_create_context_carries_no_request_headers():
|
|
logging_obj = _logging_obj(
|
|
litellm_params={
|
|
"metadata": _create_metadata(
|
|
requester_custom_headers={"x-api-key": "sk-customer-secret"},
|
|
proxy_server_request={"headers": {"x-api-key": "sk-customer-secret"}},
|
|
)
|
|
}
|
|
)
|
|
|
|
carried = _create_context(logging_obj, "gemini").metadata
|
|
|
|
assert carried["user_api_key_team_id"] == "team-1"
|
|
assert "requester_custom_headers" not in carried
|
|
assert "proxy_server_request" not in carried
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_restart_resumes_only_the_rows_no_replica_claimed():
|
|
store = InMemoryBackgroundSettlementStore()
|
|
create_context = _create_context(_logging_obj(litellm_params={"metadata": _create_metadata()}), "gemini")
|
|
for interaction_id in ("interactions/bg-orphaned", "interactions/bg-settled"):
|
|
await store.register(
|
|
PendingBackgroundInteraction(
|
|
interaction_id=interaction_id,
|
|
custom_llm_provider="gemini",
|
|
create_context=create_context,
|
|
created_at=datetime.now(timezone.utc),
|
|
)
|
|
)
|
|
assert await store.claim("interactions/bg-settled")
|
|
fetch, captured = _capturing_fetch(_response("completed", with_usage=True))
|
|
|
|
resumed = await resume_unsettled_background_interactions(store, fetch, schedule=FAST_SCHEDULE)
|
|
|
|
assert len(resumed) == 1
|
|
assert await asyncio.wait_for(resumed[0], timeout=5) == "billed"
|
|
assert [context.interaction_id for context in captured] == ["interactions/bg-orphaned"]
|
|
assert captured[0].logging_obj.model_call_details["response_cost"] > 0
|
|
assert await store.is_claimed("interactions/bg-orphaned")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_a_restart_long_after_the_create_still_retries_a_failed_fetch_before_giving_up():
|
|
"""
|
|
The regression: a resumed poll kept only what was left of the original
|
|
window, so a row older than the window got a single fetch, and one
|
|
transient error right after the restart ended it unbilled.
|
|
"""
|
|
store = InMemoryBackgroundSettlementStore()
|
|
await store.register(
|
|
PendingBackgroundInteraction(
|
|
interaction_id="interactions/bg-abc",
|
|
custom_llm_provider="gemini",
|
|
create_context=_create_context(_logging_obj(litellm_params={"metadata": _create_metadata()}), "gemini"),
|
|
created_at=datetime.now(timezone.utc) - timedelta(seconds=FAST_SCHEDULE.timeout_seconds * 100),
|
|
)
|
|
)
|
|
fetch, calls = _fetch_sequence(
|
|
RuntimeError("transient network error"),
|
|
_response("completed", with_usage=True),
|
|
)
|
|
|
|
(resumed,) = await resume_unsettled_background_interactions(store, fetch, schedule=FAST_SCHEDULE)
|
|
|
|
assert await asyncio.wait_for(resumed, timeout=5) == "billed"
|
|
assert len(calls) == 2
|
|
|
|
|
|
def _create_through_the_proxy_router(tags: list) -> LitellmLogging:
|
|
logging_obj = LitellmLogging(
|
|
model="gemini-2.5-flash",
|
|
messages=[],
|
|
stream=False,
|
|
call_type="acreate_interaction",
|
|
start_time=time.time(),
|
|
litellm_call_id="bg-interactions-call-id",
|
|
function_id="bg-interactions-fn-id",
|
|
)
|
|
logging_obj.update_from_kwargs(
|
|
kwargs={
|
|
"metadata": _create_metadata(
|
|
tags=tags, spend_logs_metadata={"job": "nightly"}, model_group="client-chosen-group"
|
|
),
|
|
"litellm_metadata": {"model_info": {"id": "deployment-1"}, "model_group": "gemini-background"},
|
|
},
|
|
model="gemini-2.5-flash",
|
|
optional_params={},
|
|
litellm_params={"litellm_call_id": "bg-interactions-call-id"},
|
|
custom_llm_provider="gemini",
|
|
)
|
|
return logging_obj
|
|
|
|
|
|
def _create_with_header_tags(tags: list) -> LitellmLogging:
|
|
return _logging_obj(
|
|
litellm_params={
|
|
"metadata": _create_metadata(tags=tags, spend_logs_metadata={"job": "nightly"}),
|
|
"proxy_server_request": {"headers": {"user-agent": "claude-cli/2.1.0"}},
|
|
}
|
|
)
|
|
|
|
|
|
async def _billed_spend_row(logging_obj: LitellmLogging) -> dict:
|
|
response = _response("completed", with_usage=True)
|
|
fetch, _ = _fetch_sequence(response)
|
|
assert await poll_and_log_background_interaction_cost(_context(logging_obj), fetch_interaction=fetch) == "billed"
|
|
settled_at = datetime.now(timezone.utc)
|
|
row = get_logging_payload(
|
|
kwargs={**logging_obj.model_call_details, "completion_start_time": settled_at},
|
|
response_obj=response,
|
|
start_time=settled_at,
|
|
end_time=settled_at,
|
|
)
|
|
return {
|
|
"request_tags": json.loads(row["request_tags"]),
|
|
"model_id": row["model_id"],
|
|
"model_group": row["model_group"],
|
|
"spend_logs_metadata": json.loads(row["metadata"]).get("spend_logs_metadata"),
|
|
}
|
|
|
|
|
|
@pytest.mark.parametrize("create", [_create_through_the_proxy_router, _create_with_header_tags])
|
|
@pytest.mark.asyncio
|
|
async def test_a_bill_settled_off_the_creating_worker_matches_the_creating_worker_spend_row(create):
|
|
"""
|
|
The regressions: the stored create context kept only the router's metadata
|
|
on the proxy path, so a row billed after a restart or by another replica
|
|
lost the request tags that tag-based spend reports group by; then it let
|
|
the client's request metadata override the router's, so a client-sent
|
|
model group relabeled the bill another replica wrote.
|
|
"""
|
|
creator_row = await _billed_spend_row(create(["team-a", "nightly-batch"]))
|
|
stored = _create_context(create(["team-a", "nightly-batch"]), "gemini")
|
|
|
|
rebuilt_row = await _billed_spend_row(_rebuild_logging_obj(stored))
|
|
|
|
assert "team-a" in creator_row["request_tags"]
|
|
assert rebuilt_row == creator_row
|
|
|
|
|
|
class _ClaimAnswersOnlyAfterTheLastFetch:
|
|
def __init__(self):
|
|
self.store = InMemoryBackgroundSettlementStore()
|
|
self.fetches = 0
|
|
self.fetches_at_last_claim = -1
|
|
|
|
async def fetch(self, context):
|
|
self.fetches += 1
|
|
return _response("completed", with_usage=True)
|
|
|
|
async def register(self, pending):
|
|
await self.store.register(pending)
|
|
|
|
async def pending(self, interaction_id):
|
|
return await self.store.pending(interaction_id)
|
|
|
|
async def is_claimed(self, interaction_id):
|
|
return await self.store.is_claimed(interaction_id)
|
|
|
|
async def claim(self, interaction_id):
|
|
if self.fetches != self.fetches_at_last_claim:
|
|
self.fetches_at_last_claim = self.fetches
|
|
raise RuntimeError("database unavailable")
|
|
return await self.store.claim(interaction_id)
|
|
|
|
async def record_outcome(self, interaction_id, outcome):
|
|
return None
|
|
|
|
async def unsettled(self):
|
|
return await self.store.unsettled()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_poller_bills_the_completed_response_it_saw_when_the_claim_only_answers_at_the_deadline():
|
|
logging_obj = _logging_obj()
|
|
store = _ClaimAnswersOnlyAfterTheLastFetch()
|
|
|
|
outcome = await poll_and_log_background_interaction_cost(
|
|
_context(logging_obj, timeout_seconds=0.01, store=store),
|
|
fetch_interaction=store.fetch,
|
|
)
|
|
|
|
assert store.fetches >= 2
|
|
assert outcome == "billed"
|
|
assert logging_obj.model_call_details["response_cost"] > 0
|
|
|
|
|
|
class _DownStore:
|
|
async def register(self, pending):
|
|
raise RuntimeError("database unavailable")
|
|
|
|
async def pending(self, interaction_id):
|
|
raise RuntimeError("database unavailable")
|
|
|
|
async def is_claimed(self, interaction_id):
|
|
raise RuntimeError("database unavailable")
|
|
|
|
async def claim(self, interaction_id):
|
|
raise RuntimeError("database unavailable")
|
|
|
|
async def record_outcome(self, interaction_id, outcome):
|
|
raise RuntimeError("database unavailable")
|
|
|
|
async def unsettled(self):
|
|
raise RuntimeError("database unavailable")
|
|
|
|
|
|
class _RegistersThenRaises:
|
|
def __init__(self):
|
|
self.store = InMemoryBackgroundSettlementStore()
|
|
|
|
async def register(self, pending):
|
|
await self.store.register(pending)
|
|
raise RuntimeError("connection reset after the row was committed")
|
|
|
|
async def pending(self, interaction_id):
|
|
return await self.store.pending(interaction_id)
|
|
|
|
async def is_claimed(self, interaction_id):
|
|
return await self.store.is_claimed(interaction_id)
|
|
|
|
async def claim(self, interaction_id):
|
|
return await self.store.claim(interaction_id)
|
|
|
|
async def record_outcome(self, interaction_id, outcome):
|
|
return None
|
|
|
|
async def unsettled(self):
|
|
return await self.store.unsettled()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_whose_registration_raised_after_landing_still_claims_the_stored_row():
|
|
"""
|
|
A registration that raises after its row committed used to move the poll
|
|
to a private in-memory gate, so the creating worker billed while the
|
|
stored row stayed unclaimed for another replica's delete or the next boot
|
|
to bill again. The row that landed is the gate every settler shares.
|
|
"""
|
|
store = _RegistersThenRaises()
|
|
logging_obj = _logging_obj(litellm_params={"metadata": _create_metadata()})
|
|
poll_fetch, _ = _fetch_sequence(_response("in_progress", with_usage=False))
|
|
task = await maybe_schedule_background_interaction_cost_polling(
|
|
response=_response("in_progress", with_usage=False),
|
|
create_kwargs={"litellm_logging_obj": logging_obj},
|
|
custom_llm_provider="gemini",
|
|
store=store,
|
|
fetch_interaction=poll_fetch,
|
|
)
|
|
fetch, _ = _capturing_fetch(_response("completed", with_usage=True))
|
|
|
|
outcome = await maybe_settle_background_interaction_before_delete(
|
|
interaction_id="interactions/bg-abc", delete_kwargs={}, fetch_interaction=fetch, store=store
|
|
)
|
|
|
|
assert outcome == "billed"
|
|
assert await store.is_claimed("interactions/bg-abc")
|
|
assert await store.unsettled() == ()
|
|
task.cancel()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await task
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_delete_on_a_worker_that_resumed_the_poll_fails_when_it_cannot_fetch():
|
|
"""
|
|
After a restart every worker resumes the unsettled rows, so none of them
|
|
is the creator whose delete may release and delete on a failed fetch. A
|
|
resumed worker's delete fails like any other replica's, and its own poll
|
|
still bills the interaction once it completes.
|
|
"""
|
|
store = InMemoryBackgroundSettlementStore()
|
|
await store.register(
|
|
PendingBackgroundInteraction(
|
|
interaction_id="interactions/bg-abc",
|
|
custom_llm_provider="gemini",
|
|
create_context=_create_context(_logging_obj(litellm_params={"metadata": _create_metadata()}), "gemini"),
|
|
created_at=datetime.now(timezone.utc),
|
|
)
|
|
)
|
|
responses = [_response("in_progress", with_usage=False)]
|
|
|
|
async def poll_fetch(context):
|
|
return responses[-1]
|
|
|
|
(resumed,) = await resume_unsettled_background_interactions(store, poll_fetch, schedule=FAST_SCHEDULE)
|
|
fetch, _ = _fetch_sequence(RuntimeError("Google API key is required"))
|
|
|
|
with pytest.raises(RuntimeError, match="Google API key is required"):
|
|
await maybe_settle_background_interaction_before_delete(
|
|
interaction_id="interactions/bg-abc", delete_kwargs={}, fetch_interaction=fetch, store=store
|
|
)
|
|
|
|
assert await store.is_claimed("interactions/bg-abc") is False
|
|
responses.append(_response("completed", with_usage=True))
|
|
assert await asyncio.wait_for(resumed, timeout=5) == "billed"
|
|
|
|
|
|
class _LandsThenGoesDown:
|
|
"""Register commits the row and loses its acknowledgement; every read fails until the store recovers."""
|
|
|
|
def __init__(self):
|
|
self.store = InMemoryBackgroundSettlementStore()
|
|
self.down = True
|
|
|
|
async def register(self, pending):
|
|
await self.store.register(pending)
|
|
raise RuntimeError("connection reset after the row was committed")
|
|
|
|
async def pending(self, interaction_id):
|
|
self._answer()
|
|
return await self.store.pending(interaction_id)
|
|
|
|
async def is_claimed(self, interaction_id):
|
|
self._answer()
|
|
return await self.store.is_claimed(interaction_id)
|
|
|
|
async def claim(self, interaction_id):
|
|
self._answer()
|
|
return await self.store.claim(interaction_id)
|
|
|
|
async def record_outcome(self, interaction_id, outcome):
|
|
return None
|
|
|
|
async def unsettled(self):
|
|
self._answer()
|
|
return await self.store.unsettled()
|
|
|
|
def _answer(self):
|
|
if self.down:
|
|
raise RuntimeError("database unavailable")
|
|
|
|
|
|
class _TableLessStore:
|
|
"""A replica whose database never got the settlement table: writes fail and reads see no rows."""
|
|
|
|
async def register(self, pending):
|
|
raise RuntimeError("the settlement table does not exist")
|
|
|
|
async def pending(self, interaction_id):
|
|
return None
|
|
|
|
async def is_claimed(self, interaction_id):
|
|
return False
|
|
|
|
async def claim(self, interaction_id):
|
|
return False
|
|
|
|
async def record_outcome(self, interaction_id, outcome):
|
|
raise RuntimeError("the settlement table does not exist")
|
|
|
|
async def unsettled(self):
|
|
raise RuntimeError("the settlement table does not exist")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_whose_registration_and_read_back_both_failed_bills_once_through_the_landed_row():
|
|
"""
|
|
A registration that raised and could not be read back used to give the
|
|
creator a private in-memory gate, so it billed while the stored row stayed
|
|
unclaimed for the next boot to resume and bill again. With the durable
|
|
state unknown, the claim waits for the store and settles through the row.
|
|
"""
|
|
store = _LandsThenGoesDown()
|
|
logging_obj = _logging_obj(litellm_params={"metadata": _create_metadata()})
|
|
poll_fetch, _ = _fetch_sequence(_response("in_progress", with_usage=False))
|
|
task = await maybe_schedule_background_interaction_cost_polling(
|
|
response=_response("in_progress", with_usage=False),
|
|
create_kwargs={"litellm_logging_obj": logging_obj},
|
|
custom_llm_provider="gemini",
|
|
store=store,
|
|
fetch_interaction=poll_fetch,
|
|
)
|
|
fetch, calls = _fetch_sequence(_response("completed", with_usage=True), _response("completed", with_usage=True))
|
|
|
|
while_down = await maybe_settle_background_interaction_before_delete(
|
|
interaction_id="interactions/bg-abc", delete_kwargs={}, fetch_interaction=fetch, store=store
|
|
)
|
|
store.down = False
|
|
recovered = await maybe_settle_background_interaction_before_delete(
|
|
interaction_id="interactions/bg-abc", delete_kwargs={}, fetch_interaction=fetch, store=store
|
|
)
|
|
|
|
assert (while_down, recovered) == (None, "billed")
|
|
assert len(calls) == 2
|
|
assert await store.is_claimed("interactions/bg-abc")
|
|
assert await store.unsettled() == ()
|
|
assert await resume_unsettled_background_interactions(store, poll_fetch, schedule=FAST_SCHEDULE) == ()
|
|
task.cancel()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await task
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_whose_store_never_answers_is_not_billed_through_a_private_gate():
|
|
logging_obj = _logging_obj()
|
|
poll_fetch, _ = _fetch_sequence(_response("in_progress", with_usage=False))
|
|
task = await maybe_schedule_background_interaction_cost_polling(
|
|
response=_response("in_progress", with_usage=False),
|
|
create_kwargs={"litellm_logging_obj": logging_obj},
|
|
custom_llm_provider="gemini",
|
|
store=_DownStore(),
|
|
fetch_interaction=poll_fetch,
|
|
)
|
|
fetch, calls = _fetch_sequence(_response("completed", with_usage=True))
|
|
|
|
outcome = await maybe_settle_background_interaction_before_delete(
|
|
interaction_id="interactions/bg-abc",
|
|
delete_kwargs={},
|
|
fetch_interaction=fetch,
|
|
store=_DownStore(),
|
|
)
|
|
|
|
assert outcome is None
|
|
assert len(calls) == 1
|
|
assert "response_cost" not in logging_obj.model_call_details
|
|
assert not task.done()
|
|
task.cancel()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await task
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create_on_a_replica_without_the_settlement_table_still_settles_in_process():
|
|
logging_obj = _logging_obj()
|
|
poll_fetch, _ = _fetch_sequence(_response("in_progress", with_usage=False))
|
|
task = await maybe_schedule_background_interaction_cost_polling(
|
|
response=_response("in_progress", with_usage=False),
|
|
create_kwargs={"litellm_logging_obj": logging_obj},
|
|
custom_llm_provider="gemini",
|
|
store=_TableLessStore(),
|
|
fetch_interaction=poll_fetch,
|
|
)
|
|
fetch, calls = _fetch_sequence(_response("completed", with_usage=True), _response("completed", with_usage=True))
|
|
|
|
outcome = await maybe_settle_background_interaction_before_delete(
|
|
interaction_id="interactions/bg-abc", delete_kwargs={}, fetch_interaction=fetch, store=_TableLessStore()
|
|
)
|
|
again = await maybe_settle_background_interaction_before_delete(
|
|
interaction_id="interactions/bg-abc", delete_kwargs={}, fetch_interaction=fetch, store=_TableLessStore()
|
|
)
|
|
|
|
assert (outcome, again) == ("billed", None)
|
|
assert len(calls) == 2
|
|
assert logging_obj.model_call_details["response_cost"] > 0
|
|
task.cancel()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await task
|