From 7fb4b427cec0db7d3bc85dd6c2055ca987f2390b Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Thu, 3 Sep 2026 18:36:05 -0400 Subject: [PATCH] feat(batches): snapshot the creating key's org on the managed object row --- .../proxy/common_utils/check_batch_cost.py | 61 +++++++++++++------ .../proxy/hooks/managed_files.py | 1 + .../migration.sql | 4 ++ .../litellm_proxy_extras/schema.prisma | 1 + litellm/models/managed_files.py | 1 + schema.prisma | 1 + .../proxy_unit_tests/test_check_batch_cost.py | 21 ++++++- ..._batch_update_db_managed_output_file_id.py | 4 +- 8 files changed, 74 insertions(+), 20 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260903230000_add_org_id_to_managed_object_table/migration.sql 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 be6e52681fb..595a8d04bed 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -112,39 +112,66 @@ class CheckBatchCost: verbose_proxy_logger.error(f"CheckBatchCost: could not look up user {user_id} for batch {batch_id}: {e}") return {} - async def _get_key_attribution(self, batch_id: str, api_key: str | None) -> tuple[str | None, str | None]: - """Resolve the creating virtual key's (alias, org_id) from its hashed token.""" + async def _get_key_alias(self, batch_id: str, api_key: str | None) -> str | None: + """Resolve the creating virtual key's alias from its hashed token.""" if not api_key: - return None, None + return None try: key_row: prisma_models.LiteLLM_VerificationToken | None = ( await self.prisma_client.db.litellm_verificationtoken.find_unique( where={"token": api_key} ) ) - if key_row is None: - return None, None - return getattr(key_row, "key_alias", None), getattr(key_row, "org_id", None) + return getattr(key_row, "key_alias", None) if key_row is not None else None except Exception as e: verbose_proxy_logger.error(f"CheckBatchCost: could not look up key alias for batch {batch_id}: {e}") - return None, None + return None - async def _get_team_attribution(self, team_id: str | None) -> tuple[str | None, str | None]: - """Resolve a team's (alias, organization_id) from its id.""" + async def _get_team_alias(self, team_id: str | None) -> str | None: + """Resolve a team's alias from its id.""" if not team_id: - return None, None + return None try: team_row: prisma_models.LiteLLM_TeamTable | None = ( await self.prisma_client.db.litellm_teamtable.find_unique( where={"team_id": team_id} ) ) - if team_row is None: - return None, None - return getattr(team_row, "team_alias", None), getattr(team_row, "organization_id", None) + return getattr(team_row, "team_alias", None) if team_row is not None else None except Exception as e: verbose_proxy_logger.error(f"CheckBatchCost: could not look up team alias for team {team_id}: {e}") - return None, None + return None + + async def _get_org_id(self, job: "LiteLLM_ManagedObjectTable", batch_id: str) -> str | None: + """Organization to bill the batch against, snapshotted on the row at creation + like team_id. Rows created before the org_id column existed carry None, so they + fall back to the creating key's org (or its team's) as resolved today.""" + org_id = getattr(job, "org_id", None) + if org_id: + return org_id + api_key = getattr(job, "api_key", None) + team_id = getattr(job, "team_id", None) + try: + if api_key: + key_row: prisma_models.LiteLLM_VerificationToken | None = ( + await self.prisma_client.db.litellm_verificationtoken.find_unique( + where={"token": api_key} + ) + ) + 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} + ) + ) + return getattr(team_row, "organization_id", None) if team_row is not None else None + return None + except Exception as e: + verbose_proxy_logger.error(f"CheckBatchCost: could not resolve org for batch {batch_id}: {e}") + return None async def _build_creator_attribution_metadata( self, job: "LiteLLM_ManagedObjectTable", batch_id: str @@ -173,13 +200,13 @@ class CheckBatchCost: **(await self._get_user_info(batch_id, job.created_by)), } - key_alias, key_org_id = await self._get_key_attribution(batch_id, api_key) + key_alias = await self._get_key_alias(batch_id, api_key) if key_alias is not None: metadata["user_api_key_alias"] = key_alias - team_alias, team_org_id = await self._get_team_attribution(team_id) + team_alias = await self._get_team_alias(team_id) if team_alias is not None: metadata["user_api_key_team_alias"] = team_alias - org_id: Final = key_org_id or team_org_id + org_id: Final = await self._get_org_id(job, batch_id) if org_id is not None: metadata["user_api_key_org_id"] = org_id if isinstance(request_tags, list) and request_tags: diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 5cfcf6129f0..4f0ca787280 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -349,6 +349,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): "file_purpose": file_purpose, "created_by": resolve_resource_owner_id(user_api_key_dict), "team_id": user_api_key_dict.team_id, + "org_id": user_api_key_dict.org_id, "updated_by": user_api_key_dict.user_id, "status": file_object.status, **attribution_columns, diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260903230000_add_org_id_to_managed_object_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260903230000_add_org_id_to_managed_object_table/migration.sql new file mode 100644 index 00000000000..bbe980bb66f --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260903230000_add_org_id_to_managed_object_table/migration.sql @@ -0,0 +1,4 @@ +-- Add org_id column to LiteLLM_ManagedObjectTable +-- Snapshots the creating key's organization at submission time, like team_id, +-- so CheckBatchCost can bill organization spend hours later without re-resolving +ALTER TABLE "LiteLLM_ManagedObjectTable" ADD COLUMN IF NOT EXISTS "org_id" TEXT; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 7604ceadf7a..e386446d5bd 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -1034,6 +1034,7 @@ model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use t created_at DateTime @default(now()) created_by String? team_id String? + org_id String? // creating key's organization at submission time; CheckBatchCost bills org spend against it api_key String? request_tags Json? @default("[]") updated_at DateTime @updatedAt diff --git a/litellm/models/managed_files.py b/litellm/models/managed_files.py index 23d70ef5c48..c90f9b535ea 100644 --- a/litellm/models/managed_files.py +++ b/litellm/models/managed_files.py @@ -32,6 +32,7 @@ class LiteLLM_ManagedObjectTable(LiteLLMPydanticObjectBase): file_object: LiteLLMBatch | LiteLLMFineTuningJob | ResponsesAPIResponse created_by: str | None = None team_id: str | None = None + org_id: str | None = None class LiteLLM_ManagedVectorStoreTable(LiteLLMPydanticObjectBase): diff --git a/schema.prisma b/schema.prisma index 7604ceadf7a..e386446d5bd 100644 --- a/schema.prisma +++ b/schema.prisma @@ -1034,6 +1034,7 @@ model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use t created_at DateTime @default(now()) created_by String? team_id String? + org_id String? // creating key's organization at submission time; CheckBatchCost bills org spend against it api_key String? request_tags Json? @default("[]") updated_at DateTime @updatedAt diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index 9ee0a36f00b..92b946e19b1 100644 --- a/tests/proxy_unit_tests/test_check_batch_cost.py +++ b/tests/proxy_unit_tests/test_check_batch_cost.py @@ -2553,10 +2553,27 @@ class TestBatchCostAttribution: assert metadata["user_api_key_alias"] == "prod-key" + @pytest.mark.asyncio + async def test_org_id_snapshotted_on_the_row_wins(self): + """The org_id column captures the creating key's organization at submission time, + like team_id, so a key later moved to another org still bills the original one.""" + from types import SimpleNamespace + + instance = self._instance( + key_row=SimpleNamespace(key_alias="prod-key", org_id="org-moved-to"), + team_row=SimpleNamespace(team_alias="Team Alpha", organization_id="org-team"), + ) + + metadata = await instance._build_creator_attribution_metadata( + self._job(org_id="org-at-creation"), "batch-1" + ) + + assert metadata["user_api_key_org_id"] == "org-at-creation" + @pytest.mark.asyncio async def test_org_id_comes_from_the_creating_key(self): - """The spend update writer increments organization spend from user_api_key_org_id, - so an org-scoped key's batch cost must carry the key's org id.""" + """The spend update writer increments organization spend from user_api_key_org_id. + A legacy row without the org_id column falls back to the creating key's org.""" from types import SimpleNamespace instance = self._instance( diff --git a/tests/test_litellm/enterprise/proxy/test_batch_update_db_managed_output_file_id.py b/tests/test_litellm/enterprise/proxy/test_batch_update_db_managed_output_file_id.py index ebd33aa2e53..63884e0a779 100644 --- a/tests/test_litellm/enterprise/proxy/test_batch_update_db_managed_output_file_id.py +++ b/tests/test_litellm/enterprise/proxy/test_batch_update_db_managed_output_file_id.py @@ -390,7 +390,7 @@ async def test_store_unified_object_id_persists_key_and_tags_on_create(): """Regression (spend loss): the batch create persists the creating key hash and tags so CheckBatchCost can write an attributed spend row instead of a blank one the DB drops.""" instance, store = _in_memory_managed_files() - creator = UserAPIKeyAuth(user_id="alice", team_id="team-alpha", api_key="hash-alice") + creator = UserAPIKeyAuth(user_id="alice", team_id="team-alpha", api_key="hash-alice", org_id="org-acme") await instance.store_unified_object_id( unified_object_id="unified-b", @@ -407,6 +407,7 @@ async def test_store_unified_object_id_persists_key_and_tags_on_create(): assert row["api_key"] == "hash-alice" assert row["created_by"] == "alice" assert row["team_id"] == "team-alpha" + assert row["org_id"] == "org-acme" assert row["request_tags"].data == ["env:prod"] @@ -471,6 +472,7 @@ async def test_store_unified_object_id_attribution_columns_are_write_once(): upsert_data = instance.prisma_client.db.litellm_managedobjecttable.upsert.call_args.kwargs["data"] assert "api_key" not in upsert_data["update"] assert "request_tags" not in upsert_data["update"] + assert "org_id" not in upsert_data["update"] @pytest.mark.asyncio