mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-29 06:33:28 +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,
|
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,
|
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
|
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:
|
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),
|
"user_api_key_user_id": _creator_user_id(job.created_by),
|
||||||
**_shared_metadata(job, rows),
|
**_shared_metadata(job, rows),
|
||||||
**({"user_api_key_end_user_id": end_user_id} if end_user_id else {}),
|
**({"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),
|
**_key_budget_fields(rows.key),
|
||||||
**_budget_fields(
|
**_budget_fields(
|
||||||
rows.team,
|
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
|
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]:
|
def _budget_fields(row: object, spend_key: str, max_budget_key: str, model_max_budget_key: str) -> Mapping[str, object]:
|
||||||
if row is None:
|
if row is None:
|
||||||
return {}
|
return {}
|
||||||
|
|
|
||||||
|
|
@ -1517,6 +1517,7 @@ async def test_the_billing_event_carries_the_creating_keys_project_and_every_cre
|
||||||
organization_id=None,
|
organization_id=None,
|
||||||
spend=0.0,
|
spend=0.0,
|
||||||
max_budget=None,
|
max_budget=None,
|
||||||
|
team_id="team-alpha",
|
||||||
project_id="project-1",
|
project_id="project-1",
|
||||||
model_max_budget=key_budget,
|
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"}}
|
_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"}}
|
_KEY_MODEL_BUDGET: Final = {"gpt-5": {"max_budget": 1.0, "budget_duration": "1d"}}
|
||||||
|
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue