scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-14 23:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 102

fix(proxy): attribute passthrough batch-cost spend to creating key/team/tags

Browse source 
Vertex and Anthropic passthrough batches built a UserAPIKeyAuth with an empty
api_key and forwarded no tags, so the managed-object row that CheckBatchCost
later reads carried no identity. At cost time the poller had a blank key, so
the batch-cost SpendLogs row was attributed to nobody (or dropped), leaving the
creating key at $0 with blank api_key/team_id/request_tags.

The request metadata already carries the authenticated hashed key, the
user/team ids and the request tags; the old handlers threw them away. Add
user_api_key and request_tags columns to LiteLLM_ManagedObjectTable (schema,
migration, Pydantic model) and persist them in store_unified_object_id, keeping
the identity in the DB row rather than the cached pydantic model. Both
passthrough handlers now go through one shared helper,
store_batch_managed_object, that reads the real identity out of the request
metadata. At cost time CheckBatchCost._build_creator_attribution_metadata
rebuilds the spend-tracking metadata from the row, resolving the key alias from
the hashed token and the team alias from the team id, so the batch-cost row is
attributed exactly like a non-batch request. _get_user_info short-circuits a
null user_id so the poller no longer issues find_unique(where={"user_id": None}),
and every lookup falls back gracefully for older rows that predate the two new
columns.
This commit is contained in:
Yassin Kortam 2026-07-15 13:21:05 -07:00
parent f1f33f560f
commit 90e941a3f8
12 changed files with 362 additions and 128 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -42,11 +42,15 @@ class CheckBatchCost:
# the guaranteed-failing primary query on every subsequent cycle.
self._has_batch_processed_column: bool = True
async def _get_user_info(self, batch_id, user_id) -> dict:
async def _get_user_info(self, batch_id: str, user_id: Optional[str]) -> dict[str, object]:
"""
Look up user email and key alias by user_id for enriching the S3 callback metadata.
Returns a dict with user_api_key_user_email and user_api_key_alias (both may be None).
Look up user email by user_id for enriching the S3 callback metadata.
Returns a dict with user_api_key_user_email (may be None). Returns an
empty dict when user_id is None (batches created by team/service-account
keys carry no user_id, and find_unique(where={"user_id": None}) raises).
"""
if not user_id:
return {}
try:
user_row = await self.prisma_client.db.litellm_usertable.find_unique(
where={"user_id": user_id}
@ -55,12 +59,69 @@ class CheckBatchCost:
return {}
return {
"user_api_key_user_email": getattr(user_row, "user_email", None),
"user_api_key_alias": getattr(user_row, "user_alias", None),
}
except Exception as e:
verbose_proxy_logger.error(f"CheckBatchCost: could not look up user {user_id} for batch {batch_id}: {e}")
return {}
async def _get_key_alias(self, user_api_key: Optional[str]) -> Optional[str]:
"""Resolve the creating virtual key's alias from its hashed token."""
if not user_api_key:
return None
try:
key_row = await self.prisma_client.db.litellm_verificationtoken.find_unique(
where={"token": user_api_key}
)
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 cost attribution: {e}")
return None
async def _get_team_alias(self, team_id: Optional[str]) -> Optional[str]:
"""Resolve a team's alias from its id."""
if not team_id:
return None
try:
team_row = await self.prisma_client.db.litellm_teamtable.find_unique(
where={"team_id": team_id}
)
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
async def _build_creator_attribution_metadata(
self, job: "LiteLLM_ManagedObjectTable", batch_id: str
) -> dict[str, object]:
"""
Rebuild the spend-tracking metadata for the key/team/tags that created the
batch so the batch-cost spend log is attributed the same way a non-batch
request is. Falls back gracefully to whatever identity the managed object
row carries (older rows created before user_api_key/request_tags were
persisted only have created_by/team_id).
"""
user_api_key = getattr(job, "user_api_key", None)
team_id = getattr(job, "team_id", None)
request_tags = getattr(job, "request_tags", None)
metadata: dict[str, object] = {
"user_api_key_user_id": job.created_by,
"user_api_key": user_api_key,
"user_api_key_team_id": team_id,
**(await self._get_user_info(batch_id, job.created_by)),
}
key_alias = await self._get_key_alias(user_api_key)
if key_alias is not None:
metadata["user_api_key_alias"] = key_alias
team_alias = await self._get_team_alias(team_id)
if team_alias is not None:
metadata["user_api_key_team_alias"] = team_alias
if isinstance(request_tags, list) and request_tags:
metadata["tags"] = request_tags
return metadata
async def _cleanup_stale_managed_objects(self) -> None:
"""
Mark managed objects older than MANAGED_OBJECT_STALENESS_CUTOFF_DAYS days
@ -458,8 +519,7 @@ class CheckBatchCost:
function_id=str(uuid.uuid4()),
)
creator_user_id = job.created_by
user_info = await self._get_user_info(batch_id, job.created_by)
attribution_metadata = await self._build_creator_attribution_metadata(job, batch_id)
logging_obj.update_environment_variables(
litellm_params={
@ -469,10 +529,7 @@ class CheckBatchCost:
"user-agent": CHECK_BATCH_COST_USER_AGENT,
}
},
"metadata": {
"user_api_key_user_id": creator_user_id,
**user_info,
},
"metadata": attribution_metadata,
},
optional_params={},
)

6
enterprise/litellm_enterprise/proxy/hooks/managed_files.py
View file

@ -165,10 +165,14 @@ 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:
verbose_logger.info(
f"Storing LiteLLM Managed {file_purpose} object with id={unified_object_id} in cache"
)
from prisma import Json
user_api_key = user_api_key_dict.api_key or None
litellm_managed_object = LiteLLM_ManagedObjectTable(
unified_object_id=unified_object_id,
model_object_id=model_object_id,
@ -191,6 +195,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,
"user_api_key": user_api_key,
"request_tags": Json(request_tags) if request_tags else None,
"updated_by": user_api_key_dict.user_id,
"status": file_object.status,
},

5
litellm-proxy-extras/litellm_proxy_extras/migrations/20260715020000_add_user_api_key_and_request_tags_to_managed_object_table/migration.sql Normal file
View file

@ -0,0 +1,5 @@
-- Add user_api_key and request_tags columns to LiteLLM_ManagedObjectTable
-- Captured at batch-create time so CheckBatchCost can attribute batch-cost spend
-- back to the creating virtual key (and its tags) even when created_by is null.
ALTER TABLE "LiteLLM_ManagedObjectTable" ADD COLUMN IF NOT EXISTS "user_api_key" TEXT;
ALTER TABLE "LiteLLM_ManagedObjectTable" ADD COLUMN IF NOT EXISTS "request_tags" JSONB;

2
litellm-proxy-extras/litellm_proxy_extras/schema.prisma
View file

@ -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?
user_api_key String?
request_tags Json?
updated_at DateTime @updatedAt
updated_by String?

2
litellm/models/managed_files.py
View file

@ -32,6 +32,8 @@ class LiteLLM_ManagedObjectTable(LiteLLMPydanticObjectBase):
file_object: Union[LiteLLMBatch, LiteLLMFineTuningJob, ResponsesAPIResponse]
created_by: Optional[str] = None
team_id: Optional[str] = None
user_api_key: Optional[str] = None
request_tags: Optional[list[str]] = None
class LiteLLM_ManagedVectorStoreTable(LiteLLMPydanticObjectBase):

69
litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py
View file

@ -19,6 +19,9 @@ from litellm.llms.anthropic.chat.handler import (
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
from litellm.proxy._types import PassThroughEndpointLoggingTypedDict
from litellm.proxy.auth.auth_utils import get_end_user_id_from_request_body
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_managed_object_utils import (
store_batch_managed_object,
)
from litellm.types.passthrough_endpoints.pass_through_endpoints import (
PassthroughStandardLoggingPayload,
)
@ -966,65 +969,13 @@ class AnthropicPassthroughLoggingHandler:
Store batch managed object for cost tracking.
This will be picked up by the check_batch_cost polling mechanism.
"""
try:
# Get the managed files hook from the logging object
# This is a bit of a hack, but we need access to the proxy logging system
from litellm.proxy.proxy_server import proxy_logging_obj
managed_files_hook = proxy_logging_obj.get_proxy_hook("managed_files")
if managed_files_hook is not None and hasattr(managed_files_hook, "store_unified_object_id"):
# Create a mock user API key dict for the managed object storage
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
_request_metadata = (kwargs.get("litellm_params", {}) or {}).get("metadata", {}) or {}
user_api_key_dict = UserAPIKeyAuth(
user_id=_request_metadata.get("user_api_key_user_id", "default-user"),
api_key="",
team_id=_request_metadata.get("user_api_key_team_id"),
team_alias=None,
user_role=LitellmUserRoles.CUSTOMER, # Use proper enum value
user_email=None,
max_budget=None,
spend=0.0, # Set to 0.0 instead of None
models=[], # Set to empty list instead of None
tpm_limit=None,
rpm_limit=None,
budget_duration=None,
budget_reset_at=None,
max_parallel_requests=None,
allowed_model_region=None,
metadata={}, # Set to empty dict instead of None
key_alias=None,
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
)
# Store the unified object for batch cost tracking
import asyncio
asyncio.create_task(
managed_files_hook.store_unified_object_id( # type: ignore
unified_object_id=unified_object_id,
file_object=batch_object,
litellm_parent_otel_span=None,
model_object_id=model_object_id,
file_purpose="batch",
user_api_key_dict=user_api_key_dict,
)
)
verbose_proxy_logger.info(
f"Stored Anthropic batch managed object with unified_object_id={unified_object_id}, batch_id={model_object_id}"
)
else:
verbose_proxy_logger.warning(
"Managed files hook not available, cannot store batch object for cost tracking"
)
except Exception as e:
verbose_proxy_logger.error(f"Error storing Anthropic batch managed object: {e}")
request_metadata = (kwargs.get("litellm_params", {}) or {}).get("metadata", {}) or {}
store_batch_managed_object(
unified_object_id=unified_object_id,
batch_object=batch_object,
model_object_id=model_object_id,
request_metadata=request_metadata,
)
@staticmethod
def get_actual_model_id_from_router(model_name: str) -> str:

96
litellm/proxy/pass_through_endpoints/llm_provider_handlers/batch_managed_object_utils.py Normal file
View file

@ -0,0 +1,96 @@
"""
Shared helper for persisting a passthrough-created batch as a managed object.
Vertex AI and Anthropic passthrough both create a provider batch job and then
need to register it for asynchronous cost tracking by CheckBatchCost. The
creating key's identity (hashed key, team, tags) is captured here so the
eventual batch-cost spend log can be attributed back to it, mirroring the
attribution a non-batch request receives.
"""
import asyncio
from typing import TYPE_CHECKING, Optional, Protocol, cast
from litellm._logging import verbose_proxy_logger
from litellm.types.utils import LiteLLMBatch
if TYPE_CHECKING:
from litellm.proxy._types import UserAPIKeyAuth
class _StoreUnifiedObjectHook(Protocol):
async def store_unified_object_id(
self,
unified_object_id: str,
file_object: LiteLLMBatch,
litellm_parent_otel_span: None,
model_object_id: str,
file_purpose: str,
user_api_key_dict: "UserAPIKeyAuth",
request_tags: Optional[list[str]],
) -> None: ...
def _optional_str(value: object) -> Optional[str]:
return value if isinstance(value, str) else None
def _optional_str_list(value: object) -> Optional[list[str]]:
if isinstance(value, list):
items = cast(list[object], value) # cast-ok: isinstance-narrowed; element type unknown
return [str(tag) for tag in items]
return None
def store_batch_managed_object(
unified_object_id: str,
batch_object: LiteLLMBatch,
model_object_id: str,
request_metadata: dict[str, object],
) -> None:
"""
Register a provider batch job as a managed object so CheckBatchCost can
later compute its cost and write an attributed spend log.
request_metadata is the request's ``litellm_params["metadata"]``, which
carries the authenticated key's ``user_api_key`` hash, user/team ids and
request tags.
"""
try:
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.proxy_server import proxy_logging_obj
managed_files_hook = proxy_logging_obj.get_proxy_hook("managed_files")
if managed_files_hook is None or not hasattr(managed_files_hook, "store_unified_object_id"):
verbose_proxy_logger.warning(
"Managed files hook not available, cannot store batch object for cost tracking"
)
return
user_api_key_dict = UserAPIKeyAuth(
user_id=_optional_str(request_metadata.get("user_api_key_user_id")) or "default-user",
api_key=_optional_str(request_metadata.get("user_api_key")),
team_id=_optional_str(request_metadata.get("user_api_key_team_id")),
team_alias=_optional_str(request_metadata.get("user_api_key_team_alias")),
key_alias=_optional_str(request_metadata.get("user_api_key_alias")),
user_role=LitellmUserRoles.CUSTOMER,
)
hook = cast(_StoreUnifiedObjectHook, managed_files_hook) # cast-ok: presence checked via hasattr above
asyncio.create_task(
hook.store_unified_object_id(
unified_object_id=unified_object_id,
file_object=batch_object,
litellm_parent_otel_span=None,
model_object_id=model_object_id,
file_purpose="batch",
user_api_key_dict=user_api_key_dict,
request_tags=_optional_str_list(request_metadata.get("tags")),
)
)
verbose_proxy_logger.info(
f"Stored batch managed object with unified_object_id={unified_object_id}, batch_id={model_object_id}"
)
except Exception as e:
verbose_proxy_logger.error(f"Error storing batch managed object: {e}")

69
litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py
View file

@ -16,6 +16,9 @@ from litellm.llms.vertex_ai.vector_stores.search_api.transformation import (
)
from litellm.llms.vertex_ai.videos.transformation import VertexAIVideoConfig
from litellm.proxy._types import PassThroughEndpointLoggingTypedDict
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_managed_object_utils import (
store_batch_managed_object,
)
from litellm.types.utils import (
Choices,
EmbeddingResponse,
@ -791,65 +794,13 @@ class VertexPassthroughLoggingHandler:
Store batch managed object for cost tracking.
This will be picked up by the check_batch_cost polling mechanism.
"""
try:
# Get the managed files hook from the logging object
# This is a bit of a hack, but we need access to the proxy logging system
from litellm.proxy.proxy_server import proxy_logging_obj
managed_files_hook = proxy_logging_obj.get_proxy_hook("managed_files")
if managed_files_hook is not None and hasattr(managed_files_hook, "store_unified_object_id"):
# Create a mock user API key dict for the managed object storage
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
_request_metadata = (kwargs.get("litellm_params", {}) or {}).get("metadata", {}) or {}
user_api_key_dict = UserAPIKeyAuth(
user_id=_request_metadata.get("user_api_key_user_id", "default-user"),
api_key="",
team_id=_request_metadata.get("user_api_key_team_id"),
team_alias=None,
user_role=LitellmUserRoles.CUSTOMER, # Use proper enum value
user_email=None,
max_budget=None,
spend=0.0, # Set to 0.0 instead of None
models=[], # Set to empty list instead of None
tpm_limit=None,
rpm_limit=None,
budget_duration=None,
budget_reset_at=None,
max_parallel_requests=None,
allowed_model_region=None,
metadata={}, # Set to empty dict instead of None
key_alias=None,
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
)
# Store the unified object for batch cost tracking
import asyncio
asyncio.create_task(
managed_files_hook.store_unified_object_id( # type: ignore
unified_object_id=unified_object_id,
file_object=batch_object,
litellm_parent_otel_span=None,
model_object_id=model_object_id,
file_purpose="batch",
user_api_key_dict=user_api_key_dict,
)
)
verbose_proxy_logger.info(
f"Stored batch managed object with unified_object_id={unified_object_id}, batch_id={model_object_id}"
)
else:
verbose_proxy_logger.warning(
"Managed files hook not available, cannot store batch object for cost tracking"
)
except Exception as e:
verbose_proxy_logger.error(f"Error storing batch managed object: {e}")
request_metadata = (kwargs.get("litellm_params", {}) or {}).get("metadata", {}) or {}
store_batch_managed_object(
unified_object_id=unified_object_id,
batch_object=batch_object,
model_object_id=model_object_id,
request_metadata=request_metadata,
)
@staticmethod
def get_actual_model_id_from_router(model_name: str) -> str:

2
litellm/proxy/schema.prisma
View file

@ -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?
user_api_key String?
request_tags Json?
updated_at DateTime @updatedAt
updated_by String?

2
schema.prisma
View file

@ -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?
user_api_key String?
request_tags Json?
updated_at DateTime @updatedAt
updated_by String?

123
tests/proxy_unit_tests/test_check_batch_cost.py
View file

@ -1221,3 +1221,126 @@ class TestUnmanagedBatchCostFlagIsGeneralized:
assert vertex_result == ("deploy-vertex", "8823717160934178816")
assert bedrock_result == ("deploy-bedrock", TestUnmanagedBedrockRouting._ARN)
class TestBatchCostAttribution:
"""A batch-cost spend log must be attributed to the creating key/team/tags the same
way a non-batch request is. Regression for batches created by keys with no user_id,
where the row was previously dropped and identity was blank."""
def _instance(self):
from litellm_enterprise.proxy.common_utils.check_batch_cost import (
CheckBatchCost,
)
instance = CheckBatchCost(
proxy_logging_obj=MagicMock(),
prisma_client=MagicMock(),
llm_router=MagicMock(),
)
instance.prisma_client.db = MagicMock()
return instance
@pytest.mark.asyncio
async def test_get_user_info_skips_query_when_user_id_none(self):
"""find_unique(where={"user_id": None}) raises "A value is required but not set";
the None user_id case (team/service-account keys) must short-circuit instead."""
instance = self._instance()
instance.prisma_client.db.litellm_usertable.find_unique = AsyncMock()
result = await instance._get_user_info("batch-1", None)
assert result == {}
instance.prisma_client.db.litellm_usertable.find_unique.assert_not_called()
@pytest.mark.asyncio
async def test_attribution_metadata_replays_key_team_and_tags(self):
from types import SimpleNamespace
instance = self._instance()
instance.prisma_client.db.litellm_usertable.find_unique = AsyncMock(
return_value=SimpleNamespace(user_email="creator@example.test")
)
instance.prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
return_value=SimpleNamespace(key_alias="prod-batch-key")
)
instance.prisma_client.db.litellm_teamtable.find_unique = AsyncMock(
return_value=SimpleNamespace(team_alias="growth-team")
)
job = SimpleNamespace(
created_by="user-1",
team_id="team-1",
user_api_key="hashed-key-abc",
request_tags=["env:prod", "team:growth"],
)
metadata = await instance._build_creator_attribution_metadata(job, "batch-1")
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["user_api_key_alias"] == "prod-batch-key"
assert metadata["user_api_key_team_alias"] == "growth-team"
assert metadata["user_api_key_user_email"] == "creator@example.test"
assert metadata["tags"] == ["env:prod", "team:growth"]
@pytest.mark.asyncio
async def test_attribution_metadata_keeps_key_when_no_user_id(self):
"""A team/service-account key has no created_by. The persisted key hash and
team_id must still flow through so _should_track_cost_callback writes the row."""
from types import SimpleNamespace
from litellm.proxy.hooks.proxy_track_cost_callback import (
_should_track_cost_callback,
)
instance = self._instance()
instance.prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
return_value=None
)
instance.prisma_client.db.litellm_teamtable.find_unique = AsyncMock(
return_value=None
)
job = SimpleNamespace(
created_by=None,
team_id="team-1",
user_api_key="hashed-key-abc",
request_tags=None,
)
metadata = await instance._build_creator_attribution_metadata(job, "batch-1")
assert metadata["user_api_key"] == "hashed-key-abc"
assert metadata["user_api_key_user_id"] is None
assert "tags" not in metadata
assert (
_should_track_cost_callback(
user_api_key=metadata["user_api_key"],
user_id=metadata.get("user_api_key_user_id"),
team_id=metadata.get("user_api_key_team_id"),
end_user_id=None,
)
is True
)
@pytest.mark.asyncio
async def test_attribution_metadata_tolerates_legacy_rows(self):
"""Older managed-object rows predate user_api_key/request_tags. Attribution must
still fall back to created_by/team_id without raising."""
from types import SimpleNamespace
instance = self._instance()
instance.prisma_client.db.litellm_usertable.find_unique = AsyncMock(
return_value=None
)
instance.prisma_client.db.litellm_teamtable.find_unique = AsyncMock(
return_value=None
)
job = SimpleNamespace(created_by="user-1", team_id="team-1")
metadata = await instance._build_creator_attribution_metadata(job, "batch-1")
assert metadata["user_api_key"] is None
assert metadata["user_api_key_user_id"] == "user-1"
assert metadata["user_api_key_team_id"] == "team-1"
assert "tags" not in metadata

37
tests/test_litellm/proxy/pass_through_endpoints/test_vertex_ai_batch_passthrough.py
View file

@ -315,6 +315,43 @@ 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_persists_key_hash_and_tags(
self, mock_logging_obj, mock_managed_files_hook
):
"""The creating key hash and request tags must be persisted so CheckBatchCost
can later attribute the batch-cost spend log. Regression: previously api_key
was hardcoded to "" and tags were never forwarded. The metadata already carries
the hashed token, so it must survive the UserAPIKeyAuth round-trip unchanged."""
hashed_key = "a" * 64
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,
"user_api_key_team_id": "team-456",
"user_api_key_alias": "prod-batch-key",
"tags": ["env:prod", "team:growth"],
}
},
)
mock_managed_files_hook.store_unified_object_id.assert_called_once()
call_kwargs = mock_managed_files_hook.store_unified_object_id.call_args[1]
assert call_kwargs["user_api_key_dict"].api_key == hashed_key
assert call_kwargs["user_api_key_dict"].key_alias == "prod-batch-key"
assert call_kwargs["request_tags"] == ["env:prod", "team:growth"]
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