diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py index 595a8d04bed..bd7f3c93ff3 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -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( diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index 92b946e19b1..59dfb67c486 100644 --- a/tests/proxy_unit_tests/test_check_batch_cost.py +++ b/tests/proxy_unit_tests/test_check_batch_cost.py @@ -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