mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-09 22:31:41 +00:00
Merge 95d721ffe6 into b1e2f5bc0b
This commit is contained in:
mubashir1osmani
GitHub
commit
c1f7e86e30
24 changed files with 965 additions and 52 deletions
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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():
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"),
|
||||
[
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
<LogDetailContent
|
||||
logEntry={createLogEntry({
|
||||
metadata: {
|
||||
status: "success",
|
||||
usage_object: { completion_tokens_details: { text_tokens: 32, reasoning_tokens: 224 } },
|
||||
},
|
||||
})}
|
||||
/>,
|
||||
);
|
||||
|
||||
expect(screen.getByText("Reasoning Tokens")).toBeInTheDocument();
|
||||
expect(screen.getByText("224")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("hides the reasoning metric when the breakout is absent or zero", () => {
|
||||
render(
|
||||
<LogDetailContent
|
||||
logEntry={createLogEntry({
|
||||
metadata: {
|
||||
status: "success",
|
||||
usage_object: { completion_tokens_details: { text_tokens: 32, reasoning_tokens: 0 } },
|
||||
},
|
||||
})}
|
||||
/>,
|
||||
);
|
||||
|
||||
expect(screen.queryByText("Reasoning Tokens")).not.toBeInTheDocument();
|
||||
});
|
||||
|
||||
describe("Batch Results section", () => {
|
||||
const batchCostEntry = (metadata: Record<string, unknown>) =>
|
||||
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(
|
||||
<LogDetailContent
|
||||
logEntry={batchCostEntry({
|
||||
batch_models: ["gemini-2.5-flash"],
|
||||
batch_successful_requests: 2,
|
||||
batch_failed_requests: 1,
|
||||
})}
|
||||
/>,
|
||||
);
|
||||
|
||||
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(<LogDetailContent logEntry={batchCostEntry({})} />);
|
||||
|
||||
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(
|
||||
<LogDetailContent
|
||||
logEntry={createLogEntry({
|
||||
metadata: { status: "success", batch_successful_requests: 2, batch_failed_requests: 1 },
|
||||
})}
|
||||
/>,
|
||||
);
|
||||
|
||||
expect(screen.queryByText("Batch Results")).not.toBeInTheDocument();
|
||||
});
|
||||
});
|
||||
|
||||
it("should show Input Tokens and Output Tokens for anthropic_messages when uncached text_tokens exist", () => {
|
||||
render(
|
||||
<LogDetailContent
|
||||
|
|
|
|||
|
|
@ -13,6 +13,13 @@ import { formatNumberWithCommas } from "@/utils/dataUtils";
|
|||
import { PROMPT_CACHE_CREATION_TOOLTIP, PROMPT_CACHE_READ_TOOLTIP } from "@/utils/promptCacheUsage";
|
||||
import GuardrailViewer from "../GuardrailViewer/GuardrailViewer";
|
||||
import EvalViewer from "../EvalViewer/EvalViewer";
|
||||
import {
|
||||
getBatchIdFromRequestId,
|
||||
getBatchModels,
|
||||
getBatchRequestCounts,
|
||||
getReasoningTokens,
|
||||
isBatchCallType,
|
||||
} from "../batchLogUtils";
|
||||
import { CostBreakdownViewer } from "../CostBreakdownViewer";
|
||||
import { ConfigInfoMessage } from "../ConfigInfoMessage";
|
||||
import { VectorStoreViewer } from "../VectorStoreViewer";
|
||||
|
|
@ -150,6 +157,9 @@ export function LogDetailContent({ logEntry, isLoadingDetails = false, accessTok
|
|||
</Card>
|
||||
</div>
|
||||
|
||||
{/* Batch Results */}
|
||||
{isBatchCallType(logEntry.call_type) && <BatchResultsSection logEntry={logEntry} metadata={metadata} />}
|
||||
|
||||
{/* Routing */}
|
||||
<RoutingDecisionCard decision={metadata?.routing_decision as RoutingDecision | undefined} />
|
||||
|
||||
|
|
@ -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<string, unknown> }) {
|
||||
const counts = getBatchRequestCounts(metadata);
|
||||
const batchId = getBatchIdFromRequestId(logEntry.request_id);
|
||||
const batchModels = getBatchModels(metadata);
|
||||
if (!counts && !batchId && !batchModels) return null;
|
||||
|
||||
return (
|
||||
<div className="bg-card rounded-lg shadow-sm w-full max-w-full overflow-hidden mb-6">
|
||||
<Card size="sm" style={{ marginBottom: 0 }}>
|
||||
<CardHeader>
|
||||
<CardTitle>Batch Results</CardTitle>
|
||||
</CardHeader>
|
||||
<CardContent>
|
||||
<DescriptionList>
|
||||
{batchId && (
|
||||
<DescriptionItem label="Batch ID">
|
||||
<TruncatedValue value={batchId} />
|
||||
</DescriptionItem>
|
||||
)}
|
||||
{counts && (
|
||||
<>
|
||||
<DescriptionItem label="Successful Requests">
|
||||
{formatNumberWithCommas(counts.successful)}
|
||||
</DescriptionItem>
|
||||
<DescriptionItem label="Failed Requests">
|
||||
{counts.failed > 0 ? (
|
||||
<Badge variant="secondary" className="bg-destructive/15 text-destructive">
|
||||
{formatNumberWithCommas(counts.failed)}
|
||||
</Badge>
|
||||
) : (
|
||||
formatNumberWithCommas(counts.failed)
|
||||
)}
|
||||
</DescriptionItem>
|
||||
</>
|
||||
)}
|
||||
{batchModels && <DescriptionItem label="Models">{batchModels.join(", ")}</DescriptionItem>}
|
||||
</DescriptionList>
|
||||
</CardContent>
|
||||
</Card>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
function MetricsSection({ logEntry, metadata }: { logEntry: LogEntry; metadata: Record<string, any> }) {
|
||||
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 (
|
||||
<div className="bg-card rounded-lg shadow-sm w-full max-w-full overflow-hidden mb-6">
|
||||
|
|
@ -416,6 +474,9 @@ function MetricsSection({ logEntry, metadata }: { logEntry: LogEntry; metadata:
|
|||
/>
|
||||
</DescriptionItem>
|
||||
)}
|
||||
{reasoningTokens !== undefined && reasoningTokens > 0 && (
|
||||
<DescriptionItem label="Reasoning Tokens">{formatNumberWithCommas(reasoningTokens)}</DescriptionItem>
|
||||
)}
|
||||
<DescriptionItem label="Cost">${formatNumberWithCommas(logEntry.spend || 0, 8)}</DescriptionItem>
|
||||
<DescriptionItem label="Duration">
|
||||
{logEntry.request_duration_ms != null ? (logEntry.request_duration_ms / 1000).toFixed(3) : "-"} s
|
||||
|
|
|
|||
|
|
@ -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<LogEntry> = {
|
||||
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 =>
|
||||
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", () => {
|
||||
|
|
|
|||
|
|
@ -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 <BatchBadge />;
|
||||
}
|
||||
if (sessionCount <= 1) {
|
||||
if (isMcp) return <McpBadge />;
|
||||
if (isAgent) return <AgentBadge />;
|
||||
|
|
@ -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 (
|
||||
<StatusBadge
|
||||
tone="warning"
|
||||
label={`${batchCounts.successful}/${total} succeeded`}
|
||||
tooltip={`${batchCounts.failed} of ${total} batch requests failed`}
|
||||
/>
|
||||
);
|
||||
}
|
||||
return <StatusBadge tone={isSuccess ? "success" : "error"} label={isSuccess ? "Success" : "Failure"} />;
|
||||
},
|
||||
},
|
||||
|
|
@ -122,7 +137,19 @@ export const getRequestLogsTableColumns = ({
|
|||
accessorKey: "request_id",
|
||||
header: "Request ID",
|
||||
enableSorting: false,
|
||||
cell: ({ row }) => <IdCell value={row.original.request_id} variant="plain" />,
|
||||
cell: ({ row }) => {
|
||||
const log = row.original;
|
||||
const batchId = isBatchCallType(log.call_type) ? getBatchIdFromRequestId(log.request_id) : undefined;
|
||||
if (batchId) {
|
||||
return (
|
||||
<div className="flex flex-col">
|
||||
<IdCell value={batchId} variant="plain" copyable tooltip={`Batch ${batchId} (row: ${log.request_id})`} />
|
||||
<span className="text-[10px] text-muted-foreground">batch cost</span>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
return <IdCell value={log.request_id} variant="plain" />;
|
||||
},
|
||||
},
|
||||
{
|
||||
id: "spend",
|
||||
|
|
|
|||
|
|
@ -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(<BatchBadge />);
|
||||
expect(screen.getByText("Batch")).toBeInTheDocument();
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -57,6 +57,25 @@ export const AgentIcon = ({ size = 12 }: { size?: number }) => (
|
|||
</svg>
|
||||
);
|
||||
|
||||
/** Stacked-layers icon for Batch API call types (Lucide Layers-style). */
|
||||
export const LayersIcon = ({ size = 12 }: { size?: number }) => (
|
||||
<svg
|
||||
width={size}
|
||||
height={size}
|
||||
viewBox="0 0 24 24"
|
||||
fill="none"
|
||||
stroke="currentColor"
|
||||
strokeWidth="2"
|
||||
strokeLinecap="round"
|
||||
strokeLinejoin="round"
|
||||
className="shrink-0"
|
||||
>
|
||||
<path d="M12 2 2 7l10 5 10-5-10-5z" />
|
||||
<path d="m2 17 10 5 10-5" />
|
||||
<path d="m2 12 10 5 10-5" />
|
||||
</svg>
|
||||
);
|
||||
|
||||
export const LlmBadge = ({ count }: { count?: number }) => (
|
||||
<span className="inline-flex items-center gap-1 px-2 py-0.5 bg-info/10 text-info border border-info/20 rounded-full text-[11px] font-medium whitespace-nowrap">
|
||||
<SparkleIcon />
|
||||
|
|
@ -77,3 +96,10 @@ export const AgentBadge = ({ count }: { count?: number }) => (
|
|||
{count != null ? count : "Agent"}
|
||||
</span>
|
||||
);
|
||||
|
||||
export const BatchBadge = () => (
|
||||
<span className="inline-flex items-center gap-1 px-2 py-0.5 bg-teal-50 text-teal-700 border border-teal-200 rounded-full text-[11px] font-medium whitespace-nowrap dark:bg-teal-950 dark:text-teal-300 dark:border-teal-800">
|
||||
<LayersIcon />
|
||||
Batch
|
||||
</span>
|
||||
);
|
||||
|
|
|
|||
|
|
@ -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();
|
||||
});
|
||||
});
|
||||
|
|
@ -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_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<string, unknown> | 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<string, unknown> | 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_id>_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<string, unknown> | 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<string, unknown> | undefined): number | undefined => {
|
||||
const readDetails = (container: unknown): number | undefined => {
|
||||
if (typeof container !== "object" || container === null) return undefined;
|
||||
const details = (container as Record<string, unknown>)["completion_tokens_details"];
|
||||
if (typeof details !== "object" || details === null) return undefined;
|
||||
const reasoning = (details as Record<string, unknown>)["reasoning_tokens"];
|
||||
return typeof reasoning === "number" && Number.isFinite(reasoning) ? reasoning : undefined;
|
||||
};
|
||||
return readDetails(metadata?.["additional_usage_values"]) ?? readDetails(metadata?.["usage_object"]);
|
||||
};
|
||||
|
|
@ -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" },
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue