From e8ed5906f25f2ca8be8fbb60bc6b37053dc35093 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 16 Jul 2026 15:58:16 +0000 Subject: [PATCH] fix(proxy): attribute vertex batch passthrough cost to creating key/team/tags --- .../proxy/common_utils/check_batch_cost.py | 49 ++++++++- .../proxy/hooks/managed_files.py | 5 + .../migration.sql | 3 + .../litellm_proxy_extras/schema.prisma | 2 + litellm/models/managed_files.py | 2 + .../anthropic_passthrough_logging_handler.py | 8 +- .../vertex_passthrough_logging_handler.py | 8 +- litellm/proxy/schema.prisma | 2 + schema.prisma | 2 + .../proxy/hooks/test_managed_files.py | 101 ++++++++++++++++++ .../proxy_unit_tests/test_check_batch_cost.py | 84 +++++++++++++++ .../test_vertex_ai_batch_passthrough.py | 40 +++++++ 12 files changed, 296 insertions(+), 10 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260716120000_add_api_key_and_request_tags_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 f209ab54f64..9e893d53d76 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -61,6 +61,45 @@ class CheckBatchCost: verbose_proxy_logger.error(f"CheckBatchCost: could not look up user {user_id} for batch {batch_id}: {e}") return {} + async def _build_spend_attribution_metadata( + self, + job: "LiteLLM_ManagedObjectTable", + creator_user_id: Optional[str], + user_info: dict, + ) -> dict: + """ + Build spend-log metadata that attributes the batch cost to the virtual key, team, + and tags captured at batch-create time. Key and team aliases are resolved via the + shared cost-callback enrichment so the row matches a non-batch request's attribution. + + Rows created before api_key was persisted (legacy) fall back to the previous + created_by/user_info-only behavior. + """ + from litellm.proxy.hooks.proxy_track_cost_callback import _ProxyDBLogger + + api_key = job.api_key if isinstance(job.api_key, str) and job.api_key else None + team_id = job.team_id if isinstance(job.team_id, str) and job.team_id else None + request_tags = job.request_tags if isinstance(job.request_tags, list) else None + + if not api_key: + return { + "user_api_key_user_id": creator_user_id, + **user_info, + } + + metadata = { + "user_api_key_user_id": creator_user_id, + "user_api_key": api_key, + **({"user_api_key_team_id": team_id} if team_id else {}), + **({"tags": request_tags} if request_tags else {}), + } + metadata = await _ProxyDBLogger._enrich_failure_metadata_with_key_info(metadata=metadata) + + user_email = user_info.get("user_api_key_user_email") + if user_email is not None and metadata.get("user_api_key_user_email") is None: + metadata = {**metadata, "user_api_key_user_email": user_email} + return metadata + async def _cleanup_stale_managed_objects(self) -> None: """ Mark managed objects older than MANAGED_OBJECT_STALENESS_CUTOFF_DAYS days @@ -460,6 +499,11 @@ class CheckBatchCost: creator_user_id = job.created_by user_info = await self._get_user_info(batch_id, job.created_by) + spend_metadata = await self._build_spend_attribution_metadata( + job=job, + creator_user_id=creator_user_id, + user_info=user_info, + ) logging_obj.update_environment_variables( litellm_params={ @@ -469,10 +513,7 @@ class CheckBatchCost: "user-agent": CHECK_BATCH_COST_USER_AGENT, } }, - "metadata": { - "user_api_key_user_id": creator_user_id, - **user_info, - }, + "metadata": spend_metadata, }, optional_params={}, ) diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 3f42867d90e..0087022b7af 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -165,7 +165,10 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): model_object_id: str, file_purpose: Literal["batch", "fine-tune", "response"], user_api_key_dict: UserAPIKeyAuth, + request_tags: Optional[List[str]] = None, ) -> None: + from prisma import Json + verbose_logger.info( f"Storing LiteLLM Managed {file_purpose} object with id={unified_object_id} in cache" ) @@ -191,6 +194,8 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): "file_purpose": file_purpose, "created_by": user_api_key_dict.user_id, "team_id": user_api_key_dict.team_id, + "api_key": user_api_key_dict.api_key or None, + "request_tags": Json(request_tags) if request_tags else None, "updated_by": user_api_key_dict.user_id, "status": file_object.status, }, diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260716120000_add_api_key_and_request_tags_to_managed_object_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260716120000_add_api_key_and_request_tags_to_managed_object_table/migration.sql new file mode 100644 index 00000000000..45a1b7e5549 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260716120000_add_api_key_and_request_tags_to_managed_object_table/migration.sql @@ -0,0 +1,3 @@ +-- AlterTable +ALTER TABLE "LiteLLM_ManagedObjectTable" ADD COLUMN IF NOT EXISTS "api_key" TEXT; +ALTER TABLE "LiteLLM_ManagedObjectTable" ADD COLUMN IF NOT EXISTS "request_tags" JSONB; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index a23cecc3911..fd670be2f21 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -942,6 +942,8 @@ model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use t created_at DateTime @default(now()) created_by String? team_id String? + api_key String? + request_tags Json? updated_at DateTime @updatedAt updated_by String? diff --git a/litellm/models/managed_files.py b/litellm/models/managed_files.py index 99ba764dd98..5830a96f6b5 100644 --- a/litellm/models/managed_files.py +++ b/litellm/models/managed_files.py @@ -32,6 +32,8 @@ class LiteLLM_ManagedObjectTable(LiteLLMPydanticObjectBase): file_object: Union[LiteLLMBatch, LiteLLMFineTuningJob, ResponsesAPIResponse] created_by: Optional[str] = None team_id: Optional[str] = None + api_key: Optional[str] = None + request_tags: Optional[list[str]] = None class LiteLLM_ManagedVectorStoreTable(LiteLLMPydanticObjectBase): diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py index 50e90699194..e3c6bc50977 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py @@ -978,11 +978,12 @@ class AnthropicPassthroughLoggingHandler: _request_metadata = (kwargs.get("litellm_params", {}) or {}).get("metadata", {}) or {} + _request_tags = _request_metadata.get("tags") user_api_key_dict = UserAPIKeyAuth( user_id=_request_metadata.get("user_api_key_user_id", "default-user"), - api_key="", + api_key=_request_metadata.get("user_api_key") or "", team_id=_request_metadata.get("user_api_key_team_id"), - team_alias=None, + team_alias=_request_metadata.get("user_api_key_team_alias"), user_role=LitellmUserRoles.CUSTOMER, # Use proper enum value user_email=None, max_budget=None, @@ -995,7 +996,7 @@ class AnthropicPassthroughLoggingHandler: max_parallel_requests=None, allowed_model_region=None, metadata={}, # Set to empty dict instead of None - key_alias=None, + key_alias=_request_metadata.get("user_api_key_alias"), permissions={}, # Set to empty dict instead of None model_max_budget={}, # Set to empty dict instead of None model_spend={}, # Set to empty dict instead of None @@ -1012,6 +1013,7 @@ class AnthropicPassthroughLoggingHandler: model_object_id=model_object_id, file_purpose="batch", user_api_key_dict=user_api_key_dict, + request_tags=_request_tags if isinstance(_request_tags, list) else None, ) ) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py index c61d48eda8c..bfa81b0919a 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py @@ -803,11 +803,12 @@ class VertexPassthroughLoggingHandler: _request_metadata = (kwargs.get("litellm_params", {}) or {}).get("metadata", {}) or {} + _request_tags = _request_metadata.get("tags") user_api_key_dict = UserAPIKeyAuth( user_id=_request_metadata.get("user_api_key_user_id", "default-user"), - api_key="", + api_key=_request_metadata.get("user_api_key") or "", team_id=_request_metadata.get("user_api_key_team_id"), - team_alias=None, + team_alias=_request_metadata.get("user_api_key_team_alias"), user_role=LitellmUserRoles.CUSTOMER, # Use proper enum value user_email=None, max_budget=None, @@ -820,7 +821,7 @@ class VertexPassthroughLoggingHandler: max_parallel_requests=None, allowed_model_region=None, metadata={}, # Set to empty dict instead of None - key_alias=None, + key_alias=_request_metadata.get("user_api_key_alias"), permissions={}, # Set to empty dict instead of None model_max_budget={}, # Set to empty dict instead of None model_spend={}, # Set to empty dict instead of None @@ -837,6 +838,7 @@ class VertexPassthroughLoggingHandler: model_object_id=model_object_id, file_purpose="batch", user_api_key_dict=user_api_key_dict, + request_tags=_request_tags if isinstance(_request_tags, list) else None, ) ) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index a23cecc3911..fd670be2f21 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -942,6 +942,8 @@ model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use t created_at DateTime @default(now()) created_by String? team_id String? + api_key String? + request_tags Json? updated_at DateTime @updatedAt updated_by String? diff --git a/schema.prisma b/schema.prisma index a23cecc3911..fd670be2f21 100644 --- a/schema.prisma +++ b/schema.prisma @@ -942,6 +942,8 @@ model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use t created_at DateTime @default(now()) created_by String? team_id String? + api_key String? + request_tags Json? updated_at DateTime @updatedAt updated_by String? diff --git a/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py b/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py index 68dc3269f34..73ab63af458 100644 --- a/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py +++ b/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py @@ -2367,3 +2367,104 @@ async def test_same_user_different_keys_can_access_batch(): assert "batch_id" in result2 # Both keys should get the same result assert result1["batch_id"] == result2["batch_id"] + + +@pytest.mark.asyncio +async def test_store_unified_object_id_persists_request_identity(): + """ + store_unified_object_id must persist the creating key's hashed token, team, and request + tags onto the managed object so CheckBatchCost can later attribute the batch-cost SpendLogs + row. Regression for batches whose cost row was dropped because identity was never carried. + """ + from prisma import Json + + from litellm.proxy._types import UserAPIKeyAuth + from litellm.types.utils import LiteLLMBatch + + prisma_client = AsyncMock() + prisma_client.db.litellm_managedobjecttable.upsert = AsyncMock( + return_value=MagicMock() + ) + internal_usage_cache = MagicMock() + internal_usage_cache.async_set_cache = AsyncMock() + + proxy_managed_files = _PROXY_LiteLLMManagedFiles( + internal_usage_cache=internal_usage_cache, + prisma_client=prisma_client, + ) + + batch_object = LiteLLMBatch( + id="batch-123", + completion_window="24h", + created_at=1, + endpoint="/v1/chat/completions", + input_file_id="gs://bucket/in.jsonl", + object="batch", + status="validating", + ) + + await proxy_managed_files.store_unified_object_id( + unified_object_id="unified-123", + file_object=batch_object, + litellm_parent_otel_span=None, + model_object_id="batch-123", + file_purpose="batch", + user_api_key_dict=UserAPIKeyAuth( + user_id="user-1", team_id="team-1", api_key="hashed-key-abc" + ), + request_tags=["tag-a", "tag-b"], + ) + + create_data = prisma_client.db.litellm_managedobjecttable.upsert.call_args.kwargs[ + "data" + ]["create"] + assert create_data["api_key"] == "hashed-key-abc" + assert create_data["created_by"] == "user-1" + assert create_data["team_id"] == "team-1" + assert create_data["request_tags"] == Json(["tag-a", "tag-b"]) + + +@pytest.mark.asyncio +async def test_store_unified_object_id_no_tags_stores_null_request_tags(): + """When no request tags are present the column is left NULL rather than an empty-list Json.""" + from litellm.proxy._types import UserAPIKeyAuth + from litellm.types.utils import LiteLLMBatch + + prisma_client = AsyncMock() + prisma_client.db.litellm_managedobjecttable.upsert = AsyncMock( + return_value=MagicMock() + ) + internal_usage_cache = MagicMock() + internal_usage_cache.async_set_cache = AsyncMock() + + proxy_managed_files = _PROXY_LiteLLMManagedFiles( + internal_usage_cache=internal_usage_cache, + prisma_client=prisma_client, + ) + + batch_object = LiteLLMBatch( + id="batch-123", + completion_window="24h", + created_at=1, + endpoint="/v1/chat/completions", + input_file_id="gs://bucket/in.jsonl", + object="batch", + status="validating", + ) + + await proxy_managed_files.store_unified_object_id( + unified_object_id="unified-123", + file_object=batch_object, + litellm_parent_otel_span=None, + model_object_id="batch-123", + file_purpose="batch", + user_api_key_dict=UserAPIKeyAuth(user_id="user-1", api_key="hashed-key-abc"), + ) + + create_data = prisma_client.db.litellm_managedobjecttable.upsert.call_args.kwargs[ + "data" + ]["create"] + assert create_data["request_tags"] is None + assert create_data["api_key"] == "hashed-key-abc" + + diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index b822799fb40..46ddab3de90 100644 --- a/tests/proxy_unit_tests/test_check_batch_cost.py +++ b/tests/proxy_unit_tests/test_check_batch_cost.py @@ -1221,3 +1221,87 @@ class TestUnmanagedBatchCostFlagIsGeneralized: assert vertex_result == ("deploy-vertex", "8823717160934178816") assert bedrock_result == ("deploy-bedrock", TestUnmanagedBedrockRouting._ARN) + + +class _Job: + """Minimal stand-in for a LiteLLM_ManagedObjectTable row with just the identity + fields _build_spend_attribution_metadata reads.""" + + def __init__(self, created_by=None, api_key=None, team_id=None, request_tags=None): + self.created_by = created_by + self.api_key = api_key + self.team_id = team_id + self.request_tags = request_tags + + +class TestSpendAttributionMetadata: + """_build_spend_attribution_metadata must carry the create-time identity into the + spend-log metadata so the batch-cost SpendLogs row is written and attributed to the + creating key/team/tags (issue #33316).""" + + def _instance(self): + from litellm_enterprise.proxy.common_utils.check_batch_cost import ( + CheckBatchCost, + ) + + return CheckBatchCost( + proxy_logging_obj=MagicMock(), + prisma_client=MagicMock(), + llm_router=MagicMock(), + ) + + @pytest.mark.asyncio + async def test_carries_key_team_and_tags_and_resolves_aliases(self): + instance = self._instance() + job = _Job( + created_by="user-1", + api_key="hashed-key-abc", + team_id="team-1", + request_tags=["tag-a", "tag-b"], + ) + + async def _fake_enrich(metadata): + return { + **metadata, + "user_api_key_alias": "my-key", + "user_api_key_team_alias": "my-team", + } + + with patch( + "litellm.proxy.hooks.proxy_track_cost_callback._ProxyDBLogger._enrich_failure_metadata_with_key_info", + side_effect=_fake_enrich, + ): + metadata = await instance._build_spend_attribution_metadata( + job=job, + creator_user_id="user-1", + user_info={"user_api_key_user_email": "u@example.com"}, + ) + + assert metadata["user_api_key"] == "hashed-key-abc" + assert metadata["user_api_key_user_id"] == "user-1" + assert metadata["user_api_key_team_id"] == "team-1" + assert metadata["tags"] == ["tag-a", "tag-b"] + assert metadata["user_api_key_alias"] == "my-key" + assert metadata["user_api_key_team_alias"] == "my-team" + assert metadata["user_api_key_user_email"] == "u@example.com" + + @pytest.mark.asyncio + async def test_legacy_row_without_api_key_falls_back(self): + """Rows created before api_key was persisted keep the previous created_by/user_info + behavior and never invoke key enrichment.""" + instance = self._instance() + job = _Job(created_by="user-1", api_key=None, team_id=None, request_tags=None) + + with patch( + "litellm.proxy.hooks.proxy_track_cost_callback._ProxyDBLogger._enrich_failure_metadata_with_key_info", + ) as mock_enrich: + metadata = await instance._build_spend_attribution_metadata( + job=job, + creator_user_id="user-1", + user_info={"user_api_key_user_email": "u@example.com"}, + ) + + mock_enrich.assert_not_called() + assert "user_api_key" not in metadata + assert metadata["user_api_key_user_id"] == "user-1" + assert metadata["user_api_key_user_email"] == "u@example.com" diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_ai_batch_passthrough.py b/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_ai_batch_passthrough.py index 5de682ec8a0..537dcf15481 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_ai_batch_passthrough.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_ai_batch_passthrough.py @@ -315,6 +315,46 @@ class TestVertexAIBatchPassthroughHandler: assert call_kwargs["user_api_key_dict"].user_id == expected_user_id assert call_kwargs["user_api_key_dict"].team_id == expected_team_id + def test_store_batch_managed_object_propagates_key_aliases_and_tags( + self, mock_logging_obj, mock_managed_files_hook + ): + """The fabricated UserAPIKeyAuth must carry the hashed key and aliases, and the request + tags must be forwarded, so CheckBatchCost can later write an attributed batch-cost + SpendLogs row. Regression for issue #33316 where api_key was hardcoded to "" and tags + were dropped, causing the cost row to be silently discarded.""" + with ( + patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_pl, + patch( + "litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_passthrough_logging_handler.verbose_proxy_logger" + ), + ): + mock_pl.get_proxy_hook.return_value = mock_managed_files_hook + + VertexPassthroughLoggingHandler._store_batch_managed_object( + unified_object_id="uoi", + batch_object={"id": "b1", "object": "batch", "status": "validating"}, + model_object_id="b1", + logging_obj=mock_logging_obj, + litellm_params={ + "metadata": { + "user_api_key": "hashed-key-abc", + "user_api_key_user_id": "real-user-123", + "user_api_key_team_id": "team-456", + "user_api_key_alias": "my-key", + "user_api_key_team_alias": "my-team", + "tags": ["tag-a", "tag-b"], + } + }, + ) + + mock_managed_files_hook.store_unified_object_id.assert_called_once() + call_kwargs = mock_managed_files_hook.store_unified_object_id.call_args[1] + user_api_key_dict = call_kwargs["user_api_key_dict"] + assert user_api_key_dict.api_key == "hashed-key-abc" + assert user_api_key_dict.key_alias == "my-key" + assert user_api_key_dict.team_alias == "my-team" + assert call_kwargs["request_tags"] == ["tag-a", "tag-b"] + def test_batch_cost_calculation_integration(self): """Single Vertex AI response → non-zero cost with correct token counts.""" from litellm.batches.batch_utils import calculate_vertex_ai_batch_cost_and_usage