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 e6f00877a26..13e9e5093a8 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -142,6 +142,40 @@ class CheckBatchCost: verbose_proxy_logger.error(f"CheckBatchCost: could not look up team alias for team {team_id}: {e}") return None + async def _get_org_id(self, job: "LiteLLM_ManagedObjectTable", batch_id: str) -> str | None: + 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) + if api_key: + try: + 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, "organization_id", None) if key_row is not None else None + if key_org_id: + return key_org_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}" + ) + 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 the team's org for batch {batch_id}: {e}") + return None + async def _build_creator_attribution_metadata( self, job: "LiteLLM_ManagedObjectTable", batch_id: str ) -> dict[str, object]: @@ -153,6 +187,10 @@ class CheckBatchCost: user_api_key_alias; when it has no alias, or the key has since been rotated or deleted, the field keeps the creating user's alias that _get_user_info filled in, because a resolvable name is more useful on the spend row than a null. + + user_api_key_org_id must be resolved here too: the spend update writer reads it + off this metadata to increment organization spend, so leaving it out silently + drops batch cost from org accounting for keys and teams that belong to one. """ api_key = getattr(job, "api_key", None) team_id = getattr(job, "team_id", None) @@ -172,6 +210,9 @@ class CheckBatchCost: 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 = 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: metadata["tags"] = [tag for tag in request_tags if isinstance(tag, str)] @@ -641,7 +682,7 @@ class CheckBatchCost: from litellm.files.main import afile_content from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging - from litellm.litellm_core_utils.litellm_logging import deployment_pricing_model_info + from litellm.litellm_core_utils.litellm_logging import deployment_pricing_model_info, mask_api_base_credentials from litellm.proxy.openai_files_endpoints.common_utils import ( _is_base64_encoded_unified_file_id, ) @@ -805,6 +846,7 @@ class CheckBatchCost: function_id=str(uuid.uuid4()), ) + deployment_api_base: Final = deployment_info.litellm_params.api_base logging_obj.update_environment_variables( litellm_params={ # set the user-agent header so that S3 callback consumers can easily identify CheckBatchCost callbacks @@ -813,9 +855,17 @@ class CheckBatchCost: "user-agent": CHECK_BATCH_COST_USER_AGENT, } }, - "metadata": await self._build_creator_attribution_metadata(job, batch_id), + **({"api_base": mask_api_base_credentials(deployment_api_base)} if deployment_api_base else {}), + "metadata": { + **(await self._build_creator_attribution_metadata(job, batch_id)), + # spend logs read the deployment identity off these metadata keys, so + # without them the batch cost row carries no model_id or model_group + "model_info": {"id": model_id}, + "model_group": deployment_info.model_name, + }, }, optional_params={}, + custom_llm_provider=str(llm_provider) if llm_provider else None, ) if not await self._claim_job_for_costing(job): @@ -833,6 +883,8 @@ class CheckBatchCost: batch_models=batch_result.models, batch_successful_requests=batch_result.successful_requests, batch_failed_requests=batch_result.failed_requests, + batch_prompt_cost=batch_result.prompt_cost, + batch_completion_cost=batch_result.completion_cost, ) except Exception: await self._release_job_claim(job) diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index bc1eb6cebc2..486904d0abe 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -280,6 +280,27 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): ) verbose_logger.debug(f"LiteLLM Managed File object with id={file_id} stored in db: {result}") + async def _resolve_creator_org_id(self, user_api_key_dict: UserAPIKeyAuth) -> Optional[str]: + if user_api_key_dict.org_id: + return user_api_key_dict.org_id + if not user_api_key_dict.team_id: + return None + from litellm.proxy.auth.auth_checks import get_team_object + from litellm.proxy.proxy_server import proxy_logging_obj, user_api_key_cache + + try: + team: Final = await get_team_object( + team_id=user_api_key_dict.team_id, + prisma_client=self.prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=user_api_key_dict.parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + return team.organization_id + except Exception as e: + verbose_logger.warning(f"could not resolve org for managed object attribution: {e}") + return None + async def store_unified_object_id( self, unified_object_id: str, @@ -352,6 +373,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": await self._resolve_creator_org_id(user_api_key_dict), "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 1c43668f227..a84ae272af6 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -1035,6 +1035,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/batches/batch_utils.py b/litellm/batches/batch_utils.py index 97be5f77d79..871fdf79182 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -10,7 +10,7 @@ from litellm._logging import verbose_logger from litellm.litellm_core_utils.get_litellm_params import AWS_CREDENTIAL_KWARGS_KEYS from litellm.litellm_core_utils.llm_cost_calc.utils import parse_prompt_tokens_details from litellm.types.llms.openai import Batch -from litellm.types.utils import CallTypes, ModelInfo, Usage +from litellm.types.utils import ModelInfo, Usage from litellm.utils import token_counter @@ -23,6 +23,8 @@ class BatchCostUsageResult: models: list[str] successful_requests: int failed_requests: int + prompt_cost: float = 0.0 + completion_cost: float = 0.0 async def calculate_batch_cost_and_usage( @@ -130,7 +132,8 @@ class _LineOutcome(Enum): @dataclass(frozen=True, slots=True) class _BatchOutputLineStats: - cost: float + prompt_cost: float + completion_cost: float prompt_tokens: int completion_tokens: int total_tokens: int @@ -193,15 +196,16 @@ def _compute_output_line_stats( raw_model: Final = response_body.get("model") response_model: Final = raw_model if isinstance(raw_model, str) and raw_model else None completion_details: Final = usage.completion_tokens_details + line_prompt_cost, line_completion_cost = _output_line_cost( + usage=usage, + custom_llm_provider=custom_llm_provider, + model_name=model_name, + response_model=response_model, + model_info=model_info, + ) return _BatchOutputLineStats( - cost=_output_line_cost( - response_body=response_body, - usage=usage, - custom_llm_provider=custom_llm_provider, - model_name=model_name, - response_model=response_model, - model_info=model_info, - ), + prompt_cost=line_prompt_cost, + completion_cost=line_completion_cost, prompt_tokens=usage.prompt_tokens, completion_tokens=usage.completion_tokens, total_tokens=usage.total_tokens, @@ -213,31 +217,24 @@ def _compute_output_line_stats( def _output_line_cost( - response_body: Mapping[str, object], usage: Usage, custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"], model_name: str | None, response_model: str | None, model_info: ModelInfo | None, -) -> float: +) -> tuple[float, float]: + """(prompt_cost, completion_cost) for one output line, priced at batch rates.""" from litellm.cost_calculator import batch_cost_calculator - if model_info is None and custom_llm_provider not in ("anthropic", "bedrock"): - return litellm.completion_cost( - completion_response=response_body, - custom_llm_provider=custom_llm_provider, - call_type=CallTypes.aretrieve_batch.value, - ) cost_model: Final = ( model_name if custom_llm_provider == "bedrock" and model_name else response_model or model_name or "" ) - prompt_cost, completion_cost = batch_cost_calculator( + return batch_cost_calculator( usage=usage, model=cost_model, custom_llm_provider=custom_llm_provider, model_info=model_info, ) - return prompt_cost + completion_cost def _aggregate_batch_cost_usage_models( @@ -270,7 +267,9 @@ def _aggregate_batch_cost_usage_models( **cache_token_params, ) batch_models: Final = [model_name] if model_name else [stats.model for stats in line_stats if stats.model] - total_cost: Final = sum((stats.cost for stats in line_stats), 0.0) + total_prompt_cost: Final = sum((stats.prompt_cost for stats in line_stats), 0.0) + total_completion_cost: Final = sum((stats.completion_cost for stats in line_stats), 0.0) + total_cost: Final = total_prompt_cost + total_completion_cost verbose_logger.debug( "batch output aggregate: cost=%s usage=%s models=%s successful=%d failed=%d", total_cost, @@ -285,6 +284,8 @@ def _aggregate_batch_cost_usage_models( models=batch_models, successful_requests=successful_requests, failed_requests=failed_requests, + prompt_cost=total_prompt_cost, + completion_cost=total_completion_cost, ) @@ -309,7 +310,8 @@ def calculate_vertex_ai_batch_cost_and_usage( """ from litellm.cost_calculator import batch_cost_calculator - total_cost = 0.0 + total_prompt_cost = 0.0 # rebind-ok: loop accumulator, matches total_tokens below + total_completion_cost = 0.0 # rebind-ok: loop accumulator, matches total_tokens below total_tokens = 0 prompt_tokens = 0 completion_tokens = 0 @@ -341,7 +343,8 @@ def calculate_vertex_ai_batch_cost_and_usage( model=actual_model_name, custom_llm_provider="vertex_ai", ) - total_cost += p_cost + c_cost + total_prompt_cost += p_cost + total_completion_cost += c_cost except Exception as e: verbose_logger.debug("vertex_ai batch cost calculation error for line: %s", str(e)) @@ -349,6 +352,7 @@ def calculate_vertex_ai_batch_cost_and_usage( completion_tokens += _completion total_tokens += _total + total_cost: Final = total_prompt_cost + total_completion_cost verbose_logger.info( "vertex_ai batch cost: cost=%s, prompt=%d, completion=%d, total=%d, successful=%d, failed=%d", total_cost, @@ -369,6 +373,8 @@ def calculate_vertex_ai_batch_cost_and_usage( models=[actual_model_name], successful_requests=successful_requests, failed_requests=failed_requests, + prompt_cost=total_prompt_cost, + completion_cost=total_completion_cost, ) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index c31c4323157..38f8fedc2b9 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -422,6 +422,13 @@ def _provider_response_id(source: object) -> str | None: return candidate if isinstance(candidate, str) and candidate else None +def mask_api_base_credentials(api_base: str) -> str: + if "key=" not in api_base: + return api_base + key_end: Final = api_base.find("key=") + 4 + return api_base[:key_end] + "*" * 5 + api_base[-4:] + + class Logging(LiteLLMLoggingBaseClass): global \ supabaseClient, \ @@ -1164,14 +1171,7 @@ class Logging(LiteLLMLoggingBaseClass): return data def _get_masked_api_base(self, api_base: str) -> str: - if "key=" in api_base: - # Find the position of "key=" in the string - key_index: Final = api_base.find("key=") + 4 - # Mask the last 5 characters after "key=" - masked_api_base = api_base[:key_index] + "*" * 5 + api_base[-4:] - else: - masked_api_base = api_base - return str(masked_api_base) + return str(mask_api_base_credentials(api_base)) def _pre_call(self, input, api_key, model=None, additional_args={}): """ @@ -2922,6 +2922,19 @@ class Logging(LiteLLMLoggingBaseClass): result._hidden_params["batch_successful_requests"] = batch_successful_requests # pyright: ignore[reportPrivateUsage] # rebind-ok: same result._hidden_params pattern as response_cost/batch_models above result._hidden_params["batch_failed_requests"] = batch_failed_requests # pyright: ignore[reportPrivateUsage] # rebind-ok: same pattern as above result.usage = batch_usage + batch_prompt_cost: Final = kwargs.get("batch_prompt_cost", None) + batch_completion_cost: Final = kwargs.get("batch_completion_cost", None) + if ( + isinstance(batch_prompt_cost, float) + and isinstance(batch_completion_cost, float) + and isinstance(batch_cost, float) + ): + self.set_cost_breakdown( + input_cost=batch_prompt_cost, + output_cost=batch_completion_cost, + total_cost=batch_cost, + cost_for_built_in_tools_cost_usd_dollar=0.0, + ) elif should_compute_batch_data: batch_result: Final = await _handle_completed_batch( @@ -2937,6 +2950,12 @@ class Logging(LiteLLMLoggingBaseClass): result._hidden_params["batch_successful_requests"] = batch_result.successful_requests # pyright: ignore[reportPrivateUsage] # rebind-ok: same pattern as above result._hidden_params["batch_failed_requests"] = batch_result.failed_requests # pyright: ignore[reportPrivateUsage] # rebind-ok: same pattern as above result.usage = batch_result.usage + self.set_cost_breakdown( + input_cost=batch_result.prompt_cost, + output_cost=batch_result.completion_cost, + total_cost=batch_result.cost, + cost_for_built_in_tools_cost_usd_dollar=0.0, + ) self.truncated_messages_for_logging = await truncate_base64_in_messages_async( StandardLoggingPayloadSetup.append_system_prompt_messages( 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/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 1c43668f227..a84ae272af6 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -1035,6 +1035,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/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 8a06bf68b81..a21d761996f 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -590,6 +590,7 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs metadata=metadata, standard_logging_payload=standard_logging_payload, omit_when_missing=_omits_session_id_when_missing(metadata), + batch_trace_session_id=_get_batch_trace_session_id(call_type=call_type, request_id=id), ), request_duration_ms=_get_request_duration_ms(start_time, end_time), status=_get_status_for_spend_log( @@ -628,20 +629,44 @@ def _omits_session_id_when_missing(metadata: Mapping[str, object] | None) -> boo return general_settings.get("missing_session_id") == "omit" +_BATCH_TRACE_CALL_TYPES: Final = frozenset( + { + CallTypes.create_batch.value, + CallTypes.acreate_batch.value, + CallTypes.retrieve_batch.value, + CallTypes.aretrieve_batch.value, + } +) + + +def _get_batch_trace_session_id(call_type: str | None, request_id: str | None) -> str | None: + """A batch's create row and its poller-written cost row both derive their request id + from the same batch id (the cost row appends BATCH_COST_REQUEST_ID_SUFFIX), so using + that id as the session groups the batch lifecycle into one trace on the logs UI. The + poller builds its own logging context, so per-request trace ids can never link them.""" + if call_type not in _BATCH_TRACE_CALL_TYPES or not request_id: + return None + return request_id.removesuffix(BATCH_COST_REQUEST_ID_SUFFIX) + + def _get_session_id_for_spend_log( kwargs: Mapping[str, object], metadata: Mapping[str, object] | None, standard_logging_payload: StandardLoggingPayload | None, omit_when_missing: bool, + batch_trace_session_id: str | None = None, ) -> str | None: """Under `omit` only `metadata.session_id`, the key Langfuse reads, counts as a session; `litellm_session_id` may - be a copied trace id.""" + be a copied trace id. Batch call types carry a deterministic session derived from the batch id, which outranks + the per-request trace ids because those differ between the create call and the cost poller's row.""" if omit_when_missing: session_id: Final = metadata.get("session_id") if metadata else None return str(session_id) if session_id else None from litellm._uuid import uuid + if batch_trace_session_id is not None: + return batch_trace_session_id if standard_logging_payload is not None and standard_logging_payload.get("trace_id") is not None: return str(standard_logging_payload.get("trace_id")) if kwargs.get("litellm_trace_id") is not None: diff --git a/schema.prisma b/schema.prisma index 1c43668f227..a84ae272af6 100644 --- a/schema.prisma +++ b/schema.prisma @@ -1035,6 +1035,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/batches_tests/test_batches_logging_unit_tests.py b/tests/batches_tests/test_batches_logging_unit_tests.py index 5bde40d90b0..73adb391481 100644 --- a/tests/batches_tests/test_batches_logging_unit_tests.py +++ b/tests/batches_tests/test_batches_logging_unit_tests.py @@ -144,18 +144,20 @@ def test_get_batch_job_total_usage_from_file_content(sample_file_content_dict): @pytest.mark.asyncio async def test_batch_cost_calculator(sample_file_content_dict): """ - mock litellm.completion_cost to return 0.5 + mock batch_cost_calculator to return (0.3, 0.2) per line we know sample_file_content_dict has 2 successful responses - so we expect the cost to be 0.5 * 2 = 1.0 + so we expect the cost to be (0.3 + 0.2) * 2 = 1.0, split 0.6 / 0.4 """ - with patch("litellm.completion_cost", return_value=0.5): + with patch("litellm.cost_calculator.batch_cost_calculator", return_value=(0.3, 0.2)): result = _aggregate_batch_cost_usage_models( entries=sample_file_content_dict, custom_llm_provider="openai", ) - assert result.cost == 1.0 # 0.5 * 2 successful responses + assert result.cost == pytest.approx(1.0) # (0.3 + 0.2) * 2 successful responses + assert result.prompt_cost == pytest.approx(0.6) + assert result.completion_cost == pytest.approx(0.4) def test_get_response_from_batch_job_output_file(sample_file_content_dict): @@ -402,6 +404,56 @@ async def test_batch_retrieve_cost_tracking_with_explicit_cost_data(): assert mock_batch.usage == explicit_usage +@pytest.mark.asyncio +async def test_batch_retrieve_explicit_cost_split_sets_cost_breakdown(): + """The poller passes the batch's prompt/completion cost split so the spend row's + cost_breakdown carries real input/output costs; without it the UI's Cost Breakdown + card renders blank for every batch. Regression for the split being dropped.""" + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.types.utils import CallTypes, LiteLLMBatch + + mock_batch = LiteLLMBatch( + id="batch-breakdown-1", + object="batch", + endpoint="/v1/chat/completions", + errors=None, + input_file_id="file-input-1", + completion_window="24h", + status="completed", + output_file_id="file-output-1", + created_at=1234567890, + ) + mock_batch._hidden_params = {} + + logging_obj = Logging( + model="gpt-5-mini", + messages=[{"role": "user", "content": "test"}], + stream=False, + call_type=CallTypes.aretrieve_batch.value, + litellm_call_id="test-call-breakdown", + function_id="test-function", + start_time=time.time(), + dynamic_success_callbacks=[], + ) + logging_obj.custom_llm_provider = "openai" + + await logging_obj.async_success_handler( + result=mock_batch, + start_time=time.time(), + end_time=time.time() + 1, + batch_cost=0.10, + batch_usage=litellm.Usage(prompt_tokens=200, completion_tokens=100, total_tokens=300), + batch_models=["gpt-5-mini"], + batch_prompt_cost=0.06, + batch_completion_cost=0.04, + ) + + assert logging_obj.cost_breakdown is not None + assert logging_obj.cost_breakdown["input_cost"] == 0.06 + assert logging_obj.cost_breakdown["output_cost"] == 0.04 + assert logging_obj.cost_breakdown["total_cost"] == 0.10 + + @pytest.mark.asyncio async def test_batch_retrieve_cost_tracking_with_unified_file_id_incomplete_batch(): """ diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index 9a6ab08e9b6..36177b44930 100644 --- a/tests/proxy_unit_tests/test_check_batch_cost.py +++ b/tests/proxy_unit_tests/test_check_batch_cost.py @@ -583,6 +583,90 @@ class TestCheckBatchCost: assert passed_model_info["input_cost_per_token_batches"] == 2e-06 assert passed_model_info["output_cost_per_token_batches"] == 4e-06 + @pytest.mark.asyncio + async def test_poller_masks_api_base_credentials_before_logging( + self, check_batch_cost_instance, mock_prisma_client, mock_llm_router + ): + """Request rows mask `key=` query credentials out of api_base before it is + logged, but the poller skips that pre-call step, so an unmasked deployment + api_base would land verbatim on the batch cost row: regression test for the + poller masking the same way. + """ + import base64 + from unittest.mock import patch + + import httpx + import respx + + from litellm.litellm_core_utils.litellm_logging import Logging + + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=1) + mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + + mock_job = MagicMock() + mock_job.id = "job-masked-api-base-1" + mock_job.unified_object_id = base64.urlsafe_b64encode( + b"litellm_proxy;model_id:model-123;llm_batch_id:batch-456" + ).decode() + mock_job.created_by = "user-1" + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) + + mock_response = MagicMock() + mock_response.status = "completed" + mock_response.output_file_id = "file-output-123" + mock_response.error_file_id = None + mock_response.model_dump_json.return_value = '{"id":"batch-1","status":"completed"}' + mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) + mock_llm_router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"}) + + mock_deployment = MagicMock() + mock_deployment.litellm_params.custom_llm_provider = "openai" + mock_deployment.litellm_params.model = "gpt-5.4-mini" + mock_deployment.litellm_params.api_base = "https://gateway.example.com/v1?key=AIzaSyVERYSECRET7890" + mock_deployment.model_info.model_dump.return_value = {} + mock_llm_router.get_deployment = MagicMock(return_value=mock_deployment) + + output_line = json.dumps( + { + "custom_id": "req-1", + "response": { + "status_code": 200, + "body": { + "id": "chatcmpl-1", + "object": "chat.completion", + "model": "gpt-5.4-mini", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "hi"}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + }, + }, + "error": None, + } + ) + + with ( + respx.mock(assert_all_called=True) as provider, + patch.object( # test-quality-ok: the poller builds Logging inline, the only seam to the row it logs + Logging, "async_success_handler", autospec=True + ) as success_handler, + ): + provider.get("https://api.openai.com/v1/files/file-output-123/content").mock( + return_value=httpx.Response(200, content=f"{output_line}\n".encode()) + ) + await check_batch_cost_instance.check_batch_cost() + + cost_row_calls = [call for call in success_handler.await_args_list if "batch_cost" in call.kwargs] + assert len(cost_row_calls) == 1 + logged_api_base = cost_row_calls[0].args[0].litellm_params["api_base"] + assert logged_api_base == "https://gateway.example.com/v1?key=*****7890" + assert "VERYSECRET" not in logged_api_base + @pytest.mark.asyncio async def test_primary_path_completion_update_includes_batch_processed( self, check_batch_cost_instance, mock_prisma_client, mock_llm_router @@ -2554,6 +2638,85 @@ 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", organization_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. + A legacy row without the org_id column falls back to the creating key's org.""" + from types import SimpleNamespace + + instance = self._instance( + key_row=SimpleNamespace(key_alias="prod-key", organization_id="org-42"), + team_row=SimpleNamespace(team_alias="Team Alpha", organization_id="org-team"), + ) + + metadata = await instance._build_creator_attribution_metadata(self._job(), "batch-1") + + assert metadata["user_api_key_org_id"] == "org-42" + + @pytest.mark.asyncio + async def test_org_id_falls_back_to_the_team_organization(self): + """A key with no org of its own still books batch spend against its team's + organization, matching how the request path resolves org attribution.""" + from types import SimpleNamespace + + instance = self._instance( + key_row=SimpleNamespace(key_alias="prod-key", organization_id=None), + team_row=SimpleNamespace(team_alias="Team Alpha", organization_id="org-team"), + ) + + 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_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 + stays skipped instead of matching an empty-string organization.""" + from types import SimpleNamespace + + instance = self._instance( + key_row=SimpleNamespace(key_alias="prod-key", organization_id=None), + team_row=SimpleNamespace(team_alias="Team Alpha", organization_id=None), + ) + + metadata = await instance._build_creator_attribution_metadata(self._job(), "batch-1") + + assert "user_api_key_org_id" not in metadata + @pytest.mark.asyncio async def test_metadata_provenance_keeps_spend_log_api_key_joinable(self): """ diff --git a/tests/test_litellm/batches/test_batch_utils.py b/tests/test_litellm/batches/test_batch_utils.py index c86c7c4df03..1fd78870481 100644 --- a/tests/test_litellm/batches/test_batch_utils.py +++ b/tests/test_litellm/batches/test_batch_utils.py @@ -489,7 +489,9 @@ def test_aggregate_counts_successful_and_failed_requests(monkeypatch): def test_aggregate_returns_batch_cost_usage_result_dataclass(monkeypatch): - monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 1.0) + import litellm.cost_calculator as cc + + monkeypatch.setattr(cc, "batch_cost_calculator", lambda **kw: (0.4, 0.6)) result = bu._aggregate_batch_cost_usage_models( entries=[_success_row(usage=_usage(10, 5))], custom_llm_provider="openai" ) @@ -500,6 +502,7 @@ def test_aggregate_returns_batch_cost_usage_result_dataclass(monkeypatch): 1, 0, ) + assert (result.prompt_cost, result.completion_cost) == (0.4, 0.6) # =========================================================================== # @@ -507,15 +510,17 @@ def test_aggregate_returns_batch_cost_usage_result_dataclass(monkeypatch): # =========================================================================== # -def test_cost_from_content_completion_cost_path(monkeypatch): - # model_info is None -> litellm.completion_cost per successful row. +def test_cost_without_model_info_prices_each_row_by_its_response_model(monkeypatch): + # model_info is None -> batch_cost_calculator per successful row, model from the response body. + import litellm.cost_calculator as cc + calls = [] - def _completion_cost(**kw): + def _batch_cost(**kw): calls.append(kw) - return 0.5 + return (0.3, 0.2) - monkeypatch.setattr(litellm, "completion_cost", _completion_cost) + monkeypatch.setattr(cc, "batch_cost_calculator", _batch_cost) rows = [ _success_row(usage=_usage(10, 5)), _failed_row(), # excluded -> not costed @@ -524,8 +529,10 @@ def test_cost_from_content_completion_cost_path(monkeypatch): result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai") - assert result.cost == 1.0 # 2 successful * 0.5 + assert result.cost == pytest.approx(1.0) # 2 successful * (0.3 + 0.2) + assert (result.prompt_cost, result.completion_cost) == (pytest.approx(0.6), pytest.approx(0.4)) assert len(calls) == 2 # failed row not costed + assert all(call["model"] == "gpt-4o" and call["model_info"] is None for call in calls) assert result.successful_requests == 2 assert result.failed_requests == 1 @@ -578,7 +585,9 @@ def test_aggregate_consumes_entries_in_a_single_pass(monkeypatch): """A one-shot generator: any implementation that iterates the entries twice (e.g. separate cost and usage passes) sees nothing on the second pass and returns wrong totals for at least one of cost/usage/models.""" - monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 0.5) + import litellm.cost_calculator as cc + + monkeypatch.setattr(cc, "batch_cost_calculator", lambda **kw: (0.25, 0.25)) one_shot = (row for row in [_success_row(usage=_usage(10, 5)), _failed_row(), _success_row(usage=_usage(20, 10))]) result = bu._aggregate_batch_cost_usage_models(entries=one_shot, custom_llm_provider="openai") @@ -753,12 +762,15 @@ def test_vertex_cost_error_in_line_is_swallowed(monkeypatch): @pytest.mark.asyncio async def test_calculate_batch_cost_and_usage_orchestration(monkeypatch): + import litellm.cost_calculator as cc + rows = [_success_row(model="gpt-4o", usage=_usage(10, 5))] - monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 2.5) + monkeypatch.setattr(cc, "batch_cost_calculator", lambda **kw: (1.5, 1.0)) result = await bu.calculate_batch_cost_and_usage(file_content_dictionary=rows, custom_llm_provider="openai") assert result.cost == 2.5 + assert (result.prompt_cost, result.completion_cost) == (1.5, 1.0) assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (10, 5, 15) assert result.models == ["gpt-4o"] @@ -1107,8 +1119,10 @@ async def test_handle_completed_batch_orchestration(monkeypatch): async def fake_fetch(batch, custom_llm_provider, litellm_params=None): return _vertex_jsonl(rows) + import litellm.cost_calculator as cc + monkeypatch.setattr(bu, "_fetch_batch_output_file_content", fake_fetch) - monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 3.3) + monkeypatch.setattr(cc, "batch_cost_calculator", lambda **kw: (2.0, 1.3)) result = await bu._handle_completed_batch(_batch("of"), custom_llm_provider="openai") 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..d3e668b8987 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 @@ -375,6 +375,7 @@ def _in_memory_managed_files(): table.upsert = AsyncMock(side_effect=_upsert) prisma = MagicMock() prisma.db.litellm_managedobjecttable = table + prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=None) cache = MagicMock() cache.async_set_cache = AsyncMock() @@ -390,7 +391,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,9 +408,70 @@ 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"] +@pytest.mark.asyncio +async def test_store_unified_object_id_resolves_org_through_the_cached_team(): + """Most keys belong to an org only through their team, so the auth object carries no + org_id. The create reads the team that auth already cached, so org spend is snapshotted + at submission time without a database query in the request path.""" + from litellm.models.team import LiteLLM_TeamTableCachedObj + from litellm.proxy.proxy_server import user_api_key_cache + + instance, store = _in_memory_managed_files() + creator = UserAPIKeyAuth(user_id="alice", team_id="team-cached", api_key="hash-alice") + await user_api_key_cache.async_set_cache( + key="team_id:team-cached", + value=LiteLLM_TeamTableCachedObj(team_id="team-cached", organization_id="org-via-team"), + model_type=LiteLLM_TeamTableCachedObj, + ) + try: + await instance.store_unified_object_id( + unified_object_id="unified-b", + file_object=_build_batch_response(batch_id="b", status="validating"), + litellm_parent_otel_span=None, + model_object_id="b", + file_purpose="batch", + user_api_key_dict=creator, + persist_attribution=True, + ) + finally: + user_api_key_cache.delete_cache(key="team_id:team-cached") + + assert store["unified-b"]["org_id"] == "org-via-team" + instance.prisma_client.db.litellm_teamtable.find_unique.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_store_unified_object_id_resolves_org_from_the_db_when_the_team_is_not_cached(): + """A team no request has run under yet is absent from the auth cache; its organization + still comes back from the table so the org is billed rather than dropped.""" + from litellm.models.team import LiteLLM_TeamTable + from litellm.proxy.proxy_server import user_api_key_cache + + instance, store = _in_memory_managed_files() + instance.prisma_client.db.litellm_teamtable.find_unique = AsyncMock( + return_value=LiteLLM_TeamTable(team_id="team-uncached", organization_id="org-via-db") + ) + creator = UserAPIKeyAuth(user_id="alice", team_id="team-uncached", api_key="hash-alice") + try: + await instance.store_unified_object_id( + unified_object_id="unified-b", + file_object=_build_batch_response(batch_id="b", status="validating"), + litellm_parent_otel_span=None, + model_object_id="b", + file_purpose="batch", + user_api_key_dict=creator, + persist_attribution=True, + ) + finally: + user_api_key_cache.delete_cache(key="team_id:team-uncached") + + assert store["unified-b"]["org_id"] == "org-via-db" + + @pytest.mark.asyncio async def test_store_unified_object_id_omits_key_and_tags_without_persist_attribution(): """Regression (spend redirect): a caller that is not the batch create (a poll, or the @@ -471,6 +533,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 diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py index 95ddc4477e1..5a79560b972 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py @@ -132,6 +132,68 @@ def test_legacy_policy_keeps_trace_id_fallback(): assert len(str(generated)) == 36 +def test_batch_lifecycle_rows_derive_the_same_session_from_the_batch_id(): + """The create call's request id IS the batch id and the poller's cost row appends + _batch_cost to it, so deriving the session from the request id lands both rows in one + trace on the logs UI even though the poller builds a fresh logging context per cycle.""" + from litellm.proxy.spend_tracking.spend_tracking_utils import _get_batch_trace_session_id + + create_session: Final = _get_batch_trace_session_id(call_type="acreate_batch", request_id="batch-uid-1") + cost_session: Final = _get_batch_trace_session_id( + call_type="aretrieve_batch", request_id="batch-uid-1_batch_cost" + ) + assert create_session == cost_session == "batch-uid-1" + + +def test_non_batch_call_types_derive_no_batch_session(): + from litellm.proxy.spend_tracking.spend_tracking_utils import _get_batch_trace_session_id + + assert _get_batch_trace_session_id(call_type="acompletion", request_id="chatcmpl-1") is None + + +def test_batch_session_outranks_the_per_request_trace_id(): + """Each batch lifecycle call carries its own auto-generated trace id, so letting the + trace id win would scatter the rows across sessions again.""" + session_id: Final = _get_session_id_for_spend_log( + kwargs={"litellm_trace_id": "trace-abc"}, + metadata={"trace_id": "trace-abc"}, + standard_logging_payload=_TRACE_ONLY_STANDARD_LOGGING, + omit_when_missing=False, + batch_trace_session_id="batch-uid-1", + ) + assert session_id == "batch-uid-1" + + +def test_omit_policy_still_suppresses_batch_sessions(): + session_id: Final = _get_session_id_for_spend_log( + kwargs={}, + metadata=None, + standard_logging_payload=None, + omit_when_missing=True, + batch_trace_session_id="batch-uid-1", + ) + assert session_id is None + + +def test_get_logging_payload_groups_batch_create_and_cost_rows_in_one_session(): + def _payload(call_type: str) -> SpendLogsPayload: + return get_logging_payload( + kwargs={ + "call_type": call_type, + "model": "gpt-4o-mini", + "litellm_params": {"metadata": {"user_api_key": "test-key"}}, + }, + response_obj=litellm.ModelResponse(id="batch-uid-1", choices=[], usage=litellm.Usage()), + start_time=datetime.datetime.now(timezone.utc), + end_time=datetime.datetime.now(timezone.utc), + ) + + create_payload: Final = _payload("acreate_batch") + cost_payload: Final = _payload("aretrieve_batch") + assert cost_payload["request_id"] == "batch-uid-1_batch_cost" + assert create_payload["session_id"] == cost_payload["session_id"] == "batch-uid-1" + + @pytest.mark.parametrize( ("request_metadata", "expected"), [ diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.test.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.test.tsx index 2e9bce5048f..e0d6a1d41e0 100644 --- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.test.tsx @@ -113,6 +113,86 @@ describe("LogDetailContent", () => { expect(screen.getAllByText("$0.00200000").length).toBeGreaterThanOrEqual(1); }); + it("shows reasoning tokens in Metrics when the usage breakout carries them", () => { + render( + , + ); + + expect(screen.getByText("Reasoning Tokens")).toBeInTheDocument(); + expect(screen.getByText("224")).toBeInTheDocument(); + }); + + it("hides the reasoning metric when the breakout is absent or zero", () => { + render( + , + ); + + expect(screen.queryByText("Reasoning Tokens")).not.toBeInTheDocument(); + }); + + describe("Batch Results section", () => { + const batchCostEntry = (metadata: Record) => + createLogEntry({ + request_id: "batch_abc123_batch_cost", + call_type: "aretrieve_batch", + metadata: { status: "success", ...metadata }, + }); + + it("renders batch id, per-request outcome counts, and batch models for a batch cost row", () => { + render( + , + ); + + expect(screen.getByText("Batch Results")).toBeInTheDocument(); + expect(screen.getByText("batch_abc123")).toBeInTheDocument(); + expect(screen.getByText("Successful Requests")).toBeInTheDocument(); + expect(screen.getByText("2")).toBeInTheDocument(); + expect(screen.getByText("Failed Requests")).toBeInTheDocument(); + expect(screen.getByText("1")).toBeInTheDocument(); + expect(screen.getByText("gemini-2.5-flash")).toBeInTheDocument(); + }); + + it("still renders the batch id when a legacy row carries no counts", () => { + render(); + + expect(screen.getByText("Batch Results")).toBeInTheDocument(); + expect(screen.getByText("batch_abc123")).toBeInTheDocument(); + expect(screen.queryByText("Successful Requests")).not.toBeInTheDocument(); + }); + + it("never renders for a non-batch call type", () => { + render( + , + ); + + expect(screen.queryByText("Batch Results")).not.toBeInTheDocument(); + }); + }); + it("should show Input Tokens and Output Tokens for anthropic_messages when uncached text_tokens exist", () => { render( + {/* Batch Results */} + {isBatchCallType(logEntry.call_type) && } + {/* Routing */} @@ -374,6 +384,53 @@ function MetricLabel({ label, tooltip, docsUrl }: { label: string; tooltip: stri ); } +/** + * Aggregate per-request outcomes for a batch cost row: batch id, success/failure counts + * from the parsed output and error files, and the models the batch actually ran on. + */ +function BatchResultsSection({ logEntry, metadata }: { logEntry: LogEntry; metadata: Record }) { + const counts = getBatchRequestCounts(metadata); + const batchId = getBatchIdFromRequestId(logEntry.request_id); + const batchModels = getBatchModels(metadata); + if (!counts && !batchId && !batchModels) return null; + + return ( +
+ + + Batch Results + + + + {batchId && ( + + + + )} + {counts && ( + <> + + {formatNumberWithCommas(counts.successful)} + + + {counts.failed > 0 ? ( + + {formatNumberWithCommas(counts.failed)} + + ) : ( + formatNumberWithCommas(counts.failed) + )} + + + )} + {batchModels && {batchModels.join(", ")}} + + + +
+ ); +} + function MetricsSection({ logEntry, metadata }: { logEntry: LogEntry; metadata: Record }) { const completionStartTime = logEntry.completionStartTime; const ttftMs = @@ -391,6 +448,7 @@ function MetricsSection({ logEntry, metadata }: { logEntry: LogEntry; metadata: const uncachedInputTokens = getUncachedInputTextTokens(metadata); const showAnthropicMessagesInputOutput = logEntry.call_type === "anthropic_messages" && uncachedInputTokens !== undefined; + const reasoningTokens = getReasoningTokens(metadata); return (
@@ -416,6 +474,9 @@ function MetricsSection({ logEntry, metadata }: { logEntry: LogEntry; metadata: /> )} + {reasoningTokens !== undefined && reasoningTokens > 0 && ( + {formatNumberWithCommas(reasoningTokens)} + )} ${formatNumberWithCommas(logEntry.spend || 0, 8)} {logEntry.request_duration_ms != null ? (logEntry.request_duration_ms / 1000).toFixed(3) : "-"} s diff --git a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.test.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.test.tsx index 3c9e6543c1c..9f0e659cb1f 100644 --- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.test.tsx @@ -134,6 +134,72 @@ describe("Type column", () => { expect(screen.getByText("MCP")).toBeInTheDocument(); }); + + it("marks a batch cost row with the Batch badge instead of LLM", () => { + renderRows([logEntry({ request_id: "batch_1_batch_cost", call_type: "aretrieve_batch" })]); + + expect(screen.getByText("Batch")).toBeInTheDocument(); + expect(screen.queryByText("LLM")).not.toBeInTheDocument(); + }); + + it("keeps the Batch label on the grouped create-plus-cost session instead of a row count", () => { + const groupedCostRow: Partial = { + request_id: "batch_1_batch_cost", + call_type: "aretrieve_batch", + session_id: "batch_1", + session_total_count: 2, + }; + renderRows([logEntry(groupedCostRow)]); + + expect(screen.getByText("Batch")).toBeInTheDocument(); + expect(screen.queryByText("2")).not.toBeInTheDocument(); + }); +}); + +describe("batch rows", () => { + const batchRow = (overrides: Partial): LogEntry => + logEntry({ + request_id: "batch_abc123_batch_cost", + call_type: "aretrieve_batch", + ...overrides, + }); + + it("rolls partial failures into the status badge instead of reporting blanket Success", async () => { + const user = userEvent.setup(); + renderRows([batchRow({ metadata: { batch_successful_requests: 2, batch_failed_requests: 1 } })]); + + expect(screen.queryByText("Success")).not.toBeInTheDocument(); + await user.hover(screen.getByText("2/3 succeeded")); + expect(await screen.findByText("1 of 3 batch requests failed")).toBeInTheDocument(); + }); + + it("keeps the Success badge when every batch request succeeded", () => { + renderRows([batchRow({ metadata: { batch_successful_requests: 3, batch_failed_requests: 0 } })]); + + expect(screen.getByText("Success")).toBeInTheDocument(); + }); + + it("keeps the Failure badge when the batch row itself failed, whatever the counts say", () => { + renderRows([batchRow({ metadata: { status: "failure", batch_successful_requests: 2, batch_failed_requests: 1 } })]); + + expect(screen.getByText("Failure")).toBeInTheDocument(); + expect(screen.queryByText("2/3 succeeded")).not.toBeInTheDocument(); + }); + + it("shows the provider batch id, not the synthetic _batch_cost request id", () => { + renderRows([batchRow({ metadata: { batch_successful_requests: 1, batch_failed_requests: 0 } })]); + + expect(screen.getByText("batch_abc123")).toBeInTheDocument(); + expect(screen.queryByText("batch_abc123_batch_cost")).not.toBeInTheDocument(); + expect(screen.getByText("batch cost")).toBeInTheDocument(); + }); + + it("leaves ordinary request ids untouched", () => { + renderRows([logEntry({ request_id: "chatcmpl-42" })]); + + expect(screen.getByText("chatcmpl-42")).toBeInTheDocument(); + expect(screen.queryByText("batch cost")).not.toBeInTheDocument(); + }); }); describe("Model column", () => { diff --git a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.tsx index 9d4dc4f7898..1ec1087a1a4 100644 --- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.tsx @@ -7,9 +7,10 @@ import { CellTooltip, DateCell, IdCell, MoneyCell, StatusBadge } from "@/compone import { getSpendString } from "@/utils/dataUtils"; import { getProviderLogoAndName } from "../provider_info_helpers"; +import { getBatchIdFromRequestId, getBatchRequestCounts, isBatchCallType } from "./batchLogUtils"; import type { LogEntry } from "./columns"; import { AGENT_CALL_TYPES, MCP_CALL_TYPES } from "./constants"; -import { AgentBadge, AgentIcon, LlmBadge, McpBadge, SparkleIcon, WrenchIcon } from "./TypeBadges"; +import { AgentBadge, AgentIcon, BatchBadge, LlmBadge, McpBadge, SparkleIcon, WrenchIcon } from "./TypeBadges"; export interface RequestLogsTableColumnsDeps { onKeyHashClick: (keyHash: string) => void; @@ -63,6 +64,9 @@ export const getRequestLogsTableColumns = ({ const sessionAgentCount = log.session_agent_count ?? (isAgent ? sessionCount : 0); const sessionMcpCount = log.mcp_tool_call_count ?? (isMcp ? sessionCount : 0); + if (isBatchCallType(log.call_type)) { + return ; + } if (sessionCount <= 1) { if (isMcp) return ; if (isAgent) return ; @@ -106,6 +110,17 @@ export const getRequestLogsTableColumns = ({ cell: ({ row }) => { const status = readMetaString(row.original.metadata, "status") ?? "Success"; const isSuccess = status.toLowerCase() !== "failure"; + const batchCounts = isSuccess ? getBatchRequestCounts(row.original.metadata) : undefined; + if (batchCounts && batchCounts.failed > 0) { + const total = batchCounts.successful + batchCounts.failed; + return ( + + ); + } return ; }, }, @@ -122,7 +137,19 @@ export const getRequestLogsTableColumns = ({ accessorKey: "request_id", header: "Request ID", enableSorting: false, - cell: ({ row }) => , + cell: ({ row }) => { + const log = row.original; + const batchId = isBatchCallType(log.call_type) ? getBatchIdFromRequestId(log.request_id) : undefined; + if (batchId) { + return ( +
+ + batch cost +
+ ); + } + return ; + }, }, { id: "spend", diff --git a/ui/litellm-dashboard/src/components/view_logs/TypeBadges.test.tsx b/ui/litellm-dashboard/src/components/view_logs/TypeBadges.test.tsx index e3467310265..9a3b53685ca 100644 --- a/ui/litellm-dashboard/src/components/view_logs/TypeBadges.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/TypeBadges.test.tsx @@ -1,6 +1,6 @@ import { render, screen } from "@testing-library/react"; import { describe, expect, it } from "vitest"; -import { LlmBadge, McpBadge, AgentBadge } from "./TypeBadges"; +import { LlmBadge, McpBadge, AgentBadge, BatchBadge } from "./TypeBadges"; describe("TypeBadges", () => { describe("LlmBadge", () => { @@ -43,4 +43,11 @@ describe("TypeBadges", () => { expect(screen.getByText("12")).toBeInTheDocument(); }); }); + + describe("BatchBadge", () => { + it("should render 'Batch'", () => { + render(); + expect(screen.getByText("Batch")).toBeInTheDocument(); + }); + }); }); diff --git a/ui/litellm-dashboard/src/components/view_logs/TypeBadges.tsx b/ui/litellm-dashboard/src/components/view_logs/TypeBadges.tsx index 1ba4365f66e..db64bfbbe73 100644 --- a/ui/litellm-dashboard/src/components/view_logs/TypeBadges.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/TypeBadges.tsx @@ -57,6 +57,25 @@ export const AgentIcon = ({ size = 12 }: { size?: number }) => ( ); +/** Stacked-layers icon for Batch API call types (Lucide Layers-style). */ +export const LayersIcon = ({ size = 12 }: { size?: number }) => ( + + + + + +); + export const LlmBadge = ({ count }: { count?: number }) => ( @@ -77,3 +96,10 @@ export const AgentBadge = ({ count }: { count?: number }) => ( {count != null ? count : "Agent"} ); + +export const BatchBadge = () => ( + + + Batch + +); diff --git a/ui/litellm-dashboard/src/components/view_logs/batchLogUtils.test.ts b/ui/litellm-dashboard/src/components/view_logs/batchLogUtils.test.ts new file mode 100644 index 00000000000..4da1e389646 --- /dev/null +++ b/ui/litellm-dashboard/src/components/view_logs/batchLogUtils.test.ts @@ -0,0 +1,90 @@ +import { describe, expect, it } from "vitest"; + +import { + getBatchIdFromRequestId, + getBatchModels, + getBatchRequestCounts, + getReasoningTokens, + isBatchCallType, +} from "./batchLogUtils"; + +/** Metadata shape the batch cost poller writes on an aretrieve_batch spend row. */ +const batchCostMetadata = { + batch_models: ["gemini-2.5-flash"], + batch_successful_requests: 2, + batch_failed_requests: 1, + usage_object: { + total_tokens: 270, + prompt_tokens: 14, + completion_tokens: 256, + completion_tokens_details: { text_tokens: 32, reasoning_tokens: 224 }, + }, +}; + +describe("isBatchCallType", () => { + it("recognizes the poller's aretrieve_batch and the create call types", () => { + for (const callType of ["aretrieve_batch", "retrieve_batch", "acreate_batch", "create_batch"]) { + expect(isBatchCallType(callType)).toBe(true); + } + expect(isBatchCallType("acompletion")).toBe(false); + }); +}); + +describe("getBatchRequestCounts", () => { + it("reads both counts off a batch cost row", () => { + expect(getBatchRequestCounts(batchCostMetadata)).toEqual({ successful: 2, failed: 1 }); + }); + + it("returns undefined for a non-batch row and for null counts, so no rollup renders", () => { + expect(getBatchRequestCounts({ status: "success" })).toBeUndefined(); + expect(getBatchRequestCounts({ batch_successful_requests: null, batch_failed_requests: null })).toBeUndefined(); + expect(getBatchRequestCounts(undefined)).toBeUndefined(); + }); + + it("treats a lone present count as the other being 0, for rows logged mid-rollout", () => { + expect(getBatchRequestCounts({ batch_successful_requests: 3 })).toEqual({ successful: 3, failed: 0 }); + }); +}); + +describe("getBatchIdFromRequestId", () => { + it("strips the poller's synthetic _batch_cost suffix down to the provider batch id", () => { + expect(getBatchIdFromRequestId("batch_abc123_batch_cost")).toBe("batch_abc123"); + }); + + it("returns undefined for ordinary request ids and a bare suffix", () => { + expect(getBatchIdFromRequestId("chatcmpl-123")).toBeUndefined(); + expect(getBatchIdFromRequestId("_batch_cost")).toBeUndefined(); + }); +}); + +describe("getBatchModels", () => { + it("returns the model list from metadata.batch_models", () => { + expect(getBatchModels(batchCostMetadata)).toEqual(["gemini-2.5-flash"]); + }); + + it("returns undefined when absent, null, or empty", () => { + expect(getBatchModels({})).toBeUndefined(); + expect(getBatchModels({ batch_models: null })).toBeUndefined(); + expect(getBatchModels({ batch_models: [] })).toBeUndefined(); + }); +}); + +describe("getReasoningTokens", () => { + it("reads reasoning tokens from usage_object on a batch cost row", () => { + expect(getReasoningTokens(batchCostMetadata)).toBe(224); + }); + + it("prefers additional_usage_values, which per-request rows carry", () => { + const metadata = { + additional_usage_values: { completion_tokens_details: { reasoning_tokens: 40 } }, + usage_object: { completion_tokens_details: { reasoning_tokens: 999 } }, + }; + expect(getReasoningTokens(metadata)).toBe(40); + }); + + it("returns undefined when the breakout is null or missing", () => { + expect(getReasoningTokens({ usage_object: { completion_tokens_details: null } })).toBeUndefined(); + expect(getReasoningTokens({})).toBeUndefined(); + expect(getReasoningTokens(undefined)).toBeUndefined(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/view_logs/batchLogUtils.ts b/ui/litellm-dashboard/src/components/view_logs/batchLogUtils.ts new file mode 100644 index 00000000000..ce7793065d2 --- /dev/null +++ b/ui/litellm-dashboard/src/components/view_logs/batchLogUtils.ts @@ -0,0 +1,67 @@ +/** + * Helpers for reading batch-specific fields off a spend log row. + * + * The proxy's batch cost poller (CheckBatchCost) writes one spend log per completed batch + * with request_id "_batch_cost" and call_type "aretrieve_batch", carrying + * batch_models / batch_successful_requests / batch_failed_requests in metadata + * (see litellm/proxy/spend_tracking/spend_tracking_utils.py). + */ + +import { BATCH_CALL_TYPES } from "./constants"; + +export const BATCH_COST_REQUEST_ID_SUFFIX = "_batch_cost"; + +export interface BatchRequestCounts { + successful: number; + failed: number; +} + +export const isBatchCallType = (callType: string): boolean => BATCH_CALL_TYPES.includes(callType); + +const readMetaNumber = (metadata: Record | undefined, key: string): number | undefined => { + const value = metadata?.[key]; + return typeof value === "number" && Number.isFinite(value) ? value : undefined; +}; + +/** + * Per-request outcome counts of a batch cost row. Undefined when the row carries neither + * count (a non-batch row, or a batch logged before counts were tracked). + */ +export const getBatchRequestCounts = ( + metadata: Record | undefined, +): BatchRequestCounts | undefined => { + const successful = readMetaNumber(metadata, "batch_successful_requests"); + const failed = readMetaNumber(metadata, "batch_failed_requests"); + if (successful === undefined && failed === undefined) return undefined; + return { successful: successful ?? 0, failed: failed ?? 0 }; +}; + +/** The provider batch id behind a poller-written "_batch_cost" spend row. */ +export const getBatchIdFromRequestId = (requestId: string): string | undefined => + requestId.endsWith(BATCH_COST_REQUEST_ID_SUFFIX) && requestId.length > BATCH_COST_REQUEST_ID_SUFFIX.length + ? requestId.slice(0, -BATCH_COST_REQUEST_ID_SUFFIX.length) + : undefined; + +/** The models the batch's requests actually ran on, from metadata.batch_models. */ +export const getBatchModels = (metadata: Record | undefined): string[] | undefined => { + const models = metadata?.["batch_models"]; + if (!Array.isArray(models)) return undefined; + const names = models.filter((model): model is string => typeof model === "string" && model !== ""); + return names.length > 0 ? names : undefined; +}; + +/** + * Reasoning tokens aggregated across the row's completion usage. Read from the same two + * metadata containers the drawer already uses for prompt-token details: per-request rows + * carry additional_usage_values, batch cost rows carry usage_object. + */ +export const getReasoningTokens = (metadata: Record | undefined): number | undefined => { + const readDetails = (container: unknown): number | undefined => { + if (typeof container !== "object" || container === null) return undefined; + const details = (container as Record)["completion_tokens_details"]; + if (typeof details !== "object" || details === null) return undefined; + const reasoning = (details as Record)["reasoning_tokens"]; + return typeof reasoning === "number" && Number.isFinite(reasoning) ? reasoning : undefined; + }; + return readDetails(metadata?.["additional_usage_values"]) ?? readDetails(metadata?.["usage_object"]); +}; diff --git a/ui/litellm-dashboard/src/components/view_logs/constants.ts b/ui/litellm-dashboard/src/components/view_logs/constants.ts index 5b0b1d0fee3..ae44d53e217 100644 --- a/ui/litellm-dashboard/src/components/view_logs/constants.ts +++ b/ui/litellm-dashboard/src/components/view_logs/constants.ts @@ -18,6 +18,9 @@ export const MCP_CALL_TYPES = ["call_mcp_tool", "list_mcp_tools"]; /** Call types that represent agent/A2A requests (e.g. asend_message). */ export const AGENT_CALL_TYPES = ["asend_message"]; +/** Call types that represent Batch API operations (creation and retrieval, sync and async). */ +export const BATCH_CALL_TYPES = ["acreate_batch", "create_batch", "aretrieve_batch", "retrieve_batch"]; + export const QUICK_SELECT_OPTIONS: { label: string; value: number; unit: string }[] = [ { label: "Last Minute", value: 1, unit: "minutes" }, { label: "Last 15 Minutes", value: 15, unit: "minutes" },