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 961795838a8..734092c4feb 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -2,7 +2,8 @@ Polls LiteLLM_ManagedObjectTable to check if the batch job is complete, and if the cost has been tracked. """ -from collections.abc import Sequence +import asyncio +from collections.abc import Mapping, Sequence from dataclasses import replace as dataclasses_replace from datetime import datetime, timedelta, timezone from types import MappingProxyType @@ -13,15 +14,22 @@ from litellm_enterprise.proxy.common_utils.managed_object_attribution import ( ManagedObjectRow, ) +import litellm from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid from litellm.constants import ( + BATCH_LINE_ITEMS_MAX_PART_BYTES, MANAGED_OBJECT_STALENESS_CUTOFF_DAYS, MAX_OBJECTS_PER_POLL_CYCLE, ) from litellm.repositories.table_repositories import ManagedObjectRepository +from litellm.repositories.team_repository import TeamRepository +from litellm.repositories.verification_token_repository import VerificationTokenRepository if TYPE_CHECKING: + from prisma import models as prisma_models + + from litellm.batches.batch_line_items import BatchFileProvider from litellm.integrations.prometheus import PrometheusLogger from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.repositories.prisma_protocols import TableActions @@ -67,6 +75,19 @@ def _managed_object_table(prisma_client: "PrismaClient") -> "TableActions[Manage return ManagedObjectRepository(prisma_client).table +def _token_table(prisma_client: "PrismaClient") -> "TableActions[prisma_models.LiteLLM_VerificationToken]": + return VerificationTokenRepository(prisma_client).table + + +def _team_table(prisma_client: "PrismaClient") -> "TableActions[prisma_models.LiteLLM_TeamTable]": + return TeamRepository(prisma_client).table + + +def _row_metadata(row: object) -> dict[str, object]: + metadata: Final = getattr(row, "metadata", None) + return cast(dict[str, object], metadata) if isinstance(metadata, dict) else {} # cast-ok: prisma Json object column + + class CheckBatchCost: def __init__( self, @@ -87,6 +108,10 @@ class CheckBatchCost: # the guaranteed-failing primary query on every subsequent cycle. self._has_batch_processed_column: bool = True self.batch_processed_support_confirmed: bool = False + self._line_item_stores: Final[set[asyncio.Task[None]]] = set() + self._line_item_store_turn: Final = asyncio.Lock() + self._poll_idle: Final = asyncio.Event() + self._poll_idle.set() @staticmethod def _is_missing_batch_processed_column_error(err: Exception) -> bool: @@ -299,9 +324,9 @@ class CheckBatchCost: async def _finalize_unbilled_terminal_job( self, job: "ManagedObjectRow", response: "LiteLLMBatch" - ) -> None: + ) -> bool: """Persist a terminal batch that has nothing billable, converting any raw - provider file ids to managed ids, and take it out of the poll page.""" + provider file ids to managed ids, and take it out of the poll page. Returns whether it was persisted.""" try: from litellm.proxy.openai_files_endpoints.common_utils import ( _is_base64_encoded_unified_file_id, @@ -329,10 +354,12 @@ class CheckBatchCost: verbose_proxy_logger.info( f"CheckBatchCost: marked job {job.id} as {response.status} in DB" ) + return True except Exception as db_err: verbose_proxy_logger.error( f"CheckBatchCost: failed to mark job {job.id} as {response.status} in DB: {db_err}" ) + return False @staticmethod def _record_error( @@ -794,7 +821,191 @@ class CheckBatchCost: return model_name, str(llm_provider) if llm_provider else None - async def check_batch_cost(self): + async def _retire_unbilled_terminal_job( + self, job: "ManagedObjectRow", response: "LiteLLMBatch", model_id: str, batch_id: str + ) -> None: + """Persist a terminal batch that has nothing billable, then store its lines, reading its input file by the id + the provider reported rather than the managed id that was persisted.""" + provider_input_file_id: Final = response.input_file_id + if await self._finalize_unbilled_terminal_job(job, response): + response.input_file_id = provider_input_file_id + self._store_batch_line_items_in_background(job, response, model_id, batch_id, billed=False) + + def _store_batch_line_items_in_background( + self, job: "ManagedObjectRow", response: "LiteLLMBatch", model_id: str, batch_id: str, *, billed: bool + ) -> None: + """Store a retired batch's lines off the poll loop, one batch at a time and only between polls, so neither + pricing the lines nor a slow or hung bucket holds up billing or retiring the other batches. While a poll + cycle's worth of batches is still waiting, the lines of the next one are dropped with an error.""" + if not litellm.batch_line_item_callbacks: + return + if len(self._line_item_stores) >= MAX_OBJECTS_PER_POLL_CYCLE: + verbose_proxy_logger.error( + f"CheckBatchCost: {len(self._line_item_stores)} batches are still waiting to store their line items, " + f"so the line items of batch {batch_id} (job {job.id}) were not stored" + ) + return + store: Final = asyncio.create_task( + self._store_batch_line_items_in_turn(job, response, model_id, batch_id, billed=billed) + ) + self._line_item_stores.add(store) + store.add_done_callback(self._line_item_stores.discard) + + async def _store_batch_line_items_in_turn( + self, job: "ManagedObjectRow", response: "LiteLLMBatch", model_id: str, batch_id: str, *, billed: bool + ) -> None: + async with self._line_item_store_turn: + await self._poll_idle.wait() + await self._store_batch_line_items(job, response, model_id, batch_id, billed=billed) + + async def _store_batch_line_items( + self, job: "ManagedObjectRow", response: "LiteLLMBatch", model_id: str, batch_id: str, *, billed: bool + ) -> None: + """Send every line of a finished batch to batch_line_item_callbacks, one part at a time. A billed batch's + output file is read again here, so nothing waiting for its turn holds a batch's files. Failures, including an + output file that can no longer be read, are logged and never raised, so the batch stays billed and retired + whether or not its lines were stored.""" + if not litellm.batch_line_item_callbacks: + return + from litellm.batches.batch_line_items import ( + BatchFileProvider, + BatchLineItemSource, + batch_line_item_parts, + dispatch_batch_line_items, + line_item_metadata, + ) + from litellm.batches.batch_utils import fetch_batch_managed_file_content + from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider + from litellm.litellm_core_utils.litellm_logging import deployment_pricing_model_info + + try: + deployment_info: Final = self.llm_router.get_deployment(model_id=model_id) + if deployment_info is None: + verbose_proxy_logger.error( + f"CheckBatchCost: deployment {model_id} is no longer in the router, so the line items of batch " + f"{batch_id} (job {job.id}) were not stored" + ) + return + litellm_model_name: Final = deployment_info.litellm_params.model + model_name, llm_provider, _, _ = get_llm_provider( + model=litellm_model_name, + custom_llm_provider=deployment_info.litellm_params.custom_llm_provider, + ) + cost_provider: Final = cast(BatchFileProvider, llm_provider) + credentials: Final = self.llm_router.get_deployment_credentials_with_provider(model_id) or {} + file_access: Final = { + **credentials, + "_litellm_internal_model_credentials": MappingProxyType(dict(credentials)), + } + output_content: Final = ( + await fetch_batch_managed_file_content( + response.output_file_id, custom_llm_provider=cost_provider, litellm_params=file_access + ) + if billed and response.output_file_id + else b"" + ) + error_content: Final = await self._batch_file_content(response.error_file_id, cost_provider, file_access) + if not output_content and not error_content: + return + redact_bodies: Final = await self._creator_turned_off_message_logging(job) + source: Final = BatchLineItemSource( + batch_id=job.unified_object_id, + finished_at=response.completed_at + or response.failed_at + or response.expired_at + or response.cancelled_at + or response.created_at, + model=model_name, + custom_llm_provider=llm_provider, + cost_provider=cost_provider, + model_info=deployment_pricing_model_info(model_id=model_id, deployment_model=litellm_model_name), + metadata=line_item_metadata(await self._creator_attribution.batch_metadata(job, batch_id)), + ) + input_content: Final = ( + b"" + if redact_bodies + else await self._batch_file_content(response.input_file_id, cost_provider, file_access) + ) + parts: Final = batch_line_item_parts( + source, input_content, (output_content, error_content), BATCH_LINE_ITEMS_MAX_PART_BYTES + ) + handed = 0 + # pricing a part's lines is CPU-bound, so each part is built off the event loop + while (payload := await asyncio.to_thread(next, parts, None)) is not None: + await dispatch_batch_line_items(payload, redact_bodies=redact_bodies) + handed += len(payload["lines"]) + if handed: + verbose_proxy_logger.info( + f"CheckBatchCost: handed the {handed} line items of batch {batch_id} to batch_line_item_callbacks" + ) + except Exception as error: # noqa: BLE001 # storing the lines is best effort, the batch is already billed + verbose_proxy_logger.exception( + f"CheckBatchCost: the line items of batch {batch_id} (job {job.id}) were not stored: {error}" + ) + + @staticmethod + async def _batch_file_content( + file_id: str | None, custom_llm_provider: "BatchFileProvider", file_access: Mapping[str, object] + ) -> bytes: + """A batch file's raw JSONL, or nothing when the batch has no such file or the provider will not serve it.""" + from litellm.batches.batch_utils import fetch_batch_managed_file_content + + if not file_id: + return b"" + try: + return await fetch_batch_managed_file_content( + file_id, custom_llm_provider=custom_llm_provider, litellm_params=file_access + ) + except Exception as error: # noqa: BLE001 # a missing input or error file leaves those lines out, not the rest + verbose_proxy_logger.warning(f"CheckBatchCost: could not read batch file {file_id}: {error}") + return b"" + + async def _creator_turned_off_message_logging(self, job: "ManagedObjectRow") -> bool: + """Whether the creating key's or team's logging settings turn off message logging, falling back to the + proxy-wide setting. A batch whose saved credential is not a virtual key's hash, such as one created under JWT + or OAuth2 auth or with the master key, has only its team's and the proxy's settings. A key or team that can no + longer be read counts as turned off.""" + from litellm.litellm_core_utils.litellm_logging import is_valid_sha256_hash + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.litellm_pre_call_utils import get_dynamic_logging_metadata + from litellm.proxy.proxy_server import proxy_config + from litellm.secret_managers.main import str_to_bool + + api_key: Final = getattr(job, "api_key", None) + team_id: Final = getattr(job, "team_id", None) + has_key_row: Final = api_key is not None and is_valid_sha256_hash(api_key) + try: + key_row: Final = ( + await _token_table(self.prisma_client).find_unique(where={"token": api_key}) if has_key_row else None + ) + team_row: Final = ( + await _team_table(self.prisma_client).find_unique(where={"team_id": team_id}) if team_id else None + ) + except Exception as error: + verbose_proxy_logger.error( + f"CheckBatchCost: could not read the logging settings of batch job {job.id}'s creator, " + f"so its line items are stored without request and response: {error}" + ) + return True + if (has_key_row and key_row is None) or (team_id and team_row is None): + return True + settings: Final = get_dynamic_logging_metadata( + UserAPIKeyAuth(metadata=_row_metadata(key_row), team_metadata=_row_metadata(team_row), team_id=team_id), + proxy_config, + ) + callback_vars: Final = (settings.callback_vars if settings is not None else None) or {} + creator_setting: Final = str_to_bool(callback_vars.get("turn_off_message_logging")) + return creator_setting if creator_setting is not None else litellm.turn_off_message_logging is True + + async def check_batch_cost(self) -> None: + """Run one poll cycle. Batch line items wait until it has billed and retired every batch it found.""" + self._poll_idle.clear() + try: + await self._poll_batches() + finally: + self._poll_idle.set() + + async def _poll_batches(self): """ Check if the batch JOB has been tracked. - get all status="validating" and file_purpose="batch" jobs @@ -908,7 +1119,7 @@ class CheckBatchCost: f"CheckBatchCost: output file {response.output_file_id} of batch {batch_id} " f"does not exist at the provider; retiring job {job.id} unbilled" ) - await self._finalize_unbilled_terminal_job(job, response) + await self._retire_unbilled_terminal_job(job, response, model_id, batch_id) continue verbose_proxy_logger.error( f"CheckBatchCost: failed to track cost for batch {batch_id} " @@ -938,6 +1149,10 @@ class CheckBatchCost: verbose_proxy_logger.error( f"CheckBatchCost: failed to mark job {job.id} complete in DB: {db_err}" ) + if not self._has_batch_processed_column: + # nothing claimed the row, so it stays in the poll page and its lines wait for that retry + continue + self._store_batch_line_items_in_background(job, response, model_id, batch_id, billed=True) elif response.status in PROVIDER_TERMINAL_BATCH_STATUSES: from litellm.proxy.openai_files_endpoints.common_utils import ( @@ -950,7 +1165,7 @@ class CheckBatchCost: f"has not appeared yet; leaving job {job.id} for the next poll cycle" ) continue - await self._finalize_unbilled_terminal_job(job, response) + await self._retire_unbilled_terminal_job(job, response, model_id, batch_id) # Record polling run metrics (always, even if nothing was processed) if prom_logger: diff --git a/litellm/__init__.py b/litellm/__init__.py index bff6588f21c..ff901ede512 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -118,6 +118,7 @@ success_callback: List[CALLBACK_TYPES] = [] failure_callback: List[CALLBACK_TYPES] = [] service_callback: List[CALLBACK_TYPES] = [] audit_log_callbacks: List[CALLBACK_TYPES] = [] +batch_line_item_callbacks: Sequence[CALLBACK_TYPES] = () # logging_callback_manager is lazy-loaded via __getattr__ _custom_logger_compatible_callbacks_literal = Literal[ "lago", diff --git a/litellm/batches/batch_line_items.py b/litellm/batches/batch_line_items.py new file mode 100644 index 00000000000..3d175b22b90 --- /dev/null +++ b/litellm/batches/batch_line_items.py @@ -0,0 +1,250 @@ +"""Per-line records of a finished batch for the loggers listed in batch_line_item_callbacks.""" + +import asyncio +import hashlib +import re +from collections import Counter +from collections.abc import Iterable, Iterator, Mapping +from dataclasses import dataclass +from itertools import accumulate, chain, count +from typing import Final, Literal, TypeAlias + +import litellm +from litellm._logging import verbose_logger +from litellm.batches.batch_utils import batch_cost_and_usage, iter_batch_input_lines, parse_batch_output_line +from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.litellm_logging import init_custom_logger_compatible_class, is_valid_sha256_hash +from litellm.llms.vertex_ai.files.transformation import vertex_batch_row_custom_id +from litellm.types.utils import ( + ModelInfo, + StandardBatchLineItem, + StandardBatchLineItemMetadata, + StandardBatchLineItemsPayload, + StandardBatchLineItemUsage, +) + +BatchFileProvider: TypeAlias = Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "mistral"] + +_VERTEX_AI_UNKNOWN_CUSTOM_ID: Final = "unknown" # what a Vertex AI batch row that carries no custom_id is read as +_LINE: Final = re.compile(rb"[^\n]+") + + +@dataclass(frozen=True, slots=True) +class BatchLineItemSource: + batch_id: str + finished_at: int + model: str + custom_llm_provider: str | None + cost_provider: BatchFileProvider + model_info: ModelInfo | None + metadata: StandardBatchLineItemMetadata + + +def line_item_metadata(attribution: Mapping[str, object]) -> StandardBatchLineItemMetadata: + return StandardBatchLineItemMetadata( + user_api_key_hash=_creator_key_hash(attribution), + user_api_key_alias=_text(attribution, "user_api_key_alias"), + user_api_key_team_id=_text(attribution, "user_api_key_team_id"), + user_api_key_team_alias=_text(attribution, "user_api_key_team_alias"), + user_api_key_user_id=_text(attribution, "user_api_key_user_id"), + user_api_key_org_id=_text(attribution, "user_api_key_org_id"), + ) + + +def _text(attribution: Mapping[str, object], name: str) -> str | None: + value: Final = attribution.get(name) + return value if isinstance(value, str) else None + + +def _creator_key_hash(attribution: Mapping[str, object]) -> str | None: + """The creating credential as a SHA-256 hash. A virtual key is saved hashed already, while a credential saved as + is, such as an OAuth2 access token, is hashed here so the object never holds it.""" + credential: Final = _text(attribution, "user_api_key_hash") + if credential is None or is_valid_sha256_hash(credential): + return credential + return hashlib.sha256(credential.encode()).hexdigest() + + +def _line_key(entry: Mapping[str, object]) -> str | None: + """A line's custom_id. Bedrock calls it recordId, and native Vertex rows carry it in the row key or request labels.""" + key: Final = entry.get("custom_id", entry.get("recordId")) + return key if isinstance(key, str) else vertex_batch_row_custom_id(entry) + + +def batch_line_item_parts( + source: BatchLineItemSource, input_content: bytes, result_contents: Iterable[bytes], max_part_bytes: int +) -> Iterator[StandardBatchLineItemsPayload]: + """Pair every output and error line with the input line that shares its custom_id, priced one line at a time, in + parts that each close once their raw request and result lines reach max_part_bytes. Input lines stay unparsed until + their result line is reached, so only one part's records are held at a time.""" + sized: Final = _sized_line_items(source, input_content, result_contents) + for part in count(): + lines = tuple(_until_full(sized, max_part_bytes)) + if not lines: + return + yield StandardBatchLineItemsPayload( + batch_id=source.batch_id, finished_at=source.finished_at, part=part, lines=lines + ) + + +def _sized_line_items( + source: BatchLineItemSource, input_content: bytes, result_contents: Iterable[bytes] +) -> Iterator[tuple[StandardBatchLineItem, int]]: + """Each output and error line's record, with the raw bytes of that line and of the input line it pairs with.""" + request_spans: Final = _request_spans_by_custom_id(input_content) + for raw_result in chain.from_iterable(map(iter_batch_input_lines, result_contents)): + entry = parse_batch_output_line(raw_result) + if entry is None: + continue + custom_id = _line_key(entry) + span = _request_span_for(source, request_spans, custom_id) + request = None if span is None else parse_batch_output_line(input_content[span[0] : span[1]]) + yield ( + _line_item(source, custom_id, request, entry), + len(raw_result) + (0 if span is None else span[1] - span[0]), + ) + + +def _until_full( + sized: Iterator[tuple[StandardBatchLineItem, int]], max_part_bytes: int +) -> Iterator[StandardBatchLineItem]: + """Take records until their raw bytes reach max_part_bytes, leaving the rest in sized for the next part.""" + for line, part_bytes in accumulate(sized, _with_bytes_so_far): + yield line + if part_bytes >= max_part_bytes: + return + + +def _with_bytes_so_far( + so_far: tuple[StandardBatchLineItem, int], line: tuple[StandardBatchLineItem, int] +) -> tuple[StandardBatchLineItem, int]: + return line[0], so_far[1] + line[1] + + +def _request_spans_by_custom_id(input_content: bytes) -> Mapping[str | None, tuple[int, int]]: + """Where each input line starts and ends, by its custom_id. A custom_id that several input lines share, as when a + Vertex AI embeddings entry is split into one row per input, is paired with none of them.""" + keyed: Final = tuple(_keyed_spans(input_content)) + counts: Final = Counter(key for key, _ in keyed) + return {key: span for key, span in keyed if counts[key] == 1} + + +def _keyed_spans(input_content: bytes) -> Iterator[tuple[str | None, tuple[int, int]]]: + for line in _LINE.finditer(input_content): + entry = parse_batch_output_line(line[0]) if line[0].strip() else None + if entry is not None: + yield _line_key(entry), line.span() + + +def _line_item( + source: BatchLineItemSource, custom_id: str | None, request: Mapping[str, object] | None, entry: dict[str, object] +) -> StandardBatchLineItem: + priced: Final = batch_cost_and_usage( + (entry,), source.cost_provider, model_name=source.model, model_info=source.model_info + ) + succeeded: Final = priced.successful_requests > 0 + return StandardBatchLineItem( + batch_id=source.batch_id, + custom_id=custom_id, + status="success" if succeeded else "failure", + request=request, + response=entry, + usage=StandardBatchLineItemUsage( + prompt_tokens=priced.usage.prompt_tokens, + completion_tokens=priced.usage.completion_tokens, + total_tokens=priced.usage.total_tokens, + ) + if succeeded + else None, + response_cost=priced.cost, + model=source.model, + custom_llm_provider=source.custom_llm_provider, + metadata=source.metadata, + ) + + +def _request_span_for( + source: BatchLineItemSource, request_spans: Mapping[str | None, tuple[int, int]], custom_id: str | None +) -> tuple[int, int] | None: + """Where the input line that shares a result line's custom_id is. A Vertex AI line read as "unknown" pairs with + none.""" + if custom_id is None or (source.cost_provider == "vertex_ai" and custom_id == _VERTEX_AI_UNKNOWN_CUSTOM_ID): + return None + return request_spans.get(custom_id) + + +def without_request_and_response(payload: StandardBatchLineItemsPayload) -> StandardBatchLineItemsPayload: + return StandardBatchLineItemsPayload( + batch_id=payload["batch_id"], + finished_at=payload["finished_at"], + part=payload["part"], + lines=tuple( + StandardBatchLineItem( + batch_id=line["batch_id"], + custom_id=line["custom_id"], + status=line["status"], + request=None, + response=None, + usage=line["usage"], + response_cost=line["response_cost"], + model=line["model"], + custom_llm_provider=line["custom_llm_provider"], + metadata=line["metadata"], + ) + for line in payload["lines"] + ), + ) + + +async def dispatch_batch_line_items(payload: StandardBatchLineItemsPayload, redact_bodies: bool) -> None: + """Hand one part's lines to every batch_line_item_callbacks logger. A failing logger is logged and never stops the + rest.""" + await asyncio.gather( + *(_deliver(callback, payload, redact_bodies) for callback in litellm.batch_line_item_callbacks) + ) + + +async def _deliver( + callback: litellm.CALLBACK_TYPES, payload: StandardBatchLineItemsPayload, redact_bodies: bool +) -> None: + try: + logger: Final = _resolve_logger(callback) + if logger is None: + verbose_logger.error( + "batch_line_item_callbacks: %r cannot store batch line items, part %d of batch %s was not stored", + callback, + payload["part"], + payload["batch_id"], + ) + return + await logger.async_log_batch_line_items( + without_request_and_response(payload) + if redact_bodies or _logger_turned_off_message_logging(logger) + else payload + ) + except Exception as error: # noqa: BLE001 # one failing sink must not stop the other sinks or the cost poller + verbose_logger.exception( + "batch_line_item_callbacks: %r failed to store the %d lines of part %d of batch %s: %s", + callback, + len(payload["lines"]), + payload["part"], + payload["batch_id"], + error, + ) + + +def _resolve_logger(callback: litellm.CALLBACK_TYPES) -> CustomLogger | None: + match callback: + case CustomLogger(): + return callback + case "s3_v2" | "gcs_bucket": + return init_custom_logger_compatible_class(callback, internal_usage_cache=None, llm_router=None) + case _: + return None + + +def _logger_turned_off_message_logging(logger: CustomLogger) -> bool: + return ( + getattr(logger, "turn_off_message_logging", False) is True + or getattr(logger, "message_logging", True) is not True + ) diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 416391d194b..555c14ae6c9 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -1,5 +1,5 @@ import json -from collections.abc import Iterable, Iterator, Mapping +from collections.abc import Iterable, Iterator, Mapping, Sequence from dataclasses import dataclass from dataclasses import replace as dataclasses_replace from enum import Enum @@ -70,7 +70,17 @@ def batch_cost_is_final(batch: Batch) -> bool: async def calculate_batch_cost_and_usage( - file_content_dictionary: list[dict[str, object]], + file_content_dictionary: Sequence[dict[str, object]], + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "mistral"], + model_name: str | None = None, + model_info: ModelInfo | None = None, +) -> BatchCostUsageResult: + """Calculate the cost and usage of a batch, see batch_cost_and_usage.""" + return batch_cost_and_usage(file_content_dictionary, custom_llm_provider, model_name, model_info) + + +def batch_cost_and_usage( + file_content_dictionary: Sequence[dict[str, object]], # mutable-ok: the rows calculate_batch_cost_and_usage takes custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "mistral"], model_name: str | None = None, model_info: ModelInfo | None = None, @@ -436,7 +446,7 @@ def _provider_output_file_id(output_file_id: str) -> str: return extracted -async def _fetch_batch_managed_file_content( +async def fetch_batch_managed_file_content( file_id: str, custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "mistral"] = "openai", litellm_params: Mapping[str, object] | None = None, @@ -467,6 +477,9 @@ async def _fetch_batch_managed_file_content( return _file_content.content +_fetch_batch_managed_file_content: Final = fetch_batch_managed_file_content + + async def _fetch_batch_output_file_content( batch: Batch, custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "mistral"] = "openai", @@ -604,12 +617,12 @@ def _iter_batch_output_entries(file_content: bytes) -> Iterator[dict[str, object whole batch's cost accounting. """ for line in iter_batch_input_lines(file_content): - entry = _parse_batch_output_line(line) + entry = parse_batch_output_line(line) if entry is not None: yield entry -def _parse_batch_output_line(line: bytes) -> dict[str, object] | None: +def parse_batch_output_line(line: bytes) -> dict[str, object] | None: try: parsed: Final[object] = json.loads(line) except ValueError as e: diff --git a/litellm/constants.py b/litellm/constants.py index 5b2bbaf95c1..847b8564507 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1862,6 +1862,10 @@ RESET_BUDGET_JOB_NAME: Final = "reset_budget_job" RESET_BUDGET_JOB_LOCK_TTL_SECONDS: Final[int] = 900 PROXY_BATCH_POLLING_INTERVAL: Final = get_env_int("PROXY_BATCH_POLLING_INTERVAL", 3600) MAX_OBJECTS_PER_POLL_CYCLE: Final = max(1, get_env_int("MAX_OBJECTS_PER_POLL_CYCLE", 50)) +# raw request and result bytes per batch line-item object, so a large batch is built and stored one part at a time +BATCH_LINE_ITEMS_MAX_PART_BYTES: Final = get_env_int_in_range( + "BATCH_LINE_ITEMS_MAX_PART_BYTES", 16 * 1024 * 1024, 1, 1024 * 1024 * 1024 +) MANAGED_OBJECT_STALENESS_CUTOFF_DAYS: Final = max(1, get_env_int("MANAGED_OBJECT_STALENESS_CUTOFF_DAYS", 7)) PENDING_RESPONSE_STATUSES: Final = ("queued", "in_progress") TERMINAL_RESPONSE_STATUSES: Final = frozenset({"completed", "failed", "cancelled", "incomplete"}) diff --git a/litellm/integrations/batch_line_item_objects.py b/litellm/integrations/batch_line_item_objects.py new file mode 100644 index 00000000000..888f1608003 --- /dev/null +++ b/litellm/integrations/batch_line_item_objects.py @@ -0,0 +1,29 @@ +"""How the line items of one finished batch are laid out as a single object in a storage bucket.""" + +import json +from datetime import datetime, timezone +from typing import Final + +from litellm.types.utils import StandardBatchLineItemsPayload + + +def batch_line_items_ndjson(payload: StandardBatchLineItemsPayload) -> str: + return "".join(f"{_ndjson_line(line)}\n" for line in payload["lines"]) + + +def _ndjson_line(line: object) -> str: + """One line as JSON, with non-ASCII text kept readable. A line holding a lone surrogate, which a provider's + JSON escapes can carry, has no UTF-8 form, so that line alone is written with every non-ASCII character escaped.""" + text: Final = json.dumps(line, default=str, ensure_ascii=False) + try: + text.encode("utf-8") + except UnicodeEncodeError: + return json.dumps(line, default=str) + return text + + +def batch_line_items_object_name(payload: StandardBatchLineItemsPayload) -> str: + """batch_line_items///part-.jsonl, the same name on every retry so a rewrite + replaces it.""" + day: Final = datetime.fromtimestamp(payload["finished_at"], tz=timezone.utc).strftime("%Y-%m-%d") + return f"batch_line_items/{day}/{payload['batch_id'].replace('/', '_')}/part-{payload['part']:05d}.jsonl" diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py index 5ca9840185b..52577d0c647 100644 --- a/litellm/integrations/custom_logger.py +++ b/litellm/integrations/custom_logger.py @@ -22,6 +22,7 @@ from litellm.types.utils import ( ModelResponse, ModelResponseStream, StandardAuditLogPayload, + StandardBatchLineItemsPayload, StandardCallbackDynamicParams, StandardLoggingPayload, ) @@ -195,6 +196,9 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac async def async_log_audit_log_event(self, audit_log: "StandardAuditLogPayload"): """Called when an audit log is created. Override in subclasses to handle.""" + async def async_log_batch_line_items(self, payload: StandardBatchLineItemsPayload) -> None: + """Called once per part of a finished batch's lines, when this logger is listed in batch_line_item_callbacks.""" + #### PROMPT MANAGEMENT HOOKS #### async def async_get_chat_completion_prompt( diff --git a/litellm/integrations/gcs_bucket/gcs_bucket.py b/litellm/integrations/gcs_bucket/gcs_bucket.py index 24a5e940e08..5bc7247e02a 100644 --- a/litellm/integrations/gcs_bucket/gcs_bucket.py +++ b/litellm/integrations/gcs_bucket/gcs_bucket.py @@ -10,13 +10,14 @@ from litellm._logging import verbose_logger from litellm._uuid import uuid from litellm.constants import LITELLM_ASYNCIO_QUEUE_MAXSIZE from litellm.integrations.additional_logging_utils import AdditionalLoggingUtils +from litellm.integrations.batch_line_item_objects import batch_line_items_ndjson, batch_line_items_object_name from litellm.integrations.gcs_bucket.gcs_bucket_base import GCSBucketBase from litellm.litellm_core_utils.cloud_storage_security import ( sanitize_cloud_object_component, ) from litellm.types.integrations.base_health_check import IntegrationHealthCheckStatus from litellm.types.integrations.gcs_bucket import * -from litellm.types.utils import StandardLoggingPayload +from litellm.types.utils import StandardBatchLineItemsPayload, StandardLoggingPayload if TYPE_CHECKING: from litellm.llms.vertex_ai.vertex_llm_base import VertexBase @@ -95,6 +96,20 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils): except Exception as e: verbose_logger.exception("GCS Bucket logging error: %s", e) + async def async_log_batch_line_items(self, payload: StandardBatchLineItemsPayload) -> None: + """Write one part of a batch's lines as an NDJSON object under batch_line_items/ in the default bucket.""" + gcs_logging_config: Final = await self.get_gcs_logging_config() + headers: Final = await self.construct_request_headers( + vertex_instance=gcs_logging_config["vertex_instance"], + service_account_json=gcs_logging_config["path_service_account"], + ) + await self._log_json_data_on_gcs( + headers=headers, + bucket_name=gcs_logging_config["bucket_name"], + object_name=batch_line_items_object_name(payload), + logging_payload=await asyncio.to_thread(batch_line_items_ndjson, payload), + ) + def _drain_queue_batch(self) -> list[GCSLogQueueItem]: """ Drain items from the queue (non-blocking), respecting batch_size limit. diff --git a/litellm/integrations/s3_v2.py b/litellm/integrations/s3_v2.py index 0b7a2609bcb..cfd07c9e2af 100644 --- a/litellm/integrations/s3_v2.py +++ b/litellm/integrations/s3_v2.py @@ -35,6 +35,7 @@ from litellm.constants import ( ) from litellm.exceptions import Timeout as LiteLLMTimeout from litellm.integrations.adaptive_concurrency import AdaptiveConcurrencyLimiter, PutSample +from litellm.integrations.batch_line_item_objects import batch_line_items_ndjson, batch_line_items_object_name from litellm.integrations.s3 import ( get_s3_object_download_filename, get_s3_object_key, @@ -60,7 +61,7 @@ from litellm.llms.custom_httpx.http_handler import ( httpxSpecialProvider, ) from litellm.types.integrations.s3_v2 import S3PartitionGranularity, s3BatchLoggingElement -from litellm.types.utils import StandardAuditLogPayload, StandardLoggingPayload +from litellm.types.utils import StandardAuditLogPayload, StandardBatchLineItemsPayload, StandardLoggingPayload from .custom_batch_logger import CustomBatchLogger @@ -548,6 +549,32 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): except Exception as e: verbose_logger.exception("S3 audit log error: %s", e) + async def async_log_batch_line_items(self, payload: StandardBatchLineItemsPayload) -> None: + """Upload one part of a batch's lines as an NDJSON object under /batch_line_items/.""" + object_name: Final = batch_line_items_object_name(payload) + s3_path: Final = cast(str | None, self.s3_path) # cast-ok: s3_path is set from untyped callback params + s3_object_key: Final = (f"{s3_path.rstrip('/')}/" if s3_path else "") + object_name + element: Final = s3BatchLoggingElement( + payload={}, + body=await asyncio.to_thread(batch_line_items_ndjson, payload), + content_type="application/x-ndjson", + s3_object_key=s3_object_key, + s3_object_download_filename="-".join(object_name.rsplit("/", 2)[-2:]), + ) + # hashing a large object is CPU-bound too, so it is also prepared off the event loop + prepared_put: Final = await asyncio.to_thread(self._prepare_put, element) + delivered: Final = await self.async_upload_data_to_s3(element, prepared_put) + if not delivered: + # only a log-queue flush reads and clears the drop marks, and this upload is not part of one + self._flush_dropped.pop(id(element), None) + verbose_logger.error( + "s3 logging: the %d line items of part %d of batch %s were not stored at %s", + len(payload["lines"]), + payload["part"], + payload["batch_id"], + s3_object_key, + ) + async def _async_log_event_base(self, kwargs, response_obj, start_time, end_time): try: verbose_logger.debug("s3 Logging - Enters logging function for model %s", kwargs) @@ -592,6 +619,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): async def async_upload_data_to_s3( self, batch_logging_element: s3BatchLoggingElement, + prepared_put: _PreparedPut | None = None, ) -> bool: try: from litellm.litellm_core_utils.asyncify import asyncify @@ -627,7 +655,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): except _TRANSIENT_ERRORS as error: return error - prepared: Final = self._prepare_put(batch_logging_element) + prepared: Final = prepared_put if prepared_put is not None else self._prepare_put(batch_logging_element) await self._put_with_retries(partial(signed_put, prepared), batch_logging_element.s3_object_key) except Exception as e: verbose_logger.exception("Error uploading to s3: %s", e) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 3d6f8616805..ebb981546c7 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -5311,6 +5311,9 @@ def init_custom_logger_compatible_class_at_startup( ) +init_custom_logger_compatible_class: Final = _init_custom_logger_compatible_class + + def _maybe_construct_otel_v2(callback_name: str, _in_memory_loggers: list[CustomLogger]) -> "OpenTelemetryV2 | None": """If ``LITELLM_OTEL_V2`` is on, build (or reuse) a single ``OpenTelemetryV2`` instance configured via the preset for ``callback_name``. diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index 794208cc2af..c392baa968b 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -424,6 +424,15 @@ def _split_vertex_batch_key(vertex_output_row: Mapping[str, object]) -> tuple[st return unquote(match["custom_id"]), int(match["index"]), int(match["total"]) +def vertex_batch_row_custom_id(vertex_row: Mapping[str, object]) -> str | None: + """The OpenAI `custom_id` a Vertex batch input or output row was written for, or None when it carries none.""" + request_data: Final = vertex_row.get("request") + has_labels: Final = isinstance(request_data, Mapping) and "labels" in request_data + if vertex_row.get(_VERTEX_BATCH_KEY_FIELD) is None and not has_labels: + return None + return _split_vertex_batch_key(vertex_row)[0] + + def _vertex_embeddings_rows_to_openai_batch_output_row( custom_id: str, vertex_output_rows: tuple[_VertexEmbeddingBatchRow, ...], diff --git a/litellm/types/utils.py b/litellm/types/utils.py index e7e4d8489d9..d83e874dee6 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3669,6 +3669,46 @@ class StandardAuditLogPayload(TypedDict): updated_values: str | None +class StandardBatchLineItemUsage(TypedDict): + prompt_tokens: ReadOnly[int] + completion_tokens: ReadOnly[int] + total_tokens: ReadOnly[int] + + +class StandardBatchLineItemMetadata(TypedDict): + user_api_key_hash: ReadOnly[str | None] + user_api_key_alias: ReadOnly[str | None] + user_api_key_team_id: ReadOnly[str | None] + user_api_key_team_alias: ReadOnly[str | None] + user_api_key_user_id: ReadOnly[str | None] + user_api_key_org_id: ReadOnly[str | None] + + +class StandardBatchLineItem(TypedDict): + """One line of a finished batch, the request line next to the output or error line the provider wrote for it.""" + + batch_id: ReadOnly[str] + custom_id: ReadOnly[str | None] + status: ReadOnly[Literal["success", "failure"]] + request: ReadOnly[Mapping[str, object] | None] + response: ReadOnly[Mapping[str, object] | None] + usage: ReadOnly[StandardBatchLineItemUsage | None] + response_cost: ReadOnly[float] + model: ReadOnly[str] + custom_llm_provider: ReadOnly[str | None] + metadata: ReadOnly[StandardBatchLineItemMetadata] + + +class StandardBatchLineItemsPayload(TypedDict): + """One part of a finished batch's lines, dispatched to batch_line_item_callbacks. A batch's lines arrive in parts + numbered from 0, each holding up to BATCH_LINE_ITEMS_MAX_PART_BYTES of raw request and result lines.""" + + batch_id: ReadOnly[str] + finished_at: ReadOnly[int] + part: ReadOnly[int] + lines: ReadOnly[tuple[StandardBatchLineItem, ...]] + + class ClassifierAudit(TypedDict, total=False): classifier_input: ReadOnly[Mapping[str, JsonValue]] originating_request_masked: ReadOnly[Mapping[str, JsonValue]] diff --git a/tests/integration/observability/test_batch_line_items_cold_storage.py b/tests/integration/observability/test_batch_line_items_cold_storage.py new file mode 100644 index 00000000000..9c2a1ecd2c4 --- /dev/null +++ b/tests/integration/observability/test_batch_line_items_cold_storage.py @@ -0,0 +1,902 @@ +"""Finished batches write their lines to cold storage. + +An owned two-worker proxy lists ``s3_v2`` and ``gcs_bucket`` under ``batch_line_item_callbacks``. Once the batch-cost +poller logs a finished batch, each sink gets the batch's lines as NDJSON part objects holding one record per line, a +single part for these small batches. The proxy runs on its own database so the shared rig's pollers never claim these +batches first, and accepts JWTs signed by a local key next to virtual keys. Four more owned proxies list an unsupported +name next to ``s3_v2``, list a null list, check OAuth2 tokens against a local token info endpoint, and close a part +after every line +""" + +import json +import threading +import time +import uuid +from collections.abc import Callable, Iterator, Mapping +from contextlib import contextmanager +from dataclasses import dataclass, field +from datetime import datetime, timezone +from hashlib import sha256 +from pathlib import Path +from types import MappingProxyType +from typing import Final +from urllib.parse import parse_qs, unquote, urlsplit + +import jwt +import pytest +import yaml +from cryptography.hazmat.primitives.asymmetric import rsa +from integration._support.client import JSON_OBJECT, Gateway, eventually, gateway_from_environment, held, string_value +from integration._support.database import read_rows, scratch_database +from integration._support.forward_proxy import tunnelling_forward_proxy +from integration._support.process import owned_proxy +from integration._support.tls import server_context, write_self_signed_cert +from integration._support.upstream import CONTROL_URL, ScenarioHandle, delete_scenario, register_scenario +from integration._support.vertex import service_account_json +from integration._support.wire import Reply, Request, wire_server +from integration.cost_calculation.cost_tracking_case import JsonResponse, RoutedResponse, TextResponse +from jwt.algorithms import RSAAlgorithm +from pydantic import JsonValue + +pytestmark: Final = pytest.mark.timeout(420) + +S3_BUCKET: Final = "line-items-bucket" +S3_PREFIX: Final = "cold" +GCS_BUCKET: Final = "line-items-gcs" +COMPLETED_AT: Final = 1791494130 +LINE_DATE: Final = datetime.fromtimestamp(COMPLETED_AT, tz=timezone.utc).strftime("%Y-%m-%d") +PROMPT_TOKENS: Final = 100 +COMPLETION_TOKENS: Final = 50 +POLLER_SECONDS: Final = 150 +JWT_KEY_ID: Final = "batch-line-items-signing-key" +TWO_ACCOUNT_MODEL: Final = "line-items-two-accounts" +FIRST_ACCOUNT: Final = f"batch-line-items-first-account-{uuid.uuid4().hex[:12]}" +SECOND_ACCOUNT: Final = f"batch-line-items-second-account-{uuid.uuid4().hex[:12]}" + + +@dataclass(slots=True) +class ObjectSink: + """Stores every accepted object body by name and answers 503 to any upload whose body holds ``reject``. An upload + whose body holds ``hang`` gets no answer until ``release``, then a 503.""" + + reject: bytes = b"\x00never" + hang: bytes = b"\x00never" + released: threading.Event = field(default_factory=threading.Event) + lock: threading.Lock = field(default_factory=threading.Lock) + held: set[str] = field(default_factory=set) # mutable-ok: written by the sink's request threads + stored: dict[str, bytes] = field(default_factory=dict) # mutable-ok: written by the sink's request threads + attempts: dict[str, int] = field(default_factory=dict) # mutable-ok: written by the sink's request threads + + def accept(self, name: str, body: bytes) -> bool: + with self.lock: + self.attempts[name] = self.attempts.get(name, 0) + 1 + hung: Final = self.hang in body + if hung: + self.held.add(name) + if hung: + self.released.wait(timeout=300) + return False + with self.lock: + if self.reject in body: + return False + self.stored[name] = body + return True + + def holding(self) -> frozenset[str]: + with self.lock: + return frozenset(self.held) + + def release(self) -> None: + self.released.set() + + def objects(self) -> Mapping[str, bytes]: + with self.lock: + return MappingProxyType(dict(self.stored)) + + def tried(self, name: str) -> int: + with self.lock: + return self.attempts.get(name, 0) + + +def s3_peer(sink: ObjectSink) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + if request.method != "PUT": + return Reply(status=404) + name: Final = unquote(urlsplit(request.target).path).removeprefix(f"/{S3_BUCKET}/") + if sink.accept(name, request.body): + return Reply() + return Reply(status=503, body=b"SlowDown", content_type="application/xml") + + return respond + + +def gcs_peer(sink: ObjectSink) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + path: Final = urlsplit(request.target).path + if request.method == "POST" and path == "/_oauth/token": + return Reply(body=b'{"access_token":"scripted-token","expires_in":3600,"token_type":"Bearer"}') + if request.method == "POST" and path == f"/upload/storage/v1/b/{GCS_BUCKET}/o": + name: Final = parse_qs(urlsplit(request.target).query)["name"][0] + if sink.accept(name, request.body): + return Reply(body=json.dumps({"name": name, "bucket": GCS_BUCKET}).encode()) + return Reply(status=503, body=b'{"error":"unavailable"}') + return Reply(status=404, body=b'{"error":"not scripted"}') + + return respond + + +@dataclass(frozen=True, slots=True) +class Rig: + gateway: Gateway + database_url: str + s3: ObjectSink + gcs: ObjectSink + signing_key: rsa.RSAPrivateKey + + +def _jwks_peer(signing_key: rsa.RSAPrivateKey) -> Callable[[Request], Reply]: + public_jwk: Final = json.loads(RSAAlgorithm.to_jwk(signing_key.public_key())) + jwks: Final = json.dumps({"keys": [{**public_jwk, "kid": JWT_KEY_ID}]}).encode() + + def respond(_request: Request) -> Reply: + return Reply(body=jwks) + + return respond + + +def _signed_token(signing_key: rsa.RSAPrivateKey, subject: str) -> str: + now: Final = int(time.time()) + return jwt.encode( + {"sub": subject, "iat": now, "exp": now + 300, "jti": uuid.uuid4().hex}, + signing_key, + algorithm="RS256", + headers={"kid": JWT_KEY_ID}, + ) + + +def _oauth2_token_info_peer() -> Callable[[Request], Reply]: + """Answers for a bearer token ``oauth__`` with that user and team, and 401 for anything else.""" + + def respond(request: Request) -> Reply: + scheme, _, token = request.headers.get("authorization", "").partition(" ") + parts: Final = token.split("_") + if scheme != "Bearer" or len(parts) != 3 or parts[0] != "oauth": + return Reply(status=401, body=b'{"error":"invalid_token"}') + return Reply(body=json.dumps({"sub": parts[2], "team_id": parts[1]}).encode()) + + return respond + + +def _oauth2_token(team_id: str, user_id: str) -> str: + return f"oauth_{team_id}_{user_id}" + + +def _config(directory: Path, s3_url: str, line_item_callbacks: JsonValue, *, oauth2: bool) -> Path: + shipped: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + configuration: Final = { + **shipped, + "model_list": [ + *shipped["model_list"], + *( + { + "model_name": TWO_ACCOUNT_MODEL, + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "sk-scripted-provider", + "api_base": f"{CONTROL_URL}/{account}", + "input_cost_per_token": 0.001, + "output_cost_per_token": 0.002, + "weight": weight, + }, + } + for account, weight in ((FIRST_ACCOUNT, 0), (SECOND_ACCOUNT, 1)) + ), + ], + "general_settings": { + **shipped["general_settings"], + "enable_jwt_auth": True, + "litellm_jwtauth": {"user_id_jwt_field": "sub", "user_id_upsert": True}, + "enable_oauth2_auth": oauth2, + }, + "litellm_settings": { + **shipped["litellm_settings"], + "callbacks": ["s3_v2"], + "batch_line_item_callbacks": line_item_callbacks, + "s3_callback_params": { + "s3_bucket_name": S3_BUCKET, + "s3_region_name": "us-east-1", + "s3_endpoint_url": s3_url, + "s3_path": S3_PREFIX, + "s3_aws_access_key_id": "AKIAIOSFODNN7EXAMPLE", + "s3_aws_secret_access_key": "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY", + "s3_flush_interval": 1, + }, + }, + } + path: Final = directory / "batch-line-items.yaml" + path.write_text(yaml.safe_dump(configuration)) + return path + + +@contextmanager +def _owned_rig( + tmp_path_factory: pytest.TempPathFactory, + name: str, + line_item_callbacks: JsonValue, + *, + oauth2: bool = False, + max_part_bytes: int | None = None, +) -> Iterator[Rig]: + directory: Final = tmp_path_factory.mktemp(name) + certificate, key = write_self_signed_cert(directory, names=("storage.googleapis.com",)) + s3: Final = ObjectSink(reject=b"sink-down-", hang=b"sink-hang-") + gcs: Final = ObjectSink() + signing_key: Final = rsa.generate_private_key(public_exponent=65537, key_size=2048) + with ( + wire_server(_jwks_peer(signing_key)) as jwks_wire, + wire_server(_oauth2_token_info_peer()) as token_info_wire, + gateway_from_environment() as environment, + scratch_database(name.replace("-", "_")) as database_url, + wire_server(s3_peer(s3)) as s3_wire, + wire_server(gcs_peer(gcs)) as token_wire, + wire_server(gcs_peer(gcs), tls=server_context(certificate, key)) as gcs_wire, + tunnelling_forward_proxy(int(urlsplit(gcs_wire.url).port or 0)) as tunnel, + ): + service_account: Final = directory / "service-account.json" + service_account.write_text(service_account_json("line-items-project", token_wire.url)) + with owned_proxy( + environment, + directory, + { + "DATABASE_URL": database_url, + "GCS_BUCKET_NAME": GCS_BUCKET, + "GCS_PATH_SERVICE_ACCOUNT": str(service_account), + "GCS_FLUSH_INTERVAL": "1", + "JWT_PUBLIC_KEY_URL": jwks_wire.url, + "OAUTH_TOKEN_INFO_ENDPOINT": token_info_wire.url, + "HTTPS_PROXY": tunnel.url, + "https_proxy": tunnel.url, + "SSL_CERT_FILE": str(certificate), + "NO_PROXY": "127.0.0.1,localhost", + "no_proxy": "127.0.0.1,localhost", + **({"BATCH_LINE_ITEMS_MAX_PART_BYTES": str(max_part_bytes)} if max_part_bytes else {}), + }, + config=_config(directory, s3_wire.url, line_item_callbacks, oauth2=oauth2), + remove_environment=("DATABASE_URL_READ_REPLICA",), + workers=2, + ) as candidate: + yield Rig(candidate, database_url, s3, gcs, signing_key) + + +@pytest.fixture(scope="module") +def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + with _owned_rig(tmp_path_factory, "batch-line-items", ["s3_v2", "gcs_bucket"]) as owned: + yield owned + + +@pytest.fixture(scope="module") +def rig_with_an_unsupported_name(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + with _owned_rig(tmp_path_factory, "batch-line-items-unsupported", ["s3", "s3_v2"]) as owned: + yield owned + + +@pytest.fixture(scope="module") +def rig_with_a_null_list(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + with _owned_rig(tmp_path_factory, "batch-line-items-null", None) as owned: + yield owned + + +@pytest.fixture(scope="module") +def rig_with_oauth2(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + with _owned_rig(tmp_path_factory, "batch-line-items-oauth2", ["s3_v2"], oauth2=True) as owned: + yield owned + + +@pytest.fixture(scope="module") +def rig_with_one_line_parts(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + with _owned_rig(tmp_path_factory, "batch-line-items-parts", ["s3_v2", "gcs_bucket"], max_part_bytes=1) as owned: + yield owned + + +def _input_line(custom_id: str, text: str) -> dict[str, JsonValue]: + return { + "custom_id": custom_id, + "method": "POST", + "url": "/v1/chat/completions", + "body": {"model": "gpt-4o-mini", "messages": [{"role": "user", "content": text}]}, + } + + +def _success_line(custom_id: str) -> dict[str, JsonValue]: + return { + "id": f"batch_req_{custom_id}", + "custom_id": custom_id, + "response": { + "status_code": 200, + "request_id": f"req-{custom_id}", + "body": { + "id": f"chatcmpl-{custom_id}", + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": { + "prompt_tokens": PROMPT_TOKENS, + "completion_tokens": COMPLETION_TOKENS, + "total_tokens": PROMPT_TOKENS + COMPLETION_TOKENS, + }, + }, + }, + "error": None, + } + + +def _error_line(custom_id: str) -> dict[str, JsonValue]: + return { + "id": f"batch_req_{custom_id}", + "custom_id": custom_id, + "response": { + "status_code": 400, + "request_id": f"req-{custom_id}", + "body": {"error": {"message": "max_tokens is too large", "type": "invalid_request_error"}}, + }, + "error": None, + } + + +def _jsonl(lines: tuple[dict[str, JsonValue], ...]) -> str: + return "".join(json.dumps(line, separators=(",", ":")) + "\n" for line in lines) + + +def _file(identity: str, purpose: str) -> JsonResponse: + return JsonResponse( + content_type="application/json", + body={ + "id": identity, + "object": "file", + "purpose": purpose, + "bytes": 100, + "created_at": 1, + "filename": f"{identity}.jsonl", + "status": "processed", + }, + ) + + +@dataclass(frozen=True, slots=True) +class Lines: + inputs: tuple[dict[str, JsonValue], ...] + outputs: tuple[dict[str, JsonValue], ...] + errors: tuple[dict[str, JsonValue], ...] + + +def _routes(lines: Lines, *, finished: bool = True) -> RoutedResponse: + has_output: Final = bool(lines.outputs) + completed: Final[dict[str, JsonValue]] = { + "id": "batch-$REQUEST_ID", + "object": "batch", + "endpoint": "/v1/chat/completions", + "errors": None, + "input_file_id": "file-in-$REQUEST_ID", + "completion_window": "24h", + "status": "completed", + "output_file_id": "file-out-$REQUEST_ID" if has_output else None, + "error_file_id": "file-err-$REQUEST_ID" if lines.errors else None, + "created_at": COMPLETED_AT - 60, + "in_progress_at": COMPLETED_AT - 30, + "completed_at": COMPLETED_AT, + "expires_at": COMPLETED_AT + 86400, + "request_counts": { + "total": len(lines.inputs), + "completed": len(lines.outputs), + "failed": len(lines.errors), + }, + "metadata": None, + } + return RoutedResponse( + content_type="application/x-routed", + routes={ + "POST /files": _file("file-in-$REQUEST_ID", "batch"), + "POST /batches": JsonResponse( + content_type="application/json", + body={**completed, "status": "validating", "output_file_id": None, "error_file_id": None}, + ), + "GET /batches/batch-$REQUEST_ID": JsonResponse( + content_type="application/json", + body=completed + if finished + else {**completed, "status": "in_progress", "output_file_id": None, "error_file_id": None}, + ), + "GET /files/file-in-$REQUEST_ID/content": TextResponse( + content_type="application/jsonl", body=_jsonl(lines.inputs) + ), + "GET /files/file-out-$REQUEST_ID": _file("file-out-$REQUEST_ID", "batch_output"), + "GET /files/file-out-$REQUEST_ID/content": TextResponse( + content_type="application/jsonl", body=_jsonl(lines.outputs) + ), + "GET /files/file-err-$REQUEST_ID": _file("file-err-$REQUEST_ID", "batch_output"), + "GET /files/file-err-$REQUEST_ID/content": TextResponse( + content_type="application/jsonl", body=_jsonl(lines.errors) + ), + }, + ) + + +@dataclass(frozen=True, slots=True) +class SubmittedBatch: + batch_id: str + key: str + handle: ScenarioHandle + + +def _submit( + rig: Rig, lines: Lines, key_fields: Mapping[str, JsonValue], *, finished: bool = True, token: str | None = None +) -> SubmittedBatch: + handle: Final = register_scenario(f"batch-line-items-{uuid.uuid4().hex}", _routes(lines, finished=finished)) + return _create_batch(rig, lines, key_fields, _add_model(rig, handle), handle, token) + + +def _add_model(rig: Rig, handle: ScenarioHandle) -> str: + return string_value( + rig.gateway.post( + "/model/new", + { + "model_name": f"line-items-{uuid.uuid4().hex[:12]}", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "sk-scripted-provider", + "api_base": handle.api_base(), + "input_cost_per_token": 0.001, + "output_cost_per_token": 0.002, + }, + }, + )["model_name"] + ) + + +def _create_batch( + rig: Rig, lines: Lines, key_fields: Mapping[str, JsonValue], model: str, handle: ScenarioHandle, token: str | None +) -> SubmittedBatch: + key: Final = ( + token + if token is not None + else string_value(rig.gateway.post("/key/generate", {"models": [model], **key_fields})["key"]) + ) + uploaded: Final = rig.gateway.request_multipart( + "/v1/files", + {"purpose": "batch", "target_model_names": model}, + {"file": ("in.jsonl", _jsonl(lines.inputs).encode(), "application/jsonl")}, + key=key, + ) + assert uploaded.is_success, uploaded.text + created: Final = rig.gateway.request( + "POST", + "/v1/batches", + { + "input_file_id": string_value(JSON_OBJECT.validate_json(uploaded.content)["id"]), + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + "model": model, + }, + key=key, + ) + assert created.is_success, created.text + return SubmittedBatch(string_value(JSON_OBJECT.validate_json(created.content)["id"]), key, handle) + + +def _quiet_logging() -> dict[str, JsonValue]: + return { + "logging": [ + { + "callback_name": "s3_v2", + "callback_type": "success", + "callback_vars": {"turn_off_message_logging": "true"}, + } + ] + } + + +def _object_name(prefix: str, batch_id: str, part: int = 0) -> str: + return f"{prefix}batch_line_items/{LINE_DATE}/{batch_id.replace('/', '_')}/part-{part:05d}.jsonl" + + +def _records(body: bytes) -> tuple[dict[str, JsonValue], ...]: + return tuple(JSON_OBJECT.validate_json(line) for line in body.splitlines() if line.strip()) + + +def _written(sink: ObjectSink, name: str) -> tuple[dict[str, JsonValue], ...]: + return _records(eventually(lambda: sink.objects().get(name), lambda body: body is not None, POLLER_SECONDS) or b"") + + +def _batch_spend(rig: Rig, key: str) -> float: + rows: Final = eventually( + lambda: read_rows( + 'SELECT spend FROM "LiteLLM_SpendLogs" WHERE api_key = %s AND call_type = %s', + (sha256(key.encode()).hexdigest(), "aretrieve_batch"), + database_url=rig.database_url, + ), + lambda found: len(found) == 1, + POLLER_SECONDS, + ) + spend: Final = rows[0]["spend"] + assert isinstance(spend, float), rows + return spend + + +def _by_custom_id(records: tuple[dict[str, JsonValue], ...]) -> Mapping[JsonValue, dict[str, JsonValue]]: + keyed: Final = {record["custom_id"]: record for record in records} + assert len(keyed) == len(records), records + return MappingProxyType(keyed) + + +def test_a_finished_batch_writes_every_line_with_its_request_and_result_to_s3_and_gcs(rig: Rig) -> None: + lines: Final = Lines( + inputs=(_input_line("ok-1", "first question"), _input_line("bad-1", "second question")), + outputs=(_success_line("ok-1"),), + errors=(_error_line("bad-1"),), + ) + batch: Final = _submit(rig, lines, {}) + try: + s3_records: Final = _written(rig.s3, _object_name(f"{S3_PREFIX}/", batch.batch_id)) + gcs_body: Final = eventually( + lambda: rig.gcs.objects().get(_object_name("", batch.batch_id)), lambda body: body is not None + ) + spend: Final = _batch_spend(rig, batch.key) + finally: + delete_scenario(batch.handle) + + assert gcs_body is not None + assert _records(gcs_body) == s3_records + records: Final = _by_custom_id(s3_records) + assert set(records) == {"ok-1", "bad-1"} + hashed_key: Final = sha256(batch.key.encode()).hexdigest() + succeeded: Final = records["ok-1"] + assert succeeded["batch_id"] == batch.batch_id + assert succeeded["status"] == "success" + assert succeeded["request"] == lines.inputs[0] + assert succeeded["response"] == lines.outputs[0] + assert succeeded["usage"] == { + "prompt_tokens": PROMPT_TOKENS, + "completion_tokens": COMPLETION_TOKENS, + "total_tokens": PROMPT_TOKENS + COMPLETION_TOKENS, + } + assert succeeded["model"] == "gpt-4o-mini" + assert succeeded["custom_llm_provider"] == "openai" + assert isinstance(succeeded["metadata"], dict) + assert succeeded["metadata"]["user_api_key_hash"] == hashed_key + failed: Final = records["bad-1"] + assert failed["status"] == "failure" + assert failed["request"] == lines.inputs[1] + assert failed["response"] == lines.errors[0] + assert failed["usage"] is None + assert failed["response_cost"] == 0 + assert isinstance(succeeded["response_cost"], float) and succeeded["response_cost"] > 0 + assert succeeded["response_cost"] == pytest.approx(spend) + assert rig.s3.tried(_object_name(f"{S3_PREFIX}/", batch.batch_id)) == 1 + + +def test_a_key_that_turned_off_message_logging_gets_its_lines_without_request_or_response(rig: Rig) -> None: + lines: Final = Lines( + inputs=(_input_line("quiet-1", "a private question"),), outputs=(_success_line("quiet-1"),), errors=() + ) + batch: Final = _submit(rig, lines, {"metadata": _quiet_logging()}) + name: Final = _object_name(f"{S3_PREFIX}/", batch.batch_id) + try: + records: Final = _written(rig.s3, name) + finally: + delete_scenario(batch.handle) + + assert len(records) == 1 + assert records[0]["custom_id"] == "quiet-1" + assert records[0]["status"] == "success" + assert records[0]["request"] is None + assert records[0]["response"] is None + assert records[0]["usage"] == { + "prompt_tokens": PROMPT_TOKENS, + "completion_tokens": COMPLETION_TOKENS, + "total_tokens": PROMPT_TOKENS + COMPLETION_TOKENS, + } + assert b"a private question" not in rig.s3.objects()[name] + + +def test_a_team_that_turned_off_message_logging_gets_its_members_lines_without_request_or_response(rig: Rig) -> None: + team_id: Final = string_value( + rig.gateway.post("/team/new", {"team_alias": f"quiet-{uuid.uuid4().hex[:8]}", "metadata": _quiet_logging()})[ + "team_id" + ] + ) + lines: Final = Lines( + inputs=(_input_line("team-1", "a team secret"),), outputs=(_success_line("team-1"),), errors=() + ) + batch: Final = _submit(rig, lines, {"team_id": team_id}) + name: Final = _object_name(f"{S3_PREFIX}/", batch.batch_id) + try: + records: Final = _written(rig.s3, name) + finally: + delete_scenario(batch.handle) + + assert [(record["custom_id"], record["request"], record["response"]) for record in records] == [ + ("team-1", None, None) + ] + assert isinstance(records[0]["metadata"], dict) + assert records[0]["metadata"]["user_api_key_team_id"] == team_id + assert b"a team secret" not in rig.s3.objects()[name] + + +def test_a_batch_whose_creating_key_was_deleted_gets_its_lines_without_request_or_response(rig: Rig) -> None: + lines: Final = Lines( + inputs=(_input_line("orphan-1", "asked before the key was deleted"),), + outputs=(_success_line("orphan-1"),), + errors=(), + ) + batch: Final = _submit(rig, lines, {}, finished=False) + name: Final = _object_name(f"{S3_PREFIX}/", batch.batch_id) + try: + rig.gateway.post("/key/delete", {"keys": [batch.key]}) + register_scenario(batch.handle.scenario_id, _routes(lines)) + records: Final = _written(rig.s3, name) + spend: Final = _batch_spend(rig, batch.key) + finally: + delete_scenario(batch.handle) + + assert [(record["custom_id"], record["status"], record["request"], record["response"]) for record in records] == [ + ("orphan-1", "success", None, None) + ] + assert records[0]["response_cost"] == pytest.approx(spend) + assert b"asked before the key was deleted" not in rig.s3.objects()[name] + + +def test_a_batch_created_with_a_jwt_and_no_key_gets_its_lines_with_request_and_response(rig: Rig) -> None: + subject: Final = f"jwt-{uuid.uuid4().hex[:12]}" + lines: Final = Lines( + inputs=(_input_line("jwt-1", "asked with a signed token"),), outputs=(_success_line("jwt-1"),), errors=() + ) + batch: Final = _submit(rig, lines, {}, token=_signed_token(rig.signing_key, subject)) + try: + records: Final = _written(rig.s3, _object_name(f"{S3_PREFIX}/", batch.batch_id)) + finally: + delete_scenario(batch.handle) + + assert read_rows( + 'SELECT api_key, created_by FROM "LiteLLM_ManagedObjectTable" WHERE unified_object_id = %s', + (batch.batch_id,), + database_url=rig.database_url, + ) == [{"api_key": None, "created_by": subject}] + assert [(record["custom_id"], record["request"], record["response"]) for record in records] == [ + ("jwt-1", lines.inputs[0], lines.outputs[0]) + ] + + +def test_a_batch_created_with_an_oauth2_token_follows_its_teams_logging_setting_and_records_only_its_hash( + rig_with_oauth2: Rig, +) -> None: + rig: Final = rig_with_oauth2 + open_team: Final = string_value( + rig.gateway.post("/team/new", {"team_alias": f"open-{uuid.uuid4().hex[:8]}"})["team_id"] + ) + quiet_team: Final = string_value( + rig.gateway.post("/team/new", {"team_alias": f"quiet-{uuid.uuid4().hex[:8]}", "metadata": _quiet_logging()})[ + "team_id" + ] + ) + user_id: Final = f"oauth-user-{uuid.uuid4().hex[:12]}" + open_lines: Final = Lines( + inputs=(_input_line("oauth-open", "asked with an oauth token"),), + outputs=(_success_line("oauth-open"),), + errors=(), + ) + quiet_lines: Final = Lines( + inputs=(_input_line("oauth-quiet", "a quiet team's oauth secret"),), + outputs=(_success_line("oauth-quiet"),), + errors=(), + ) + opened: Final = _submit(rig, open_lines, {}, token=_oauth2_token(open_team, user_id)) + quieted: Final = _submit(rig, quiet_lines, {}, token=_oauth2_token(quiet_team, user_id)) + quiet_name: Final = _object_name(f"{S3_PREFIX}/", quieted.batch_id) + try: + open_records: Final = _written(rig.s3, _object_name(f"{S3_PREFIX}/", opened.batch_id)) + quiet_records: Final = _written(rig.s3, quiet_name) + finally: + delete_scenario(opened.handle) + delete_scenario(quieted.handle) + + assert read_rows( + 'SELECT api_key, created_by, team_id FROM "LiteLLM_ManagedObjectTable" WHERE unified_object_id = %s', + (opened.batch_id,), + database_url=rig.database_url, + ) == [{"api_key": _oauth2_token(open_team, user_id), "created_by": user_id, "team_id": open_team}] + assert [(record["custom_id"], record["request"], record["response"]) for record in open_records] == [ + ("oauth-open", open_lines.inputs[0], open_lines.outputs[0]) + ] + assert [(record["custom_id"], record["request"], record["response"]) for record in quiet_records] == [ + ("oauth-quiet", None, None) + ] + assert b"a quiet team's oauth secret" not in rig.s3.objects()[quiet_name] + open_token: Final = _oauth2_token(open_team, user_id) + assert [record["metadata"] for record in open_records] == [ + {**open_records[0]["metadata"], "user_api_key_hash": sha256(open_token.encode()).hexdigest()} + ] + assert open_token.encode() not in rig.s3.objects()[_object_name(f"{S3_PREFIX}/", opened.batch_id)] + + +def test_a_batch_bigger_than_one_part_is_written_as_numbered_parts_holding_each_line_once( + rig_with_one_line_parts: Rig, +) -> None: + rig: Final = rig_with_one_line_parts + lines: Final = Lines( + inputs=(_input_line("part-a", "first"), _input_line("part-b", "second"), _input_line("part-c", "third")), + outputs=(_success_line("part-a"), _success_line("part-b")), + errors=(_error_line("part-c"),), + ) + batch: Final = _submit(rig, lines, {}) + try: + s3_parts: Final = tuple( + _written(rig.s3, _object_name(f"{S3_PREFIX}/", batch.batch_id, part)) for part in range(3) + ) + gcs_parts: Final = tuple(_written(rig.gcs, _object_name("", batch.batch_id, part)) for part in range(3)) + finally: + delete_scenario(batch.handle) + + assert [[record["custom_id"] for record in part] for part in s3_parts] == [["part-a"], ["part-b"], ["part-c"]] + assert [record["request"] for part in s3_parts for record in part] == list(lines.inputs) + assert gcs_parts == s3_parts + assert _object_name(f"{S3_PREFIX}/", batch.batch_id, 3) not in rig.s3.objects() + + +def test_a_batch_where_every_line_failed_still_writes_its_error_lines(rig: Rig) -> None: + lines: Final = Lines( + inputs=(_input_line("bad-a", "first"), _input_line("bad-b", "second")), + outputs=(), + errors=(_error_line("bad-a"), _error_line("bad-b")), + ) + batch: Final = _submit(rig, lines, {}) + try: + records: Final = _by_custom_id(_written(rig.s3, _object_name(f"{S3_PREFIX}/", batch.batch_id))) + finally: + delete_scenario(batch.handle) + + assert set(records) == {"bad-a", "bad-b"} + assert records["bad-a"]["status"] == "failure" + assert records["bad-a"]["request"] == lines.inputs[0] + assert records["bad-b"]["response"] == lines.errors[1] + assert all(record["response_cost"] == 0 for record in records.values()) + + +def test_a_failed_batch_on_the_second_deployment_of_its_model_group_keeps_each_lines_request(rig: Rig) -> None: + lines: Final = Lines( + inputs=(_input_line("second-a", "first"), _input_line("second-b", "second")), + outputs=(), + errors=(_error_line("second-a"), _error_line("second-b")), + ) + first: Final = register_scenario(FIRST_ACCOUNT, _routes(lines)) + second: Final = register_scenario(SECOND_ACCOUNT, _routes(lines)) + try: + batch: Final = _create_batch(rig, lines, {}, TWO_ACCOUNT_MODEL, second, None) + records: Final = _by_custom_id(_written(rig.s3, _object_name(f"{S3_PREFIX}/", batch.batch_id))) + finally: + delete_scenario(first) + delete_scenario(second) + + assert read_rows( + 'SELECT model_object_id FROM "LiteLLM_ManagedObjectTable" WHERE unified_object_id = %s', + (batch.batch_id,), + database_url=rig.database_url, + ) == [{"model_object_id": f"batch-{SECOND_ACCOUNT}"}] + assert {custom_id: record["request"] for custom_id, record in records.items()} == { + "second-a": lines.inputs[0], + "second-b": lines.inputs[1], + } + + +def test_a_failing_bucket_still_bills_the_batch_and_the_other_sink_keeps_its_lines(rig: Rig) -> None: + custom_id: Final = f"sink-down-{uuid.uuid4().hex[:8]}" + lines: Final = Lines( + inputs=(_input_line(custom_id, "billed anyway"),), outputs=(_success_line(custom_id),), errors=() + ) + batch: Final = _submit(rig, lines, {}) + try: + gcs_body: Final = eventually( + lambda: rig.gcs.objects().get(_object_name("", batch.batch_id)), + lambda body: body is not None, + POLLER_SECONDS, + ) + spend: Final = _batch_spend(rig, batch.key) + processed: Final = eventually( + lambda: read_rows( + 'SELECT batch_processed FROM "LiteLLM_ManagedObjectTable" WHERE unified_object_id = %s', + (batch.batch_id,), + database_url=rig.database_url, + ), + lambda rows: rows == [{"batch_processed": True}], + ) + finally: + delete_scenario(batch.handle) + + assert processed == [{"batch_processed": True}] + assert gcs_body is not None + assert [record["custom_id"] for record in _records(gcs_body)] == [custom_id] + assert spend > 0 + s3_name: Final = _object_name(f"{S3_PREFIX}/", batch.batch_id) + assert rig.s3.tried(s3_name) >= 1 + assert s3_name not in rig.s3.objects() + + +def _hanging_lines() -> Lines: + custom_id: Final = f"sink-hang-{uuid.uuid4().hex[:8]}" + return Lines(inputs=(_input_line(custom_id, "hangs the bucket"),), outputs=(_success_line(custom_id),), errors=()) + + +def test_a_hung_bucket_never_holds_up_billing_the_next_batch(rig: Rig) -> None: + """The first hung write holds one worker's poller if the poll waits on it, and the second batch, billed by the + other worker, holds that one too, so only a poll that never waits on a write bills the third batch.""" + first: Final = _submit(rig, _hanging_lines(), {}) + second: SubmittedBatch | None = None + following: SubmittedBatch | None = None + try: + eventually(rig.s3.holding, lambda names: _object_name(f"{S3_PREFIX}/", first.batch_id) in names, POLLER_SECONDS) + second = _submit(rig, _hanging_lines(), {}) + _batch_spend(rig, second.key) + following = _submit( + rig, + Lines(inputs=(_input_line("after-1", "asked later"),), outputs=(_success_line("after-1"),), errors=()), + {}, + ) + spend: Final = _batch_spend(rig, following.key) + still_holding: Final = rig.s3.holding() + finally: + rig.s3.release() + for batch in (first, second, following): + if batch is not None: + delete_scenario(batch.handle) + + assert spend > 0 + assert _object_name(f"{S3_PREFIX}/", first.batch_id) in still_holding + assert _object_name(f"{S3_PREFIX}/", first.batch_id) not in rig.s3.objects() + + +def _line_item_objects(rig: Rig) -> frozenset[str]: + return frozenset(name for sink in (rig.s3, rig.gcs) for name in sink.objects() if "batch_line_items/" in name) + + +def test_an_unsupported_name_in_the_list_leaves_the_supported_sink_writing(rig_with_an_unsupported_name: Rig) -> None: + rig: Final = rig_with_an_unsupported_name + lines: Final = Lines( + inputs=(_input_line("listed-1", "asked with s3 listed"),), outputs=(_success_line("listed-1"),), errors=() + ) + batch: Final = _submit(rig, lines, {}) + try: + records: Final = _written(rig.s3, _object_name(f"{S3_PREFIX}/", batch.batch_id)) + spend: Final = _batch_spend(rig, batch.key) + finally: + delete_scenario(batch.handle) + + assert [(record["custom_id"], record["request"], record["response"]) for record in records] == [ + ("listed-1", lines.inputs[0], lines.outputs[0]) + ] + assert records[0]["response_cost"] == pytest.approx(spend) + assert _line_item_objects(rig) == {_object_name(f"{S3_PREFIX}/", batch.batch_id)} + + +def test_a_null_list_bills_the_batch_and_writes_no_line_items(rig_with_a_null_list: Rig) -> None: + rig: Final = rig_with_a_null_list + lines: Final = Lines( + inputs=(_input_line("null-1", "asked with a null list"),), outputs=(_success_line("null-1"),), errors=() + ) + batch: Final = _submit(rig, lines, {}) + try: + spend: Final = _batch_spend(rig, batch.key) + eventually( + lambda: read_rows( + 'SELECT batch_processed FROM "LiteLLM_ManagedObjectTable" WHERE unified_object_id = %s', + (batch.batch_id,), + database_url=rig.database_url, + ), + lambda rows: rows == [{"batch_processed": True}], + POLLER_SECONDS, + ) + stored: Final = held(lambda: _line_item_objects(rig), lambda names: not names, holding=10) + finally: + delete_scenario(batch.handle) + + assert spend > 0 + assert stored == frozenset() diff --git a/tests/unit/batches/test_batch_line_items.py b/tests/unit/batches/test_batch_line_items.py new file mode 100644 index 00000000000..d181a2faa86 --- /dev/null +++ b/tests/unit/batches/test_batch_line_items.py @@ -0,0 +1,434 @@ +import json +import logging +from hashlib import sha256 + +import pytest + +import litellm +from litellm.batches.batch_line_items import ( + BatchLineItemSource, + batch_line_item_parts, + dispatch_batch_line_items, + line_item_metadata, + without_request_and_response, +) +from litellm.batches.batch_utils import calculate_batch_cost_and_usage, get_file_content_as_dictionary +from litellm.constants import BATCH_LINE_ITEMS_MAX_PART_BYTES +from litellm.integrations.batch_line_item_objects import batch_line_items_ndjson, batch_line_items_object_name +from litellm.integrations.custom_logger import CustomLogger +from litellm.llms.vertex_ai.files.transformation import VertexAIFilesConfig, _OpenAIToVertexBatchUploadStream +from litellm.types.utils import ModelInfo, StandardBatchLineItemsPayload + +INPUT_RATE = 0.001 +OUTPUT_RATE = 0.002 +FINISHED_AT = 1791494130 +ALICE_KEY_HASH = sha256(b"sk-alice").hexdigest() + + +def _jsonl(*lines: dict) -> bytes: + return "".join(json.dumps(line) + "\n" for line in lines).encode() + + +def _request(custom_id: str, text: str) -> dict: + return { + "custom_id": custom_id, + "method": "POST", + "url": "/v1/chat/completions", + "body": {"model": "gpt-4o-mini", "messages": [{"role": "user", "content": text}]}, + } + + +def _success(custom_id: str, prompt_tokens: int, completion_tokens: int) -> dict: + return { + "id": f"batch_req_{custom_id}", + "custom_id": custom_id, + "response": { + "status_code": 200, + "body": { + "model": "gpt-4o-mini", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}}], + "usage": { + "prompt_tokens": prompt_tokens, + "completion_tokens": completion_tokens, + "total_tokens": prompt_tokens + completion_tokens, + }, + }, + }, + "error": None, + } + + +def _error(custom_id: str) -> dict: + return { + "id": f"batch_req_{custom_id}", + "custom_id": custom_id, + "response": {"status_code": 400, "body": {"error": {"message": "max_tokens is too large"}}}, + "error": None, + } + + +def _source(**overrides) -> BatchLineItemSource: + fields = { + "batch_id": "bGl0ZWxsbV9wcm94eTtiYXRjaA", + "finished_at": FINISHED_AT, + "model": "gpt-4o-mini", + "custom_llm_provider": "openai", + "cost_provider": "openai", + "model_info": ModelInfo(input_cost_per_token_batches=INPUT_RATE, output_cost_per_token_batches=OUTPUT_RATE), + "metadata": line_item_metadata( + {"user_api_key_hash": ALICE_KEY_HASH, "user_api_key_team_id": "team-alpha", "tags": ["env:prod"]} + ), + } + return BatchLineItemSource(**{**fields, **overrides}) + + +def _whole_batch(source: BatchLineItemSource, input_content: bytes, result_contents: tuple[bytes, ...]): + """Every line of a batch small enough for one part.""" + (payload,) = batch_line_item_parts(source, input_content, result_contents, BATCH_LINE_ITEMS_MAX_PART_BYTES) + return payload + + +def _payload() -> StandardBatchLineItemsPayload: + return _whole_batch( + _source(), + _jsonl(_request("ok-1", "first question"), _request("bad-1", "second question")), + (_jsonl(_success("ok-1", 100, 50)), _jsonl(_error("bad-1"))), + ) + + +class _RecordingLogger(CustomLogger): + def __init__(self, **kwargs) -> None: + super().__init__(**kwargs) + self.received: list[StandardBatchLineItemsPayload] = [] + + async def async_log_batch_line_items(self, payload: StandardBatchLineItemsPayload) -> None: + self.received.append(payload) + + +class _FailingLogger(CustomLogger): + async def async_log_batch_line_items(self, payload: StandardBatchLineItemsPayload) -> None: + raise RuntimeError("bucket is gone") + + +@pytest.mark.asyncio +async def test_each_result_line_is_paired_with_the_request_that_shares_its_custom_id(): + payload = _whole_batch( + _source(), + _jsonl(_request("ok-1", "first"), _request("bad-1", "second"), _request("never-ran", "third")), + (_jsonl(_success("ok-1", 100, 50)), _jsonl(_error("bad-1"))), + ) + + assert payload["batch_id"] == "bGl0ZWxsbV9wcm94eTtiYXRjaA" + assert payload["finished_at"] == FINISHED_AT + succeeded, failed = payload["lines"] + assert succeeded == { + "batch_id": "bGl0ZWxsbV9wcm94eTtiYXRjaA", + "custom_id": "ok-1", + "status": "success", + "request": _request("ok-1", "first"), + "response": _success("ok-1", 100, 50), + "usage": {"prompt_tokens": 100, "completion_tokens": 50, "total_tokens": 150}, + "response_cost": pytest.approx(100 * INPUT_RATE + 50 * OUTPUT_RATE), + "model": "gpt-4o-mini", + "custom_llm_provider": "openai", + "metadata": { + "user_api_key_hash": ALICE_KEY_HASH, + "user_api_key_alias": None, + "user_api_key_team_id": "team-alpha", + "user_api_key_team_alias": None, + "user_api_key_user_id": None, + "user_api_key_org_id": None, + }, + } + assert failed["custom_id"] == "bad-1" + assert failed["status"] == "failure" + assert failed["request"] == _request("bad-1", "second") + assert failed["response"] == _error("bad-1") + assert failed["usage"] is None + assert failed["response_cost"] == 0 + + +@pytest.mark.asyncio +async def test_line_costs_add_up_to_the_batch_cost_the_spend_row_records(): + outputs = _jsonl(_success("a", 100, 50), _success("b", 7, 3), _error("c")) + + payload = _whole_batch(_source(), b"", (outputs,)) + aggregate = await calculate_batch_cost_and_usage( + get_file_content_as_dictionary(outputs), "openai", model_name="gpt-4o-mini", model_info=_source().model_info + ) + + assert [line["response_cost"] for line in payload["lines"]] == pytest.approx( + [100 * INPUT_RATE + 50 * OUTPUT_RATE, 7 * INPUT_RATE + 3 * OUTPUT_RATE, 0] + ) + assert sum(line["response_cost"] for line in payload["lines"]) == pytest.approx(aggregate.cost) + + +@pytest.mark.asyncio +async def test_anthropic_results_without_an_input_file_keep_their_status_and_usage(): + succeeded = { + "custom_id": "a-1", + "result": { + "type": "succeeded", + "message": {"model": "claude-sonnet-5-5", "usage": {"input_tokens": 10, "output_tokens": 4}}, + }, + } + errored = {"custom_id": "a-2", "result": {"type": "errored", "error": {"type": "invalid_request_error"}}} + + payload = _whole_batch( + _source(model="claude-sonnet-5-5", custom_llm_provider="anthropic", cost_provider="anthropic"), + b"", + (_jsonl(succeeded, errored),), + ) + + first, second = payload["lines"] + assert (first["custom_id"], first["status"], first["request"]) == ("a-1", "success", None) + assert first["usage"] == {"prompt_tokens": 10, "completion_tokens": 4, "total_tokens": 14} + assert (second["custom_id"], second["status"], second["usage"]) == ("a-2", "failure", None) + + +@pytest.mark.asyncio +async def test_bedrock_lines_pair_by_record_id(): + request = {"recordId": "r-1", "modelInput": {"messages": [{"role": "user", "content": "hi"}]}} + result = {**request, "modelOutput": {"usage": {"input_tokens": 3, "output_tokens": 2}}} + + payload = _whole_batch( + _source(custom_llm_provider="bedrock", cost_provider="bedrock"), _jsonl(request), (_jsonl(result),) + ) + + (line,) = payload["lines"] + assert (line["custom_id"], line["status"], line["request"]) == ("r-1", "success", request) + + +def test_vertex_lines_pair_with_the_input_row_that_carries_their_custom_id(): + vertex_input = b"\n".join( + _OpenAIToVertexBatchUploadStream( + _jsonl(_request("vx-0", "first"), _request("vx-1", "second")), + VertexAIFilesConfig()._map_openai_to_vertex_params, + ).iter_bytes() + ) + vertex_rows = get_file_content_as_dictionary(vertex_input) + assert all("custom_id" not in row for row in vertex_rows) + + payload = _whole_batch( + _source(model="gemini-2.5-flash", custom_llm_provider="vertex_ai", cost_provider="vertex_ai"), + vertex_input, + (_jsonl(_success("vx-1", 4, 2), _success("vx-0", 3, 1)),), + ) + + assert [(line["custom_id"], line["request"]) for line in payload["lines"]] == [ + ("vx-1", vertex_rows[1]), + ("vx-0", vertex_rows[0]), + ] + + +def test_a_vertex_embeddings_entry_split_into_several_rows_pairs_with_none_of_them(): + embeddings = { + "custom_id": "emb-1", + "method": "POST", + "url": "/v1/embeddings", + "body": {"model": "gemini-embedding-001", "input": ["first", "second"]}, + } + vertex_input = b"\n".join( + _OpenAIToVertexBatchUploadStream( + _jsonl(embeddings, _request("vx-0", "chat")), VertexAIFilesConfig()._map_openai_to_vertex_params + ).iter_bytes() + ) + vertex_rows = get_file_content_as_dictionary(vertex_input) + assert [row.get("key") for row in vertex_rows] == ["emb-1#0/2", "emb-1#1/2", None] + + payload = _whole_batch( + _source(model="gemini-2.5-flash", custom_llm_provider="vertex_ai", cost_provider="vertex_ai"), + vertex_input, + (_jsonl(_success("emb-1", 2, 0), _success("vx-0", 3, 1)),), + ) + + assert [(line["custom_id"], line["request"]) for line in payload["lines"]] == [ + ("emb-1", None), + ("vx-0", vertex_rows[2]), + ] + + +def test_a_vertex_input_row_without_a_custom_id_pairs_with_nothing(): + unlabeled = {"request": {"contents": [{"role": "user", "parts": [{"text": "hi"}]}]}} + + payload = _whole_batch( + _source(model="gemini-2.5-flash", custom_llm_provider="vertex_ai", cost_provider="vertex_ai"), + _jsonl(unlabeled), + (_jsonl({**_success("x", 3, 1), "custom_id": None}),), + ) + + (line,) = payload["lines"] + assert (line["custom_id"], line["request"]) == (None, None) + + +@pytest.mark.parametrize( + "vertex_input", + [ + pytest.param( + b"\n".join( + _OpenAIToVertexBatchUploadStream( + _jsonl(_request("unknown", "first"), {**_request("vx-1", "second"), "custom_id": None}), + VertexAIFilesConfig()._map_openai_to_vertex_params, + ).iter_bytes() + ), + id="custom-id-unknown-next-to-a-null-one", + ), + pytest.param( + _jsonl( + {"request": {"contents": [{"role": "user", "parts": [{"text": "first"}]}], "labels": {"team": "a"}}}, + {"request": {"contents": [{"role": "user", "parts": [{"text": "second"}]}]}}, + ), + id="native-rows-with-other-labels-next-to-unlabeled-ones", + ), + ], +) +def test_a_vertex_line_read_as_unknown_pairs_with_no_input_row(vertex_input): + payload = _whole_batch( + _source(model="gemini-2.5-flash", custom_llm_provider="vertex_ai", cost_provider="vertex_ai"), + vertex_input, + (_jsonl(_success("unknown", 3, 1), _success("unknown", 4, 2)),), + ) + + assert [(line["custom_id"], line["request"]) for line in payload["lines"]] == [("unknown", None)] * 2 + + +@pytest.mark.asyncio +async def test_redaction_drops_only_the_request_and_response(): + payload = _payload() + + redacted = without_request_and_response(payload) + + assert redacted["batch_id"] == payload["batch_id"] + assert redacted["finished_at"] == payload["finished_at"] + assert redacted["part"] == payload["part"] + assert [{**line, "request": None, "response": None} for line in payload["lines"]] == list(redacted["lines"]) + assert all(line["request"] is not None for line in payload["lines"]) + + +@pytest.mark.asyncio +async def test_the_object_holds_one_json_line_per_batch_line_under_its_finished_day(): + payload = _payload() + + ndjson = batch_line_items_ndjson(payload) + + assert [json.loads(line) for line in ndjson.splitlines()] == list(payload["lines"]) + assert ndjson.endswith("}\n") + assert ( + batch_line_items_object_name(payload) + == "batch_line_items/2026-10-08/bGl0ZWxsbV9wcm94eTtiYXRjaA/part-00000.jsonl" + ) + assert ( + batch_line_items_object_name( + {**payload, "batch_id": "projects/p/locations/l/batchPredictionJobs/9", "part": 12} + ) + == "batch_line_items/2026-10-08/projects_p_locations_l_batchPredictionJobs_9/part-00012.jsonl" + ) + + +def _raw_size(request: dict, result: dict) -> int: + return len(json.dumps(request)) + len(json.dumps(result)) + + +@pytest.mark.parametrize( + "extra_bytes, parts", + [(0, [["a", "b"], ["c"]]), (1, [["a", "b", "c"]])], + ids=["closes once its raw lines reach the size", "stays open below the size"], +) +def test_a_part_closes_once_its_raw_request_and_result_lines_reach_the_part_size(extra_bytes, parts): + requests = [_request(custom_id, f"question {custom_id}") for custom_id in "abc"] + results = [_success(custom_id, 10, 5) for custom_id in "abc"] + max_part_bytes = _raw_size(requests[0], results[0]) + _raw_size(requests[1], results[1]) + extra_bytes + + payloads = list(batch_line_item_parts(_source(), _jsonl(*requests), (_jsonl(*results),), max_part_bytes)) + + assert [payload["part"] for payload in payloads] == list(range(len(parts))) + assert [[line["custom_id"] for line in payload["lines"]] for payload in payloads] == parts + assert [line["request"] for payload in payloads for line in payload["lines"]] == requests + + +def test_blank_and_unreadable_input_lines_leave_the_other_lines_paired(): + input_content = b"\n".join( + [json.dumps(_request("a", "first")).encode(), b"", b"not json", json.dumps(_request("b", "second")).encode()] + ) + + payload = _whole_batch(_source(), input_content, (_jsonl(_success("a", 1, 1), _success("b", 1, 1)),)) + + assert [line["request"] for line in payload["lines"]] == [_request("a", "first"), _request("b", "second")] + + +def test_files_with_no_readable_line_make_no_part(): + assert list(batch_line_item_parts(_source(), b"", (b"not json\n", b""), BATCH_LINE_ITEMS_MAX_PART_BYTES)) == [] + + +@pytest.mark.parametrize( + "saved, recorded", + [ + (ALICE_KEY_HASH, ALICE_KEY_HASH), + ("ya29.opaque-oauth2-access-token", sha256(b"ya29.opaque-oauth2-access-token").hexdigest()), + ("litellm_proxy_master_key", sha256(b"litellm_proxy_master_key").hexdigest()), + (None, None), + ], + ids=["virtual key hash", "oauth2 token", "master key alias", "nothing saved"], +) +def test_the_creator_credential_is_recorded_only_as_a_hash(saved, recorded): + assert line_item_metadata({"user_api_key_hash": saved})["user_api_key_hash"] == recorded + + +def test_a_line_holding_a_lone_surrogate_is_escaped_so_the_object_still_encodes_as_utf8(): + payload = _payload() + (line, *_) = payload["lines"] + surrogate_line = {**line, "custom_id": "emoji-\ud83d", "response": {"body": {"content": "caf\u00e9 \ud83d"}}} + readable_line = {**line, "custom_id": "caf\u00e9"} + + ndjson = batch_line_items_ndjson({**payload, "lines": (surrogate_line, readable_line)}) + + surrogate_text, readable_text = ndjson.encode("utf-8").decode("utf-8").splitlines() + assert json.loads(surrogate_text) == surrogate_line + assert "\\ud83d" in surrogate_text + assert json.loads(readable_text) == readable_line + assert '"custom_id": "caf\u00e9"' in readable_text + + +@pytest.mark.asyncio +async def test_a_logger_gets_the_request_and_response_unless_message_logging_is_off(monkeypatch): + payload = _payload() + open_logger, quiet_logger, legacy_logger = ( + _RecordingLogger(), + _RecordingLogger(turn_off_message_logging=True), + _RecordingLogger(message_logging=False), + ) + monkeypatch.setattr(litellm, "batch_line_item_callbacks", (open_logger, quiet_logger, legacy_logger)) + + await dispatch_batch_line_items(payload, redact_bodies=False) + + assert open_logger.received == [payload] + assert quiet_logger.received == [without_request_and_response(payload)] + assert legacy_logger.received == [without_request_and_response(payload)] + + +@pytest.mark.asyncio +async def test_redacted_dispatch_never_hands_a_logger_the_bodies(monkeypatch): + payload = _payload() + logger = _RecordingLogger() + monkeypatch.setattr(litellm, "batch_line_item_callbacks", (logger,)) + + await dispatch_batch_line_items(payload, redact_bodies=True) + + assert logger.received == [without_request_and_response(payload)] + + +@pytest.mark.asyncio +async def test_a_failing_or_unknown_sink_is_logged_and_the_other_sinks_still_get_the_lines(monkeypatch, caplog): + payload = _payload() + logger = _RecordingLogger() + monkeypatch.setattr(litellm, "batch_line_item_callbacks", (_FailingLogger(), "langfuse", logger)) + + with caplog.at_level(logging.ERROR, logger="LiteLLM"): + await dispatch_batch_line_items(payload, redact_bodies=False) + + assert logger.received == [payload] + failures = [record.getMessage() for record in caplog.records if record.levelno == logging.ERROR] + assert len(failures) == 2 + assert all(payload["batch_id"] in message for message in failures) + assert any("bucket is gone" in message for message in failures) + assert any("'langfuse'" in message for message in failures) diff --git a/tests/unit/integrations/gcs_bucket/test_gcs_bucket.py b/tests/unit/integrations/gcs_bucket/test_gcs_bucket.py index 58b0f26ac3b..6f143a244f6 100644 --- a/tests/unit/integrations/gcs_bucket/test_gcs_bucket.py +++ b/tests/unit/integrations/gcs_bucket/test_gcs_bucket.py @@ -330,3 +330,67 @@ async def test_unbatched_gcs_logs_upload_individual_success_and_failure_objects( flush_task.cancel() with pytest.raises(asyncio.CancelledError): await flush_task + + +@pytest.mark.asyncio +async def test_batch_line_items_are_one_ndjson_object_in_the_default_bucket(monkeypatch: pytest.MonkeyPatch) -> None: + logger, uploads = _gcs_logger_with_upload_status(monkeypatch, 200) + + def stored(request: httpx.Request) -> httpx.Response: + uploads(request=request) + return httpx.Response(200, json={"name": request.url.params["name"]}) + + logger.async_httpx_client = AsyncHTTPHandler(transport=httpx.MockTransport(stored)) + lines: Final = ( + {"batch_id": "bGl0ZWxsbV9wcm94eTtiYXRjaA", "custom_id": "ok-1", "status": "success", "response_cost": 0.5}, + {"batch_id": "bGl0ZWxsbV9wcm94eTtiYXRjaA", "custom_id": "bad-1", "status": "failure", "response_cost": 0.0}, + ) + + await logger.async_log_batch_line_items( + {"batch_id": "bGl0ZWxsbV9wcm94eTtiYXRjaA", "finished_at": 1791494130, "part": 3, "lines": lines} + ) + + request: Final = uploads.call_args.kwargs["request"] + assert uploads.call_count == 1 + assert request.url.path == "/upload/storage/v1/b/test-bucket/o" + assert request.url.params["name"] == "batch_line_items/2026-10-08/bGl0ZWxsbV9wcm94eTtiYXRjaA/part-00003.jsonl" + assert request.headers["Authorization"] == "Bearer test-access-token" + assert [json.loads(line) for line in request.content.decode().splitlines()] == list(lines) + + +@pytest.mark.asyncio +async def test_a_rejected_batch_line_items_write_raises_for_the_dispatcher(monkeypatch: pytest.MonkeyPatch) -> None: + logger, uploads = _gcs_logger_with_upload_status(monkeypatch, 503) + + with pytest.raises(httpx.HTTPStatusError): + await logger.async_log_batch_line_items( + {"batch_id": "b", "finished_at": 1791494130, "part": 0, "lines": ({"custom_id": "ok-1"},)} + ) + + assert uploads.call_count == 1 + + +@pytest.mark.asyncio +async def test_batch_line_items_are_serialized_off_the_event_loop(monkeypatch: pytest.MonkeyPatch) -> None: + import threading + + import litellm.integrations.gcs_bucket.gcs_bucket as gcs_bucket_module + from litellm.types.utils import StandardBatchLineItemsPayload + + logger, _ = _gcs_logger_with_upload_status(monkeypatch, 200) + logger.async_httpx_client = AsyncHTTPHandler( + transport=httpx.MockTransport(lambda request: httpx.Response(200, json={"name": request.url.params["name"]})) + ) + serialize: Final = gcs_bucket_module.batch_line_items_ndjson + threads: Final[list[threading.Thread]] = [] + + def _serialize(payload: StandardBatchLineItemsPayload) -> str: + threads.append(threading.current_thread()) + return serialize(payload) + + monkeypatch.setattr(gcs_bucket_module, "batch_line_items_ndjson", _serialize) + + await logger.async_log_batch_line_items({"batch_id": "b", "finished_at": 1791494130, "part": 0, "lines": ()}) + + (thread,) = threads + assert thread is not threading.main_thread() diff --git a/tests/unit/integrations/test_s3_v2.py b/tests/unit/integrations/test_s3_v2.py index fcd2c5673af..80f79517fab 100644 --- a/tests/unit/integrations/test_s3_v2.py +++ b/tests/unit/integrations/test_s3_v2.py @@ -5848,3 +5848,126 @@ class TestS3Logger(S3Logger): self.recorded_requests[response_obj["id"]] = start_time self.logged_standard_logging_payload = kwargs["standard_logging_object"] return await super().async_log_success_event(kwargs, response_obj, start_time, end_time) + + +def _batch_line_items_payload() -> dict: + line = { + "batch_id": "bGl0ZWxsbV9wcm94eTtiYXRjaA", + "custom_id": "ok-1", + "status": "success", + "request": {"custom_id": "ok-1", "body": {"messages": [{"role": "user", "content": "hi"}]}}, + "response": {"custom_id": "ok-1", "response": {"status_code": 200, "body": {"id": "chatcmpl-1"}}}, + "usage": {"prompt_tokens": 3, "completion_tokens": 1, "total_tokens": 4}, + "response_cost": 0.005, + "model": "gpt-4o-mini", + "custom_llm_provider": "openai", + "metadata": {"user_api_key_hash": "hash-alice"}, + } + return { + "batch_id": "bGl0ZWxsbV9wcm94eTtiYXRjaA", + "finished_at": 1791494130, + "part": 0, + "lines": (line, {**line, "custom_id": "ok-2"}), + } + + +def _s3_logger_answering( + status: int, s3_path: str | None, content: bytes = b"" +) -> tuple[S3Logger, list[httpx.Request]]: + requests: list[httpx.Request] = [] + + def respond(request: httpx.Request) -> httpx.Response: + requests.append(request) + return httpx.Response(status, content=content, request=request) + + logger = S3Logger( + s3_bucket_name="test-bucket", + s3_region_name="us-east-1", + s3_path=s3_path, + s3_aws_access_key_id="AKIAIOSFODNN7EXAMPLE", + s3_aws_secret_access_key="wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY", + s3_flush_interval=3600, + ) + logger.async_httpx_client = AsyncHTTPHandler(transport=httpx.MockTransport(respond)) + return logger, requests + + +@pytest.mark.asyncio +async def test_batch_line_items_are_one_ndjson_object_under_the_configured_path(): + logger, requests = _s3_logger_answering(200, "cold/") + payload = _batch_line_items_payload() + + await logger.async_log_batch_line_items(payload) + + (request,) = requests + assert request.method == "PUT" + assert request.url.path == "/cold/batch_line_items/2026-10-08/bGl0ZWxsbV9wcm94eTtiYXRjaA/part-00000.jsonl" + assert request.headers["Content-Type"] == "application/x-ndjson" + assert request.headers["Content-Disposition"] == 'inline; filename="bGl0ZWxsbV9wcm94eTtiYXRjaA-part-00000.jsonl"' + assert [json.loads(line) for line in request.content.decode().splitlines()] == list(payload["lines"]) + assert logger.log_queue == [] + + +@pytest.mark.asyncio +async def test_batch_line_items_without_an_s3_path_sit_at_the_bucket_root(): + logger, requests = _s3_logger_answering(200, None) + + await logger.async_log_batch_line_items(_batch_line_items_payload()) + + assert [request.url.path for request in requests] == [ + "/batch_line_items/2026-10-08/bGl0ZWxsbV9wcm94eTtiYXRjaA/part-00000.jsonl" + ] + + +@pytest.mark.asyncio +async def test_a_rejected_batch_line_items_upload_is_logged_with_the_batch_id(caplog): + logger, requests = _s3_logger_answering(400, "cold") + + with caplog.at_level(logging.ERROR, logger="LiteLLM"): + await logger.async_log_batch_line_items(_batch_line_items_payload()) + + assert len(requests) == 1 + assert any( + "bGl0ZWxsbV9wcm94eTtiYXRjaA" in record.getMessage() and "were not stored" in record.getMessage() + for record in caplog.records + ) + + +@pytest.mark.asyncio +async def test_a_batch_line_items_object_the_bucket_refuses_for_good_is_not_kept_in_memory(): + logger, requests = _s3_logger_answering(400, "cold", b"InvalidArgument") + + await logger.async_log_batch_line_items(_batch_line_items_payload()) + await logger.flush_queue() + + assert len(requests) == 1 + assert logger._flush_dropped == {} + + +@pytest.mark.asyncio +async def test_batch_line_items_are_serialized_and_hashed_off_the_event_loop(monkeypatch): + import threading + + import litellm.integrations.s3_v2 as s3_v2_module + + logger, requests = _s3_logger_answering(200, "cold/") + serialize = s3_v2_module.batch_line_items_ndjson + prepare = logger._prepare_put + threads = {} + + def _serialize(payload): + threads["serialize"] = threading.current_thread() + return serialize(payload) + + def _prepare(element): + threads["hash"] = threading.current_thread() + return prepare(element) + + monkeypatch.setattr(s3_v2_module, "batch_line_items_ndjson", _serialize) + monkeypatch.setattr(logger, "_prepare_put", _prepare) + + await logger.async_log_batch_line_items(_batch_line_items_payload()) + + assert len(requests) == 1 + assert set(threads) == {"serialize", "hash"} + assert threading.main_thread() not in threads.values() diff --git a/tests/unit/proxy/common_utils/test_check_batch_cost.py b/tests/unit/proxy/common_utils/test_check_batch_cost.py index 3b9e7a30302..a5ee7fe23bb 100644 --- a/tests/unit/proxy/common_utils/test_check_batch_cost.py +++ b/tests/unit/proxy/common_utils/test_check_batch_cost.py @@ -8,7 +8,9 @@ ARN unified_object_id) batches with no managed unified id. import asyncio import json +import logging from contextlib import contextmanager +from hashlib import sha256 from typing import TYPE_CHECKING from unittest.mock import AsyncMock, MagicMock, patch @@ -3242,3 +3244,703 @@ class TestMultiPodBatchCostClaim: assert self._claim_calls(prisma) == [] assert journal == ["fetch", "bill", "finalize"] logging_obj.async_success_handler.assert_awaited_once() + + +class TestBatchLineItemsFromThePoller: + """When batch_line_item_callbacks is set, the poller hands every line of a finished batch to those loggers + once the batch is billed and retired, in the background, so a slow, hung or failing sink never unbills, re-polls + or holds up a batch.""" + + _INPUT = b'{"custom_id":"ok-1","body":{"messages":[{"role":"user","content":"hi"}]}}\n' + _OUTPUT = ( + b'{"custom_id":"ok-1","response":{"status_code":200,"body":{"model":"gpt-4",' + b'"usage":{"prompt_tokens":10,"completion_tokens":5,"total_tokens":15}}}}\n' + ) + _ERROR = b'{"custom_id":"ok-1","response":{"status_code":400,"body":{"error":{"message":"expired"}}}}\n' + + @staticmethod + def _recording_logger(journal: list, error: Exception | None = None): + from litellm.integrations.custom_logger import CustomLogger + + class _Recorder(CustomLogger): + def __init__(self): + super().__init__() + self.payloads = [] + + async def async_log_batch_line_items(self, payload): + journal.append("lines") + self.payloads.append(payload) + if error is not None: + raise error + + return _Recorder() + + @staticmethod + def _setup(journal: list, *, status="completed", output_file_id=_CLAIM_OUTPUT_FILE_ID, error_file_id=None): + row = _FakeManagedObjectRow() + prisma = TestMultiPodBatchCostClaim._prisma(row, journal) + router = TestMultiPodBatchCostClaim._router() + response = router.aretrieve_batch.return_value + response.status = status + response.output_file_id = output_file_id + response.error_file_id = error_file_id + response.input_file_id = "file-input-1" + return row, TestMultiPodBatchCostClaim._instance(prisma, router) + + @staticmethod + async def _stored(instance) -> None: + """Wait for the line items the poll handed to the background.""" + await asyncio.gather(*instance._line_item_stores) + + @classmethod + @contextmanager + def _provider(cls, journal: list): + """The provider's batch files by id, and a spend-log writer that journals each write.""" + files = {"file-input-1": cls._INPUT, _CLAIM_OUTPUT_FILE_ID: cls._OUTPUT, "file-err-1": cls._ERROR} + + async def _afile_content(**kwargs): + journal.append(f"read {kwargs['file_id']}") + if kwargs["file_id"] not in files: + from litellm.exceptions import NotFoundError + + raise NotFoundError( + message=f"No file found with id '{kwargs['file_id']}'.", model="gpt-4", llm_provider="openai" + ) + content = MagicMock() + content.content = files[kwargs["file_id"]] + return content + + def _is_b64(file_id): + if file_id == _CLAIM_UNIFIED_BATCH_ID: + return "llm_model_id,model-123;llm_batch_id,batch-456;" + return False + + logging_obj = MagicMock() + logging_obj.async_success_handler = AsyncMock(side_effect=lambda **kwargs: journal.append("bill")) + with ( + patch(_IS_B64, side_effect=_is_b64), + patch( + "litellm.proxy.openai_files_endpoints.common_utils.get_model_id_from_unified_batch_id", + return_value="model-123", + ), + patch( + "litellm.proxy.openai_files_endpoints.common_utils.get_batch_id_from_unified_batch_id", + return_value="batch-456", + ), + patch("litellm.files.main.afile_content", new=AsyncMock(side_effect=_afile_content)), + patch("litellm.litellm_core_utils.litellm_logging.Logging", return_value=logging_obj), + ): + yield + + @pytest.mark.asyncio + async def test_lines_reach_the_callbacks_after_the_claimed_batch_is_billed_and_retired(self, monkeypatch): + import litellm + + journal: list = [] + row, instance = self._setup(journal) + logger = self._recording_logger(journal) + monkeypatch.setattr(litellm, "batch_line_item_callbacks", (logger,)) + + with self._provider(journal): + await instance.check_batch_cost() + await self._stored(instance) + + assert journal == [ + f"read {_CLAIM_OUTPUT_FILE_ID}", + "claim", + "bill", + "finalize", + f"read {_CLAIM_OUTPUT_FILE_ID}", + "read file-input-1", + "lines", + ] + (payload,) = logger.payloads + assert payload["batch_id"] == _CLAIM_UNIFIED_BATCH_ID + assert payload["finished_at"] == 2 + assert payload["part"] == 0 + (line,) = payload["lines"] + assert (line["custom_id"], line["status"]) == ("ok-1", "success") + assert line["usage"] == {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15} + assert line["metadata"]["user_api_key_user_id"] == "user-1" + assert row.batch_processed is True + + @pytest.mark.asyncio + async def test_the_lines_are_priced_off_the_event_loop(self, monkeypatch): + import threading + + import litellm + from litellm.batches import batch_line_items + + journal: list = [] + row, instance = self._setup(journal) + monkeypatch.setattr(litellm, "batch_line_item_callbacks", (self._recording_logger(journal),)) + price = batch_line_items._line_item + threads = [] + + def _price(*args): + threads.append(threading.current_thread()) + return price(*args) + + monkeypatch.setattr(batch_line_items, "_line_item", _price) + with self._provider(journal): + await instance.check_batch_cost() + await self._stored(instance) + + assert journal[-1] == "lines" + (thread,) = threads + assert thread is not threading.main_thread() + + @pytest.mark.asyncio + async def test_a_creator_that_cannot_be_looked_up_gets_lines_without_bodies(self, monkeypatch): + import litellm + + journal: list = [] + row, instance = self._setup(journal) + row.api_key = sha256(b"sk-deleted-key").hexdigest() + logger = self._recording_logger(journal) + monkeypatch.setattr(litellm, "batch_line_item_callbacks", (logger,)) + + with self._provider(journal): + await instance.check_batch_cost() + await self._stored(instance) + + (payload,) = logger.payloads + assert [(line["request"], line["response"]) for line in payload["lines"]] == [(None, None)] + assert "read file-input-1" not in journal + + @pytest.mark.asyncio + async def test_a_creator_with_logging_on_gets_the_request_and_response(self, monkeypatch): + import litellm + + journal: list = [] + row, instance = self._setup(journal) + row.api_key = "cli-session-user-1" + logger = self._recording_logger(journal) + monkeypatch.setattr(litellm, "batch_line_item_callbacks", (logger,)) + monkeypatch.setattr(litellm, "turn_off_message_logging", False) + + with self._provider(journal): + await instance.check_batch_cost() + await self._stored(instance) + + assert journal[-4:] == ["finalize", f"read {_CLAIM_OUTPUT_FILE_ID}", "read file-input-1", "lines"] + ((line,),) = [payload["lines"] for payload in logger.payloads] + assert line["request"] == json.loads(self._INPUT) + assert line["response"] == json.loads(self._OUTPUT) + + @pytest.mark.asyncio + async def test_a_failing_sink_leaves_the_batch_billed_and_retired(self, monkeypatch): + import litellm + + journal: list = [] + row, instance = self._setup(journal) + monkeypatch.setattr( + litellm, "batch_line_item_callbacks", (self._recording_logger(journal, RuntimeError("bucket gone")),) + ) + + with self._provider(journal): + await instance.check_batch_cost() + await self._stored(instance) + + assert [step for step in journal if not step.startswith("read")] == ["claim", "bill", "finalize", "lines"] + assert row.batch_processed is True + assert row.status == "complete" + + @pytest.mark.asyncio + async def test_without_callbacks_only_the_output_file_is_read(self, monkeypatch): + import litellm + + journal: list = [] + _, instance = self._setup(journal) + monkeypatch.setattr(litellm, "batch_line_item_callbacks", ()) + + with self._provider(journal): + await instance.check_batch_cost() + await self._stored(instance) + + assert journal == [f"read {_CLAIM_OUTPUT_FILE_ID}", "claim", "bill", "finalize"] + + @pytest.mark.asyncio + async def test_a_batch_where_every_line_failed_sends_its_error_lines_after_it_is_retired(self, monkeypatch): + import litellm + + journal: list = [] + row, instance = self._setup(journal, status="expired", output_file_id=None, error_file_id="file-err-1") + logger = self._recording_logger(journal) + monkeypatch.setattr(litellm, "batch_line_item_callbacks", (logger,)) + + with self._provider(journal): + await instance.check_batch_cost() + await self._stored(instance) + + assert journal == ["finalize", "read file-err-1", "read file-input-1", "lines"] + (payload,) = logger.payloads + assert [(line["custom_id"], line["status"], line["response_cost"]) for line in payload["lines"]] == [ + ("ok-1", "failure", 0.0) + ] + assert row.batch_processed is True + + @pytest.mark.asyncio + async def test_a_batch_retired_with_nothing_billable_reads_its_input_file_by_the_id_the_provider_reported( + self, monkeypatch + ): + import litellm + + journal: list = [] + _, instance = self._setup(journal, status="expired", output_file_id=None, error_file_id="file-err-1") + logger = self._recording_logger(journal) + monkeypatch.setattr(litellm, "batch_line_item_callbacks", (logger,)) + + async def _persist_managed_ids(*, response, **_): + response.input_file_id = "managed-id-of-the-copy-on-another-deployment" + + with ( + self._provider(journal), + patch( + "litellm.proxy.openai_files_endpoints.common_utils.ensure_batch_response_managed_file_ids", + new=AsyncMock(side_effect=_persist_managed_ids), + ), + ): + await instance.check_batch_cost() + await self._stored(instance) + + assert journal == ["finalize", "read file-err-1", "read file-input-1", "lines"] + ((line,),) = [payload["lines"] for payload in logger.payloads] + assert line["request"] == json.loads(self._INPUT) + + @pytest.mark.asyncio + async def test_a_batch_whose_files_hold_no_readable_line_sends_no_lines(self, monkeypatch): + import litellm + + journal: list = [] + _, instance = self._setup(journal, status="expired", output_file_id=None, error_file_id="file-err-1") + logger = self._recording_logger(journal) + monkeypatch.setattr(litellm, "batch_line_item_callbacks", (logger,)) + monkeypatch.setattr(type(self), "_ERROR", b"not json\n") + + with self._provider(journal): + await instance.check_batch_cost() + await self._stored(instance) + + assert journal == ["finalize", "read file-err-1", "read file-input-1"] + assert logger.payloads == [] + + @pytest.mark.asyncio + async def test_a_batch_that_could_not_be_retired_sends_no_lines_until_it_is(self, monkeypatch): + import litellm + + journal: list = [] + row, instance = self._setup(journal, status="expired", output_file_id=None, error_file_id="file-err-1") + instance.prisma_client.db.litellm_managedobjecttable.update = AsyncMock(side_effect=RuntimeError("db down")) + logger = self._recording_logger(journal) + monkeypatch.setattr(litellm, "batch_line_item_callbacks", (logger,)) + + with self._provider(journal): + await instance.check_batch_cost() + await self._stored(instance) + + assert journal == [] + assert logger.payloads == [] + assert row.batch_processed is False + + @pytest.mark.asyncio + @pytest.mark.parametrize("has_batch_processed_column, stored", [(True, 1), (False, 0)]) + async def test_a_billed_batch_that_could_not_be_marked_complete_sends_lines_only_if_its_claim_retired_it( + self, monkeypatch, has_batch_processed_column, stored + ): + import litellm + + journal: list = [] + row, instance = self._setup(journal) + instance._has_batch_processed_column = has_batch_processed_column + instance.prisma_client.db.litellm_managedobjecttable.update = AsyncMock(side_effect=RuntimeError("db down")) + logger = self._recording_logger(journal) + monkeypatch.setattr(litellm, "batch_line_item_callbacks", (logger,)) + + with self._provider(journal): + await instance.check_batch_cost() + await self._stored(instance) + + assert journal.count("bill") == 1 + assert len(logger.payloads) == stored + + @pytest.mark.asyncio + async def test_a_batch_whose_output_file_is_gone_still_sends_its_error_lines_after_it_is_retired( + self, monkeypatch + ): + import litellm + + journal: list = [] + row, instance = self._setup(journal, output_file_id="file-out-deleted", error_file_id="file-err-1") + logger = self._recording_logger(journal) + monkeypatch.setattr(litellm, "batch_line_item_callbacks", (logger,)) + + with self._provider(journal): + await instance.check_batch_cost() + await self._stored(instance) + + assert journal == ["read file-out-deleted", "finalize", "read file-err-1", "read file-input-1", "lines"] + (payload,) = logger.payloads + assert [(line["custom_id"], line["status"]) for line in payload["lines"]] == [("ok-1", "failure")] + assert row.batch_processed is True + + @pytest.mark.asyncio + async def test_a_terminal_batch_with_no_lines_sends_nothing(self, monkeypatch): + import litellm + + journal: list = [] + _, instance = self._setup(journal, status="failed", output_file_id=None, error_file_id=None) + logger = self._recording_logger(journal) + monkeypatch.setattr(litellm, "batch_line_item_callbacks", (logger,)) + + with self._provider(journal): + await instance.check_batch_cost() + await self._stored(instance) + + assert journal == ["finalize"] + assert logger.payloads == [] + + @pytest.mark.asyncio + async def test_a_hung_sink_never_holds_up_billing_and_retiring_the_batch(self, monkeypatch): + import litellm + from litellm.integrations.custom_logger import CustomLogger + + journal: list = [] + row, instance = self._setup(journal) + released = asyncio.Event() + + class _HungSink(CustomLogger): + async def async_log_batch_line_items(self, payload): + journal.append("lines started") + await released.wait() + journal.append("lines") + + monkeypatch.setattr(litellm, "batch_line_item_callbacks", (_HungSink(),)) + + async def _sink_started(): + while "lines started" not in journal: + await asyncio.sleep(0) + + with self._provider(journal): + await asyncio.wait_for(instance.check_batch_cost(), timeout=5) + await asyncio.wait_for(_sink_started(), timeout=5) + billed_while_hung = (row.batch_processed, row.status, list(journal)) + released.set() + await self._stored(instance) + + assert billed_while_hung == ( + True, + "complete", + [ + f"read {_CLAIM_OUTPUT_FILE_ID}", + "claim", + "bill", + "finalize", + f"read {_CLAIM_OUTPUT_FILE_ID}", + "read file-input-1", + "lines started", + ], + ) + assert journal[-1] == "lines" + + @pytest.mark.asyncio + async def test_lines_are_dropped_with_an_error_while_a_poll_cycle_of_batches_is_still_waiting( + self, monkeypatch, caplog + ): + import litellm + from litellm.constants import MAX_OBJECTS_PER_POLL_CYCLE + + journal: list = [] + row, instance = self._setup(journal) + logger = self._recording_logger(journal) + monkeypatch.setattr(litellm, "batch_line_item_callbacks", (logger,)) + waiting = [asyncio.create_task(asyncio.sleep(3600)) for _ in range(MAX_OBJECTS_PER_POLL_CYCLE)] + instance._line_item_stores.update(waiting) + + try: + with self._provider(journal), caplog.at_level(logging.ERROR): + await instance.check_batch_cost() + finally: + for task in waiting: + task.cancel() + + assert logger.payloads == [] + assert row.batch_processed is True + assert f"so the line items of batch batch-456 (job {row.id}) were not stored" in caplog.text + + + @pytest.mark.asyncio + async def test_batches_store_their_lines_one_at_a_time_in_the_order_they_were_retired(self, monkeypatch): + import litellm + + journal: list = [] + _, instance = self._setup(journal) + monkeypatch.setattr(litellm, "batch_line_item_callbacks", (self._recording_logger(journal),)) + running: list = [] + + async def _store(job, response, model_id, batch_id, *, billed): + running.append(batch_id) + journal.append(f"start {batch_id} with {len(running)} running") + await asyncio.sleep(0.01) + running.remove(batch_id) + + monkeypatch.setattr(instance, "_store_batch_line_items", _store) + for batch_id in ("batch-a", "batch-b", "batch-c"): + instance._store_batch_line_items_in_background(MagicMock(), MagicMock(), "model-123", batch_id, billed=True) + await self._stored(instance) + + assert journal == [f"start batch-{name} with 1 running" for name in "abc"] + + @pytest.mark.asyncio + async def test_lines_are_stored_only_once_the_poll_has_billed_and_retired_every_batch_it_found(self, monkeypatch): + import litellm + + journal: list = [] + _, instance = self._setup(journal) + monkeypatch.setattr(litellm, "batch_line_item_callbacks", (self._recording_logger(journal),)) + + async def _an_earlier_batch_waits_to_store_its_lines(): + instance._store_batch_line_items_in_background( + MagicMock(), MagicMock(), "model-123", "batch-earlier", billed=True + ) + await asyncio.sleep(0.01) + + async def _store(job, response, model_id, batch_id, *, billed): + journal.append(f"store {batch_id}") + + monkeypatch.setattr(instance, "_cleanup_stale_managed_objects", _an_earlier_batch_waits_to_store_its_lines) + monkeypatch.setattr(instance, "_store_batch_line_items", _store) + with self._provider(journal): + await instance.check_batch_cost() + await self._stored(instance) + + assert journal == [ + f"read {_CLAIM_OUTPUT_FILE_ID}", + "claim", + "bill", + "finalize", + "store batch-earlier", + "store batch-456", + ] + + @pytest.mark.asyncio + async def test_an_output_file_that_cannot_be_read_again_leaves_the_batch_billed_and_stores_nothing( + self, monkeypatch, caplog + ): + import litellm + from litellm.batches import batch_utils + from litellm.exceptions import NotFoundError + + journal: list = [] + row, instance = self._setup(journal) + logger = self._recording_logger(journal) + monkeypatch.setattr(litellm, "batch_line_item_callbacks", (logger,)) + fetch = batch_utils.fetch_batch_managed_file_content + + async def _output_gone_after_billing(file_id, **kwargs): + if file_id == _CLAIM_OUTPUT_FILE_ID: + raise NotFoundError(message="No file found.", model="gpt-4", llm_provider="openai") + return await fetch(file_id, **kwargs) + + monkeypatch.setattr(batch_utils, "fetch_batch_managed_file_content", _output_gone_after_billing) + with self._provider(journal), caplog.at_level(logging.ERROR): + await instance.check_batch_cost() + await self._stored(instance) + + assert journal == [f"read {_CLAIM_OUTPUT_FILE_ID}", "claim", "bill", "finalize"] + assert logger.payloads == [] + assert (row.batch_processed, row.status) == (True, "complete") + assert f"the line items of batch batch-456 (job {row.id}) were not stored" in caplog.text + + @pytest.mark.asyncio + async def test_a_deployment_removed_before_the_lines_are_stored_leaves_the_batch_billed_and_logs_an_error( + self, monkeypatch, caplog + ): + import litellm + + journal: list = [] + row, instance = self._setup(journal) + logger = self._recording_logger(journal) + monkeypatch.setattr(litellm, "batch_line_item_callbacks", (logger,)) + deployment = instance.llm_router.get_deployment.return_value + instance.llm_router.get_deployment = MagicMock( + side_effect=lambda **_: None if instance._poll_idle.is_set() else deployment + ) + + with self._provider(journal), caplog.at_level(logging.ERROR): + await instance.check_batch_cost() + await self._stored(instance) + + assert journal == [f"read {_CLAIM_OUTPUT_FILE_ID}", "claim", "bill", "finalize"] + assert logger.payloads == [] + assert (row.batch_processed, row.status) == (True, "complete") + assert f"the line items of batch batch-456 (job {row.id}) were not stored" in caplog.text + + @pytest.mark.asyncio + async def test_a_batch_bigger_than_one_part_reaches_the_callbacks_as_numbered_parts(self, monkeypatch): + import litellm + from litellm_enterprise.proxy.common_utils import check_batch_cost + + journal: list = [] + _, instance = self._setup(journal) + logger = self._recording_logger(journal) + monkeypatch.setattr(litellm, "batch_line_item_callbacks", (logger,)) + monkeypatch.setattr(litellm, "turn_off_message_logging", False) + monkeypatch.setattr(check_batch_cost, "BATCH_LINE_ITEMS_MAX_PART_BYTES", 1) + monkeypatch.setattr(type(self), "_INPUT", self._INPUT + self._INPUT.replace(b"ok-1", b"ok-2")) + monkeypatch.setattr(type(self), "_OUTPUT", self._OUTPUT + self._OUTPUT.replace(b"ok-1", b"ok-2")) + + with self._provider(journal): + await instance.check_batch_cost() + await self._stored(instance) + + assert [(payload["part"], [line["custom_id"] for line in payload["lines"]]) for payload in logger.payloads] == [ + (0, ["ok-1"]), + (1, ["ok-2"]), + ] + assert [payload["lines"][0]["request"]["custom_id"] for payload in logger.payloads] == ["ok-1", "ok-2"] + + @pytest.mark.asyncio + async def test_an_oauth2_creators_token_reaches_the_callbacks_only_as_its_hash(self, monkeypatch): + import litellm + + journal: list = [] + row, instance = self._setup(journal) + row.api_key = "ya29.opaque-oauth2-access-token" + logger = self._recording_logger(journal) + monkeypatch.setattr(litellm, "batch_line_item_callbacks", (logger,)) + monkeypatch.setattr(litellm, "turn_off_message_logging", False) + + with self._provider(journal): + await instance.check_batch_cost() + await self._stored(instance) + + (payload,) = logger.payloads + (line,) = payload["lines"] + assert line["metadata"]["user_api_key_hash"] == sha256(b"ya29.opaque-oauth2-access-token").hexdigest() + assert line["request"] == json.loads(self._INPUT) + assert "ya29.opaque-oauth2-access-token" not in json.dumps(payload) + + +class TestCreatorTurnedOffMessageLogging: + """Batch line items drop the request and response when the creating key or team turned off message logging, + when the proxy did, or when the creator can no longer be read.""" + + @staticmethod + def _instance(key_row=None, team_row=None, error: Exception | None = None): + from litellm_enterprise.proxy.common_utils.check_batch_cost import CheckBatchCost + + prisma = MagicMock() + prisma.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=key_row, side_effect=error) + prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row) + return CheckBatchCost(proxy_logging_obj=MagicMock(), prisma_client=prisma, llm_router=MagicMock()) + + @staticmethod + def _job(api_key=sha256(b"sk-alice").hexdigest(), team_id=None, created_by="alice"): + from types import SimpleNamespace + + return SimpleNamespace(id="job-1", api_key=api_key, team_id=team_id, created_by=created_by) + + @staticmethod + def _logging(value: str) -> dict: + return { + "logging": [ + { + "callback_name": "s3_v2", + "callback_type": "success", + "callback_vars": {"turn_off_message_logging": value}, + } + ] + } + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("key_value", "proxy_wide", "expected"), + [("true", False, True), ("false", True, False), (None, True, True), (None, False, False)], + ) + async def test_the_key_setting_wins_and_the_proxy_setting_fills_in( + self, monkeypatch, key_value, proxy_wide, expected + ): + from types import SimpleNamespace + + import litellm + + monkeypatch.setattr(litellm, "turn_off_message_logging", proxy_wide) + metadata = self._logging(key_value) if key_value is not None else {} + instance = self._instance(key_row=SimpleNamespace(metadata=metadata)) + + assert await instance._creator_turned_off_message_logging(self._job()) is expected + + @pytest.mark.asyncio + async def test_a_team_that_turned_off_message_logging_covers_its_keys(self, monkeypatch): + from types import SimpleNamespace + + import litellm + + monkeypatch.setattr(litellm, "turn_off_message_logging", False) + instance = self._instance( + key_row=SimpleNamespace(metadata={}), team_row=SimpleNamespace(metadata=self._logging("true")) + ) + + assert await instance._creator_turned_off_message_logging(self._job(team_id="team-alpha")) is True + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("job_fields", "key_row", "team_row"), + [ + ({}, None, None), + ({"team_id": "team-gone"}, "present", None), + ({"api_key": None, "team_id": "team-gone"}, None, None), + ], + ids=["deleted-key", "deleted-team", "deleted-team-of-a-batch-created-without-a-key"], + ) + async def test_a_creator_that_is_gone_counts_as_turned_off(self, monkeypatch, job_fields, key_row, team_row): + from types import SimpleNamespace + + import litellm + + monkeypatch.setattr(litellm, "turn_off_message_logging", False) + instance = self._instance(key_row=SimpleNamespace(metadata={}) if key_row else None, team_row=team_row) + + assert await instance._creator_turned_off_message_logging(self._job(**job_fields)) is True + + @pytest.mark.asyncio + @pytest.mark.parametrize("api_key", [None, "opaque-oauth2-access-token"], ids=["jwt", "oauth2"]) + @pytest.mark.parametrize( + ("team_value", "proxy_wide", "expected"), + [(None, False, False), (None, True, True), ("true", False, True), ("false", True, False)], + ) + async def test_a_batch_created_without_a_virtual_key_follows_its_team_and_the_proxy_setting( + self, monkeypatch, api_key, team_value, proxy_wide, expected + ): + from types import SimpleNamespace + + import litellm + + monkeypatch.setattr(litellm, "turn_off_message_logging", proxy_wide) + team_row = SimpleNamespace(metadata=self._logging(team_value)) if team_value is not None else None + instance = self._instance(team_row=team_row) + job = self._job(api_key=api_key, team_id="team-alpha" if team_row is not None else None, created_by="alice") + + assert await instance._creator_turned_off_message_logging(job) is expected + instance.prisma_client.db.litellm_verificationtoken.find_unique.assert_not_called() + + @pytest.mark.asyncio + async def test_a_failed_lookup_counts_as_turned_off(self, monkeypatch): + import litellm + + monkeypatch.setattr(litellm, "turn_off_message_logging", False) + instance = self._instance(error=Exception("db down")) + + assert await instance._creator_turned_off_message_logging(self._job()) is True + + @pytest.mark.asyncio + @pytest.mark.parametrize("api_key", ["litellm_proxy_master_key", "cli-session-alice"]) + async def test_the_master_key_and_cli_sessions_have_no_key_row_to_read(self, monkeypatch, api_key): + import litellm + + monkeypatch.setattr(litellm, "turn_off_message_logging", False) + instance = self._instance() + + assert await instance._creator_turned_off_message_logging(self._job(api_key=api_key)) is False + instance.prisma_client.db.litellm_verificationtoken.find_unique.assert_not_called()