mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-28 06:23:07 +00:00
fix(responses): bill a background response's project only while the key stays on its team
This commit is contained in:
parent
a8ac9678e2
commit
ded4f084c5
2 changed files with 53 additions and 2 deletions
|
|
@ -110,7 +110,9 @@ class ManagedObjectCreatorAttribution:
|
|||
A background response also carries its end user, the creating key's project, and the key,
|
||||
team, and user budgets and per-model budgets, the way the read that billed it before did,
|
||||
and never names a key-scoped creator as its user. A key that sets no budget of its own
|
||||
takes its budget tier's, the same rule auth applies.
|
||||
takes its budget tier's, the same rule auth applies. The project is billed only while the
|
||||
key is still on the team the object was created under, so one charge never splits across
|
||||
two teams.
|
||||
"""
|
||||
|
||||
def __init__(self, prisma_client: PrismaClient, poller_name: str) -> None:
|
||||
|
|
@ -128,7 +130,7 @@ class ManagedObjectCreatorAttribution:
|
|||
"user_api_key_user_id": _creator_user_id(job.created_by),
|
||||
**_shared_metadata(job, rows),
|
||||
**({"user_api_key_end_user_id": end_user_id} if end_user_id else {}),
|
||||
"user_api_key_project_id": getattr(rows.key, "project_id", None),
|
||||
"user_api_key_project_id": _project_on_creating_team(rows.key, getattr(job, "team_id", None)),
|
||||
**_key_budget_fields(rows.key),
|
||||
**_budget_fields(
|
||||
rows.team,
|
||||
|
|
@ -227,6 +229,10 @@ def _organization_id(row: object) -> str | None:
|
|||
return cast("str | None", getattr(row, "organization_id", None)) if row is not None else None
|
||||
|
||||
|
||||
def _project_on_creating_team(key: object, creating_team_id: str | None) -> object:
|
||||
return getattr(key, "project_id", None) if getattr(key, "team_id", None) == creating_team_id else None
|
||||
|
||||
|
||||
def _budget_fields(row: object, spend_key: str, max_budget_key: str, model_max_budget_key: str) -> Mapping[str, object]:
|
||||
if row is None:
|
||||
return {}
|
||||
|
|
|
|||
|
|
@ -1517,6 +1517,7 @@ async def test_the_billing_event_carries_the_creating_keys_project_and_every_cre
|
|||
organization_id=None,
|
||||
spend=0.0,
|
||||
max_budget=None,
|
||||
team_id="team-alpha",
|
||||
project_id="project-1",
|
||||
model_max_budget=key_budget,
|
||||
)
|
||||
|
|
@ -1551,6 +1552,50 @@ async def test_the_billing_event_carries_the_creating_keys_project_and_every_cre
|
|||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("key_team_id", "billed_project_id"),
|
||||
[("team-alpha", "project-1"), ("team-beta", None)],
|
||||
ids=["key_still_on_the_creating_team", "key_moved_to_another_team"],
|
||||
)
|
||||
async def test_the_billing_event_bills_the_keys_project_only_while_the_key_stays_on_the_creating_team(
|
||||
key_team_id: str, billed_project_id: str | None
|
||||
):
|
||||
table: Final = _InMemoryManagedObjectTable(
|
||||
rows=(
|
||||
_StoredResponse(
|
||||
id="job-moved",
|
||||
unified_object_id="resp_moved",
|
||||
status="queued",
|
||||
created_by="alice",
|
||||
api_key="hash-alice",
|
||||
team_id="team-alpha",
|
||||
),
|
||||
)
|
||||
)
|
||||
billed_logger: Final = _RecordingBilledLogger()
|
||||
poller: Final = _poller_over(table, billed_logger)
|
||||
poller.prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
|
||||
return_value=SimpleNamespace(
|
||||
key_alias=None,
|
||||
organization_id=None,
|
||||
spend=0.0,
|
||||
max_budget=None,
|
||||
team_id=key_team_id,
|
||||
project_id="project-1",
|
||||
model_max_budget=None,
|
||||
)
|
||||
)
|
||||
|
||||
await poller.bill_if_finished(response_id="resp_moved", response=_terminal_response())
|
||||
|
||||
(billed,) = billed_logger.billed
|
||||
assert (billed.metadata["user_api_key_team_id"], billed.metadata["user_api_key_project_id"]) == (
|
||||
"team-alpha",
|
||||
billed_project_id,
|
||||
)
|
||||
|
||||
|
||||
_TIER_MODEL_BUDGET: Final = {"gpt-5": {"max_budget": 4.0, "budget_duration": "1d"}}
|
||||
_KEY_MODEL_BUDGET: Final = {"gpt-5": {"max_budget": 1.0, "budget_duration": "1d"}}
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue