scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-10 22:41:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 129

fix(batches): keep team org attribution when the key lookup fails

Browse source
This commit is contained in:
mubashir1osmani 2026-09-03 19:09:45 -04:00
parent 0d7976116e
commit 35371a34c1
2 changed files with 32 additions and 9 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

24
enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
View file

@ -151,8 +151,8 @@ class CheckBatchCost:
return org_id
api_key = getattr(job, "api_key", None)
team_id = getattr(job, "team_id", None)
try:
if api_key:
if api_key:
try:
key_row: prisma_models.LiteLLM_VerificationToken | None = (
await self.prisma_client.db.litellm_verificationtoken.find_unique(
where={"token": api_key}
@ -161,16 +161,22 @@ class CheckBatchCost:
key_org_id = getattr(key_row, "org_id", None) if key_row is not None else None
if key_org_id:
return key_org_id
if team_id:
team_row: prisma_models.LiteLLM_TeamTable | None = (
await self.prisma_client.db.litellm_teamtable.find_unique(
where={"team_id": team_id}
)
except Exception as e:
verbose_proxy_logger.error(
f"CheckBatchCost: could not resolve the key's org for batch {batch_id}, "
f"still trying the team's: {e}"
)
return getattr(team_row, "organization_id", None) if team_row is not None else None
if not team_id:
return None
try:
team_row: prisma_models.LiteLLM_TeamTable | None = (
await self.prisma_client.db.litellm_teamtable.find_unique(
where={"team_id": team_id}
)
)
return getattr(team_row, "organization_id", None) if team_row is not None else None
except Exception as e:
verbose_proxy_logger.error(f"CheckBatchCost: could not resolve org for batch {batch_id}: {e}")
verbose_proxy_logger.error(f"CheckBatchCost: could not resolve the team's org for batch {batch_id}: {e}")
return None
async def _build_creator_attribution_metadata(

17
tests/proxy_unit_tests/test_check_batch_cost.py
View file

@ -2600,6 +2600,23 @@ class TestBatchCostAttribution:
assert metadata["user_api_key_org_id"] == "org-team"
@pytest.mark.asyncio
async def test_key_lookup_failure_still_bills_the_team_org(self):
"""A key-table error while resolving a legacy row's org must not drop the team's
organization: the two lookups fail independently, so org spend still lands."""
from types import SimpleNamespace
instance = self._instance(
team_row=SimpleNamespace(team_alias="Team Alpha", organization_id="org-team"),
)
instance.prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
side_effect=Exception("db down")
)
metadata = await instance._build_creator_attribution_metadata(self._job(), "batch-1")
assert metadata["user_api_key_org_id"] == "org-team"
@pytest.mark.asyncio
async def test_no_org_leaves_the_key_unset(self):
"""Without any org the key is absent entirely, so the spend writer's org update