From b38b8679189e96c8f20be9d572b1b14318b402d3 Mon Sep 17 00:00:00 2001 From: shivam Date: Thu, 17 Sep 2026 22:48:28 +0000 Subject: [PATCH 01/89] feat(batches): emit per-line JSONL batch records to callbacks Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/common_utils/check_batch_cost.py | 1 + litellm/__init__.py | 1 + litellm/batches/batch_line_item_logging.py | 341 ++++++++++++++++++ litellm/litellm_core_utils/litellm_logging.py | 18 +- litellm/proxy/_types.py | 4 + .../proxy/hooks/proxy_track_cost_callback.py | 7 +- litellm/proxy/proxy_server.py | 20 + litellm/types/utils.py | 3 + .../batches/test_batch_line_item_logging.py | 233 ++++++++++++ .../hooks/test_proxy_track_cost_callback.py | 18 + 10 files changed, 643 insertions(+), 3 deletions(-) create mode 100644 litellm/batches/batch_line_item_logging.py create mode 100644 tests/test_litellm/batches/test_batch_line_item_logging.py 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 13e9e5093a80..457bcf8c1479 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -856,6 +856,7 @@ async def _track_completed_batch_cost( } }, **({"api_base": mask_api_base_credentials(deployment_api_base)} if deployment_api_base else {}), + "_litellm_internal_model_credentials": MappingProxyType({**credentials}), "metadata": { **(await self._build_creator_attribution_metadata(job, batch_id)), # spend logs read the deployment identity off these metadata keys, so diff --git a/litellm/__init__.py b/litellm/__init__.py index 71857877e532..7aaef9b91d55 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -346,6 +346,7 @@ def _dev_env_hot_reload_enabled() -> bool: ) openai_system_messages_first: bool = False disable_vertex_batch_output_transformation: bool = False +store_batch_line_items_in_callbacks: bool = False extra_spend_tag_headers: Optional[List[str]] = None in_memory_llm_clients_cache: "LLMClientCache" safe_memory_mode: bool = False diff --git a/litellm/batches/batch_line_item_logging.py b/litellm/batches/batch_line_item_logging.py new file mode 100644 index 000000000000..077457dd2995 --- /dev/null +++ b/litellm/batches/batch_line_item_logging.py @@ -0,0 +1,341 @@ +import json +import uuid +from collections.abc import Iterator, Mapping +from datetime import datetime +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, Literal, TypeAlias + +from litellm._logging import verbose_logger +from litellm.batches.batch_utils import ( + _batch_response_was_successful, # pyright: ignore[reportPrivateUsage] # batch-internal helper shared with the aggregate cost path by design + _fetch_batch_managed_file_content, # pyright: ignore[reportPrivateUsage, reportUnknownVariableType] # same reuse; helper is untyped upstream + _get_response_from_batch_job_output_file, # pyright: ignore[reportPrivateUsage] # same reuse + _iter_batch_output_entries, # pyright: ignore[reportPrivateUsage, reportUnknownVariableType] # same reuse; helper is untyped upstream + _safe_output_line_stats, # pyright: ignore[reportPrivateUsage] # same reuse +) +from litellm.types.llms.openai import ResponsesAPIResponse +from litellm.types.utils import ( + EmbeddingResponse, + LiteLLMBatch, + ModelInfo, + ModelResponse, + Usage, +) + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging + +_BatchLineProvider: TypeAlias = Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"] + +_CALL_TYPE_BY_BATCH_URL: Final = MappingProxyType( + { + "/v1/chat/completions": "acompletion", + "/v1/embeddings": "aembedding", + "/v1/responses": "aresponses", + } +) + +_EMPTY_BODY: Final[Mapping[str, object]] = MappingProxyType({}) + + +class _BatchLineFailure(Exception): + """A provider-reported per-line batch failure; carries the batch's hidden + params so the failure logging payload can attribute the line.""" + + def __init__(self, error_payload: object) -> None: + super().__init__(json.dumps(error_payload)) + self._hidden_params: dict[str, object] = {} # mutable-ok: mirrors the plain-dict _hidden_params contract on litellm response objects + + +def _as_object_mapping(value: object) -> Mapping[str, object] | None: + if isinstance(value, Mapping) and all(isinstance(key, str) for key in value): # pyright: ignore[reportUnknownVariableType] # keys of an unparameterized Mapping are unknown until checked here + return value # pyright: ignore[reportUnknownVariableType, reportReturnType] # every key was verified str above + return None + + +def _output_entries(file_content: bytes) -> Iterator[Mapping[str, object]]: + for entry in _iter_batch_output_entries(file_content): # pyright: ignore[reportUnknownVariableType] # entries are validated into typed mappings below + mapping = _as_object_mapping(entry) # pyright: ignore[reportUnknownArgumentType] # raw entry is unknown until validated here + if mapping is not None: + yield mapping + + +def _requests_by_custom_id(input_file_content: bytes) -> Mapping[str, Mapping[str, object]]: + """Parse the batch input JSONL into {custom_id: request line}, skipping + malformed lines and lines without a custom_id.""" + return MappingProxyType( + { + custom_id: entry + for entry in _output_entries(input_file_content) + if isinstance((custom_id := entry.get("custom_id")), str) and custom_id + } + ) + + +def _request_body_for_entry( + entry: Mapping[str, object], request_line: Mapping[str, object] | None +) -> Mapping[str, object]: + if request_line is not None: + body: Final = _as_object_mapping(request_line.get("body")) + if body: + return body + params: Final = _as_object_mapping(request_line.get("params")) + if params: + return params + model_input: Final = _as_object_mapping(entry.get("modelInput")) + return model_input if model_input else _EMPTY_BODY + + +def _line_status_code(entry: Mapping[str, object], custom_llm_provider: str) -> int | None: + response: Final = _as_object_mapping(entry.get("response")) + status: Final = response.get("status_code") if response is not None else None + if isinstance(status, int): + return status + if custom_llm_provider == "anthropic": + result: Final = _as_object_mapping(entry.get("result")) + if result is not None and result.get("type") == "succeeded": + return 200 + return None + + +def _call_type_for_request(request_line: Mapping[str, object] | None) -> str: + url: Final = request_line.get("url") if request_line is not None else None + return _CALL_TYPE_BY_BATCH_URL.get(url if isinstance(url, str) else "", "acompletion") + + +def _line_messages(request_body: Mapping[str, object]) -> object: + return request_body.get("messages") or request_body.get("input") or () + + +def _line_model( + response_body: Mapping[str, object], + request_body: Mapping[str, object], + parent: "Logging", +) -> str: + for candidate in (response_body.get("model"), request_body.get("model"), parent.model): + if isinstance(candidate, str) and candidate: + return candidate + return "" + + +_BatchLineResult: TypeAlias = "ModelResponse | EmbeddingResponse | ResponsesAPIResponse" + + +def _line_result(call_type: str, response_body: Mapping[str, object]) -> _BatchLineResult: + if call_type == "aembedding": + return EmbeddingResponse(**response_body) # pyright: ignore[reportArgumentType] # provider output bodies are dicts expanded as response ctor kwargs + if call_type == "aresponses": + return ResponsesAPIResponse(**response_body) # pyright: ignore[reportArgumentType] # same as above + return ModelResponse(**response_body) # pyright: ignore[reportArgumentType] # same as above + + +def _new_child_logging( + parent: "Logging", + model: str, + messages: object, + call_type: str, + start_time: datetime, +) -> "Logging": + from litellm.litellm_core_utils.litellm_logging import Logging + + return Logging( + model=model, + messages=messages, + stream=False, + call_type=call_type, + start_time=start_time, + litellm_call_id=str(uuid.uuid4()), + function_id=str(uuid.uuid4()), + litellm_trace_id=parent.litellm_trace_id, + dynamic_success_callbacks=parent.dynamic_success_callbacks, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # Logging ctor takes these dynamic callback lists untyped + dynamic_async_success_callbacks=parent.dynamic_async_success_callbacks, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # same as above + dynamic_failure_callbacks=parent.dynamic_failure_callbacks, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # same as above + dynamic_async_failure_callbacks=parent.dynamic_async_failure_callbacks, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # same as above + kwargs={"litellm_session_id": parent.litellm_session_id}, # mutable-ok: Logging's kwargs param takes a plain dict + ) + + +def _line_hidden_params( + batch: LiteLLMBatch, + custom_id: object, + status_code: int | None, + response_cost: float | None = None, +) -> dict[str, object]: # mutable-ok: response objects declare _hidden_params as a plain dict + return { # mutable-ok: same contract + "batch_id": batch.id, + "batch_custom_id": custom_id, + "batch_line_status_code": status_code, + "response_cost": response_cost, + } + + +def _optional_params_for_body(request_body: Mapping[str, object]) -> dict[str, object]: # mutable-ok: update_environment_variables takes a plain dict + return { # mutable-ok: same contract + key: value for key, value in request_body.items() if key not in ("model", "messages", "input") + } + + +async def _emit_line_event( + entry: Mapping[str, object], + request_line: Mapping[str, object] | None, + batch: LiteLLMBatch, + custom_llm_provider: _BatchLineProvider, + parent: "Logging", + model_name: str | None, + model_info: ModelInfo | None, +) -> bool: + custom_id: Final = entry.get("custom_id") or entry.get("recordId") + request_body: Final = _request_body_for_entry(entry, request_line) + status_code: Final = _line_status_code(entry, custom_llm_provider) + call_type: Final = _call_type_for_request(request_line) + response_body: Final = _get_response_from_batch_job_output_file(entry, custom_llm_provider) + parent_start_time: Final = parent.start_time # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # Logging.start_time is untyped upstream + start_time: Final = parent_start_time if isinstance(parent_start_time, datetime) else datetime.now() # noqa: DTZ005 # naive to match the logging pipeline start_time + parent_params: Final = _as_object_mapping(parent.litellm_params) or _EMPTY_BODY # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # Logging.litellm_params is untyped upstream + + child: Final = _new_child_logging( + parent=parent, + model=_line_model(response_body, request_body, parent), + messages=_line_messages(request_body), + call_type=call_type, + start_time=start_time, + ) + child.update_environment_variables( # pyright: ignore[reportUnknownMemberType] # Logging.update_environment_variables is untyped upstream + litellm_params={ # mutable-ok: update_environment_variables takes a plain dict + **parent_params, + "batch_parent_id": batch.id, + "metadata": dict(_as_object_mapping(parent_params.get("metadata")) or {}), # mutable-ok: copy of the parent's metadata dict + }, + optional_params=_optional_params_for_body(request_body), + model=child.model, + custom_llm_provider=custom_llm_provider, + ) + + now: Final = datetime.now() # noqa: DTZ005 # naive to match the logging pipeline start_time + if not _batch_response_was_successful(entry, custom_llm_provider): + exception: Final = _BatchLineFailure( + entry.get("error") or entry.get("response") or {} # mutable-ok: fallback payload dict passed to Exception + ) + exception._hidden_params = _line_hidden_params(batch, custom_id, status_code) # pyright: ignore[reportPrivateUsage] # _hidden_params is set on the exception instance itself + await child.async_failure_handler( + exception=exception, + traceback_exception="", + start_time=start_time, + end_time=now, + ) + return True + + stats: Final = _safe_output_line_stats(entry, custom_llm_provider, model_name, model_info) + try: + result: Final = _line_result(call_type, response_body) + except Exception: # noqa: BLE001 # one unparseable line must not drop the rest of the batch's line events + verbose_logger.warning( + "batch output line could not be reconstructed as a %s response, skipping it. custom_id=%s", + call_type, + custom_id, + ) + return False + + result._hidden_params = _line_hidden_params( # pyright: ignore[reportPrivateUsage] # same hidden_params channel the aggregate batch event uses + batch, + custom_id, + status_code, + response_cost=stats.prompt_cost + stats.completion_cost if stats is not None else None, + ) + if stats is not None and not response_body.get("usage") and isinstance(result, (ModelResponse, EmbeddingResponse)): + setattr( # noqa: B010 # ModelResponse.usage is set dynamically by its ctor, so setattr keeps parity for both response types + result, + "usage", + Usage( + prompt_tokens=stats.prompt_tokens, + completion_tokens=stats.completion_tokens, + total_tokens=stats.total_tokens, + ), + ) + await child.async_success_handler( + result=result, + start_time=start_time, + end_time=now, + cache_hit=False, + ) + return True + + +async def _fetch_managed_file_or_empty( + file_id: str | None, + custom_llm_provider: _BatchLineProvider, + fetch_params: dict[str, object] | None, # mutable-ok: batch_utils file fetch takes the shared litellm_params dict +) -> bytes: + if file_id is None: + return b"" + return await _fetch_batch_managed_file_content( + file_id, + custom_llm_provider=custom_llm_provider, + litellm_params=fetch_params, # pyright: ignore[reportArgumentType] # batch_utils types this param as an unparameterized dict + ) + + +async def log_batch_line_items( + batch: LiteLLMBatch, + custom_llm_provider: _BatchLineProvider, + parent: "Logging", + model_name: str | None, + litellm_params: dict[str, object] | None, # mutable-ok: the logging object's shared litellm_params dict + model_info: ModelInfo | None, +) -> int: + """Emit one callback event per JSONL line of a completed batch (request + paired with its response/error), behind the opt-in + ``litellm.store_batch_line_items_in_callbacks`` flag. The aggregate + aretrieve_batch event still bills the batch, so per-line events carry + ``batch_parent_id`` and never update spend themselves. Any failure here + is logged and swallowed: aggregate accounting must be unaffected.""" + emitted = 0 # rebind-ok: loop accumulator for emitted line count + try: + internal_credentials: Final = ( + litellm_params.get("_litellm_internal_model_credentials") if litellm_params else None + ) + internal_mapping: Final = _as_object_mapping(internal_credentials) + fetch_params: Final[dict[str, object] | None] = ( # mutable-ok: _fetch_batch_managed_file_content requires a plain dict + dict(internal_mapping) # mutable-ok: the file fetcher reads credential kwargs off a plain dict + if internal_mapping is not None + else litellm_params + ) + + input_file_content: Final = await _fetch_managed_file_or_empty( + batch.input_file_id, custom_llm_provider, fetch_params + ) + requests_by_id: Final = _requests_by_custom_id(input_file_content) + + output_content: Final = await _fetch_managed_file_or_empty( + batch.output_file_id, custom_llm_provider, fetch_params + ) + error_content: Final = await _fetch_managed_file_or_empty( + batch.error_file_id, custom_llm_provider, fetch_params + ) + for content in (output_content, error_content): + for entry in _output_entries(content): + try: + entry_key = entry.get("custom_id") or entry.get("recordId") # rebind-ok: per-iteration binding inside a loop cannot carry Final + request_line = requests_by_id.get(entry_key if isinstance(entry_key, str) else "") # rebind-ok: per-iteration binding inside a loop cannot carry Final + line_emitted = await _emit_line_event( # rebind-ok: same per-iteration binding + entry=entry, + request_line=request_line, + batch=batch, + custom_llm_provider=custom_llm_provider, + parent=parent, + model_name=model_name, + model_info=model_info, + ) + if line_emitted: + emitted += 1 + except Exception: # noqa: BLE001 # one bad line must not drop the rest of the batch's line events + verbose_logger.exception( + "batch line item logging failed for entry, continuing with remaining lines. batch_id=%s", + batch.id, + ) + except Exception: # noqa: BLE001 # line-item logging must never break the aggregate aretrieve_batch accounting + verbose_logger.exception( + "batch line item logging failed for batch_id=%s; aggregate logging unaffected", + batch.id, + ) + return emitted diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 40621a2f68d5..5e8139c1e8db 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -3027,6 +3027,18 @@ async def _async_success_handler_body( cost_for_built_in_tools_cost_usd_dollar=0.0, ) + if litellm.store_batch_line_items_in_callbacks: + from litellm.batches.batch_line_item_logging import log_batch_line_items + + await log_batch_line_items( + batch=result, + custom_llm_provider=self.custom_llm_provider, + parent=self, + model_name=self.get_deployment_model_for_cost(), + litellm_params=self.litellm_params, + model_info=self.get_router_deployment_model_info(), + ) + self.truncated_messages_for_logging = await truncate_base64_in_messages_async( StandardLoggingPayloadSetup.append_system_prompt_messages( kwargs=self.model_call_details, messages=self.model_call_details.get("messages") @@ -6102,8 +6114,10 @@ def _extract_response_obj_and_hidden_params( response_obj = {} if original_exception is not None and hidden_params is None: - response_headers: Final = _get_response_headers(original_exception) - if response_headers is not None: + exception_hidden_params: Final = getattr(original_exception, "_hidden_params", None) + if isinstance(exception_hidden_params, dict) and exception_hidden_params: + hidden_params = dict(exception_hidden_params) # mutable-ok: hidden_params downstream expects a plain dict + elif (response_headers := _get_response_headers(original_exception)) is not None: hidden_params = dict( StandardLoggingHiddenParams( additional_headers=StandardLoggingPayloadSetup.get_additional_headers(dict(response_headers)), diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 680c63393e82..dacbfb8ec815 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2797,6 +2797,10 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): None, description="If True, stores request messages and responses in spend logs. Default is False.", ) + store_batch_line_items_in_callbacks: bool | None = Field( + None, + description="If True, a completed batch logged via aretrieve_batch also emits one callback event per JSONL line item (request paired with its response or error). The aggregate batch callback is unchanged. Default is False.", + ) disable_auto_add_proxy_admin_to_teams: bool | None = Field( None, description="By default, the user calling /team/new is automatically added to the new team as a team admin. If True, proxy admins are no longer auto-added; members explicitly listed in members_with_roles are unaffected. Default is False.", diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 1ae106be3903..8ad4de38095b 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -1,6 +1,6 @@ import asyncio import traceback -from collections.abc import Callable, Sequence +from collections.abc import Callable, Mapping, Sequence from datetime import datetime from typing import TYPE_CHECKING, Any, Final, cast @@ -105,6 +105,11 @@ def __init__( async def async_log_success_event( self, kwargs: ObjectMapping, response_obj: object, start_time: datetime, end_time: datetime ) -> None: + # Per-line batch events emitted under store_batch_line_items_in_callbacks + # never touch spend: the aggregate aretrieve_batch event already bills the batch. + litellm_params: Final = kwargs.get("litellm_params") + if isinstance(litellm_params, Mapping) and litellm_params.get("batch_parent_id"): + return if self.spend_event_producer is None or not is_offloadable_success(response_obj): await self._PROXY_track_cost_callback(kwargs, response_obj, start_time, end_time) return diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index d7d8413d2ce7..e47361d1de5f 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -6044,6 +6044,13 @@ async def load_config(self, router: litellm.Router | None, config_file_path: str litellm.use_legacy_interactions_schema = _use_legacy_interactions_schema.lower() == "true" else: litellm.use_legacy_interactions_schema = bool(_use_legacy_interactions_schema) + ### BATCH LINE ITEM CALLBACKS ### + _store_batch_line_items: Final = general_settings.get("store_batch_line_items_in_callbacks") + if _store_batch_line_items is not None: + if isinstance(_store_batch_line_items, str): + litellm.store_batch_line_items_in_callbacks = _store_batch_line_items.lower() == "true" + else: + litellm.store_batch_line_items_in_callbacks = bool(_store_batch_line_items) # Health-check-driven routing (opt-in, passes through to Router later) _enable_hc_routing = general_settings.get("enable_health_check_routing", False) _hc_staleness = general_settings.get("health_check_staleness_threshold", None) @@ -7158,6 +7165,19 @@ async def _update_general_settings(self, db_general_settings: Json | None): # For other types, convert to bool general_settings["store_prompts_in_spend_logs"] = bool(value) + if "store_batch_line_items_in_callbacks" in _general_settings: + store_line_items_value: Final = ( + general_settings.get("store_batch_line_items_in_callbacks") + if "store_batch_line_items_in_callbacks" in self._yaml_general_settings_keys + else _general_settings["store_batch_line_items_in_callbacks"] + ) + if store_line_items_value is not None: + litellm.store_batch_line_items_in_callbacks = ( + store_line_items_value.lower() == "true" + if isinstance(store_line_items_value, str) + else bool(store_line_items_value) + ) + if "disable_auto_add_proxy_admin_to_teams" in _general_settings: value = _general_settings["disable_auto_add_proxy_admin_to_teams"] if isinstance(value, str): diff --git a/litellm/types/utils.py b/litellm/types/utils.py index aaa16fd2d44e..8db91ffcb527 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3098,6 +3098,9 @@ class StandardLoggingHiddenParams(TypedDict): batch_failed_requests: ReadOnly[int | None] litellm_model_name: str | None # the model name sent to the provider by litellm usage_object: dict | None + batch_id: NotRequired[ReadOnly[str | None]] + batch_custom_id: NotRequired[ReadOnly[str | None]] + batch_line_status_code: NotRequired[ReadOnly[int | None]] class StandardLoggingModelInformation(TypedDict): diff --git a/tests/test_litellm/batches/test_batch_line_item_logging.py b/tests/test_litellm/batches/test_batch_line_item_logging.py new file mode 100644 index 000000000000..7fc60ddba778 --- /dev/null +++ b/tests/test_litellm/batches/test_batch_line_item_logging.py @@ -0,0 +1,233 @@ +""" +Tests for litellm/batches/batch_line_item_logging.py and its hook in +Logging._async_success_handler_body. + +When ``litellm.store_batch_line_items_in_callbacks`` is on and a completed +batch is logged (call_type aretrieve_batch), litellm emits one callback event +per JSONL line (request paired with response/error) in addition to the +aggregate batch event. These tests run the real Logging.async_success_handler +the way the CheckBatchCost poller invokes it, with a recording CustomLogger on +the async success/failure lists, so a regression in pairing, hidden params, +cost, or error propagation fails here. +""" + +import json +import uuid +from datetime import datetime +from types import SimpleNamespace +from unittest.mock import AsyncMock, patch + +import pytest + +import litellm +from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.types.utils import LiteLLMBatch, Usage + +INPUT_JSONL = b"\n".join( + [ + json.dumps( + { + "custom_id": "a", + "method": "POST", + "url": "/v1/chat/completions", + "body": { + "model": "gpt-4o", + "messages": [{"role": "user", "content": "hi a"}], + "temperature": 0.2, + }, + } + ).encode(), + json.dumps( + { + "custom_id": "b", + "method": "POST", + "url": "/v1/chat/completions", + "body": { + "model": "gpt-4o", + "messages": [{"role": "user", "content": "hi b"}], + }, + } + ).encode(), + ] +) + +OUTPUT_JSONL = json.dumps( + { + "custom_id": "a", + "response": { + "status_code": 200, + "body": { + "id": "chatcmpl-1", + "model": "gpt-4o", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "hello back"}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + }, + }, + } +).encode() + +ERROR_JSONL = json.dumps( + { + "custom_id": "b", + "response": {"status_code": 400, "body": {"error": {"message": "bad request boom"}}}, + "error": {"message": "bad request boom"}, + } +).encode() + +_FILE_BYTES = { + "input-file-1": INPUT_JSONL, + "output-file-1": OUTPUT_JSONL, + "error-file-1": ERROR_JSONL, +} + + +def _batch() -> LiteLLMBatch: + return LiteLLMBatch( + id="batch_1", + object="batch", + endpoint="/v1/chat/completions", + input_file_id="input-file-1", + output_file_id="output-file-1", + error_file_id="error-file-1", + status="completed", + completion_window="24h", + created_at=1, + ) + + +def _file_content(file_id: str, **_kwargs): + return SimpleNamespace(content=_FILE_BYTES[file_id]) + + +class _RecordingLogger(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.success_events = [] + self.failure_events = [] + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + self.success_events.append(kwargs) + + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + self.failure_events.append(kwargs) + + +@pytest.fixture +def recorder(): + logger = _RecordingLogger() + saved_flag = litellm.store_batch_line_items_in_callbacks # test-quality-ok: process-wide opt-in flag; restored in teardown + saved_success = list(litellm._async_success_callback) # test-quality-ok: the feature dispatches through this global list; restored in teardown + saved_failure = list(litellm._async_failure_callback) # test-quality-ok: same dispatch seam, restored in teardown + litellm._async_success_callback = [logger] # test-quality-ok: there is no injection seam for callback lists; teardown restores + litellm._async_failure_callback = [logger] # test-quality-ok: same dispatch seam, restored in teardown + yield logger + litellm.store_batch_line_items_in_callbacks = saved_flag # test-quality-ok: teardown restoring the value set above + litellm._async_success_callback = saved_success # test-quality-ok: teardown restoring the value set above + litellm._async_failure_callback = saved_failure # test-quality-ok: teardown restoring the value set above + + +def _parent_logging() -> Logging: + logging_obj = Logging( + model="gpt-4o", + messages=[{"role": "user", "content": ""}], + stream=False, + call_type="aretrieve_batch", + start_time=datetime.now(), + litellm_call_id=str(uuid.uuid4()), + function_id=str(uuid.uuid4()), + ) + logging_obj.update_environment_variables( + litellm_params={"metadata": {"model_info": {"id": "dep-1"}, "model_group": "gpt-4o"}}, + optional_params={}, + custom_llm_provider="openai", + ) + return logging_obj + + +async def _log_completed_batch(logging_obj: Logging, batch: LiteLLMBatch) -> None: + await logging_obj.async_success_handler( + result=batch, + batch_cost=1.5, + batch_usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15), + batch_models=["gpt-4o"], + batch_successful_requests=1, + batch_failed_requests=1, + batch_prompt_cost=1.0, + batch_completion_cost=0.5, + ) + + +def _payload(event: dict) -> dict: + return event["standard_logging_object"] + + +def _hidden(event: dict) -> dict: + return _payload(event)["hidden_params"] + + +@pytest.mark.asyncio +async def test_line_items_emitted_alongside_aggregate(recorder): + litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: the flag under test is a module global; fixture restores it + batch = _batch() + with ( + patch("litellm.files.main.afile_content", new_callable=AsyncMock, side_effect=_file_content), # test-quality-ok: afile_content is the provider boundary; there is no HTTP transport or injection seam for managed file fetch + patch("litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02)), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch + ): + await _log_completed_batch(_parent_logging(), batch) + + assert len(recorder.success_events) == 2 + assert len(recorder.failure_events) == 1 + + aggregate = next(e for e in recorder.success_events if _hidden(e).get("batch_custom_id") is None) + assert _payload(aggregate)["response_cost"] == 1.5 + assert "batch_custom_id" not in _hidden(aggregate) + + line = next(e for e in recorder.success_events if _hidden(e).get("batch_custom_id") == "a") + hidden = _hidden(line) + assert hidden["batch_id"] == batch.id + assert hidden["batch_line_status_code"] == 200 + assert line["litellm_params"]["batch_parent_id"] == batch.id + + payload = _payload(line) + assert payload["response_cost"] == pytest.approx(0.03) + assert payload["prompt_tokens"] == 10 + assert payload["completion_tokens"] == 5 + assert payload["model_parameters"]["temperature"] == 0.2 + assert any(m.get("content") == "hi a" for m in payload["messages"]) + assert payload["response"]["choices"][0]["message"]["content"] == "hello back" + + failure = recorder.failure_events[0] + assert _hidden(failure)["batch_custom_id"] == "b" + assert _hidden(failure)["batch_id"] == batch.id + assert _hidden(failure)["batch_line_status_code"] == 400 + assert "bad request boom" in _payload(failure)["error_str"] + + +@pytest.mark.asyncio +async def test_flag_off_emits_only_aggregate(recorder): + assert litellm.store_batch_line_items_in_callbacks is False + file_mock = AsyncMock(side_effect=_file_content) + with patch("litellm.files.main.afile_content", file_mock): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + await _log_completed_batch(_parent_logging(), _batch()) + + assert len(recorder.success_events) == 1 + assert len(recorder.failure_events) == 0 + file_mock.assert_not_called() + + +@pytest.mark.asyncio +async def test_input_fetch_failure_still_emits_aggregate(recorder): + litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: the flag under test is a module global; fixture restores it + with patch("litellm.files.main.afile_content", new_callable=AsyncMock, side_effect=ValueError("boom")): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + await _log_completed_batch(_parent_logging(), _batch()) + + assert len(recorder.success_events) == 1 + assert len(recorder.failure_events) == 0 + assert _payload(recorder.success_events[0])["response_cost"] == 1.5 diff --git a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py index dfc95db3e143..b42dc41f097f 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py @@ -2536,3 +2536,21 @@ async def test_async_post_call_failure_hook_persists_no_raw_model_on_an_unknown_ == "/chat/completions: Invalid model name passed in. Call `/v1/models` to view available models for your key." ) assert error_information["error_class"] == "ProxyModelNotFoundError" + + +@pytest.mark.asyncio +async def test_batch_line_item_event_never_updates_spend(): # test-quality-ok: the observable contract is exactly that the spend path is never invoked + logger: Final = _ProxyDBLogger() + kwargs: Final = { + "litellm_params": {"batch_parent_id": "batch_1", "metadata": {}}, + "model": "gpt-4o", + "call_type": "acompletion", + } + with patch.object(logger, "_PROXY_track_cost_callback", new_callable=AsyncMock) as mock_track: # test-quality-ok: asserts the callback's own method is skipped; the DB writer is never reached + await logger.async_log_success_event( + kwargs=kwargs, + response_obj=ModelResponse(), + start_time=datetime.now(), + end_time=datetime.now(), + ) + mock_track.assert_not_awaited() From 0ea8ce740c671171ee07959c8519454270ce9902 Mon Sep 17 00:00:00 2001 From: shivam Date: Thu, 17 Sep 2026 22:50:34 +0000 Subject: [PATCH 02/89] fix(batches): only emit line items for final batches Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/litellm_core_utils/litellm_logging.py | 2 +- litellm/llms/bedrock/common_utils.py | 10 +++++--- .../batches/test_batch_line_item_logging.py | 24 +++++++++++++++++++ .../chat/test_converse_transformation.py | 18 ++++++++++---- 4 files changed, 46 insertions(+), 8 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 5e8139c1e8db..acba0eb10252 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -3027,7 +3027,7 @@ async def _async_success_handler_body( cost_for_built_in_tools_cost_usd_dollar=0.0, ) - if litellm.store_batch_line_items_in_callbacks: + if litellm.store_batch_line_items_in_callbacks and (has_explicit_batch_data or should_compute_batch_data): from litellm.batches.batch_line_item_logging import log_batch_line_items await log_batch_line_items( diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index cb2c70e74c83..20c8258e4404 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -34,6 +34,7 @@ _ERROR_REQUEST_URL: Final = "https://docs.litellm.ai/docs" +_OPENAI_FAMILY_MODEL_RE: Final = re.compile(r"(^|[./])openai\.") def error_response_text(response: httpx.Response) -> str: @@ -878,9 +879,10 @@ def bedrock_model_accepts_cache_points(model: str | None) -> bool: """ Whether Converse ``cachePoint`` blocks may be sent to this model. - Bedrock rejects requests carrying cachePoint blocks for models without prompt - caching support ("You invoked an unsupported model or your request did not allow - prompt caching"), so a model whose cost-map entry does not declare + OpenAI-family models only support implicit caching and never accept explicit + ``cachePoint`` blocks. Bedrock rejects requests carrying cachePoint blocks for + models without prompt caching support ("You invoked an unsupported model or your + request did not allow prompt caching"), so a model whose cost-map entry does not declare ``supports_prompt_caching`` must not receive them. A model absent from the map (an application inference profile ARN, a model newer than the map) keeps emitting so existing caching setups never silently degrade. ``litellm.utils.supports_prompt_caching`` @@ -888,6 +890,8 @@ def bedrock_model_accepts_cache_points(model: str | None) -> bool: """ if model is None: return True + if _OPENAI_FAMILY_MODEL_RE.search(model): + return False entries: Final = tuple( entry for candidate in (model, get_bedrock_base_model(model)) diff --git a/tests/test_litellm/batches/test_batch_line_item_logging.py b/tests/test_litellm/batches/test_batch_line_item_logging.py index 7fc60ddba778..78265c2fea3d 100644 --- a/tests/test_litellm/batches/test_batch_line_item_logging.py +++ b/tests/test_litellm/batches/test_batch_line_item_logging.py @@ -15,6 +15,7 @@ import uuid from datetime import datetime from types import SimpleNamespace +from typing import Final from unittest.mock import AsyncMock, patch import pytest @@ -222,6 +223,29 @@ async def test_flag_off_emits_only_aggregate(recorder): file_mock.assert_not_called() +@pytest.mark.asyncio +async def test_in_progress_batch_poll_emits_no_line_events(recorder): + litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: the flag under test is a module global; fixture restores it + in_progress: Final = LiteLLMBatch( + id="batch_wip", + object="batch", + endpoint="/v1/chat/completions", + input_file_id="input-file-1", + output_file_id=None, + error_file_id=None, + status="in_progress", + completion_window="24h", + created_at=1, + ) + file_mock: Final = AsyncMock(side_effect=_file_content) + with patch("litellm.files.main.afile_content", file_mock): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + await _parent_logging().async_success_handler(result=in_progress) + + file_mock.assert_not_called() + assert all(_hidden(e).get("batch_custom_id") is None for e in recorder.success_events) + assert len(recorder.failure_events) == 0 + + @pytest.mark.asyncio async def test_input_fetch_failure_still_emits_aggregate(recorder): litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: the flag under test is a module global; fixture restores it diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index eba8d912fe07..9b28e42f93b9 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -1189,17 +1189,24 @@ def test_get_supported_openai_params_bedrock_converse(): @pytest.mark.parametrize( - "tools, expected_marker", + "tools, model, expected_marker", [ pytest.param( [{"type": "function", "function": {"name": "f", "parameters": {"type": "object", "properties": {}}}}], + "anthropic.claude-sonnet-4-5-20250929-v1:0", "dep-bedrock", id="tools-present-so-the-cachepoint-is-placed", ), - pytest.param(None, None, id="no-tools-so-nothing-is-placed"), + pytest.param(None, "anthropic.claude-sonnet-4-5-20250929-v1:0", None, id="no-tools-so-nothing-is-placed"), + pytest.param( + [{"type": "function", "function": {"name": "f", "parameters": {"type": "object", "properties": {}}}}], + "global.openai.gpt-6-astra", + None, + id="openai-family-implicit-caching-only", + ), ], ) -def test_tool_config_cachepoint_is_credited_only_where_it_is_placed(tools, expected_marker): +def test_tool_config_cachepoint_is_credited_only_where_it_is_placed(tools, model, expected_marker): """Spend attribution credits the gateway for breakpoints it placed, and a tool_config point becomes one here or nowhere. @@ -1213,7 +1220,7 @@ def test_tool_config_cachepoint_is_credited_only_where_it_is_placed(tools, expec optional_params["tools"] = tools data = AmazonConverseConfig()._transform_request_helper( - model="anthropic.claude-sonnet-4-5-20250929-v1:0", + model=model, system_content_blocks=[], optional_params=optional_params, messages=[{"role": "user", "content": "hi"}], @@ -5591,6 +5598,9 @@ def test_cache_control_injection_tool_config_drops_ttl_for_unsupported_model(): True, id="unmapped-arn-keeps-emitting", ), + pytest.param("global.openai.gpt-6-astra", False, id="openai-family-implicit-caching-only"), + pytest.param("openai.gpt-oss-120b-1:0", False, id="openai-gpt-oss"), + pytest.param("us.openai.gpt-99-unmapped", False, id="unmapped-openai-family-still-suppressed"), ], ) def test_cache_points_emitted_only_for_models_that_support_prompt_caching(model, expects_cache_points, monkeypatch): From 7fa8c92a049587651b1344d5900c0adb20f945e5 Mon Sep 17 00:00:00 2001 From: shivam Date: Thu, 17 Sep 2026 22:53:05 +0000 Subject: [PATCH 03/89] chore(deps): bump soupsieve to 2.9.2 for GHSA-gjv8-xp57-g29c and GHSA-j934-xhv5-fg8f Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- pyproject.toml | 2 +- uv.lock | 8 ++++---- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index dfe84a28d52a..14c980486194 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -309,7 +309,7 @@ constraint-dependencies = [ "tornado>=6.5.8", "aiohttp>=3.14.2,<4.0", "packaging>=24.0", - "soupsieve>=2.8.4", + "soupsieve>=2.9.0", "httplib2>=0.32.0", "setuptools>=83.0.0", ] diff --git a/uv.lock b/uv.lock index 35eaa20c39e9..e7e46cdec1d5 100644 --- a/uv.lock +++ b/uv.lock @@ -24,7 +24,7 @@ constraints = [ { name = "httplib2", specifier = ">=0.32.0" }, { name = "packaging", specifier = ">=24.0" }, { name = "setuptools", specifier = ">=83.0.0" }, - { name = "soupsieve", specifier = ">=2.8.4" }, + { name = "soupsieve", specifier = ">=2.9.0" }, { name = "tornado", specifier = ">=6.5.8" }, ] overrides = [ @@ -9262,11 +9262,11 @@ wheels = [ [[package]] name = "soupsieve" -version = "2.8.4" +version = "2.9.2" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/47/2c/0a5f6f8ee0d5589e48c7640213ed5175d52cf540a06725b628cc1a45d6ce/soupsieve-2.8.4.tar.gz", hash = "sha256:e121fd02e975c695e4e9e8774a5ee35d74714b59307868dcc5319ad2d9e3328e", size = 121110, upload-time = "2026-05-24T13:55:57.154Z" } +sdist = { url = "https://files.pythonhosted.org/packages/69/99/a6ca3beb3ccacb41fb3321d8a60e5566f9e6467601ef8eba6a17e1b89778/soupsieve-2.9.2.tar.gz", hash = "sha256:4a55d8cf158a9c2e587fa4922f1bbb91d68ac829e2d6f25403a85747c71daf74", size = 122445, upload-time = "2026-08-07T00:57:24.801Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/5e/f5/0c41cb68dcae6b7de4fac4188a3a9589e21fb31df21ea3a2e888db95e6c9/soupsieve-2.8.4-py3-none-any.whl", hash = "sha256:e7e6b0769c8f51ed59acab6e994b00621096cfb1c640a7509295987388fbaf65", size = 37304, upload-time = "2026-05-24T13:55:55.406Z" }, + { url = "https://files.pythonhosted.org/packages/eb/dc/ad025c1ee131eba60c69f4dd5779b18fcf1e6b21a343e2162a84d5d133c7/soupsieve-2.9.2-py3-none-any.whl", hash = "sha256:8089a26fd974ca7a1f30276d3d8492ab266ab15af581642dfe8aa162e0c1c823", size = 37370, upload-time = "2026-08-07T00:57:23.524Z" }, ] [[package]] From c7edcce3c91be229fe6ce5ed516ff5df3e6aed7c Mon Sep 17 00:00:00 2001 From: shivam Date: Thu, 17 Sep 2026 22:57:12 +0000 Subject: [PATCH 04/89] style(batches): satisfy ruff format for line item logging Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/batches/batch_line_item_logging.py | 28 ++++++++++++---------- 1 file changed, 16 insertions(+), 12 deletions(-) diff --git a/litellm/batches/batch_line_item_logging.py b/litellm/batches/batch_line_item_logging.py index 077457dd2995..65227c24c8ee 100644 --- a/litellm/batches/batch_line_item_logging.py +++ b/litellm/batches/batch_line_item_logging.py @@ -44,7 +44,7 @@ class _BatchLineFailure(Exception): def __init__(self, error_payload: object) -> None: super().__init__(json.dumps(error_payload)) - self._hidden_params: dict[str, object] = {} # mutable-ok: mirrors the plain-dict _hidden_params contract on litellm response objects + self._hidden_params: dict[str, object] = {} # mutable-ok: plain-dict contract like response _hidden_params def _as_object_mapping(value: object) -> Mapping[str, object] | None: @@ -151,7 +151,7 @@ def _new_child_logging( dynamic_async_success_callbacks=parent.dynamic_async_success_callbacks, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # same as above dynamic_failure_callbacks=parent.dynamic_failure_callbacks, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # same as above dynamic_async_failure_callbacks=parent.dynamic_async_failure_callbacks, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # same as above - kwargs={"litellm_session_id": parent.litellm_session_id}, # mutable-ok: Logging's kwargs param takes a plain dict + kwargs={"litellm_session_id": parent.litellm_session_id}, # mutable-ok: kwargs takes a plain dict ) @@ -169,15 +169,22 @@ def _line_hidden_params( } -def _optional_params_for_body(request_body: Mapping[str, object]) -> dict[str, object]: # mutable-ok: update_environment_variables takes a plain dict +def _optional_params_for_body( + request_body: Mapping[str, object], +) -> dict[str, object]: # mutable-ok: update_environment_variables takes a plain dict return { # mutable-ok: same contract key: value for key, value in request_body.items() if key not in ("model", "messages", "input") } +def _metadata_copy(params: Mapping[str, object]) -> dict[str, object]: # mutable-ok: dict out for litellm_params + metadata: Final = _as_object_mapping(params.get("metadata")) or _EMPTY_BODY + return {**metadata} # mutable-ok: plain-dict copy + + async def _emit_line_event( entry: Mapping[str, object], - request_line: Mapping[str, object] | None, + requests_by_id: Mapping[str, Mapping[str, object]], batch: LiteLLMBatch, custom_llm_provider: _BatchLineProvider, parent: "Logging", @@ -185,6 +192,7 @@ async def _emit_line_event( model_info: ModelInfo | None, ) -> bool: custom_id: Final = entry.get("custom_id") or entry.get("recordId") + request_line: Final = requests_by_id.get(custom_id if isinstance(custom_id, str) else "") request_body: Final = _request_body_for_entry(entry, request_line) status_code: Final = _line_status_code(entry, custom_llm_provider) call_type: Final = _call_type_for_request(request_line) @@ -204,7 +212,7 @@ async def _emit_line_event( litellm_params={ # mutable-ok: update_environment_variables takes a plain dict **parent_params, "batch_parent_id": batch.id, - "metadata": dict(_as_object_mapping(parent_params.get("metadata")) or {}), # mutable-ok: copy of the parent's metadata dict + "metadata": _metadata_copy(parent_params), }, optional_params=_optional_params_for_body(request_body), model=child.model, @@ -295,7 +303,7 @@ async def log_batch_line_items( litellm_params.get("_litellm_internal_model_credentials") if litellm_params else None ) internal_mapping: Final = _as_object_mapping(internal_credentials) - fetch_params: Final[dict[str, object] | None] = ( # mutable-ok: _fetch_batch_managed_file_content requires a plain dict + fetch_params: Final[dict[str, object] | None] = ( # mutable-ok: file fetcher requires a plain dict dict(internal_mapping) # mutable-ok: the file fetcher reads credential kwargs off a plain dict if internal_mapping is not None else litellm_params @@ -315,19 +323,15 @@ async def log_batch_line_items( for content in (output_content, error_content): for entry in _output_entries(content): try: - entry_key = entry.get("custom_id") or entry.get("recordId") # rebind-ok: per-iteration binding inside a loop cannot carry Final - request_line = requests_by_id.get(entry_key if isinstance(entry_key, str) else "") # rebind-ok: per-iteration binding inside a loop cannot carry Final - line_emitted = await _emit_line_event( # rebind-ok: same per-iteration binding + emitted += await _emit_line_event( entry=entry, - request_line=request_line, + requests_by_id=requests_by_id, batch=batch, custom_llm_provider=custom_llm_provider, parent=parent, model_name=model_name, model_info=model_info, ) - if line_emitted: - emitted += 1 except Exception: # noqa: BLE001 # one bad line must not drop the rest of the batch's line events verbose_logger.exception( "batch line item logging failed for entry, continuing with remaining lines. batch_id=%s", From 7f067086f7599e1423e4041c3b833b1b30629a04 Mon Sep 17 00:00:00 2001 From: shivam Date: Thu, 17 Sep 2026 23:00:41 +0000 Subject: [PATCH 05/89] chore(ui): regenerate schema.d.ts for store_batch_line_items_in_callbacks Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ui/litellm-dashboard/src/lib/http/schema.d.ts | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 872875cc5352..a74b996eaa04 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -26712,6 +26712,11 @@ export interface components { reject_clientside_metadata_tags?: boolean | null; /** @description Spreads the proxy's scheduled background jobs (spend flushes, budget resets, config reloads, exports) across a window instead of firing them together on every replica. On by default; set to tune the window, pin a job, or turn it off. */ scheduled_job_stagger?: components["schemas"]["ScheduledJobStaggerSettings"] | null; + /** + * Store Batch Line Items In Callbacks + * @description If True, a completed batch logged via aretrieve_batch also emits one callback event per JSONL line item (request paired with its response or error). The aggregate batch callback is unchanged. Default is False. + */ + store_batch_line_items_in_callbacks?: boolean | null; /** * Store Model In Db * @description If True, models and config are stored in and loaded from the database. Default is False. From 890802326a12ffbd9bf9244d0359dbbe645b5a68 Mon Sep 17 00:00:00 2001 From: shivam Date: Thu, 17 Sep 2026 23:16:07 +0000 Subject: [PATCH 06/89] test(batches): cover line-item edge shapes and general settings propagation Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../batches/test_batch_line_item_logging.py | 226 ++++++++++++++++++ tests/test_litellm/proxy/test_proxy_server.py | 31 +++ 2 files changed, 257 insertions(+) diff --git a/tests/test_litellm/batches/test_batch_line_item_logging.py b/tests/test_litellm/batches/test_batch_line_item_logging.py index 78265c2fea3d..32eeb326c606 100644 --- a/tests/test_litellm/batches/test_batch_line_item_logging.py +++ b/tests/test_litellm/batches/test_batch_line_item_logging.py @@ -255,3 +255,229 @@ async def test_input_fetch_failure_still_emits_aggregate(recorder): assert len(recorder.success_events) == 1 assert len(recorder.failure_events) == 0 assert _payload(recorder.success_events[0])["response_cost"] == 1.5 + +EDGE_INPUT_JSONL = b"\n".join( + [ + json.dumps( + { + "custom_id": "e", + "method": "POST", + "url": "/v1/embeddings", + "body": {"model": "text-embedding-3-small", "input": "embed me", "encoding_format": "float"}, + } + ).encode(), + json.dumps( + { + "custom_id": "r", + "method": "POST", + "url": "/v1/responses", + "body": {"model": "gpt-4o", "input": "respond to me"}, + } + ).encode(), + json.dumps( + { + "custom_id": "badresp", + "method": "POST", + "url": "/v1/responses", + "body": {"model": "gpt-4o", "input": "unreconstructable"}, + } + ).encode(), + json.dumps( + { + "custom_id": "n", + "method": "POST", + "url": "/v1/chat/completions", + "body": {"model": "gpt-4o", "messages": [{"role": "user", "content": "hi n"}]}, + } + ).encode(), + ] +) + +EDGE_OUTPUT_JSONL = b"\n".join( + [ + json.dumps( + { + "custom_id": "e", + "response": { + "status_code": 200, + "body": { + "object": "list", + "data": [{"object": "embedding", "embedding": [0.1], "index": 0}], + "model": "text-embedding-3-small", + "usage": {"prompt_tokens": 3, "total_tokens": 3}, + }, + }, + } + ).encode(), + json.dumps( + { + "custom_id": "r", + "response": { + "status_code": 200, + "body": { + "id": "resp_1", + "object": "response", + "created_at": 1, + "status": "completed", + "output": [], + "model": "gpt-4o", + }, + }, + } + ).encode(), + json.dumps({"custom_id": "badresp", "response": {"status_code": 200, "body": {}}}).encode(), + json.dumps({"custom_id": "n", "response": {"body": {"error": {"message": "no status here"}}}}).encode(), + json.dumps({"custom_id": "boom", "response": "not-a-dict"}).encode(), + json.dumps( + { + "custom_id": "mi", + "modelInput": {"messages": [{"role": "user", "content": "hi mi"}]}, + "response": { + "status_code": 200, + "body": { + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "mi back"}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 2, "completion_tokens": 1, "total_tokens": 3}, + }, + }, + } + ).encode(), + ] +) + +ANTHROPIC_INPUT_JSONL = json.dumps( + { + "custom_id": "b2", + "params": {"model": "claude-3", "max_tokens": 5, "messages": [{"role": "user", "content": "hi b2"}]}, + } +).encode() + +ANTHROPIC_OUTPUT_JSONL = json.dumps( + { + "custom_id": "b2", + "result": { + "type": "succeeded", + "message": { + "id": "msg_1", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "hello b2"}], + "model": "claude-3", + "usage": {"input_tokens": 1, "output_tokens": 2}, + }, + }, + } +).encode() + + +def _edge_file_content(file_id: str, **_kwargs): + return SimpleNamespace( + content={ + "input-2": EDGE_INPUT_JSONL, + "output-2": EDGE_OUTPUT_JSONL, + "input-anth": ANTHROPIC_INPUT_JSONL, + "output-anth": ANTHROPIC_OUTPUT_JSONL, + }[file_id] + ) + + +def _parent_logging_with_params(litellm_params: dict) -> Logging: + logging_obj = _parent_logging() + logging_obj.update_environment_variables( + litellm_params=litellm_params, + optional_params={}, + custom_llm_provider="openai", + ) + return logging_obj + + +@pytest.mark.asyncio +async def test_line_items_edge_shapes_and_edge_cases(recorder): + litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: the flag under test is a module global; fixture restores it + batch = LiteLLMBatch( + id="batch_edge", + object="batch", + endpoint="/v1/chat/completions", + input_file_id="input-2", + output_file_id="output-2", + error_file_id=None, + status="completed", + completion_window="24h", + created_at=1, + ) + file_mock: Final = AsyncMock(side_effect=_edge_file_content) + parent: Final = _parent_logging_with_params( + { + "metadata": {"model_info": {"id": "dep-1"}, "model_group": "gpt-4o"}, + "_litellm_internal_model_credentials": {"api_key": "sk-line-items-marker"}, + } + ) + with ( + patch("litellm.files.main.afile_content", file_mock), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + patch("litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02)), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch + ): + await _log_completed_batch(parent, batch) + + by_custom_id = {_hidden(e).get("batch_custom_id"): e for e in recorder.success_events} + aggregate = next(e for e in recorder.success_events if _hidden(e).get("batch_custom_id") is None) + assert _payload(aggregate)["response_cost"] == 1.5 + + assert by_custom_id["e"]["litellm_params"]["batch_parent_id"] == batch.id + assert by_custom_id["e"]["call_type"] == "aembedding" + assert _payload(by_custom_id["e"])["response"]["data"][0]["embedding"] == [0.1] + + assert by_custom_id["r"]["call_type"] == "aresponses" + assert _payload(by_custom_id["r"])["response"]["id"] == "resp_1" + + assert by_custom_id["mi"]["call_type"] == "acompletion" + assert _hidden(by_custom_id["mi"])["batch_line_status_code"] == 200 + + assert "badresp" not in by_custom_id + assert "boom" not in by_custom_id + + failure = recorder.failure_events[0] + assert _hidden(failure)["batch_custom_id"] == "n" + assert _hidden(failure)["batch_line_status_code"] is None + assert "no status here" in _payload(failure)["error_str"] + + assert "sk-line-items-marker" in str(file_mock.call_args_list) + file_ids_fetched = [call.kwargs.get("file_id") or call.args[0] for call in file_mock.call_args_list] + assert "input-2" in file_ids_fetched and "output-2" in file_ids_fetched + assert not any(file_id is None for file_id in file_ids_fetched) + + +@pytest.mark.asyncio +async def test_line_items_anthropic_shapes(recorder): + litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: the flag under test is a module global; fixture restores it + batch = LiteLLMBatch( + id="batch_anth", + object="batch", + endpoint="/v1/messages", + input_file_id="input-anth", + output_file_id="output-anth", + error_file_id=None, + status="completed", + completion_window="24h", + created_at=1, + ) + file_mock: Final = AsyncMock(side_effect=_edge_file_content) + logging_obj = _parent_logging() + logging_obj.update_environment_variables( + litellm_params={"metadata": {"model_info": {"id": "dep-1"}, "model_group": "claude-3"}}, + optional_params={}, + custom_llm_provider="anthropic", + ) + with ( + patch("litellm.files.main.afile_content", file_mock), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + patch("litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02)), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch + ): + await _log_completed_batch(logging_obj, batch) + + line = next(e for e in recorder.success_events if _hidden(e).get("batch_custom_id") == "b2") + assert _hidden(line)["batch_line_status_code"] == 200 + assert line["litellm_params"]["batch_parent_id"] == batch.id diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 41c4956dba6b..7503a4b30ffe 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -7291,6 +7291,37 @@ async def test_update_general_settings_store_model_in_db_false(): assert ps.general_settings["store_model_in_db"] is False +@pytest.mark.asyncio +async def test_update_general_settings_store_batch_line_items_in_callbacks(): + """ + Verify _update_general_settings sets the litellm module flag when the DB + general_settings carries store_batch_line_items_in_callbacks, and that a + YAML-explicit value wins over the DB value. + """ + from litellm.proxy.proxy_server import ProxyConfig + + proxy_config = ProxyConfig() + saved_flag = litellm.store_batch_line_items_in_callbacks + try: + with patch("litellm.proxy.proxy_server.general_settings", {}): # test-quality-ok: module-global seam + await proxy_config._update_general_settings( + db_general_settings={"store_batch_line_items_in_callbacks": True} + ) + assert litellm.store_batch_line_items_in_callbacks is True + + proxy_config._yaml_general_settings_keys = {"store_batch_line_items_in_callbacks"} + with patch( # test-quality-ok: module-global seam + "litellm.proxy.proxy_server.general_settings", + {"store_batch_line_items_in_callbacks": "false"}, + ): + await proxy_config._update_general_settings( + db_general_settings={"store_batch_line_items_in_callbacks": True} + ) + assert litellm.store_batch_line_items_in_callbacks is False + finally: + litellm.store_batch_line_items_in_callbacks = saved_flag + + @pytest.mark.asyncio async def test_update_general_settings_propagates_apply_user_budget_to_team_keys(): """The Admin UI toggle writes to the DB config, so the flag has to be in the From 8aca0c37fd7c2cd00830dbe3ddd5e6322bd1c122 Mon Sep 17 00:00:00 2001 From: shivam Date: Thu, 17 Sep 2026 23:31:18 +0000 Subject: [PATCH 07/89] fix(batches): keep deployment credentials out of callback-visible params Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/common_utils/check_batch_cost.py | 6 +++++- litellm/batches/batch_line_item_logging.py | 2 +- tests/test_litellm/batches/test_batch_line_item_logging.py | 3 ++- 3 files changed, 8 insertions(+), 3 deletions(-) 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 457bcf8c1479..b41b3a6e1ad5 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -856,7 +856,6 @@ async def _track_completed_batch_cost( } }, **({"api_base": mask_api_base_credentials(deployment_api_base)} if deployment_api_base else {}), - "_litellm_internal_model_credentials": MappingProxyType({**credentials}), "metadata": { **(await self._build_creator_attribution_metadata(job, batch_id)), # spend logs read the deployment identity off these metadata keys, so @@ -868,6 +867,11 @@ async def _track_completed_batch_cost( optional_params={}, custom_llm_provider=str(llm_provider) if llm_provider else None, ) + # deployment credentials stay off litellm_params because every callback sees those; the + # line-item logger reads them back off this private attribute to fetch the batch files + setattr( # noqa: B010 # Logging has no declared attribute for trusted credentials transport + logging_obj, "_litellm_internal_model_credentials", MappingProxyType(dict(credentials)) + ) if not await self._claim_job_for_costing(job): verbose_proxy_logger.info( diff --git a/litellm/batches/batch_line_item_logging.py b/litellm/batches/batch_line_item_logging.py index 65227c24c8ee..8a4206d3606d 100644 --- a/litellm/batches/batch_line_item_logging.py +++ b/litellm/batches/batch_line_item_logging.py @@ -299,7 +299,7 @@ async def log_batch_line_items( is logged and swallowed: aggregate accounting must be unaffected.""" emitted = 0 # rebind-ok: loop accumulator for emitted line count try: - internal_credentials: Final = ( + internal_credentials: Final = getattr(parent, "_litellm_internal_model_credentials", None) or ( litellm_params.get("_litellm_internal_model_credentials") if litellm_params else None ) internal_mapping: Final = _as_object_mapping(internal_credentials) diff --git a/tests/test_litellm/batches/test_batch_line_item_logging.py b/tests/test_litellm/batches/test_batch_line_item_logging.py index 32eeb326c606..fe03e348ce06 100644 --- a/tests/test_litellm/batches/test_batch_line_item_logging.py +++ b/tests/test_litellm/batches/test_batch_line_item_logging.py @@ -414,9 +414,9 @@ async def test_line_items_edge_shapes_and_edge_cases(recorder): parent: Final = _parent_logging_with_params( { "metadata": {"model_info": {"id": "dep-1"}, "model_group": "gpt-4o"}, - "_litellm_internal_model_credentials": {"api_key": "sk-line-items-marker"}, } ) + parent._litellm_internal_model_credentials = {"api_key": "sk-line-items-marker"} # test-quality-ok: private transport attribute, same channel the batch cost tracker uses with ( patch("litellm.files.main.afile_content", file_mock), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch patch("litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02)), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch @@ -446,6 +446,7 @@ async def test_line_items_edge_shapes_and_edge_cases(recorder): assert "no status here" in _payload(failure)["error_str"] assert "sk-line-items-marker" in str(file_mock.call_args_list) + assert "sk-line-items-marker" not in str(by_custom_id["e"]["litellm_params"]) file_ids_fetched = [call.kwargs.get("file_id") or call.args[0] for call in file_mock.call_args_list] assert "input-2" in file_ids_fetched and "output-2" in file_ids_fetched assert not any(file_id is None for file_id in file_ids_fetched) From ff9d30bc7992f77dbc7352079fb942cfe095ccb2 Mon Sep 17 00:00:00 2001 From: shivam Date: Thu, 17 Sep 2026 23:44:39 +0000 Subject: [PATCH 08/89] refactor(batches): declare credential transport attribute on Logging Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/common_utils/check_batch_cost.py | 6 +----- litellm/batches/batch_line_item_logging.py | 4 +--- litellm/litellm_core_utils/litellm_logging.py | 1 + 3 files changed, 3 insertions(+), 8 deletions(-) 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 b41b3a6e1ad5..d3bd830a9c0b 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -867,11 +867,7 @@ async def _track_completed_batch_cost( optional_params={}, custom_llm_provider=str(llm_provider) if llm_provider else None, ) - # deployment credentials stay off litellm_params because every callback sees those; the - # line-item logger reads them back off this private attribute to fetch the batch files - setattr( # noqa: B010 # Logging has no declared attribute for trusted credentials transport - logging_obj, "_litellm_internal_model_credentials", MappingProxyType(dict(credentials)) - ) + logging_obj._litellm_internal_model_credentials = MappingProxyType(dict(credentials)) # pyright: ignore[reportPrivateUsage] # trusted credentials transport, consumed by batch_line_item_logging if not await self._claim_job_for_costing(job): verbose_proxy_logger.info( diff --git a/litellm/batches/batch_line_item_logging.py b/litellm/batches/batch_line_item_logging.py index 8a4206d3606d..911a4d478ec1 100644 --- a/litellm/batches/batch_line_item_logging.py +++ b/litellm/batches/batch_line_item_logging.py @@ -299,9 +299,7 @@ async def log_batch_line_items( is logged and swallowed: aggregate accounting must be unaffected.""" emitted = 0 # rebind-ok: loop accumulator for emitted line count try: - internal_credentials: Final = getattr(parent, "_litellm_internal_model_credentials", None) or ( - litellm_params.get("_litellm_internal_model_credentials") if litellm_params else None - ) + internal_credentials: Final = parent._litellm_internal_model_credentials # pyright: ignore[reportPrivateUsage] # declared transport attribute on Logging internal_mapping: Final = _as_object_mapping(internal_credentials) fetch_params: Final[dict[str, object] | None] = ( # mutable-ok: file fetcher requires a plain dict dict(internal_mapping) # mutable-ok: the file fetcher reads credential kwargs off a plain dict diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index acba0eb10252..64cbf8eff4ba 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -574,6 +574,7 @@ def __init__( self.sync_streaming_chunks: list[Any] = [] # for generating complete stream response self.log_raw_request_response = log_raw_request_response self._native_callback_fast_path: bool = False + self._litellm_internal_model_credentials: Mapping[str, object] | None = None # Initialize dynamic callbacks self.dynamic_input_callbacks: list[str | Callable | CustomLogger] | None = dynamic_input_callbacks From 8484b2fb705fd1254a44c753ac219f86b735770f Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 23 Sep 2026 07:06:42 +0000 Subject: [PATCH 09/89] test(batches): integration contract for per-line batch callback events Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/integration/contracts.json | 4 + .../spend/test_batch_line_item_callbacks.py | 274 ++++++++++++++++++ 2 files changed, 278 insertions(+) create mode 100644 tests/integration/spend/test_batch_line_item_callbacks.py diff --git a/tests/integration/contracts.json b/tests/integration/contracts.json index a8d9cf1df8b9..828e9d8cab18 100644 --- a/tests/integration/contracts.json +++ b/tests/integration/contracts.json @@ -2027,6 +2027,10 @@ "mgmt.key.reset_spend.returns_promptly_with_wedged_coordination_redis", "mgmt.auth_cache_invalidation.publish_parked_by_short_redis_wedge_lands_after_recovery", "mgmt.auth_cache_invalidation.burst_with_worker_kill_keeps_serving_while_redis_wedged" + ], + "tests/integration/spend/test_batch_line_item_callbacks.py::test_completed_batch_emits_paired_request_response_callback_events_per_jsonl_line": [ + "spend.batches.completed_batch_emits_one_callback_event_per_jsonl_line_beside_the_aggregate", + "spend.batches.line_item_callback_events_do_not_bill_spend_twice" ] }, "browser": { diff --git a/tests/integration/spend/test_batch_line_item_callbacks.py b/tests/integration/spend/test_batch_line_item_callbacks.py new file mode 100644 index 000000000000..0407f68cc847 --- /dev/null +++ b/tests/integration/spend/test_batch_line_item_callbacks.py @@ -0,0 +1,274 @@ +from __future__ import annotations + +import json +import uuid +from hashlib import sha256 +from pathlib import Path +from typing import Final + +import pytest +import yaml +from integration._support.client import JSON_OBJECT, Gateway, eventually, object_value, string_value +from integration._support.database import read_rows +from integration._support.process import owned_proxy +from integration._support.upstream import delete_scenario, register_scenario +from integration._support.wire import Reply, Request, wire_server +from integration.cost_calculation.cost_tracking_case import JsonResponse, RoutedResponse, TextResponse +from pydantic import JsonValue + +PROMPT_TOKENS: Final = 10 +COMPLETION_TOKENS: Final = 100 +OUTPUT_SUCCESS_IDS: Final = ("r1", "r3") +OUTPUT_FAILED_IDS: Final = ("r2",) +ERROR_FILE_IDS: Final = ("r4", "r5") +ALL_CUSTOM_IDS: Final = ("r1", "r2", "r3", "r4", "r5") + + +def _successful_line(custom_id: str) -> str: + return json.dumps( + { + "id": f"batch_req_{custom_id}", + "custom_id": custom_id, + "response": { + "status_code": 200, + "request_id": f"$REQUEST_ID-{custom_id}", + "body": { + "id": f"chatcmpl-$REQUEST_ID-{custom_id}", + "object": "chat.completion", + "model": "gpt-4o-mini", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": f"answer {custom_id}"}, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": PROMPT_TOKENS, + "completion_tokens": COMPLETION_TOKENS, + "total_tokens": PROMPT_TOKENS + COMPLETION_TOKENS, + }, + }, + }, + "error": None, + }, + separators=(",", ":"), + ) + + +def _failed_line(custom_id: str) -> str: + return json.dumps( + { + "id": f"batch_req_{custom_id}", + "custom_id": custom_id, + "response": { + "status_code": 400, + "request_id": f"$REQUEST_ID-{custom_id}", + "body": {"error": {"message": f"synthetic line failure {custom_id}", "code": "bad_request"}}, + }, + "error": {"code": "bad_request", "message": f"synthetic line failure {custom_id}"}, + }, + separators=(",", ":"), + ) + + +def _batch(status: str, *, files_ready: bool) -> dict[str, JsonValue]: + return { + "id": "batch-$REQUEST_ID", + "object": "batch", + "endpoint": "/v1/chat/completions", + "errors": None, + "input_file_id": "file-in-$REQUEST_ID", + "completion_window": "24h", + "status": status, + "output_file_id": "file-out-$REQUEST_ID" if files_ready else None, + "error_file_id": "file-err-$REQUEST_ID" if files_ready else None, + "created_at": 1, + "in_progress_at": 1, + "completed_at": 1 if files_ready else None, + "expires_at": 1, + "request_counts": {"total": 5, "completed": 2, "failed": 3}, + "metadata": None, + } + + +def _input_lines(model_name: str, marker: str) -> tuple[str, ...]: + return tuple( + json.dumps( + { + "custom_id": custom_id, + "method": "POST", + "url": "/v1/chat/completions", + "body": {"model": model_name, "messages": [{"role": "user", "content": f"{marker} {custom_id}"}]}, + }, + separators=(",", ":"), + ) + for custom_id in ALL_CUSTOM_IDS + ) + + +def _provider_routes(input_lines: tuple[str, ...]) -> RoutedResponse: + output_lines: Final = (_successful_line("r1"), _failed_line("r2"), _successful_line("r3")) + error_lines: Final = tuple(_failed_line(custom_id) for custom_id in ERROR_FILE_IDS) + return RoutedResponse( + content_type="application/x-routed", + routes={ + "POST /files": JsonResponse( + content_type="application/json", + body={ + "id": "file-in-$REQUEST_ID", + "object": "file", + "purpose": "batch", + "bytes": 100, + "created_at": 1, + "filename": "in.jsonl", + "status": "processed", + }, + ), + "POST /batches": JsonResponse( + content_type="application/json", body=_batch("validating", files_ready=False) + ), + "GET /batches/batch-$REQUEST_ID": JsonResponse( + content_type="application/json", body=_batch("completed", files_ready=True) + ), + "GET /files/file-in-$REQUEST_ID/content": TextResponse( + content_type="application/jsonl", body="\n".join(input_lines) + "\n" + ), + "GET /files/file-out-$REQUEST_ID/content": TextResponse( + content_type="application/jsonl", body="\n".join(output_lines) + "\n" + ), + "GET /files/file-err-$REQUEST_ID/content": TextResponse( + content_type="application/jsonl", body="\n".join(error_lines) + "\n" + ), + }, + ) + + +def _line_item_config(base: Path, destination: Path) -> Path: + config: Final = yaml.safe_load(base.read_text()) + config["litellm_settings"].update({"callbacks": ["generic_api"], "DEFAULT_FLUSH_INTERVAL_SECONDS": 1}) + config["general_settings"].update({"store_batch_line_items_in_callbacks": True}) + destination.write_text(yaml.safe_dump(config)) + return destination + + +def _spend_rows(key: str) -> tuple[dict[str, JsonValue], ...]: + return tuple( + read_rows( + 'SELECT call_type, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" WHERE api_key=%s', + (sha256(key.encode()).hexdigest(),), + ) + ) + + +def _hidden(event: dict[str, JsonValue]) -> dict[str, JsonValue]: + return object_value(event["hidden_params"]) + + +@pytest.mark.covers( + "spend.batches.completed_batch_emits_one_callback_event_per_jsonl_line_beside_the_aggregate", + "spend.batches.line_item_callback_events_do_not_bill_spend_twice", +) +def test_completed_batch_emits_paired_request_response_callback_events_per_jsonl_line( + gateway: Gateway, tmp_path: Path +) -> None: + marker: Final = "batch-line-items-" + uuid.uuid4().hex[:12] + sink_secret: Final = "synthetic-sink-secret-" + marker + + def sink(request: Request) -> Reply: + assert request.headers["authorization"] == f"Bearer {sink_secret}" + return Reply() + + with ( + wire_server(sink) as endpoint, + owned_proxy( + gateway, + tmp_path, + {"GENERIC_LOGGER_ENDPOINT": endpoint.url, "GENERIC_LOGGER_HEADERS": f"Authorization=Bearer {sink_secret}"}, + config=_line_item_config(Path("tests/integration/proxy_config.yaml"), tmp_path / "line_items.yaml"), + ) as candidate, + candidate.scenario() as scenario, + ): + routed_model: Final = scenario.model(api_base=f"{gateway.upstream_url}/{marker}") + input_lines: Final = _input_lines(routed_model, marker) + handle: Final = register_scenario(marker, _provider_routes(input_lines)) + scenario.cleanups.callback(delete_scenario, handle) + key: Final = scenario.key(models=[routed_model]) + file_response: Final = candidate.request_multipart( + "/v1/files", + {"purpose": "batch", "model": routed_model}, + {"file": ("in.jsonl", ("\n".join(input_lines) + "\n").encode(), "application/jsonl")}, + key=key, + ) + assert file_response.status_code == 200, file_response.text + input_file_id: Final = string_value(JSON_OBJECT.validate_json(file_response.content)["id"]) + batch_response: Final = candidate.request( + "POST", + "/v1/batches", + { + "input_file_id": input_file_id, + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + "model": routed_model, + }, + key=key, + ) + assert batch_response.status_code == 200, batch_response.text + batch_id: Final = string_value(JSON_OBJECT.validate_json(batch_response.content)["id"]) + retrieval: Final = eventually( + lambda: candidate.request("GET", f"/v1/batches/{batch_id}", key=key), + lambda response: response.status_code == 200 and response.json()["status"] == "completed", + seconds=30, + ) + assert retrieval.status_code == 200, retrieval.text + key_hash: Final = sha256(key.encode()).hexdigest() + batches: Final[list[Request]] = [] # mutable-ok: drain() consumes the queue, later polls must keep earlier batches + + def delivered() -> tuple[dict[str, JsonValue], ...]: + batches.extend(endpoint.drain()) + return tuple( + event + for batch in batches + for event in json.loads(batch.body) + if object_value(event["metadata"]).get("user_api_key_hash") == key_hash + ) + + events: Final = eventually( + delivered, + lambda values: len([e for e in values if e["call_type"] == "aretrieve_batch"]) >= 1 + and len([e for e in values if _hidden(e).get("batch_custom_id") is not None]) >= len(ALL_CUSTOM_IDS), + seconds=40, + ) + aggregate_events: Final = tuple(event for event in events if event["call_type"] == "aretrieve_batch") + line_events: Final = tuple(event for event in events if event["call_type"] != "aretrieve_batch") + assert len(aggregate_events) == 1, [event["call_type"] for event in events] + assert sorted(string_value(_hidden(event)["batch_custom_id"]) for event in line_events) == list( + ALL_CUSTOM_IDS + ), line_events + by_custom_id: Final = {string_value(_hidden(event)["batch_custom_id"]): event for event in line_events} + for custom_id, line in zip(ALL_CUSTOM_IDS, input_lines, strict=True): + event: Final = by_custom_id[custom_id] + hidden: Final = _hidden(event) + assert hidden["batch_id"] == batch_id, hidden + assert event["call_type"] == "acompletion", event["call_type"] + assert event["messages"] == json.loads(line)["body"]["messages"], (custom_id, event["messages"]) + if custom_id in OUTPUT_SUCCESS_IDS: + assert event["status"] == "success", event + assert hidden["batch_line_status_code"] == 200, hidden + assert event["prompt_tokens"] == PROMPT_TOKENS and event["completion_tokens"] == COMPLETION_TOKENS + assert object_value(event["response"])["id"] == f"chatcmpl-{marker}-{custom_id}", event["response"] + assert object_value(object_value(event["response"])["choices"][0]["message"])["content"] == ( + f"answer {custom_id}" + ) + else: + assert event["status"] == "failure", event + assert hidden["batch_line_status_code"] == 400, hidden + assert f"synthetic line failure {custom_id}" in json.dumps(event["error_information"]), event + aggregate: Final = aggregate_events[0] + assert aggregate["prompt_tokens"] == len(OUTPUT_SUCCESS_IDS) * PROMPT_TOKENS, aggregate + assert aggregate["completion_tokens"] == len(OUTPUT_SUCCESS_IDS) * COMPLETION_TOKENS, aggregate + rows: Final = eventually( + lambda: _spend_rows(key), lambda values: any(r["call_type"] == "aretrieve_batch" for r in values), seconds=70 + ) + assert [row["call_type"] for row in rows] == ["aretrieve_batch"], rows + assert rows[0]["prompt_tokens"] == len(OUTPUT_SUCCESS_IDS) * PROMPT_TOKENS, rows From b4cb884d6017c60f06b4332c43c8a649f2c3424d Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 23 Sep 2026 07:10:34 +0000 Subject: [PATCH 10/89] fix(proxy): resolve store_batch_line_items_in_callbacks through the settings store Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/proxy_server.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index a7b875dc71cb..d8b6d0537f19 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -7593,7 +7593,7 @@ async def _apply_boolean_settings(self, db_values: Mapping[str, SettingsJsonValu async def _apply_batch_line_items_setting(self, db_values: Mapping[str, SettingsJsonValue]) -> None: key: Final = "store_batch_line_items_in_callbacks" - if key not in db_values or self.settings.owned_by_config(key): + if key not in db_values and not self.settings.owned_by_config(key): return value: Final = coerce_bool(self.settings.get(key)) if value is not None: From 61d5ab03d79df3b196a39b852d0ff42cc35022c0 Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 23 Sep 2026 07:19:04 +0000 Subject: [PATCH 11/89] test(batches): assert per-line events beside the aggregate and create events Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/integration/spend/test_batch_line_item_callbacks.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/integration/spend/test_batch_line_item_callbacks.py b/tests/integration/spend/test_batch_line_item_callbacks.py index 0407f68cc847..3e1e3ad4e089 100644 --- a/tests/integration/spend/test_batch_line_item_callbacks.py +++ b/tests/integration/spend/test_batch_line_item_callbacks.py @@ -240,8 +240,10 @@ def delivered() -> tuple[dict[str, JsonValue], ...]: seconds=40, ) aggregate_events: Final = tuple(event for event in events if event["call_type"] == "aretrieve_batch") - line_events: Final = tuple(event for event in events if event["call_type"] != "aretrieve_batch") + line_events: Final = tuple(event for event in events if _hidden(event).get("batch_custom_id") is not None) assert len(aggregate_events) == 1, [event["call_type"] for event in events] + other_events: Final = tuple(event for event in events if event not in aggregate_events + line_events) + assert sorted(event["call_type"] for event in other_events) == ["acreate_batch", "acreate_file"], other_events assert sorted(string_value(_hidden(event)["batch_custom_id"]) for event in line_events) == list( ALL_CUSTOM_IDS ), line_events From bc53e231ba9b33b108373fb97644ef77c4747c5e Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 23 Sep 2026 07:20:37 +0000 Subject: [PATCH 12/89] fix(proxy): drop resurrected use_legacy_interactions_schema shim from merge Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/proxy_server.py | 7 ------- 1 file changed, 7 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index d8b6d0537f19..c08a7b8156db 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -6516,13 +6516,6 @@ async def load_config(self, router: litellm.Router | None, config_file_path: str health_check_interval = general_settings.get("health_check_interval", DEFAULT_HEALTH_CHECK_INTERVAL) health_check_concurrency = general_settings.get("health_check_concurrency", None) health_check_details = general_settings.get("health_check_details", True) - ### INTERACTIONS API SCHEMA ### - _use_legacy_interactions_schema: Final = general_settings.get("use_legacy_interactions_schema") - if _use_legacy_interactions_schema is not None: - if isinstance(_use_legacy_interactions_schema, str): - litellm.use_legacy_interactions_schema = _use_legacy_interactions_schema.lower() == "true" - else: - litellm.use_legacy_interactions_schema = bool(_use_legacy_interactions_schema) ### BATCH LINE ITEM CALLBACKS ### _store_batch_line_items: Final = general_settings.get("store_batch_line_items_in_callbacks") if _store_batch_line_items is not None: From 4b6efef9deca2ed6a4739e38bd1d4d2ed736be19 Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 23 Sep 2026 07:23:22 +0000 Subject: [PATCH 13/89] test(batches): assert line item events carry no credentials Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/integration/spend/test_batch_line_item_callbacks.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/integration/spend/test_batch_line_item_callbacks.py b/tests/integration/spend/test_batch_line_item_callbacks.py index 3e1e3ad4e089..e0f328736651 100644 --- a/tests/integration/spend/test_batch_line_item_callbacks.py +++ b/tests/integration/spend/test_batch_line_item_callbacks.py @@ -174,6 +174,7 @@ def test_completed_batch_emits_paired_request_response_callback_events_per_jsonl ) -> None: marker: Final = "batch-line-items-" + uuid.uuid4().hex[:12] sink_secret: Final = "synthetic-sink-secret-" + marker + provider_secret: Final = "synthetic-provider-secret-" + marker def sink(request: Request) -> Reply: assert request.headers["authorization"] == f"Bearer {sink_secret}" @@ -189,7 +190,7 @@ def sink(request: Request) -> Reply: ) as candidate, candidate.scenario() as scenario, ): - routed_model: Final = scenario.model(api_base=f"{gateway.upstream_url}/{marker}") + routed_model: Final = scenario.model(api_base=f"{gateway.upstream_url}/{marker}", api_key=provider_secret) input_lines: Final = _input_lines(routed_model, marker) handle: Final = register_scenario(marker, _provider_routes(input_lines)) scenario.cleanups.callback(delete_scenario, handle) @@ -239,6 +240,9 @@ def delivered() -> tuple[dict[str, JsonValue], ...]: and len([e for e in values if _hidden(e).get("batch_custom_id") is not None]) >= len(ALL_CUSTOM_IDS), seconds=40, ) + for batch in batches: + for credential in (provider_secret, sink_secret, key): + assert credential.encode() not in batch.body aggregate_events: Final = tuple(event for event in events if event["call_type"] == "aretrieve_batch") line_events: Final = tuple(event for event in events if _hidden(event).get("batch_custom_id") is not None) assert len(aggregate_events) == 1, [event["call_type"] for event in events] From ed23e2e82dd208c6ef4f0a59374225ef3d3ffc4e Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 23 Sep 2026 08:00:25 +0000 Subject: [PATCH 14/89] fix(proxy): reset batch line-item flag when the setting is absent Removing store_batch_line_items_in_callbacks from the DB general settings left the module flag at its previous value because the effect early-returned when the key was neither in the DB nor owned by YAML. The proxy owns the global, so an absent resolved value now writes False. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/proxy_server.py | 5 +---- tests/test_litellm/proxy/test_proxy_server.py | 6 ++++++ 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index c08a7b8156db..9eab5478e3fc 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -7586,11 +7586,8 @@ async def _apply_boolean_settings(self, db_values: Mapping[str, SettingsJsonValu async def _apply_batch_line_items_setting(self, db_values: Mapping[str, SettingsJsonValue]) -> None: key: Final = "store_batch_line_items_in_callbacks" - if key not in db_values and not self.settings.owned_by_config(key): - return value: Final = coerce_bool(self.settings.get(key)) - if value is not None: - litellm.store_batch_line_items_in_callbacks = value if isinstance(value, bool) else bool(value) + litellm.store_batch_line_items_in_callbacks = bool(value) if value is not None else False async def _apply_cache_size_setting( self, diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 7fa63a952647..a52b119e02b5 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -7890,6 +7890,12 @@ async def test_update_general_settings_store_batch_line_items_in_callbacks(): db_general_settings={"store_batch_line_items_in_callbacks": True} ) assert litellm.store_batch_line_items_in_callbacks is False + + proxy_config._yaml_general_settings_keys = set() + litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: seed the prior opt-in so its removal is observable + with patch("litellm.proxy.proxy_server.general_settings", {}): # test-quality-ok: module-global seam + await proxy_config._update_general_settings(db_general_settings={}) + assert litellm.store_batch_line_items_in_callbacks is False finally: litellm.store_batch_line_items_in_callbacks = saved_flag From 096ba84f00636bf3e317e1f8fac42a4c3f3834ce Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 23 Sep 2026 08:00:26 +0000 Subject: [PATCH 15/89] fix(batches): map anthropic batch lines through the anthropic message transformer Anthropic result lines carry a message-shaped body, so ModelResponse(**body) produced empty choices and a usage key that blocked the batch-rate usage attach. Extract the existing AnthropicResponse -> ModelResponse conversion used for /v1/messages logging into anthropic_message_to_model_response and reuse it for anthropic line items. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/batches/batch_line_item_logging.py | 10 ++++-- litellm/litellm_core_utils/litellm_logging.py | 33 +++++++++++-------- .../batches/test_batch_line_item_logging.py | 6 ++++ 3 files changed, 34 insertions(+), 15 deletions(-) diff --git a/litellm/batches/batch_line_item_logging.py b/litellm/batches/batch_line_item_logging.py index 911a4d478ec1..b91ef3dccd34 100644 --- a/litellm/batches/batch_line_item_logging.py +++ b/litellm/batches/batch_line_item_logging.py @@ -121,11 +121,17 @@ def _line_model( _BatchLineResult: TypeAlias = "ModelResponse | EmbeddingResponse | ResponsesAPIResponse" -def _line_result(call_type: str, response_body: Mapping[str, object]) -> _BatchLineResult: +def _line_result( + call_type: str, custom_llm_provider: _BatchLineProvider, response_body: Mapping[str, object] +) -> _BatchLineResult: if call_type == "aembedding": return EmbeddingResponse(**response_body) # pyright: ignore[reportArgumentType] # provider output bodies are dicts expanded as response ctor kwargs if call_type == "aresponses": return ResponsesAPIResponse(**response_body) # pyright: ignore[reportArgumentType] # same as above + if custom_llm_provider == "anthropic": + from litellm.litellm_core_utils.litellm_logging import anthropic_message_to_model_response + + return anthropic_message_to_model_response(response_body, None) return ModelResponse(**response_body) # pyright: ignore[reportArgumentType] # same as above @@ -235,7 +241,7 @@ async def _emit_line_event( stats: Final = _safe_output_line_stats(entry, custom_llm_provider, model_name, model_info) try: - result: Final = _line_result(call_type, response_body) + result: Final = _line_result(call_type, custom_llm_provider, response_body) except Exception: # noqa: BLE001 # one unparseable line must not drop the rest of the batch's line events verbose_logger.warning( "batch output line could not be reconstructed as a %s response, skipping it. custom_id=%s", diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 06c8218eea5f..11b1d4725f77 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -535,6 +535,24 @@ def _provider_response_id(source: object) -> str | None: return candidate if isinstance(candidate, str) and candidate else None +def anthropic_message_to_model_response(result: Mapping[str, object], speed: str | None) -> ModelResponse: + import httpx + + from litellm.types.llms.anthropic import AnthropicResponse + + pydantic_result: Final = AnthropicResponse.model_validate(result) + return litellm.AnthropicConfig().transform_parsed_response( + completion_response=pydantic_result.model_dump(), + raw_response=httpx.Response( + status_code=200, + headers={}, + ), + model_response=litellm.ModelResponse(id=_provider_response_id(result)), + json_mode=None, + speed=speed, + ) + + def mask_api_base_credentials(api_base: str) -> str: if "key=" not in api_base: return api_base @@ -4237,19 +4255,8 @@ def _handle_anthropic_messages_response_logging(self, result: Any) -> ModelRespo litellm_params={}, ) else: - from litellm.types.llms.anthropic import AnthropicResponse - - pydantic_result: Final = AnthropicResponse.model_validate(result) - import httpx - - result = litellm.AnthropicConfig().transform_parsed_response( - completion_response=pydantic_result.model_dump(), - raw_response=httpx.Response( - status_code=200, - headers={}, - ), - model_response=litellm.ModelResponse(id=provider_response_id), - json_mode=None, + result = anthropic_message_to_model_response( + cast(Mapping[str, object], result), speed=self.optional_params.get("speed") if self.optional_params else None, ) return result diff --git a/tests/test_litellm/batches/test_batch_line_item_logging.py b/tests/test_litellm/batches/test_batch_line_item_logging.py index fe03e348ce06..62c6bd359f26 100644 --- a/tests/test_litellm/batches/test_batch_line_item_logging.py +++ b/tests/test_litellm/batches/test_batch_line_item_logging.py @@ -368,6 +368,8 @@ async def test_input_fetch_failure_still_emits_aggregate(recorder): "role": "assistant", "content": [{"type": "text", "text": "hello b2"}], "model": "claude-3", + "stop_reason": "end_turn", + "stop_sequence": None, "usage": {"input_tokens": 1, "output_tokens": 2}, }, }, @@ -482,3 +484,7 @@ async def test_line_items_anthropic_shapes(recorder): line = next(e for e in recorder.success_events if _hidden(e).get("batch_custom_id") == "b2") assert _hidden(line)["batch_line_status_code"] == 200 assert line["litellm_params"]["batch_parent_id"] == batch.id + payload = _payload(line) + assert payload["response"]["choices"][0]["message"]["content"] == "hello b2" + assert payload["prompt_tokens"] == 1 + assert payload["completion_tokens"] == 2 From a8fafaf75032bf6919c33530310d1bfe527381d0 Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 23 Sep 2026 08:06:27 +0000 Subject: [PATCH 16/89] style: tag the anthropic result cast with a cast-ok reason Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/litellm_core_utils/litellm_logging.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 11b1d4725f77..2cf1b7e46735 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -4256,7 +4256,7 @@ def _handle_anthropic_messages_response_logging(self, result: Any) -> ModelRespo ) else: result = anthropic_message_to_model_response( - cast(Mapping[str, object], result), + cast(Mapping[str, object], result), # cast-ok: this handler receives heterogeneous provider results typed as Any upstream speed=self.optional_params.get("speed") if self.optional_params else None, ) return result From c746e5720baf7116e32bf63a95bc7a075eb01e3e Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 23 Sep 2026 08:27:51 +0000 Subject: [PATCH 17/89] style: shorten cast-ok reason so ruff format passes Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/litellm_core_utils/litellm_logging.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 2cf1b7e46735..de740fdeca7a 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -4256,7 +4256,7 @@ def _handle_anthropic_messages_response_logging(self, result: Any) -> ModelRespo ) else: result = anthropic_message_to_model_response( - cast(Mapping[str, object], result), # cast-ok: this handler receives heterogeneous provider results typed as Any upstream + cast(Mapping[str, object], result), # cast-ok: handler result is typed Any upstream speed=self.optional_params.get("speed") if self.optional_params else None, ) return result From 45b7509692ad7f57b7d9440883973afdc736e080 Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 23 Sep 2026 09:00:38 +0000 Subject: [PATCH 18/89] fix(batches): skip line-item callbacks for providers without a line reconstruction Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/batches/batch_line_item_logging.py | 35 ++++++++++++------- .../batches/test_batch_line_item_logging.py | 18 ++++++++-- 2 files changed, 39 insertions(+), 14 deletions(-) diff --git a/litellm/batches/batch_line_item_logging.py b/litellm/batches/batch_line_item_logging.py index b91ef3dccd34..a414c4197105 100644 --- a/litellm/batches/batch_line_item_logging.py +++ b/litellm/batches/batch_line_item_logging.py @@ -3,7 +3,7 @@ from collections.abc import Iterator, Mapping from datetime import datetime from types import MappingProxyType -from typing import TYPE_CHECKING, Final, Literal, TypeAlias +from typing import TYPE_CHECKING, Final, Literal, TypeAlias, cast, get_args from litellm._logging import verbose_logger from litellm.batches.batch_utils import ( @@ -27,6 +27,15 @@ _BatchLineProvider: TypeAlias = Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"] +_SUPPORTED_LINE_PROVIDERS: Final = frozenset(get_args(_BatchLineProvider)) + + +def _supported_line_provider(value: str) -> _BatchLineProvider | None: + if value in _SUPPORTED_LINE_PROVIDERS: + return cast("_BatchLineProvider", value) # cast-ok: membership in the literal's args was just checked + return None + + _CALL_TYPE_BY_BATCH_URL: Final = MappingProxyType( { "/v1/chat/completions": "acompletion", @@ -291,7 +300,7 @@ async def _fetch_managed_file_or_empty( async def log_batch_line_items( batch: LiteLLMBatch, - custom_llm_provider: _BatchLineProvider, + custom_llm_provider: str, parent: "Logging", model_name: str | None, litellm_params: dict[str, object] | None, # mutable-ok: the logging object's shared litellm_params dict @@ -303,6 +312,14 @@ async def log_batch_line_items( aretrieve_batch event still bills the batch, so per-line events carry ``batch_parent_id`` and never update spend themselves. Any failure here is logged and swallowed: aggregate accounting must be unaffected.""" + line_provider: Final = _supported_line_provider(custom_llm_provider) + if line_provider is None: + verbose_logger.warning( + "batch line-item callbacks are not supported for provider %s, skipping. batch_id=%s", + custom_llm_provider, + batch.id, + ) + return 0 emitted = 0 # rebind-ok: loop accumulator for emitted line count try: internal_credentials: Final = parent._litellm_internal_model_credentials # pyright: ignore[reportPrivateUsage] # declared transport attribute on Logging @@ -313,17 +330,11 @@ async def log_batch_line_items( else litellm_params ) - input_file_content: Final = await _fetch_managed_file_or_empty( - batch.input_file_id, custom_llm_provider, fetch_params - ) + input_file_content: Final = await _fetch_managed_file_or_empty(batch.input_file_id, line_provider, fetch_params) requests_by_id: Final = _requests_by_custom_id(input_file_content) - output_content: Final = await _fetch_managed_file_or_empty( - batch.output_file_id, custom_llm_provider, fetch_params - ) - error_content: Final = await _fetch_managed_file_or_empty( - batch.error_file_id, custom_llm_provider, fetch_params - ) + output_content: Final = await _fetch_managed_file_or_empty(batch.output_file_id, line_provider, fetch_params) + error_content: Final = await _fetch_managed_file_or_empty(batch.error_file_id, line_provider, fetch_params) for content in (output_content, error_content): for entry in _output_entries(content): try: @@ -331,7 +342,7 @@ async def log_batch_line_items( entry=entry, requests_by_id=requests_by_id, batch=batch, - custom_llm_provider=custom_llm_provider, + custom_llm_provider=line_provider, parent=parent, model_name=model_name, model_info=model_info, diff --git a/tests/test_litellm/batches/test_batch_line_item_logging.py b/tests/test_litellm/batches/test_batch_line_item_logging.py index 62c6bd359f26..1c59eac24bb2 100644 --- a/tests/test_litellm/batches/test_batch_line_item_logging.py +++ b/tests/test_litellm/batches/test_batch_line_item_logging.py @@ -134,7 +134,7 @@ def recorder(): litellm._async_failure_callback = saved_failure # test-quality-ok: teardown restoring the value set above -def _parent_logging() -> Logging: +def _parent_logging(custom_llm_provider: str = "openai") -> Logging: logging_obj = Logging( model="gpt-4o", messages=[{"role": "user", "content": ""}], @@ -147,7 +147,7 @@ def _parent_logging() -> Logging: logging_obj.update_environment_variables( litellm_params={"metadata": {"model_info": {"id": "dep-1"}, "model_group": "gpt-4o"}}, optional_params={}, - custom_llm_provider="openai", + custom_llm_provider=custom_llm_provider, ) return logging_obj @@ -223,6 +223,20 @@ async def test_flag_off_emits_only_aggregate(recorder): file_mock.assert_not_called() +@pytest.mark.asyncio +async def test_line_items_skipped_for_unsupported_provider(recorder): + litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: the flag under test is a module global; fixture restores it + file_mock: Final = AsyncMock(side_effect=_file_content) + with patch("litellm.files.main.afile_content", file_mock): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + await _log_completed_batch(_parent_logging(custom_llm_provider="bedrock"), _batch()) + + file_mock.assert_not_awaited() + assert len(recorder.success_events) == 1 + assert _payload(recorder.success_events[0])["response_cost"] == 1.5 + assert "batch_custom_id" not in _hidden(recorder.success_events[0]) + assert len(recorder.failure_events) == 0 + + @pytest.mark.asyncio async def test_in_progress_batch_poll_emits_no_line_events(recorder): litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: the flag under test is a module global; fixture restores it From bef2f0489fdf5ac8b2b2d87c96f0cc5ab3e9aaaa Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 23 Sep 2026 10:40:51 +0000 Subject: [PATCH 19/89] fix(batches): reconstruct bedrock and mistral batch lines through their provider transformations Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/batches/batch_line_item_logging.py | 23 ++- litellm/batches/batch_utils.py | 10 +- litellm/litellm_core_utils/litellm_logging.py | 20 +-- litellm/llms/anthropic/chat/transformation.py | 13 ++ .../llms/bedrock/batches/transformation.py | 37 ++++- .../batches/test_batch_line_item_logging.py | 156 +++++++++++++++++- 6 files changed, 229 insertions(+), 30 deletions(-) diff --git a/litellm/batches/batch_line_item_logging.py b/litellm/batches/batch_line_item_logging.py index a414c4197105..02dd9cceaeb8 100644 --- a/litellm/batches/batch_line_item_logging.py +++ b/litellm/batches/batch_line_item_logging.py @@ -25,7 +25,9 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging -_BatchLineProvider: TypeAlias = Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"] +_BatchLineProvider: TypeAlias = Literal[ + "openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral" +] _SUPPORTED_LINE_PROVIDERS: Final = frozenset(get_args(_BatchLineProvider)) @@ -131,14 +133,24 @@ def _line_model( def _line_result( - call_type: str, custom_llm_provider: _BatchLineProvider, response_body: Mapping[str, object] + call_type: str, + custom_llm_provider: _BatchLineProvider, + model: str, + response_body: Mapping[str, object], ) -> _BatchLineResult: + if custom_llm_provider == "bedrock": + from litellm.llms.bedrock.batches.transformation import bedrock_batch_line_to_response + + bedrock_result: Final = bedrock_batch_line_to_response(response_body, model) + if bedrock_result is None: + raise ValueError(f"unrecognized bedrock batch output line shape. keys={sorted(response_body)}") + return bedrock_result if call_type == "aembedding": return EmbeddingResponse(**response_body) # pyright: ignore[reportArgumentType] # provider output bodies are dicts expanded as response ctor kwargs if call_type == "aresponses": return ResponsesAPIResponse(**response_body) # pyright: ignore[reportArgumentType] # same as above if custom_llm_provider == "anthropic": - from litellm.litellm_core_utils.litellm_logging import anthropic_message_to_model_response + from litellm.llms.anthropic.chat.transformation import anthropic_message_to_model_response return anthropic_message_to_model_response(response_body, None) return ModelResponse(**response_body) # pyright: ignore[reportArgumentType] # same as above @@ -216,9 +228,10 @@ async def _emit_line_event( start_time: Final = parent_start_time if isinstance(parent_start_time, datetime) else datetime.now() # noqa: DTZ005 # naive to match the logging pipeline start_time parent_params: Final = _as_object_mapping(parent.litellm_params) or _EMPTY_BODY # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # Logging.litellm_params is untyped upstream + model: Final = _line_model(response_body, request_body, parent) child: Final = _new_child_logging( parent=parent, - model=_line_model(response_body, request_body, parent), + model=model, messages=_line_messages(request_body), call_type=call_type, start_time=start_time, @@ -250,7 +263,7 @@ async def _emit_line_event( stats: Final = _safe_output_line_stats(entry, custom_llm_provider, model_name, model_info) try: - result: Final = _line_result(call_type, custom_llm_provider, response_body) + result: Final = _line_result(call_type, custom_llm_provider, model, response_body) except Exception: # noqa: BLE001 # one unparseable line must not drop the rest of the batch's line events verbose_logger.warning( "batch output line could not be reconstructed as a %s response, skipping it. custom_id=%s", diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 7209ac6a1e73..468a0c63eaee 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -53,7 +53,7 @@ def batch_cost_is_final(batch: Batch) -> bool: async def calculate_batch_cost_and_usage( file_content_dictionary: list[dict], - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "mistral"], + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral"], model_name: str | None = None, model_info: ModelInfo | None = None, ) -> BatchCostUsageResult: @@ -83,7 +83,7 @@ async def calculate_batch_cost_and_usage( async def _handle_completed_batch( batch: Batch, - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "mistral"], + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral"], model_name: str | None = None, litellm_params: dict | None = None, model_info: ModelInfo | None = None, @@ -449,7 +449,7 @@ def _provider_output_file_id(output_file_id: str) -> str: async def _fetch_batch_managed_file_content( file_id: str, - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "mistral"] = "openai", + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral"] = "openai", litellm_params: dict | None = None, ) -> bytes: """ @@ -479,7 +479,7 @@ async def _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", + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral"] = "openai", litellm_params: dict | None = None, ) -> bytes: """ @@ -501,7 +501,7 @@ async def _fetch_batch_output_file_content( async def count_error_file_failed_requests( batch: Batch, - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "mistral"], + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral"], litellm_params: dict | None, ) -> int: """Count failed requests reported only in the batch's separate error file. diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index de740fdeca7a..ea410d6d6ba2 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -535,24 +535,6 @@ def _provider_response_id(source: object) -> str | None: return candidate if isinstance(candidate, str) and candidate else None -def anthropic_message_to_model_response(result: Mapping[str, object], speed: str | None) -> ModelResponse: - import httpx - - from litellm.types.llms.anthropic import AnthropicResponse - - pydantic_result: Final = AnthropicResponse.model_validate(result) - return litellm.AnthropicConfig().transform_parsed_response( - completion_response=pydantic_result.model_dump(), - raw_response=httpx.Response( - status_code=200, - headers={}, - ), - model_response=litellm.ModelResponse(id=_provider_response_id(result)), - json_mode=None, - speed=speed, - ) - - def mask_api_base_credentials(api_base: str) -> str: if "key=" not in api_base: return api_base @@ -4255,6 +4237,8 @@ def _handle_anthropic_messages_response_logging(self, result: Any) -> ModelRespo litellm_params={}, ) else: + from litellm.llms.anthropic.chat.transformation import anthropic_message_to_model_response + result = anthropic_message_to_model_response( cast(Mapping[str, object], result), # cast-ok: handler result is typed Any upstream speed=self.optional_params.get("speed") if self.optional_params else None, diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index b7c2ce3c5689..7f9294c58dfa 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -48,6 +48,7 @@ AnthropicMessagesToolChoice, AnthropicOutputSchema, AnthropicOutputTokensDetails, + AnthropicResponse, AnthropicSystemMessageContent, AnthropicThinkingParam, AnthropicWebSearchTool, @@ -2789,3 +2790,15 @@ def _valid_user_id(user_id: str) -> bool: return False return True + + +def anthropic_message_to_model_response(result: Mapping[str, object], speed: str | None) -> ModelResponse: + pydantic_result: Final = AnthropicResponse.model_validate(result) + result_id: Final = result.get("id") + return AnthropicConfig().transform_parsed_response( + completion_response=pydantic_result.model_dump(), + raw_response=httpx.Response(status_code=200, headers={}), + model_response=ModelResponse(id=result_id if isinstance(result_id, str) and result_id else None), + json_mode=None, + speed=speed, + ) diff --git a/litellm/llms/bedrock/batches/transformation.py b/litellm/llms/bedrock/batches/transformation.py index e4001566b8c0..6d43e9ee7f90 100644 --- a/litellm/llms/bedrock/batches/transformation.py +++ b/litellm/llms/bedrock/batches/transformation.py @@ -27,7 +27,7 @@ AllMessageValues, CreateBatchRequest, ) -from litellm.types.utils import LiteLLMBatch, LlmProviders, Usage +from litellm.types.utils import EmbeddingResponse, LiteLLMBatch, LlmProviders, ModelResponse, Usage from ..base_aws_llm import BaseAWSLLM from ..common_utils import ( @@ -96,6 +96,41 @@ def titan_embedding_usage_from_batch_output(model_output: Mapping[str, object]) ) +def bedrock_batch_line_to_response( + model_output: Mapping[str, object], model: str +) -> ModelResponse | EmbeddingResponse | None: + """Reconstruct a Bedrock batch output line (the ``modelOutput`` object) into + the litellm response type its shape implies, or None when the shape is + unrecognized.""" + if "embedding" in model_output: + embedding: Final = model_output.get("embedding") + return EmbeddingResponse( + model=model, + data=[{"object": "embedding", "index": 0, "embedding": embedding if isinstance(embedding, list) else []}], + usage=titan_embedding_usage_from_batch_output(model_output), + ) + if "output" in model_output: + from ..chat.converse_transformation import AmazonConverseConfig + + return AmazonConverseConfig()._transform_response( # pyright: ignore[reportPrivateUsage] # same reconstruction the converse chat path performs on the live response + model=model, + response=Response(200, json=dict(model_output)), + model_response=ModelResponse(), + stream=False, + logging_obj=None, + optional_params={}, + api_key=None, + data="", + messages=[], + encoding=None, + ) + if "content" in model_output: + from litellm.llms.anthropic.chat.transformation import anthropic_message_to_model_response + + return anthropic_message_to_model_response(model_output, None) + return None + + class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): """ Config for Bedrock Batches - handles batch job creation and management for Bedrock diff --git a/tests/test_litellm/batches/test_batch_line_item_logging.py b/tests/test_litellm/batches/test_batch_line_item_logging.py index 1c59eac24bb2..0be7f76b5947 100644 --- a/tests/test_litellm/batches/test_batch_line_item_logging.py +++ b/tests/test_litellm/batches/test_batch_line_item_logging.py @@ -228,7 +228,7 @@ async def test_line_items_skipped_for_unsupported_provider(recorder): litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: the flag under test is a module global; fixture restores it file_mock: Final = AsyncMock(side_effect=_file_content) with patch("litellm.files.main.afile_content", file_mock): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch - await _log_completed_batch(_parent_logging(custom_llm_provider="bedrock"), _batch()) + await _log_completed_batch(_parent_logging(custom_llm_provider="cohere"), _batch()) file_mock.assert_not_awaited() assert len(recorder.success_events) == 1 @@ -391,6 +391,94 @@ async def test_input_fetch_failure_still_emits_aggregate(recorder): ).encode() +BEDROCK_INPUT_JSONL = b"\n".join( + [ + json.dumps( + { + "recordId": "br-anth", + "modelInput": {"messages": [{"role": "user", "content": "hi br-anth"}]}, + } + ).encode(), + json.dumps( + { + "recordId": "br-titan", + "modelInput": {"inputText": "embed me"}, + } + ).encode(), + json.dumps( + { + "recordId": "br-conv", + "modelInput": {"messages": [{"role": "user", "content": "hi br-conv"}]}, + } + ).encode(), + ] +) + +BEDROCK_OUTPUT_JSONL = b"\n".join( + [ + json.dumps( + { + "recordId": "br-anth", + "modelOutput": { + "id": "msg_br", + "type": "message", + "role": "assistant", + "content": [{"type": "text", "text": "bedrock claude hi"}], + "model": "anthropic.claude-3", + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 2, "output_tokens": 3}, + }, + } + ).encode(), + json.dumps( + { + "recordId": "br-titan", + "modelOutput": {"embedding": [0.1, 0.2], "inputTextTokenCount": 4}, + } + ).encode(), + json.dumps( + { + "recordId": "br-conv", + "modelOutput": { + "output": {"message": {"role": "assistant", "content": [{"text": "converse hi"}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 5, "outputTokens": 6, "totalTokens": 11}, + "metrics": {"latencyMs": 12}, + }, + } + ).encode(), + ] +) + +MISTRAL_INPUT_JSONL = json.dumps( + { + "custom_id": "mis", + "body": {"model": "mistral-small", "messages": [{"role": "user", "content": "hi mis"}]}, + } +).encode() + +MISTRAL_OUTPUT_JSONL = json.dumps( + { + "custom_id": "mis", + "response": { + "status_code": 200, + "body": { + "id": "mis-1", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "mistral hi"}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5}, + }, + }, + } +).encode() + + def _edge_file_content(file_id: str, **_kwargs): return SimpleNamespace( content={ @@ -398,6 +486,10 @@ def _edge_file_content(file_id: str, **_kwargs): "output-2": EDGE_OUTPUT_JSONL, "input-anth": ANTHROPIC_INPUT_JSONL, "output-anth": ANTHROPIC_OUTPUT_JSONL, + "input-bed": BEDROCK_INPUT_JSONL, + "output-bed": BEDROCK_OUTPUT_JSONL, + "input-mist": MISTRAL_INPUT_JSONL, + "output-mist": MISTRAL_OUTPUT_JSONL, }[file_id] ) @@ -502,3 +594,65 @@ async def test_line_items_anthropic_shapes(recorder): assert payload["response"]["choices"][0]["message"]["content"] == "hello b2" assert payload["prompt_tokens"] == 1 assert payload["completion_tokens"] == 2 + + +def _provider_batch(batch_id: str, input_file_id: str, output_file_id: str) -> LiteLLMBatch: + return LiteLLMBatch( + id=batch_id, + object="batch", + endpoint="/v1/chat/completions", + input_file_id=input_file_id, + output_file_id=output_file_id, + error_file_id=None, + status="completed", + completion_window="24h", + created_at=1, + ) + + +@pytest.mark.asyncio +async def test_line_items_bedrock_shapes(recorder): + litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: the flag under test is a module global; fixture restores it + file_mock: Final = AsyncMock(side_effect=_edge_file_content) + with ( + patch("litellm.files.main.afile_content", file_mock), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + patch("litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02)), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch + ): + await _log_completed_batch( + _parent_logging(custom_llm_provider="bedrock"), + _provider_batch("batch_bed", "input-bed", "output-bed"), + ) + + by_id = {_hidden(e).get("batch_custom_id"): e for e in recorder.success_events} + + anthropic_line = _payload(by_id["br-anth"]) + assert anthropic_line["response"]["choices"][0]["message"]["content"] == "bedrock claude hi" + assert anthropic_line["prompt_tokens"] == 2 + assert anthropic_line["completion_tokens"] == 3 + + titan_line = _payload(by_id["br-titan"]) + assert titan_line["response"]["data"][0]["embedding"] == [0.1, 0.2] + assert titan_line["prompt_tokens"] == 4 + + converse_line = _payload(by_id["br-conv"]) + assert converse_line["response"]["choices"][0]["message"]["content"] == "converse hi" + assert converse_line["prompt_tokens"] == 5 + assert converse_line["completion_tokens"] == 6 + + +@pytest.mark.asyncio +async def test_line_items_mistral_shape(recorder): + litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: the flag under test is a module global; fixture restores it + file_mock: Final = AsyncMock(side_effect=_edge_file_content) + with ( + patch("litellm.files.main.afile_content", file_mock), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + patch("litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02)), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch + ): + await _log_completed_batch( + _parent_logging(custom_llm_provider="mistral"), + _provider_batch("batch_mist", "input-mist", "output-mist"), + ) + + line = next(e for e in recorder.success_events if _hidden(e).get("batch_custom_id") == "mis") + assert _hidden(line)["batch_line_status_code"] == 200 + assert _payload(line)["response"]["choices"][0]["message"]["content"] == "mistral hi" From 6cbad5cbab4d83be599c3f54b34c1f675d0b4899 Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 23 Sep 2026 11:02:31 +0000 Subject: [PATCH 20/89] style: wrap widened provider literals so ruff format passes Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/batches/batch_utils.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 468a0c63eaee..240ed048b632 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -449,7 +449,9 @@ def _provider_output_file_id(output_file_id: str) -> str: async def _fetch_batch_managed_file_content( file_id: str, - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral"] = "openai", + custom_llm_provider: Literal[ + "openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral" + ] = "openai", litellm_params: dict | None = None, ) -> bytes: """ @@ -479,7 +481,9 @@ async def _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", "bedrock", "mistral"] = "openai", + custom_llm_provider: Literal[ + "openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral" + ] = "openai", litellm_params: dict | None = None, ) -> bytes: """ From a760991a7a7041c68775008168af847589e629d3 Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 23 Sep 2026 13:12:25 +0000 Subject: [PATCH 21/89] fix(batches): pair bedrock lines by recordId and type them from the reconstructed response Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/batches/batch_line_item_logging.py | 93 +++++++++++++------ .../llms/base_llm/batches/transformation.py | 11 ++- .../llms/bedrock/batches/transformation.py | 5 + .../batches/test_batch_line_item_logging.py | 17 +++- .../bedrock/batches/test_transformation.py | 24 ++++- 5 files changed, 116 insertions(+), 34 deletions(-) diff --git a/litellm/batches/batch_line_item_logging.py b/litellm/batches/batch_line_item_logging.py index 02dd9cceaeb8..3d54625d3fee 100644 --- a/litellm/batches/batch_line_item_logging.py +++ b/litellm/batches/batch_line_item_logging.py @@ -71,15 +71,16 @@ def _output_entries(file_content: bytes) -> Iterator[Mapping[str, object]]: yield mapping +def _line_id(entry: Mapping[str, object]) -> str | None: + line_id: Final = entry.get("custom_id") or entry.get("recordId") + return line_id if isinstance(line_id, str) and line_id else None + + def _requests_by_custom_id(input_file_content: bytes) -> Mapping[str, Mapping[str, object]]: - """Parse the batch input JSONL into {custom_id: request line}, skipping - malformed lines and lines without a custom_id.""" + """Parse the batch input JSONL into {line id: request line}, keyed by + custom_id or recordId, skipping malformed lines and lines without either.""" return MappingProxyType( - { - custom_id: entry - for entry in _output_entries(input_file_content) - if isinstance((custom_id := entry.get("custom_id")), str) and custom_id - } + {line_id: entry for entry in _output_entries(input_file_content) if (line_id := _line_id(entry)) is not None} ) @@ -93,6 +94,9 @@ def _request_body_for_entry( params: Final = _as_object_mapping(request_line.get("params")) if params: return params + request_model_input: Final = _as_object_mapping(request_line.get("modelInput")) + if request_model_input: + return request_model_input model_input: Final = _as_object_mapping(entry.get("modelInput")) return model_input if model_input else _EMPTY_BODY @@ -114,6 +118,16 @@ def _call_type_for_request(request_line: Mapping[str, object] | None) -> str: return _CALL_TYPE_BY_BATCH_URL.get(url if isinstance(url, str) else "", "acompletion") +def _call_type_for_line(request_line: Mapping[str, object] | None, result: "_BatchLineResult | None") -> str: + if isinstance(result, EmbeddingResponse): + return "aembedding" + if isinstance(result, ResponsesAPIResponse): + return "aresponses" + if isinstance(result, ModelResponse): + return "acompletion" + return _call_type_for_request(request_line) + + def _line_messages(request_body: Mapping[str, object]) -> object: return request_body.get("messages") or request_body.get("input") or () @@ -138,13 +152,16 @@ def _line_result( model: str, response_body: Mapping[str, object], ) -> _BatchLineResult: - if custom_llm_provider == "bedrock": - from litellm.llms.bedrock.batches.transformation import bedrock_batch_line_to_response - - bedrock_result: Final = bedrock_batch_line_to_response(response_body, model) - if bedrock_result is None: + from litellm.types.utils import LlmProviders + from litellm.utils import ProviderConfigManager + + provider_config: Final = ProviderConfigManager.get_provider_batches_config(model, LlmProviders(custom_llm_provider)) + if provider_config is not None: + transformed: Final = provider_config.transform_batch_output_line(response_body, model) + if transformed is not None: + return transformed + if custom_llm_provider == "bedrock": raise ValueError(f"unrecognized bedrock batch output line shape. keys={sorted(response_body)}") - return bedrock_result if call_type == "aembedding": return EmbeddingResponse(**response_body) # pyright: ignore[reportArgumentType] # provider output bodies are dicts expanded as response ctor kwargs if call_type == "aresponses": @@ -156,6 +173,24 @@ def _line_result( return ModelResponse(**response_body) # pyright: ignore[reportArgumentType] # same as above +def _line_result_or_none( + request_call_type: str, + custom_llm_provider: _BatchLineProvider, + model: str, + response_body: Mapping[str, object], + custom_id: object, +) -> "_BatchLineResult | None": + try: + return _line_result(request_call_type, custom_llm_provider, model, response_body) + except Exception: # noqa: BLE001 # one unparseable line must not drop the rest of the batch's line events + verbose_logger.warning( + "batch output line could not be reconstructed as a %s response, skipping it. custom_id=%s", + request_call_type, + custom_id, + ) + return None + + def _new_child_logging( parent: "Logging", model: str, @@ -218,17 +253,28 @@ async def _emit_line_event( model_name: str | None, model_info: ModelInfo | None, ) -> bool: - custom_id: Final = entry.get("custom_id") or entry.get("recordId") - request_line: Final = requests_by_id.get(custom_id if isinstance(custom_id, str) else "") + custom_id: Final = _line_id(entry) + request_line: Final = requests_by_id.get(custom_id or "") request_body: Final = _request_body_for_entry(entry, request_line) status_code: Final = _line_status_code(entry, custom_llm_provider) - call_type: Final = _call_type_for_request(request_line) + request_call_type: Final = _call_type_for_request(request_line) response_body: Final = _get_response_from_batch_job_output_file(entry, custom_llm_provider) parent_start_time: Final = parent.start_time # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # Logging.start_time is untyped upstream start_time: Final = parent_start_time if isinstance(parent_start_time, datetime) else datetime.now() # noqa: DTZ005 # naive to match the logging pipeline start_time parent_params: Final = _as_object_mapping(parent.litellm_params) or _EMPTY_BODY # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # Logging.litellm_params is untyped upstream - model: Final = _line_model(response_body, request_body, parent) + + successful: Final = _batch_response_was_successful(entry, custom_llm_provider) + stats: Final = _safe_output_line_stats(entry, custom_llm_provider, model_name, model_info) if successful else None + result: Final[_BatchLineResult | None] = ( + _line_result_or_none(request_call_type, custom_llm_provider, model, response_body, custom_id) + if successful + else None + ) + if successful and result is None: + return False + + call_type: Final = _call_type_for_line(request_line, result) child: Final = _new_child_logging( parent=parent, model=model, @@ -248,7 +294,7 @@ async def _emit_line_event( ) now: Final = datetime.now() # noqa: DTZ005 # naive to match the logging pipeline start_time - if not _batch_response_was_successful(entry, custom_llm_provider): + if result is None: exception: Final = _BatchLineFailure( entry.get("error") or entry.get("response") or {} # mutable-ok: fallback payload dict passed to Exception ) @@ -261,17 +307,6 @@ async def _emit_line_event( ) return True - stats: Final = _safe_output_line_stats(entry, custom_llm_provider, model_name, model_info) - try: - result: Final = _line_result(call_type, custom_llm_provider, model, response_body) - except Exception: # noqa: BLE001 # one unparseable line must not drop the rest of the batch's line events - verbose_logger.warning( - "batch output line could not be reconstructed as a %s response, skipping it. custom_id=%s", - call_type, - custom_id, - ) - return False - result._hidden_params = _line_hidden_params( # pyright: ignore[reportPrivateUsage] # same hidden_params channel the aggregate batch event uses batch, custom_id, diff --git a/litellm/llms/base_llm/batches/transformation.py b/litellm/llms/base_llm/batches/transformation.py index 34c622d4cf6b..904eae2f4f27 100644 --- a/litellm/llms/base_llm/batches/transformation.py +++ b/litellm/llms/base_llm/batches/transformation.py @@ -1,5 +1,6 @@ import types from abc import ABC, abstractmethod +from collections.abc import Mapping from typing import TYPE_CHECKING, Any import httpx @@ -9,7 +10,7 @@ AllMessageValues, CreateBatchRequest, ) -from litellm.types.utils import LiteLLMBatch, LlmProviders +from litellm.types.utils import EmbeddingResponse, LiteLLMBatch, LlmProviders, ModelResponse if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj @@ -39,6 +40,14 @@ def __init__(self): def custom_llm_provider(self) -> LlmProviders: """Return the LLM provider type for this configuration.""" + def transform_batch_output_line( + self, model_output: Mapping[str, object], model: str + ) -> ModelResponse | EmbeddingResponse | None: + """Reconstruct one provider batch output line into a litellm response, or + None when the line is OpenAI-shaped and the caller can use the generic + reconstruction.""" + return None + @classmethod def get_config(cls): """Get configuration dictionary for this class.""" diff --git a/litellm/llms/bedrock/batches/transformation.py b/litellm/llms/bedrock/batches/transformation.py index 6d43e9ee7f90..f8d71bab2c72 100644 --- a/litellm/llms/bedrock/batches/transformation.py +++ b/litellm/llms/bedrock/batches/transformation.py @@ -144,6 +144,11 @@ def __init__(self): def custom_llm_provider(self) -> LlmProviders: return LlmProviders.BEDROCK + def transform_batch_output_line( + self, model_output: Mapping[str, object], model: str + ) -> ModelResponse | EmbeddingResponse | None: + return bedrock_batch_line_to_response(model_output, model) + @classmethod def _get_bare_model_name_from_s3_key(cls, object_key: str) -> str | None: if not object_key.startswith(BEDROCK_MANAGED_S3_BATCH_PREFIX): diff --git a/tests/test_litellm/batches/test_batch_line_item_logging.py b/tests/test_litellm/batches/test_batch_line_item_logging.py index 0be7f76b5947..1f37969bcfb6 100644 --- a/tests/test_litellm/batches/test_batch_line_item_logging.py +++ b/tests/test_litellm/batches/test_batch_line_item_logging.py @@ -625,16 +625,27 @@ async def test_line_items_bedrock_shapes(recorder): by_id = {_hidden(e).get("batch_custom_id"): e for e in recorder.success_events} - anthropic_line = _payload(by_id["br-anth"]) + anthropic_event = by_id["br-anth"] + assert anthropic_event["call_type"] == "acompletion" + assert _hidden(anthropic_event)["batch_custom_id"] == "br-anth" + anthropic_line = _payload(anthropic_event) + assert any(m.get("content") == "hi br-anth" for m in anthropic_line["messages"]) assert anthropic_line["response"]["choices"][0]["message"]["content"] == "bedrock claude hi" assert anthropic_line["prompt_tokens"] == 2 assert anthropic_line["completion_tokens"] == 3 - titan_line = _payload(by_id["br-titan"]) + titan_event = by_id["br-titan"] + assert titan_event["call_type"] == "aembedding" + assert _hidden(titan_event)["batch_custom_id"] == "br-titan" + titan_line = _payload(titan_event) + assert titan_event["optional_params"]["inputText"] == "embed me" assert titan_line["response"]["data"][0]["embedding"] == [0.1, 0.2] assert titan_line["prompt_tokens"] == 4 - converse_line = _payload(by_id["br-conv"]) + converse_event = by_id["br-conv"] + assert converse_event["call_type"] == "acompletion" + converse_line = _payload(converse_event) + assert any(m.get("content") == "hi br-conv" for m in converse_line["messages"]) assert converse_line["response"]["choices"][0]["message"]["content"] == "converse hi" assert converse_line["prompt_tokens"] == 5 assert converse_line["completion_tokens"] == 6 diff --git a/tests/test_litellm/llms/bedrock/batches/test_transformation.py b/tests/test_litellm/llms/bedrock/batches/test_transformation.py index 347c459a3693..4bce3c37504b 100644 --- a/tests/test_litellm/llms/bedrock/batches/test_transformation.py +++ b/tests/test_litellm/llms/bedrock/batches/test_transformation.py @@ -20,7 +20,7 @@ import pytest from litellm.llms.bedrock.batches.transformation import BedrockBatchesConfig -from litellm.types.utils import LlmProviders +from litellm.types.utils import EmbeddingResponse, LlmProviders, ModelResponse # AWS JobStatus -> OpenAI BatchJobStatus, exactly as encoded in transformation.py # (both transform_create_batch_response and transform_retrieve_batch_response). @@ -941,3 +941,25 @@ def test_retrieve_request_accepts_partition_arns(config: BedrockBatchesConfig, a batch_id=arn, optional_params={}, litellm_params={} ) assert result["url"].startswith(expected_prefix) + + +def test_transform_batch_output_line_dispatches_on_shape(config: BedrockBatchesConfig) -> None: + titan = config.transform_batch_output_line( + {"embedding": [0.1, 0.2], "inputTextTokenCount": 4}, model="amazon.titan-embed" + ) + assert isinstance(titan, EmbeddingResponse) + assert titan.data[0]["embedding"] == [0.1, 0.2] + assert titan.usage.prompt_tokens == 4 + + converse = config.transform_batch_output_line( + { + "output": {"message": {"role": "assistant", "content": [{"text": "hi"}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 5, "outputTokens": 6, "totalTokens": 11}, + }, + model="amazon.nova-lite", + ) + assert isinstance(converse, ModelResponse) + assert converse.choices[0].message.content == "hi" + + assert config.transform_batch_output_line({"foo": 1}, model="x") is None From c0b69aedb92e395872a6a9ff582314538dc2aa5d Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 23 Sep 2026 15:38:04 +0000 Subject: [PATCH 22/89] fix(batches): keep line item events out of proxy limiters and strip credentials from child params Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/batches/batch_line_item_logging.py | 18 ++++++++++ litellm/litellm_core_utils/core_helpers.py | 15 +++++++- litellm/llms/anthropic/chat/transformation.py | 4 ++- .../llms/bedrock/batches/transformation.py | 10 +++--- .../proxy/hooks/dynamic_rate_limiter_v3.py | 3 ++ .../hooks/max_budget_per_session_limiter.py | 3 ++ .../proxy/hooks/model_max_budget_limiter.py | 3 ++ .../proxy/hooks/parallel_request_limiter.py | 7 +++- .../hooks/parallel_request_limiter_v3.py | 3 ++ .../proxy/hooks/proxy_track_cost_callback.py | 6 ++-- .../batches/test_batch_line_item_logging.py | 35 ++++++++++++++++++- .../litellm_core_utils/test_core_helpers.py | 11 ++++++ .../hooks/test_model_max_budget_limiter.py | 14 ++++++++ .../hooks/test_parallel_request_limiter_v3.py | 27 ++++++++++++++ 14 files changed, 147 insertions(+), 12 deletions(-) diff --git a/litellm/batches/batch_line_item_logging.py b/litellm/batches/batch_line_item_logging.py index 3d54625d3fee..98b6f0a17571 100644 --- a/litellm/batches/batch_line_item_logging.py +++ b/litellm/batches/batch_line_item_logging.py @@ -31,6 +31,20 @@ _SUPPORTED_LINE_PROVIDERS: Final = frozenset(get_args(_BatchLineProvider)) +_SECRET_PARAM_KEYS: Final = frozenset( + { + "api_key", + "_litellm_internal_model_credentials", + "azure_ad_token", + "azure_ad_token_provider", + "vertex_credentials", + "aws_access_key_id", + "aws_secret_access_key", + "aws_session_token", + "aws_web_identity_token", + } +) + def _supported_line_provider(value: str) -> _BatchLineProvider | None: if value in _SUPPORTED_LINE_PROVIDERS: @@ -292,6 +306,10 @@ async def _emit_line_event( model=child.model, custom_llm_provider=custom_llm_provider, ) + for secret_key in _SECRET_PARAM_KEYS: + child.litellm_params.pop( + secret_key, None + ) # mutable-ok: model_call_details holds this same dict # pyright: ignore[reportUnknownMemberType] # Logging.litellm_params is untyped upstream now: Final = datetime.now() # noqa: DTZ005 # naive to match the logging pipeline start_time if result is None: diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index d7fbe9f7e098..088a4a414283 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -5,7 +5,7 @@ import re from collections.abc import Collection, Iterable, Mapping from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Literal, Protocol +from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, cast import httpx from pydantic import TypeAdapter, ValidationError @@ -765,3 +765,16 @@ def set_response_cost_in_hidden_params(response: _CarriesHiddenParams, cost: flo RESPONSE_COST_HEADER: cost, } hidden_params["additional_headers"] = merged # rebind-ok: the caller's record is the point + + +def is_batch_line_item_event(kwargs: object) -> bool: + if not isinstance(kwargs, Mapping): + return False + litellm_params: Final = cast(Mapping[str, object], kwargs).get( + "litellm_params" + ) # cast-ok: isinstance narrows only to unparameterized Mapping + if not isinstance(litellm_params, Mapping): + return False + return bool( + cast(Mapping[str, object], litellm_params).get("batch_parent_id") + ) # cast-ok: same narrowing limit as above diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index 7f9294c58dfa..14c28b0ce68b 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -2797,7 +2797,9 @@ def anthropic_message_to_model_response(result: Mapping[str, object], speed: str result_id: Final = result.get("id") return AnthropicConfig().transform_parsed_response( completion_response=pydantic_result.model_dump(), - raw_response=httpx.Response(status_code=200, headers={}), + raw_response=httpx.Response( + status_code=200, headers={} + ), # mutable-ok: httpx.Response wants a plain dict of headers model_response=ModelResponse(id=result_id if isinstance(result_id, str) and result_id else None), json_mode=None, speed=speed, diff --git a/litellm/llms/bedrock/batches/transformation.py b/litellm/llms/bedrock/batches/transformation.py index f8d71bab2c72..c61cd0dfeef2 100644 --- a/litellm/llms/bedrock/batches/transformation.py +++ b/litellm/llms/bedrock/batches/transformation.py @@ -106,7 +106,9 @@ def bedrock_batch_line_to_response( embedding: Final = model_output.get("embedding") return EmbeddingResponse( model=model, - data=[{"object": "embedding", "index": 0, "embedding": embedding if isinstance(embedding, list) else []}], + data=[ + {"object": "embedding", "index": 0, "embedding": embedding if isinstance(embedding, list) else []} + ], # mutable-ok: EmbeddingResponse takes a plain data list usage=titan_embedding_usage_from_batch_output(model_output), ) if "output" in model_output: @@ -114,14 +116,14 @@ def bedrock_batch_line_to_response( return AmazonConverseConfig()._transform_response( # pyright: ignore[reportPrivateUsage] # same reconstruction the converse chat path performs on the live response model=model, - response=Response(200, json=dict(model_output)), + response=Response(200, json=dict(model_output)), # mutable-ok: httpx.Response json= wants a plain dict model_response=ModelResponse(), stream=False, logging_obj=None, - optional_params={}, + optional_params={}, # mutable-ok: the converse transform signature takes a plain dict api_key=None, data="", - messages=[], + messages=[], # mutable-ok: the converse transform signature takes a plain list encoding=None, ) if "content" in model_output: diff --git a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py index 0339cf4dfea7..6777e0c88f38 100644 --- a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py +++ b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py @@ -14,6 +14,7 @@ from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.core_helpers import is_batch_line_item_event from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.proxy_rate_limit_error import ( ProxyRateLimitError, @@ -693,6 +694,8 @@ async def async_log_success_event(self, kwargs, response_obj, start_time, end_ti - model_saturation_check: Model-wide token tracking - priority_model: Priority-specific token tracking """ + if is_batch_line_item_event(kwargs): + return from litellm.litellm_core_utils.core_helpers import ( _get_parent_otel_span_from_kwargs, ) diff --git a/litellm/proxy/hooks/max_budget_per_session_limiter.py b/litellm/proxy/hooks/max_budget_per_session_limiter.py index e07b96e5773e..a1f60fc4d874 100644 --- a/litellm/proxy/hooks/max_budget_per_session_limiter.py +++ b/litellm/proxy/hooks/max_budget_per_session_limiter.py @@ -23,6 +23,7 @@ from litellm.caching.redis_cache import log_redis_failure from litellm.exceptions import RateLimitType from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.core_helpers import is_batch_line_item_event from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit @@ -131,6 +132,8 @@ async def async_log_success_event(self, kwargs, response_obj, start_time, end_ti """ After a successful LLM call, increment the session spend by the response cost. """ + if is_batch_line_item_event(kwargs): # pyright: ignore[reportUnknownArgumentType] # hook kwargs arrive untyped from the logging dispatcher + return try: litellm_params: Final = kwargs.get("litellm_params") or {} metadata: Final = litellm_params.get("metadata") or {} diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index cfa54ae01a2b..de3d72ddb5b4 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/litellm/proxy/hooks/model_max_budget_limiter.py @@ -11,6 +11,7 @@ from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import Span +from litellm.litellm_core_utils.core_helpers import is_batch_line_item_event from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.llms.bedrock.common_utils import get_bedrock_base_model from litellm.proxy._types import Litellm_EntityType, UserAPIKeyAuth @@ -486,6 +487,8 @@ async def async_log_success_event(self, kwargs, response_obj, start_time, end_ti Example: key=sk-1234567890, model=gpt-4o, max_budget=100, time_period=1d """ + if is_batch_line_item_event(kwargs): + return verbose_proxy_logger.debug("in RouterBudgetLimiting.async_log_success_event") standard_logging_payload: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object", None) if standard_logging_payload is None: diff --git a/litellm/proxy/hooks/parallel_request_limiter.py b/litellm/proxy/hooks/parallel_request_limiter.py index d41acadc4dda..a6dde5853091 100644 --- a/litellm/proxy/hooks/parallel_request_limiter.py +++ b/litellm/proxy/hooks/parallel_request_limiter.py @@ -11,7 +11,10 @@ from litellm._logging import verbose_proxy_logger from litellm.exceptions import RateLimitType from litellm.integrations.custom_logger import CustomLogger -from litellm.litellm_core_utils.core_helpers import _get_parent_otel_span_from_kwargs +from litellm.litellm_core_utils.core_helpers import ( + _get_parent_otel_span_from_kwargs, + is_batch_line_item_event, +) from litellm.proxy._types import CommonProxyErrors, CurrentItemRateLimit, UserAPIKeyAuth from litellm.proxy.auth.auth_utils import ( get_key_model_rpm_limit, @@ -489,6 +492,8 @@ async def async_pre_call_hook( ) async def async_log_success_event(self, kwargs, response_obj: object, start_time, end_time): + if is_batch_line_item_event(kwargs): + return from litellm.proxy.common_utils.callback_utils import ( get_model_group_from_litellm_kwargs, ) diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 33744906b13e..d5b31624064a 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -33,6 +33,7 @@ from litellm.caching.redis_cache import log_redis_failure from litellm.constants import DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE, INTERNAL_CALL_ORIGIN_METADATA_KEY from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.core_helpers import is_batch_line_item_event from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_str_from_messages, ) @@ -4557,6 +4558,8 @@ async def async_log_success_event(self, kwargs, response_obj, start_time, end_ti """ Update TPM usage on successful API calls by incrementing counters using pipeline """ + if is_batch_line_item_event(kwargs): + return from litellm.litellm_core_utils.core_helpers import ( _get_parent_otel_span_from_kwargs, ) diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 66df6578bc89..e8d508feb07b 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -13,6 +13,7 @@ _get_parent_otel_span_from_kwargs, budget_reservation_from_metadata, get_litellm_metadata_from_kwargs, + is_batch_line_item_event, ) from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import guardrail_information_cost @@ -106,10 +107,7 @@ def __init__( async def async_log_success_event( self, kwargs: ObjectMapping, response_obj: object, start_time: datetime, end_time: datetime ) -> None: - # Per-line batch events emitted under store_batch_line_items_in_callbacks - # never touch spend: the aggregate aretrieve_batch event already bills the batch. - litellm_params: Final = kwargs.get("litellm_params") - if isinstance(litellm_params, Mapping) and litellm_params.get("batch_parent_id"): + if is_batch_line_item_event(kwargs): return if self.spend_event_producer is None or not is_offloadable_success(response_obj): await self._PROXY_track_cost_callback(kwargs, response_obj, start_time, end_time) diff --git a/tests/test_litellm/batches/test_batch_line_item_logging.py b/tests/test_litellm/batches/test_batch_line_item_logging.py index 1f37969bcfb6..eb8ae26222db 100644 --- a/tests/test_litellm/batches/test_batch_line_item_logging.py +++ b/tests/test_litellm/batches/test_batch_line_item_logging.py @@ -14,7 +14,7 @@ import json import uuid from datetime import datetime -from types import SimpleNamespace +from types import MappingProxyType, SimpleNamespace from typing import Final from unittest.mock import AsyncMock, patch @@ -560,6 +560,39 @@ async def test_line_items_edge_shapes_and_edge_cases(recorder): assert not any(file_id is None for file_id in file_ids_fetched) +@pytest.mark.asyncio +async def test_line_items_child_params_drop_parent_credentials(recorder): + litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: the flag under test is a module global; fixture restores it + file_mock: Final = AsyncMock(side_effect=_file_content) + parent: Final = _parent_logging_with_params( + { + "api_key": "sk-parent-secret", + "_litellm_internal_model_credentials": MappingProxyType({"api_key": "sk-parent-secret"}), + "api_base": "https://api.openai.com/v1", + "metadata": {"model_info": {"id": "dep-1"}, "model_group": "gpt-4o"}, + } + ) + batch: Final = _batch() + with ( + patch("litellm.files.main.afile_content", file_mock), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + patch("litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02)), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch + ): + await _log_completed_batch(parent, batch) + + line_events = [ + e + for e in [*recorder.success_events, *recorder.failure_events] + if _hidden(e).get("batch_custom_id") is not None + ] + assert len(line_events) == 2 + for event in line_events: + params = event["litellm_params"] + assert "api_key" not in params + assert "_litellm_internal_model_credentials" not in params + assert params["batch_parent_id"] == batch.id + assert params["metadata"]["model_group"] == "gpt-4o" + + @pytest.mark.asyncio async def test_line_items_anthropic_shapes(recorder): litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: the flag under test is a module global; fixture restores it diff --git a/tests/test_litellm/litellm_core_utils/test_core_helpers.py b/tests/test_litellm/litellm_core_utils/test_core_helpers.py index 6eeea271127d..368248fd9567 100644 --- a/tests/test_litellm/litellm_core_utils/test_core_helpers.py +++ b/tests/test_litellm/litellm_core_utils/test_core_helpers.py @@ -11,6 +11,7 @@ drop_params_env_flag, drop_params_flag, get_or_create_metadata_bucket, + is_batch_line_item_event, map_finish_reason, normalize_drop_params, reconstruct_model_name, @@ -489,3 +490,13 @@ def test_budget_rejection_decorated_with_provider_is_expected(self): category=RateLimitErrorCategory.VENDOR_RATE_LIMIT, ) assert is_expected_client_error(vendor_limit) is False + + +def test_is_batch_line_item_event(): + assert is_batch_line_item_event({"litellm_params": {"batch_parent_id": "batch_1"}}) is True + assert is_batch_line_item_event({"litellm_params": {"batch_parent_id": "batch_1", "metadata": {}}}) is True + assert is_batch_line_item_event({"litellm_params": {"metadata": {}}}) is False + assert is_batch_line_item_event({"litellm_params": {"batch_parent_id": None}}) is False + assert is_batch_line_item_event({}) is False + assert is_batch_line_item_event({"litellm_params": "not-a-mapping"}) is False + assert is_batch_line_item_event({"litellm_params": None}) is False diff --git a/tests/test_litellm/proxy/hooks/test_model_max_budget_limiter.py b/tests/test_litellm/proxy/hooks/test_model_max_budget_limiter.py index ffb60fb46517..b062efb3e81b 100644 --- a/tests/test_litellm/proxy/hooks/test_model_max_budget_limiter.py +++ b/tests/test_litellm/proxy/hooks/test_model_max_budget_limiter.py @@ -196,3 +196,17 @@ async def test_a_batch_polled_within_every_budget_window_is_never_charged_again( await _poll(limiter, finished, BATCH_COST) assert _local_spend(limiter, KEY_SPEND_KEY) == pytest.approx(BATCH_COST) + + +@pytest.mark.asyncio +async def test_batch_line_item_events_do_not_charge_the_model_budget(): + """Line events carry batch_parent_id; the aggregate aretrieve_batch event is + the one that already bills the batch, so children must not double-charge.""" + limiter: Final = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache()) + + line_event = _event("acompletion", CHAT_COST) + line_event["litellm_params"]["batch_parent_id"] = "batch_first" + await limiter.async_log_success_event(line_event, response_obj=None, start_time=None, end_time=None) + + assert await _spend(limiter, KEY_SPEND_KEY) == 0.0 + assert await _spend(limiter, USER_SPEND_KEY) == 0.0 diff --git a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py index 8c0dcd3383c1..904b5d67556d 100644 --- a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py @@ -7150,3 +7150,30 @@ async def test_success_tpm_accounting_keeps_the_admission_target_after_an_alias_ charged: Final = {op["key"]: op["increment_value"] for op in ops} assert charged[admission_bucket] == 150 - stash.reserved_tokens assert not any(":target-b" in key for key in charged) + + +@pytest.mark.asyncio +async def test_async_log_success_event_skips_batch_line_item_events(): + """Per-line batch events already ran through the limiter as the aggregate + aretrieve_batch; children must not increment TPM or request counters.""" + local_cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + + await handler.async_log_success_event( + kwargs={ + "standard_logging_object": {"metadata": {"user_api_key_hash": hash_token("sk-line-item")}}, + "litellm_params": { + "batch_parent_id": "batch_1", + "metadata": {"user_api_key_hash": hash_token("sk-line-item"), "model_group": "gpt-3.5-turbo"}, + }, + "model": "gpt-3.5-turbo", + }, + response_obj=ModelResponse( + id="x", object="chat.completion", created=1, model="gpt-3.5-turbo", + usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15), choices=[], + ), + start_time=datetime.now(), + end_time=datetime.now(), + ) + + assert local_cache.in_memory_cache.cache_dict == {} From 67d4471a378d12bce2f061ac39dcc2633742acf1 Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 23 Sep 2026 15:43:47 +0000 Subject: [PATCH 23/89] fix(batches): keep cast-ok tags on the flagged cast lines for the discipline gate Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/litellm_core_utils/core_helpers.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index 088a4a414283..0bcf222ec6b6 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -770,11 +770,9 @@ def set_response_cost_in_hidden_params(response: _CarriesHiddenParams, cost: flo def is_batch_line_item_event(kwargs: object) -> bool: if not isinstance(kwargs, Mapping): return False - litellm_params: Final = cast(Mapping[str, object], kwargs).get( - "litellm_params" - ) # cast-ok: isinstance narrows only to unparameterized Mapping + typed_kwargs: Final = cast(Mapping[str, object], kwargs) # cast-ok: isinstance leaves Mapping unparameterized + litellm_params: Final = typed_kwargs.get("litellm_params") if not isinstance(litellm_params, Mapping): return False - return bool( - cast(Mapping[str, object], litellm_params).get("batch_parent_id") - ) # cast-ok: same narrowing limit as above + typed_params: Final = cast(Mapping[str, object], litellm_params) # cast-ok: same narrowing limit + return bool(typed_params.get("batch_parent_id")) From ee8dcfc20da1d7d9938795804e1a7ec48b746b16 Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 23 Sep 2026 15:48:58 +0000 Subject: [PATCH 24/89] fix(batches): skip failed line item events in parallel request limiter failure hooks Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/hooks/parallel_request_limiter.py | 2 ++ .../hooks/parallel_request_limiter_v3.py | 2 ++ .../hooks/test_parallel_request_limiter.py | 26 +++++++++++++++++ .../hooks/test_parallel_request_limiter_v3.py | 29 +++++++++++++++++++ 4 files changed, 59 insertions(+) diff --git a/litellm/proxy/hooks/parallel_request_limiter.py b/litellm/proxy/hooks/parallel_request_limiter.py index a6dde5853091..4ad2062a4e2c 100644 --- a/litellm/proxy/hooks/parallel_request_limiter.py +++ b/litellm/proxy/hooks/parallel_request_limiter.py @@ -699,6 +699,8 @@ async def async_log_success_event(self, kwargs, response_obj: object, start_time self.print_verbose(e) async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + if is_batch_line_item_event(kwargs): + return try: self.print_verbose("Inside Max Parallel Request Failure Hook") litellm_parent_otel_span: Final[Span | None] = _get_parent_otel_span_from_kwargs(kwargs=kwargs) diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index d5b31624064a..981c0b81a078 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -4681,6 +4681,8 @@ async def async_log_failure_event(self, kwargs, response_obj, start_time, end_ti whose partial usage was recovered settles the reservation at that usage instead of refunding it. """ + if is_batch_line_item_event(kwargs): + return from litellm.litellm_core_utils.core_helpers import ( _get_parent_otel_span_from_kwargs, ) diff --git a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter.py b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter.py index 0e2683dcbfd6..cbc16d734a1f 100644 --- a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter.py +++ b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter.py @@ -84,3 +84,29 @@ async def test_async_log_success_event_counts_non_chat_response_tokens(response_ f"expected 50 tokens counted for {scope_id}, " f"got {current['current_tpm']}" ) + + +@pytest.mark.asyncio +async def test_async_log_failure_event_skips_batch_line_item_events(): + """Failed line children were never admitted by the limiter, so the failure + hook must not decrement request counters they never incremented.""" + parallel_request_handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(DualCache()) + ) + local_cache = parallel_request_handler.internal_usage_cache.dual_cache.in_memory_cache + + await parallel_request_handler.async_log_failure_event( + kwargs={ + "exception": "litellm.APIError: upstream 500", + "litellm_params": { + "batch_parent_id": "batch_1", + "metadata": {"user_api_key": hash_token("sk-line-item")}, + }, + "model": "gpt-3.5-turbo", + }, + response_obj=None, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + assert local_cache.cache_dict == {} diff --git a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py index 904b5d67556d..ea8d50164631 100644 --- a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py @@ -7177,3 +7177,32 @@ async def test_async_log_success_event_skips_batch_line_item_events(): ) assert local_cache.in_memory_cache.cache_dict == {} + + +@pytest.mark.asyncio +async def test_async_log_failure_event_skips_batch_line_item_events(): + """Failed line children were never admitted by the limiter, so the failure + hook must not release slots or refund TPM they never reserved.""" + local_cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) + stash = get_or_create_request_stash() + stash.owner_litellm_call_id = "call-line-item" + stash.reserved_tokens = 50 + stash.reserved_scopes = frozenset({("api_key", hash_token("sk-line-item"))}) + + await handler.async_log_failure_event( + kwargs={ + "litellm_call_id": "call-line-item", + "standard_logging_object": {"metadata": {"user_api_key_hash": hash_token("sk-line-item")}}, + "litellm_params": { + "batch_parent_id": "batch_1", + "metadata": {"user_api_key_hash": hash_token("sk-line-item"), "model_group": "gpt-3.5-turbo"}, + }, + "model": "gpt-3.5-turbo", + }, + response_obj=None, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + assert local_cache.in_memory_cache.cache_dict == {} From a2fd2a8d63b5f1f9d1a8a288f27bfdd2aafc88c6 Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 23 Sep 2026 18:02:08 +0000 Subject: [PATCH 25/89] test(batches): tolerate lifecycle spend rows when asserting no line item rows Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../spend/test_batch_line_item_callbacks.py | 20 +++++++++++++------ 1 file changed, 14 insertions(+), 6 deletions(-) diff --git a/tests/integration/spend/test_batch_line_item_callbacks.py b/tests/integration/spend/test_batch_line_item_callbacks.py index e0f328736651..97cad137b516 100644 --- a/tests/integration/spend/test_batch_line_item_callbacks.py +++ b/tests/integration/spend/test_batch_line_item_callbacks.py @@ -223,7 +223,9 @@ def sink(request: Request) -> Reply: ) assert retrieval.status_code == 200, retrieval.text key_hash: Final = sha256(key.encode()).hexdigest() - batches: Final[list[Request]] = [] # mutable-ok: drain() consumes the queue, later polls must keep earlier batches + batches: Final[ + list[Request] + ] = [] # mutable-ok: drain() consumes the queue, later polls must keep earlier batches def delivered() -> tuple[dict[str, JsonValue], ...]: batches.extend(endpoint.drain()) @@ -236,8 +238,10 @@ def delivered() -> tuple[dict[str, JsonValue], ...]: events: Final = eventually( delivered, - lambda values: len([e for e in values if e["call_type"] == "aretrieve_batch"]) >= 1 - and len([e for e in values if _hidden(e).get("batch_custom_id") is not None]) >= len(ALL_CUSTOM_IDS), + lambda values: ( + len([e for e in values if e["call_type"] == "aretrieve_batch"]) >= 1 + and len([e for e in values if _hidden(e).get("batch_custom_id") is not None]) >= len(ALL_CUSTOM_IDS) + ), seconds=40, ) for batch in batches: @@ -274,7 +278,11 @@ def delivered() -> tuple[dict[str, JsonValue], ...]: assert aggregate["prompt_tokens"] == len(OUTPUT_SUCCESS_IDS) * PROMPT_TOKENS, aggregate assert aggregate["completion_tokens"] == len(OUTPUT_SUCCESS_IDS) * COMPLETION_TOKENS, aggregate rows: Final = eventually( - lambda: _spend_rows(key), lambda values: any(r["call_type"] == "aretrieve_batch" for r in values), seconds=70 + lambda: _spend_rows(key), + lambda values: any(r["call_type"] == "aretrieve_batch" for r in values), + seconds=70, ) - assert [row["call_type"] for row in rows] == ["aretrieve_batch"], rows - assert rows[0]["prompt_tokens"] == len(OUTPUT_SUCCESS_IDS) * PROMPT_TOKENS, rows + batch_rows: Final = tuple(row for row in rows if row["call_type"] == "aretrieve_batch") + assert len(batch_rows) == 1, rows + assert not any(row["call_type"] == "acompletion" for row in rows), rows + assert batch_rows[0]["prompt_tokens"] == len(OUTPUT_SUCCESS_IDS) * PROMPT_TOKENS, rows From 05d6977052aac902a53fa14d14c9eddb7c173019 Mon Sep 17 00:00:00 2001 From: yucheng Date: Thu, 24 Sep 2026 01:03:31 +0000 Subject: [PATCH 26/89] fix(batches): place mutable-ok tags on the lines the discipline gate flags Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/batches/batch_line_item_logging.py | 4 +--- litellm/llms/anthropic/chat/transformation.py | 5 +++-- litellm/llms/bedrock/batches/transformation.py | 12 +++++++++--- 3 files changed, 13 insertions(+), 8 deletions(-) diff --git a/litellm/batches/batch_line_item_logging.py b/litellm/batches/batch_line_item_logging.py index 98b6f0a17571..cf4de4863289 100644 --- a/litellm/batches/batch_line_item_logging.py +++ b/litellm/batches/batch_line_item_logging.py @@ -307,9 +307,7 @@ async def _emit_line_event( custom_llm_provider=custom_llm_provider, ) for secret_key in _SECRET_PARAM_KEYS: - child.litellm_params.pop( - secret_key, None - ) # mutable-ok: model_call_details holds this same dict # pyright: ignore[reportUnknownMemberType] # Logging.litellm_params is untyped upstream + child.litellm_params.pop(secret_key, None) # pyright: ignore[reportUnknownMemberType] # Logging.litellm_params is untyped upstream now: Final = datetime.now() # noqa: DTZ005 # naive to match the logging pipeline start_time if result is None: diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index 14c28b0ce68b..77b8f906cbe9 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -2798,8 +2798,9 @@ def anthropic_message_to_model_response(result: Mapping[str, object], speed: str return AnthropicConfig().transform_parsed_response( completion_response=pydantic_result.model_dump(), raw_response=httpx.Response( - status_code=200, headers={} - ), # mutable-ok: httpx.Response wants a plain dict of headers + status_code=200, + headers={}, # mutable-ok: httpx.Response wants a plain dict of headers + ), model_response=ModelResponse(id=result_id if isinstance(result_id, str) and result_id else None), json_mode=None, speed=speed, diff --git a/litellm/llms/bedrock/batches/transformation.py b/litellm/llms/bedrock/batches/transformation.py index c61cd0dfeef2..b4c16f010cc2 100644 --- a/litellm/llms/bedrock/batches/transformation.py +++ b/litellm/llms/bedrock/batches/transformation.py @@ -106,9 +106,15 @@ def bedrock_batch_line_to_response( embedding: Final = model_output.get("embedding") return EmbeddingResponse( model=model, - data=[ - {"object": "embedding", "index": 0, "embedding": embedding if isinstance(embedding, list) else []} - ], # mutable-ok: EmbeddingResponse takes a plain data list + data=[ # mutable-ok: EmbeddingResponse takes a plain data list + { # mutable-ok: plain row dict for EmbeddingResponse.data + "object": "embedding", + "index": 0, + "embedding": embedding + if isinstance(embedding, list) + else [], # mutable-ok: empty fallback for the row + } + ], usage=titan_embedding_usage_from_batch_output(model_output), ) if "output" in model_output: From f5670e087167603f9ab3e23102d445caa9b83217 Mon Sep 17 00:00:00 2001 From: yucheng Date: Thu, 24 Sep 2026 07:42:31 +0000 Subject: [PATCH 27/89] fix(batches): skip router budget accounting for batch line item events Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/router_strategy/budget_limiter.py | 3 + .../test_budget_limiter_hotpath.py | 55 +++++++++++-------- 2 files changed, 36 insertions(+), 22 deletions(-) diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index 3e094df7ac80..2691610a2aa8 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/litellm/router_strategy/budget_limiter.py @@ -32,6 +32,7 @@ from litellm.integrations.custom_logger import CustomLogger, Span from litellm.litellm_core_utils.core_helpers import ( get_metadata_variable_name_from_kwargs, + is_batch_line_item_event, ) from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.router_strategy.tag_based_routing import _get_tags_from_request_kwargs @@ -407,6 +408,8 @@ async def _increment_spend_in_current_window(self, spend_key: str, response_cost async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): """Original method now uses helper functions""" verbose_router_logger.debug("in RouterBudgetLimiting.async_log_success_event") + if is_batch_line_item_event(kwargs): + return # WS session wrappers fire with result=None; per-turn costs tracked by inner calls. if kwargs.get("call_type") in ("_aresponses_websocket", "_arealtime"): return diff --git a/tests/test_litellm/router_strategy/test_budget_limiter_hotpath.py b/tests/test_litellm/router_strategy/test_budget_limiter_hotpath.py index 4cc8fe78811e..af64b18388f1 100644 --- a/tests/test_litellm/router_strategy/test_budget_limiter_hotpath.py +++ b/tests/test_litellm/router_strategy/test_budget_limiter_hotpath.py @@ -30,9 +30,7 @@ async def test_get_llm_provider_for_deployment_dict_does_not_require_litellm_par ): class RaiseOnInit: def __init__(self, *args, **kwargs): - raise AssertionError( - "LiteLLM_Params should not be instantiated in hot path" - ) + raise AssertionError("LiteLLM_Params should not be instantiated in hot path") monkeypatch.setattr( "litellm.router_strategy.budget_limiter.LiteLLM_Params", @@ -99,9 +97,7 @@ def _future_style_get_llm_provider( @pytest.mark.asyncio -async def test_async_filter_deployments_resolves_provider_once_per_deployment( - disable_budget_sync, monkeypatch -): +async def test_async_filter_deployments_resolves_provider_once_per_deployment(disable_budget_sync, monkeypatch): provider_budget = RouterBudgetLimiting( dual_cache=DualCache(), provider_budget_config={ @@ -207,9 +203,7 @@ def _legacy_provider_resolution(deployment): Reference implementation used before hot-path optimization. """ try: - _litellm_params = LiteLLM_Params( - **deployment.get("litellm_params", {"model": ""}) - ) + _litellm_params = LiteLLM_Params(**deployment.get("litellm_params", {"model": ""})) _, custom_llm_provider, _, _ = litellm.get_llm_provider( model=_litellm_params.model, litellm_params=_litellm_params, @@ -228,9 +222,7 @@ def _legacy_provider_resolution(deployment): ], ) @pytest.mark.asyncio -async def test_get_llm_provider_for_deployment_matches_legacy_behavior( - disable_budget_sync, deployment -): +async def test_get_llm_provider_for_deployment_matches_legacy_behavior(disable_budget_sync, deployment): provider_budget = RouterBudgetLimiting( dual_cache=DualCache(), provider_budget_config={}, @@ -242,9 +234,7 @@ async def test_get_llm_provider_for_deployment_matches_legacy_behavior( assert current_provider == legacy_provider -def test_register_deployment_budget_for_runtime_added_deployment( - disable_budget_sync, monkeypatch -): +def test_register_deployment_budget_for_runtime_added_deployment(disable_budget_sync, monkeypatch): import asyncio monkeypatch.setattr(asyncio, "create_task", lambda coro: None) @@ -274,9 +264,7 @@ def test_register_deployment_budget_for_runtime_added_deployment( assert budget_limiter._get_budget_config_for_deployment(model_id) is None -def test_router_add_deployment_registers_deployment_budget( - disable_budget_sync, monkeypatch -): +def test_router_add_deployment_registers_deployment_budget(disable_budget_sync, monkeypatch): import asyncio from litellm import Router @@ -304,9 +292,7 @@ def test_router_add_deployment_registers_deployment_budget( budget_limiter = router._get_router_deployment_budget_limiter() assert budget_limiter is not None - config = budget_limiter._get_budget_config_for_deployment( - "runtime-budget-deployment" - ) + config = budget_limiter._get_budget_config_for_deployment("runtime-budget-deployment") assert config is not None assert config.max_budget == 0.000000000001 @@ -378,7 +364,9 @@ async def wait_for_redis(**_: object) -> None: async def test_push_task_failure_is_logged_once_and_not_leaked(disable_budget_sync, caplog): """A real Redis failure on the background push must surface as one error line, never as an unretrieved task exception.""" redis_cache = MagicMock(spec=RedisCache) - redis_cache.async_increment_pipeline = AsyncMock(side_effect=ConnectionError("Error 61 connecting to 127.0.0.1:6379")) + redis_cache.async_increment_pipeline = AsyncMock( + side_effect=ConnectionError("Error 61 connecting to 127.0.0.1:6379") + ) limiter = await _limiter_with_redis(redis_cache) loop = asyncio.get_running_loop() unretrieved = MagicMock() @@ -396,3 +384,26 @@ async def test_push_task_failure_is_logged_once_and_not_leaked(disable_budget_sy "Error syncing in-memory cache with Redis: Error 61 connecting to 127.0.0.1:6379" ] unretrieved.assert_not_called() + + +@pytest.mark.asyncio +async def test_batch_line_item_events_do_not_charge_the_provider_budget(disable_budget_sync): + limiter = RouterBudgetLimiting( + dual_cache=DualCache(), + provider_budget_config={"openai": BudgetConfig(max_budget=100.0, budget_duration="1d")}, + ) + await asyncio.gather(*(task for task in asyncio.all_tasks() if task is not asyncio.current_task())) + + await limiter.async_log_success_event( + kwargs={ + "call_type": "acompletion", + "litellm_params": {"custom_llm_provider": "openai", "batch_parent_id": "batch_x"}, + "standard_logging_object": {"response_cost": 1.0, "model_id": "dep-1"}, + }, + response_obj=None, + start_time=None, + end_time=None, + ) + + assert limiter.dual_cache.in_memory_cache.get_cache("provider_spend:openai:1d") in (None, 0) + assert limiter.redis_increment_operation_queue == [] From 4c86f0b5caa883d60390230858d22d75eb2c8927 Mon Sep 17 00:00:00 2001 From: yucheng Date: Fri, 25 Sep 2026 01:10:11 +0000 Subject: [PATCH 28/89] fix(batches): keep anthropic line errors and skip native vertex batch rows Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/batches/batch_line_item_logging.py | 20 +- .../batches/test_batch_line_item_logging.py | 226 +++++++++++++++--- 2 files changed, 211 insertions(+), 35 deletions(-) diff --git a/litellm/batches/batch_line_item_logging.py b/litellm/batches/batch_line_item_logging.py index cf4de4863289..3be093657ec6 100644 --- a/litellm/batches/batch_line_item_logging.py +++ b/litellm/batches/batch_line_item_logging.py @@ -12,6 +12,7 @@ _get_response_from_batch_job_output_file, # pyright: ignore[reportPrivateUsage] # same reuse _iter_batch_output_entries, # pyright: ignore[reportPrivateUsage, reportUnknownVariableType] # same reuse; helper is untyped upstream _safe_output_line_stats, # pyright: ignore[reportPrivateUsage] # same reuse + _uses_native_vertex_output, # pyright: ignore[reportPrivateUsage] # same reuse ) from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.utils import ( @@ -127,6 +128,14 @@ def _line_status_code(entry: Mapping[str, object], custom_llm_provider: str) -> return None +def _line_error_payload(entry: Mapping[str, object], custom_llm_provider: _BatchLineProvider) -> object: + if custom_llm_provider == "anthropic": + return ( + (_as_object_mapping(entry.get("result")) or _EMPTY_BODY).get("error") or entry.get("result") or _EMPTY_BODY + ) + return entry.get("error") or entry.get("response") or _EMPTY_BODY + + def _call_type_for_request(request_line: Mapping[str, object] | None) -> str: url: Final = request_line.get("url") if request_line is not None else None return _CALL_TYPE_BY_BATCH_URL.get(url if isinstance(url, str) else "", "acompletion") @@ -311,9 +320,7 @@ async def _emit_line_event( now: Final = datetime.now() # noqa: DTZ005 # naive to match the logging pipeline start_time if result is None: - exception: Final = _BatchLineFailure( - entry.get("error") or entry.get("response") or {} # mutable-ok: fallback payload dict passed to Exception - ) + exception: Final = _BatchLineFailure(_line_error_payload(entry, custom_llm_provider)) exception._hidden_params = _line_hidden_params(batch, custom_id, status_code) # pyright: ignore[reportPrivateUsage] # _hidden_params is set on the exception instance itself await child.async_failure_handler( exception=exception, @@ -398,6 +405,13 @@ async def log_batch_line_items( requests_by_id: Final = _requests_by_custom_id(input_file_content) output_content: Final = await _fetch_managed_file_or_empty(batch.output_file_id, line_provider, fetch_params) + first_row: Final = next(_output_entries(output_content), None) + if _uses_native_vertex_output(line_provider, model_name, first_row): + verbose_logger.warning( + "batch line-item callbacks do not support native vertex_ai batch output rows yet, skipping. batch_id=%s", + batch.id, + ) + return 0 error_content: Final = await _fetch_managed_file_or_empty(batch.error_file_id, line_provider, fetch_params) for content in (output_content, error_content): for entry in _output_entries(content): diff --git a/tests/test_litellm/batches/test_batch_line_item_logging.py b/tests/test_litellm/batches/test_batch_line_item_logging.py index eb8ae26222db..2d126fc3f298 100644 --- a/tests/test_litellm/batches/test_batch_line_item_logging.py +++ b/tests/test_litellm/batches/test_batch_line_item_logging.py @@ -123,10 +123,16 @@ async def async_log_failure_event(self, kwargs, response_obj, start_time, end_ti @pytest.fixture def recorder(): logger = _RecordingLogger() - saved_flag = litellm.store_batch_line_items_in_callbacks # test-quality-ok: process-wide opt-in flag; restored in teardown - saved_success = list(litellm._async_success_callback) # test-quality-ok: the feature dispatches through this global list; restored in teardown + saved_flag = ( + litellm.store_batch_line_items_in_callbacks + ) # test-quality-ok: process-wide opt-in flag; restored in teardown + saved_success = list( + litellm._async_success_callback + ) # test-quality-ok: the feature dispatches through this global list; restored in teardown saved_failure = list(litellm._async_failure_callback) # test-quality-ok: same dispatch seam, restored in teardown - litellm._async_success_callback = [logger] # test-quality-ok: there is no injection seam for callback lists; teardown restores + litellm._async_success_callback = [ + logger + ] # test-quality-ok: there is no injection seam for callback lists; teardown restores litellm._async_failure_callback = [logger] # test-quality-ok: same dispatch seam, restored in teardown yield logger litellm.store_batch_line_items_in_callbacks = saved_flag # test-quality-ok: teardown restoring the value set above @@ -175,11 +181,17 @@ def _hidden(event: dict) -> dict: @pytest.mark.asyncio async def test_line_items_emitted_alongside_aggregate(recorder): - litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: the flag under test is a module global; fixture restores it + litellm.store_batch_line_items_in_callbacks = ( + True # test-quality-ok: the flag under test is a module global; fixture restores it + ) batch = _batch() with ( - patch("litellm.files.main.afile_content", new_callable=AsyncMock, side_effect=_file_content), # test-quality-ok: afile_content is the provider boundary; there is no HTTP transport or injection seam for managed file fetch - patch("litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02)), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch + patch( + "litellm.files.main.afile_content", new_callable=AsyncMock, side_effect=_file_content + ), # test-quality-ok: afile_content is the provider boundary; there is no HTTP transport or injection seam for managed file fetch + patch( + "litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02) + ), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch ): await _log_completed_batch(_parent_logging(), batch) @@ -215,7 +227,9 @@ async def test_line_items_emitted_alongside_aggregate(recorder): async def test_flag_off_emits_only_aggregate(recorder): assert litellm.store_batch_line_items_in_callbacks is False file_mock = AsyncMock(side_effect=_file_content) - with patch("litellm.files.main.afile_content", file_mock): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + with patch( + "litellm.files.main.afile_content", file_mock + ): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch await _log_completed_batch(_parent_logging(), _batch()) assert len(recorder.success_events) == 1 @@ -225,9 +239,13 @@ async def test_flag_off_emits_only_aggregate(recorder): @pytest.mark.asyncio async def test_line_items_skipped_for_unsupported_provider(recorder): - litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: the flag under test is a module global; fixture restores it + litellm.store_batch_line_items_in_callbacks = ( + True # test-quality-ok: the flag under test is a module global; fixture restores it + ) file_mock: Final = AsyncMock(side_effect=_file_content) - with patch("litellm.files.main.afile_content", file_mock): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + with patch( + "litellm.files.main.afile_content", file_mock + ): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch await _log_completed_batch(_parent_logging(custom_llm_provider="cohere"), _batch()) file_mock.assert_not_awaited() @@ -239,7 +257,9 @@ async def test_line_items_skipped_for_unsupported_provider(recorder): @pytest.mark.asyncio async def test_in_progress_batch_poll_emits_no_line_events(recorder): - litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: the flag under test is a module global; fixture restores it + litellm.store_batch_line_items_in_callbacks = ( + True # test-quality-ok: the flag under test is a module global; fixture restores it + ) in_progress: Final = LiteLLMBatch( id="batch_wip", object="batch", @@ -252,7 +272,9 @@ async def test_in_progress_batch_poll_emits_no_line_events(recorder): created_at=1, ) file_mock: Final = AsyncMock(side_effect=_file_content) - with patch("litellm.files.main.afile_content", file_mock): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + with patch( + "litellm.files.main.afile_content", file_mock + ): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch await _parent_logging().async_success_handler(result=in_progress) file_mock.assert_not_called() @@ -262,14 +284,19 @@ async def test_in_progress_batch_poll_emits_no_line_events(recorder): @pytest.mark.asyncio async def test_input_fetch_failure_still_emits_aggregate(recorder): - litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: the flag under test is a module global; fixture restores it - with patch("litellm.files.main.afile_content", new_callable=AsyncMock, side_effect=ValueError("boom")): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + litellm.store_batch_line_items_in_callbacks = ( + True # test-quality-ok: the flag under test is a module global; fixture restores it + ) + with patch( + "litellm.files.main.afile_content", new_callable=AsyncMock, side_effect=ValueError("boom") + ): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch await _log_completed_batch(_parent_logging(), _batch()) assert len(recorder.success_events) == 1 assert len(recorder.failure_events) == 0 assert _payload(recorder.success_events[0])["response_cost"] == 1.5 + EDGE_INPUT_JSONL = b"\n".join( [ json.dumps( @@ -506,7 +533,9 @@ def _parent_logging_with_params(litellm_params: dict) -> Logging: @pytest.mark.asyncio async def test_line_items_edge_shapes_and_edge_cases(recorder): - litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: the flag under test is a module global; fixture restores it + litellm.store_batch_line_items_in_callbacks = ( + True # test-quality-ok: the flag under test is a module global; fixture restores it + ) batch = LiteLLMBatch( id="batch_edge", object="batch", @@ -524,10 +553,16 @@ async def test_line_items_edge_shapes_and_edge_cases(recorder): "metadata": {"model_info": {"id": "dep-1"}, "model_group": "gpt-4o"}, } ) - parent._litellm_internal_model_credentials = {"api_key": "sk-line-items-marker"} # test-quality-ok: private transport attribute, same channel the batch cost tracker uses + parent._litellm_internal_model_credentials = { + "api_key": "sk-line-items-marker" + } # test-quality-ok: private transport attribute, same channel the batch cost tracker uses with ( - patch("litellm.files.main.afile_content", file_mock), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch - patch("litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02)), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch + patch( + "litellm.files.main.afile_content", file_mock + ), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + patch( + "litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02) + ), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch ): await _log_completed_batch(parent, batch) @@ -562,7 +597,9 @@ async def test_line_items_edge_shapes_and_edge_cases(recorder): @pytest.mark.asyncio async def test_line_items_child_params_drop_parent_credentials(recorder): - litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: the flag under test is a module global; fixture restores it + litellm.store_batch_line_items_in_callbacks = ( + True # test-quality-ok: the flag under test is a module global; fixture restores it + ) file_mock: Final = AsyncMock(side_effect=_file_content) parent: Final = _parent_logging_with_params( { @@ -574,15 +611,17 @@ async def test_line_items_child_params_drop_parent_credentials(recorder): ) batch: Final = _batch() with ( - patch("litellm.files.main.afile_content", file_mock), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch - patch("litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02)), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch + patch( + "litellm.files.main.afile_content", file_mock + ), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + patch( + "litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02) + ), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch ): await _log_completed_batch(parent, batch) line_events = [ - e - for e in [*recorder.success_events, *recorder.failure_events] - if _hidden(e).get("batch_custom_id") is not None + e for e in [*recorder.success_events, *recorder.failure_events] if _hidden(e).get("batch_custom_id") is not None ] assert len(line_events) == 2 for event in line_events: @@ -595,7 +634,9 @@ async def test_line_items_child_params_drop_parent_credentials(recorder): @pytest.mark.asyncio async def test_line_items_anthropic_shapes(recorder): - litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: the flag under test is a module global; fixture restores it + litellm.store_batch_line_items_in_callbacks = ( + True # test-quality-ok: the flag under test is a module global; fixture restores it + ) batch = LiteLLMBatch( id="batch_anth", object="batch", @@ -615,8 +656,12 @@ async def test_line_items_anthropic_shapes(recorder): custom_llm_provider="anthropic", ) with ( - patch("litellm.files.main.afile_content", file_mock), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch - patch("litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02)), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch + patch( + "litellm.files.main.afile_content", file_mock + ), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + patch( + "litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02) + ), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch ): await _log_completed_batch(logging_obj, batch) @@ -645,11 +690,17 @@ def _provider_batch(batch_id: str, input_file_id: str, output_file_id: str) -> L @pytest.mark.asyncio async def test_line_items_bedrock_shapes(recorder): - litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: the flag under test is a module global; fixture restores it + litellm.store_batch_line_items_in_callbacks = ( + True # test-quality-ok: the flag under test is a module global; fixture restores it + ) file_mock: Final = AsyncMock(side_effect=_edge_file_content) with ( - patch("litellm.files.main.afile_content", file_mock), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch - patch("litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02)), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch + patch( + "litellm.files.main.afile_content", file_mock + ), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + patch( + "litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02) + ), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch ): await _log_completed_batch( _parent_logging(custom_llm_provider="bedrock"), @@ -686,11 +737,17 @@ async def test_line_items_bedrock_shapes(recorder): @pytest.mark.asyncio async def test_line_items_mistral_shape(recorder): - litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: the flag under test is a module global; fixture restores it + litellm.store_batch_line_items_in_callbacks = ( + True # test-quality-ok: the flag under test is a module global; fixture restores it + ) file_mock: Final = AsyncMock(side_effect=_edge_file_content) with ( - patch("litellm.files.main.afile_content", file_mock), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch - patch("litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02)), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch + patch( + "litellm.files.main.afile_content", file_mock + ), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + patch( + "litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02) + ), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch ): await _log_completed_batch( _parent_logging(custom_llm_provider="mistral"), @@ -700,3 +757,108 @@ async def test_line_items_mistral_shape(recorder): line = next(e for e in recorder.success_events if _hidden(e).get("batch_custom_id") == "mis") assert _hidden(line)["batch_line_status_code"] == 200 assert _payload(line)["response"]["choices"][0]["message"]["content"] == "mistral hi" + + +ANTHROPIC_ERR_INPUT_JSONL = json.dumps( + { + "custom_id": "aerr", + "params": {"model": "claude-3", "max_tokens": 5, "messages": [{"role": "user", "content": "hi aerr"}]}, + } +).encode() + +ANTHROPIC_ERR_OUTPUT_JSONL = json.dumps( + { + "custom_id": "aerr", + "result": { + "type": "errored", + "error": { + "type": "error", + "error": {"type": "invalid_request_error", "message": "max_tokens must be positive"}, + }, + }, + } +).encode() + +VERTEX_NATIVE_OUTPUT_JSONL = b"\n".join( + [ + json.dumps( + { + "request": {"contents": [{"role": "user", "parts": [{"text": "hi v1"}]}]}, + "response": { + "candidates": [{"content": {"role": "model", "parts": [{"text": "hi back"}]}}], + "usageMetadata": {"promptTokenCount": 3, "candidatesTokenCount": 2, "totalTokenCount": 5}, + }, + } + ).encode(), + json.dumps( + { + "request": {"contents": [{"role": "user", "parts": [{"text": "hi v2"}]}]}, + "status": "Internal error", + } + ).encode(), + ] +) + + +@pytest.mark.asyncio +async def test_line_items_anthropic_failure_keeps_provider_error(recorder): + litellm.store_batch_line_items_in_callbacks = ( + True # test-quality-ok: the flag under test is a module global; fixture restores it + ) + file_mock: Final = AsyncMock( + side_effect=lambda file_id, **kwargs: SimpleNamespace( + content={ + "input-anth-err": ANTHROPIC_ERR_INPUT_JSONL, + "output-anth-err": ANTHROPIC_ERR_OUTPUT_JSONL, + }[file_id] + ) + ) + with ( + patch( + "litellm.files.main.afile_content", file_mock + ), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + patch( + "litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02) + ), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch + ): + await _log_completed_batch( + _parent_logging(custom_llm_provider="anthropic"), + _provider_batch("batch_anth_err", "input-anth-err", "output-anth-err"), + ) + + assert len(recorder.failure_events) == 1 + failure = recorder.failure_events[0] + assert _hidden(failure)["batch_custom_id"] == "aerr" + assert _hidden(failure)["batch_line_status_code"] is None + assert "max_tokens must be positive" in _payload(failure)["error_str"] + + +@pytest.mark.asyncio +async def test_line_items_native_vertex_rows_are_skipped(recorder): + litellm.store_batch_line_items_in_callbacks = ( + True # test-quality-ok: the flag under test is a module global; fixture restores it + ) + file_mock: Final = AsyncMock( + side_effect=lambda file_id, **kwargs: SimpleNamespace( + content={ + "input-vtx": b"", + "output-vtx": VERTEX_NATIVE_OUTPUT_JSONL, + }[file_id] + ) + ) + with ( + patch( + "litellm.files.main.afile_content", file_mock + ), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + patch( + "litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02) + ), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch + ): + await _log_completed_batch( + _parent_logging(custom_llm_provider="vertex_ai"), + _provider_batch("batch_vtx", "input-vtx", "output-vtx"), + ) + + assert len(recorder.success_events) == 1 + assert len(recorder.failure_events) == 0 + assert _hidden(recorder.success_events[0]).get("batch_custom_id") is None From 5777f7e72dbbc32206bfa639ce9bc08d6a907068 Mon Sep 17 00:00:00 2001 From: yucheng Date: Fri, 25 Sep 2026 01:11:27 +0000 Subject: [PATCH 29/89] test(batches): restore test-quality tags on line item tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../batches/test_batch_line_item_logging.py | 161 ++++++------------ 1 file changed, 48 insertions(+), 113 deletions(-) diff --git a/tests/test_litellm/batches/test_batch_line_item_logging.py b/tests/test_litellm/batches/test_batch_line_item_logging.py index 2d126fc3f298..97442f59a0da 100644 --- a/tests/test_litellm/batches/test_batch_line_item_logging.py +++ b/tests/test_litellm/batches/test_batch_line_item_logging.py @@ -123,16 +123,10 @@ async def async_log_failure_event(self, kwargs, response_obj, start_time, end_ti @pytest.fixture def recorder(): logger = _RecordingLogger() - saved_flag = ( - litellm.store_batch_line_items_in_callbacks - ) # test-quality-ok: process-wide opt-in flag; restored in teardown - saved_success = list( - litellm._async_success_callback - ) # test-quality-ok: the feature dispatches through this global list; restored in teardown + saved_flag = litellm.store_batch_line_items_in_callbacks # test-quality-ok: process-wide opt-in flag; restored in teardown + saved_success = list(litellm._async_success_callback) # test-quality-ok: the feature dispatches through this global list; restored in teardown saved_failure = list(litellm._async_failure_callback) # test-quality-ok: same dispatch seam, restored in teardown - litellm._async_success_callback = [ - logger - ] # test-quality-ok: there is no injection seam for callback lists; teardown restores + litellm._async_success_callback = [logger] # test-quality-ok: there is no injection seam for callback lists; teardown restores litellm._async_failure_callback = [logger] # test-quality-ok: same dispatch seam, restored in teardown yield logger litellm.store_batch_line_items_in_callbacks = saved_flag # test-quality-ok: teardown restoring the value set above @@ -181,17 +175,11 @@ def _hidden(event: dict) -> dict: @pytest.mark.asyncio async def test_line_items_emitted_alongside_aggregate(recorder): - litellm.store_batch_line_items_in_callbacks = ( - True # test-quality-ok: the flag under test is a module global; fixture restores it - ) + litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: the flag under test is a module global; fixture restores it batch = _batch() with ( - patch( - "litellm.files.main.afile_content", new_callable=AsyncMock, side_effect=_file_content - ), # test-quality-ok: afile_content is the provider boundary; there is no HTTP transport or injection seam for managed file fetch - patch( - "litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02) - ), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch + patch("litellm.files.main.afile_content", new_callable=AsyncMock, side_effect=_file_content), # test-quality-ok: afile_content is the provider boundary; there is no HTTP transport or injection seam for managed file fetch + patch("litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02)), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch ): await _log_completed_batch(_parent_logging(), batch) @@ -227,9 +215,7 @@ async def test_line_items_emitted_alongside_aggregate(recorder): async def test_flag_off_emits_only_aggregate(recorder): assert litellm.store_batch_line_items_in_callbacks is False file_mock = AsyncMock(side_effect=_file_content) - with patch( - "litellm.files.main.afile_content", file_mock - ): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + with patch("litellm.files.main.afile_content", file_mock): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch await _log_completed_batch(_parent_logging(), _batch()) assert len(recorder.success_events) == 1 @@ -239,13 +225,9 @@ async def test_flag_off_emits_only_aggregate(recorder): @pytest.mark.asyncio async def test_line_items_skipped_for_unsupported_provider(recorder): - litellm.store_batch_line_items_in_callbacks = ( - True # test-quality-ok: the flag under test is a module global; fixture restores it - ) + litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: the flag under test is a module global; fixture restores it file_mock: Final = AsyncMock(side_effect=_file_content) - with patch( - "litellm.files.main.afile_content", file_mock - ): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + with patch("litellm.files.main.afile_content", file_mock): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch await _log_completed_batch(_parent_logging(custom_llm_provider="cohere"), _batch()) file_mock.assert_not_awaited() @@ -257,9 +239,7 @@ async def test_line_items_skipped_for_unsupported_provider(recorder): @pytest.mark.asyncio async def test_in_progress_batch_poll_emits_no_line_events(recorder): - litellm.store_batch_line_items_in_callbacks = ( - True # test-quality-ok: the flag under test is a module global; fixture restores it - ) + litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: the flag under test is a module global; fixture restores it in_progress: Final = LiteLLMBatch( id="batch_wip", object="batch", @@ -272,9 +252,7 @@ async def test_in_progress_batch_poll_emits_no_line_events(recorder): created_at=1, ) file_mock: Final = AsyncMock(side_effect=_file_content) - with patch( - "litellm.files.main.afile_content", file_mock - ): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + with patch("litellm.files.main.afile_content", file_mock): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch await _parent_logging().async_success_handler(result=in_progress) file_mock.assert_not_called() @@ -284,19 +262,14 @@ async def test_in_progress_batch_poll_emits_no_line_events(recorder): @pytest.mark.asyncio async def test_input_fetch_failure_still_emits_aggregate(recorder): - litellm.store_batch_line_items_in_callbacks = ( - True # test-quality-ok: the flag under test is a module global; fixture restores it - ) - with patch( - "litellm.files.main.afile_content", new_callable=AsyncMock, side_effect=ValueError("boom") - ): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: the flag under test is a module global; fixture restores it + with patch("litellm.files.main.afile_content", new_callable=AsyncMock, side_effect=ValueError("boom")): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch await _log_completed_batch(_parent_logging(), _batch()) assert len(recorder.success_events) == 1 assert len(recorder.failure_events) == 0 assert _payload(recorder.success_events[0])["response_cost"] == 1.5 - EDGE_INPUT_JSONL = b"\n".join( [ json.dumps( @@ -533,9 +506,7 @@ def _parent_logging_with_params(litellm_params: dict) -> Logging: @pytest.mark.asyncio async def test_line_items_edge_shapes_and_edge_cases(recorder): - litellm.store_batch_line_items_in_callbacks = ( - True # test-quality-ok: the flag under test is a module global; fixture restores it - ) + litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: the flag under test is a module global; fixture restores it batch = LiteLLMBatch( id="batch_edge", object="batch", @@ -553,16 +524,10 @@ async def test_line_items_edge_shapes_and_edge_cases(recorder): "metadata": {"model_info": {"id": "dep-1"}, "model_group": "gpt-4o"}, } ) - parent._litellm_internal_model_credentials = { - "api_key": "sk-line-items-marker" - } # test-quality-ok: private transport attribute, same channel the batch cost tracker uses + parent._litellm_internal_model_credentials = {"api_key": "sk-line-items-marker"} # test-quality-ok: private transport attribute, same channel the batch cost tracker uses with ( - patch( - "litellm.files.main.afile_content", file_mock - ), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch - patch( - "litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02) - ), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch + patch("litellm.files.main.afile_content", file_mock), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + patch("litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02)), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch ): await _log_completed_batch(parent, batch) @@ -597,9 +562,7 @@ async def test_line_items_edge_shapes_and_edge_cases(recorder): @pytest.mark.asyncio async def test_line_items_child_params_drop_parent_credentials(recorder): - litellm.store_batch_line_items_in_callbacks = ( - True # test-quality-ok: the flag under test is a module global; fixture restores it - ) + litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: the flag under test is a module global; fixture restores it file_mock: Final = AsyncMock(side_effect=_file_content) parent: Final = _parent_logging_with_params( { @@ -611,17 +574,15 @@ async def test_line_items_child_params_drop_parent_credentials(recorder): ) batch: Final = _batch() with ( - patch( - "litellm.files.main.afile_content", file_mock - ), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch - patch( - "litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02) - ), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch + patch("litellm.files.main.afile_content", file_mock), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + patch("litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02)), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch ): await _log_completed_batch(parent, batch) line_events = [ - e for e in [*recorder.success_events, *recorder.failure_events] if _hidden(e).get("batch_custom_id") is not None + e + for e in [*recorder.success_events, *recorder.failure_events] + if _hidden(e).get("batch_custom_id") is not None ] assert len(line_events) == 2 for event in line_events: @@ -634,9 +595,7 @@ async def test_line_items_child_params_drop_parent_credentials(recorder): @pytest.mark.asyncio async def test_line_items_anthropic_shapes(recorder): - litellm.store_batch_line_items_in_callbacks = ( - True # test-quality-ok: the flag under test is a module global; fixture restores it - ) + litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: the flag under test is a module global; fixture restores it batch = LiteLLMBatch( id="batch_anth", object="batch", @@ -656,12 +615,8 @@ async def test_line_items_anthropic_shapes(recorder): custom_llm_provider="anthropic", ) with ( - patch( - "litellm.files.main.afile_content", file_mock - ), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch - patch( - "litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02) - ), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch + patch("litellm.files.main.afile_content", file_mock), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + patch("litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02)), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch ): await _log_completed_batch(logging_obj, batch) @@ -690,17 +645,11 @@ def _provider_batch(batch_id: str, input_file_id: str, output_file_id: str) -> L @pytest.mark.asyncio async def test_line_items_bedrock_shapes(recorder): - litellm.store_batch_line_items_in_callbacks = ( - True # test-quality-ok: the flag under test is a module global; fixture restores it - ) + litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: the flag under test is a module global; fixture restores it file_mock: Final = AsyncMock(side_effect=_edge_file_content) with ( - patch( - "litellm.files.main.afile_content", file_mock - ), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch - patch( - "litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02) - ), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch + patch("litellm.files.main.afile_content", file_mock), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + patch("litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02)), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch ): await _log_completed_batch( _parent_logging(custom_llm_provider="bedrock"), @@ -737,17 +686,11 @@ async def test_line_items_bedrock_shapes(recorder): @pytest.mark.asyncio async def test_line_items_mistral_shape(recorder): - litellm.store_batch_line_items_in_callbacks = ( - True # test-quality-ok: the flag under test is a module global; fixture restores it - ) + litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: the flag under test is a module global; fixture restores it file_mock: Final = AsyncMock(side_effect=_edge_file_content) with ( - patch( - "litellm.files.main.afile_content", file_mock - ), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch - patch( - "litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02) - ), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch + patch("litellm.files.main.afile_content", file_mock), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + patch("litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02)), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch ): await _log_completed_batch( _parent_logging(custom_llm_provider="mistral"), @@ -800,26 +743,24 @@ async def test_line_items_mistral_shape(recorder): ) +def _scoped_file_content(file_map): + return lambda file_id, **kwargs: SimpleNamespace(content=file_map[file_id]) + + @pytest.mark.asyncio async def test_line_items_anthropic_failure_keeps_provider_error(recorder): - litellm.store_batch_line_items_in_callbacks = ( - True # test-quality-ok: the flag under test is a module global; fixture restores it - ) + litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: the flag under test is a module global; fixture restores it file_mock: Final = AsyncMock( - side_effect=lambda file_id, **kwargs: SimpleNamespace( - content={ + side_effect=_scoped_file_content( + { "input-anth-err": ANTHROPIC_ERR_INPUT_JSONL, "output-anth-err": ANTHROPIC_ERR_OUTPUT_JSONL, - }[file_id] + } ) ) with ( - patch( - "litellm.files.main.afile_content", file_mock - ), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch - patch( - "litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02) - ), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch + patch("litellm.files.main.afile_content", file_mock), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + patch("litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02)), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch ): await _log_completed_batch( _parent_logging(custom_llm_provider="anthropic"), @@ -835,24 +776,18 @@ async def test_line_items_anthropic_failure_keeps_provider_error(recorder): @pytest.mark.asyncio async def test_line_items_native_vertex_rows_are_skipped(recorder): - litellm.store_batch_line_items_in_callbacks = ( - True # test-quality-ok: the flag under test is a module global; fixture restores it - ) + litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: the flag under test is a module global; fixture restores it file_mock: Final = AsyncMock( - side_effect=lambda file_id, **kwargs: SimpleNamespace( - content={ + side_effect=_scoped_file_content( + { "input-vtx": b"", "output-vtx": VERTEX_NATIVE_OUTPUT_JSONL, - }[file_id] + } ) ) with ( - patch( - "litellm.files.main.afile_content", file_mock - ), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch - patch( - "litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02) - ), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch + patch("litellm.files.main.afile_content", file_mock), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + patch("litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02)), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch ): await _log_completed_batch( _parent_logging(custom_llm_provider="vertex_ai"), From f3ba65d04ef969b0cd8809104a1a81a034b81b94 Mon Sep 17 00:00:00 2001 From: yucheng Date: Sat, 26 Sep 2026 22:24:54 +0000 Subject: [PATCH 30/89] test: move batch line-item unit test to tests/unit after merging main Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../batches/test_batch_line_item_logging.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename tests/{test_litellm => unit}/batches/test_batch_line_item_logging.py (100%) diff --git a/tests/test_litellm/batches/test_batch_line_item_logging.py b/tests/unit/batches/test_batch_line_item_logging.py similarity index 100% rename from tests/test_litellm/batches/test_batch_line_item_logging.py rename to tests/unit/batches/test_batch_line_item_logging.py From f65b5b12811adad97efaeb770da03bc93b0c0297 Mon Sep 17 00:00:00 2001 From: yucheng Date: Sun, 27 Sep 2026 00:18:14 +0000 Subject: [PATCH 31/89] fix(batches): emit batch line-item callbacks once per batch Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/batches/batch_line_item_logging.py | 12 +++ litellm/proxy/proxy_server.py | 12 ++- .../spend/test_batch_line_item_callbacks.py | 17 +++- tests/test_litellm/proxy/test_proxy_server.py | 1 + .../proxy/test_redis_auth_cache_flag.py | 2 + .../batches/test_batch_line_item_logging.py | 95 +++++++++++++++++++ 6 files changed, 136 insertions(+), 3 deletions(-) diff --git a/litellm/batches/batch_line_item_logging.py b/litellm/batches/batch_line_item_logging.py index 3be093657ec6..f24ae9fb4320 100644 --- a/litellm/batches/batch_line_item_logging.py +++ b/litellm/batches/batch_line_item_logging.py @@ -14,6 +14,7 @@ _safe_output_line_stats, # pyright: ignore[reportPrivateUsage] # same reuse _uses_native_vertex_output, # pyright: ignore[reportPrivateUsage] # same reuse ) +from litellm.caching.caching import DualCache from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.utils import ( EmbeddingResponse, @@ -63,6 +64,10 @@ def _supported_line_provider(value: str) -> _BatchLineProvider | None: _EMPTY_BODY: Final[Mapping[str, object]] = MappingProxyType({}) +_LINE_ITEM_CLAIM_TTL_SECONDS: Final = 30 * 24 * 60 * 60 + +batch_line_item_claim_cache: Final = DualCache() + class _BatchLineFailure(Exception): """A provider-reported per-line batch failure; carries the batch's hidden @@ -376,6 +381,7 @@ async def log_batch_line_items( model_name: str | None, litellm_params: dict[str, object] | None, # mutable-ok: the logging object's shared litellm_params dict model_info: ModelInfo | None, + claim_cache: DualCache = batch_line_item_claim_cache, ) -> int: """Emit one callback event per JSONL line of a completed batch (request paired with its response/error), behind the opt-in @@ -391,6 +397,12 @@ async def log_batch_line_items( batch.id, ) return 0 + claim: Final = await claim_cache.async_increment_cache( + f"batch_line_items_emitted:{batch.id}", 1, ttl=_LINE_ITEM_CLAIM_TTL_SECONDS + ) + if claim is not None and claim > 1: + verbose_logger.debug("batch line items already emitted for batch_id=%s, skipping", batch.id) + return 0 emitted = 0 # rebind-ok: loop accumulator for emitted line count try: internal_credentials: Final = parent._litellm_internal_model_credentials # pyright: ignore[reportPrivateUsage] # declared transport attribute on Logging diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 73c993f57daa..bea6a9d24391 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -267,6 +267,10 @@ def generate_feedback_box(): import litellm._redis from litellm import Router from litellm._logging import _redact_string, verbose_proxy_logger, verbose_router_logger +from litellm.batches.batch_line_item_logging import ( + _LINE_ITEM_CLAIM_TTL_SECONDS, # pyright: ignore[reportPrivateUsage] # the claim window is defined next to the cache it expires + batch_line_item_claim_cache, +) from litellm.caching.caching import DualCache, RedisCache from litellm.caching.redis_cache import RedisCircuitBreakerOpenError, is_redis_timeout_failure from litellm.caching.redis_cluster_cache import RedisClusterCache @@ -4733,8 +4737,8 @@ def _attach_redis_usage_cache(redis_cache: RedisCache, enable_redis_auth_cache: """ Wires an established coordination Redis into the proxy-level caches that consume it directly: the spend counter cache, the CLI SSO login-session - cache, the cluster-wide config cache, and (only when opted in) the - virtual-key auth cache. + cache, the batch line-item claim cache, the cluster-wide config cache, + and (only when opted in) the virtual-key auth cache. The CLI SSO login-session cache is always backed by Redis when available so that the browser SSO flow behind `lite login` survives landing on different @@ -4748,6 +4752,10 @@ def _attach_redis_usage_cache(redis_cache: RedisCache, enable_redis_auth_cache: redis_cache, default_redis_ttl=CLI_SSO_SESSION_TTL_SECONDS, ) + batch_line_item_claim_cache.attach_redis_cache( + redis_cache, + default_redis_ttl=_LINE_ITEM_CLAIM_TTL_SECONDS, + ) if enable_redis_auth_cache is True: user_api_key_cache.attach_redis_cache( redis_cache, diff --git a/tests/integration/spend/test_batch_line_item_callbacks.py b/tests/integration/spend/test_batch_line_item_callbacks.py index 97cad137b516..a707a7a4f7a6 100644 --- a/tests/integration/spend/test_batch_line_item_callbacks.py +++ b/tests/integration/spend/test_batch_line_item_callbacks.py @@ -239,7 +239,8 @@ def delivered() -> tuple[dict[str, JsonValue], ...]: events: Final = eventually( delivered, lambda values: ( - len([e for e in values if e["call_type"] == "aretrieve_batch"]) >= 1 + len([e for e in values if e["call_type"] == "acreate_batch"]) >= 1 + and len([e for e in values if e["call_type"] == "aretrieve_batch"]) >= 1 and len([e for e in values if _hidden(e).get("batch_custom_id") is not None]) >= len(ALL_CUSTOM_IDS) ), seconds=40, @@ -286,3 +287,17 @@ def delivered() -> tuple[dict[str, JsonValue], ...]: assert len(batch_rows) == 1, rows assert not any(row["call_type"] == "acompletion" for row in rows), rows assert batch_rows[0]["prompt_tokens"] == len(OUTPUT_SUCCESS_IDS) * PROMPT_TOKENS, rows + + repeated_gets: Final = [candidate.request("GET", f"/v1/batches/{batch_id}", key=key) for _ in range(2)] + assert all(response.status_code == 200 for response in repeated_gets) + events_after_repeats: Final = eventually( + delivered, + lambda values: len([e for e in values if e["call_type"] == "aretrieve_batch"]) >= 1 + len(repeated_gets), + seconds=40, + ) + repeat_line_events: Final = tuple( + event for event in events_after_repeats if _hidden(event).get("batch_custom_id") is not None + ) + assert len(repeat_line_events) == len(ALL_CUSTOM_IDS), [ + (event["call_type"], _hidden(event).get("batch_custom_id")) for event in events_after_repeats + ] diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index ffbcd1fba242..e3737b0b7615 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -13121,6 +13121,7 @@ def _patched_coordination_redis_module_state( patch.object(proxy_server_module, "spend_counter_cache", spend_cache), patch.object(proxy_server_module, "user_api_key_cache", DualCache()), patch.object(proxy_server_module, "cli_sso_session_cache", DualCache()), + patch.object(proxy_server_module, "batch_line_item_claim_cache", DualCache()), patch.object(proxy_server_module, "llm_router", None), patch.object(proxy_server_module, "litellm_config_cache", config_cache), patch.object(proxy_server_module, "RedisCache", redis_cache_class), diff --git a/tests/test_litellm/proxy/test_redis_auth_cache_flag.py b/tests/test_litellm/proxy/test_redis_auth_cache_flag.py index 573bfc40c96d..c59c4f3422d6 100644 --- a/tests/test_litellm/proxy/test_redis_auth_cache_flag.py +++ b/tests/test_litellm/proxy/test_redis_auth_cache_flag.py @@ -182,10 +182,12 @@ def test_flag_absent_still_shares_the_model_budget_counters_over_redis(self): ps.cli_sso_session_cache, ps.user_api_key_cache, ps.litellm_config_cache, + ps.batch_line_item_claim_cache, ) with ExitStack() as detached: for cache in touched_caches: detached.enter_context(patch.object(cache, "redis_cache", None)) ps._attach_redis_usage_cache(fake_redis, enable_redis_auth_cache=False) assert limiter_cache.redis_cache is fake_redis + assert ps.batch_line_item_claim_cache.redis_cache is fake_redis assert ps.user_api_key_cache.redis_cache is None diff --git a/tests/unit/batches/test_batch_line_item_logging.py b/tests/unit/batches/test_batch_line_item_logging.py index 97442f59a0da..9a855b602e5d 100644 --- a/tests/unit/batches/test_batch_line_item_logging.py +++ b/tests/unit/batches/test_batch_line_item_logging.py @@ -21,6 +21,8 @@ import pytest import litellm +from litellm.batches.batch_line_item_logging import batch_line_item_claim_cache, log_batch_line_items +from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging from litellm.types.utils import LiteLLMBatch, Usage @@ -120,6 +122,12 @@ async def async_log_failure_event(self, kwargs, response_obj, start_time, end_ti self.failure_events.append(kwargs) +@pytest.fixture(autouse=True) +def _fresh_line_item_claim_cache(): + batch_line_item_claim_cache.in_memory_cache.flush_cache() + yield + + @pytest.fixture def recorder(): logger = _RecordingLogger() @@ -797,3 +805,90 @@ async def test_line_items_native_vertex_rows_are_skipped(recorder): assert len(recorder.success_events) == 1 assert len(recorder.failure_events) == 0 assert _hidden(recorder.success_events[0]).get("batch_custom_id") is None + + +class _NeverClaimsCache(DualCache): + async def async_increment_cache(self, *_args: object, **_kwargs: object) -> None: + return None + + +@pytest.mark.asyncio +async def test_line_items_emit_once_per_batch_id(recorder): + file_mock: Final = AsyncMock(side_effect=_file_content) + claim_cache: Final = DualCache() + parent: Final = _parent_logging() + batch: Final = _batch() + with patch("litellm.files.main.afile_content", file_mock): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + first: Final = await log_batch_line_items( + batch=batch, + custom_llm_provider="openai", + parent=parent, + model_name="gpt-4o", + litellm_params=None, + model_info=None, + claim_cache=claim_cache, + ) + second: Final = await log_batch_line_items( + batch=batch, + custom_llm_provider="openai", + parent=parent, + model_name="gpt-4o", + litellm_params=None, + model_info=None, + claim_cache=claim_cache, + ) + + assert first == 2 + assert second == 0 + assert len(recorder.success_events) == 1 + assert len(recorder.failure_events) == 1 + + +@pytest.mark.asyncio +async def test_line_items_claim_is_scoped_to_the_batch_id(recorder): + file_mock: Final = AsyncMock(side_effect=_file_content) + claim_cache: Final = DualCache() + parent: Final = _parent_logging() + with patch("litellm.files.main.afile_content", file_mock): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + first: Final = await log_batch_line_items( + batch=_batch(), + custom_llm_provider="openai", + parent=parent, + model_name="gpt-4o", + litellm_params=None, + model_info=None, + claim_cache=claim_cache, + ) + second: Final = await log_batch_line_items( + batch=_provider_batch("batch_other", "input-file-1", "output-file-1"), + custom_llm_provider="openai", + parent=parent, + model_name="gpt-4o", + litellm_params=None, + model_info=None, + claim_cache=claim_cache, + ) + + assert first == 2 + assert second == 1 + assert len(recorder.success_events) == 2 + assert len(recorder.failure_events) == 1 + + +@pytest.mark.asyncio +async def test_line_items_emit_when_the_claim_backend_returns_nothing(recorder): + file_mock: Final = AsyncMock(side_effect=_file_content) + with patch("litellm.files.main.afile_content", file_mock): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + emitted: Final = await log_batch_line_items( + batch=_batch(), + custom_llm_provider="openai", + parent=_parent_logging(), + model_name="gpt-4o", + litellm_params=None, + model_info=None, + claim_cache=_NeverClaimsCache(), + ) + + assert emitted == 2 + assert len(recorder.success_events) == 1 + assert len(recorder.failure_events) == 1 From cffa6405b4c9074ac95edf29831ea83e483ba302 Mon Sep 17 00:00:00 2001 From: yucheng Date: Sun, 27 Sep 2026 00:52:17 +0000 Subject: [PATCH 32/89] fix(batches): release line-item claim when fan-out emitted nothing Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/batches/batch_line_item_logging.py | 13 +++- .../batches/test_batch_line_item_logging.py | 72 +++++++++++++++++++ 2 files changed, 82 insertions(+), 3 deletions(-) diff --git a/litellm/batches/batch_line_item_logging.py b/litellm/batches/batch_line_item_logging.py index f24ae9fb4320..21dd23bdb599 100644 --- a/litellm/batches/batch_line_item_logging.py +++ b/litellm/batches/batch_line_item_logging.py @@ -397,9 +397,8 @@ async def log_batch_line_items( batch.id, ) return 0 - claim: Final = await claim_cache.async_increment_cache( - f"batch_line_items_emitted:{batch.id}", 1, ttl=_LINE_ITEM_CLAIM_TTL_SECONDS - ) + claim_key: Final = f"batch_line_items_emitted:{batch.id}" + claim: Final = await claim_cache.async_increment_cache(claim_key, 1, ttl=_LINE_ITEM_CLAIM_TTL_SECONDS) if claim is not None and claim > 1: verbose_logger.debug("batch line items already emitted for batch_id=%s, skipping", batch.id) return 0 @@ -447,4 +446,12 @@ async def log_batch_line_items( "batch line item logging failed for batch_id=%s; aggregate logging unaffected", batch.id, ) + if emitted == 0: + try: + await claim_cache.async_delete_cache(claim_key) + except Exception: # noqa: BLE001 # the claim release must never raise; worst case the batch stays claimed + verbose_logger.debug( + "batch line item claim release failed for batch_id=%s, claim persists until ttl", + batch.id, + ) return emitted diff --git a/tests/unit/batches/test_batch_line_item_logging.py b/tests/unit/batches/test_batch_line_item_logging.py index 9a855b602e5d..f5514c5814c4 100644 --- a/tests/unit/batches/test_batch_line_item_logging.py +++ b/tests/unit/batches/test_batch_line_item_logging.py @@ -892,3 +892,75 @@ async def test_line_items_emit_when_the_claim_backend_returns_nothing(recorder): assert emitted == 2 assert len(recorder.success_events) == 1 assert len(recorder.failure_events) == 1 + + +_BROKEN_ERROR_FILE: Final = {**_FILE_BYTES, "error-file-1": "not bytes"} + + +def _broken_error_file_content(file_id: str, **_kwargs): + return SimpleNamespace(content=_BROKEN_ERROR_FILE[file_id]) + + +@pytest.mark.asyncio +async def test_line_items_retry_after_a_failed_fanout(recorder): + claim_cache: Final = DualCache() + batch: Final = _batch() + parent: Final = _parent_logging() + with patch("litellm.files.main.afile_content", new_callable=AsyncMock, side_effect=ValueError("boom")): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + first: Final = await log_batch_line_items( + batch=batch, + custom_llm_provider="openai", + parent=parent, + model_name="gpt-4o", + litellm_params=None, + model_info=None, + claim_cache=claim_cache, + ) + file_mock: Final = AsyncMock(side_effect=_file_content) + with patch("litellm.files.main.afile_content", file_mock): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + second: Final = await log_batch_line_items( + batch=batch, + custom_llm_provider="openai", + parent=parent, + model_name="gpt-4o", + litellm_params=None, + model_info=None, + claim_cache=claim_cache, + ) + + assert first == 0 + assert second == 2 + assert len(recorder.success_events) == 1 + assert len(recorder.failure_events) == 1 + + +@pytest.mark.asyncio +async def test_line_items_claim_kept_after_partial_emission(recorder): + claim_cache: Final = DualCache() + batch: Final = _batch() + parent: Final = _parent_logging() + file_mock: Final = AsyncMock(side_effect=_broken_error_file_content) + with patch("litellm.files.main.afile_content", file_mock): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + first: Final = await log_batch_line_items( + batch=batch, + custom_llm_provider="openai", + parent=parent, + model_name="gpt-4o", + litellm_params=None, + model_info=None, + claim_cache=claim_cache, + ) + second: Final = await log_batch_line_items( + batch=batch, + custom_llm_provider="openai", + parent=parent, + model_name="gpt-4o", + litellm_params=None, + model_info=None, + claim_cache=claim_cache, + ) + + assert first == 1 + assert second == 0 + assert len(recorder.success_events) == 1 + assert len(recorder.failure_events) == 0 From a7c9dedbf5a3dc77c7b9ce1f373477e4b0b8e9a8 Mon Sep 17 00:00:00 2001 From: yucheng Date: Tue, 29 Sep 2026 20:45:36 +0000 Subject: [PATCH 33/89] fix(batches): make the line-item claim release ownership-safe Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/batches/batch_line_item_logging.py | 64 ++++++-- .../batches/test_batch_line_item_logging.py | 151 +++++++++++++++++- 2 files changed, 203 insertions(+), 12 deletions(-) diff --git a/litellm/batches/batch_line_item_logging.py b/litellm/batches/batch_line_item_logging.py index 21dd23bdb599..94c2bb0bb95e 100644 --- a/litellm/batches/batch_line_item_logging.py +++ b/litellm/batches/batch_line_item_logging.py @@ -5,6 +5,8 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Final, Literal, TypeAlias, cast, get_args +from typing_extensions import assert_never + from litellm._logging import verbose_logger from litellm.batches.batch_utils import ( _batch_response_was_successful, # pyright: ignore[reportPrivateUsage] # batch-internal helper shared with the aggregate cost path by design @@ -68,6 +70,40 @@ def _supported_line_provider(value: str) -> _BatchLineProvider | None: batch_line_item_claim_cache: Final = DualCache() +_ClaimResult: TypeAlias = Literal["claimed", "already_claimed", "unavailable"] + + +async def _claim_line_items(claim_cache: DualCache, claim_key: str, token: str) -> _ClaimResult: + redis_cache: Final = claim_cache.redis_cache + if redis_cache is None: + count: Final = await claim_cache.async_increment_cache(claim_key, 1, ttl=_LINE_ITEM_CLAIM_TTL_SECONDS) # pyright: ignore[reportUnknownMemberType] # DualCache.increment is untyped upstream + if count is None or count == 1: + return "claimed" + return "already_claimed" + try: + ok: Final = await redis_cache.async_set_cache(claim_key, token, ttl=_LINE_ITEM_CLAIM_TTL_SECONDS, nx=True) # pyright: ignore[reportUnknownMemberType] # RedisCache.set is untyped upstream + except Exception: # noqa: BLE001 # a redis outage must not emit duplicates; the next retrieve retries + return "unavailable" + if ok: + return "claimed" + return "already_claimed" + + +async def _release_line_item_claim(claim_cache: DualCache, claim_key: str, token: str) -> None: + try: + redis_cache: Final = claim_cache.redis_cache + if redis_cache is None: + await claim_cache.async_delete_cache(claim_key) + return + owner: Final = await redis_cache.async_get_cache(claim_key) # pyright: ignore[reportUnknownMemberType] # RedisCache.get is untyped upstream + if owner == token: + await redis_cache.async_delete_cache(claim_key) + except Exception: # noqa: BLE001 # the claim release must never raise; worst case the batch stays claimed + verbose_logger.debug( + "batch line item claim release failed for %s, claim persists until ttl", + claim_key, + ) + class _BatchLineFailure(Exception): """A provider-reported per-line batch failure; carries the batch's hidden @@ -398,10 +434,22 @@ async def log_batch_line_items( ) return 0 claim_key: Final = f"batch_line_items_emitted:{batch.id}" - claim: Final = await claim_cache.async_increment_cache(claim_key, 1, ttl=_LINE_ITEM_CLAIM_TTL_SECONDS) - if claim is not None and claim > 1: - verbose_logger.debug("batch line items already emitted for batch_id=%s, skipping", batch.id) - return 0 + token: Final = uuid.uuid4().hex + claim: Final = await _claim_line_items(claim_cache, claim_key, token) + match claim: + case "already_claimed": + verbose_logger.debug("batch line items already emitted for batch_id=%s, skipping", batch.id) + return 0 + case "unavailable": + verbose_logger.warning( + "batch line item claim backend unavailable for batch_id=%s, line items will be retried on the next retrieve", + batch.id, + ) + return 0 + case "claimed": + pass + case _: + assert_never(claim) emitted = 0 # rebind-ok: loop accumulator for emitted line count try: internal_credentials: Final = parent._litellm_internal_model_credentials # pyright: ignore[reportPrivateUsage] # declared transport attribute on Logging @@ -447,11 +495,5 @@ async def log_batch_line_items( batch.id, ) if emitted == 0: - try: - await claim_cache.async_delete_cache(claim_key) - except Exception: # noqa: BLE001 # the claim release must never raise; worst case the batch stays claimed - verbose_logger.debug( - "batch line item claim release failed for batch_id=%s, claim persists until ttl", - batch.id, - ) + await _release_line_item_claim(claim_cache, claim_key, token) return emitted diff --git a/tests/unit/batches/test_batch_line_item_logging.py b/tests/unit/batches/test_batch_line_item_logging.py index f5514c5814c4..e78c0a1c3070 100644 --- a/tests/unit/batches/test_batch_line_item_logging.py +++ b/tests/unit/batches/test_batch_line_item_logging.py @@ -21,8 +21,13 @@ import pytest import litellm -from litellm.batches.batch_line_item_logging import batch_line_item_claim_cache, log_batch_line_items +from litellm.batches.batch_line_item_logging import ( + _release_line_item_claim, + batch_line_item_claim_cache, + log_batch_line_items, +) from litellm.caching.caching import DualCache +from litellm.caching.redis_cache import RedisCache from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging from litellm.types.utils import LiteLLMBatch, Usage @@ -964,3 +969,147 @@ async def test_line_items_claim_kept_after_partial_emission(recorder): assert second == 0 assert len(recorder.success_events) == 1 assert len(recorder.failure_events) == 0 + + +class _FakeRedisCache(RedisCache): + """Dict-backed RedisCache stand-in honoring SET NX so the tokenized + line-item claim can be exercised without a Redis process.""" + + def __init__(self) -> None: + self._store: dict[str, object] = {} # mutable-ok: a dict-backed fake needs a mutable store + self.fail_next_set = False + self.swap_owner_on_get: str | None = None + + async def async_set_cache(self, key, value, nx=False, **_kwargs): + if self.fail_next_set: + self.fail_next_set = False + raise ConnectionError("redis down") + if nx and key in self._store: + return None + self._store[key] = value + return True + + async def async_get_cache(self, key, **_kwargs): + if self.swap_owner_on_get is not None and key in self._store: + self._store[key] = self.swap_owner_on_get + return self._store.get(key) + + async def async_delete_cache(self, key): + self._store.pop(key, None) + + +@pytest.mark.asyncio +async def test_line_items_emit_once_per_batch_id_over_redis(recorder): + claim_cache: Final = DualCache(redis_cache=_FakeRedisCache()) + file_mock: Final = AsyncMock(side_effect=_file_content) + parent: Final = _parent_logging() + batch: Final = _batch() + with patch("litellm.files.main.afile_content", file_mock): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + first: Final = await log_batch_line_items( + batch=batch, + custom_llm_provider="openai", + parent=parent, + model_name="gpt-4o", + litellm_params=None, + model_info=None, + claim_cache=claim_cache, + ) + second: Final = await log_batch_line_items( + batch=batch, + custom_llm_provider="openai", + parent=parent, + model_name="gpt-4o", + litellm_params=None, + model_info=None, + claim_cache=claim_cache, + ) + + assert first == 2 + assert second == 0 + assert len(recorder.success_events) == 1 + assert len(recorder.failure_events) == 1 + + +@pytest.mark.asyncio +async def test_line_items_skip_when_redis_claim_backend_is_down(recorder): + fake: Final = _FakeRedisCache() + fake.fail_next_set = True + claim_cache: Final = DualCache(redis_cache=fake) + file_mock: Final = AsyncMock(side_effect=_file_content) + parent: Final = _parent_logging() + batch: Final = _batch() + with patch("litellm.files.main.afile_content", file_mock): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + first: Final = await log_batch_line_items( + batch=batch, + custom_llm_provider="openai", + parent=parent, + model_name="gpt-4o", + litellm_params=None, + model_info=None, + claim_cache=claim_cache, + ) + second: Final = await log_batch_line_items( + batch=batch, + custom_llm_provider="openai", + parent=parent, + model_name="gpt-4o", + litellm_params=None, + model_info=None, + claim_cache=claim_cache, + ) + + assert first == 0 + assert second == 2 + assert len(recorder.success_events) == 1 + assert len(recorder.failure_events) == 1 + + +@pytest.mark.asyncio +async def test_release_line_item_claim_only_deletes_owned_claims(): + fake: Final = _FakeRedisCache() + claim_cache: Final = DualCache(redis_cache=fake) + key: Final = "batch_line_items_emitted:batch_1" + fake._store[key] = "someone-else" # test-quality-ok: seeding the fake's store is the arrangement, like writing Redis directly + + await _release_line_item_claim(claim_cache, key, "my-token") + assert fake._store.get(key) == "someone-else" + + await _release_line_item_claim(claim_cache, key, "someone-else") + assert key not in fake._store + + +@pytest.mark.asyncio +async def test_line_items_failed_fanout_does_not_delete_another_workers_claim(recorder): + fake: Final = _FakeRedisCache() + fake.swap_owner_on_get = "other-worker" + claim_cache: Final = DualCache(redis_cache=fake) + key: Final = "batch_line_items_emitted:batch_1" + batch: Final = _batch() + parent: Final = _parent_logging() + with patch("litellm.files.main.afile_content", new_callable=AsyncMock, side_effect=ValueError("boom")): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + first: Final = await log_batch_line_items( + batch=batch, + custom_llm_provider="openai", + parent=parent, + model_name="gpt-4o", + litellm_params=None, + model_info=None, + claim_cache=claim_cache, + ) + file_mock: Final = AsyncMock(side_effect=_file_content) + with patch("litellm.files.main.afile_content", file_mock): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + second: Final = await log_batch_line_items( + batch=batch, + custom_llm_provider="openai", + parent=parent, + model_name="gpt-4o", + litellm_params=None, + model_info=None, + claim_cache=claim_cache, + ) + + assert first == 0 + assert second == 0 + assert fake._store.get(key) == "other-worker" + assert len(recorder.success_events) == 0 + assert len(recorder.failure_events) == 0 From f69d5c56bb60b28a86e907f3adc18a066d31e678 Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 30 Sep 2026 00:07:27 +0000 Subject: [PATCH 34/89] fix(batches): keep bodyless failed batch lines reaching callbacks Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/batches/batch_line_item_logging.py | 7 +- .../batches/test_batch_line_item_logging.py | 225 ++++++++++++++---- 2 files changed, 182 insertions(+), 50 deletions(-) diff --git a/litellm/batches/batch_line_item_logging.py b/litellm/batches/batch_line_item_logging.py index 94c2bb0bb95e..171c03d6f003 100644 --- a/litellm/batches/batch_line_item_logging.py +++ b/litellm/batches/batch_line_item_logging.py @@ -105,12 +105,17 @@ async def _release_line_item_claim(claim_cache: DualCache, claim_key: str, token ) +def _json_fallback(value: object) -> dict[str, object] | str: + mapping: Final = _as_object_mapping(value) + return dict(mapping) if mapping is not None else str(value) + + class _BatchLineFailure(Exception): """A provider-reported per-line batch failure; carries the batch's hidden params so the failure logging payload can attribute the line.""" def __init__(self, error_payload: object) -> None: - super().__init__(json.dumps(error_payload)) + super().__init__(json.dumps(error_payload, default=_json_fallback)) self._hidden_params: dict[str, object] = {} # mutable-ok: plain-dict contract like response _hidden_params diff --git a/tests/unit/batches/test_batch_line_item_logging.py b/tests/unit/batches/test_batch_line_item_logging.py index e78c0a1c3070..ca99f6229aae 100644 --- a/tests/unit/batches/test_batch_line_item_logging.py +++ b/tests/unit/batches/test_batch_line_item_logging.py @@ -136,10 +136,16 @@ def _fresh_line_item_claim_cache(): @pytest.fixture def recorder(): logger = _RecordingLogger() - saved_flag = litellm.store_batch_line_items_in_callbacks # test-quality-ok: process-wide opt-in flag; restored in teardown - saved_success = list(litellm._async_success_callback) # test-quality-ok: the feature dispatches through this global list; restored in teardown + saved_flag = ( + litellm.store_batch_line_items_in_callbacks + ) # test-quality-ok: process-wide opt-in flag; restored in teardown + saved_success = list( + litellm._async_success_callback + ) # test-quality-ok: the feature dispatches through this global list; restored in teardown saved_failure = list(litellm._async_failure_callback) # test-quality-ok: same dispatch seam, restored in teardown - litellm._async_success_callback = [logger] # test-quality-ok: there is no injection seam for callback lists; teardown restores + litellm._async_success_callback = [ + logger + ] # test-quality-ok: there is no injection seam for callback lists; teardown restores litellm._async_failure_callback = [logger] # test-quality-ok: same dispatch seam, restored in teardown yield logger litellm.store_batch_line_items_in_callbacks = saved_flag # test-quality-ok: teardown restoring the value set above @@ -188,11 +194,17 @@ def _hidden(event: dict) -> dict: @pytest.mark.asyncio async def test_line_items_emitted_alongside_aggregate(recorder): - litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: the flag under test is a module global; fixture restores it + litellm.store_batch_line_items_in_callbacks = ( + True # test-quality-ok: the flag under test is a module global; fixture restores it + ) batch = _batch() with ( - patch("litellm.files.main.afile_content", new_callable=AsyncMock, side_effect=_file_content), # test-quality-ok: afile_content is the provider boundary; there is no HTTP transport or injection seam for managed file fetch - patch("litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02)), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch + patch( + "litellm.files.main.afile_content", new_callable=AsyncMock, side_effect=_file_content + ), # test-quality-ok: afile_content is the provider boundary; there is no HTTP transport or injection seam for managed file fetch + patch( + "litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02) + ), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch ): await _log_completed_batch(_parent_logging(), batch) @@ -228,7 +240,9 @@ async def test_line_items_emitted_alongside_aggregate(recorder): async def test_flag_off_emits_only_aggregate(recorder): assert litellm.store_batch_line_items_in_callbacks is False file_mock = AsyncMock(side_effect=_file_content) - with patch("litellm.files.main.afile_content", file_mock): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + with patch( + "litellm.files.main.afile_content", file_mock + ): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch await _log_completed_batch(_parent_logging(), _batch()) assert len(recorder.success_events) == 1 @@ -238,9 +252,13 @@ async def test_flag_off_emits_only_aggregate(recorder): @pytest.mark.asyncio async def test_line_items_skipped_for_unsupported_provider(recorder): - litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: the flag under test is a module global; fixture restores it + litellm.store_batch_line_items_in_callbacks = ( + True # test-quality-ok: the flag under test is a module global; fixture restores it + ) file_mock: Final = AsyncMock(side_effect=_file_content) - with patch("litellm.files.main.afile_content", file_mock): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + with patch( + "litellm.files.main.afile_content", file_mock + ): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch await _log_completed_batch(_parent_logging(custom_llm_provider="cohere"), _batch()) file_mock.assert_not_awaited() @@ -252,7 +270,9 @@ async def test_line_items_skipped_for_unsupported_provider(recorder): @pytest.mark.asyncio async def test_in_progress_batch_poll_emits_no_line_events(recorder): - litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: the flag under test is a module global; fixture restores it + litellm.store_batch_line_items_in_callbacks = ( + True # test-quality-ok: the flag under test is a module global; fixture restores it + ) in_progress: Final = LiteLLMBatch( id="batch_wip", object="batch", @@ -265,7 +285,9 @@ async def test_in_progress_batch_poll_emits_no_line_events(recorder): created_at=1, ) file_mock: Final = AsyncMock(side_effect=_file_content) - with patch("litellm.files.main.afile_content", file_mock): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + with patch( + "litellm.files.main.afile_content", file_mock + ): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch await _parent_logging().async_success_handler(result=in_progress) file_mock.assert_not_called() @@ -275,14 +297,19 @@ async def test_in_progress_batch_poll_emits_no_line_events(recorder): @pytest.mark.asyncio async def test_input_fetch_failure_still_emits_aggregate(recorder): - litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: the flag under test is a module global; fixture restores it - with patch("litellm.files.main.afile_content", new_callable=AsyncMock, side_effect=ValueError("boom")): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + litellm.store_batch_line_items_in_callbacks = ( + True # test-quality-ok: the flag under test is a module global; fixture restores it + ) + with patch( + "litellm.files.main.afile_content", new_callable=AsyncMock, side_effect=ValueError("boom") + ): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch await _log_completed_batch(_parent_logging(), _batch()) assert len(recorder.success_events) == 1 assert len(recorder.failure_events) == 0 assert _payload(recorder.success_events[0])["response_cost"] == 1.5 + EDGE_INPUT_JSONL = b"\n".join( [ json.dumps( @@ -519,7 +546,9 @@ def _parent_logging_with_params(litellm_params: dict) -> Logging: @pytest.mark.asyncio async def test_line_items_edge_shapes_and_edge_cases(recorder): - litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: the flag under test is a module global; fixture restores it + litellm.store_batch_line_items_in_callbacks = ( + True # test-quality-ok: the flag under test is a module global; fixture restores it + ) batch = LiteLLMBatch( id="batch_edge", object="batch", @@ -537,10 +566,16 @@ async def test_line_items_edge_shapes_and_edge_cases(recorder): "metadata": {"model_info": {"id": "dep-1"}, "model_group": "gpt-4o"}, } ) - parent._litellm_internal_model_credentials = {"api_key": "sk-line-items-marker"} # test-quality-ok: private transport attribute, same channel the batch cost tracker uses + parent._litellm_internal_model_credentials = { + "api_key": "sk-line-items-marker" + } # test-quality-ok: private transport attribute, same channel the batch cost tracker uses with ( - patch("litellm.files.main.afile_content", file_mock), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch - patch("litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02)), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch + patch( + "litellm.files.main.afile_content", file_mock + ), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + patch( + "litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02) + ), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch ): await _log_completed_batch(parent, batch) @@ -575,7 +610,9 @@ async def test_line_items_edge_shapes_and_edge_cases(recorder): @pytest.mark.asyncio async def test_line_items_child_params_drop_parent_credentials(recorder): - litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: the flag under test is a module global; fixture restores it + litellm.store_batch_line_items_in_callbacks = ( + True # test-quality-ok: the flag under test is a module global; fixture restores it + ) file_mock: Final = AsyncMock(side_effect=_file_content) parent: Final = _parent_logging_with_params( { @@ -587,15 +624,17 @@ async def test_line_items_child_params_drop_parent_credentials(recorder): ) batch: Final = _batch() with ( - patch("litellm.files.main.afile_content", file_mock), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch - patch("litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02)), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch + patch( + "litellm.files.main.afile_content", file_mock + ), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + patch( + "litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02) + ), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch ): await _log_completed_batch(parent, batch) line_events = [ - e - for e in [*recorder.success_events, *recorder.failure_events] - if _hidden(e).get("batch_custom_id") is not None + e for e in [*recorder.success_events, *recorder.failure_events] if _hidden(e).get("batch_custom_id") is not None ] assert len(line_events) == 2 for event in line_events: @@ -608,7 +647,9 @@ async def test_line_items_child_params_drop_parent_credentials(recorder): @pytest.mark.asyncio async def test_line_items_anthropic_shapes(recorder): - litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: the flag under test is a module global; fixture restores it + litellm.store_batch_line_items_in_callbacks = ( + True # test-quality-ok: the flag under test is a module global; fixture restores it + ) batch = LiteLLMBatch( id="batch_anth", object="batch", @@ -628,8 +669,12 @@ async def test_line_items_anthropic_shapes(recorder): custom_llm_provider="anthropic", ) with ( - patch("litellm.files.main.afile_content", file_mock), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch - patch("litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02)), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch + patch( + "litellm.files.main.afile_content", file_mock + ), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + patch( + "litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02) + ), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch ): await _log_completed_batch(logging_obj, batch) @@ -658,11 +703,17 @@ def _provider_batch(batch_id: str, input_file_id: str, output_file_id: str) -> L @pytest.mark.asyncio async def test_line_items_bedrock_shapes(recorder): - litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: the flag under test is a module global; fixture restores it + litellm.store_batch_line_items_in_callbacks = ( + True # test-quality-ok: the flag under test is a module global; fixture restores it + ) file_mock: Final = AsyncMock(side_effect=_edge_file_content) with ( - patch("litellm.files.main.afile_content", file_mock), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch - patch("litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02)), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch + patch( + "litellm.files.main.afile_content", file_mock + ), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + patch( + "litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02) + ), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch ): await _log_completed_batch( _parent_logging(custom_llm_provider="bedrock"), @@ -699,11 +750,17 @@ async def test_line_items_bedrock_shapes(recorder): @pytest.mark.asyncio async def test_line_items_mistral_shape(recorder): - litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: the flag under test is a module global; fixture restores it + litellm.store_batch_line_items_in_callbacks = ( + True # test-quality-ok: the flag under test is a module global; fixture restores it + ) file_mock: Final = AsyncMock(side_effect=_edge_file_content) with ( - patch("litellm.files.main.afile_content", file_mock), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch - patch("litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02)), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch + patch( + "litellm.files.main.afile_content", file_mock + ), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + patch( + "litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02) + ), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch ): await _log_completed_batch( _parent_logging(custom_llm_provider="mistral"), @@ -762,7 +819,9 @@ def _scoped_file_content(file_map): @pytest.mark.asyncio async def test_line_items_anthropic_failure_keeps_provider_error(recorder): - litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: the flag under test is a module global; fixture restores it + litellm.store_batch_line_items_in_callbacks = ( + True # test-quality-ok: the flag under test is a module global; fixture restores it + ) file_mock: Final = AsyncMock( side_effect=_scoped_file_content( { @@ -772,8 +831,12 @@ async def test_line_items_anthropic_failure_keeps_provider_error(recorder): ) ) with ( - patch("litellm.files.main.afile_content", file_mock), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch - patch("litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02)), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch + patch( + "litellm.files.main.afile_content", file_mock + ), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + patch( + "litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02) + ), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch ): await _log_completed_batch( _parent_logging(custom_llm_provider="anthropic"), @@ -787,9 +850,47 @@ async def test_line_items_anthropic_failure_keeps_provider_error(recorder): assert "max_tokens must be positive" in _payload(failure)["error_str"] +BODYLESS_ERR_INPUT_JSONL = json.dumps( + {"custom_id": "bare", "url": "/v1/chat/completions", "body": {"model": "gpt-4.1-nano", "messages": []}} +).encode() + +BODYLESS_ERR_OUTPUT_JSONL = json.dumps({"custom_id": "bare", "response": None, "error": None}).encode() + + +@pytest.mark.asyncio +async def test_line_items_failure_without_error_or_response_body_still_reaches_callbacks(recorder): + litellm.store_batch_line_items_in_callbacks = ( + True # test-quality-ok: the flag under test is a module global; fixture restores it + ) + file_mock: Final = AsyncMock( + side_effect=_scoped_file_content( + { + "input-bare": BODYLESS_ERR_INPUT_JSONL, + "output-bare": BODYLESS_ERR_OUTPUT_JSONL, + } + ) + ) + with ( + patch( + "litellm.files.main.afile_content", file_mock + ), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + patch( + "litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02) + ), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch + ): + await _log_completed_batch(_parent_logging(), _provider_batch("batch_bare", "input-bare", "output-bare")) + + assert len(recorder.failure_events) == 1 + failure = recorder.failure_events[0] + assert _hidden(failure)["batch_custom_id"] == "bare" + assert _payload(failure)["error_str"] == "{}" + + @pytest.mark.asyncio async def test_line_items_native_vertex_rows_are_skipped(recorder): - litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: the flag under test is a module global; fixture restores it + litellm.store_batch_line_items_in_callbacks = ( + True # test-quality-ok: the flag under test is a module global; fixture restores it + ) file_mock: Final = AsyncMock( side_effect=_scoped_file_content( { @@ -799,8 +900,12 @@ async def test_line_items_native_vertex_rows_are_skipped(recorder): ) ) with ( - patch("litellm.files.main.afile_content", file_mock), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch - patch("litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02)), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch + patch( + "litellm.files.main.afile_content", file_mock + ), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + patch( + "litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02) + ), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch ): await _log_completed_batch( _parent_logging(custom_llm_provider="vertex_ai"), @@ -823,7 +928,9 @@ async def test_line_items_emit_once_per_batch_id(recorder): claim_cache: Final = DualCache() parent: Final = _parent_logging() batch: Final = _batch() - with patch("litellm.files.main.afile_content", file_mock): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + with patch( + "litellm.files.main.afile_content", file_mock + ): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch first: Final = await log_batch_line_items( batch=batch, custom_llm_provider="openai", @@ -854,7 +961,9 @@ async def test_line_items_claim_is_scoped_to_the_batch_id(recorder): file_mock: Final = AsyncMock(side_effect=_file_content) claim_cache: Final = DualCache() parent: Final = _parent_logging() - with patch("litellm.files.main.afile_content", file_mock): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + with patch( + "litellm.files.main.afile_content", file_mock + ): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch first: Final = await log_batch_line_items( batch=_batch(), custom_llm_provider="openai", @@ -883,7 +992,9 @@ async def test_line_items_claim_is_scoped_to_the_batch_id(recorder): @pytest.mark.asyncio async def test_line_items_emit_when_the_claim_backend_returns_nothing(recorder): file_mock: Final = AsyncMock(side_effect=_file_content) - with patch("litellm.files.main.afile_content", file_mock): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + with patch( + "litellm.files.main.afile_content", file_mock + ): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch emitted: Final = await log_batch_line_items( batch=_batch(), custom_llm_provider="openai", @@ -911,7 +1022,9 @@ async def test_line_items_retry_after_a_failed_fanout(recorder): claim_cache: Final = DualCache() batch: Final = _batch() parent: Final = _parent_logging() - with patch("litellm.files.main.afile_content", new_callable=AsyncMock, side_effect=ValueError("boom")): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + with patch( + "litellm.files.main.afile_content", new_callable=AsyncMock, side_effect=ValueError("boom") + ): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch first: Final = await log_batch_line_items( batch=batch, custom_llm_provider="openai", @@ -922,7 +1035,9 @@ async def test_line_items_retry_after_a_failed_fanout(recorder): claim_cache=claim_cache, ) file_mock: Final = AsyncMock(side_effect=_file_content) - with patch("litellm.files.main.afile_content", file_mock): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + with patch( + "litellm.files.main.afile_content", file_mock + ): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch second: Final = await log_batch_line_items( batch=batch, custom_llm_provider="openai", @@ -945,7 +1060,9 @@ async def test_line_items_claim_kept_after_partial_emission(recorder): batch: Final = _batch() parent: Final = _parent_logging() file_mock: Final = AsyncMock(side_effect=_broken_error_file_content) - with patch("litellm.files.main.afile_content", file_mock): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + with patch( + "litellm.files.main.afile_content", file_mock + ): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch first: Final = await log_batch_line_items( batch=batch, custom_llm_provider="openai", @@ -1004,7 +1121,9 @@ async def test_line_items_emit_once_per_batch_id_over_redis(recorder): file_mock: Final = AsyncMock(side_effect=_file_content) parent: Final = _parent_logging() batch: Final = _batch() - with patch("litellm.files.main.afile_content", file_mock): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + with patch( + "litellm.files.main.afile_content", file_mock + ): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch first: Final = await log_batch_line_items( batch=batch, custom_llm_provider="openai", @@ -1038,7 +1157,9 @@ async def test_line_items_skip_when_redis_claim_backend_is_down(recorder): file_mock: Final = AsyncMock(side_effect=_file_content) parent: Final = _parent_logging() batch: Final = _batch() - with patch("litellm.files.main.afile_content", file_mock): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + with patch( + "litellm.files.main.afile_content", file_mock + ): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch first: Final = await log_batch_line_items( batch=batch, custom_llm_provider="openai", @@ -1069,7 +1190,9 @@ async def test_release_line_item_claim_only_deletes_owned_claims(): fake: Final = _FakeRedisCache() claim_cache: Final = DualCache(redis_cache=fake) key: Final = "batch_line_items_emitted:batch_1" - fake._store[key] = "someone-else" # test-quality-ok: seeding the fake's store is the arrangement, like writing Redis directly + fake._store[key] = ( + "someone-else" # test-quality-ok: seeding the fake's store is the arrangement, like writing Redis directly + ) await _release_line_item_claim(claim_cache, key, "my-token") assert fake._store.get(key) == "someone-else" @@ -1086,7 +1209,9 @@ async def test_line_items_failed_fanout_does_not_delete_another_workers_claim(re key: Final = "batch_line_items_emitted:batch_1" batch: Final = _batch() parent: Final = _parent_logging() - with patch("litellm.files.main.afile_content", new_callable=AsyncMock, side_effect=ValueError("boom")): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + with patch( + "litellm.files.main.afile_content", new_callable=AsyncMock, side_effect=ValueError("boom") + ): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch first: Final = await log_batch_line_items( batch=batch, custom_llm_provider="openai", @@ -1097,7 +1222,9 @@ async def test_line_items_failed_fanout_does_not_delete_another_workers_claim(re claim_cache=claim_cache, ) file_mock: Final = AsyncMock(side_effect=_file_content) - with patch("litellm.files.main.afile_content", file_mock): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + with patch( + "litellm.files.main.afile_content", file_mock + ): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch second: Final = await log_batch_line_items( batch=batch, custom_llm_provider="openai", From 5f35c25ce1726c4a79a28e09363abcb19d13c186 Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 30 Sep 2026 09:20:37 +0000 Subject: [PATCH 35/89] fix(logging): seed batch line hidden params in get_hidden_params defaults Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/litellm_core_utils/litellm_logging.py | 3 +++ tests/unit/batches/test_batch_line_item_logging.py | 4 ++-- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 1365d22d0384..83485d1efca4 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -5951,6 +5951,9 @@ def get_hidden_params( batch_failed_requests=None, litellm_model_name=None, usage_object=None, + batch_id=None, + batch_custom_id=None, + batch_line_status_code=None, ) if hidden_params is not None: for key in StandardLoggingHiddenParams.__annotations__: diff --git a/tests/unit/batches/test_batch_line_item_logging.py b/tests/unit/batches/test_batch_line_item_logging.py index ca99f6229aae..877b6b67854f 100644 --- a/tests/unit/batches/test_batch_line_item_logging.py +++ b/tests/unit/batches/test_batch_line_item_logging.py @@ -213,7 +213,7 @@ async def test_line_items_emitted_alongside_aggregate(recorder): aggregate = next(e for e in recorder.success_events if _hidden(e).get("batch_custom_id") is None) assert _payload(aggregate)["response_cost"] == 1.5 - assert "batch_custom_id" not in _hidden(aggregate) + assert _hidden(aggregate)["batch_custom_id"] is None line = next(e for e in recorder.success_events if _hidden(e).get("batch_custom_id") == "a") hidden = _hidden(line) @@ -264,7 +264,7 @@ async def test_line_items_skipped_for_unsupported_provider(recorder): file_mock.assert_not_awaited() assert len(recorder.success_events) == 1 assert _payload(recorder.success_events[0])["response_cost"] == 1.5 - assert "batch_custom_id" not in _hidden(recorder.success_events[0]) + assert _hidden(recorder.success_events[0])["batch_custom_id"] is None assert len(recorder.failure_events) == 0 From ff77170749b28833fe5ebe061557ed45aeab2c48 Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 30 Sep 2026 23:35:59 +0000 Subject: [PATCH 36/89] perf(batches): reuse fetched batch result files for line-item logging Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/common_utils/check_batch_cost.py | 8 +- litellm/batches/batch_line_item_logging.py | 18 +- litellm/batches/batch_utils.py | 104 +++++++--- litellm/litellm_core_utils/litellm_logging.py | 40 +++- .../test_batches_logging_unit_tests.py | 44 ++-- .../batches/test_batch_line_item_logging.py | 47 +++++ tests/unit/batches/test_batch_utils.py | 54 +++++ .../test_litellm_logging.py | 195 ++++++++++++++++-- .../common_utils/test_check_batch_cost.py | 8 +- 9 files changed, 433 insertions(+), 85 deletions(-) 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 db28317ae2c5..cb41eba32eba 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -705,9 +705,10 @@ async def _track_completed_batch_cost( later poll. """ from litellm.batches.batch_utils import ( - count_error_file_failed_requests, _get_file_content_as_dictionary, calculate_batch_cost_and_usage, + count_error_file_failed_requests_from_content, + fetch_batch_error_file_content, ) from litellm.files.main import afile_content from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider @@ -850,7 +851,7 @@ async def _track_completed_batch_cost( model_name=model_name, model_info=deployment_model_info, ) - error_file_failed_requests: Final = await count_error_file_failed_requests( + error_file_content = await fetch_batch_error_file_content( response, custom_llm_provider=batch_file_provider, litellm_params={ @@ -858,6 +859,7 @@ async def _track_completed_batch_cost( "_litellm_internal_model_credentials": MappingProxyType(dict(credentials)), }, ) + error_file_failed_requests: Final = count_error_file_failed_requests_from_content(error_file_content) batch_result: Final = ( output_file_result if not error_file_failed_requests @@ -916,6 +918,8 @@ async def _track_completed_batch_cost( batch_failed_requests=batch_result.failed_requests, batch_prompt_cost=batch_result.prompt_cost, batch_completion_cost=batch_result.completion_cost, + batch_output_file_content=content_bytes, + batch_error_file_content=error_file_content, ) except Exception: await self._release_job_claim(job) diff --git a/litellm/batches/batch_line_item_logging.py b/litellm/batches/batch_line_item_logging.py index 171c03d6f003..db9f0c9b6b18 100644 --- a/litellm/batches/batch_line_item_logging.py +++ b/litellm/batches/batch_line_item_logging.py @@ -9,6 +9,7 @@ from litellm._logging import verbose_logger from litellm.batches.batch_utils import ( + BatchResultFiles, _batch_response_was_successful, # pyright: ignore[reportPrivateUsage] # batch-internal helper shared with the aggregate cost path by design _fetch_batch_managed_file_content, # pyright: ignore[reportPrivateUsage, reportUnknownVariableType] # same reuse; helper is untyped upstream _get_response_from_batch_job_output_file, # pyright: ignore[reportPrivateUsage] # same reuse @@ -422,6 +423,7 @@ async def log_batch_line_items( model_name: str | None, litellm_params: dict[str, object] | None, # mutable-ok: the logging object's shared litellm_params dict model_info: ModelInfo | None, + result_files: BatchResultFiles | None = None, claim_cache: DualCache = batch_line_item_claim_cache, ) -> int: """Emit one callback event per JSONL line of a completed batch (request @@ -429,7 +431,9 @@ async def log_batch_line_items( ``litellm.store_batch_line_items_in_callbacks`` flag. The aggregate aretrieve_batch event still bills the batch, so per-line events carry ``batch_parent_id`` and never update spend themselves. Any failure here - is logged and swallowed: aggregate accounting must be unaffected.""" + is logged and swallowed: aggregate accounting must be unaffected. + ``result_files`` carries the output/error bytes the aggregate path already + fetched, so they are reused instead of refetched.""" line_provider: Final = _supported_line_provider(custom_llm_provider) if line_provider is None: verbose_logger.warning( @@ -468,7 +472,11 @@ async def log_batch_line_items( input_file_content: Final = await _fetch_managed_file_or_empty(batch.input_file_id, line_provider, fetch_params) requests_by_id: Final = _requests_by_custom_id(input_file_content) - output_content: Final = await _fetch_managed_file_or_empty(batch.output_file_id, line_provider, fetch_params) + output_content: Final = ( + result_files.output + if result_files is not None and result_files.output is not None + else await _fetch_managed_file_or_empty(batch.output_file_id, line_provider, fetch_params) + ) first_row: Final = next(_output_entries(output_content), None) if _uses_native_vertex_output(line_provider, model_name, first_row): verbose_logger.warning( @@ -476,7 +484,11 @@ async def log_batch_line_items( batch.id, ) return 0 - error_content: Final = await _fetch_managed_file_or_empty(batch.error_file_id, line_provider, fetch_params) + error_content: Final = ( + result_files.error + if result_files is not None and result_files.error is not None + else await _fetch_managed_file_or_empty(batch.error_file_id, line_provider, fetch_params) + ) for content in (output_content, error_content): for entry in _output_entries(content): try: diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 8160618f2c54..3660dfed94be 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -33,6 +33,14 @@ class BatchCostUsageResult: completion_cost: float = 0.0 +@dataclass(frozen=True, slots=True) +class BatchResultFiles: + """Raw JSONL bytes of a completed batch's result files; None means not fetched.""" + + output: bytes | None + error: bytes | None + + _COMPLETED_BATCH_STATUSES: Final = frozenset({"completed", "complete"}) @@ -115,29 +123,44 @@ async def _handle_completed_batch( threaded through so a deployment's configured rates win over the global cost map. """ - # A completed batch whose request lines all failed has no output file - the - # results are written to a separate error_file_id and output_file_id is None. - # There is nothing to price or measure, so report an empty result set instead - # of calling _fetch_batch_output_file_content, which raises on a missing - # output file. Without this guard the logging worker crashes on every - # aretrieve_batch poll and the completed batch's zero-cost accounting is lost. - # The generic retrieval helper keeps raising for callers that explicitly ask - # for a missing output file. + return ( + await _handle_completed_batch_with_files( + batch, + custom_llm_provider, + model_name=model_name, + litellm_params=litellm_params, + model_info=model_info, + ) + )[0] + + +async def _handle_completed_batch_with_files( + batch: Batch, + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral"], + model_name: str | None = None, + litellm_params: dict | None = None, + model_info: ModelInfo | None = None, +) -> "tuple[BatchCostUsageResult, BatchResultFiles]": + """_handle_completed_batch plus the raw output/error bytes it fetched, so + downstream line-item logging reuses them instead of refetching.""" + error_file_content: Final = await fetch_batch_error_file_content( + batch, custom_llm_provider=custom_llm_provider, litellm_params=litellm_params + ) + error_file_failed_requests: Final = count_error_file_failed_requests_from_content(error_file_content) + if batch.output_file_id is None: - return BatchCostUsageResult( - cost=0.0, - usage=Usage(prompt_tokens=0, completion_tokens=0, total_tokens=0), - models=[], # mutable-ok: no output file means no model was ever priced; BatchCostUsageResult.models requires list[str] - successful_requests=0, - failed_requests=await count_error_file_failed_requests( - batch, custom_llm_provider=custom_llm_provider, litellm_params=litellm_params + return ( + BatchCostUsageResult( + cost=0.0, + usage=Usage(prompt_tokens=0, completion_tokens=0, total_tokens=0), + models=[], # mutable-ok: same contract as _handle_completed_batch above + successful_requests=0, + failed_requests=error_file_failed_requests, ), + BatchResultFiles(output=None, error=error_file_content), ) file_content = await _fetch_batch_output_file_content(batch, custom_llm_provider, litellm_params=litellm_params) - error_file_failed_requests: Final = await count_error_file_failed_requests( - batch, custom_llm_provider=custom_llm_provider, litellm_params=litellm_params - ) output_file_result: Final = ( calculate_vertex_ai_batch_cost_and_usage( @@ -154,11 +177,14 @@ async def _handle_completed_batch( ) ) - if not error_file_failed_requests: - return output_file_result - return dataclasses_replace( - output_file_result, failed_requests=output_file_result.failed_requests + error_file_failed_requests + result: Final = ( + output_file_result + if not error_file_failed_requests + else dataclasses_replace( + output_file_result, failed_requests=output_file_result.failed_requests + error_file_failed_requests + ) ) + return result, BatchResultFiles(output=file_content, error=error_file_content) class _LineOutcome(Enum): @@ -488,6 +514,27 @@ async def _fetch_batch_output_file_content( ) +async def fetch_batch_error_file_content( + batch: Batch, + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral"], + litellm_params: dict | None, +) -> bytes | None: + """Fetch the batch's separate error file bytes; None when it has none or the fetch fails.""" + if batch.error_file_id is None: + return None + try: + return await _fetch_batch_managed_file_content( + batch.error_file_id, custom_llm_provider=custom_llm_provider, litellm_params=litellm_params + ) + except Exception as e: # noqa: BLE001 # a failed/missing error file must not abort cost tracking for the batch + verbose_logger.debug("Failed to fetch batch error file %s: %s", batch.error_file_id, e) + return None + + +def count_error_file_failed_requests_from_content(content: bytes | None) -> int: + return 0 if content is None else sum(1 for _ in _iter_batch_input_lines(content)) + + async def count_error_file_failed_requests( batch: Batch, custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral"], @@ -500,16 +547,11 @@ async def count_error_file_failed_requests( ``error_file_id`` - they never appear in the output file at all, so counting failures from the output file alone silently undercounts them. """ - if batch.error_file_id is None: - return 0 - try: - error_file_content = await _fetch_batch_managed_file_content( - batch.error_file_id, custom_llm_provider=custom_llm_provider, litellm_params=litellm_params + return count_error_file_failed_requests_from_content( + await fetch_batch_error_file_content( + batch, custom_llm_provider=custom_llm_provider, litellm_params=litellm_params ) - except Exception as e: # noqa: BLE001 # a failed/missing error file must not abort cost tracking for the batch - verbose_logger.debug("Failed to fetch batch error file %s: %s", batch.error_file_id, e) - return 0 - return sum(1 for _ in _iter_batch_input_lines(error_file_content)) + ) def _extract_file_access_credentials(litellm_params: dict | None) -> dict: diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index de573d04f98c..d5b111a41cf9 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -32,7 +32,11 @@ verbose_logger, ) from litellm._uuid import uuid -from litellm.batches.batch_utils import _handle_completed_batch, batch_cost_is_final +from litellm.batches.batch_utils import ( + BatchResultFiles, + _handle_completed_batch_with_files, + batch_cost_is_final, +) from litellm.caching.caching import DualCache from litellm.caching.caching_handler import LLMCachingHandler from litellm.caching.redis_batch import flush_post_call_redis_batches @@ -3219,9 +3223,12 @@ async def _async_success_handler_body( batch_models = kwargs.get("batch_models", None) batch_successful_requests: Final = kwargs.get("batch_successful_requests", None) batch_failed_requests: Final = kwargs.get("batch_failed_requests", None) + batch_output_file_content: Final = kwargs.get("batch_output_file_content", None) + batch_error_file_content: Final = kwargs.get("batch_error_file_content", None) has_explicit_batch_data: Final = all(x is not None for x in (batch_cost, batch_usage, batch_models)) should_compute_batch_data: Final = not has_explicit_batch_data and batch_cost_is_final(result) + result_files: BatchResultFiles | None = None # rebind-ok: one batch-data branch below supplies the bytes if has_explicit_batch_data: result._hidden_params["response_cost"] = batch_cost result._hidden_params["batch_models"] = batch_models @@ -3241,15 +3248,21 @@ async def _async_success_handler_body( total_cost=batch_cost, cost_for_built_in_tools_cost_usd_dollar=0.0, ) + if batch_output_file_content is not None or batch_error_file_content is not None: + result_files = BatchResultFiles( + output=batch_output_file_content if isinstance(batch_output_file_content, bytes) else None, + error=batch_error_file_content if isinstance(batch_error_file_content, bytes) else None, + ) elif should_compute_batch_data: - batch_result: Final = await _handle_completed_batch( + batch_result, fetched_result_files = await _handle_completed_batch_with_files( batch=result, custom_llm_provider=self.custom_llm_provider, model_name=self.get_deployment_model_for_cost(), litellm_params=self.litellm_params, model_info=self.get_router_deployment_model_info(), ) + result_files = fetched_result_files result._hidden_params["response_cost"] = batch_result.cost result._hidden_params["batch_models"] = batch_result.models @@ -3266,14 +3279,21 @@ async def _async_success_handler_body( if litellm.store_batch_line_items_in_callbacks and (has_explicit_batch_data or should_compute_batch_data): from litellm.batches.batch_line_item_logging import log_batch_line_items - await log_batch_line_items( - batch=result, - custom_llm_provider=self.custom_llm_provider, - parent=self, - model_name=self.get_deployment_model_for_cost(), - litellm_params=self.litellm_params, - model_info=self.get_router_deployment_model_info(), - ) + try: + await log_batch_line_items( + batch=result, + custom_llm_provider=self.custom_llm_provider, + parent=self, + model_name=self.get_deployment_model_for_cost(), + litellm_params=self.litellm_params, + model_info=self.get_router_deployment_model_info(), + result_files=result_files, + ) + except Exception: # noqa: BLE001 # line-item logging (claim step included) must never reach the aggregate aretrieve_batch path + verbose_logger.exception( + "batch line item logging failed for batch_id=%s; aggregate logging unaffected", + result.id, + ) self.truncated_messages_for_logging = await truncate_base64_in_messages_async( StandardLoggingPayloadSetup.append_system_prompt_messages( diff --git a/tests/batches_tests/test_batches_logging_unit_tests.py b/tests/batches_tests/test_batches_logging_unit_tests.py index 73adb3914816..e8a90b385f21 100644 --- a/tests/batches_tests/test_batches_logging_unit_tests.py +++ b/tests/batches_tests/test_batches_logging_unit_tests.py @@ -222,8 +222,8 @@ async def test_batch_retrieve_cost_tracking_with_completed_batch_no_explicit_cos ) logging_obj.custom_llm_provider = "openai" - # Mock _handle_completed_batch to return cost data - from litellm.batches.batch_utils import BatchCostUsageResult + # Mock _handle_completed_batch_with_files to return cost data + from litellm.batches.batch_utils import BatchCostUsageResult, BatchResultFiles expected_cost = 0.05 expected_usage = litellm.Usage( @@ -234,14 +234,17 @@ async def test_batch_retrieve_cost_tracking_with_completed_batch_no_explicit_cos expected_models = ["gpt-5-mini"] with patch( - "litellm.litellm_core_utils.litellm_logging._handle_completed_batch", + "litellm.litellm_core_utils.litellm_logging._handle_completed_batch_with_files", new=AsyncMock( - return_value=BatchCostUsageResult( - cost=expected_cost, - usage=expected_usage, - models=expected_models, - successful_requests=10, - failed_requests=0, + return_value=( + BatchCostUsageResult( + cost=expected_cost, + usage=expected_usage, + models=expected_models, + successful_requests=10, + failed_requests=0, + ), + BatchResultFiles(output=None, error=None), ) ), ) as mock_handle_batch: @@ -382,7 +385,7 @@ async def test_batch_retrieve_cost_tracking_with_explicit_cost_data(): explicit_models = ["gpt-5-mini", "gpt-5.5"] with patch( - "litellm.litellm_core_utils.litellm_logging._handle_completed_batch", + "litellm.litellm_core_utils.litellm_logging._handle_completed_batch_with_files", new=AsyncMock(), ) as mock_handle_batch: # Call async_success_handler with explicit cost data @@ -517,7 +520,7 @@ async def test_batch_retrieve_cost_tracking_with_unified_file_id_incomplete_batc logging_obj.custom_llm_provider = "openai" with patch( - "litellm.litellm_core_utils.litellm_logging._handle_completed_batch", + "litellm.litellm_core_utils.litellm_logging._handle_completed_batch_with_files", new=AsyncMock(), ) as mock_handle_batch: # Call async_success_handler with in_progress batch (unified file ID) @@ -603,17 +606,20 @@ async def test_batch_retrieve_cost_tracking_with_partial_explicit_data(): ) expected_models = ["gpt-5-mini"] - from litellm.batches.batch_utils import BatchCostUsageResult + from litellm.batches.batch_utils import BatchCostUsageResult, BatchResultFiles with patch( - "litellm.litellm_core_utils.litellm_logging._handle_completed_batch", + "litellm.litellm_core_utils.litellm_logging._handle_completed_batch_with_files", new=AsyncMock( - return_value=BatchCostUsageResult( - cost=expected_cost, - usage=expected_usage, - models=expected_models, - successful_requests=8, - failed_requests=0, + return_value=( + BatchCostUsageResult( + cost=expected_cost, + usage=expected_usage, + models=expected_models, + successful_requests=8, + failed_requests=0, + ), + BatchResultFiles(output=None, error=None), ) ), ) as mock_handle_batch: diff --git a/tests/unit/batches/test_batch_line_item_logging.py b/tests/unit/batches/test_batch_line_item_logging.py index 877b6b67854f..a8edb920bf7b 100644 --- a/tests/unit/batches/test_batch_line_item_logging.py +++ b/tests/unit/batches/test_batch_line_item_logging.py @@ -26,6 +26,7 @@ batch_line_item_claim_cache, log_batch_line_items, ) +from litellm.batches.batch_utils import BatchResultFiles from litellm.caching.caching import DualCache from litellm.caching.redis_cache import RedisCache from litellm.integrations.custom_logger import CustomLogger @@ -1240,3 +1241,49 @@ async def test_line_items_failed_fanout_does_not_delete_another_workers_claim(re assert fake._store.get(key) == "other-worker" assert len(recorder.success_events) == 0 assert len(recorder.failure_events) == 0 + + +@pytest.mark.asyncio +async def test_supplied_result_files_skip_the_output_and_error_fetch(recorder): + file_mock: Final = AsyncMock(side_effect=_file_content) + with patch( + "litellm.files.main.afile_content", file_mock + ): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + emitted: Final = await log_batch_line_items( + batch=_batch(), + custom_llm_provider="openai", + parent=_parent_logging(), + model_name="gpt-4o", + litellm_params=None, + model_info=None, + result_files=BatchResultFiles(output=OUTPUT_JSONL, error=ERROR_JSONL), + claim_cache=DualCache(), + ) + + assert emitted == 2 + assert len(recorder.success_events) == 1 + assert len(recorder.failure_events) == 1 + assert file_mock.await_count == 1 + assert file_mock.await_args.kwargs["file_id"] == "input-file-1" + + +@pytest.mark.asyncio +async def test_empty_result_files_fall_back_to_fetching_everything(recorder): + file_mock: Final = AsyncMock(side_effect=_file_content) + with patch( + "litellm.files.main.afile_content", file_mock + ): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + emitted: Final = await log_batch_line_items( + batch=_batch(), + custom_llm_provider="openai", + parent=_parent_logging(), + model_name="gpt-4o", + litellm_params=None, + model_info=None, + result_files=BatchResultFiles(output=None, error=None), + claim_cache=DualCache(), + ) + + assert emitted == 2 + fetched: Final = sorted(call.kwargs["file_id"] for call in file_mock.await_args_list) + assert fetched == ["error-file-1", "input-file-1", "output-file-1"] diff --git a/tests/unit/batches/test_batch_utils.py b/tests/unit/batches/test_batch_utils.py index b8b922f72a7b..367039cc0b47 100644 --- a/tests/unit/batches/test_batch_utils.py +++ b/tests/unit/batches/test_batch_utils.py @@ -2451,3 +2451,57 @@ async def test_flag_sends_every_vertex_row_down_the_native_path_when_a_model_is_ assert calls == [] assert (result.successful_requests, result.failed_requests) == (0, 1) + + +@pytest.mark.asyncio +async def test_handle_completed_batch_with_files_returns_bytes_and_fetches_error_once(monkeypatch): + rows = [_success_row(model="gpt-4o", usage=_usage(10, 5))] + output_bytes = _vertex_jsonl(rows) + error_bytes = _vertex_jsonl([{"custom_id": "req-bad", "error": {"message": "rejected"}}]) + + fetched = [] + + async def fake_fetch(batch, custom_llm_provider, litellm_params=None): + return output_bytes + + async def fake_afile_content(**kw): + fetched.append(kw["file_id"]) + return type("R", (), {"content": error_bytes})() + + import litellm.cost_calculator as cc + import litellm.files.main as files_main + + monkeypatch.setattr(bu, "_fetch_batch_output_file_content", fake_fetch) + monkeypatch.setattr(files_main, "afile_content", fake_afile_content) + monkeypatch.setattr(cc, "batch_cost_calculator", lambda **kw: (2.0, 1.3)) + + batch = _batch("of").model_copy(update={"error_file_id": "ef"}) + + result, files = await bu._handle_completed_batch_with_files(batch, custom_llm_provider="openai") + + assert result.cost == 3.3 + assert result.failed_requests == 1 + assert files.output == output_bytes + assert files.error == error_bytes + assert fetched == ["ef"] + + +@pytest.mark.asyncio +async def test_handle_completed_batch_with_files_no_output_returns_error_bytes(monkeypatch): + error_bytes = _vertex_jsonl([{"custom_id": "req-bad", "error": {"message": "rejected"}}]) + + async def fake_afile_content(**kw): + return type("R", (), {"content": error_bytes})() + + import litellm.files.main as files_main + + monkeypatch.setattr(files_main, "afile_content", fake_afile_content) + + batch = _batch(None).model_copy(update={"error_file_id": "ef"}) + + result, files = await bu._handle_completed_batch_with_files(batch, custom_llm_provider="openai") + + assert result.failed_requests == 1 + assert result.usage.total_tokens == 0 + assert files.output is None + assert files.error == error_bytes diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index 60b7ed32399a..9239e47ee7f8 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -1185,19 +1185,24 @@ async def test_forwards_deployment_model_and_pricing(self, monkeypatch) -> None: captured: dict[str, object] = {} - from litellm.batches.batch_utils import BatchCostUsageResult + from litellm.batches.batch_utils import BatchCostUsageResult, BatchResultFiles - async def fake_handle_completed_batch(**kwargs: object) -> BatchCostUsageResult: + async def fake_handle_completed_batch( + **kwargs: object, + ) -> tuple[BatchCostUsageResult, BatchResultFiles]: captured.update(kwargs) - return BatchCostUsageResult( - cost=1.25, - usage=Usage(prompt_tokens=1800, completion_tokens=1000, total_tokens=2800), - models=["m"], - successful_requests=1, - failed_requests=0, + return ( + BatchCostUsageResult( + cost=1.25, + usage=Usage(prompt_tokens=1800, completion_tokens=1000, total_tokens=2800), + models=["m"], + successful_requests=1, + failed_requests=0, + ), + BatchResultFiles(output=None, error=None), ) - monkeypatch.setattr(logging_module, "_handle_completed_batch", fake_handle_completed_batch) + monkeypatch.setattr(logging_module, "_handle_completed_batch_with_files", fake_handle_completed_batch) obj = LitellmLogging( model="bedrock/global.anthropic.claude-sonnet-4-6", @@ -1280,7 +1285,7 @@ async def test_non_final_batch_is_not_priced(self, monkeypatch, status, output_f from litellm.litellm_core_utils import litellm_logging as logging_module handle_completed_batch = AsyncMock() - monkeypatch.setattr(logging_module, "_handle_completed_batch", handle_completed_batch) + monkeypatch.setattr(logging_module, "_handle_completed_batch_with_files", handle_completed_batch) batch = self._batch(status, output_file_id) await self._logging_obj()._async_success_handler_body(result=batch, start_time=None, end_time=None) @@ -1290,20 +1295,23 @@ async def test_non_final_batch_is_not_priced(self, monkeypatch, status, output_f @pytest.mark.asyncio async def test_completed_batch_with_output_is_priced(self, monkeypatch) -> None: - from litellm.batches.batch_utils import BatchCostUsageResult + from litellm.batches.batch_utils import BatchCostUsageResult, BatchResultFiles from litellm.litellm_core_utils import litellm_logging as logging_module from litellm.types.utils import Usage handle_completed_batch = AsyncMock( - return_value=BatchCostUsageResult( - cost=8e-06, - usage=Usage(prompt_tokens=26, completion_tokens=9, total_tokens=35), - models=["gpt-5.6-luna"], - successful_requests=2, - failed_requests=0, + return_value=( + BatchCostUsageResult( + cost=8e-06, + usage=Usage(prompt_tokens=26, completion_tokens=9, total_tokens=35), + models=["gpt-5.6-luna"], + successful_requests=2, + failed_requests=0, + ), + BatchResultFiles(output=None, error=None), ) ) - monkeypatch.setattr(logging_module, "_handle_completed_batch", handle_completed_batch) + monkeypatch.setattr(logging_module, "_handle_completed_batch_with_files", handle_completed_batch) batch = self._batch("completed", "file-out") await self._logging_obj()._async_success_handler_body(result=batch, start_time=None, end_time=None) @@ -9081,3 +9089,154 @@ def test_signoz_dispatch_requires_an_endpoint(monkeypatch): logging_module._in_memory_loggers.clear() monkeypatch.delenv("LITELLM_OTEL_V2", raising=False) is_otel_v2_enabled.cache_clear() + + +class TestRetrieveBatchReusesFetchedResultFiles: + """The aggregate path and log_batch_line_items must share one fetch of each + result file: output and error are fetched once total (not once each per + consumer), and forwarded bytes keep line-item logging off the wire.""" + + @staticmethod + def _batch_file_bytes() -> dict[str, bytes]: + input_jsonl = json.dumps( + { + "custom_id": "a", + "method": "POST", + "url": "/v1/chat/completions", + "body": {"model": "gpt-4o", "messages": [{"role": "user", "content": "hi"}]}, + } + ).encode() + output_jsonl = json.dumps( + { + "custom_id": "a", + "response": { + "status_code": 200, + "body": { + "id": "chatcmpl-1", + "model": "gpt-4o", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "hi back"}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + }, + }, + } + ).encode() + error_jsonl = json.dumps({"custom_id": "b", "error": {"message": "rejected"}}).encode() + return {"file-in": input_jsonl, "file-out": output_jsonl, "file-err": error_jsonl} + + @staticmethod + def _logging_obj() -> LitellmLogging: + obj = LitellmLogging( + model="gpt-4o", + messages=[{"role": "user", "content": "Hey"}], + stream=False, + call_type="aretrieve_batch", + start_time=time.time(), + litellm_call_id="batch-call-reuse", + function_id="f", + ) + obj.custom_llm_provider = "openai" + obj.update_environment_variables( + litellm_params={"metadata": {}}, + optional_params={}, + custom_llm_provider="openai", + ) + return obj + + @staticmethod + def _batch(): + from litellm.types.utils import LiteLLMBatch + + return LiteLLMBatch( + id="batch_reuse_fetched_files", + completion_window="24h", + created_at=1, + endpoint="/v1/chat/completions", + input_file_id="file-in", + object="batch", + status="completed", + output_file_id="file-out", + error_file_id="file-err", + ) + + @pytest.mark.asyncio + async def test_compute_path_fetches_each_result_file_once(self, monkeypatch) -> None: + file_bytes = self._batch_file_bytes() + + async def fake_afile_content(**kwargs): + from types import SimpleNamespace + + return SimpleNamespace(content=file_bytes[kwargs["file_id"]]) + + file_mock = AsyncMock(side_effect=fake_afile_content) + monkeypatch.setattr(litellm, "store_batch_line_items_in_callbacks", True, raising=False) + monkeypatch.setattr("litellm.files.main.afile_content", file_mock) + monkeypatch.setattr("litellm.cost_calculator.batch_cost_calculator", lambda **kw: (0.01, 0.02)) + + await self._logging_obj()._async_success_handler_body( + result=self._batch(), start_time=None, end_time=None + ) + + fetched = sorted(call.kwargs["file_id"] for call in file_mock.await_args_list) + assert fetched == ["file-err", "file-in", "file-out"], ( + "each result file must be fetched exactly once across aggregate costing and line-item logging" + ) + + @pytest.mark.asyncio + async def test_explicit_kwargs_path_forwards_result_files(self, monkeypatch) -> None: + from litellm.types.utils import Usage + + file_bytes = self._batch_file_bytes() + + async def fake_afile_content(**kwargs): + from types import SimpleNamespace + + return SimpleNamespace(content=file_bytes[kwargs["file_id"]]) + + file_mock = AsyncMock(side_effect=fake_afile_content) + monkeypatch.setattr(litellm, "store_batch_line_items_in_callbacks", True, raising=False) + monkeypatch.setattr("litellm.files.main.afile_content", file_mock) + + await self._logging_obj().async_success_handler( + result=self._batch(), + batch_cost=1.5, + batch_usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15), + batch_models=["gpt-4o"], + batch_successful_requests=1, + batch_failed_requests=1, + batch_output_file_content=file_bytes["file-out"], + batch_error_file_content=file_bytes["file-err"], + ) + + assert file_mock.await_count == 1 + assert file_mock.await_args.kwargs["file_id"] == "file-in", ( + "forwarded output/error bytes must leave only the input file to fetch" + ) + + @pytest.mark.asyncio + async def test_line_item_logging_failure_leaves_aggregate_result_intact(self, monkeypatch) -> None: + from litellm.types.utils import Usage + + monkeypatch.setattr(litellm, "store_batch_line_items_in_callbacks", True, raising=False) + monkeypatch.setattr( + "litellm.batches.batch_line_item_logging.log_batch_line_items", + AsyncMock(side_effect=RuntimeError("claim backend exploded")), + ) + + batch = self._batch() + await self._logging_obj().async_success_handler( + result=batch, + batch_cost=1.5, + batch_usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15), + batch_models=["gpt-4o"], + batch_successful_requests=1, + batch_failed_requests=1, + ) + + assert batch._hidden_params["response_cost"] == 1.5 + assert batch.usage.total_tokens == 15 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 4e7effad3bf5..41470cc0d2e8 100644 --- a/tests/unit/proxy/common_utils/test_check_batch_cost.py +++ b/tests/unit/proxy/common_utils/test_check_batch_cost.py @@ -1546,10 +1546,10 @@ async def test_error_file_failures_add_to_failed_request_count( Logging, "async_success_handler", new_callable=AsyncMock ) as success_handler, ): - provider.get("https://api.openai.com/v1/files/file-output-123/content").mock( + output_route = provider.get("https://api.openai.com/v1/files/file-output-123/content").mock( return_value=httpx.Response(200, content=f"{succeeded_line}\n{rejected_line}\n".encode()) ) - provider.get("https://api.openai.com/v1/files/file-error-456/content").mock( + error_route = provider.get("https://api.openai.com/v1/files/file-error-456/content").mock( return_value=httpx.Response(200, content=f"{error_file_lines}\n\n".encode()) ) await check_batch_cost_instance.check_batch_cost() @@ -1558,6 +1558,10 @@ async def test_error_file_failures_add_to_failed_request_count( assert len(spend_log_calls) == 1 handler_kwargs = spend_log_calls[0] assert handler_kwargs["batch_successful_requests"] == 1 + assert output_route.call_count == 1 + assert error_route.call_count == 1 + assert handler_kwargs["batch_output_file_content"] == f"{succeeded_line}\n{rejected_line}\n".encode() + assert handler_kwargs["batch_error_file_content"] == f"{error_file_lines}\n\n".encode() assert handler_kwargs["batch_failed_requests"] == 3, ( "2 error-file lines must add to the output file's 1 rejected request" ) From 1388dbca5e0116aa93820c4194a099de30dff9ad Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 30 Sep 2026 23:38:37 +0000 Subject: [PATCH 37/89] fix(batches): release line-item claim when nothing was emitted Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/batches/batch_line_item_logging.py | 1 + .../batches/test_batch_line_item_logging.py | 86 +++++++++++++++++++ .../test_litellm_logging.py | 9 ++ 3 files changed, 96 insertions(+) diff --git a/litellm/batches/batch_line_item_logging.py b/litellm/batches/batch_line_item_logging.py index db9f0c9b6b18..580100b736fe 100644 --- a/litellm/batches/batch_line_item_logging.py +++ b/litellm/batches/batch_line_item_logging.py @@ -511,6 +511,7 @@ async def log_batch_line_items( "batch line item logging failed for batch_id=%s; aggregate logging unaffected", batch.id, ) + finally: if emitted == 0: await _release_line_item_claim(claim_cache, claim_key, token) return emitted diff --git a/tests/unit/batches/test_batch_line_item_logging.py b/tests/unit/batches/test_batch_line_item_logging.py index a8edb920bf7b..e338cad5d214 100644 --- a/tests/unit/batches/test_batch_line_item_logging.py +++ b/tests/unit/batches/test_batch_line_item_logging.py @@ -1287,3 +1287,89 @@ async def test_empty_result_files_fall_back_to_fetching_everything(recorder): assert emitted == 2 fetched: Final = sorted(call.kwargs["file_id"] for call in file_mock.await_args_list) assert fetched == ["error-file-1", "input-file-1", "output-file-1"] + + +@pytest.mark.asyncio +async def test_native_vertex_skip_releases_the_claim(recorder): + file_mock: Final = AsyncMock( + side_effect=_scoped_file_content( + { + "input-vtx-rel": b"", + "output-vtx-rel": VERTEX_NATIVE_OUTPUT_JSONL, + } + ) + ) + claim_cache: Final = DualCache() + batch: Final = _provider_batch("batch_vtx_rel", "input-vtx-rel", "output-vtx-rel") + parent: Final = _parent_logging(custom_llm_provider="vertex_ai") + with patch( + "litellm.files.main.afile_content", file_mock + ): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + first: Final = await log_batch_line_items( + batch=batch, + custom_llm_provider="vertex_ai", + parent=parent, + model_name="vertex_ai/gemini-2.5-flash", + litellm_params=None, + model_info=None, + claim_cache=claim_cache, + ) + second: Final = await log_batch_line_items( + batch=batch, + custom_llm_provider="vertex_ai", + parent=parent, + model_name="vertex_ai/gemini-2.5-flash", + litellm_params=None, + model_info=None, + claim_cache=claim_cache, + ) + + assert first == 0 + assert second == 0 + assert file_mock.await_count == 4, ( + "a released claim must let the next retrieve reach the fetch step again; already_claimed would return early at 2 fetches" + ) + + +@pytest.mark.asyncio +async def test_fanout_emitting_nothing_releases_the_claim(recorder): + file_mock: Final = AsyncMock( + side_effect=_scoped_file_content( + { + "input-empty": b"", + "output-empty": b"", + "error-empty": b"", + } + ) + ) + claim_cache: Final = DualCache() + batch: Final = _provider_batch("batch_empty", "input-empty", "output-empty") + batch = batch.model_copy(update={"error_file_id": "error-empty"}) + parent: Final = _parent_logging() + with patch( + "litellm.files.main.afile_content", file_mock + ): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + first: Final = await log_batch_line_items( + batch=batch, + custom_llm_provider="openai", + parent=parent, + model_name="gpt-4o", + litellm_params=None, + model_info=None, + claim_cache=claim_cache, + ) + second: Final = await log_batch_line_items( + batch=batch, + custom_llm_provider="openai", + parent=parent, + model_name="gpt-4o", + litellm_params=None, + model_info=None, + claim_cache=claim_cache, + ) + + assert first == 0 + assert second == 0 + assert file_mock.await_count == 6, ( + "a completed fan-out that emits nothing must release the claim so the next retrieve retries" + ) diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index 9239e47ee7f8..1fd8a6818731 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -9166,6 +9166,9 @@ def _batch(): @pytest.mark.asyncio async def test_compute_path_fetches_each_result_file_once(self, monkeypatch) -> None: + from litellm.batches.batch_line_item_logging import batch_line_item_claim_cache + + batch_line_item_claim_cache.in_memory_cache.flush_cache() file_bytes = self._batch_file_bytes() async def fake_afile_content(**kwargs): @@ -9189,6 +9192,9 @@ async def fake_afile_content(**kwargs): @pytest.mark.asyncio async def test_explicit_kwargs_path_forwards_result_files(self, monkeypatch) -> None: + from litellm.batches.batch_line_item_logging import batch_line_item_claim_cache + + batch_line_item_claim_cache.in_memory_cache.flush_cache() from litellm.types.utils import Usage file_bytes = self._batch_file_bytes() @@ -9220,6 +9226,9 @@ async def fake_afile_content(**kwargs): @pytest.mark.asyncio async def test_line_item_logging_failure_leaves_aggregate_result_intact(self, monkeypatch) -> None: + from litellm.batches.batch_line_item_logging import batch_line_item_claim_cache + + batch_line_item_claim_cache.in_memory_cache.flush_cache() from litellm.types.utils import Usage monkeypatch.setattr(litellm, "store_batch_line_items_in_callbacks", True, raising=False) From 64c514cda71392452360f252af511b647475b6d5 Mon Sep 17 00:00:00 2001 From: yucheng Date: Thu, 1 Oct 2026 17:55:25 +0000 Subject: [PATCH 38/89] fix(batches): type batch line-item callback sites to satisfy basedpyright budget Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../litellm_enterprise/proxy/common_utils/check_batch_cost.py | 2 +- litellm/litellm_core_utils/litellm_logging.py | 4 ++-- litellm/proxy/hooks/dynamic_rate_limiter_v3.py | 4 ++-- litellm/proxy/hooks/model_max_budget_limiter.py | 4 ++-- litellm/proxy/hooks/parallel_request_limiter.py | 4 ++-- litellm/proxy/hooks/parallel_request_limiter_v3.py | 4 ++-- litellm/proxy/proxy_server.py | 2 +- litellm/router_strategy/budget_limiter.py | 2 +- 8 files changed, 13 insertions(+), 13 deletions(-) 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 cb41eba32eba..d6ef044ef4cb 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -918,7 +918,7 @@ async def _track_completed_batch_cost( batch_failed_requests=batch_result.failed_requests, batch_prompt_cost=batch_result.prompt_cost, batch_completion_cost=batch_result.completion_cost, - batch_output_file_content=content_bytes, + batch_output_file_content=content_bytes if isinstance(content_bytes, bytes) else None, batch_error_file_content=error_file_content, ) except Exception: diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 00c940a8116f..f2694c6c2989 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -3290,7 +3290,7 @@ async def _async_success_handler_body( custom_llm_provider=self.custom_llm_provider, parent=self, model_name=self.get_deployment_model_for_cost(), - litellm_params=self.litellm_params, + litellm_params=cast("dict[str, object] | None", self.litellm_params), model_info=self.get_router_deployment_model_info(), result_files=result_files, ) @@ -6437,7 +6437,7 @@ def _extract_response_obj_and_hidden_params( if original_exception is not None and hidden_params is None: exception_hidden_params: Final = getattr(original_exception, "_hidden_params", None) if isinstance(exception_hidden_params, dict) and exception_hidden_params: - hidden_params = dict(exception_hidden_params) # mutable-ok: hidden_params downstream expects a plain dict + hidden_params = dict(cast("dict[str, object]", exception_hidden_params)) # cast-ok: isinstance(dict) leaves the value type unparameterized; mutable-ok: hidden_params downstream expects a plain dict elif (response_headers := _get_response_headers(original_exception)) is not None: hidden_params = dict( StandardLoggingHiddenParams( diff --git a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py index 6777e0c88f38..09cf8af9681a 100644 --- a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py +++ b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py @@ -5,7 +5,7 @@ import os from collections.abc import Callable from datetime import datetime -from typing import TYPE_CHECKING, Final, Literal +from typing import TYPE_CHECKING, Any, Final, Literal from fastapi import HTTPException @@ -686,7 +686,7 @@ async def async_post_call_success_hook(self, data: dict, user_api_key_dict: User verbose_proxy_logger.exception("Error in dynamic rate limiter v3 post-call hook: %s", e) return response - async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + async def async_log_success_event(self, kwargs: dict[str, Any], response_obj, start_time, end_time): """ Update token usage for priority-based rate limiting after successful API calls. diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index cd28a5771d62..cb3413564521 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/litellm/proxy/hooks/model_max_budget_limiter.py @@ -4,7 +4,7 @@ from collections.abc import Iterable, Mapping, Sequence from dataclasses import dataclass from types import MappingProxyType -from typing import Final +from typing import Any, Final from openai.types import Batch @@ -485,7 +485,7 @@ async def async_filter_deployments( ) -> list[dict]: return healthy_deployments - async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + async def async_log_success_event(self, kwargs: dict[str, Any], response_obj, start_time, end_time): """ Track spend for virtual key + model in DualCache diff --git a/litellm/proxy/hooks/parallel_request_limiter.py b/litellm/proxy/hooks/parallel_request_limiter.py index b04a4823dd59..a628f055f805 100644 --- a/litellm/proxy/hooks/parallel_request_limiter.py +++ b/litellm/proxy/hooks/parallel_request_limiter.py @@ -492,7 +492,7 @@ async def async_pre_call_hook( ) # don't block execution for cache updates ) - async def async_log_success_event(self, kwargs, response_obj: object, start_time, end_time): + async def async_log_success_event(self, kwargs: dict[str, Any], response_obj: object, start_time, end_time): if is_batch_line_item_event(kwargs): return from litellm.proxy.common_utils.callback_utils import ( @@ -699,7 +699,7 @@ async def async_log_success_event(self, kwargs, response_obj: object, start_time except Exception as e: self.print_verbose(e) - async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + async def async_log_failure_event(self, kwargs: dict[str, Any], response_obj, start_time, end_time): if is_batch_line_item_event(kwargs): return try: diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index e47dcacd91e2..603d3be53c94 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -5008,7 +5008,7 @@ def _build_success_event_pipeline_operations( return pipeline_operations - async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + async def async_log_success_event(self, kwargs: dict[str, Any], response_obj, start_time, end_time): """ Update TPM usage on successful API calls by incrementing counters using pipeline """ @@ -5126,7 +5126,7 @@ def _recovered_partial_usage_tokens(self, source: Mapping[str, object]) -> tuple completion_tokens, ) - async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + async def async_log_failure_event(self, kwargs: dict[str, Any], response_obj, start_time, end_time): """ On failure: decrement max_parallel_requests and refund the upfront TPM reservation only against the scopes the reservation actually diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index e5b2550328ae..d2ff21ca3ace 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -6829,7 +6829,7 @@ async def load_config(self, router: litellm.Router | None, config_file_path: str health_check_concurrency = general_settings.get("health_check_concurrency", None) health_check_details = general_settings.get("health_check_details", True) ### BATCH LINE ITEM CALLBACKS ### - _store_batch_line_items: Final = general_settings.get("store_batch_line_items_in_callbacks") + _store_batch_line_items: Final[object] = general_settings.get("store_batch_line_items_in_callbacks") if _store_batch_line_items is not None: if isinstance(_store_batch_line_items, str): litellm.store_batch_line_items_in_callbacks = _store_batch_line_items.lower() == "true" diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index cb58afef0447..7e00793aea07 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/litellm/router_strategy/budget_limiter.py @@ -491,7 +491,7 @@ async def _write_queued_increment_operations(self, redis_cache: RedisCache) -> b await self._clear_detached_increment_operations() return True - async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + async def async_log_success_event(self, kwargs: dict[str, Any], response_obj, start_time, end_time): """Original method now uses helper functions""" verbose_router_logger.debug("in RouterBudgetLimiting.async_log_success_event") if is_batch_line_item_event(kwargs): From aebde2a31f87cb38522624a39f31055fc00f9d5e Mon Sep 17 00:00:00 2001 From: yucheng Date: Thu, 1 Oct 2026 18:18:19 +0000 Subject: [PATCH 39/89] style: wrap batch hidden-params cast for ruff format Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/litellm_core_utils/litellm_logging.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index f2694c6c2989..59a0cf83bb59 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -6437,7 +6437,9 @@ def _extract_response_obj_and_hidden_params( if original_exception is not None and hidden_params is None: exception_hidden_params: Final = getattr(original_exception, "_hidden_params", None) if isinstance(exception_hidden_params, dict) and exception_hidden_params: - hidden_params = dict(cast("dict[str, object]", exception_hidden_params)) # cast-ok: isinstance(dict) leaves the value type unparameterized; mutable-ok: hidden_params downstream expects a plain dict + # cast-ok: isinstance(dict) leaves the value type unparameterized + # mutable-ok: hidden_params downstream expects a plain dict + hidden_params = dict(cast("dict[str, object]", exception_hidden_params)) elif (response_headers := _get_response_headers(original_exception)) is not None: hidden_params = dict( StandardLoggingHiddenParams( From eac1ed00c31b9a39e89970697029724009199433 Mon Sep 17 00:00:00 2001 From: yucheng Date: Thu, 1 Oct 2026 18:48:03 +0000 Subject: [PATCH 40/89] fix(batches): reason suppressions for batch line-item lint deltas Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/batches/batch_line_item_logging.py | 4 +++- litellm/litellm_core_utils/litellm_logging.py | 13 ++++++++----- 2 files changed, 11 insertions(+), 6 deletions(-) diff --git a/litellm/batches/batch_line_item_logging.py b/litellm/batches/batch_line_item_logging.py index 580100b736fe..3fb3dd1cb89f 100644 --- a/litellm/batches/batch_line_item_logging.py +++ b/litellm/batches/batch_line_item_logging.py @@ -108,7 +108,9 @@ async def _release_line_item_claim(claim_cache: DualCache, claim_key: str, token def _json_fallback(value: object) -> dict[str, object] | str: mapping: Final = _as_object_mapping(value) - return dict(mapping) if mapping is not None else str(value) + if mapping is None: + return str(value) + return dict(mapping) # mutable-ok: json.dumps default hook must return a plain dict class _BatchLineFailure(Exception): diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 7af16f712f74..1a0a3c145e4f 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -3285,13 +3285,16 @@ async def _async_success_handler_body( if litellm.store_batch_line_items_in_callbacks and (has_explicit_batch_data or should_compute_batch_data): from litellm.batches.batch_line_item_logging import log_batch_line_items + _line_item_litellm_params: Final = cast( # cast-ok: shared attribute is an untyped dict + "dict[str, object] | None", self.litellm_params + ) try: await log_batch_line_items( batch=result, custom_llm_provider=self.custom_llm_provider, parent=self, model_name=self.get_deployment_model_for_cost(), - litellm_params=cast("dict[str, object] | None", self.litellm_params), + litellm_params=_line_item_litellm_params, model_info=self.get_router_deployment_model_info(), result_files=result_files, ) @@ -6440,11 +6443,11 @@ def _extract_response_obj_and_hidden_params( response_obj = {} if original_exception is not None and hidden_params is None: - exception_hidden_params: Final = getattr(original_exception, "_hidden_params", None) + exception_hidden_params: Final[Mapping[str, object] | None] = getattr( + original_exception, "_hidden_params", None + ) if isinstance(exception_hidden_params, dict) and exception_hidden_params: - # cast-ok: isinstance(dict) leaves the value type unparameterized - # mutable-ok: hidden_params downstream expects a plain dict - hidden_params = dict(cast("dict[str, object]", exception_hidden_params)) + hidden_params = dict(exception_hidden_params) # mutable-ok: hidden_params downstream expects a plain dict elif (response_headers := _get_response_headers(original_exception)) is not None: hidden_params = dict( StandardLoggingHiddenParams( From 60a47e2900ee972a34fbc736870a91128f9f356c Mon Sep 17 00:00:00 2001 From: yucheng Date: Thu, 1 Oct 2026 19:29:09 +0000 Subject: [PATCH 41/89] chore: drop stale mutable-ok suppressions after LIT002 removal Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/batches/batch_line_item_logging.py | 18 +++++++----------- litellm/litellm_core_utils/litellm_logging.py | 2 +- litellm/llms/anthropic/chat/transformation.py | 2 +- litellm/llms/bedrock/batches/transformation.py | 14 ++++++-------- 4 files changed, 15 insertions(+), 21 deletions(-) diff --git a/litellm/batches/batch_line_item_logging.py b/litellm/batches/batch_line_item_logging.py index 3fb3dd1cb89f..fe6a027da979 100644 --- a/litellm/batches/batch_line_item_logging.py +++ b/litellm/batches/batch_line_item_logging.py @@ -110,7 +110,7 @@ def _json_fallback(value: object) -> dict[str, object] | str: mapping: Final = _as_object_mapping(value) if mapping is None: return str(value) - return dict(mapping) # mutable-ok: json.dumps default hook must return a plain dict + return dict(mapping) class _BatchLineFailure(Exception): @@ -285,7 +285,7 @@ def _new_child_logging( dynamic_async_success_callbacks=parent.dynamic_async_success_callbacks, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # same as above dynamic_failure_callbacks=parent.dynamic_failure_callbacks, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # same as above dynamic_async_failure_callbacks=parent.dynamic_async_failure_callbacks, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # same as above - kwargs={"litellm_session_id": parent.litellm_session_id}, # mutable-ok: kwargs takes a plain dict + kwargs={"litellm_session_id": parent.litellm_session_id}, ) @@ -295,7 +295,7 @@ def _line_hidden_params( status_code: int | None, response_cost: float | None = None, ) -> dict[str, object]: # mutable-ok: response objects declare _hidden_params as a plain dict - return { # mutable-ok: same contract + return { "batch_id": batch.id, "batch_custom_id": custom_id, "batch_line_status_code": status_code, @@ -306,14 +306,12 @@ def _line_hidden_params( def _optional_params_for_body( request_body: Mapping[str, object], ) -> dict[str, object]: # mutable-ok: update_environment_variables takes a plain dict - return { # mutable-ok: same contract - key: value for key, value in request_body.items() if key not in ("model", "messages", "input") - } + return {key: value for key, value in request_body.items() if key not in ("model", "messages", "input")} def _metadata_copy(params: Mapping[str, object]) -> dict[str, object]: # mutable-ok: dict out for litellm_params metadata: Final = _as_object_mapping(params.get("metadata")) or _EMPTY_BODY - return {**metadata} # mutable-ok: plain-dict copy + return {**metadata} async def _emit_line_event( @@ -355,7 +353,7 @@ async def _emit_line_event( start_time=start_time, ) child.update_environment_variables( # pyright: ignore[reportUnknownMemberType] # Logging.update_environment_variables is untyped upstream - litellm_params={ # mutable-ok: update_environment_variables takes a plain dict + litellm_params={ **parent_params, "batch_parent_id": batch.id, "metadata": _metadata_copy(parent_params), @@ -466,9 +464,7 @@ async def log_batch_line_items( internal_credentials: Final = parent._litellm_internal_model_credentials # pyright: ignore[reportPrivateUsage] # declared transport attribute on Logging internal_mapping: Final = _as_object_mapping(internal_credentials) fetch_params: Final[dict[str, object] | None] = ( # mutable-ok: file fetcher requires a plain dict - dict(internal_mapping) # mutable-ok: the file fetcher reads credential kwargs off a plain dict - if internal_mapping is not None - else litellm_params + dict(internal_mapping) if internal_mapping is not None else litellm_params ) input_file_content: Final = await _fetch_managed_file_or_empty(batch.input_file_id, line_provider, fetch_params) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 2e1b9cd58250..591383d2b81c 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -6445,7 +6445,7 @@ def _extract_response_obj_and_hidden_params( original_exception, "_hidden_params", None ) if isinstance(exception_hidden_params, dict) and exception_hidden_params: - hidden_params = dict(exception_hidden_params) # mutable-ok: hidden_params downstream expects a plain dict + hidden_params = dict(exception_hidden_params) elif (response_headers := _get_response_headers(original_exception)) is not None: hidden_params = dict( StandardLoggingHiddenParams( diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index 2039e3fb9548..b9b22fa0da4a 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -2816,7 +2816,7 @@ def anthropic_message_to_model_response(result: Mapping[str, object], speed: str completion_response=pydantic_result.model_dump(), raw_response=httpx.Response( status_code=200, - headers={}, # mutable-ok: httpx.Response wants a plain dict of headers + headers={}, ), model_response=ModelResponse(id=result_id if isinstance(result_id, str) and result_id else None), json_mode=None, diff --git a/litellm/llms/bedrock/batches/transformation.py b/litellm/llms/bedrock/batches/transformation.py index b4c16f010cc2..32fdf7e89973 100644 --- a/litellm/llms/bedrock/batches/transformation.py +++ b/litellm/llms/bedrock/batches/transformation.py @@ -106,13 +106,11 @@ def bedrock_batch_line_to_response( embedding: Final = model_output.get("embedding") return EmbeddingResponse( model=model, - data=[ # mutable-ok: EmbeddingResponse takes a plain data list - { # mutable-ok: plain row dict for EmbeddingResponse.data + data=[ + { "object": "embedding", "index": 0, - "embedding": embedding - if isinstance(embedding, list) - else [], # mutable-ok: empty fallback for the row + "embedding": embedding if isinstance(embedding, list) else [], } ], usage=titan_embedding_usage_from_batch_output(model_output), @@ -122,14 +120,14 @@ def bedrock_batch_line_to_response( return AmazonConverseConfig()._transform_response( # pyright: ignore[reportPrivateUsage] # same reconstruction the converse chat path performs on the live response model=model, - response=Response(200, json=dict(model_output)), # mutable-ok: httpx.Response json= wants a plain dict + response=Response(200, json=dict(model_output)), model_response=ModelResponse(), stream=False, logging_obj=None, - optional_params={}, # mutable-ok: the converse transform signature takes a plain dict + optional_params={}, api_key=None, data="", - messages=[], # mutable-ok: the converse transform signature takes a plain list + messages=[], encoding=None, ) if "content" in model_output: From a7f197862ee1467e14304d300ddb68a6b5b2c36e Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Sun, 4 Oct 2026 03:07:02 -0700 Subject: [PATCH 42/89] fix(batches): skip batch line-item events in ClickHouse spend sink and router quota counters (#44486) * fix(batches): skip batch line-item events in ClickHouse spend sink and router quota counters Batch line-item callbacks (call_type=acompletion + litellm_params.batch_parent_id) bypassed the call-type-only batch guards: - ClickHouseSpendLogger wrote one spend row per line on top of the aggregate aretrieve_batch row, billing every batch twice with the flag on. - Router deployment callbacks and the routing-strategy usage handlers charged the historical batch's tokens/rpm to the live TPM/RPM counters, so after Router.aretrieve_batch the next chat raised RouterRateLimitError. Skip these events wherever the aggregate batch retrieve is already skipped. * fix(logging): keep internal litellm_params authoritative in model_call_details A request body key named litellm_params lands in optional_params, and update_environment_variables merged **optional_params after the internal entry, so a caller could overwrite the logging kwargs' litellm_params and forge markers like batch_parent_id. Live-reproduced: a Gemini chat whose body carried {"litellm_params": {"batch_parent_id": ...}} returned 200 with no ClickHouse spend row. Re-assert the internal litellm_params after the merge so every sink (ClickHouse, Postgres, router quota counters) reads the same trusted value, and assert the aggregate aretrieve_batch row still writes in the ClickHouse line-item test (review finding). --- .../clickhouse/clickhouse_spend_logger.py | 5 + litellm/litellm_core_utils/litellm_logging.py | 4 + litellm/router.py | 7 +- litellm/router_strategy/least_busy.py | 3 +- litellm/router_strategy/lowest_cost.py | 5 +- litellm/router_strategy/lowest_latency.py | 12 +- litellm/router_strategy/lowest_tpm_rpm.py | 5 +- litellm/router_strategy/lowest_tpm_rpm_v2.py | 6 +- .../test_clickhouse_spend_logger.py | 59 +++++++++ .../test_litellm_logging.py | 16 +++ tests/unit/router_strategy/test_least_busy.py | 25 ++++ .../router_strategy/test_lowest_tpm_rpm.py | 37 ++++++ tests/unit/test_router/test_router.py | 122 ++++++++++++++++++ 13 files changed, 291 insertions(+), 15 deletions(-) diff --git a/litellm/integrations/clickhouse/clickhouse_spend_logger.py b/litellm/integrations/clickhouse/clickhouse_spend_logger.py index c1401e111bb0..431512ed33ae 100644 --- a/litellm/integrations/clickhouse/clickhouse_spend_logger.py +++ b/litellm/integrations/clickhouse/clickhouse_spend_logger.py @@ -16,6 +16,7 @@ from litellm.integrations.clickhouse.clickhouse_batch_logger import ClickHouseBatchLogger from litellm.integrations.clickhouse.context import is_lens_analysis from litellm.integrations.clickhouse.schema import SPEND_LOGS_TABLE +from litellm.litellm_core_utils.core_helpers import is_batch_line_item_event from litellm.tracing.types import SpendLogRecord from litellm.types.utils import StandardLoggingPayload @@ -154,6 +155,10 @@ class ClickHouseSpendLogger(ClickHouseBatchLogger): table = SPEND_LOGS_TABLE async def async_log_success_event(self, kwargs, response_obj, start_time, end_time) -> None: + # Batch line items are billed by the aggregate aretrieve_batch row; a per-line + # spend row here would bill the batch twice. + if is_batch_line_item_event(kwargs): + return self._log(kwargs) async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time) -> None: diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 591383d2b81c..32fd00c4f4d2 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -947,6 +947,10 @@ def update_environment_variables( **additional_params, } ) + # Provider params / additional kwargs are caller-influenced; they must not + # overwrite the internal litellm_params, or a request body carrying + # `litellm_params` could forge logging markers (e.g. batch_parent_id). + self.model_call_details["litellm_params"] = self.litellm_params ## check if stream options is set ## - used by CustomStreamWrapper for easy instrumentation if "stream_options" in additional_params: diff --git a/litellm/router.py b/litellm/router.py index 24554e615167..4cc0bdcbe8c6 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -78,6 +78,7 @@ get_litellm_metadata_from_kwargs, get_metadata_variable_name_from_kwargs, get_or_create_metadata_bucket, + is_batch_line_item_event, ) from litellm.litellm_core_utils.coroutine_checker import coroutine_checker from litellm.litellm_core_utils.credential_accessor import CredentialAccessor @@ -8050,7 +8051,7 @@ async def deployment_callback_on_success( # WS session wrappers fire with result=None; per-turn costs tracked by inner calls. if kwargs.get("call_type") in ("_aresponses_websocket", "_arealtime"): return - if is_batch_retrieve_call_type(kwargs.get("call_type")): + if is_batch_retrieve_call_type(kwargs.get("call_type")) or is_batch_line_item_event(kwargs): return standard_logging_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object", None) if standard_logging_object is None: @@ -8198,7 +8199,7 @@ def sync_deployment_callback_on_success( - key: str - The key used to increment the cache - None: if no key is found """ - if is_batch_retrieve_call_type(kwargs.get("call_type")): + if is_batch_retrieve_call_type(kwargs.get("call_type")) or is_batch_line_item_event(kwargs): return None id = None if kwargs["litellm_params"].get("metadata") is None: @@ -8328,7 +8329,7 @@ async def async_deployment_callback_on_failure( """ Update RPM usage for a deployment """ - if is_batch_retrieve_call_type(kwargs.get("call_type")): + if is_batch_retrieve_call_type(kwargs.get("call_type")) or is_batch_line_item_event(kwargs): return deployment_name: Final = kwargs["litellm_params"]["metadata"].get( "deployment", None diff --git a/litellm/router_strategy/least_busy.py b/litellm/router_strategy/least_busy.py index 9ab670e4b954..12afa512e103 100644 --- a/litellm/router_strategy/least_busy.py +++ b/litellm/router_strategy/least_busy.py @@ -9,6 +9,7 @@ from litellm.caching.caching import DualCache from litellm.caching.redis_cache import log_redis_failure from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.core_helpers import is_batch_line_item_event from litellm.router_utils.batch_utils import is_batch_retrieve_call_type IN_FLIGHT_COUNT_TTL_SECONDS: Final = 60 * 60 @@ -49,7 +50,7 @@ def _request_count_key(model_group: str, deployment_id: str) -> str: def _deployment_ref(kwargs: Mapping[str, object]) -> tuple[str, str] | None: - if is_batch_retrieve_call_type(kwargs.get("call_type")): + if is_batch_retrieve_call_type(kwargs.get("call_type")) or is_batch_line_item_event(kwargs): return None try: call: Final = _CALL_KWARGS.validate_python(kwargs) diff --git a/litellm/router_strategy/lowest_cost.py b/litellm/router_strategy/lowest_cost.py index 22c321c65fbb..d47584ebbfc0 100644 --- a/litellm/router_strategy/lowest_cost.py +++ b/litellm/router_strategy/lowest_cost.py @@ -8,6 +8,7 @@ from litellm._logging import verbose_router_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.core_helpers import is_batch_line_item_event from litellm.router_utils.batch_utils import is_batch_retrieve_call_type @@ -20,7 +21,7 @@ def __init__(self, router_cache: DualCache, routing_args: dict = {}): self.router_cache = router_cache def log_success_event(self, kwargs, response_obj, start_time, end_time): - if is_batch_retrieve_call_type(kwargs.get("call_type")): + if is_batch_retrieve_call_type(kwargs.get("call_type")) or is_batch_line_item_event(kwargs): return try: """ @@ -95,7 +96,7 @@ def log_success_event(self, kwargs, response_obj, start_time, end_time): ) async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): - if is_batch_retrieve_call_type(kwargs.get("call_type")): + if is_batch_retrieve_call_type(kwargs.get("call_type")) or is_batch_line_item_event(kwargs): return try: """ diff --git a/litellm/router_strategy/lowest_latency.py b/litellm/router_strategy/lowest_latency.py index 66c8227195d7..a03fcde74e5c 100644 --- a/litellm/router_strategy/lowest_latency.py +++ b/litellm/router_strategy/lowest_latency.py @@ -12,7 +12,11 @@ from litellm import ModelResponse, token_counter, verbose_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger -from litellm.litellm_core_utils.core_helpers import _get_parent_otel_span_from_kwargs, safe_divide_seconds +from litellm.litellm_core_utils.core_helpers import ( + _get_parent_otel_span_from_kwargs, + is_batch_line_item_event, + safe_divide_seconds, +) from litellm.router_utils.batch_utils import is_batch_retrieve_call_type from litellm.types.utils import LiteLLMPydanticObjectBase @@ -59,7 +63,7 @@ def __init__(self, router_cache: DualCache, routing_args: dict = {}): self.routing_args = RoutingArgs(**routing_args) def log_success_event(self, kwargs, response_obj, start_time, end_time): - if is_batch_retrieve_call_type(kwargs.get("call_type")): + if is_batch_retrieve_call_type(kwargs.get("call_type")) or is_batch_line_item_event(kwargs): return try: """ @@ -185,7 +189,7 @@ async def async_log_failure_event(self, kwargs, response_obj, start_time, end_ti """ Check if Timeout Error, if timeout set deployment latency -> 100 """ - if is_batch_retrieve_call_type(kwargs.get("call_type")): + if is_batch_retrieve_call_type(kwargs.get("call_type")) or is_batch_line_item_event(kwargs): return try: metadata_field: Final = self._select_metadata_field(kwargs) @@ -241,7 +245,7 @@ async def async_log_failure_event(self, kwargs, response_obj, start_time, end_ti ) async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): - if is_batch_retrieve_call_type(kwargs.get("call_type")): + if is_batch_retrieve_call_type(kwargs.get("call_type")) or is_batch_line_item_event(kwargs): return try: """ diff --git a/litellm/router_strategy/lowest_tpm_rpm.py b/litellm/router_strategy/lowest_tpm_rpm.py index d4abf1f8f708..3575db0aa53d 100644 --- a/litellm/router_strategy/lowest_tpm_rpm.py +++ b/litellm/router_strategy/lowest_tpm_rpm.py @@ -8,6 +8,7 @@ from litellm._logging import verbose_router_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.core_helpers import is_batch_line_item_event from litellm.router_utils.batch_utils import is_batch_retrieve_call_type from litellm.types.utils import LiteLLMPydanticObjectBase from litellm.utils import print_verbose @@ -28,7 +29,7 @@ def __init__(self, router_cache: DualCache, routing_args: dict = {}): self.routing_args = RoutingArgs(**routing_args) def log_success_event(self, kwargs, response_obj, start_time, end_time): - if is_batch_retrieve_call_type(kwargs.get("call_type")): + if is_batch_retrieve_call_type(kwargs.get("call_type")) or is_batch_line_item_event(kwargs): return try: """ @@ -82,7 +83,7 @@ def log_success_event(self, kwargs, response_obj, start_time, end_time): verbose_router_logger.debug(traceback.format_exc()) async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): - if is_batch_retrieve_call_type(kwargs.get("call_type")): + if is_batch_retrieve_call_type(kwargs.get("call_type")) or is_batch_line_item_event(kwargs): return try: """ diff --git a/litellm/router_strategy/lowest_tpm_rpm_v2.py b/litellm/router_strategy/lowest_tpm_rpm_v2.py index 25564a80e0a2..02cea3209359 100644 --- a/litellm/router_strategy/lowest_tpm_rpm_v2.py +++ b/litellm/router_strategy/lowest_tpm_rpm_v2.py @@ -14,7 +14,7 @@ from litellm._logging import verbose_logger, verbose_router_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger -from litellm.litellm_core_utils.core_helpers import _get_parent_otel_span_from_kwargs +from litellm.litellm_core_utils.core_helpers import _get_parent_otel_span_from_kwargs, is_batch_line_item_event from litellm.router_utils.batch_utils import is_batch_retrieve_call_type from litellm.types.router import RouterErrors from litellm.types.utils import LiteLLMPydanticObjectBase, StandardLoggingPayload @@ -250,7 +250,7 @@ async def async_pre_call_check(self, deployment: dict, parent_otel_span: Span | return deployment # don't fail calls if eg. redis fails to connect def log_success_event(self, kwargs, response_obj, start_time, end_time): - if is_batch_retrieve_call_type(kwargs.get("call_type")): + if is_batch_retrieve_call_type(kwargs.get("call_type")) or is_batch_line_item_event(kwargs): return try: """ @@ -292,7 +292,7 @@ def log_success_event(self, kwargs, response_obj, start_time, end_time): ) async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): - if is_batch_retrieve_call_type(kwargs.get("call_type")): + if is_batch_retrieve_call_type(kwargs.get("call_type")) or is_batch_line_item_event(kwargs): return try: """ diff --git a/tests/test_litellm/integrations/clickhouse/test_clickhouse_spend_logger.py b/tests/test_litellm/integrations/clickhouse/test_clickhouse_spend_logger.py index b183bf84ea4a..a2eed7ce3e99 100644 --- a/tests/test_litellm/integrations/clickhouse/test_clickhouse_spend_logger.py +++ b/tests/test_litellm/integrations/clickhouse/test_clickhouse_spend_logger.py @@ -330,3 +330,62 @@ async def test_trace_ingest_and_invalid_payload_do_not_write_spend(): assert logger.log_queue == [] storage.ensure_schema.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_batch_line_item_success_event_does_not_write_spend(monkeypatch: pytest.MonkeyPatch): + # A batch line item carries call_type=acompletion + litellm_params.batch_parent_id. + # The aggregate aretrieve_batch row already bills the batch, so a per-line spend row + # would bill it twice. + monkeypatch.setattr(litellm, "store_batch_line_items_in_callbacks", True, raising=False) + storage = MagicMock() + logger = ClickHouseSpendLogger(storage=storage) + now = datetime.now(timezone.utc) + + await logger.async_log_success_event( + { + "standard_logging_object": _minimal_payload("chatcmpl-line-1", status="success", cost=0.25), + "call_type": "acompletion", + "litellm_params": {"batch_parent_id": "batch-1"}, + }, + None, + now, + now, + ) + + assert logger.log_queue == [] + + await logger.async_log_success_event( + { + "standard_logging_object": _minimal_payload("chatcmpl-live-1", status="success", cost=0.25), + "call_type": "acompletion", + "litellm_params": {}, + }, + None, + now, + now, + ) + + # The aggregate aretrieve_batch event bills the batch; it must still write one + # row with the batch's full cost even while line items are skipped. + await logger.async_log_success_event( + { + "standard_logging_object": { + **_minimal_payload("batch-1", status="success", cost=0.0001032), + "call_type": "aretrieve_batch", + }, + "call_type": "aretrieve_batch", + "litellm_params": {}, + }, + None, + now, + now, + ) + + assert len(logger.log_queue) == 2 + assert logger.log_queue[0]["request_id"] == "chatcmpl-live-1" + assert logger.log_queue[1]["call_type"] == "aretrieve_batch" + assert logger.log_queue[1]["request_id"] == "batch-1" + assert logger.log_queue[1]["spend"] == 0.0001032 + if logger._flush_task is not None: + logger._flush_task.cancel() diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index deeb9c198d23..2690dd06b58e 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -9321,3 +9321,19 @@ async def test_line_item_logging_failure_leaves_aggregate_result_intact(self, mo assert batch._hidden_params["response_cost"] == 1.5 assert batch.usage.total_tokens == 15 + + +def test_caller_supplied_litellm_params_cannot_forge_logging_markers(logging_obj): + """A request body key `litellm_params` flows into optional_params; the logging + kwargs must keep the internal litellm_params or a caller could forge markers + (batch_parent_id) that spend sinks and router counters use to skip events.""" + from litellm.litellm_core_utils.core_helpers import is_batch_line_item_event + + logging_obj.update_environment_variables( + litellm_params={"metadata": {"model_group": "gpt-4o"}}, + optional_params={"litellm_params": {"batch_parent_id": "fake-batch"}}, + ) + + assert logging_obj.model_call_details["litellm_params"] is logging_obj.litellm_params + assert "batch_parent_id" not in logging_obj.model_call_details["litellm_params"] + assert not is_batch_line_item_event(logging_obj.model_call_details) diff --git a/tests/unit/router_strategy/test_least_busy.py b/tests/unit/router_strategy/test_least_busy.py index c2fa41f4ca84..cd0537dd8f56 100644 --- a/tests/unit/router_strategy/test_least_busy.py +++ b/tests/unit/router_strategy/test_least_busy.py @@ -208,3 +208,28 @@ async def test_an_open_circuit_breaker_falls_back_without_a_warning_per_request( assert picked is DEPLOYMENT_B assert [record.getMessage() for record in caplog.records if record.levelno >= logging.WARNING] == [] assert sum("circuit breaker is open" in record.getMessage() for record in caplog.records) == 2 + + +@pytest.mark.asyncio +async def test_batch_line_item_success_does_not_release_in_flight_slots(monkeypatch: pytest.MonkeyPatch) -> None: + """A batch line item is historical batch traffic (call_type=acompletion + + litellm_params.batch_parent_id): it never started a live request, so its success + callback must not decrement the in-flight count a live request still holds.""" + import litellm + + monkeypatch.setattr(litellm, "store_batch_line_items_in_callbacks", True, raising=False) + shared: Final = SharedRedisCounters() + worker: Final = _worker(shared) + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + + line_item_kwargs: Final = { + "call_type": "acompletion", + "litellm_params": { + "batch_parent_id": "batch-1", + "metadata": {"model_group": GROUP}, + "model_info": {"id": "dep-a"}, + }, + } + await worker.async_log_success_event(line_item_kwargs, None, None, None) + + assert shared.count(f"{GROUP}_request_count:dep-a") == 1 diff --git a/tests/unit/router_strategy/test_lowest_tpm_rpm.py b/tests/unit/router_strategy/test_lowest_tpm_rpm.py index 625f648bec46..55147b316dbc 100644 --- a/tests/unit/router_strategy/test_lowest_tpm_rpm.py +++ b/tests/unit/router_strategy/test_lowest_tpm_rpm.py @@ -117,3 +117,40 @@ async def async_get_available_deployments( f"from {HIGH_USAGE_DEPLOYMENT_ID}", f"from {LOW_USAGE_DEPLOYMENT_ID}", } + + +@pytest.mark.asyncio +async def test_usage_based_routing_handlers_skip_batch_line_items(monkeypatch: pytest.MonkeyPatch) -> None: + """ + Batch line-item callbacks carry call_type=acompletion plus + litellm_params.batch_parent_id, so the call-type-only batch guard does not fire. + They are historical batch traffic and must not update the TPM/RPM usage counters. + """ + import litellm + from litellm.router_strategy.lowest_tpm_rpm import LowestTPMLoggingHandler + + monkeypatch.setattr(litellm, "store_batch_line_items_in_callbacks", True, raising=False) + router: Final = Router( + model_list=[_deployment(HIGH_USAGE_DEPLOYMENT_ID)], + routing_strategy="usage-based-routing", + ) + handler: Final = LowestTPMLoggingHandler(router_cache=router.cache, routing_args={}) + line_item_kwargs: Final = { + "call_type": "acompletion", + "litellm_params": { + "batch_parent_id": "batch-1", + "metadata": {"model_group": MODEL_GROUP}, + "model_info": {"id": HIGH_USAGE_DEPLOYMENT_ID}, + }, + } + response_obj: Final = {"usage": {"total_tokens": 600}} + + handler.log_success_event(line_item_kwargs, response_obj, None, None) + await handler.async_log_success_event(line_item_kwargs, response_obj, None, None) + + moved: Final = sorted( + f"{key}={router.cache.in_memory_cache.cache_dict[key]}" + for key in router.cache.in_memory_cache.cache_dict + if ":tpm:" in key or ":rpm:" in key + ) + assert moved == [] diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 96dddf158699..d04495e12d9c 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -1263,6 +1263,53 @@ def test_sync_deployment_callback_on_success_skips_batch_retrieves( == expected_successes ) + +@pytest.mark.asyncio +async def test_deployment_callbacks_skip_batch_line_items(monkeypatch: pytest.MonkeyPatch): + """ + Batch line-item callbacks carry call_type=acompletion plus + litellm_params.batch_parent_id, so the call-type-only batch guards do not fire. + They describe historical batch traffic already reported by the aggregate + aretrieve_batch event and must not consume live TPM/RPM quota. + """ + monkeypatch.setattr(litellm, "store_batch_line_items_in_callbacks", True, raising=False) + router = litellm.Router( + model_list=[ + { + "model_name": _BATCH_GROUP, + "litellm_params": {"model": _BATCH_DEPLOYMENT_MODEL, "api_base": _BATCH_API_BASE, "api_key": "sk-fake"}, + "model_info": {"id": "batch-dep"}, + } + ] + ) + now = datetime.now() + line_item_kwargs = { + "call_type": "acompletion", + "standard_logging_object": {"total_tokens": _BATCH_TOKENS_PER_ROW}, + "litellm_params": { + "batch_parent_id": _BATCH_ID, + "metadata": {"model_group": _BATCH_GROUP, "deployment": _BATCH_DEPLOYMENT_MODEL}, + "model_info": {"id": "batch-dep"}, + }, + } + + await router.deployment_callback_on_success( + kwargs=line_item_kwargs, completion_response=None, start_time=now, end_time=now + ) + sync_key = router.sync_deployment_callback_on_success( + kwargs=line_item_kwargs, completion_response=None, start_time=now, end_time=now + ) + await router.async_deployment_callback_on_failure( + kwargs=line_item_kwargs, completion_response=None, start_time=now, end_time=now + ) + + assert sync_key is None + assert ( + get_deployment_successes_for_current_minute(litellm_router_instance=router, deployment_id="batch-dep") == 0 + ) + assert await _moved_routing_counters(router) == [] + assert await _router_usage_keys(router) == [] + _ROUTING_STRATEGY_CACHE_MARKERS = ("_map", "_request_count", ":tpm:", ":rpm:") @@ -1350,6 +1397,81 @@ async def test_arouter_aretrieve_batch_does_not_feed_routing_strategies( assert moved_counters == [] +_BATCH_INPUT_JSONL = "\n".join( + json.dumps( + { + "custom_id": f"row-{row}", + "method": "POST", + "url": "/v1/chat/completions", + "body": {"model": "gpt-4o-mini", "messages": [{"role": "user", "content": "hi"}]}, + } + ) + for row in range(_BATCH_ROWS) +) + + +def _mock_batch_input_file(respx_mock): + respx_mock.get(f"{_BATCH_API_BASE}/files/file-in-1/content").mock( + return_value=httpx.Response(200, text=_BATCH_INPUT_JSONL) + ) + + +async def _line_item_payloads(collector: "_BatchPayloadCollector", minimum: int, timeout: float = 5.0) -> list: + loop = asyncio.get_event_loop() + deadline = loop.time() + timeout + while loop.time() < deadline: + line_items = [p for p in collector.payloads if p and p.get("call_type") == "acompletion"] + if len(line_items) >= minimum: + return line_items + await asyncio.sleep(0.05) + raise AssertionError(f"expected at least {minimum} batch line-item payloads, saw {len(collector.payloads)} total") + + +@pytest.mark.parametrize( + "routing_strategy", + [ + "usage-based-routing", + "usage-based-routing-v2", + "latency-based-routing", + "cost-based-routing", + "least-busy", + ], +) +@pytest.mark.asyncio +async def test_arouter_aretrieve_batch_line_items_do_not_feed_routing_strategies( + monkeypatch: pytest.MonkeyPatch, routing_strategy: str +): + """ + With the line-item flag on, retrieving a completed batch also emits one child + callback per JSONL line (call_type=acompletion + litellm_params.batch_parent_id, + inheriting the parent deployment's model_info). Those children are historical + batch traffic: like the aggregate poll, they must not move the counters that + decide where the next live chat request goes. + """ + from litellm.batches.batch_line_item_logging import batch_line_item_claim_cache + + batch_line_item_claim_cache.in_memory_cache.flush_cache() + monkeypatch.setattr(litellm, "store_batch_line_items_in_callbacks", True, raising=False) + collector = _BatchPayloadCollector() + monkeypatch.setattr(litellm, "callbacks", [collector]) + monkeypatch.setattr(litellm, "input_callback", []) + router = _batch_fan_out_router(routing_strategy) + + with respx.mock(assert_all_called=True) as respx_mock: + _mock_batch_provider(respx_mock) + _mock_batch_input_file(respx_mock) + respx_mock.get(f"{_UNRELATED_BATCH_API_BASE}/batches/{_BATCH_ID}").mock( + return_value=httpx.Response(404, json=_BATCH_NOT_FOUND) + ) + response = await router.aretrieve_batch(batch_id=_BATCH_ID) + line_items = await _line_item_payloads(collector, minimum=_BATCH_ROWS) + moved_counters = await _moved_routing_counters(router) + + assert response.id == _BATCH_ID + assert len(line_items) == _BATCH_ROWS + assert moved_counters == [] + + @pytest.mark.asyncio async def test_arouter_aretrieve_file_content(): """ From 9c71dae02e8577c29f2b5a52d6b963200c8e3271 Mon Sep 17 00:00:00 2001 From: Yucheng He Date: Mon, 5 Oct 2026 01:07:58 -0700 Subject: [PATCH 43/89] fix(clickhouse): skip failed batch line items in the spend sink too A failed JSONL line is part of the aggregate aretrieve_batch row's request counts; writing it its own spend_logs row inflated request counts with the line-item flag on. Extend the batch line-item guard to the failure handler (review finding) and pin both directions: a failed line item writes no row, a genuine live failure still does. --- .../clickhouse/clickhouse_spend_logger.py | 4 +++ .../test_clickhouse_spend_logger.py | 27 +++++++++++++++++++ 2 files changed, 31 insertions(+) diff --git a/litellm/integrations/clickhouse/clickhouse_spend_logger.py b/litellm/integrations/clickhouse/clickhouse_spend_logger.py index 3659248ad3d8..4a10e29c7021 100644 --- a/litellm/integrations/clickhouse/clickhouse_spend_logger.py +++ b/litellm/integrations/clickhouse/clickhouse_spend_logger.py @@ -205,6 +205,10 @@ async def async_log_success_event(self, kwargs, response_obj, start_time, end_ti self._log(kwargs) async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time) -> None: + # Failed batch lines are part of the same aggregate row's request counts; + # per-line rows here would inflate spend-log request counts. + if is_batch_line_item_event(kwargs): + return self._log(kwargs) def _log(self, kwargs: Mapping[str, Any]) -> None: diff --git a/tests/test_litellm/integrations/clickhouse/test_clickhouse_spend_logger.py b/tests/test_litellm/integrations/clickhouse/test_clickhouse_spend_logger.py index 0bb7735477a7..37716f067e83 100644 --- a/tests/test_litellm/integrations/clickhouse/test_clickhouse_spend_logger.py +++ b/tests/test_litellm/integrations/clickhouse/test_clickhouse_spend_logger.py @@ -579,6 +579,33 @@ async def test_batch_line_item_success_event_does_not_write_spend(monkeypatch: p assert logger.log_queue[1]["call_type"] == "aretrieve_batch" assert logger.log_queue[1]["request_id"] == "batch-1" assert logger.log_queue[1]["spend"] == 0.0001032 + + # A FAILED batch line is still part of the aggregate row's request counts and + # must not add its own spend row; a genuine live failure still writes one. + await logger.async_log_failure_event( + { + "standard_logging_object": _minimal_payload("chatcmpl-line-2", status="failure", cost=0.0), + "call_type": "acompletion", + "litellm_params": {"batch_parent_id": "batch-1"}, + }, + None, + now, + now, + ) + await logger.async_log_failure_event( + { + "standard_logging_object": _minimal_payload("chatcmpl-live-2", status="failure", cost=0.0), + "call_type": "acompletion", + "litellm_params": {}, + }, + None, + now, + now, + ) + + assert len(logger.log_queue) == 3 + assert logger.log_queue[-1]["request_id"] == "chatcmpl-live-2" + assert logger.log_queue[-1]["status"] == "failure" if logger._flush_task is not None: logger._flush_task.cancel() From 8ec50a22f61551c3021d9e53b2576d01af7a6cdc Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 5 Oct 2026 08:12:18 +0000 Subject: [PATCH 44/89] fix(batches): declare batch_line_items service family on claim calls Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/batches/batch_line_item_logging.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/litellm/batches/batch_line_item_logging.py b/litellm/batches/batch_line_item_logging.py index df8034d87647..a53159eb6d57 100644 --- a/litellm/batches/batch_line_item_logging.py +++ b/litellm/batches/batch_line_item_logging.py @@ -7,6 +7,7 @@ from typing_extensions import assert_never +from litellm._internal_context import with_service_target from litellm._logging import verbose_logger from litellm.batches.batch_utils import ( BatchResultFiles, @@ -74,6 +75,7 @@ def _supported_line_provider(value: str) -> _BatchLineProvider | None: _ClaimResult: TypeAlias = Literal["claimed", "already_claimed", "unavailable"] +@with_service_target("batch_line_items") async def _claim_line_items(claim_cache: DualCache, claim_key: str, token: str) -> _ClaimResult: redis_cache: Final = claim_cache.redis_cache if redis_cache is None: @@ -90,6 +92,7 @@ async def _claim_line_items(claim_cache: DualCache, claim_key: str, token: str) return "already_claimed" +@with_service_target("batch_line_items") async def _release_line_item_claim(claim_cache: DualCache, claim_key: str, token: str) -> None: try: redis_cache: Final = claim_cache.redis_cache From a35d4ee464b94969f6a7e2ff7a765b3784cb46c9 Mon Sep 17 00:00:00 2001 From: Yucheng He Date: Mon, 5 Oct 2026 12:26:15 -0700 Subject: [PATCH 45/89] fix(integrations): keep batch line items out of the built-in metering sinks With store_batch_line_items_in_callbacks on, per-line batch events carry a real response_cost and reached six metering sinks that lacked the is_batch_line_item_event guard, so a completed batch metered aggregate + per-line (roughly 2x true spend): prometheus (litellm_spend_metric and request counters), openmeter, lago, datadog cost management (FOCUS BilledCost), newrelic metrics, and the OTel v1/v2 gen_ai.usage.cost metrics. Skip line-item events in each sink's success/failure metering entry points, matching the guards already on the Postgres and ClickHouse spend sinks. OTel spans for line items are still emitted; only their cost/token metrics are skipped. Tests drive real sink instances on the real dispatch lists through a completed 2-line batch and assert each sink meters exactly the aggregate. --- .../datadog/datadog_cost_management.py | 5 + litellm/integrations/lago.py | 6 + .../integrations/newrelic/newrelic_metrics.py | 6 + litellm/integrations/openmeter.py | 6 + litellm/integrations/opentelemetry.py | 6 + litellm/integrations/otel/plumbing/metrics.py | 9 +- litellm/integrations/prometheus.py | 7 + .../test_batch_line_item_metering_sinks.py | 376 ++++++++++++++++++ 8 files changed, 419 insertions(+), 2 deletions(-) create mode 100644 tests/unit/integrations/test_batch_line_item_metering_sinks.py diff --git a/litellm/integrations/datadog/datadog_cost_management.py b/litellm/integrations/datadog/datadog_cost_management.py index 538dd95abdd2..73017b79f7b3 100644 --- a/litellm/integrations/datadog/datadog_cost_management.py +++ b/litellm/integrations/datadog/datadog_cost_management.py @@ -14,6 +14,7 @@ get_datadog_service, normalize_datadog_tag_value, ) +from litellm.litellm_core_utils.core_helpers import is_batch_line_item_event from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, @@ -72,6 +73,10 @@ def __init__(self, cost_tag_keys: list[str] | None = None, **kwargs): super().__init__(**kwargs) async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + # A batch line item is billed by the aggregate aretrieve_batch event; a + # per-line FOCUS BilledCost row would double-count cloud spend. + if is_batch_line_item_event(kwargs): + return try: standard_logging_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object", None) diff --git a/litellm/integrations/lago.py b/litellm/integrations/lago.py index 594427b1e0a0..193f809a5fd7 100644 --- a/litellm/integrations/lago.py +++ b/litellm/integrations/lago.py @@ -11,6 +11,7 @@ from litellm._logging import verbose_logger from litellm._uuid import uuid from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.core_helpers import is_batch_line_item_event from litellm.llms.custom_httpx.http_handler import ( HTTPHandler, get_async_httpx_client, @@ -121,6 +122,9 @@ def _common_logic(self, kwargs: dict, response_obj) -> dict: return returned_val def log_success_event(self, kwargs, response_obj, start_time, end_time): + # A batch line item is billed by the aggregate aretrieve_batch event. + if is_batch_line_item_event(kwargs): + return _url = os.getenv("LAGO_API_BASE") assert _url is not None and isinstance(_url, str), ( f"LAGO_API_BASE missing or not set correctly. LAGO_API_BASE={_url}" @@ -153,6 +157,8 @@ def log_success_event(self, kwargs, response_obj, start_time, end_time): raise e async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + if is_batch_line_item_event(kwargs): + return try: verbose_logger.debug("ENTERS LAGO CALLBACK") _url = os.getenv("LAGO_API_BASE") diff --git a/litellm/integrations/newrelic/newrelic_metrics.py b/litellm/integrations/newrelic/newrelic_metrics.py index a2cedeba0f52..8bd77c981f63 100644 --- a/litellm/integrations/newrelic/newrelic_metrics.py +++ b/litellm/integrations/newrelic/newrelic_metrics.py @@ -34,6 +34,7 @@ from litellm._logging import verbose_logger from litellm.integrations.custom_batch_logger import CustomBatchLogger +from litellm.litellm_core_utils.core_helpers import is_batch_line_item_event from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, @@ -301,12 +302,17 @@ async def periodic_flush(self) -> None: await self._final_drain() async def async_log_success_event(self, kwargs, response_obj, start_time, end_time) -> None: + # A batch line item is metered by the aggregate aretrieve_batch event. + if is_batch_line_item_event(kwargs): + return try: await self._log_async_event(standard_logging_object=kwargs.get("standard_logging_object", None)) except Exception as e: # noqa: BLE001 # logging must never break the request path verbose_logger.exception("New Relic Metrics Layer Error - %s\n%s", e, traceback.format_exc()) async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time) -> None: + if is_batch_line_item_event(kwargs): + return try: await self._log_async_event(standard_logging_object=kwargs.get("standard_logging_object", None)) except Exception as e: # noqa: BLE001 # logging must never break the request path diff --git a/litellm/integrations/openmeter.py b/litellm/integrations/openmeter.py index db2fe386dec0..029c684f9693 100644 --- a/litellm/integrations/openmeter.py +++ b/litellm/integrations/openmeter.py @@ -9,6 +9,7 @@ import litellm from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.core_helpers import is_batch_line_item_event from litellm.llms.custom_httpx.http_handler import ( HTTPHandler, get_async_httpx_client, @@ -97,6 +98,9 @@ def _common_logic(self, kwargs: dict, response_obj): } def log_success_event(self, kwargs, response_obj, start_time, end_time): + # A batch line item is billed by the aggregate aretrieve_batch event. + if is_batch_line_item_event(kwargs): + return _url = os.getenv("OPENMETER_API_ENDPOINT", "https://openmeter.cloud") if _url.endswith("/"): _url += "api/v1/events" @@ -123,6 +127,8 @@ def log_success_event(self, kwargs, response_obj, start_time, end_time): raise e async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + if is_batch_line_item_event(kwargs): + return _url = os.getenv("OPENMETER_API_ENDPOINT", "https://openmeter.cloud") if _url.endswith("/"): _url += "api/v1/events" diff --git a/litellm/integrations/opentelemetry.py b/litellm/integrations/opentelemetry.py index 8d588896b2f3..bd9d1057a9d9 100644 --- a/litellm/integrations/opentelemetry.py +++ b/litellm/integrations/opentelemetry.py @@ -28,6 +28,7 @@ from litellm.integrations.otel.model.metadata import flatten_metadata from litellm.integrations.otel.model.semconv import LiteLLM, Metric from litellm.integrations.otel.plumbing.otlp_tls import resolve_otlp_http_tls +from litellm.litellm_core_utils.core_helpers import is_batch_line_item_event from litellm.litellm_core_utils.internal_call_metadata import is_unbilled_non_inference_call_from_params from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.secret_redaction import redact_string @@ -1655,6 +1656,11 @@ def _metric_attribute_allowed(self, key: str) -> bool: return True def _record_metrics(self, kwargs, response_obj, start_time, end_time): + # A batch line item's tokens and cost are metered by the aggregate + # aretrieve_batch event; per-line samples here would double-count them. + # Spans for line items are still emitted by _handle_success. + if is_batch_line_item_event(kwargs): + return duration_s: Final = (end_time - start_time).total_seconds() params: Final = kwargs.get("litellm_params") or {} provider: Final = _provider_label(params.get("custom_llm_provider")) diff --git a/litellm/integrations/otel/plumbing/metrics.py b/litellm/integrations/otel/plumbing/metrics.py index e1623f4697f0..005658edb458 100644 --- a/litellm/integrations/otel/plumbing/metrics.py +++ b/litellm/integrations/otel/plumbing/metrics.py @@ -33,6 +33,7 @@ resolve_provider, ) from litellm.integrations.otel.model.utils import to_seconds +from litellm.litellm_core_utils.core_helpers import is_batch_line_item_event from litellm.litellm_core_utils.internal_call_metadata import is_unbilled_non_inference_call_from_params from litellm.litellm_core_utils.safe_json_dumps import safe_dumps @@ -226,13 +227,17 @@ def record( usage_is_replayed: Final = is_unbilled_non_inference_call_from_params( kwargs.get("call_type"), kwargs.get("litellm_params"), response_obj ) + # A batch line item's tokens and cost are metered by the aggregate + # aretrieve_batch event; per-line samples here would double-count them. + # Duration samples for line items are still recorded. + batch_line_item: Final = is_batch_line_item_event(kwargs) self._metrics.operation_duration.record(duration_s, attributes=common_attrs) - if not usage_is_replayed: + if not usage_is_replayed and not batch_line_item: self._record_token_usage(response_obj, common_attrs) cost: Final = kwargs.get("response_cost") - if cost: + if cost and not batch_line_item: self._metrics.token_cost.record(cost, attributes=common_attrs) self._record_time_to_first_token(kwargs, common_attrs) diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 6468dc41ea5f..47b5cae63fef 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -35,6 +35,7 @@ from litellm.litellm_core_utils.core_helpers import ( get_litellm_metadata_from_kwargs, get_metadata_variable_name_from_kwargs, + is_batch_line_item_event, ) from litellm.litellm_core_utils.service_tier_utils import ( get_service_tier_from_standard_logging_payload, @@ -1345,6 +1346,10 @@ def _inc_labeled_counter( self._track_end_user_metric_series(counter, metric_name, _labels) async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + # A batch line item is metered by the aggregate aretrieve_batch event; a + # per-line sample here would double-count spend and requests. + if is_batch_line_item_event(kwargs): + return # Define prometheus client verbose_logger.debug( "prometheus Logging - Enters success logging function (kwargs keys: %s)", @@ -2365,6 +2370,8 @@ def _set_latency_metrics( ) async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + if is_batch_line_item_event(kwargs): + return verbose_logger.debug( "prometheus Logging - Enters failure logging function (kwargs keys: %s)", list(kwargs.keys()) if isinstance(kwargs, dict) else type(kwargs).__name__, diff --git a/tests/unit/integrations/test_batch_line_item_metering_sinks.py b/tests/unit/integrations/test_batch_line_item_metering_sinks.py new file mode 100644 index 000000000000..d31c25aa6213 --- /dev/null +++ b/tests/unit/integrations/test_batch_line_item_metering_sinks.py @@ -0,0 +1,376 @@ +""" +Guards that keep batch line-item callback events out of the built-in metering sinks. + +With ``litellm.store_batch_line_items_in_callbacks`` on, a completed batch emits +one child callback per JSONL line on top of the aggregate ``aretrieve_batch`` +event. Every line carries a real ``response_cost``, so any billing/metering sink +without the ``is_batch_line_item_event`` guard meters aggregate + per-line and +reports roughly 2x the true spend. + +These tests put REAL sink instances on the REAL dispatch lists, drive the REAL +``Logging.async_success_handler`` for a completed 2-line batch (1 success line + +1 error line, aggregate cost $1.50, per-line cost $0.03), and assert each sink +meters exactly the aggregate. OTel spans for line items must still be emitted: +the guard belongs on cost/token metrics, not on tracing. +""" + +import io +import json +import os +import time +import uuid +from contextlib import redirect_stdout +from datetime import datetime +from types import SimpleNamespace +from typing import Any, Final +from unittest.mock import AsyncMock, patch + +import pytest + +import litellm +from litellm.batches.batch_line_item_logging import batch_line_item_claim_cache +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.types.utils import LiteLLMBatch, Usage + +AGGREGATE_COST: Final[float] = 1.5 + +INPUT_JSONL: Final[bytes] = b"\n".join( + [ + json.dumps( + { + "custom_id": "a", + "method": "POST", + "url": "/v1/chat/completions", + "body": {"model": "gpt-4o", "messages": [{"role": "user", "content": "hi a"}]}, + } + ).encode(), + json.dumps( + { + "custom_id": "b", + "method": "POST", + "url": "/v1/chat/completions", + "body": {"model": "gpt-4o", "messages": [{"role": "user", "content": "hi b"}]}, + } + ).encode(), + ] +) + +OUTPUT_JSONL: Final[bytes] = json.dumps( + { + "custom_id": "a", + "response": { + "status_code": 200, + "body": { + "id": "chatcmpl-line-1", + "model": "gpt-4o", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + }, + }, + } +).encode() + +ERROR_JSONL: Final[bytes] = json.dumps( + { + "custom_id": "b", + "response": {"status_code": 400, "body": {"error": {"message": "boom"}}}, + "error": {"message": "boom"}, + } +).encode() + +_FILE_BYTES: Final[dict[str, bytes]] = { + "input-file-1": INPUT_JSONL, + "output-file-1": OUTPUT_JSONL, + "error-file-1": ERROR_JSONL, +} + + +def _file_content(file_id: str, **_kwargs: Any) -> SimpleNamespace: + return SimpleNamespace(content=_FILE_BYTES[file_id]) + + +def _batch() -> LiteLLMBatch: + return LiteLLMBatch( + id=f"batch_{uuid.uuid4().hex[:8]}", + object="batch", + endpoint="/v1/chat/completions", + input_file_id="input-file-1", + output_file_id="output-file-1", + error_file_id="error-file-1", + status="completed", + completion_window="24h", + created_at=1, + ) + + +def _parent_logging() -> Logging: + logging_obj = Logging( + model="gpt-4o", + messages=[{"role": "user", "content": ""}], + stream=False, + call_type="aretrieve_batch", + start_time=datetime.now(), + litellm_call_id=str(uuid.uuid4()), + function_id=str(uuid.uuid4()), + ) + logging_obj.update_environment_variables( + litellm_params={ + "metadata": { + "model_info": {"id": "dep-1"}, + "model_group": "gpt-4o", + "user_api_key_user_id": "user-77", + "user_api_key_team_id": "team-7", + "user_api_key_team_alias": "team-seven", + "user_api_key_alias": "key-alias-1", + } + }, + optional_params={}, + custom_llm_provider="openai", + ) + return logging_obj + + +class RecordingHTTP: + """Stands in for a sink's HTTP egress object only; the sink logic is real.""" + + def __init__(self) -> None: + self.posts: list[dict[str, Any]] = [] + + async def post(self, url: str, data: Any = None, content: Any = None, **_kw: Any) -> SimpleNamespace: + self.posts.append({"url": url, "body": data if data is not None else content}) + return SimpleNamespace(status_code=200, text="ok", raise_for_status=lambda: None) + + async def put(self, url: str, content: Any = None, **_kw: Any) -> SimpleNamespace: + return SimpleNamespace(status_code=202, text="ok", raise_for_status=lambda: None) + + +async def _log_completed_batch(monkeypatch: pytest.MonkeyPatch, loggers: list) -> None: + batch_line_item_claim_cache.in_memory_cache.flush_cache() + monkeypatch.setattr(litellm, "store_batch_line_items_in_callbacks", True, raising=False) + saved_success = list(litellm._async_success_callback) + saved_failure = list(litellm._async_failure_callback) + litellm._async_success_callback = list(loggers) + litellm._async_failure_callback = list(loggers) + buf = io.StringIO() + try: + with ( + patch("litellm.files.main.afile_content", new_callable=AsyncMock, side_effect=_file_content), + patch("litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02)), + redirect_stdout(buf), + ): + await _parent_logging().async_success_handler( + result=_batch(), + batch_cost=1.5, + batch_usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15), + batch_models=["gpt-4o"], + batch_successful_requests=1, + batch_failed_requests=1, + batch_prompt_cost=1.0, + batch_completion_cost=0.5, + ) + finally: + litellm._async_success_callback = saved_success + litellm._async_failure_callback = saved_failure + + +def _openmeter_sinks() -> tuple[Any, RecordingHTTP]: + from litellm.integrations.openmeter import OpenMeterLogger + + recorder = RecordingHTTP() + logger = OpenMeterLogger() + logger.async_http_handler = recorder + return logger, recorder + + +def _lago_sinks() -> tuple[Any, RecordingHTTP]: + from litellm.integrations.lago import LagoLogger + + recorder = RecordingHTTP() + logger = LagoLogger() + logger.async_http_handler = recorder + return logger, recorder + + +@pytest.mark.asyncio +async def test_openmeter_meters_only_the_aggregate_batch_event(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("OPENMETER_API_KEY", "test-openmeter-key") + logger, recorder = _openmeter_sinks() + + await _log_completed_batch(monkeypatch, [logger]) + + costs = [json.loads(p["body"])["data"]["cost"] for p in recorder.posts] + assert costs == [AGGREGATE_COST], ( + f"OpenMeter metered per-line costs on top of the aggregate: {costs}; " + "line items must be billed only by the aggregate aretrieve_batch event" + ) + + +@pytest.mark.asyncio +async def test_lago_bills_only_the_aggregate_batch_event(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("LAGO_API_KEY", "test-lago-key") + monkeypatch.setenv("LAGO_API_BASE", "http://lago.invalid") + monkeypatch.setenv("LAGO_API_EVENT_CODE", "litellm-usage") + monkeypatch.setenv("LAGO_API_CHARGE_BY", "user_id") + logger, recorder = _lago_sinks() + + await _log_completed_batch(monkeypatch, [logger]) + + costs = [json.loads(p["body"])["event"]["properties"]["response_cost"] for p in recorder.posts] + assert costs == [AGGREGATE_COST], ( + f"Lago billed per-line costs on top of the aggregate: {costs}; " + "line items must be billed only by the aggregate aretrieve_batch event" + ) + + +@pytest.mark.asyncio +async def test_datadog_cost_management_bills_only_the_aggregate_batch_event( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.integrations.datadog.datadog_cost_management import DatadogCostManagementLogger + + monkeypatch.setenv("DD_API_KEY", "test-dd-key") + monkeypatch.setenv("DD_APP_KEY", "test-dd-app-key") + logger = DatadogCostManagementLogger(cost_tag_keys=[]) + + await _log_completed_batch(monkeypatch, [logger]) + + entries = list(logger.log_queue) + costs = [e.get("response_cost", 0) for e in entries] + assert costs == [AGGREGATE_COST], ( + f"Datadog FOCUS BilledCost queued per-line entries on top of the aggregate: {costs}; " + "line items must be billed only by the aggregate aretrieve_batch event" + ) + + +@pytest.mark.asyncio +async def test_newrelic_metrics_meters_only_the_aggregate_batch_event(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.integrations.newrelic.newrelic_metrics import NewRelicMetricsLogger, build_metric_payload + + logger = NewRelicMetricsLogger(newrelic_api_key="test-nr-key") + + await _log_completed_batch(monkeypatch, [logger]) + + records = tuple(logger.log_queue) + costs = [r.response_cost for r in records] + assert costs == [AGGREGATE_COST], ( + f"New Relic metered per-line costs on top of the aggregate: {costs}; " + "line items must be metered only by the aggregate aretrieve_batch event" + ) + now = time.time() + envelopes = build_metric_payload(records=records, window_start=now - 1, now=now) + cost_sum = 0.0 + for envelope in envelopes: + for metric in envelope["metrics"]: + if "cost" in metric["name"]: + value = metric["value"]["sum"] if isinstance(metric["value"], dict) else metric["value"] + cost_sum += value + assert cost_sum == pytest.approx(AGGREGATE_COST) + + +@pytest.mark.asyncio +async def test_prometheus_meters_only_the_aggregate_batch_event(monkeypatch: pytest.MonkeyPatch) -> None: + prometheus_client = pytest.importorskip("prometheus_client") + from litellm.integrations.prometheus import PrometheusLogger + from prometheus_client import REGISTRY + + for collector in list(REGISTRY._collector_to_names.keys()): + REGISTRY.unregister(collector) + logger = PrometheusLogger() + + await _log_completed_batch(monkeypatch, [logger]) + + spend_samples = [] + for metric in REGISTRY.collect(): + if metric.name == "litellm_spend_metric": + spend_samples = [sample.value for sample in metric.samples if sample.name.endswith("_total")] + assert spend_samples, "litellm_spend_metric saw no samples at all" + assert sum(spend_samples) == pytest.approx(AGGREGATE_COST), ( + f"Prometheus spend metric double-counted batch line items: {spend_samples}; " + "line items must be metered only by the aggregate aretrieve_batch event" + ) + + +def _otel_v1(monkeypatch: pytest.MonkeyPatch): + otel_sdk = pytest.importorskip("opentelemetry.sdk") + from opentelemetry.sdk.metrics import MeterProvider + from opentelemetry.sdk.metrics.export import InMemoryMetricReader + from opentelemetry.sdk.trace import TracerProvider + from opentelemetry.sdk.trace.export import SimpleSpanProcessor + from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + + from litellm.integrations.opentelemetry import OpenTelemetry as OTelV1, OpenTelemetryConfig + + reader = InMemoryMetricReader() + span_exporter = InMemorySpanExporter() + tracer_provider = TracerProvider() + tracer_provider.add_span_processor(SimpleSpanProcessor(span_exporter)) + logger = OTelV1( + config=OpenTelemetryConfig(exporter="console", enable_metrics=True), + callback_name="batch_sink_guard_v1", + tracer_provider=tracer_provider, + meter_provider=MeterProvider(metric_readers=[reader]), + ) + return logger, reader, span_exporter + + +def _cost_points(reader: Any) -> list[float]: + data = reader.get_metrics_data() + points: list[float] = [] + if data is None: + return points + for resource_metrics in data.resource_metrics: + for scope_metrics in resource_metrics.scope_metrics: + for metric in scope_metrics.metrics: + if not metric.name.endswith("cost"): + continue + for data_point in metric.data.data_points: + points.append(getattr(data_point, "sum", getattr(data_point, "value", None))) + return points + + +@pytest.mark.asyncio +async def test_otel_v1_meter_cost_skips_line_items_but_spans_stay( + monkeypatch: pytest.MonkeyPatch, +) -> None: + logger, reader, span_exporter = _otel_v1(monkeypatch) + + await _log_completed_batch(monkeypatch, [logger]) + + costs = _cost_points(reader) + assert sum(costs) == pytest.approx(AGGREGATE_COST), ( + f"OTel v1 gen_ai.usage.cost double-counted batch line items: {costs}; " + "line items must be metered only by the aggregate aretrieve_batch event" + ) + finished = span_exporter.get_finished_spans() + assert finished, "OTel v1 emitted no spans at all" + assert any("chatcmpl-line-1" in str(span.attributes) for span in finished), ( + f"OTel v1 dropped the per-line span the feature exists to deliver: " + f"{[span.name for span in finished]}" + ) + + +@pytest.mark.asyncio +async def test_otel_v2_meter_cost_skips_line_items(monkeypatch: pytest.MonkeyPatch) -> None: + pytest.importorskip("opentelemetry.sdk") + from opentelemetry.sdk.metrics import MeterProvider + from opentelemetry.sdk.metrics.export import InMemoryMetricReader + from opentelemetry.sdk.trace import TracerProvider + + from litellm.integrations.otel.logger import OpenTelemetryV2 + from litellm.integrations.otel.model.config import OpenTelemetryV2Config + + reader = InMemoryMetricReader() + logger = OpenTelemetryV2( + config=OpenTelemetryV2Config(exporter="console", enable_metrics=True), + callback_name="batch_sink_guard_v2", + tracer_provider=TracerProvider(), + meter_provider=MeterProvider(metric_readers=[reader]), + ) + + await _log_completed_batch(monkeypatch, [logger]) + + costs = _cost_points(reader) + assert sum(costs) == pytest.approx(AGGREGATE_COST), ( + f"OTel v2 gen_ai.usage.cost double-counted batch line items: {costs}; " + "line items must be metered only by the aggregate aretrieve_batch event" + ) From c9f67ff43e354939b3f360bda5a6648ecf84842c Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 5 Oct 2026 19:31:35 +0000 Subject: [PATCH 46/89] style: format metering-sinks test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/unit/integrations/test_batch_line_item_metering_sinks.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tests/unit/integrations/test_batch_line_item_metering_sinks.py b/tests/unit/integrations/test_batch_line_item_metering_sinks.py index d31c25aa6213..a89d8ea9a531 100644 --- a/tests/unit/integrations/test_batch_line_item_metering_sinks.py +++ b/tests/unit/integrations/test_batch_line_item_metering_sinks.py @@ -344,8 +344,7 @@ async def test_otel_v1_meter_cost_skips_line_items_but_spans_stay( finished = span_exporter.get_finished_spans() assert finished, "OTel v1 emitted no spans at all" assert any("chatcmpl-line-1" in str(span.attributes) for span in finished), ( - f"OTel v1 dropped the per-line span the feature exists to deliver: " - f"{[span.name for span in finished]}" + f"OTel v1 dropped the per-line span the feature exists to deliver: {[span.name for span in finished]}" ) From 1b02cb64fb1799784f1ea9cb58d9644c074bb087 Mon Sep 17 00:00:00 2001 From: Yucheng He Date: Mon, 5 Oct 2026 12:47:25 -0700 Subject: [PATCH 47/89] fix(otel): skip synthetic line-item durations in v2 metrics; isolate prometheus test Follow-ups from review on the sink guards: record_failure() now skips batch line items too (a failed line's interval measures retrieval and emission, not its model call), and the prometheus guard test snapshots and restores the process-wide registry collectors instead of leaving them unregistered. The OTel v2 test is now a flag on/off A/B over the whole metric snapshot: cost and usage values must match exactly and no metric may gain datapoints with the flag on. --- litellm/integrations/otel/plumbing/metrics.py | 18 +++-- .../test_batch_line_item_metering_sinks.py | 78 +++++++++++++++---- 2 files changed, 73 insertions(+), 23 deletions(-) diff --git a/litellm/integrations/otel/plumbing/metrics.py b/litellm/integrations/otel/plumbing/metrics.py index 005658edb458..3af54f7f4594 100644 --- a/litellm/integrations/otel/plumbing/metrics.py +++ b/litellm/integrations/otel/plumbing/metrics.py @@ -222,22 +222,24 @@ def record( start_time: datetime, end_time: datetime, ) -> None: + # A batch line item is metered by the aggregate aretrieve_batch event, and + # its child interval (parent retrieve start -> line emission) measures + # retrieval and callback processing rather than the line's model call, so + # skip every metric for it; spans for line items are still emitted. + if is_batch_line_item_event(kwargs): + return common_attrs: Final = self._filter_attributes(self._bounded_attributes(kwargs)) duration_s: Final = (end_time - start_time).total_seconds() usage_is_replayed: Final = is_unbilled_non_inference_call_from_params( kwargs.get("call_type"), kwargs.get("litellm_params"), response_obj ) - # A batch line item's tokens and cost are metered by the aggregate - # aretrieve_batch event; per-line samples here would double-count them. - # Duration samples for line items are still recorded. - batch_line_item: Final = is_batch_line_item_event(kwargs) self._metrics.operation_duration.record(duration_s, attributes=common_attrs) - if not usage_is_replayed and not batch_line_item: + if not usage_is_replayed: self._record_token_usage(response_obj, common_attrs) cost: Final = kwargs.get("response_cost") - if cost and not batch_line_item: + if cost: self._metrics.token_cost.record(cost, attributes=common_attrs) self._record_time_to_first_token(kwargs, common_attrs) @@ -251,6 +253,10 @@ def record_failure( start_time: datetime, end_time: datetime, ) -> None: + # A batch line item's interval measures retrieval and emission, not the + # line's model call, so its duration would be synthetic here too. + if is_batch_line_item_event(kwargs): + return """Record the one metric a failed request can honestly report: the operation's duration, tagged with ``error.type``. diff --git a/tests/unit/integrations/test_batch_line_item_metering_sinks.py b/tests/unit/integrations/test_batch_line_item_metering_sinks.py index a89d8ea9a531..525f6ae15d04 100644 --- a/tests/unit/integrations/test_batch_line_item_metering_sinks.py +++ b/tests/unit/integrations/test_batch_line_item_metering_sinks.py @@ -144,9 +144,9 @@ async def put(self, url: str, content: Any = None, **_kw: Any) -> SimpleNamespac return SimpleNamespace(status_code=202, text="ok", raise_for_status=lambda: None) -async def _log_completed_batch(monkeypatch: pytest.MonkeyPatch, loggers: list) -> None: +async def _log_completed_batch(monkeypatch: pytest.MonkeyPatch, loggers: list, flag: bool = True) -> None: batch_line_item_claim_cache.in_memory_cache.flush_cache() - monkeypatch.setattr(litellm, "store_batch_line_items_in_callbacks", True, raising=False) + monkeypatch.setattr(litellm, "store_batch_line_items_in_callbacks", flag, raising=False) saved_success = list(litellm._async_success_callback) saved_failure = list(litellm._async_failure_callback) litellm._async_success_callback = list(loggers) @@ -273,16 +273,26 @@ async def test_prometheus_meters_only_the_aggregate_batch_event(monkeypatch: pyt from litellm.integrations.prometheus import PrometheusLogger from prometheus_client import REGISTRY - for collector in list(REGISTRY._collector_to_names.keys()): + saved_collectors = list(REGISTRY._collector_to_names.keys()) + for collector in saved_collectors: REGISTRY.unregister(collector) logger = PrometheusLogger() - await _log_completed_batch(monkeypatch, [logger]) + try: + await _log_completed_batch(monkeypatch, [logger]) + spend_samples = [] + for metric in REGISTRY.collect(): + if metric.name == "litellm_spend_metric": + spend_samples = [sample.value for sample in metric.samples if sample.name.endswith("_total")] + finally: + for collector in list(REGISTRY._collector_to_names.keys()): + REGISTRY.unregister(collector) + for collector in saved_collectors: + try: + REGISTRY.register(collector) + except Exception: # noqa: BLE001 # already re-registered by another holder + pass - spend_samples = [] - for metric in REGISTRY.collect(): - if metric.name == "litellm_spend_metric": - spend_samples = [sample.value for sample in metric.samples if sample.name.endswith("_total")] assert spend_samples, "litellm_spend_metric saw no samples at all" assert sum(spend_samples) == pytest.approx(AGGREGATE_COST), ( f"Prometheus spend metric double-counted batch line items: {spend_samples}; " @@ -348,9 +358,25 @@ async def test_otel_v1_meter_cost_skips_line_items_but_spans_stay( ) -@pytest.mark.asyncio -async def test_otel_v2_meter_cost_skips_line_items(monkeypatch: pytest.MonkeyPatch) -> None: - pytest.importorskip("opentelemetry.sdk") +def _reader_snapshot(reader: Any) -> dict[str, list[float]]: + """Every metric's datapoint values, keyed by metric name (sums for histograms).""" + data = reader.get_metrics_data() + out: dict[str, list[float]] = {} + if data is None: + return out + for resource_metrics in data.resource_metrics: + for scope_metrics in resource_metrics.scope_metrics: + for metric in scope_metrics.metrics: + for data_point in metric.data.data_points: + value = getattr(data_point, "sum", None) + if value is None: + value = getattr(data_point, "value", None) + if value is not None: + out.setdefault(metric.name, []).append(value) + return out + + +async def _run_batch_with_otel_v2(monkeypatch: pytest.MonkeyPatch, flag: bool) -> dict[str, list[float]]: from opentelemetry.sdk.metrics import MeterProvider from opentelemetry.sdk.metrics.export import InMemoryMetricReader from opentelemetry.sdk.trace import TracerProvider @@ -365,11 +391,29 @@ async def test_otel_v2_meter_cost_skips_line_items(monkeypatch: pytest.MonkeyPat tracer_provider=TracerProvider(), meter_provider=MeterProvider(metric_readers=[reader]), ) + await _log_completed_batch(monkeypatch, [logger], flag=flag) + return _reader_snapshot(reader) - await _log_completed_batch(monkeypatch, [logger]) - costs = _cost_points(reader) - assert sum(costs) == pytest.approx(AGGREGATE_COST), ( - f"OTel v2 gen_ai.usage.cost double-counted batch line items: {costs}; " - "line items must be metered only by the aggregate aretrieve_batch event" - ) +@pytest.mark.asyncio +async def test_otel_v2_metrics_are_identical_with_the_flag_on_and_off( + monkeypatch: pytest.MonkeyPatch, +) -> None: + pytest.importorskip("opentelemetry.sdk") + + off = await _run_batch_with_otel_v2(monkeypatch, flag=False) + on = await _run_batch_with_otel_v2(monkeypatch, flag=True) + + for name in ("gen_ai.usage.cost", "gen_ai.client.token.usage"): + assert on.get(name) == off.get(name), ( + f"OTel v2 {name} changed when the line-item flag was turned on; line items " + "must be metered only by the aggregate aretrieve_batch event. " + f"off={off.get(name)} on={on.get(name)}" + ) + for name, values in off.items(): + on_counts = [len(on.get(name, [])), len(values)] + assert on_counts[0] == on_counts[1], ( + f"OTel v2 {name} gained synthetic per-line samples with the flag on: " + f"{on_counts[0]} datapoints on vs {on_counts[1]} off" + ) + assert sum(off.get("gen_ai.usage.cost", [])) == pytest.approx(AGGREGATE_COST) From 1b63787eaa39a5ecee3e38b104002f7df9f0fdf0 Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 7 Oct 2026 00:13:38 +0000 Subject: [PATCH 48/89] chore: finish main merge conflict resolution for batch line items Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/batches/batch_utils.py | 8 +++-- tests/unit/test_router/test_router.py | 49 +++++++++++++++++++++++++-- 2 files changed, 53 insertions(+), 4 deletions(-) diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 03bca29f47d3..86dfb682fe19 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -464,7 +464,9 @@ def _provider_output_file_id(output_file_id: str) -> str: async def _fetch_batch_managed_file_content( file_id: str, - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral"] = "openai", + custom_llm_provider: Literal[ + "openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral" + ] = "openai", litellm_params: Mapping[str, object] | None = None, ) -> bytes: """ @@ -495,7 +497,9 @@ async def _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", "bedrock", "mistral"] = "openai", + custom_llm_provider: Literal[ + "openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral" + ] = "openai", litellm_params: Mapping[str, object] | None = None, ) -> bytes: """ diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index d7ddd5c33269..354178b6f210 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -1246,8 +1246,53 @@ def test_sync_deployment_callback_on_success_skips_batch_retrieves( == expected_successes ) - get_parent_otel_span_from_kwargs, - is_batch_line_item_event, + +@pytest.mark.asyncio +async def test_deployment_callbacks_skip_batch_line_items(monkeypatch: pytest.MonkeyPatch): + """ + Batch line-item callbacks carry call_type=acompletion plus + litellm_params.batch_parent_id, so the call-type-only batch guards do not fire. + They describe historical batch traffic already reported by the aggregate + aretrieve_batch event and must not consume live TPM/RPM quota. + """ + monkeypatch.setattr(litellm, "store_batch_line_items_in_callbacks", True, raising=False) + router = litellm.Router( + model_list=[ + { + "model_name": _BATCH_GROUP, + "litellm_params": {"model": _BATCH_DEPLOYMENT_MODEL, "api_base": _BATCH_API_BASE, "api_key": "sk-fake"}, + "model_info": {"id": "batch-dep"}, + } + ] + ) + now = datetime.now() + line_item_kwargs = { + "call_type": "acompletion", + "standard_logging_object": {"total_tokens": _BATCH_TOKENS_PER_ROW}, + "litellm_params": { + "batch_parent_id": _BATCH_ID, + "metadata": {"model_group": _BATCH_GROUP, "deployment": _BATCH_DEPLOYMENT_MODEL}, + "model_info": {"id": "batch-dep"}, + }, + } + + await router.deployment_callback_on_success( + kwargs=line_item_kwargs, completion_response=None, start_time=now, end_time=now + ) + sync_key = router.sync_deployment_callback_on_success( + kwargs=line_item_kwargs, completion_response=None, start_time=now, end_time=now + ) + await router.async_deployment_callback_on_failure( + kwargs=line_item_kwargs, completion_response=None, start_time=now, end_time=now + ) + + assert sync_key is None + assert ( + get_deployment_successes_for_current_minute(litellm_router_instance=router, deployment_id="batch-dep") == 0 + ) + assert await _moved_routing_counters(router) == [] + assert await _router_usage_keys(router) == [] + _ROUTING_STRATEGY_CACHE_MARKERS = ("_map", "_request_count", ":tpm:", ":rpm:") From 235019dbdc00c06f36f15ab6a01580110a6cf85b Mon Sep 17 00:00:00 2001 From: yucheng Date: Wed, 7 Oct 2026 00:34:42 +0000 Subject: [PATCH 49/89] style: sort imports in litellm_logging after main merge Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/litellm_core_utils/litellm_logging.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 4e2444b42e14..1a4e7810de73 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -35,8 +35,8 @@ from litellm._uuid import uuid from litellm.batches.batch_utils import ( BatchResultFiles, - handle_completed_batch_with_files, batch_cost_is_final, + handle_completed_batch_with_files, ) from litellm.caching.caching import DualCache from litellm.caching.caching_handler import LLMCachingHandler From b57e4f06e82e9a5d9e840e5f40b71bfe3c293ce6 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Thu, 8 Oct 2026 13:35:44 -0700 Subject: [PATCH 50/89] test(batches): use the public hook class names main now exports --- tests/unit/proxy/hooks/test_model_max_budget_limiter.py | 2 +- tests/unit/proxy/hooks/test_proxy_track_cost_callback.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/unit/proxy/hooks/test_model_max_budget_limiter.py b/tests/unit/proxy/hooks/test_model_max_budget_limiter.py index ed564bd552e7..b0ace15a84a5 100644 --- a/tests/unit/proxy/hooks/test_model_max_budget_limiter.py +++ b/tests/unit/proxy/hooks/test_model_max_budget_limiter.py @@ -202,7 +202,7 @@ async def test_a_batch_polled_within_every_budget_window_is_never_charged_again( async def test_batch_line_item_events_do_not_charge_the_model_budget(): """Line events carry batch_parent_id; the aggregate aretrieve_batch event is the one that already bills the batch, so children must not double-charge.""" - limiter: Final = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache()) + limiter: Final = PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache()) line_event = _event("acompletion", CHAT_COST) line_event["litellm_params"]["batch_parent_id"] = "batch_first" diff --git a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py index 29d41d7842bd..899584520d6d 100644 --- a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py @@ -2774,7 +2774,7 @@ async def test_async_post_call_failure_hook_persists_no_raw_model_on_an_unknown_ @pytest.mark.asyncio async def test_batch_line_item_event_never_updates_spend(): # test-quality-ok: the observable contract is exactly that the spend path is never invoked - logger: Final = _ProxyDBLogger() + logger: Final = ProxyDBLogger() kwargs: Final = { "litellm_params": {"batch_parent_id": "batch_1", "metadata": {}}, "model": "gpt-4o", From c653e2e771d46a32adc936ad604ba2aac1ee7184 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Thu, 8 Oct 2026 14:13:56 -0700 Subject: [PATCH 51/89] fix(batches): only the batch writer can mark a callback event as a line item amoderation, transcription and speech build litellm_params from every request kwarg, so a body field named batch_parent_id reached Logging.litellm_params and every spend, budget and rate-limit hook skipped the request. update_environment_variables now drops an incoming batch_parent_id; the line-item writer sets it through Logging.mark_batch_line_item. Batch line hidden params are emitted only when set. --- litellm/batches/batch_line_item_logging.py | 2 +- litellm/litellm_core_utils/core_helpers.py | 5 +- litellm/litellm_core_utils/litellm_logging.py | 15 ++- .../spend/test_batch_parent_id_body_field.py | 127 ++++++++++++++++++ .../batches/test_batch_line_item_logging.py | 4 +- .../test_litellm_logging.py | 44 +++++- 6 files changed, 186 insertions(+), 11 deletions(-) create mode 100644 tests/integration/spend/test_batch_parent_id_body_field.py diff --git a/litellm/batches/batch_line_item_logging.py b/litellm/batches/batch_line_item_logging.py index a53159eb6d57..9565614f78f2 100644 --- a/litellm/batches/batch_line_item_logging.py +++ b/litellm/batches/batch_line_item_logging.py @@ -358,13 +358,13 @@ async def _emit_line_event( child.update_environment_variables( # pyright: ignore[reportUnknownMemberType] # Logging.update_environment_variables is untyped upstream litellm_params={ **parent_params, - "batch_parent_id": batch.id, "metadata": _metadata_copy(parent_params), }, optional_params=_optional_params_for_body(request_body), model=child.model, custom_llm_provider=custom_llm_provider, ) + child.mark_batch_line_item(batch.id) for secret_key in _SECRET_PARAM_KEYS: child.litellm_params.pop(secret_key, None) # pyright: ignore[reportUnknownMemberType] # Logging.litellm_params is untyped upstream diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index d10a7a1a46d6..0a18fc19cd10 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -793,13 +793,16 @@ def set_response_cost_in_hidden_params(response: _CarriesHiddenParams, cost: flo hidden_params["additional_headers"] = merged +BATCH_PARENT_ID_KEY: Final = "batch_parent_id" + + def is_batch_line_item_event(kwargs: object) -> bool: if not isinstance(kwargs, Mapping): return False litellm_params: Final[object] = kwargs.get("litellm_params") if not isinstance(litellm_params, Mapping): return False - return bool(litellm_params.get("batch_parent_id")) + return bool(litellm_params.get(BATCH_PARENT_ID_KEY)) _HIDDEN_PARAMS_ADAPTER: Final = TypeAdapter(Mapping[str, object]) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 62cc69a75183..923059d01b98 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -75,6 +75,7 @@ is_classifier_call, ) from litellm.litellm_core_utils.core_helpers import ( + BATCH_PARENT_ID_KEY, get_provider_response_headers_from_hidden_params, is_expected_client_error, proxy_stamped_used_client_oauth_token, @@ -962,10 +963,17 @@ def update_environment_variables( if model is not None: self.model = model self.user = user + marked_batch_parent_id: Final[object] = self.litellm_params.get(BATCH_PARENT_ID_KEY) self.litellm_params = { **self.litellm_params, **scrub_sensitive_keys_in_metadata(litellm_params), } + # Every spend, budget and metering hook skips an event carrying batch_parent_id, so only + # mark_batch_line_item may set it; the incoming params can hold caller request fields + if marked_batch_parent_id is None: + self.litellm_params.pop(BATCH_PARENT_ID_KEY, None) + else: + self.litellm_params[BATCH_PARENT_ID_KEY] = marked_batch_parent_id self.litellm_request_debug = litellm_params.get("litellm_request_debug", False) self.logger_fn = litellm_params.get("logger_fn", None) if is_debugging_on() or self.litellm_request_debug: @@ -1003,6 +1011,10 @@ def update_environment_variables( if "custom_llm_provider" in self.model_call_details: self.custom_llm_provider = self.model_call_details["custom_llm_provider"] + def mark_batch_line_item(self, batch_id: str) -> None: + """Mark this logging object as one JSONL line of the completed batch ``batch_id``""" + self.litellm_params[BATCH_PARENT_ID_KEY] = batch_id + def update_from_kwargs( self, kwargs: dict, @@ -6148,9 +6160,6 @@ def get_hidden_params( batch_failed_requests=None, litellm_model_name=None, usage_object=None, - batch_id=None, - batch_custom_id=None, - batch_line_status_code=None, ) if hidden_params is not None: for key in StandardLoggingHiddenParams.__annotations__: diff --git a/tests/integration/spend/test_batch_parent_id_body_field.py b/tests/integration/spend/test_batch_parent_id_body_field.py new file mode 100644 index 000000000000..46d36020274b --- /dev/null +++ b/tests/integration/spend/test_batch_parent_id_body_field.py @@ -0,0 +1,127 @@ +from __future__ import annotations + +import io +import os +import uuid +import wave +from collections.abc import Callable +from hashlib import sha256 +from typing import Final + +import httpx +import pytest +from integration._support.client import Gateway, Scenario, eventually +from integration._support.database import read_rows +from integration._support.upstream import delete_scenario, register_scenario +from integration.cost_calculation.cost_tracking_case import BinaryResponse, JsonResponse +from pydantic import JsonValue + +FORGED_PARENT: Final = "batch_forged_by_caller" + +_Reply = Callable[[], JsonResponse | BinaryResponse] +_Send = Callable[[Gateway, str, str, dict[str, JsonValue]], httpx.Response] + + +def _unique_wav() -> bytes: + output: Final = io.BytesIO() + with wave.open(output, "wb") as wav: + wav.setnchannels(1) + wav.setsampwidth(2) + wav.setframerate(16000) + wav.writeframes(os.urandom(2 * 16000)) + return output.getvalue() + + +def _moderation_reply() -> JsonResponse: + return JsonResponse( + content_type="application/json", + body={ + "id": f"modr-{uuid.uuid4().hex}", + "model": "omni-moderation-latest", + "results": [{"flagged": False, "categories": {}, "category_scores": {}}], + }, + ) + + +def _speech_reply() -> BinaryResponse: + return BinaryResponse(content_type="audio/mpeg", length=2048) + + +def _transcription_reply() -> JsonResponse: + return JsonResponse(content_type="application/json", body={"text": "hello"}) + + +def _moderation(gateway: Gateway, model: str, key: str, extra: dict[str, JsonValue]) -> httpx.Response: + return gateway.request("POST", "/v1/moderations", {"model": model, "input": "safe text", **extra}, key=key) + + +def _speech(gateway: Gateway, model: str, key: str, extra: dict[str, JsonValue]) -> httpx.Response: + return gateway.request( + "POST", "/v1/audio/speech", {"model": model, "input": "hello world", "voice": "alloy", **extra}, key=key + ) + + +def _transcription(gateway: Gateway, model: str, key: str, extra: dict[str, JsonValue]) -> httpx.Response: + fields: Final = {"model": model, **{name: str(value) for name, value in extra.items()}} + return gateway.request_multipart( + "/v1/audio/transcriptions", fields, {"file": ("audio.wav", _unique_wav(), "audio/wav")}, key=key + ) + + +ROUTES: Final = ( + pytest.param("openai/omni-moderation-latest", _moderation_reply, _moderation, "amoderation", id="moderations"), + pytest.param("openai/tts-1", _speech_reply, _speech, "aspeech", id="audio-speech"), + pytest.param("openai/whisper-1", _transcription_reply, _transcription, "atranscription", id="audio-transcriptions"), +) + + +def _scripted_model(gateway: Gateway, scenario: Scenario, provider_model: str, reply: _Reply) -> str: + marker: Final = "batch-parent-body-" + uuid.uuid4().hex[:12] + handle: Final = register_scenario(marker, reply()) + scenario.cleanups.callback(delete_scenario, handle) + return scenario.model(model=provider_model, api_base=f"{gateway.upstream_url}/{marker}") + + +def _spend_rows(key: str) -> tuple[dict[str, JsonValue], ...]: + return tuple( + read_rows( + 'SELECT call_type, spend FROM "LiteLLM_SpendLogs" WHERE api_key=%s ORDER BY "startTime"', + (sha256(key.encode()).hexdigest(),), + ) + ) + + +def _key_spend(key: str) -> float: + rows: Final = read_rows( + 'SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s', (sha256(key.encode()).hexdigest(),) + ) + assert len(rows) == 1, rows + return float(str(rows[0]["spend"])) + + +@pytest.mark.parametrize(("provider_model", "reply", "send", "call_type"), ROUTES) +def test_a_caller_supplied_batch_parent_id_is_still_tracked_like_any_request( + gateway: Gateway, + provider_model: str, + reply: _Reply, + send: _Send, + call_type: str, +) -> None: + with gateway.scenario() as scenario: + control_model: Final = _scripted_model(gateway, scenario, provider_model, reply) + forged_model: Final = _scripted_model(gateway, scenario, provider_model, reply) + key: Final = scenario.key(models=[control_model, forged_model]) + + control: Final = send(gateway, control_model, key, {}) + assert control.status_code == 200, control.text + control_rows: Final = eventually(lambda: _spend_rows(key), lambda values: len(values) == 1, seconds=70) + assert control_rows[0]["call_type"] == call_type, control_rows + control_spend: Final = float(str(control_rows[0]["spend"])) + eventually(lambda: _key_spend(key), lambda value: value == pytest.approx(control_spend), seconds=70) + + forged: Final = send(gateway, forged_model, key, {"batch_parent_id": FORGED_PARENT}) + assert forged.status_code == 200, forged.text + rows: Final = eventually(lambda: _spend_rows(key), lambda values: len(values) == 2, seconds=70) + assert [row["call_type"] for row in rows] == [call_type, call_type], rows + assert float(str(rows[1]["spend"])) == pytest.approx(control_spend), rows + eventually(lambda: _key_spend(key), lambda value: value == pytest.approx(2 * control_spend), seconds=70) diff --git a/tests/unit/batches/test_batch_line_item_logging.py b/tests/unit/batches/test_batch_line_item_logging.py index e338cad5d214..43da11528b03 100644 --- a/tests/unit/batches/test_batch_line_item_logging.py +++ b/tests/unit/batches/test_batch_line_item_logging.py @@ -214,7 +214,7 @@ async def test_line_items_emitted_alongside_aggregate(recorder): aggregate = next(e for e in recorder.success_events if _hidden(e).get("batch_custom_id") is None) assert _payload(aggregate)["response_cost"] == 1.5 - assert _hidden(aggregate)["batch_custom_id"] is None + assert "batch_custom_id" not in _hidden(aggregate) line = next(e for e in recorder.success_events if _hidden(e).get("batch_custom_id") == "a") hidden = _hidden(line) @@ -265,7 +265,7 @@ async def test_line_items_skipped_for_unsupported_provider(recorder): file_mock.assert_not_awaited() assert len(recorder.success_events) == 1 assert _payload(recorder.success_events[0])["response_cost"] == 1.5 - assert _hidden(recorder.success_events[0])["batch_custom_id"] is None + assert "batch_custom_id" not in _hidden(recorder.success_events[0]) assert len(recorder.failure_events) == 0 diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index 84aa3e695620..07d1b1024b13 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -32,6 +32,7 @@ from litellm.constants import LOGGING_WORKER_MAX_TIME_PER_COROUTINE, REDACTED_BY_LITELLM, SENTRY_PII_DENYLIST from litellm.cost_calculator import ocr_batch_cost from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.core_helpers import BATCH_PARENT_ID_KEY, is_batch_line_item_event from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging from litellm.litellm_core_utils.litellm_logging import ( Logging, @@ -9769,6 +9770,41 @@ def test_caller_supplied_litellm_params_cannot_forge_logging_markers(logging_obj assert not is_batch_line_item_event(logging_obj.model_call_details) +def test_a_top_level_batch_parent_id_request_field_is_not_a_line_item_marker(logging_obj): + """moderation, transcription and speech build litellm_params from every request kwarg, + so a body field named batch_parent_id lands in the litellm_params argument itself""" + logging_obj.update_environment_variables( + litellm_params={"metadata": {"model_group": "tts-1"}, BATCH_PARENT_ID_KEY: "batch_forged_by_caller"}, + optional_params={}, + ) + + assert BATCH_PARENT_ID_KEY not in logging_obj.model_call_details["litellm_params"] + assert logging_obj.model_call_details["litellm_params"]["metadata"] == {"model_group": "tts-1"} + assert not is_batch_line_item_event(logging_obj.model_call_details) + + +def test_a_marked_batch_line_item_keeps_its_batch_through_later_updates(logging_obj): + logging_obj.update_environment_variables(litellm_params={}, optional_params={}) + logging_obj.mark_batch_line_item("batch_abc") + + logging_obj.update_environment_variables( + litellm_params={BATCH_PARENT_ID_KEY: "batch_forged_by_caller"}, optional_params={} + ) + + assert logging_obj.model_call_details["litellm_params"][BATCH_PARENT_ID_KEY] == "batch_abc" + assert is_batch_line_item_event(logging_obj.model_call_details) + + +def test_hidden_params_carry_batch_line_fields_only_for_batch_line_items(): + plain: Final = StandardLoggingPayloadSetup.get_hidden_params({"model_id": "m"}) + line: Final = StandardLoggingPayloadSetup.get_hidden_params( + {"batch_id": "batch_abc", "batch_custom_id": "r1", "batch_line_status_code": 200} + ) + + assert {"batch_id", "batch_custom_id", "batch_line_status_code"}.isdisjoint(plain) + assert (line["batch_id"], line["batch_custom_id"], line["batch_line_status_code"]) == ("batch_abc", "r1", 200) + + def test_enterprise_alerting_loggers_resolves_real_classes(): from litellm_enterprise.enterprise_callbacks.pagerduty.pagerduty import PagerDutyAlerting from litellm_enterprise.enterprise_callbacks.send_emails.resend_email import ResendEmailLogger @@ -10388,8 +10424,8 @@ def test_get_hidden_params(): assert result["response_cost"] is None assert result["additional_headers"] is None - # assert all fields in StandardLoggingHiddenParams are present - assert all(field in result for field in StandardLoggingHiddenParams.__annotations__) + # assert all required fields in StandardLoggingHiddenParams are present + assert all(field in result for field in StandardLoggingHiddenParams.__required_keys__) # Test with valid params hidden_params = { @@ -10409,8 +10445,8 @@ def test_get_hidden_params(): assert result["response_cost"] == 0.001 assert result["additional_headers"] is not None assert result["additional_headers"]["x_ratelimit_limit_requests"] == 2000 - # assert all fields in StandardLoggingHiddenParams are present - assert all(field in result for field in StandardLoggingHiddenParams.__annotations__) + # assert all required fields in StandardLoggingHiddenParams are present + assert all(field in result for field in StandardLoggingHiddenParams.__required_keys__) @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") def test_get_final_response_obj(): From e00471b86a74acb7ab97efe3095617a727bca651 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Thu, 8 Oct 2026 15:08:46 -0700 Subject: [PATCH 52/89] chore(batches): keep the lint gates within main's ceilings after the merge Typed the helpers main now types (Mapping litellm_params, Any Span), replaced the banned cast with a literal match, dropped isinstance checks main made redundant, and moved test-quality reasons onto the lines the gate reads. --- .../proxy/common_utils/check_batch_cost.py | 2 +- litellm/batches/batch_line_item_logging.py | 14 ++--- litellm/batches/batch_utils.py | 4 +- litellm/litellm_core_utils/core_helpers.py | 24 ++++----- litellm/litellm_core_utils/litellm_logging.py | 14 ++--- litellm/llms/anthropic/chat/transformation.py | 2 +- .../llms/bedrock/batches/transformation.py | 2 +- litellm/proxy/proxy_server.py | 7 +-- .../batches/test_batch_line_item_logging.py | 52 +++++-------------- .../test_batch_line_item_metering_sinks.py | 8 +-- ...test_proxy_server_endpoints_and_startup.py | 2 +- 11 files changed, 51 insertions(+), 80 deletions(-) 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 167f4dac82c6..c885fde6ae46 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -940,7 +940,7 @@ async def _track_completed_batch_cost( batch_failed_requests=batch_result.failed_requests, batch_prompt_cost=batch_result.prompt_cost, batch_completion_cost=batch_result.completion_cost, - batch_output_file_content=content_bytes if isinstance(content_bytes, bytes) else None, + batch_output_file_content=content_bytes, batch_error_file_content=error_file_content, ) except Exception: diff --git a/litellm/batches/batch_line_item_logging.py b/litellm/batches/batch_line_item_logging.py index 9565614f78f2..854f8c2ef34b 100644 --- a/litellm/batches/batch_line_item_logging.py +++ b/litellm/batches/batch_line_item_logging.py @@ -3,7 +3,7 @@ from collections.abc import Iterator, Mapping from datetime import datetime from types import MappingProxyType -from typing import TYPE_CHECKING, Final, Literal, TypeAlias, cast, get_args +from typing import TYPE_CHECKING, Final, Literal, TypeAlias from typing_extensions import assert_never @@ -35,7 +35,6 @@ "openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral" ] -_SUPPORTED_LINE_PROVIDERS: Final = frozenset(get_args(_BatchLineProvider)) _SECRET_PARAM_KEYS: Final = frozenset( { @@ -53,9 +52,11 @@ def _supported_line_provider(value: str) -> _BatchLineProvider | None: - if value in _SUPPORTED_LINE_PROVIDERS: - return cast("_BatchLineProvider", value) # cast-ok: membership in the literal's args was just checked - return None + match value: + case "openai" | "azure" | "vertex_ai" | "hosted_vllm" | "anthropic" | "bedrock" | "mistral": + return value + case _: + return None _CALL_TYPE_BY_BATCH_URL: Final = MappingProxyType( @@ -332,8 +333,7 @@ async def _emit_line_event( status_code: Final = _line_status_code(entry, custom_llm_provider) request_call_type: Final = _call_type_for_request(request_line) response_body: Final = _get_response_from_batch_job_output_file(entry, custom_llm_provider) - parent_start_time: Final = parent.start_time # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # Logging.start_time is untyped upstream - start_time: Final = parent_start_time if isinstance(parent_start_time, datetime) else datetime.now() # noqa: DTZ005 # naive to match the logging pipeline start_time + start_time: Final = parent.start_time parent_params: Final = _as_object_mapping(parent.litellm_params) or _EMPTY_BODY # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # Logging.litellm_params is untyped upstream model: Final = _line_model(response_body, request_body, parent) diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 26ae92526e05..38736b432242 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -139,7 +139,7 @@ async def handle_completed_batch_with_files( batch: Batch, custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral"], model_name: str | None = None, - litellm_params: dict[str, object] | None = None, # mutable-ok: file fetchers require a plain dict + litellm_params: Mapping[str, object] | None = None, model_info: ModelInfo | None = None, ) -> "tuple[BatchCostUsageResult, BatchResultFiles]": """_handle_completed_batch plus the raw output/error bytes it fetched, so @@ -522,7 +522,7 @@ async def _fetch_batch_output_file_content( async def fetch_batch_error_file_content( batch: Batch, custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral"], - litellm_params: dict[str, object] | None, # mutable-ok: file fetchers require a plain dict + litellm_params: Mapping[str, object] | None, ) -> bytes | None: """Fetch the batch's separate error file bytes; None when it has none or the fetch fails.""" if batch.error_file_id is None: diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index 0a18fc19cd10..8ffcc4939dcc 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -19,9 +19,9 @@ from litellm.types.utils import ModelResponseStream - Span = _Span | object + Span = _Span | Any else: - Span = object + Span = Any _CODEX_CLIENT_PREFIX_RE: Final = re.compile(r"^codex[-_ /]", re.IGNORECASE) @@ -271,7 +271,7 @@ def remove_index_from_tool_calls( tool_call.pop("index", None) -def remove_items_at_indices(items: list[object] | None, indices: Iterable[int]) -> None: +def remove_items_at_indices(items: list[Any] | None, indices: Iterable[int]) -> None: """Remove items from a list in-place by index""" if items is None: return @@ -728,9 +728,9 @@ def filter_internal_params(data: dict, additional_internal_params: set | None = def redact_nested_match_and_regex_keys( - payload: dict | list[object] | str | None, + payload: dict | list[Any] | str | None, keys: Collection[str] = ("match", "regex"), -) -> dict | list[object] | str | None: +) -> dict | list[Any] | str | None: """ Deep-copy `payload` and replace every configured string field with "[REDACTED]" anywhere in nested dict/list structures. @@ -796,13 +796,13 @@ def set_response_cost_in_hidden_params(response: _CarriesHiddenParams, cost: flo BATCH_PARENT_ID_KEY: Final = "batch_parent_id" -def is_batch_line_item_event(kwargs: object) -> bool: - if not isinstance(kwargs, Mapping): - return False - litellm_params: Final[object] = kwargs.get("litellm_params") - if not isinstance(litellm_params, Mapping): - return False - return bool(litellm_params.get(BATCH_PARENT_ID_KEY)) +def is_batch_line_item_event(kwargs: Any) -> bool: # noqa: ANN401 # every hook passes its own unannotated callback kwargs + """Whether a callback event is one JSONL line of a completed batch rather than a request""" + match kwargs: + case {"litellm_params": {"batch_parent_id": str() as batch_parent_id}}: + return bool(batch_parent_id) + case _: + return False _HIDDEN_PARAMS_ADAPTER: Final = TypeAdapter(Mapping[str, object]) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 923059d01b98..26d9035e7a24 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -963,7 +963,7 @@ def update_environment_variables( if model is not None: self.model = model self.user = user - marked_batch_parent_id: Final[object] = self.litellm_params.get(BATCH_PARENT_ID_KEY) + marked_batch_parent_id: Final[object] = self.litellm_params.get(BATCH_PARENT_ID_KEY) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # Logging.litellm_params is untyped upstream self.litellm_params = { **self.litellm_params, **scrub_sensitive_keys_in_metadata(litellm_params), @@ -971,9 +971,9 @@ def update_environment_variables( # Every spend, budget and metering hook skips an event carrying batch_parent_id, so only # mark_batch_line_item may set it; the incoming params can hold caller request fields if marked_batch_parent_id is None: - self.litellm_params.pop(BATCH_PARENT_ID_KEY, None) + self.litellm_params.pop(BATCH_PARENT_ID_KEY, None) # pyright: ignore[reportUnknownMemberType] # Logging.litellm_params is untyped upstream else: - self.litellm_params[BATCH_PARENT_ID_KEY] = marked_batch_parent_id + self.litellm_params[BATCH_PARENT_ID_KEY] = marked_batch_parent_id # pyright: ignore[reportUnknownMemberType] # Logging.litellm_params is untyped upstream self.litellm_request_debug = litellm_params.get("litellm_request_debug", False) self.logger_fn = litellm_params.get("logger_fn", None) if is_debugging_on() or self.litellm_request_debug: @@ -1013,7 +1013,7 @@ def update_environment_variables( def mark_batch_line_item(self, batch_id: str) -> None: """Mark this logging object as one JSONL line of the completed batch ``batch_id``""" - self.litellm_params[BATCH_PARENT_ID_KEY] = batch_id + self.litellm_params[BATCH_PARENT_ID_KEY] = batch_id # pyright: ignore[reportUnknownMemberType] # Logging.litellm_params is untyped upstream def update_from_kwargs( self, @@ -3336,7 +3336,7 @@ async def _async_success_handler_body( cost_for_built_in_tools_cost_usd_dollar=0.0, ) if batch_output_file_content is not None or batch_error_file_content is not None: - result_files = BatchResultFiles( + result_files = BatchResultFiles( # rebind-ok: the caller supplied the files output=batch_output_file_content if isinstance(batch_output_file_content, bytes) else None, error=batch_error_file_content if isinstance(batch_error_file_content, bytes) else None, ) @@ -3349,7 +3349,7 @@ async def _async_success_handler_body( litellm_params=self.litellm_params, model_info=self.get_router_deployment_model_info(), ) - result_files = fetched_result_files + result_files = fetched_result_files # rebind-ok: the files this branch fetched set_hidden_param(result, "response_cost", batch_result.cost) set_hidden_param(result, "batch_models", batch_result.models) @@ -6582,7 +6582,7 @@ def _extract_response_obj_and_hidden_params( original_exception, "_hidden_params", None ) if isinstance(exception_hidden_params, dict) and exception_hidden_params: - hidden_params = dict(exception_hidden_params) + hidden_params = dict(exception_hidden_params) # rebind-ok: the exception carries the hidden params elif (response_headers := _get_response_headers(original_exception)) is not None: hidden_params = dict( StandardLoggingHiddenParams( diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index 0633280ebf46..f65cfcb18fa0 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -2887,7 +2887,7 @@ def _valid_user_id(user_id: str) -> bool: def anthropic_message_to_model_response(result: Mapping[str, object], speed: str | None) -> ModelResponse: pydantic_result: Final = AnthropicResponse.model_validate(result) result_id: Final = result.get("id") - return AnthropicConfig().transform_parsed_response( + return AnthropicConfig().transform_parsed_response( # pyright: ignore[reportUnknownMemberType] # completion_response is a bare dict upstream completion_response=pydantic_result.model_dump(), raw_response=httpx.Response( status_code=200, diff --git a/litellm/llms/bedrock/batches/transformation.py b/litellm/llms/bedrock/batches/transformation.py index 51608ad3a05e..136cf8f1b99b 100644 --- a/litellm/llms/bedrock/batches/transformation.py +++ b/litellm/llms/bedrock/batches/transformation.py @@ -118,7 +118,7 @@ def bedrock_batch_line_to_response( if "output" in model_output: from ..chat.converse_transformation import AmazonConverseConfig - return AmazonConverseConfig()._transform_response( # pyright: ignore[reportPrivateUsage] # same reconstruction the converse chat path performs on the live response + return AmazonConverseConfig()._transform_response( # pyright: ignore[reportPrivateUsage, reportUnknownMemberType] # same reconstruction the converse chat path performs on the live response model=model, response=Response(200, json=dict(model_output)), model_response=ModelResponse(), diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 707ab48219b8..f07656bc51de 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -7087,12 +7087,9 @@ async def load_config(self, router: litellm.Router | None, config_file_path: str health_check_concurrency = general_settings.get("health_check_concurrency", None) health_check_details = general_settings.get("health_check_details", True) ### BATCH LINE ITEM CALLBACKS ### - _store_batch_line_items: Final[object] = general_settings.get("store_batch_line_items_in_callbacks") + _store_batch_line_items: Final = coerce_bool(general_settings.get("store_batch_line_items_in_callbacks")) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # general_settings is untyped upstream if _store_batch_line_items is not None: - if isinstance(_store_batch_line_items, str): - litellm.store_batch_line_items_in_callbacks = _store_batch_line_items.lower() == "true" - else: - litellm.store_batch_line_items_in_callbacks = bool(_store_batch_line_items) + litellm.store_batch_line_items_in_callbacks = bool(_store_batch_line_items) # Health-check-driven routing (opt-in, passes through to Router later) _enable_hc_routing = general_settings.get("enable_health_check_routing", False) _hc_staleness = general_settings.get("health_check_staleness_threshold", None) diff --git a/tests/unit/batches/test_batch_line_item_logging.py b/tests/unit/batches/test_batch_line_item_logging.py index 43da11528b03..d383546b7bf3 100644 --- a/tests/unit/batches/test_batch_line_item_logging.py +++ b/tests/unit/batches/test_batch_line_item_logging.py @@ -144,9 +144,7 @@ def recorder(): litellm._async_success_callback ) # test-quality-ok: the feature dispatches through this global list; restored in teardown saved_failure = list(litellm._async_failure_callback) # test-quality-ok: same dispatch seam, restored in teardown - litellm._async_success_callback = [ - logger - ] # test-quality-ok: there is no injection seam for callback lists; teardown restores + litellm._async_success_callback = [logger] # test-quality-ok: no injection seam; teardown restores litellm._async_failure_callback = [logger] # test-quality-ok: same dispatch seam, restored in teardown yield logger litellm.store_batch_line_items_in_callbacks = saved_flag # test-quality-ok: teardown restoring the value set above @@ -195,9 +193,7 @@ def _hidden(event: dict) -> dict: @pytest.mark.asyncio async def test_line_items_emitted_alongside_aggregate(recorder): - litellm.store_batch_line_items_in_callbacks = ( - True # test-quality-ok: the flag under test is a module global; fixture restores it - ) + litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: flag under test; fixture restores it batch = _batch() with ( patch( @@ -253,9 +249,7 @@ async def test_flag_off_emits_only_aggregate(recorder): @pytest.mark.asyncio async def test_line_items_skipped_for_unsupported_provider(recorder): - litellm.store_batch_line_items_in_callbacks = ( - True # test-quality-ok: the flag under test is a module global; fixture restores it - ) + litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: flag under test; fixture restores it file_mock: Final = AsyncMock(side_effect=_file_content) with patch( "litellm.files.main.afile_content", file_mock @@ -271,9 +265,7 @@ async def test_line_items_skipped_for_unsupported_provider(recorder): @pytest.mark.asyncio async def test_in_progress_batch_poll_emits_no_line_events(recorder): - litellm.store_batch_line_items_in_callbacks = ( - True # test-quality-ok: the flag under test is a module global; fixture restores it - ) + litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: flag under test; fixture restores it in_progress: Final = LiteLLMBatch( id="batch_wip", object="batch", @@ -298,9 +290,7 @@ async def test_in_progress_batch_poll_emits_no_line_events(recorder): @pytest.mark.asyncio async def test_input_fetch_failure_still_emits_aggregate(recorder): - litellm.store_batch_line_items_in_callbacks = ( - True # test-quality-ok: the flag under test is a module global; fixture restores it - ) + litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: flag under test; fixture restores it with patch( "litellm.files.main.afile_content", new_callable=AsyncMock, side_effect=ValueError("boom") ): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch @@ -547,9 +537,7 @@ def _parent_logging_with_params(litellm_params: dict) -> Logging: @pytest.mark.asyncio async def test_line_items_edge_shapes_and_edge_cases(recorder): - litellm.store_batch_line_items_in_callbacks = ( - True # test-quality-ok: the flag under test is a module global; fixture restores it - ) + litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: flag under test; fixture restores it batch = LiteLLMBatch( id="batch_edge", object="batch", @@ -611,9 +599,7 @@ async def test_line_items_edge_shapes_and_edge_cases(recorder): @pytest.mark.asyncio async def test_line_items_child_params_drop_parent_credentials(recorder): - litellm.store_batch_line_items_in_callbacks = ( - True # test-quality-ok: the flag under test is a module global; fixture restores it - ) + litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: flag under test; fixture restores it file_mock: Final = AsyncMock(side_effect=_file_content) parent: Final = _parent_logging_with_params( { @@ -648,9 +634,7 @@ async def test_line_items_child_params_drop_parent_credentials(recorder): @pytest.mark.asyncio async def test_line_items_anthropic_shapes(recorder): - litellm.store_batch_line_items_in_callbacks = ( - True # test-quality-ok: the flag under test is a module global; fixture restores it - ) + litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: flag under test; fixture restores it batch = LiteLLMBatch( id="batch_anth", object="batch", @@ -704,9 +688,7 @@ def _provider_batch(batch_id: str, input_file_id: str, output_file_id: str) -> L @pytest.mark.asyncio async def test_line_items_bedrock_shapes(recorder): - litellm.store_batch_line_items_in_callbacks = ( - True # test-quality-ok: the flag under test is a module global; fixture restores it - ) + litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: flag under test; fixture restores it file_mock: Final = AsyncMock(side_effect=_edge_file_content) with ( patch( @@ -751,9 +733,7 @@ async def test_line_items_bedrock_shapes(recorder): @pytest.mark.asyncio async def test_line_items_mistral_shape(recorder): - litellm.store_batch_line_items_in_callbacks = ( - True # test-quality-ok: the flag under test is a module global; fixture restores it - ) + litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: flag under test; fixture restores it file_mock: Final = AsyncMock(side_effect=_edge_file_content) with ( patch( @@ -820,9 +800,7 @@ def _scoped_file_content(file_map): @pytest.mark.asyncio async def test_line_items_anthropic_failure_keeps_provider_error(recorder): - litellm.store_batch_line_items_in_callbacks = ( - True # test-quality-ok: the flag under test is a module global; fixture restores it - ) + litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: flag under test; fixture restores it file_mock: Final = AsyncMock( side_effect=_scoped_file_content( { @@ -860,9 +838,7 @@ async def test_line_items_anthropic_failure_keeps_provider_error(recorder): @pytest.mark.asyncio async def test_line_items_failure_without_error_or_response_body_still_reaches_callbacks(recorder): - litellm.store_batch_line_items_in_callbacks = ( - True # test-quality-ok: the flag under test is a module global; fixture restores it - ) + litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: flag under test; fixture restores it file_mock: Final = AsyncMock( side_effect=_scoped_file_content( { @@ -889,9 +865,7 @@ async def test_line_items_failure_without_error_or_response_body_still_reaches_c @pytest.mark.asyncio async def test_line_items_native_vertex_rows_are_skipped(recorder): - litellm.store_batch_line_items_in_callbacks = ( - True # test-quality-ok: the flag under test is a module global; fixture restores it - ) + litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: flag under test; fixture restores it file_mock: Final = AsyncMock( side_effect=_scoped_file_content( { diff --git a/tests/unit/integrations/test_batch_line_item_metering_sinks.py b/tests/unit/integrations/test_batch_line_item_metering_sinks.py index 525f6ae15d04..9cec64b12736 100644 --- a/tests/unit/integrations/test_batch_line_item_metering_sinks.py +++ b/tests/unit/integrations/test_batch_line_item_metering_sinks.py @@ -149,8 +149,8 @@ async def _log_completed_batch(monkeypatch: pytest.MonkeyPatch, loggers: list, f monkeypatch.setattr(litellm, "store_batch_line_items_in_callbacks", flag, raising=False) saved_success = list(litellm._async_success_callback) saved_failure = list(litellm._async_failure_callback) - litellm._async_success_callback = list(loggers) - litellm._async_failure_callback = list(loggers) + litellm._async_success_callback = list(loggers) # test-quality-ok: no injection seam; finally restores + litellm._async_failure_callback = list(loggers) # test-quality-ok: same dispatch seam buf = io.StringIO() try: with ( @@ -169,8 +169,8 @@ async def _log_completed_batch(monkeypatch: pytest.MonkeyPatch, loggers: list, f batch_completion_cost=0.5, ) finally: - litellm._async_success_callback = saved_success - litellm._async_failure_callback = saved_failure + litellm._async_success_callback = saved_success # test-quality-ok: restores the seam + litellm._async_failure_callback = saved_failure # test-quality-ok: restores the seam def _openmeter_sinks() -> tuple[Any, RecordingHTTP]: diff --git a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py index 038c991040f7..0b0708955bb4 100644 --- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py +++ b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py @@ -8153,7 +8153,7 @@ async def test_update_general_settings_store_batch_line_items_in_callbacks(): await proxy_config._update_general_settings(db_general_settings={}) assert litellm.store_batch_line_items_in_callbacks is False finally: - litellm.store_batch_line_items_in_callbacks = saved_flag + litellm.store_batch_line_items_in_callbacks = saved_flag # test-quality-ok: restores the prior flag @pytest.mark.asyncio From 317764c3a150e120d580c3f789a450e678adfca4 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Thu, 8 Oct 2026 16:13:31 -0700 Subject: [PATCH 53/89] fix(anthropic): read the message id off the validated message The anthropic_messages logging path can hand over an already parsed message object, which has no .get; main read its id through getattr, so take it from the validated AnthropicResponse instead. --- litellm/llms/anthropic/chat/transformation.py | 3 +-- .../test_litellm_logging.py | 26 +++++++++++++++++++ 2 files changed, 27 insertions(+), 2 deletions(-) diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index f65cfcb18fa0..3943cb823fd0 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -2886,14 +2886,13 @@ def _valid_user_id(user_id: str) -> bool: def anthropic_message_to_model_response(result: Mapping[str, object], speed: str | None) -> ModelResponse: pydantic_result: Final = AnthropicResponse.model_validate(result) - result_id: Final = result.get("id") return AnthropicConfig().transform_parsed_response( # pyright: ignore[reportUnknownMemberType] # completion_response is a bare dict upstream completion_response=pydantic_result.model_dump(), raw_response=httpx.Response( status_code=200, headers={}, ), - model_response=ModelResponse(id=result_id if isinstance(result_id, str) and result_id else None), + model_response=ModelResponse(id=pydantic_result.id or None), json_mode=None, speed=speed, ) diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index 07d1b1024b13..9bb6d4ec85d3 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -5806,6 +5806,32 @@ def test_handle_anthropic_messages_response_logging_keeps_the_served_response_id assert result.service_tier == "flex" +@pytest.mark.parametrize("as_model", [False, True], ids=["dict", "pydantic"]) +def test_handle_anthropic_messages_response_logging_keeps_the_message_id_of_a_parsed_message(as_model): + from litellm.types.llms.anthropic import AnthropicResponse + + message: Final = { + "id": "msg_served", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-5", + "content": [{"type": "text", "text": "hi"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 3, "output_tokens": 1}, + } + logging_obj = _anthropic_messages_logging_obj() + logging_obj.update_environment_variables(litellm_params={}, optional_params={}) + + result = logging_obj._handle_anthropic_messages_response_logging( + result=AnthropicResponse.model_validate(message) if as_model else message + ) + + assert isinstance(result, ModelResponse) + assert result.id == "msg_served" + assert result.choices[0].message.content == "hi" # type: ignore[union-attr] + + def test_handle_anthropic_messages_response_logging_passes_model_response_through(): """Anthropic-native path already yields a ModelResponse; it must be returned unchanged.""" logging_obj = _anthropic_messages_logging_obj() From 85fc4498cc5eed967d5ed812a8e0d95f1b0fa330 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Thu, 8 Oct 2026 16:22:04 -0700 Subject: [PATCH 54/89] fix(batches): an unreadable error file no longer drops the output line events Costing counts an error file it could not fetch as no failed lines; line logging refetched it and the raise aborted the whole fan-out, so the successful lines were never emitted. The error file now reads as empty with a warning. --- litellm/batches/batch_line_item_logging.py | 18 +++++++++++- .../batches/test_batch_line_item_logging.py | 29 +++++++++++++++++++ 2 files changed, 46 insertions(+), 1 deletion(-) diff --git a/litellm/batches/batch_line_item_logging.py b/litellm/batches/batch_line_item_logging.py index 854f8c2ef34b..bde63efb30f3 100644 --- a/litellm/batches/batch_line_item_logging.py +++ b/litellm/batches/batch_line_item_logging.py @@ -419,6 +419,22 @@ async def _fetch_managed_file_or_empty( ) +async def _fetch_error_file_or_empty( + batch: LiteLLMBatch, + custom_llm_provider: _BatchLineProvider, + fetch_params: dict[str, object] | None, # mutable-ok: batch_utils file fetch takes the shared litellm_params dict +) -> bytes: + try: + return await _fetch_managed_file_or_empty(batch.error_file_id, custom_llm_provider, fetch_params) + except Exception as e: # noqa: BLE001 # costing counts an unreadable error file as no failed lines, so the output lines still go out + verbose_logger.warning( + "batch line-item callbacks could not fetch the error file, emitting the output lines only. batch_id=%s error=%s", + batch.id, + e, + ) + return b"" + + async def log_batch_line_items( batch: LiteLLMBatch, custom_llm_provider: str, @@ -488,7 +504,7 @@ async def log_batch_line_items( error_content: Final = ( result_files.error if result_files is not None and result_files.error is not None - else await _fetch_managed_file_or_empty(batch.error_file_id, line_provider, fetch_params) + else await _fetch_error_file_or_empty(batch, line_provider, fetch_params) ) for content in (output_content, error_content): for entry in _output_entries(content): diff --git a/tests/unit/batches/test_batch_line_item_logging.py b/tests/unit/batches/test_batch_line_item_logging.py index d383546b7bf3..aea0ae63ebec 100644 --- a/tests/unit/batches/test_batch_line_item_logging.py +++ b/tests/unit/batches/test_batch_line_item_logging.py @@ -1241,6 +1241,35 @@ async def test_supplied_result_files_skip_the_output_and_error_fetch(recorder): assert file_mock.await_args.kwargs["file_id"] == "input-file-1" +def _error_file_unreadable(file_id: str, **_kwargs): + if file_id == "error-file-1": + raise ValueError("error file gone") + return SimpleNamespace(content=_FILE_BYTES[file_id]) + + +@pytest.mark.asyncio +async def test_an_unreadable_error_file_still_emits_the_output_lines(recorder): + file_mock: Final = AsyncMock(side_effect=_error_file_unreadable) + with patch( + "litellm.files.main.afile_content", file_mock + ): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + emitted: Final = await log_batch_line_items( + batch=_batch(), + custom_llm_provider="openai", + parent=_parent_logging(), + model_name="gpt-4o", + litellm_params=None, + model_info=None, + result_files=BatchResultFiles(output=OUTPUT_JSONL, error=None), + claim_cache=DualCache(), + ) + + assert emitted == 1 + assert len(recorder.success_events) == 1 + assert len(recorder.failure_events) == 0 + assert sorted(call.kwargs["file_id"] for call in file_mock.await_args_list) == ["error-file-1", "input-file-1"] + + @pytest.mark.asyncio async def test_empty_result_files_fall_back_to_fetching_everything(recorder): file_mock: Final = AsyncMock(side_effect=_file_content) From 21cd892e3ddb17b1f74b5b9e934cbb292cd18674 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Thu, 8 Oct 2026 16:22:04 -0700 Subject: [PATCH 55/89] test(batches): pin billing for batch_parent_id on every endpoint and value shape Chat completions, embeddings, responses and messages join moderations, speech and transcriptions, and speech is driven with an empty, int, list, object and 5 KB value. --- .../spend/test_batch_parent_id_body_field.py | 123 ++++++++++++++++-- 1 file changed, 115 insertions(+), 8 deletions(-) diff --git a/tests/integration/spend/test_batch_parent_id_body_field.py b/tests/integration/spend/test_batch_parent_id_body_field.py index 46d36020274b..ea93f27075dd 100644 --- a/tests/integration/spend/test_batch_parent_id_body_field.py +++ b/tests/integration/spend/test_batch_parent_id_body_field.py @@ -51,6 +51,89 @@ def _transcription_reply() -> JsonResponse: return JsonResponse(content_type="application/json", body={"text": "hello"}) +def _chat_reply() -> JsonResponse: + return JsonResponse( + content_type="application/json", + body={ + "id": f"chatcmpl-{uuid.uuid4().hex}", + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 9, "completion_tokens": 3, "total_tokens": 12}, + }, + ) + + +def _embedding_reply() -> JsonResponse: + return JsonResponse( + content_type="application/json", + body={ + "object": "list", + "data": [{"object": "embedding", "embedding": [0.1, 0.2], "index": 0}], + "model": "text-embedding-3-small", + "usage": {"prompt_tokens": 7, "total_tokens": 7}, + }, + ) + + +def _responses_reply() -> JsonResponse: + return JsonResponse( + content_type="application/json", + body={ + "id": f"resp_{uuid.uuid4().hex}", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-4o-mini", + "output": [], + "usage": {"input_tokens": 9, "output_tokens": 3, "total_tokens": 12}, + }, + ) + + +def _messages_reply() -> JsonResponse: + return JsonResponse( + content_type="application/json", + body={ + "id": f"msg_{uuid.uuid4().hex}", + "type": "message", + "role": "assistant", + "model": "claude-haiku-4-5", + "content": [{"type": "text", "text": "ok"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 9, "output_tokens": 3}, + }, + ) + + +def _chat(gateway: Gateway, model: str, key: str, extra: dict[str, JsonValue]) -> httpx.Response: + return gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": "hi"}], **extra}, + key=key, + ) + + +def _embedding(gateway: Gateway, model: str, key: str, extra: dict[str, JsonValue]) -> httpx.Response: + return gateway.request("POST", "/v1/embeddings", {"model": model, "input": "hi", **extra}, key=key) + + +def _responses(gateway: Gateway, model: str, key: str, extra: dict[str, JsonValue]) -> httpx.Response: + return gateway.request("POST", "/v1/responses", {"model": model, "input": "hi", **extra}, key=key) + + +def _messages(gateway: Gateway, model: str, key: str, extra: dict[str, JsonValue]) -> httpx.Response: + return gateway.request( + "POST", + "/v1/messages", + {"model": model, "max_tokens": 16, "messages": [{"role": "user", "content": "hi"}], **extra}, + key=key, + ) + + def _moderation(gateway: Gateway, model: str, key: str, extra: dict[str, JsonValue]) -> httpx.Response: return gateway.request("POST", "/v1/moderations", {"model": model, "input": "safe text", **extra}, key=key) @@ -72,6 +155,17 @@ def _transcription(gateway: Gateway, model: str, key: str, extra: dict[str, Json pytest.param("openai/omni-moderation-latest", _moderation_reply, _moderation, "amoderation", id="moderations"), pytest.param("openai/tts-1", _speech_reply, _speech, "aspeech", id="audio-speech"), pytest.param("openai/whisper-1", _transcription_reply, _transcription, "atranscription", id="audio-transcriptions"), + pytest.param("openai/gpt-4o-mini", _chat_reply, _chat, "acompletion", id="chat-completions"), + pytest.param("openai/text-embedding-3-small", _embedding_reply, _embedding, "aembedding", id="embeddings"), + pytest.param("openai/gpt-4o-mini", _responses_reply, _responses, "aresponses", id="responses"), + pytest.param("anthropic/claude-haiku-4-5", _messages_reply, _messages, "anthropic_messages", id="messages"), +) +FORGED_VALUES: Final = ( + pytest.param("", id="empty-string"), + pytest.param(5, id="int"), + pytest.param(["batch_forged_by_caller"], id="list"), + pytest.param({"id": "batch_forged_by_caller"}, id="object"), + pytest.param("b" * 5000, id="5kb-string"), ) @@ -99,13 +193,8 @@ def _key_spend(key: str) -> float: return float(str(rows[0]["spend"])) -@pytest.mark.parametrize(("provider_model", "reply", "send", "call_type"), ROUTES) -def test_a_caller_supplied_batch_parent_id_is_still_tracked_like_any_request( - gateway: Gateway, - provider_model: str, - reply: _Reply, - send: _Send, - call_type: str, +def _assert_forged_request_bills_like_control( + gateway: Gateway, provider_model: str, reply: _Reply, send: _Send, call_type: str, forged_value: JsonValue ) -> None: with gateway.scenario() as scenario: control_model: Final = _scripted_model(gateway, scenario, provider_model, reply) @@ -119,9 +208,27 @@ def test_a_caller_supplied_batch_parent_id_is_still_tracked_like_any_request( control_spend: Final = float(str(control_rows[0]["spend"])) eventually(lambda: _key_spend(key), lambda value: value == pytest.approx(control_spend), seconds=70) - forged: Final = send(gateway, forged_model, key, {"batch_parent_id": FORGED_PARENT}) + forged: Final = send(gateway, forged_model, key, {"batch_parent_id": forged_value}) assert forged.status_code == 200, forged.text rows: Final = eventually(lambda: _spend_rows(key), lambda values: len(values) == 2, seconds=70) assert [row["call_type"] for row in rows] == [call_type, call_type], rows assert float(str(rows[1]["spend"])) == pytest.approx(control_spend), rows eventually(lambda: _key_spend(key), lambda value: value == pytest.approx(2 * control_spend), seconds=70) + + +@pytest.mark.parametrize(("provider_model", "reply", "send", "call_type"), ROUTES) +def test_a_caller_supplied_batch_parent_id_is_still_tracked_like_any_request( + gateway: Gateway, + provider_model: str, + reply: _Reply, + send: _Send, + call_type: str, +) -> None: + _assert_forged_request_bills_like_control(gateway, provider_model, reply, send, call_type, FORGED_PARENT) + + +@pytest.mark.parametrize("forged_value", FORGED_VALUES) +def test_any_shape_of_batch_parent_id_on_a_speech_request_is_still_tracked( + gateway: Gateway, forged_value: JsonValue +) -> None: + _assert_forged_request_bills_like_control(gateway, "openai/tts-1", _speech_reply, _speech, "aspeech", forged_value) From 7cf3024ae6c1fb728403b78e361be7bdac791e8c Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Thu, 8 Oct 2026 16:22:38 -0700 Subject: [PATCH 56/89] test(batches): bill every request of a concurrent burst mixing forged and plain bodies --- .../spend/test_batch_parent_id_body_field.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/tests/integration/spend/test_batch_parent_id_body_field.py b/tests/integration/spend/test_batch_parent_id_body_field.py index ea93f27075dd..c5502e4957d9 100644 --- a/tests/integration/spend/test_batch_parent_id_body_field.py +++ b/tests/integration/spend/test_batch_parent_id_body_field.py @@ -5,6 +5,7 @@ import uuid import wave from collections.abc import Callable +from concurrent.futures import ThreadPoolExecutor from hashlib import sha256 from typing import Final @@ -232,3 +233,21 @@ def test_any_shape_of_batch_parent_id_on_a_speech_request_is_still_tracked( gateway: Gateway, forged_value: JsonValue ) -> None: _assert_forged_request_bills_like_control(gateway, "openai/tts-1", _speech_reply, _speech, "aspeech", forged_value) + + +def test_a_concurrent_burst_mixing_forged_and_plain_requests_bills_every_request(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + models: Final = tuple( + _scripted_model(gateway, scenario, "openai/omni-moderation-latest", _moderation_reply) for _ in range(20) + ) + key: Final = scenario.key(models=list(models)) + extras: Final = tuple({"batch_parent_id": FORGED_PARENT} if index % 2 else {} for index in range(20)) + + with ThreadPoolExecutor(max_workers=10) as pool: + responses: Final = tuple( + pool.map(lambda pair: _moderation(gateway, pair[0], key, pair[1]), zip(models, extras)) + ) + + assert [response.status_code for response in responses] == [200] * 20, [r.text for r in responses] + rows: Final = eventually(lambda: _spend_rows(key), lambda values: len(values) == 20, seconds=90) + assert {row["call_type"] for row in rows} == {"amoderation"}, rows From 4bd5762784a9e0a9602377131940c2a58403c42a Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Thu, 8 Oct 2026 16:30:10 -0700 Subject: [PATCH 57/89] revert: keep the line-item fan-out all-or-nothing on an unreadable error file 85fc4498cc emitted the output lines and kept the batch claim when the error file could not be read, so its failed lines were never retried. Raising releases the claim and the next retrieve retries the whole fan-out. --- litellm/batches/batch_line_item_logging.py | 18 +----------- .../batches/test_batch_line_item_logging.py | 29 ------------------- 2 files changed, 1 insertion(+), 46 deletions(-) diff --git a/litellm/batches/batch_line_item_logging.py b/litellm/batches/batch_line_item_logging.py index bde63efb30f3..854f8c2ef34b 100644 --- a/litellm/batches/batch_line_item_logging.py +++ b/litellm/batches/batch_line_item_logging.py @@ -419,22 +419,6 @@ async def _fetch_managed_file_or_empty( ) -async def _fetch_error_file_or_empty( - batch: LiteLLMBatch, - custom_llm_provider: _BatchLineProvider, - fetch_params: dict[str, object] | None, # mutable-ok: batch_utils file fetch takes the shared litellm_params dict -) -> bytes: - try: - return await _fetch_managed_file_or_empty(batch.error_file_id, custom_llm_provider, fetch_params) - except Exception as e: # noqa: BLE001 # costing counts an unreadable error file as no failed lines, so the output lines still go out - verbose_logger.warning( - "batch line-item callbacks could not fetch the error file, emitting the output lines only. batch_id=%s error=%s", - batch.id, - e, - ) - return b"" - - async def log_batch_line_items( batch: LiteLLMBatch, custom_llm_provider: str, @@ -504,7 +488,7 @@ async def log_batch_line_items( error_content: Final = ( result_files.error if result_files is not None and result_files.error is not None - else await _fetch_error_file_or_empty(batch, line_provider, fetch_params) + else await _fetch_managed_file_or_empty(batch.error_file_id, line_provider, fetch_params) ) for content in (output_content, error_content): for entry in _output_entries(content): diff --git a/tests/unit/batches/test_batch_line_item_logging.py b/tests/unit/batches/test_batch_line_item_logging.py index aea0ae63ebec..d383546b7bf3 100644 --- a/tests/unit/batches/test_batch_line_item_logging.py +++ b/tests/unit/batches/test_batch_line_item_logging.py @@ -1241,35 +1241,6 @@ async def test_supplied_result_files_skip_the_output_and_error_fetch(recorder): assert file_mock.await_args.kwargs["file_id"] == "input-file-1" -def _error_file_unreadable(file_id: str, **_kwargs): - if file_id == "error-file-1": - raise ValueError("error file gone") - return SimpleNamespace(content=_FILE_BYTES[file_id]) - - -@pytest.mark.asyncio -async def test_an_unreadable_error_file_still_emits_the_output_lines(recorder): - file_mock: Final = AsyncMock(side_effect=_error_file_unreadable) - with patch( - "litellm.files.main.afile_content", file_mock - ): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch - emitted: Final = await log_batch_line_items( - batch=_batch(), - custom_llm_provider="openai", - parent=_parent_logging(), - model_name="gpt-4o", - litellm_params=None, - model_info=None, - result_files=BatchResultFiles(output=OUTPUT_JSONL, error=None), - claim_cache=DualCache(), - ) - - assert emitted == 1 - assert len(recorder.success_events) == 1 - assert len(recorder.failure_events) == 0 - assert sorted(call.kwargs["file_id"] for call in file_mock.await_args_list) == ["error-file-1", "input-file-1"] - - @pytest.mark.asyncio async def test_empty_result_files_fall_back_to_fetching_everything(recorder): file_mock: Final = AsyncMock(side_effect=_file_content) From 4cc9f561bc2306fa98518c5083e19090ab968d47 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Thu, 8 Oct 2026 16:31:54 -0700 Subject: [PATCH 58/89] test(batches): pin the all-or-nothing retry and warm the burst model first An unreadable error file emits no lines and the next retrieve emits all of them. The burst cell drives one speech model after a warm-up request so a cold worker cannot fail it, and waits for the batched key spend. --- .../spend/test_batch_parent_id_body_field.py | 19 ++++++----- .../batches/test_batch_line_item_logging.py | 34 +++++++++++++++++++ 2 files changed, 44 insertions(+), 9 deletions(-) diff --git a/tests/integration/spend/test_batch_parent_id_body_field.py b/tests/integration/spend/test_batch_parent_id_body_field.py index c5502e4957d9..d99d68d033f6 100644 --- a/tests/integration/spend/test_batch_parent_id_body_field.py +++ b/tests/integration/spend/test_batch_parent_id_body_field.py @@ -237,17 +237,18 @@ def test_any_shape_of_batch_parent_id_on_a_speech_request_is_still_tracked( def test_a_concurrent_burst_mixing_forged_and_plain_requests_bills_every_request(gateway: Gateway) -> None: with gateway.scenario() as scenario: - models: Final = tuple( - _scripted_model(gateway, scenario, "openai/omni-moderation-latest", _moderation_reply) for _ in range(20) - ) - key: Final = scenario.key(models=list(models)) + model: Final = _scripted_model(gateway, scenario, "openai/tts-1", _speech_reply) + key: Final = scenario.key(models=[model]) + warm_up: Final = _speech(gateway, model, key, {}) + assert warm_up.status_code == 200, warm_up.text + eventually(lambda: _spend_rows(key), lambda values: len(values) == 1, seconds=70) extras: Final = tuple({"batch_parent_id": FORGED_PARENT} if index % 2 else {} for index in range(20)) with ThreadPoolExecutor(max_workers=10) as pool: - responses: Final = tuple( - pool.map(lambda pair: _moderation(gateway, pair[0], key, pair[1]), zip(models, extras)) - ) + responses: Final = tuple(pool.map(lambda extra: _speech(gateway, model, key, extra), extras)) assert [response.status_code for response in responses] == [200] * 20, [r.text for r in responses] - rows: Final = eventually(lambda: _spend_rows(key), lambda values: len(values) == 20, seconds=90) - assert {row["call_type"] for row in rows} == {"amoderation"}, rows + rows: Final = eventually(lambda: _spend_rows(key), lambda values: len(values) == 21, seconds=90) + assert {row["call_type"] for row in rows} == {"aspeech"}, rows + expected_spend: Final = 21 * float(str(rows[0]["spend"])) + eventually(lambda: _key_spend(key), lambda value: value == pytest.approx(expected_spend), seconds=90) diff --git a/tests/unit/batches/test_batch_line_item_logging.py b/tests/unit/batches/test_batch_line_item_logging.py index d383546b7bf3..f014c6e6fdea 100644 --- a/tests/unit/batches/test_batch_line_item_logging.py +++ b/tests/unit/batches/test_batch_line_item_logging.py @@ -1241,6 +1241,40 @@ async def test_supplied_result_files_skip_the_output_and_error_fetch(recorder): assert file_mock.await_args.kwargs["file_id"] == "input-file-1" +def _error_file_unreadable(file_id: str, **_kwargs): + if file_id == "error-file-1": + raise ValueError("error file gone") + return SimpleNamespace(content=_FILE_BYTES[file_id]) + + +@pytest.mark.asyncio +async def test_an_unreadable_error_file_defers_every_line_to_the_next_retrieve(recorder): + claim_cache: Final = DualCache() + batch: Final = _batch() + parent: Final = _parent_logging() + emitted: list[int] = [] # mutable-ok: collects one count per retrieve + for side_effect in (_error_file_unreadable, _file_content): + with patch( + "litellm.files.main.afile_content", AsyncMock(side_effect=side_effect) + ): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + emitted.append( + await log_batch_line_items( + batch=batch, + custom_llm_provider="openai", + parent=parent, + model_name="gpt-4o", + litellm_params=None, + model_info=None, + result_files=BatchResultFiles(output=OUTPUT_JSONL, error=None), + claim_cache=claim_cache, + ) + ) + + assert emitted == [0, 2] + assert len(recorder.success_events) == 1 + assert len(recorder.failure_events) == 1 + + @pytest.mark.asyncio async def test_empty_result_files_fall_back_to_fetching_everything(recorder): file_mock: Final = AsyncMock(side_effect=_file_content) From 07d543cf00bfcb1e73fcd3d2a28df480780e17d1 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Fri, 9 Oct 2026 16:32:54 -0700 Subject: [PATCH 59/89] fix(batches): run the line fan-out outside the aggregate event's deadline The logging worker cancels a callback coroutine after its per-callback deadline, and the line fan-out ran inside the aggregate aretrieve_batch event before its callbacks, so a large batch or a slow sink could cancel the aggregate billing event. The fan-out now runs in its own task, built from a copy of the parent's litellm_params taken before the aggregate callbacks write their request-local markers. Anthropic keeps no input file and reports input_file_id "None", which the fan-out fetched and failed on, so no Anthropic line ever arrived. That id is now treated as no input file. --- litellm/batches/batch_line_item_logging.py | 68 ++++++++++++- litellm/litellm_core_utils/litellm_logging.py | 6 +- .../batches/test_batch_line_item_logging.py | 95 +++++++++++++++++++ .../test_batch_line_item_metering_sinks.py | 4 +- .../test_litellm_logging.py | 6 +- 5 files changed, 170 insertions(+), 9 deletions(-) diff --git a/litellm/batches/batch_line_item_logging.py b/litellm/batches/batch_line_item_logging.py index 854f8c2ef34b..e0a44b40ddf2 100644 --- a/litellm/batches/batch_line_item_logging.py +++ b/litellm/batches/batch_line_item_logging.py @@ -1,3 +1,4 @@ +import asyncio import json import uuid from collections.abc import Iterator, Mapping @@ -73,6 +74,12 @@ def _supported_line_provider(value: str) -> _BatchLineProvider | None: batch_line_item_claim_cache: Final = DualCache() +# asyncio holds tasks weakly, so each spawned line fan-out stays referenced here until it finishes +_LINE_ITEM_TASKS: Final[set["asyncio.Task[int]"]] = set() # mutable-ok: strong refs to running fan-outs + +# Anthropic keeps no input file, so its retrieved batch reports input_file_id as the string "None" +_NO_FILE_IDS: Final = frozenset(("", "None")) + _ClaimResult: TypeAlias = Literal["claimed", "already_claimed", "unavailable"] @@ -324,6 +331,7 @@ async def _emit_line_event( batch: LiteLLMBatch, custom_llm_provider: _BatchLineProvider, parent: "Logging", + parent_params: Mapping[str, object], model_name: str | None, model_info: ModelInfo | None, ) -> bool: @@ -334,7 +342,6 @@ async def _emit_line_event( request_call_type: Final = _call_type_for_request(request_line) response_body: Final = _get_response_from_batch_job_output_file(entry, custom_llm_provider) start_time: Final = parent.start_time - parent_params: Final = _as_object_mapping(parent.litellm_params) or _EMPTY_BODY # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # Logging.litellm_params is untyped upstream model: Final = _line_model(response_body, request_body, parent) successful: Final = _batch_response_was_successful(entry, custom_llm_provider) @@ -410,7 +417,7 @@ async def _fetch_managed_file_or_empty( custom_llm_provider: _BatchLineProvider, fetch_params: dict[str, object] | None, # mutable-ok: batch_utils file fetch takes the shared litellm_params dict ) -> bytes: - if file_id is None: + if file_id is None or file_id in _NO_FILE_IDS: return b"" return await _fetch_batch_managed_file_content( file_id, @@ -419,6 +426,57 @@ async def _fetch_managed_file_or_empty( ) +def spawn_batch_line_items( + batch: LiteLLMBatch, + custom_llm_provider: str, + parent: "Logging", + model_name: str | None, + litellm_params: dict[str, object] | None, # mutable-ok: the logging object's shared litellm_params dict + model_info: ModelInfo | None, + result_files: BatchResultFiles | None, +) -> "asyncio.Task[int]": + """Run ``log_batch_line_items`` in its own task, so the fan-out never shares + the logging worker deadline the aggregate aretrieve_batch event runs under. + The parent's litellm_params are copied first, because the aggregate event's + callbacks keep writing request-local markers into its metadata.""" + task: Final = asyncio.get_running_loop().create_task( + log_batch_line_items( + batch=batch, + custom_llm_provider=custom_llm_provider, + parent=parent, + model_name=model_name, + litellm_params=litellm_params, + model_info=model_info, + result_files=result_files, + parent_params=parent_params_snapshot(parent), + ) + ) + _LINE_ITEM_TASKS.add(task) + task.add_done_callback(_LINE_ITEM_TASKS.discard) + task.add_done_callback(_log_line_item_task_failure) + return task + + +def parent_params_snapshot(parent: "Logging") -> Mapping[str, object]: + """The parent's litellm_params with their own copy of the metadata dict.""" + params: Final = _as_object_mapping(parent.litellm_params) or _EMPTY_BODY # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # Logging.litellm_params is untyped upstream + return MappingProxyType({**params, "metadata": _metadata_copy(params)}) + + +def pending_batch_line_item_tasks() -> tuple["asyncio.Task[int]", ...]: + """The spawned line fan-outs that have not finished yet.""" + return tuple(_LINE_ITEM_TASKS) + + +def _log_line_item_task_failure(task: "asyncio.Task[int]") -> None: + if task.cancelled(): + verbose_logger.warning("batch line item logging was cancelled before it finished") + return + error: Final = task.exception() + if error is not None: + verbose_logger.error("batch line item logging failed; aggregate logging unaffected", exc_info=error) + + async def log_batch_line_items( batch: LiteLLMBatch, custom_llm_provider: str, @@ -428,6 +486,7 @@ async def log_batch_line_items( model_info: ModelInfo | None, result_files: BatchResultFiles | None = None, claim_cache: DualCache = batch_line_item_claim_cache, + parent_params: Mapping[str, object] | None = None, ) -> int: """Emit one callback event per JSONL line of a completed batch (request paired with its response/error), behind the opt-in @@ -436,7 +495,9 @@ async def log_batch_line_items( ``batch_parent_id`` and never update spend themselves. Any failure here is logged and swallowed: aggregate accounting must be unaffected. ``result_files`` carries the output/error bytes the aggregate path already - fetched, so they are reused instead of refetched.""" + fetched, so they are reused instead of refetched. ``parent_params`` is the + parent's litellm_params as of the aggregate event, read now when omitted.""" + line_params: Final = parent_params if parent_params is not None else parent_params_snapshot(parent) line_provider: Final = _supported_line_provider(custom_llm_provider) if line_provider is None: verbose_logger.warning( @@ -499,6 +560,7 @@ async def log_batch_line_items( batch=batch, custom_llm_provider=line_provider, parent=parent, + parent_params=line_params, model_name=model_name, model_info=model_info, ) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 26d9035e7a24..29156c3561b8 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -3364,13 +3364,13 @@ async def _async_success_handler_body( ) if litellm.store_batch_line_items_in_callbacks and (has_explicit_batch_data or should_compute_batch_data): - from litellm.batches.batch_line_item_logging import log_batch_line_items + from litellm.batches.batch_line_item_logging import spawn_batch_line_items _line_item_litellm_params: Final = cast( # cast-ok: shared attribute is an untyped dict "dict[str, object] | None", self.litellm_params ) try: - await log_batch_line_items( + spawn_batch_line_items( batch=result, custom_llm_provider=self.custom_llm_provider, parent=self, @@ -3379,7 +3379,7 @@ async def _async_success_handler_body( model_info=self.get_router_deployment_model_info(), result_files=result_files, ) - except Exception: # noqa: BLE001 # line-item logging (claim step included) must never reach the aggregate aretrieve_batch path + except Exception: # noqa: BLE001 # starting the line fan-out must never break the aggregate aretrieve_batch path verbose_logger.exception( "batch line item logging failed for batch_id=%s; aggregate logging unaffected", result.id, diff --git a/tests/unit/batches/test_batch_line_item_logging.py b/tests/unit/batches/test_batch_line_item_logging.py index f014c6e6fdea..e48c1b941b03 100644 --- a/tests/unit/batches/test_batch_line_item_logging.py +++ b/tests/unit/batches/test_batch_line_item_logging.py @@ -11,6 +11,7 @@ cost, or error propagation fails here. """ +import asyncio import json import uuid from datetime import datetime @@ -25,6 +26,7 @@ _release_line_item_claim, batch_line_item_claim_cache, log_batch_line_items, + pending_batch_line_item_tasks, ) from litellm.batches.batch_utils import BatchResultFiles from litellm.caching.caching import DualCache @@ -181,6 +183,7 @@ async def _log_completed_batch(logging_obj: Logging, batch: LiteLLMBatch) -> Non batch_prompt_cost=1.0, batch_completion_cost=0.5, ) + await asyncio.gather(*pending_batch_line_item_tasks()) def _payload(event: dict) -> dict: @@ -288,6 +291,65 @@ async def test_in_progress_batch_poll_emits_no_line_events(recorder): assert len(recorder.failure_events) == 0 +class _GatedLineLogger(CustomLogger): + """Delivers the aggregate event at once and holds every line event until the gate opens.""" + + def __init__(self) -> None: + super().__init__() + self.gate = asyncio.Event() + self.aggregate_events: list[dict] = [] + self.line_events: list[dict] = [] + + async def _record(self, kwargs: dict) -> None: + if kwargs["litellm_params"].get("batch_parent_id") is None: + self.aggregate_events.append(kwargs) + return + await self.gate.wait() + self.line_events.append(kwargs) + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + await self._record(kwargs) + + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + await self._record(kwargs) + + +@pytest.mark.asyncio +async def test_a_slow_line_fan_out_does_not_hold_back_the_aggregate_event(recorder, monkeypatch): + """The logging worker cancels a callback coroutine after its deadline, so the + aggregate event, which bills the batch, must not wait on the line fan-out.""" + litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: flag under test; fixture restores it + gated: Final = _GatedLineLogger() + monkeypatch.setattr(litellm, "_async_success_callback", [gated]) + monkeypatch.setattr(litellm, "_async_failure_callback", [gated]) + with ( + patch( + "litellm.files.main.afile_content", new_callable=AsyncMock, side_effect=_file_content + ), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + patch( + "litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02) + ), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch + ): + await asyncio.wait_for( + _parent_logging().async_success_handler( + result=_batch(), + batch_cost=1.5, + batch_usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15), + batch_models=["gpt-4o"], + batch_successful_requests=1, + batch_failed_requests=1, + ), + timeout=5, + ) + assert [_payload(e)["response_cost"] for e in gated.aggregate_events] == [1.5] + assert gated.line_events == [] + + gated.gate.set() + await asyncio.gather(*pending_batch_line_item_tasks()) + + assert sorted(_hidden(e)["batch_custom_id"] for e in gated.line_events) == ["a", "b"] + + @pytest.mark.asyncio async def test_input_fetch_failure_still_emits_aggregate(recorder): litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: flag under test; fixture restores it @@ -672,6 +734,39 @@ async def test_line_items_anthropic_shapes(recorder): assert payload["completion_tokens"] == 2 +@pytest.mark.asyncio +async def test_anthropic_line_items_arrive_without_an_input_file(recorder): + """Anthropic keeps no input file, so a retrieved Anthropic batch reports input_file_id "None".""" + litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: flag under test; fixture restores it + batch = LiteLLMBatch( + id="batch_anth_no_input", + object="batch", + endpoint="/v1/messages", + input_file_id="None", + output_file_id="output-anth", + error_file_id=None, + status="completed", + completion_window="24h", + created_at=1, + ) + file_mock: Final = AsyncMock(side_effect=_edge_file_content) + logging_obj = _parent_logging(custom_llm_provider="anthropic") + with ( + patch( + "litellm.files.main.afile_content", file_mock + ), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + patch( + "litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02) + ), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch + ): + await _log_completed_batch(logging_obj, batch) + + assert "None" not in [call.kwargs["file_id"] for call in file_mock.await_args_list] + line = next(e for e in recorder.success_events if _hidden(e).get("batch_custom_id") == "b2") + assert line["litellm_params"]["batch_parent_id"] == batch.id + assert _payload(line)["response"]["choices"][0]["message"]["content"] == "hello b2" + + def _provider_batch(batch_id: str, input_file_id: str, output_file_id: str) -> LiteLLMBatch: return LiteLLMBatch( id=batch_id, diff --git a/tests/unit/integrations/test_batch_line_item_metering_sinks.py b/tests/unit/integrations/test_batch_line_item_metering_sinks.py index 9cec64b12736..9d658a1bbc5f 100644 --- a/tests/unit/integrations/test_batch_line_item_metering_sinks.py +++ b/tests/unit/integrations/test_batch_line_item_metering_sinks.py @@ -14,6 +14,7 @@ the guard belongs on cost/token metrics, not on tracing. """ +import asyncio import io import json import os @@ -28,7 +29,7 @@ import pytest import litellm -from litellm.batches.batch_line_item_logging import batch_line_item_claim_cache +from litellm.batches.batch_line_item_logging import batch_line_item_claim_cache, pending_batch_line_item_tasks from litellm.litellm_core_utils.litellm_logging import Logging from litellm.types.utils import LiteLLMBatch, Usage @@ -168,6 +169,7 @@ async def _log_completed_batch(monkeypatch: pytest.MonkeyPatch, loggers: list, f batch_prompt_cost=1.0, batch_completion_cost=0.5, ) + await asyncio.gather(*pending_batch_line_item_tasks()) finally: litellm._async_success_callback = saved_success # test-quality-ok: restores the seam litellm._async_failure_callback = saved_failure # test-quality-ok: restores the seam diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index 234c68f96edc..dd5c8e0d717c 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -9697,7 +9697,7 @@ def _batch(): @pytest.mark.asyncio async def test_compute_path_fetches_each_result_file_once(self, monkeypatch) -> None: - from litellm.batches.batch_line_item_logging import batch_line_item_claim_cache + from litellm.batches.batch_line_item_logging import batch_line_item_claim_cache, pending_batch_line_item_tasks batch_line_item_claim_cache.in_memory_cache.flush_cache() file_bytes = self._batch_file_bytes() @@ -9713,6 +9713,7 @@ async def fake_afile_content(**kwargs): monkeypatch.setattr("litellm.cost_calculator.batch_cost_calculator", lambda **kw: (0.01, 0.02)) await self._logging_obj()._async_success_handler_body(result=self._batch(), start_time=None, end_time=None) + await asyncio.gather(*pending_batch_line_item_tasks()) fetched = sorted(call.kwargs["file_id"] for call in file_mock.await_args_list) assert fetched == ["file-err", "file-in", "file-out"], ( @@ -9721,7 +9722,7 @@ async def fake_afile_content(**kwargs): @pytest.mark.asyncio async def test_explicit_kwargs_path_forwards_result_files(self, monkeypatch) -> None: - from litellm.batches.batch_line_item_logging import batch_line_item_claim_cache + from litellm.batches.batch_line_item_logging import batch_line_item_claim_cache, pending_batch_line_item_tasks batch_line_item_claim_cache.in_memory_cache.flush_cache() from litellm.types.utils import Usage @@ -9747,6 +9748,7 @@ async def fake_afile_content(**kwargs): batch_output_file_content=file_bytes["file-out"], batch_error_file_content=file_bytes["file-err"], ) + await asyncio.gather(*pending_batch_line_item_tasks()) assert file_mock.await_count == 1 assert file_mock.await_args.kwargs["file_id"] == "file-in", ( From 6e826b463b2007f24ca3bd0495a3cf5c0f0792ca Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Fri, 9 Oct 2026 17:23:44 -0700 Subject: [PATCH 60/89] fix(batches): let running line fan-outs finish before proxy shutdown closes their clients A started fan-out keeps its batch claim, so a shutdown that cut it short skipped the unsent lines until the claim expired. proxy_shutdown_event now waits up to MAX_TIME_TO_CLEAR_QUEUE for them first. Also replace a test's type: ignore with an isinstance check. --- litellm/batches/batch_line_item_logging.py | 19 +++++ litellm/proxy/proxy_server.py | 6 ++ .../batches/test_batch_line_item_logging.py | 76 ++++++++++++++++--- .../test_litellm_logging.py | 4 +- .../unit/proxy/proxy_server/test_lifecycle.py | 32 ++++++++ 5 files changed, 125 insertions(+), 12 deletions(-) diff --git a/litellm/batches/batch_line_item_logging.py b/litellm/batches/batch_line_item_logging.py index e0a44b40ddf2..aad1a0a7c6a2 100644 --- a/litellm/batches/batch_line_item_logging.py +++ b/litellm/batches/batch_line_item_logging.py @@ -468,6 +468,25 @@ def pending_batch_line_item_tasks() -> tuple["asyncio.Task[int]", ...]: return tuple(_LINE_ITEM_TASKS) +async def drain_batch_line_items(timeout: float) -> int: + """Wait up to ``timeout`` seconds for this loop's running line fan-outs, so a + shutdown lets them finish before it closes the clients they write through. + Returns how many are still running when the wait ends.""" + loop: Final = asyncio.get_running_loop() + running: Final = [task for task in pending_batch_line_item_tasks() if task.get_loop() is loop and not task.done()] + if not running: + return 0 + _, still_running = await asyncio.wait(running, timeout=timeout) + if still_running: + verbose_logger.warning( + "%d batch line item fan-out(s) still running after waiting %ss at shutdown; lines they have not " + "sent yet are skipped until the batch claim expires", + len(still_running), + timeout, + ) + return len(still_running) + + def _log_line_item_task_failure(task: "asyncio.Task[int]") -> None: if task.cancelled(): verbose_logger.warning("batch line item logging was cancelled before it finished") diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 0e13bc6d8808..63fe2bf090b6 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -73,6 +73,7 @@ LITELLM_SETTINGS_SAFE_DB_OVERRIDES, LITELLM_UI_ALLOW_HEADERS, LITELLM_UI_SESSION_DURATION, + MAX_TIME_TO_CLEAR_QUEUE, RUNTIME_UPDATABLE_ROUTER_SETTINGS, ) from litellm.litellm_core_utils.asyncify import asyncify @@ -274,6 +275,7 @@ def generate_feedback_box(): from litellm.batches.batch_line_item_logging import ( _LINE_ITEM_CLAIM_TTL_SECONDS, # pyright: ignore[reportPrivateUsage] # the claim window is defined next to the cache it expires batch_line_item_claim_cache, + drain_batch_line_items, ) from litellm.caching.caching import DualCache, RedisCache from litellm.caching.dual_cache import DeclaredBatchRead @@ -1195,6 +1197,10 @@ async def _flush_spend_logs_queue_on_shutdown() -> None: async def proxy_shutdown_event(worker_heartbeat: ProxyWorkerHeartbeat | None = None) -> None: global prisma_client, master_key, user_custom_auth, user_custom_key_generate, user_custom_key_update verbose_proxy_logger.info("Shutting down LiteLLM Proxy Server") + # Batch line fan-outs run outside the logging queue and keep their batch claim + # once started, so let them finish while the claim cache, database and + # callback clients are still open. + await drain_batch_line_items(timeout=MAX_TIME_TO_CLEAR_QUEUE) if worker_heartbeat is not None and prisma_client: await worker_heartbeat.deregister() if prisma_client: diff --git a/tests/unit/batches/test_batch_line_item_logging.py b/tests/unit/batches/test_batch_line_item_logging.py index e48c1b941b03..149620e638ea 100644 --- a/tests/unit/batches/test_batch_line_item_logging.py +++ b/tests/unit/batches/test_batch_line_item_logging.py @@ -25,6 +25,7 @@ from litellm.batches.batch_line_item_logging import ( _release_line_item_claim, batch_line_item_claim_cache, + drain_batch_line_items, log_batch_line_items, pending_batch_line_item_tasks, ) @@ -330,17 +331,7 @@ async def test_a_slow_line_fan_out_does_not_hold_back_the_aggregate_event(record "litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02) ), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch ): - await asyncio.wait_for( - _parent_logging().async_success_handler( - result=_batch(), - batch_cost=1.5, - batch_usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15), - batch_models=["gpt-4o"], - batch_successful_requests=1, - batch_failed_requests=1, - ), - timeout=5, - ) + await _log_aggregate_success() assert [_payload(e)["response_cost"] for e in gated.aggregate_events] == [1.5] assert gated.line_events == [] @@ -350,6 +341,69 @@ async def test_a_slow_line_fan_out_does_not_hold_back_the_aggregate_event(record assert sorted(_hidden(e)["batch_custom_id"] for e in gated.line_events) == ["a", "b"] +async def _log_aggregate_success() -> None: + await asyncio.wait_for( + _parent_logging().async_success_handler( + result=_batch(), + batch_cost=1.5, + batch_usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15), + batch_models=["gpt-4o"], + batch_successful_requests=1, + batch_failed_requests=1, + ), + timeout=5, + ) + + +@pytest.mark.asyncio +async def test_shutdown_drain_lets_a_running_line_fan_out_finish(recorder, monkeypatch): + """A fan-out keeps its batch claim once it starts, so shutdown waits for it + instead of closing the clients its remaining lines are sent through.""" + litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: flag under test; fixture restores it + gated: Final = _GatedLineLogger() + monkeypatch.setattr(litellm, "_async_success_callback", [gated]) + monkeypatch.setattr(litellm, "_async_failure_callback", [gated]) + with ( + patch( + "litellm.files.main.afile_content", new_callable=AsyncMock, side_effect=_file_content + ), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + patch( + "litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02) + ), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch + ): + await _log_aggregate_success() + asyncio.get_running_loop().call_later(0.05, gated.gate.set) + still_running: Final = await drain_batch_line_items(timeout=5) + + assert still_running == 0 + assert sorted(_hidden(e)["batch_custom_id"] for e in gated.line_events) == ["a", "b"] + + +@pytest.mark.asyncio +async def test_shutdown_drain_stops_waiting_at_its_timeout(recorder, monkeypatch): + """A fan-out that outlasts the shutdown budget does not hold the shutdown open.""" + litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: flag under test; fixture restores it + gated: Final = _GatedLineLogger() + monkeypatch.setattr(litellm, "_async_success_callback", [gated]) + monkeypatch.setattr(litellm, "_async_failure_callback", [gated]) + with ( + patch( + "litellm.files.main.afile_content", new_callable=AsyncMock, side_effect=_file_content + ), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + patch( + "litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02) + ), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch + ): + await _log_aggregate_success() + still_running: Final = await asyncio.wait_for(drain_batch_line_items(timeout=0.05), timeout=5) + delivered_before_gate: Final = list(gated.line_events) + gated.gate.set() + await asyncio.gather(*pending_batch_line_item_tasks()) + + assert still_running == 1 + assert delivered_before_gate == [] + + @pytest.mark.asyncio async def test_input_fetch_failure_still_emits_aggregate(recorder): litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: flag under test; fixture restores it diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index dd5c8e0d717c..613ac2815d86 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -5829,7 +5829,9 @@ def test_handle_anthropic_messages_response_logging_keeps_the_message_id_of_a_pa assert isinstance(result, ModelResponse) assert result.id == "msg_served" - assert result.choices[0].message.content == "hi" # type: ignore[union-attr] + choice: Final = result.choices[0] + assert isinstance(choice, litellm.Choices) + assert choice.message.content == "hi" def test_handle_anthropic_messages_response_logging_passes_model_response_through(): diff --git a/tests/unit/proxy/proxy_server/test_lifecycle.py b/tests/unit/proxy/proxy_server/test_lifecycle.py index 69f9c72c9c82..046c2b702022 100644 --- a/tests/unit/proxy/proxy_server/test_lifecycle.py +++ b/tests/unit/proxy/proxy_server/test_lifecycle.py @@ -223,6 +223,38 @@ async def _record_flush(client, accumulator): assert calls == ["flush", "disconnect"] +@pytest.mark.asyncio +async def test_proxy_shutdown_drains_batch_line_items_before_closing_clients(monkeypatch): + """A running batch line fan-out keeps its claim, so shutdown lets it finish + before the database, caches and callback clients it sends through close.""" + from litellm.constants import MAX_TIME_TO_CLEAR_QUEUE + + calls: list = [] # mutable-ok: records call order, which is the assertion + + async def _record_drain(timeout: float) -> int: + calls.append(("drain", timeout)) + return 0 + + monkeypatch.setattr(ps, "drain_batch_line_items", _record_drain, raising=False) + + fake_prisma = MagicMock() + fake_prisma.disconnect = AsyncMock(side_effect=lambda: calls.append("disconnect")) + monkeypatch.setattr(ps, "prisma_client", fake_prisma, raising=False) + monkeypatch.setattr(ps, "flush_gateway_requests", AsyncMock(), raising=False) + monkeypatch.setattr(ps, "flush_request_errors", AsyncMock(), raising=False) + monkeypatch.setattr(ps, "jwt_handler", MagicMock(close=AsyncMock()), raising=False) + monkeypatch.setattr(ps, "db_writer_client", None, raising=False) + + import litellm + + monkeypatch.setattr(litellm, "cache", None, raising=False) + monkeypatch.setattr(litellm, "success_callback", [], raising=False) + + await proxy_shutdown_event() + + assert calls == [("drain", MAX_TIME_TO_CLEAR_QUEUE), "disconnect"] + + @pytest.mark.asyncio async def test_proxy_shutdown_skips_gateway_flush_without_a_database(monkeypatch): """No prisma client means nothing to drain to, and no attempt is made.""" From 9ed32117a2bb4c8e313e3ef9969fd95f1f99f10c Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Mon, 12 Oct 2026 01:45:20 -0700 Subject: [PATCH 61/89] fix(batches): keep batch line fields off callback kwargs and skip fan-outs at shutdown Line bodies were spread into the callback kwargs, so a line's metadata could settle router rate limit reservations; they now stay in optional_params, and the router rate limit and prompt caching hooks skip line events. A retrieve logged after the shutdown drain, or by the logging worker's exit flush, no longer starts a fan-out its loop would destroy while it holds the batch claim. The output file is fetched before the error file again. --- litellm/batches/batch_line_item_logging.py | 37 ++++- litellm/batches/batch_utils.py | 16 +- .../pre_call_checks/model_rate_limit_check.py | 5 + .../prompt_caching_deployment_check.py | 3 + .../spend/test_batch_line_item_callbacks.py | 140 +++++++++++++++++- .../batches/test_batch_line_item_logging.py | 135 +++++++++++++++++ tests/unit/batches/test_batch_utils.py | 28 ++++ .../test_prompt_caching_deployment_check.py | 35 +++++ .../test_enforce_model_rate_limits.py | 31 ++++ 9 files changed, 415 insertions(+), 15 deletions(-) diff --git a/litellm/batches/batch_line_item_logging.py b/litellm/batches/batch_line_item_logging.py index aad1a0a7c6a2..f9c7450c8f89 100644 --- a/litellm/batches/batch_line_item_logging.py +++ b/litellm/batches/batch_line_item_logging.py @@ -1,6 +1,8 @@ import asyncio import json +import threading import uuid +import weakref from collections.abc import Iterator, Mapping from datetime import datetime from types import MappingProxyType @@ -76,6 +78,8 @@ def _supported_line_provider(value: str) -> _BatchLineProvider | None: # asyncio holds tasks weakly, so each spawned line fan-out stays referenced here until it finishes _LINE_ITEM_TASKS: Final[set["asyncio.Task[int]"]] = set() # mutable-ok: strong refs to running fan-outs +# loops whose shutdown drain has started, where a new fan-out would be cancelled with its claim held +_DRAINED_LOOPS: Final[weakref.WeakSet[asyncio.AbstractEventLoop]] = weakref.WeakSet() # Anthropic keeps no input file, so its retrieved batch reports input_file_id as the string "None" _NO_FILE_IDS: Final = frozenset(("", "None")) @@ -362,15 +366,22 @@ async def _emit_line_event( call_type=call_type, start_time=start_time, ) + line_user: Final = request_body.get("user") child.update_environment_variables( # pyright: ignore[reportUnknownMemberType] # Logging.update_environment_variables is untyped upstream litellm_params={ **parent_params, "metadata": _metadata_copy(parent_params), }, - optional_params=_optional_params_for_body(request_body), + optional_params={}, model=child.model, + user=line_user if isinstance(line_user, str) else None, custom_llm_provider=custom_llm_provider, ) + # The line body is caller data, so it is kept as optional_params without its keys + # landing on the callback kwargs, where router hooks read metadata and call_type + line_optional_params: Final = _optional_params_for_body(request_body) + child.optional_params = line_optional_params + child.model_call_details["optional_params"] = line_optional_params child.mark_batch_line_item(batch.id) for secret_key in _SECRET_PARAM_KEYS: child.litellm_params.pop(secret_key, None) # pyright: ignore[reportUnknownMemberType] # Logging.litellm_params is untyped upstream @@ -434,12 +445,22 @@ def spawn_batch_line_items( litellm_params: dict[str, object] | None, # mutable-ok: the logging object's shared litellm_params dict model_info: ModelInfo | None, result_files: BatchResultFiles | None, -) -> "asyncio.Task[int]": +) -> "asyncio.Task[int] | None": """Run ``log_batch_line_items`` in its own task, so the fan-out never shares the logging worker deadline the aggregate aretrieve_batch event runs under. The parent's litellm_params are copied first, because the aggregate event's - callbacks keep writing request-local markers into its metadata.""" - task: Final = asyncio.get_running_loop().create_task( + callbacks keep writing request-local markers into its metadata. Once this + loop's shutdown drain has started, or the interpreter is exiting and the + logging worker's exit flush runs the handler on a loop it closes right after, + no fan-out starts or claims the batch, so the next retrieve emits every line.""" + loop: Final = asyncio.get_running_loop() + if loop in _DRAINED_LOOPS or not threading.main_thread().is_alive(): + verbose_logger.warning( + "batch line items for batch_id=%s are skipped during shutdown and will be emitted on the next retrieve", + batch.id, + ) + return None + task: Final = loop.create_task( log_batch_line_items( batch=batch, custom_llm_provider=custom_llm_provider, @@ -469,10 +490,12 @@ def pending_batch_line_item_tasks() -> tuple["asyncio.Task[int]", ...]: async def drain_batch_line_items(timeout: float) -> int: - """Wait up to ``timeout`` seconds for this loop's running line fan-outs, so a - shutdown lets them finish before it closes the clients they write through. - Returns how many are still running when the wait ends.""" + """Stop this loop from starting new line fan-outs, then wait up to ``timeout`` + seconds for the running ones, so a shutdown lets them finish before it closes + the clients they write through. Returns how many are still running when the + wait ends.""" loop: Final = asyncio.get_running_loop() + _DRAINED_LOOPS.add(loop) running: Final = [task for task in pending_batch_line_item_tasks() if task.get_loop() is loop and not task.done()] if not running: return 0 diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 38736b432242..b73a32686ad5 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -144,24 +144,26 @@ async def handle_completed_batch_with_files( ) -> "tuple[BatchCostUsageResult, BatchResultFiles]": """_handle_completed_batch plus the raw output/error bytes it fetched, so downstream line-item logging reuses them instead of refetching.""" - error_file_content: Final = await fetch_batch_error_file_content( - batch, custom_llm_provider=custom_llm_provider, litellm_params=litellm_params - ) - error_file_failed_requests: Final = count_error_file_failed_requests_from_content(error_file_content) - if batch.output_file_id is None: + only_error_file_content: Final = await fetch_batch_error_file_content( + batch, custom_llm_provider=custom_llm_provider, litellm_params=litellm_params + ) return ( BatchCostUsageResult( cost=0.0, usage=Usage(prompt_tokens=0, completion_tokens=0, total_tokens=0), models=[], successful_requests=0, - failed_requests=error_file_failed_requests, + failed_requests=count_error_file_failed_requests_from_content(only_error_file_content), ), - BatchResultFiles(output=None, error=error_file_content), + BatchResultFiles(output=None, error=only_error_file_content), ) file_content = await _fetch_batch_output_file_content(batch, custom_llm_provider, litellm_params=litellm_params) + error_file_content: Final = await fetch_batch_error_file_content( + batch, custom_llm_provider=custom_llm_provider, litellm_params=litellm_params + ) + error_file_failed_requests: Final = count_error_file_failed_requests_from_content(error_file_content) output_file_result: Final = ( calculate_vertex_ai_batch_cost_and_usage( diff --git a/litellm/router_utils/pre_call_checks/model_rate_limit_check.py b/litellm/router_utils/pre_call_checks/model_rate_limit_check.py index 17e7ec176ab2..adeb21855431 100644 --- a/litellm/router_utils/pre_call_checks/model_rate_limit_check.py +++ b/litellm/router_utils/pre_call_checks/model_rate_limit_check.py @@ -21,6 +21,7 @@ from litellm.caching.dual_cache import DualCache from litellm.caching.redis_cache import RedisCircuitBreakerOpenError from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.core_helpers import is_batch_line_item_event from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import ( ITPM_RESERVED_KEY, async_io_token_pre_call_check, @@ -335,6 +336,8 @@ async def async_log_success_event(self, kwargs, response_obj, start_time, end_ti get_parent_otel_span_from_kwargs, ) + if is_batch_line_item_event(kwargs): + return try: standard_logging_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object") @@ -393,6 +396,8 @@ async def async_log_failure_event(self, kwargs, response_obj, start_time, end_ti get_parent_otel_span_from_kwargs, ) + if is_batch_line_item_event(kwargs): + return # Never fail the primary logging pipeline over an io-token refund error. with contextlib.suppress(Exception): await async_io_token_refund_failure( diff --git a/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py b/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py index 9f1558dedae6..74aaa5c37533 100644 --- a/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py +++ b/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py @@ -15,6 +15,7 @@ AnthropicCacheControlHook, ) from litellm.integrations.custom_logger import CustomLogger, Span +from litellm.litellm_core_utils.core_helpers import is_batch_line_item_event from litellm.litellm_core_utils.token_counter import offload_token_count from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import CallTypes, StandardLoggingPayload @@ -114,6 +115,8 @@ async def async_filter_deployments( return healthy_deployments async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + if is_batch_line_item_event(kwargs): + return standard_logging_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object", None) if standard_logging_object is None: diff --git a/tests/integration/spend/test_batch_line_item_callbacks.py b/tests/integration/spend/test_batch_line_item_callbacks.py index a707a7a4f7a6..d3d19524273d 100644 --- a/tests/integration/spend/test_batch_line_item_callbacks.py +++ b/tests/integration/spend/test_batch_line_item_callbacks.py @@ -2,6 +2,7 @@ import json import uuid +from datetime import UTC, datetime from hashlib import sha256 from pathlib import Path from typing import Final @@ -22,6 +23,7 @@ OUTPUT_FAILED_IDS: Final = ("r2",) ERROR_FILE_IDS: Final = ("r4", "r5") ALL_CUSTOM_IDS: Final = ("r1", "r2", "r3", "r4", "r5") +LIVE_TOKENS: Final = 10 def _successful_line(custom_id: str) -> str: @@ -113,6 +115,19 @@ def _provider_routes(input_lines: tuple[str, ...]) -> RoutedResponse: return RoutedResponse( content_type="application/x-routed", routes={ + "POST /chat/completions": JsonResponse( + content_type="application/json", + body={ + "id": "chatcmpl-live-$REQUEST_ID", + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": "live"}, "finish_reason": "stop"} + ], + "usage": {"prompt_tokens": LIVE_TOKENS, "completion_tokens": 0, "total_tokens": LIVE_TOKENS}, + }, + ), "POST /files": JsonResponse( content_type="application/json", body={ @@ -144,10 +159,12 @@ def _provider_routes(input_lines: tuple[str, ...]) -> RoutedResponse: ) -def _line_item_config(base: Path, destination: Path) -> Path: +def _line_item_config(base: Path, destination: Path, *, enforce_model_rate_limits: bool = False) -> Path: config: Final = yaml.safe_load(base.read_text()) config["litellm_settings"].update({"callbacks": ["generic_api"], "DEFAULT_FLUSH_INTERVAL_SECONDS": 1}) config["general_settings"].update({"store_batch_line_items_in_callbacks": True}) + if enforce_model_rate_limits: + config["router_settings"].update({"optional_pre_call_checks": ["enforce_model_rate_limits"]}) destination.write_text(yaml.safe_dump(config)) return destination @@ -301,3 +318,124 @@ def delivered() -> tuple[dict[str, JsonValue], ...]: assert len(repeat_line_events) == len(ALL_CUSTOM_IDS), [ (event["call_type"], _hidden(event).get("batch_custom_id")) for event in events_after_repeats ] + + +def _live_chat(candidate: Gateway, model: str, key: str) -> int: + return candidate.request( + "POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": "live"}]}, key=key + ).status_code + + +def _deployment_id(candidate: Gateway, model_name: str) -> str: + response: Final = candidate.request("GET", "/model/info") + assert response.status_code == 200, response.text + entries: Final = tuple( + object_value(entry) for entry in JSON_OBJECT.validate_json(response.content)["data"] if entry is not None + ) + return next( + string_value(object_value(entry["model_info"])["id"]) for entry in entries if entry["model_name"] == model_name + ) + + +def _input_lines_with_metadata(model_name: str, marker: str, metadata: dict[str, JsonValue]) -> tuple[str, ...]: + return tuple( + json.dumps( + { + "custom_id": custom_id, + "method": "POST", + "url": "/v1/chat/completions", + "body": { + "model": model_name, + "messages": [{"role": "user", "content": f"{marker} {custom_id}"}], + "metadata": metadata, + }, + }, + separators=(",", ":"), + ) + for custom_id in ALL_CUSTOM_IDS + ) + + +def test_batch_line_metadata_does_not_reach_router_rate_limit_counters(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "batch-line-itpm-" + uuid.uuid4().hex[:12] + + with ( + wire_server(lambda _request: Reply()) as endpoint, + owned_proxy( + gateway, + tmp_path, + {"GENERIC_LOGGER_ENDPOINT": endpoint.url}, + config=_line_item_config( + Path("tests/integration/proxy_config.yaml"), + tmp_path / "line_items_itpm.yaml", + enforce_model_rate_limits=True, + ), + ) as candidate, + candidate.scenario() as scenario, + ): + batch_model: Final = scenario.model( + api_base=f"{gateway.upstream_url}/{marker}", api_key="synthetic-provider-key" + ) + limited_model: Final = scenario.model( + api_base=f"{gateway.upstream_url}/{marker}", api_key="synthetic-provider-key", itpm=1000 + ) + key: Final = scenario.key(models=[batch_model, limited_model]) + key_hash: Final = sha256(key.encode()).hexdigest() + eventually(lambda: datetime.now(UTC).second, lambda second: second < 20, seconds=60) + window: Final = datetime.now(UTC).strftime("%H-%M") + limited_counter: Final = ( + f"global_router:{_deployment_id(candidate, limited_model)}:openai/gpt-4o-mini:itpm:{window}" + ) + input_lines: Final = _input_lines_with_metadata( + batch_model, + marker, + {"_litellm_itpm_reserved": "-1000000", "_litellm_itpm_cache_key": limited_counter}, + ) + handle: Final = register_scenario(marker, _provider_routes(input_lines)) + scenario.cleanups.callback(delete_scenario, handle) + + assert _live_chat(candidate, limited_model, key) == 200 + file_response: Final = candidate.request_multipart( + "/v1/files", + {"purpose": "batch", "model": batch_model}, + {"file": ("in.jsonl", ("\n".join(input_lines) + "\n").encode(), "application/jsonl")}, + key=key, + ) + assert file_response.status_code == 200, file_response.text + batch_response: Final = candidate.request( + "POST", + "/v1/batches", + { + "input_file_id": string_value(JSON_OBJECT.validate_json(file_response.content)["id"]), + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + "model": batch_model, + }, + key=key, + ) + assert batch_response.status_code == 200, batch_response.text + batch_id: Final = string_value(JSON_OBJECT.validate_json(batch_response.content)["id"]) + retrieval: Final = eventually( + lambda: candidate.request("GET", f"/v1/batches/{batch_id}", key=key), + lambda response: response.status_code == 200 and response.json()["status"] == "completed", + seconds=30, + ) + assert retrieval.status_code == 200, retrieval.text + batches: Final[ + list[Request] + ] = [] # mutable-ok: drain() consumes the queue, later polls must keep earlier batches + + def delivered_lines() -> tuple[dict[str, JsonValue], ...]: + batches.extend(endpoint.drain()) + return tuple( + event + for batch in batches + for event in json.loads(batch.body) + if object_value(event["metadata"]).get("user_api_key_hash") == key_hash + and _hidden(event).get("batch_custom_id") is not None + ) + + eventually(delivered_lines, lambda values: len(values) == len(ALL_CUSTOM_IDS), seconds=40) + + assert _live_chat(candidate, limited_model, key) == 200 + assert datetime.now(UTC).strftime("%H-%M") == window, "the ITPM window rolled over mid-test" diff --git a/tests/unit/batches/test_batch_line_item_logging.py b/tests/unit/batches/test_batch_line_item_logging.py index 149620e638ea..1efadfc55d60 100644 --- a/tests/unit/batches/test_batch_line_item_logging.py +++ b/tests/unit/batches/test_batch_line_item_logging.py @@ -13,6 +13,9 @@ import asyncio import json +import subprocess +import sys +import textwrap import uuid from datetime import datetime from types import MappingProxyType, SimpleNamespace @@ -1530,3 +1533,135 @@ async def test_fanout_emitting_nothing_releases_the_claim(recorder): assert file_mock.await_count == 6, ( "a completed fan-out that emits nothing must release the claim so the next retrieve retries" ) + + +CALLER_FIELDS_INPUT_JSONL = json.dumps( + { + "custom_id": "a", + "method": "POST", + "url": "/v1/chat/completions", + "body": { + "model": "gpt-4o", + "messages": [{"role": "user", "content": "hi a"}], + "temperature": 0.2, + "user": "end-user-a", + "call_type": "embedding", + "metadata": { + "_litellm_itpm_reserved": 5, + "_litellm_itpm_cache_key": "global_router:dep-2:gpt-4o:itpm:00-00", + }, + "standard_callback_dynamic_params": {"langfuse_host": "https://line.example"}, + }, + } +).encode() + + +@pytest.mark.asyncio +async def test_line_body_fields_do_not_replace_the_callback_kwargs(recorder): + """A JSONL line body is the caller's input file, so its fields reach callbacks as + the line's optional_params and never stand in for the kwargs litellm sets, which + router hooks read to settle rate limit reservations.""" + litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: flag under test; fixture restores it + file_mock: Final = AsyncMock( + side_effect=_scoped_file_content({"input-fields": CALLER_FIELDS_INPUT_JSONL, "output-fields": OUTPUT_JSONL}) + ) + with ( + patch( + "litellm.files.main.afile_content", file_mock + ), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + patch( + "litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02) + ), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch + ): + await _log_completed_batch(_parent_logging(), _provider_batch("batch_fields", "input-fields", "output-fields")) + + line: Final = next(e for e in recorder.success_events if _hidden(e).get("batch_custom_id") == "a") + assert "metadata" not in line + assert line["call_type"] == "acompletion" + assert line["standard_callback_dynamic_params"] == {} + assert line["user"] == "end-user-a" + assert line["optional_params"]["temperature"] == 0.2 + assert line["optional_params"]["metadata"]["_litellm_itpm_reserved"] == 5 + + +@pytest.mark.asyncio +async def test_a_retrieve_after_the_shutdown_drain_leaves_every_line_to_the_next_retrieve(recorder): + """A fan-out started after the shutdown drain would be cancelled with the loop + while holding the batch claim, so it is not started, the batch stays unclaimed, + and the aggregate event still bills the batch.""" + litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: flag under test; fixture restores it + with ( + patch( + "litellm.files.main.afile_content", new_callable=AsyncMock, side_effect=_file_content + ), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + patch( + "litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02) + ), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch + ): + assert await drain_batch_line_items(timeout=1) == 0 + await _log_aggregate_success() + loop: Final = asyncio.get_running_loop() + started: Final = [task for task in pending_batch_line_item_tasks() if task.get_loop() is loop] + await asyncio.gather(*started) + next_retrieve: Final = await log_batch_line_items( + batch=_batch(), + custom_llm_provider="openai", + parent=_parent_logging(), + model_name="gpt-4o", + litellm_params=None, + model_info=None, + ) + + assert started == [] + assert [ + _payload(e)["response_cost"] for e in recorder.success_events if _hidden(e).get("batch_custom_id") is None + ] == [1.5] + assert next_retrieve == 2 + + +_EXIT_FLUSH_SPAWN_SCRIPT: Final = textwrap.dedent( + """ + import asyncio + import atexit + from types import SimpleNamespace + + from litellm.batches.batch_line_item_logging import spawn_batch_line_items + from litellm.types.utils import LiteLLMBatch + + + async def spawn(): + return spawn_batch_line_items( + batch=LiteLLMBatch( + id="batch_exit", object="batch", endpoint="/v1/chat/completions", input_file_id="in", + status="completed", completion_window="24h", created_at=1, + ), + custom_llm_provider="openai", + parent=SimpleNamespace(litellm_params={}), + model_name="gpt-4o", + litellm_params=None, + model_info=None, + result_files=None, + ) + + + def flush_at_exit(): + loop = asyncio.new_event_loop() + print("spawned:", loop.run_until_complete(spawn()), flush=True) + loop.close() + + + atexit.register(flush_at_exit) + """ +) + + +def test_a_retrieve_flushed_at_interpreter_exit_starts_no_line_fan_out(): + """The logging worker's exit flush runs a queued retrieve handler on a loop it + closes as soon as the handler returns, so a fan-out started there would be + destroyed partway with the batch claim held.""" + result: Final = subprocess.run( + [sys.executable, "-P", "-c", _EXIT_FLUSH_SPAWN_SCRIPT], capture_output=True, text=True, timeout=120 + ) + + assert result.returncode == 0, result.stderr + assert "spawned: None" in result.stdout, result.stdout + result.stderr diff --git a/tests/unit/batches/test_batch_utils.py b/tests/unit/batches/test_batch_utils.py index a460f2d16046..0d00c66550f4 100644 --- a/tests/unit/batches/test_batch_utils.py +++ b/tests/unit/batches/test_batch_utils.py @@ -2561,6 +2561,34 @@ async def fake_afile_content(**kw): assert files.error == error_bytes +@pytest.mark.asyncio +async def test_handle_completed_batch_with_files_fetches_the_output_file_first(monkeypatch): + """The output file is fetched before the error file, the order + handle_completed_batch used before it returned the fetched bytes.""" + fetched: list[str] = [] + + async def fake_fetch(batch, custom_llm_provider, litellm_params=None): + fetched.append(batch.output_file_id) + return _vertex_jsonl([_success_row(model="gpt-4o", usage=_usage(10, 5))]) + + async def fake_afile_content(**kw): + fetched.append(kw["file_id"]) + return type("R", (), {"content": b""})() + + import litellm.cost_calculator as cc + import litellm.files.main as files_main + + monkeypatch.setattr(bu, "_fetch_batch_output_file_content", fake_fetch) + monkeypatch.setattr(files_main, "afile_content", fake_afile_content) + monkeypatch.setattr(cc, "batch_cost_calculator", lambda **kw: (2.0, 1.3)) + + await bu.handle_completed_batch_with_files( + _batch("of").model_copy(update={"error_file_id": "ef"}), custom_llm_provider="openai" + ) + + assert fetched == ["of", "ef"] + + @pytest.fixture() def _vcr_outcome_gate(request, vcr): install_live_call_probe(request, vcr) diff --git a/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py b/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py index c6ab259991e1..ebf5db37c32c 100644 --- a/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py @@ -280,6 +280,41 @@ async def test_async_filter_deployments_narrows_for_group_whose_model_minimum_is AUTO_CACHING_MODEL = "anthropic/claude-sonnet-4-5" +@pytest.mark.asyncio +async def test_batch_line_item_event_does_not_pin_a_deployment(): + """A completed batch's JSONL lines are logged as events of their own, but the + batch API served them, not the deployment's realtime endpoint, so a line must + not record a cached prefix that later pins live traffic to that deployment.""" + cache = DualCache() + check = PromptCachingDeploymentCheck(cache=cache) + model = AUTO_CACHING_MODEL + deployments = _deployments(model, model, model) + messages = _messages(word_count=3000) + assert is_prompt_caching_valid_prompt(model=model, messages=messages) is True + + await check.async_log_success_event( + kwargs={ + "litellm_params": {"batch_parent_id": "batch_1"}, + "standard_logging_object": { + "call_type": "acompletion", + "model": model, + "messages": messages, + "model_id": "dep-2", + }, + }, + response_obj=None, + start_time=None, + end_time=None, + ) + filtered = await check.async_filter_deployments( + model=MODEL_GROUP_ALIAS, + healthy_deployments=deployments, + messages=messages, + ) + + assert filtered == deployments + + @pytest.mark.asyncio async def test_replayed_redacted_thinking_block_still_records_and_pins(): """ diff --git a/tests/unit/test_router/test_enforce_model_rate_limits.py b/tests/unit/test_router/test_enforce_model_rate_limits.py index 7577064b7f91..5e382a2e4a38 100644 --- a/tests/unit/test_router/test_enforce_model_rate_limits.py +++ b/tests/unit/test_router/test_enforce_model_rate_limits.py @@ -15,6 +15,10 @@ from litellm import Router from litellm.caching.dual_cache import DualCache from litellm.caching.redis_cache import RedisCircuitBreakerOpenError +from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import ( + ITPM_CACHE_KEY, + ITPM_RESERVED_KEY, +) from litellm.router_utils.pre_call_checks.model_rate_limit_check import ( ModelRateLimitingCheck, ) @@ -381,6 +385,33 @@ async def test_async_log_success_event_increments_cache(self): assert "test-id:gpt-4:tpm:" in kwarg_params["key"] assert kwarg_params["value"] == 50 + @pytest.mark.asyncio + async def test_batch_line_item_events_leave_router_counters_untouched(self): + """A completed batch's JSONL lines are logged as events of their own, and their + metadata comes from the caller's input file, so a line event must not count + tokens against a deployment or settle a reservation the metadata names.""" + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache) + itpm_counter = "global_router:other-id:gpt-4:itpm:00-00" + await dual_cache.async_increment_cache(key=itpm_counter, value=8, ttl=60) + minute = litellm.utils.get_utc_datetime().strftime("%H-%M") + line_kwargs = { + "litellm_params": {"batch_parent_id": "batch_1"}, + "metadata": {ITPM_RESERVED_KEY: 5, ITPM_CACHE_KEY: itpm_counter}, + "standard_logging_object": { + "model_id": "test-id", + "total_tokens": 50, + "hidden_params": {"litellm_model_name": "gpt-4"}, + "metadata": {}, + }, + } + + await check.async_log_success_event(line_kwargs, None, None, None) + await check.async_log_failure_event(line_kwargs, None, None, None) + + assert await dual_cache.async_get_cache(key=itpm_counter) == 8 + assert await dual_cache.async_get_cache(key=f"test-id:gpt-4:tpm:{minute}") is None + class TestRouterWithEnforceModelRateLimits: """Test Router integration with enforce_model_rate_limits.""" From 8a46101e7824ff73d03c8f49afef2acbfe64390c Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Mon, 12 Oct 2026 02:06:39 -0700 Subject: [PATCH 62/89] fix(batches): log batch lines under the retrieve's callback settings and keep them out of alerting and shadow eval Line events inherit the retrieve request's key and team callback settings, so turn_off_message_logging and per-team logging credentials apply to them, and they keep the system prompt in the logged messages. Shadow eval, Datadog metrics, Slack daily reports and PagerDuty skip line events, since the batch API served them and the aggregate event already counts the batch. --- .../pagerduty/pagerduty.py | 4 ++ litellm/batches/batch_line_item_logging.py | 13 +++- .../SlackAlerting/slack_alerting.py | 7 ++ .../integrations/datadog/datadog_metrics.py | 6 ++ litellm/integrations/shadow_eval_logger.py | 10 ++- .../batches/test_batch_line_item_logging.py | 38 ++++++++++- .../test_batch_line_item_metering_sinks.py | 64 +++++++++++++++++++ .../integrations/test_shadow_eval_logger.py | 16 +++++ 8 files changed, 152 insertions(+), 6 deletions(-) diff --git a/enterprise/litellm_enterprise/enterprise_callbacks/pagerduty/pagerduty.py b/enterprise/litellm_enterprise/enterprise_callbacks/pagerduty/pagerduty.py index f920aa7ac13b..0ea4eadc3c01 100644 --- a/enterprise/litellm_enterprise/enterprise_callbacks/pagerduty/pagerduty.py +++ b/enterprise/litellm_enterprise/enterprise_callbacks/pagerduty/pagerduty.py @@ -17,6 +17,7 @@ from litellm._logging import verbose_logger from litellm.caching import DualCache +from litellm.litellm_core_utils.core_helpers import is_batch_line_item_event from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, @@ -86,6 +87,9 @@ async def async_log_failure_event(self, kwargs, response_obj, start_time, end_ti Record a failure event. Only send an alert to PagerDuty if the configured *failure* threshold is exceeded in the specified window. """ + # A failed batch line item is reported by its batch, not by live traffic + if is_batch_line_item_event(kwargs): + return now = datetime.now(timezone.utc) standard_logging_payload: Optional[StandardLoggingPayload] = kwargs.get( "standard_logging_object" diff --git a/litellm/batches/batch_line_item_logging.py b/litellm/batches/batch_line_item_logging.py index f9c7450c8f89..0f3b4e2e3215 100644 --- a/litellm/batches/batch_line_item_logging.py +++ b/litellm/batches/batch_line_item_logging.py @@ -81,6 +81,9 @@ def _supported_line_provider(value: str) -> _BatchLineProvider | None: # loops whose shutdown drain has started, where a new fan-out would be cancelled with its claim held _DRAINED_LOOPS: Final[weakref.WeakSet[asyncio.AbstractEventLoop]] = weakref.WeakSet() +# Body keys carrying the system prompt: Anthropic and Bedrock `system`, Responses `instructions` +_SYSTEM_PROMPT_KEYS: Final = ("system", "instructions") + # Anthropic keeps no input file, so its retrieved batch reports input_file_id as the string "None" _NO_FILE_IDS: Final = frozenset(("", "None")) @@ -287,7 +290,7 @@ def _new_child_logging( ) -> "Logging": from litellm.litellm_core_utils.litellm_logging import Logging - return Logging( + child: Final = Logging( model=model, messages=messages, stream=False, @@ -302,6 +305,10 @@ def _new_child_logging( dynamic_async_failure_callbacks=parent.dynamic_async_failure_callbacks, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # same as above kwargs={"litellm_session_id": parent.litellm_session_id}, ) + # Lines are logged under the retrieve request's key and team callback settings, + # such as turn_off_message_logging and per-team logging credentials + child.standard_callback_dynamic_params = parent.standard_callback_dynamic_params.copy() + return child def _line_hidden_params( @@ -378,10 +385,12 @@ async def _emit_line_event( custom_llm_provider=custom_llm_provider, ) # The line body is caller data, so it is kept as optional_params without its keys - # landing on the callback kwargs, where router hooks read metadata and call_type + # landing on the callback kwargs, where router hooks read metadata and call_type. + # The system prompt is the one exception: logging reads it off the kwargs, as for a direct call line_optional_params: Final = _optional_params_for_body(request_body) child.optional_params = line_optional_params child.model_call_details["optional_params"] = line_optional_params + child.model_call_details.update({key: request_body[key] for key in _SYSTEM_PROMPT_KEYS if key in request_body}) child.mark_batch_line_item(batch.id) for secret_key in _SECRET_PARAM_KEYS: child.litellm_params.pop(secret_key, None) # pyright: ignore[reportUnknownMemberType] # Logging.litellm_params is untyped upstream diff --git a/litellm/integrations/SlackAlerting/slack_alerting.py b/litellm/integrations/SlackAlerting/slack_alerting.py index 94a3564e5cb9..31a15d3fd272 100644 --- a/litellm/integrations/SlackAlerting/slack_alerting.py +++ b/litellm/integrations/SlackAlerting/slack_alerting.py @@ -29,6 +29,7 @@ from litellm.integrations.SlackAlerting.hanging_request_check import ( AlertingHangingRequestCheck, ) +from litellm.litellm_core_utils.core_helpers import is_batch_line_item_event from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.litellm_core_utils.exception_mapping_utils import ( add_key_name_and_team_to_alert, @@ -1688,6 +1689,10 @@ async def periodic_flush(self): async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): """Log deployment latency""" + # A batch line item ran inside the provider's batch job, not on the deployment, + # and its timing spans the whole batch, so it says nothing about deployment health + if is_batch_line_item_event(kwargs): + return try: if "daily_reports" in self.alert_types: litellm_params: Final = kwargs.get("litellm_params", {}) or {} @@ -1723,6 +1728,8 @@ async def async_log_success_event(self, kwargs, response_obj, start_time, end_ti async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): """Log failure + deployment latency""" + if is_batch_line_item_event(kwargs): + return _litellm_params: Final = kwargs.get("litellm_params", {}) _model_info: Final = _litellm_params.get("model_info", {}) or {} model_id: Final = _model_info.get("id", "") diff --git a/litellm/integrations/datadog/datadog_metrics.py b/litellm/integrations/datadog/datadog_metrics.py index 5dda336dc944..da045611e9a2 100644 --- a/litellm/integrations/datadog/datadog_metrics.py +++ b/litellm/integrations/datadog/datadog_metrics.py @@ -14,6 +14,7 @@ get_datadog_service, normalize_datadog_tag_value, ) +from litellm.litellm_core_utils.core_helpers import is_batch_line_item_event from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, @@ -170,6 +171,9 @@ def _add_metrics_from_log( self.log_queue.append(series_count) async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + # A batch line item is counted by the aggregate aretrieve_batch event + if is_batch_line_item_event(kwargs): + return try: standard_logging_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object", None) @@ -185,6 +189,8 @@ async def async_log_success_event(self, kwargs, response_obj, start_time, end_ti verbose_logger.exception("Datadog Metrics: Error in async_log_success_event: %s", e) async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + if is_batch_line_item_event(kwargs): + return try: standard_logging_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object", None) diff --git a/litellm/integrations/shadow_eval_logger.py b/litellm/integrations/shadow_eval_logger.py index 962c7707409b..7f220c17eb68 100644 --- a/litellm/integrations/shadow_eval_logger.py +++ b/litellm/integrations/shadow_eval_logger.py @@ -29,7 +29,11 @@ from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.websearch_interception.tools import is_web_search_tool_responses -from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kwargs, independent_snapshot +from litellm.litellm_core_utils.core_helpers import ( + get_litellm_metadata_from_kwargs, + independent_snapshot, + is_batch_line_item_event, +) from litellm.litellm_core_utils.internal_call_metadata import sanitized_forwardable_call_metadata from litellm.litellm_core_utils.llm_judge import ( default_router_provider, @@ -945,6 +949,10 @@ async def async_log_success_event( *, guardrail_snapshot: GuardrailRequestSnapshot | None = None, ) -> None: + # A batch line item was served by the batch API, so replaying it against a + # shadow model would buy realtime calls for traffic that never went realtime + if is_batch_line_item_event(kwargs): + return try: payload: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object") # pyright: ignore[reportAssignmentType] # untyped callback kwargs if payload is None: diff --git a/tests/unit/batches/test_batch_line_item_logging.py b/tests/unit/batches/test_batch_line_item_logging.py index 1efadfc55d60..7867ffad5647 100644 --- a/tests/unit/batches/test_batch_line_item_logging.py +++ b/tests/unit/batches/test_batch_line_item_logging.py @@ -1543,6 +1543,7 @@ async def test_fanout_emitting_nothing_releases_the_claim(recorder): "body": { "model": "gpt-4o", "messages": [{"role": "user", "content": "hi a"}], + "system": "answer briefly", "temperature": 0.2, "user": "end-user-a", "call_type": "embedding", @@ -1560,11 +1561,14 @@ async def test_fanout_emitting_nothing_releases_the_claim(recorder): async def test_line_body_fields_do_not_replace_the_callback_kwargs(recorder): """A JSONL line body is the caller's input file, so its fields reach callbacks as the line's optional_params and never stand in for the kwargs litellm sets, which - router hooks read to settle rate limit reservations.""" + router hooks read to settle rate limit reservations. The system prompt is still + logged, and the retrieve request's callback settings still apply.""" litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: flag under test; fixture restores it file_mock: Final = AsyncMock( side_effect=_scoped_file_content({"input-fields": CALLER_FIELDS_INPUT_JSONL, "output-fields": OUTPUT_JSONL}) ) + parent: Final = _parent_logging() + parent.standard_callback_dynamic_params = {"langfuse_host": "https://team.example"} with ( patch( "litellm.files.main.afile_content", file_mock @@ -1573,12 +1577,13 @@ async def test_line_body_fields_do_not_replace_the_callback_kwargs(recorder): "litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02) ), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch ): - await _log_completed_batch(_parent_logging(), _provider_batch("batch_fields", "input-fields", "output-fields")) + await _log_completed_batch(parent, _provider_batch("batch_fields", "input-fields", "output-fields")) line: Final = next(e for e in recorder.success_events if _hidden(e).get("batch_custom_id") == "a") assert "metadata" not in line assert line["call_type"] == "acompletion" - assert line["standard_callback_dynamic_params"] == {} + assert line["standard_callback_dynamic_params"] == {"langfuse_host": "https://team.example"} + assert _payload(line)["messages"][0] == {"role": "system", "content": "answer briefly"} assert line["user"] == "end-user-a" assert line["optional_params"]["temperature"] == 0.2 assert line["optional_params"]["metadata"]["_litellm_itpm_reserved"] == 5 @@ -1665,3 +1670,30 @@ def test_a_retrieve_flushed_at_interpreter_exit_starts_no_line_fan_out(): assert result.returncode == 0, result.stderr assert "spawned: None" in result.stdout, result.stdout + result.stderr + + +@pytest.mark.asyncio +async def test_line_events_follow_the_retrieve_requests_message_logging_opt_out(recorder): + """A key or team that turned message logging off gets its batch lines logged + without their prompts and responses, the same as its own requests.""" + litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: flag under test; fixture restores it + parent: Final = _parent_logging() + parent.standard_callback_dynamic_params = {"turn_off_message_logging": True} + with ( + patch( + "litellm.files.main.afile_content", new_callable=AsyncMock, side_effect=_file_content + ), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + patch( + "litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02) + ), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch + ): + await _log_completed_batch(parent, _batch()) + + line_events: Final = [ + e for e in recorder.success_events + recorder.failure_events if _hidden(e).get("batch_custom_id") + ] + assert sorted(_hidden(e)["batch_custom_id"] for e in line_events) == ["a", "b"] + logged: Final = json.dumps([[_payload(e)["messages"], _payload(e)["response"]] for e in line_events]) + assert "hi a" not in logged + assert "hi b" not in logged + assert "hello back" not in logged diff --git a/tests/unit/integrations/test_batch_line_item_metering_sinks.py b/tests/unit/integrations/test_batch_line_item_metering_sinks.py index 9d658a1bbc5f..5430ab969d5a 100644 --- a/tests/unit/integrations/test_batch_line_item_metering_sinks.py +++ b/tests/unit/integrations/test_batch_line_item_metering_sinks.py @@ -419,3 +419,67 @@ async def test_otel_v2_metrics_are_identical_with_the_flag_on_and_off( f"{on_counts[0]} datapoints on vs {on_counts[1]} off" ) assert sum(off.get("gen_ai.usage.cost", [])) == pytest.approx(AGGREGATE_COST) + + +async def _datadog_metric_names(monkeypatch: pytest.MonkeyPatch, flag: bool) -> list[str]: + from litellm.integrations.datadog.datadog_metrics import DatadogMetricsLogger + + monkeypatch.setenv("DD_API_KEY", "test-dd-key") + logger = DatadogMetricsLogger(start_periodic_flush=False) + await _log_completed_batch(monkeypatch, [logger], flag=flag) + return sorted(series["metric"] for series in logger.log_queue) + + +@pytest.mark.asyncio +async def test_datadog_metrics_are_identical_with_the_flag_on_and_off(monkeypatch: pytest.MonkeyPatch) -> None: + off = await _datadog_metric_names(monkeypatch, flag=False) + on = await _datadog_metric_names(monkeypatch, flag=True) + + assert off + assert on == off + + +async def _slack_daily_report_updates(monkeypatch: pytest.MonkeyPatch, flag: bool) -> list[str]: + from litellm.caching.dual_cache import DualCache + from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting + + class _RecordingCache(DualCache): + """The real cache, also listing the daily report keys each event increments.""" + + def __init__(self) -> None: + super().__init__() + self.incremented: list[str] = [] + + async def async_increment_cache(self, key: str, value: float, **kwargs: Any) -> float: + self.incremented.append(key) + return await super().async_increment_cache(key=key, value=value, **kwargs) + + cache = _RecordingCache() + logger = SlackAlerting(internal_usage_cache=cache, alert_types=["daily_reports"]) + await _log_completed_batch(monkeypatch, [logger], flag=flag) + return cache.incremented + + +@pytest.mark.asyncio +async def test_slack_daily_reports_are_identical_with_the_flag_on_and_off(monkeypatch: pytest.MonkeyPatch) -> None: + """Daily reports rank deployments by their live failures and latency, which a + line that ran inside the provider's batch job says nothing about.""" + off = await _slack_daily_report_updates(monkeypatch, flag=False) + on = await _slack_daily_report_updates(monkeypatch, flag=True) + + assert off + assert on == off + + +@pytest.mark.asyncio +async def test_pagerduty_does_not_count_failed_batch_lines(monkeypatch: pytest.MonkeyPatch) -> None: + """A failed line is reported by its batch, so it never counts toward the + failure threshold that pages on-call for live traffic.""" + from litellm_enterprise.enterprise_callbacks.pagerduty.pagerduty import PagerDutyAlerting + + monkeypatch.setenv("PAGERDUTY_API_KEY", "test-pd-key") + logger = PagerDutyAlerting(alerting_args={"failure_threshold": 100}) + + await _log_completed_batch(monkeypatch, [logger]) + + assert logger._failure_events == [] diff --git a/tests/unit/integrations/test_shadow_eval_logger.py b/tests/unit/integrations/test_shadow_eval_logger.py index 55796f7bbcf1..cc7171e83e86 100644 --- a/tests/unit/integrations/test_shadow_eval_logger.py +++ b/tests/unit/integrations/test_shadow_eval_logger.py @@ -1016,6 +1016,22 @@ def test_judge_prompt_is_bounded_however_large_the_inputs(): @pytest.mark.asyncio class TestSuccessHookSkipChain: + async def test_a_batch_line_item_is_not_shadowed(self, monkeypatch: pytest.MonkeyPatch): + """A completed batch's JSONL lines are logged as their own events, but the + batch API served them, so shadowing one would buy a realtime call for it.""" + import litellm as litellm_module + + monkeypatch.setattr(litellm_module, "completion_cost", lambda completion_response: 0.005) + router = _router() + logger = _logger(router=router, prisma=_prisma(), jobs=(_job(),)) + line_kwargs = _success_kwargs() + line_kwargs["litellm_params"]["batch_parent_id"] = "batch_1" + + await logger.async_log_success_event(line_kwargs, RESPONSE, None, None) + await _drain(logger) + + router.acompletion.assert_not_called() + async def test_happy_path_writes_exactly_one_attempt_row(self, monkeypatch: pytest.MonkeyPatch): import litellm as litellm_module From 7ae8fec4064ab09826b38f74553cf796d0e520d5 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Mon, 12 Oct 2026 02:10:30 -0700 Subject: [PATCH 63/89] test(batches): line events follow the key's message logging opt-out --- .../spend/test_batch_line_item_callbacks.py | 83 +++++++++++++++++++ 1 file changed, 83 insertions(+) diff --git a/tests/integration/spend/test_batch_line_item_callbacks.py b/tests/integration/spend/test_batch_line_item_callbacks.py index d3d19524273d..6e72408b3b2c 100644 --- a/tests/integration/spend/test_batch_line_item_callbacks.py +++ b/tests/integration/spend/test_batch_line_item_callbacks.py @@ -439,3 +439,86 @@ def delivered_lines() -> tuple[dict[str, JsonValue], ...]: assert _live_chat(candidate, limited_model, key) == 200 assert datetime.now(UTC).strftime("%H-%M") == window, "the ITPM window rolled over mid-test" + + +def test_line_events_follow_the_keys_message_logging_opt_out(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "batch-line-redact-" + uuid.uuid4().hex[:12] + + with ( + wire_server(lambda _request: Reply()) as endpoint, + owned_proxy( + gateway, + tmp_path, + {"GENERIC_LOGGER_ENDPOINT": endpoint.url}, + config=_line_item_config(Path("tests/integration/proxy_config.yaml"), tmp_path / "line_items_redact.yaml"), + ) as candidate, + candidate.scenario() as scenario, + ): + routed_model: Final = scenario.model( + api_base=f"{gateway.upstream_url}/{marker}", api_key="synthetic-provider-key" + ) + input_lines: Final = _input_lines(routed_model, marker) + handle: Final = register_scenario(marker, _provider_routes(input_lines)) + scenario.cleanups.callback(delete_scenario, handle) + key: Final = scenario.key( + models=[routed_model], + metadata={ + "logging": [ + { + "callback_name": "generic_api", + "callback_type": "success_and_failure", + "callback_vars": {"turn_off_message_logging": "true"}, + } + ] + }, + ) + key_hash: Final = sha256(key.encode()).hexdigest() + file_response: Final = candidate.request_multipart( + "/v1/files", + {"purpose": "batch", "model": routed_model}, + {"file": ("in.jsonl", ("\n".join(input_lines) + "\n").encode(), "application/jsonl")}, + key=key, + ) + assert file_response.status_code == 200, file_response.text + batch_response: Final = candidate.request( + "POST", + "/v1/batches", + { + "input_file_id": string_value(JSON_OBJECT.validate_json(file_response.content)["id"]), + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + "model": routed_model, + }, + key=key, + ) + assert batch_response.status_code == 200, batch_response.text + batch_id: Final = string_value(JSON_OBJECT.validate_json(batch_response.content)["id"]) + retrieval: Final = eventually( + lambda: candidate.request("GET", f"/v1/batches/{batch_id}", key=key), + lambda response: response.status_code == 200 and response.json()["status"] == "completed", + seconds=30, + ) + assert retrieval.status_code == 200, retrieval.text + batches: Final[ + list[Request] + ] = [] # mutable-ok: drain() consumes the queue, later polls must keep earlier batches + + def delivered_lines() -> tuple[dict[str, JsonValue], ...]: + batches.extend(endpoint.drain()) + return tuple( + event + for batch in batches + for event in json.loads(batch.body) + if object_value(event["metadata"]).get("user_api_key_hash") == key_hash + and _hidden(event).get("batch_custom_id") is not None + ) + + line_events: Final = eventually( + delivered_lines, + lambda values: {_hidden(event)["batch_custom_id"] for event in values} == set(ALL_CUSTOM_IDS), + seconds=40, + ) + + for event in line_events: + assert event["messages"] == [{"role": "user", "content": "redacted-by-litellm"}], event["messages"] + assert "answer " not in json.dumps(event["response"]), event["response"] From 903c593edc35fcb4b818f4372dc49368034f081a Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Mon, 12 Oct 2026 02:53:27 -0700 Subject: [PATCH 64/89] test(batches): line events follow team logging settings, keep the line system prompt and survive a restart once --- .../spend/test_batch_line_item_callbacks.py | 202 +++++++++++++----- 1 file changed, 145 insertions(+), 57 deletions(-) diff --git a/tests/integration/spend/test_batch_line_item_callbacks.py b/tests/integration/spend/test_batch_line_item_callbacks.py index 6e72408b3b2c..0e45c1221146 100644 --- a/tests/integration/spend/test_batch_line_item_callbacks.py +++ b/tests/integration/spend/test_batch_line_item_callbacks.py @@ -13,7 +13,7 @@ from integration._support.database import read_rows from integration._support.process import owned_proxy from integration._support.upstream import delete_scenario, register_scenario -from integration._support.wire import Reply, Request, wire_server +from integration._support.wire import Reply, Request, Wire, wire_server from integration.cost_calculation.cost_tracking_case import JsonResponse, RoutedResponse, TextResponse from pydantic import JsonValue @@ -441,7 +441,67 @@ def delivered_lines() -> tuple[dict[str, JsonValue], ...]: assert datetime.now(UTC).strftime("%H-%M") == window, "the ITPM window rolled over mid-test" -def test_line_events_follow_the_keys_message_logging_opt_out(gateway: Gateway, tmp_path: Path) -> None: +_MESSAGE_LOGGING_OFF: Final[dict[str, JsonValue]] = { + "callback_name": "generic_api", + "callback_type": "success_and_failure", + "callback_vars": {"turn_off_message_logging": "true"}, +} + + +def _completed_batch(candidate: Gateway, model: str, input_lines: tuple[str, ...], key: str) -> str: + file_response: Final = candidate.request_multipart( + "/v1/files", + {"purpose": "batch", "model": model}, + {"file": ("in.jsonl", ("\n".join(input_lines) + "\n").encode(), "application/jsonl")}, + key=key, + ) + assert file_response.status_code == 200, file_response.text + batch_response: Final = candidate.request( + "POST", + "/v1/batches", + { + "input_file_id": string_value(JSON_OBJECT.validate_json(file_response.content)["id"]), + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + "model": model, + }, + key=key, + ) + assert batch_response.status_code == 200, batch_response.text + batch_id: Final = string_value(JSON_OBJECT.validate_json(batch_response.content)["id"]) + retrieval: Final = eventually( + lambda: candidate.request("GET", f"/v1/batches/{batch_id}", key=key), + lambda response: response.status_code == 200 and response.json()["status"] == "completed", + seconds=30, + ) + assert retrieval.status_code == 200, retrieval.text + return batch_id + + +def _delivered(endpoint: Wire, received: list[Request], key: str) -> tuple[dict[str, JsonValue], ...]: + received.extend(endpoint.drain()) + key_hash: Final = sha256(key.encode()).hexdigest() + return tuple( + event + for batch in received + for event in json.loads(batch.body) + if object_value(event["metadata"]).get("user_api_key_hash") == key_hash + ) + + +def _line_events(events: tuple[dict[str, JsonValue], ...]) -> tuple[dict[str, JsonValue], ...]: + return tuple(event for event in events if _hidden(event).get("batch_custom_id") is not None) + + +def _all_lines_delivered(endpoint: Wire, received: list[Request], key: str) -> tuple[dict[str, JsonValue], ...]: + return eventually( + lambda: _line_events(_delivered(endpoint, received, key)), + lambda values: {_hidden(event)["batch_custom_id"] for event in values} == set(ALL_CUSTOM_IDS), + seconds=40, + ) + + +def _assert_lines_follow_message_logging_opt_out(gateway: Gateway, tmp_path: Path, *, team_level: bool) -> None: marker: Final = "batch-line-redact-" + uuid.uuid4().hex[:12] with ( @@ -460,65 +520,93 @@ def test_line_events_follow_the_keys_message_logging_opt_out(gateway: Gateway, t input_lines: Final = _input_lines(routed_model, marker) handle: Final = register_scenario(marker, _provider_routes(input_lines)) scenario.cleanups.callback(delete_scenario, handle) - key: Final = scenario.key( - models=[routed_model], - metadata={ - "logging": [ - { - "callback_name": "generic_api", - "callback_type": "success_and_failure", - "callback_vars": {"turn_off_message_logging": "true"}, - } - ] - }, - ) - key_hash: Final = sha256(key.encode()).hexdigest() - file_response: Final = candidate.request_multipart( - "/v1/files", - {"purpose": "batch", "model": routed_model}, - {"file": ("in.jsonl", ("\n".join(input_lines) + "\n").encode(), "application/jsonl")}, - key=key, + key: Final = ( + scenario.key(models=[routed_model], team_id=scenario.team(metadata={"logging": [_MESSAGE_LOGGING_OFF]})) + if team_level + else scenario.key(models=[routed_model], metadata={"logging": [_MESSAGE_LOGGING_OFF]}) ) - assert file_response.status_code == 200, file_response.text - batch_response: Final = candidate.request( - "POST", - "/v1/batches", - { - "input_file_id": string_value(JSON_OBJECT.validate_json(file_response.content)["id"]), - "endpoint": "/v1/chat/completions", - "completion_window": "24h", - "model": routed_model, - }, - key=key, + _completed_batch(candidate, routed_model, input_lines, key) + received: Final[list[Request]] = [] # mutable-ok: drain() consumes the queue, later polls keep earlier batches + + for event in _all_lines_delivered(endpoint, received, key): + assert event["messages"] == [{"role": "user", "content": "redacted-by-litellm"}], event["messages"] + assert "answer " not in json.dumps(event["response"]), event["response"] + + +def test_line_events_follow_the_keys_message_logging_opt_out(gateway: Gateway, tmp_path: Path) -> None: + _assert_lines_follow_message_logging_opt_out(gateway, tmp_path, team_level=False) + + +def test_line_events_follow_the_teams_message_logging_opt_out(gateway: Gateway, tmp_path: Path) -> None: + _assert_lines_follow_message_logging_opt_out(gateway, tmp_path, team_level=True) + + +def test_a_line_body_system_prompt_leads_the_logged_messages(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "batch-line-system-" + uuid.uuid4().hex[:12] + system_prompt: Final = f"{marker} answer tersely" + + with ( + wire_server(lambda _request: Reply()) as endpoint, + owned_proxy( + gateway, + tmp_path, + {"GENERIC_LOGGER_ENDPOINT": endpoint.url}, + config=_line_item_config(Path("tests/integration/proxy_config.yaml"), tmp_path / "line_items_system.yaml"), + ) as candidate, + candidate.scenario() as scenario, + ): + routed_model: Final = scenario.model( + api_base=f"{gateway.upstream_url}/{marker}", api_key="synthetic-provider-key" ) - assert batch_response.status_code == 200, batch_response.text - batch_id: Final = string_value(JSON_OBJECT.validate_json(batch_response.content)["id"]) - retrieval: Final = eventually( - lambda: candidate.request("GET", f"/v1/batches/{batch_id}", key=key), - lambda response: response.status_code == 200 and response.json()["status"] == "completed", - seconds=30, + input_lines: Final = tuple( + json.dumps({**line, "body": {**line["body"], "system": system_prompt}}, separators=(",", ":")) + for line in (json.loads(raw) for raw in _input_lines(routed_model, marker)) ) - assert retrieval.status_code == 200, retrieval.text - batches: Final[ - list[Request] - ] = [] # mutable-ok: drain() consumes the queue, later polls must keep earlier batches + handle: Final = register_scenario(marker, _provider_routes(input_lines)) + scenario.cleanups.callback(delete_scenario, handle) + key: Final = scenario.key(models=[routed_model]) + _completed_batch(candidate, routed_model, input_lines, key) + received: Final[list[Request]] = [] # mutable-ok: drain() consumes the queue, later polls keep earlier batches + + for event in _all_lines_delivered(endpoint, received, key): + custom_id = string_value(_hidden(event)["batch_custom_id"]) + assert event["messages"] == [ + {"role": "system", "content": system_prompt}, + {"role": "user", "content": f"{marker} {custom_id}"}, + ], event["messages"] + assert event["call_type"] == "acompletion", event["call_type"] + assert "system" not in object_value(event["metadata"]), event["metadata"] - def delivered_lines() -> tuple[dict[str, JsonValue], ...]: - batches.extend(endpoint.drain()) - return tuple( - event - for batch in batches - for event in json.loads(batch.body) - if object_value(event["metadata"]).get("user_api_key_hash") == key_hash - and _hidden(event).get("batch_custom_id") is not None - ) - line_events: Final = eventually( - delivered_lines, - lambda values: {_hidden(event)["batch_custom_id"] for event in values} == set(ALL_CUSTOM_IDS), - seconds=40, +def test_a_restarted_proxy_does_not_emit_a_completed_batchs_lines_again(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "batch-line-restart-" + uuid.uuid4().hex[:12] + config: Final = _line_item_config(Path("tests/integration/proxy_config.yaml"), tmp_path / "line_items_restart.yaml") + received: Final[list[Request]] = [] # mutable-ok: drain() consumes the queue, later polls keep earlier batches + + with wire_server(lambda _request: Reply()) as endpoint, gateway.scenario() as scenario: + routed_model: Final = scenario.model( + api_base=f"{gateway.upstream_url}/{marker}", api_key="synthetic-provider-key" ) + input_lines: Final = _input_lines(routed_model, marker) + handle: Final = register_scenario(marker, _provider_routes(input_lines)) + scenario.cleanups.callback(delete_scenario, handle) + key: Final = scenario.key(models=[routed_model]) + with owned_proxy(gateway, tmp_path, {"GENERIC_LOGGER_ENDPOINT": endpoint.url}, config=config) as first: + batch_id: Final = _completed_batch(first, routed_model, input_lines, key) + first_lines: Final = _all_lines_delivered(endpoint, received, key) + assert sorted(string_value(_hidden(event)["batch_custom_id"]) for event in first_lines) == list(ALL_CUSTOM_IDS) + + with owned_proxy(gateway, tmp_path, {"GENERIC_LOGGER_ENDPOINT": endpoint.url}, config=config) as second: + retrieval: Final = second.request("GET", f"/v1/batches/{batch_id}", key=key) + assert retrieval.status_code == 200 and retrieval.json()["status"] == "completed", retrieval.text + events_after_restart: Final = eventually( + lambda: _delivered(endpoint, received, key), + lambda values: len([e for e in values if e["call_type"] == "aretrieve_batch"]) >= 2, + seconds=40, + ) - for event in line_events: - assert event["messages"] == [{"role": "user", "content": "redacted-by-litellm"}], event["messages"] - assert "answer " not in json.dumps(event["response"]), event["response"] + assert sorted( + string_value(_hidden(event)["batch_custom_id"]) for event in _line_events(events_after_restart) + ) == list(ALL_CUSTOM_IDS), [ + (event["call_type"], _hidden(event).get("batch_custom_id")) for event in events_after_restart + ] From 2b30a092fafbbac7bba5b8912703b9a2d20eb7f0 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Mon, 12 Oct 2026 03:15:14 -0700 Subject: [PATCH 65/89] fix(batches): keep in-memory line claims for 10,000 batches and pin managed batch emission across poller and retrieves --- litellm/batches/batch_line_item_logging.py | 11 ++- .../spend/test_batch_line_item_callbacks.py | 86 ++++++++++++++++++- .../batches/test_batch_line_item_logging.py | 11 +++ 3 files changed, 106 insertions(+), 2 deletions(-) diff --git a/litellm/batches/batch_line_item_logging.py b/litellm/batches/batch_line_item_logging.py index 0f3b4e2e3215..d46548f082d2 100644 --- a/litellm/batches/batch_line_item_logging.py +++ b/litellm/batches/batch_line_item_logging.py @@ -22,6 +22,7 @@ _uses_native_vertex_output, # pyright: ignore[reportPrivateUsage] # same reuse ) from litellm.caching.caching import DualCache +from litellm.caching.in_memory_cache import InMemoryCache from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.utils import ( EmbeddingResponse, @@ -74,7 +75,15 @@ def _supported_line_provider(value: str) -> _BatchLineProvider | None: _LINE_ITEM_CLAIM_TTL_SECONDS: Final = 30 * 24 * 60 * 60 -batch_line_item_claim_cache: Final = DualCache() +# Without Redis each worker holds its claims in memory for the claim ttl, so the +# in-memory tier keeps far more than its 200-entry default before evicting a claim +_IN_MEMORY_CLAIM_CAPACITY: Final = 10_000 + +batch_line_item_claim_cache: Final = DualCache( + in_memory_cache=InMemoryCache( + max_size_in_memory=_IN_MEMORY_CLAIM_CAPACITY, default_ttl=_LINE_ITEM_CLAIM_TTL_SECONDS + ) +) # asyncio holds tasks weakly, so each spawned line fan-out stays referenced here until it finishes _LINE_ITEM_TASKS: Final[set["asyncio.Task[int]"]] = set() # mutable-ok: strong refs to running fan-outs diff --git a/tests/integration/spend/test_batch_line_item_callbacks.py b/tests/integration/spend/test_batch_line_item_callbacks.py index 0e45c1221146..30b80f84bbc6 100644 --- a/tests/integration/spend/test_batch_line_item_callbacks.py +++ b/tests/integration/spend/test_batch_line_item_callbacks.py @@ -10,7 +10,7 @@ import pytest import yaml from integration._support.client import JSON_OBJECT, Gateway, eventually, object_value, string_value -from integration._support.database import read_rows +from integration._support.database import read_rows, scratch_database from integration._support.process import owned_proxy from integration._support.upstream import delete_scenario, register_scenario from integration._support.wire import Reply, Request, Wire, wire_server @@ -24,6 +24,9 @@ ERROR_FILE_IDS: Final = ("r4", "r5") ALL_CUSTOM_IDS: Final = ("r1", "r2", "r3", "r4", "r5") LIVE_TOKENS: Final = 10 +_PROCESSED_BATCH_SQL: Final = ( + 'SELECT unified_object_id FROM "LiteLLM_ManagedObjectTable" WHERE unified_object_id=%s AND batch_processed' +) def _successful_line(custom_id: str) -> str: @@ -610,3 +613,84 @@ def test_a_restarted_proxy_does_not_emit_a_completed_batchs_lines_again(gateway: ) == list(ALL_CUSTOM_IDS), [ (event["call_type"], _hidden(event).get("batch_custom_id")) for event in events_after_restart ] + + +def test_a_managed_batch_emits_its_lines_once_across_the_poller_and_retrieves(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "batch-line-managed-" + uuid.uuid4().hex[:12] + received: Final[list[Request]] = [] # mutable-ok: drain() consumes the queue, later polls keep earlier batches + + with ( + scratch_database() as database_url, + wire_server(lambda _request: Reply()) as endpoint, + owned_proxy( + gateway, + tmp_path, + {"GENERIC_LOGGER_ENDPOINT": endpoint.url, "DATABASE_URL": database_url}, + config=_line_item_config(Path("tests/integration/proxy_config.yaml"), tmp_path / "line_items_managed.yaml"), + ) as candidate, + candidate.scenario() as scenario, + ): + routed_model: Final = scenario.model( + api_base=f"{gateway.upstream_url}/{marker}", api_key="synthetic-provider-key" + ) + input_lines: Final = _input_lines(routed_model, marker) + handle: Final = register_scenario(marker, _provider_routes(input_lines)) + scenario.cleanups.callback(delete_scenario, handle) + key: Final = scenario.key(user_id=scenario.user(user_role="internal_user"), models=[routed_model]) + uploaded: Final = candidate.request_multipart( + "/v1/files", + {"purpose": "batch", "target_model_names": routed_model}, + {"file": ("in.jsonl", ("\n".join(input_lines) + "\n").encode(), "application/jsonl")}, + key=key, + ) + assert uploaded.status_code == 200, uploaded.text + created: Final = candidate.request( + "POST", + "/v1/batches", + { + "input_file_id": string_value(JSON_OBJECT.validate_json(uploaded.content)["id"]), + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + "model": routed_model, + }, + key=key, + ) + assert created.status_code == 200, created.text + batch_id: Final = string_value(JSON_OBJECT.validate_json(created.content)["id"]) + + def line_events() -> tuple[dict[str, JsonValue], ...]: + received.extend(endpoint.drain()) + return tuple( + event + for batch in received + for event in json.loads(batch.body) + if _hidden(event).get("batch_custom_id") is not None and marker in json.dumps(event["messages"]) + ) + + retrievals: Final = [candidate.request("GET", f"/v1/batches/{batch_id}", key=key) for _ in range(3)] + assert all(response.status_code == 200 for response in retrievals), [r.text for r in retrievals] + eventually( + line_events, + lambda values: {_hidden(event)["batch_custom_id"] for event in values} == set(ALL_CUSTOM_IDS), + seconds=60, + ) + later_retrievals: Final = [candidate.request("GET", f"/v1/batches/{batch_id}", key=key) for _ in range(3)] + assert all(response.status_code == 200 for response in later_retrievals), [r.text for r in later_retrievals] + eventually(lambda: read_rows(_PROCESSED_BATCH_SQL, (batch_id,), database_url=database_url), bool, seconds=30) + assert _live_chat(candidate, routed_model, key) == 200 + live_events: Final = eventually( + lambda: tuple( + event + for event in _delivered(endpoint, received, key) + if event["call_type"] == "acompletion" and _hidden(event).get("batch_custom_id") is None + ), + bool, + seconds=30, + ) + assert len(live_events) == 1, live_events + settled: Final = line_events() + + assert sorted(string_value(_hidden(event)["batch_custom_id"]) for event in settled) == list(ALL_CUSTOM_IDS), [ + (event["call_type"], _hidden(event).get("batch_custom_id")) for event in settled + ] + assert {_hidden(event)["batch_id"] for event in settled} == {batch_id}, settled diff --git a/tests/unit/batches/test_batch_line_item_logging.py b/tests/unit/batches/test_batch_line_item_logging.py index 7867ffad5647..c6cdfdd583e9 100644 --- a/tests/unit/batches/test_batch_line_item_logging.py +++ b/tests/unit/batches/test_batch_line_item_logging.py @@ -26,6 +26,7 @@ import litellm from litellm.batches.batch_line_item_logging import ( + _claim_line_items, _release_line_item_claim, batch_line_item_claim_cache, drain_batch_line_items, @@ -1083,6 +1084,16 @@ async def test_line_items_emit_once_per_batch_id(recorder): assert len(recorder.failure_events) == 1 +@pytest.mark.asyncio +async def test_an_in_memory_claim_survives_a_thousand_newer_batches(): + first: Final = await _claim_line_items(batch_line_item_claim_cache, "batch_line_items_emitted:batch_old", "old") + for index in range(1_000): + await _claim_line_items(batch_line_item_claim_cache, f"batch_line_items_emitted:batch_{index}", str(index)) + again: Final = await _claim_line_items(batch_line_item_claim_cache, "batch_line_items_emitted:batch_old", "again") + + assert (first, again) == ("claimed", "already_claimed") + + @pytest.mark.asyncio async def test_line_items_claim_is_scoped_to_the_batch_id(recorder): file_mock: Final = AsyncMock(side_effect=_file_content) From 08651232eda4e9841d8cb825007043fcb43345dd Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Mon, 12 Oct 2026 03:30:54 -0700 Subject: [PATCH 66/89] Keep batch line unit tests in-process and drop repeated per-sink comments --- .../pagerduty/pagerduty.py | 1 - litellm/batches/batch_line_item_logging.py | 1 - .../SlackAlerting/slack_alerting.py | 2 - .../clickhouse/clickhouse_spend_logger.py | 4 -- .../datadog/datadog_cost_management.py | 2 - .../integrations/datadog/datadog_metrics.py | 1 - litellm/integrations/lago.py | 1 - .../integrations/newrelic/newrelic_metrics.py | 1 - litellm/integrations/openmeter.py | 1 - litellm/integrations/opentelemetry.py | 3 - litellm/integrations/otel/plumbing/metrics.py | 6 -- litellm/integrations/prometheus.py | 2 - litellm/integrations/shadow_eval_logger.py | 2 - litellm/proxy/proxy_server.py | 1 - .../batches/test_batch_line_item_logging.py | 64 +++++++------------ .../test_clickhouse_spend_logger.py | 2 +- .../test_batch_line_item_metering_sinks.py | 5 +- .../test_litellm_logging.py | 12 ++-- .../hooks/test_parallel_request_limiter.py | 4 +- .../hooks/test_parallel_request_limiter_v3.py | 8 +-- .../hooks/test_proxy_track_cost_callback.py | 4 +- tests/unit/test_router/test_router.py | 2 +- 22 files changed, 41 insertions(+), 88 deletions(-) diff --git a/enterprise/litellm_enterprise/enterprise_callbacks/pagerduty/pagerduty.py b/enterprise/litellm_enterprise/enterprise_callbacks/pagerduty/pagerduty.py index 0ea4eadc3c01..c437345a98d7 100644 --- a/enterprise/litellm_enterprise/enterprise_callbacks/pagerduty/pagerduty.py +++ b/enterprise/litellm_enterprise/enterprise_callbacks/pagerduty/pagerduty.py @@ -87,7 +87,6 @@ async def async_log_failure_event(self, kwargs, response_obj, start_time, end_ti Record a failure event. Only send an alert to PagerDuty if the configured *failure* threshold is exceeded in the specified window. """ - # A failed batch line item is reported by its batch, not by live traffic if is_batch_line_item_event(kwargs): return now = datetime.now(timezone.utc) diff --git a/litellm/batches/batch_line_item_logging.py b/litellm/batches/batch_line_item_logging.py index d46548f082d2..f508b50ea5f4 100644 --- a/litellm/batches/batch_line_item_logging.py +++ b/litellm/batches/batch_line_item_logging.py @@ -90,7 +90,6 @@ def _supported_line_provider(value: str) -> _BatchLineProvider | None: # loops whose shutdown drain has started, where a new fan-out would be cancelled with its claim held _DRAINED_LOOPS: Final[weakref.WeakSet[asyncio.AbstractEventLoop]] = weakref.WeakSet() -# Body keys carrying the system prompt: Anthropic and Bedrock `system`, Responses `instructions` _SYSTEM_PROMPT_KEYS: Final = ("system", "instructions") # Anthropic keeps no input file, so its retrieved batch reports input_file_id as the string "None" diff --git a/litellm/integrations/SlackAlerting/slack_alerting.py b/litellm/integrations/SlackAlerting/slack_alerting.py index 31a15d3fd272..43c6a7a77a74 100644 --- a/litellm/integrations/SlackAlerting/slack_alerting.py +++ b/litellm/integrations/SlackAlerting/slack_alerting.py @@ -1689,8 +1689,6 @@ async def periodic_flush(self): async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): """Log deployment latency""" - # A batch line item ran inside the provider's batch job, not on the deployment, - # and its timing spans the whole batch, so it says nothing about deployment health if is_batch_line_item_event(kwargs): return try: diff --git a/litellm/integrations/clickhouse/clickhouse_spend_logger.py b/litellm/integrations/clickhouse/clickhouse_spend_logger.py index 3eab03bd76f0..cb2f4f0bcf8b 100644 --- a/litellm/integrations/clickhouse/clickhouse_spend_logger.py +++ b/litellm/integrations/clickhouse/clickhouse_spend_logger.py @@ -206,15 +206,11 @@ class ClickHouseSpendLogger(ClickHouseBatchLogger): table = SPEND_LOGS_TABLE async def async_log_success_event(self, kwargs, response_obj, start_time, end_time) -> None: - # Batch line items are billed by the aggregate aretrieve_batch row; a per-line - # spend row here would bill the batch twice. if is_batch_line_item_event(kwargs): return self._log(kwargs) async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time) -> None: - # Failed batch lines are part of the same aggregate row's request counts; - # per-line rows here would inflate spend-log request counts. if is_batch_line_item_event(kwargs): return self._log(kwargs) diff --git a/litellm/integrations/datadog/datadog_cost_management.py b/litellm/integrations/datadog/datadog_cost_management.py index 73017b79f7b3..1584cfca1c43 100644 --- a/litellm/integrations/datadog/datadog_cost_management.py +++ b/litellm/integrations/datadog/datadog_cost_management.py @@ -73,8 +73,6 @@ def __init__(self, cost_tag_keys: list[str] | None = None, **kwargs): super().__init__(**kwargs) async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): - # A batch line item is billed by the aggregate aretrieve_batch event; a - # per-line FOCUS BilledCost row would double-count cloud spend. if is_batch_line_item_event(kwargs): return try: diff --git a/litellm/integrations/datadog/datadog_metrics.py b/litellm/integrations/datadog/datadog_metrics.py index da045611e9a2..eae6bda0dc35 100644 --- a/litellm/integrations/datadog/datadog_metrics.py +++ b/litellm/integrations/datadog/datadog_metrics.py @@ -171,7 +171,6 @@ def _add_metrics_from_log( self.log_queue.append(series_count) async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): - # A batch line item is counted by the aggregate aretrieve_batch event if is_batch_line_item_event(kwargs): return try: diff --git a/litellm/integrations/lago.py b/litellm/integrations/lago.py index 193f809a5fd7..5b858f5901da 100644 --- a/litellm/integrations/lago.py +++ b/litellm/integrations/lago.py @@ -122,7 +122,6 @@ def _common_logic(self, kwargs: dict, response_obj) -> dict: return returned_val def log_success_event(self, kwargs, response_obj, start_time, end_time): - # A batch line item is billed by the aggregate aretrieve_batch event. if is_batch_line_item_event(kwargs): return _url = os.getenv("LAGO_API_BASE") diff --git a/litellm/integrations/newrelic/newrelic_metrics.py b/litellm/integrations/newrelic/newrelic_metrics.py index 8bd77c981f63..c8b2c6be3ba8 100644 --- a/litellm/integrations/newrelic/newrelic_metrics.py +++ b/litellm/integrations/newrelic/newrelic_metrics.py @@ -302,7 +302,6 @@ async def periodic_flush(self) -> None: await self._final_drain() async def async_log_success_event(self, kwargs, response_obj, start_time, end_time) -> None: - # A batch line item is metered by the aggregate aretrieve_batch event. if is_batch_line_item_event(kwargs): return try: diff --git a/litellm/integrations/openmeter.py b/litellm/integrations/openmeter.py index 029c684f9693..d0f1dab988ce 100644 --- a/litellm/integrations/openmeter.py +++ b/litellm/integrations/openmeter.py @@ -98,7 +98,6 @@ def _common_logic(self, kwargs: dict, response_obj): } def log_success_event(self, kwargs, response_obj, start_time, end_time): - # A batch line item is billed by the aggregate aretrieve_batch event. if is_batch_line_item_event(kwargs): return _url = os.getenv("OPENMETER_API_ENDPOINT", "https://openmeter.cloud") diff --git a/litellm/integrations/opentelemetry.py b/litellm/integrations/opentelemetry.py index 82bd43b477b2..0af657442777 100644 --- a/litellm/integrations/opentelemetry.py +++ b/litellm/integrations/opentelemetry.py @@ -1670,9 +1670,6 @@ def _metric_attribute_allowed(self, key: str) -> bool: return True def _record_metrics(self, kwargs, response_obj, start_time, end_time): - # A batch line item's tokens and cost are metered by the aggregate - # aretrieve_batch event; per-line samples here would double-count them. - # Spans for line items are still emitted by _handle_success. if is_batch_line_item_event(kwargs): return duration_s: Final = (end_time - start_time).total_seconds() diff --git a/litellm/integrations/otel/plumbing/metrics.py b/litellm/integrations/otel/plumbing/metrics.py index 9c0794576810..e94f6ab9f25a 100644 --- a/litellm/integrations/otel/plumbing/metrics.py +++ b/litellm/integrations/otel/plumbing/metrics.py @@ -222,10 +222,6 @@ def record( start_time: datetime, end_time: datetime, ) -> None: - # A batch line item is metered by the aggregate aretrieve_batch event, and - # its child interval (parent retrieve start -> line emission) measures - # retrieval and callback processing rather than the line's model call, so - # skip every metric for it; spans for line items are still emitted. if is_batch_line_item_event(kwargs): return common_attrs: Final = self._filter_attributes(self._bounded_attributes(kwargs)) @@ -253,8 +249,6 @@ def record_failure( start_time: datetime, end_time: datetime, ) -> None: - # A batch line item's interval measures retrieval and emission, not the - # line's model call, so its duration would be synthetic here too. if is_batch_line_item_event(kwargs): return """Record the one metric a failed request can honestly report: the diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index fe7a0e622a83..b90e571a5c9c 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -1507,8 +1507,6 @@ def _inc_labeled_counter( self._track_end_user_metric_series(counter, metric_name, _labels) async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): - # A batch line item is metered by the aggregate aretrieve_batch event; a - # per-line sample here would double-count spend and requests. if is_batch_line_item_event(kwargs): return # Define prometheus client diff --git a/litellm/integrations/shadow_eval_logger.py b/litellm/integrations/shadow_eval_logger.py index 7f220c17eb68..d4fbc5d97e8f 100644 --- a/litellm/integrations/shadow_eval_logger.py +++ b/litellm/integrations/shadow_eval_logger.py @@ -949,8 +949,6 @@ async def async_log_success_event( *, guardrail_snapshot: GuardrailRequestSnapshot | None = None, ) -> None: - # A batch line item was served by the batch API, so replaying it against a - # shadow model would buy realtime calls for traffic that never went realtime if is_batch_line_item_event(kwargs): return try: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 3842b43a9703..edb852a6becb 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -7067,7 +7067,6 @@ async def load_config(self, router: litellm.Router | None, config_file_path: str health_check_interval = general_settings.get("health_check_interval", DEFAULT_HEALTH_CHECK_INTERVAL) health_check_concurrency = general_settings.get("health_check_concurrency", None) health_check_details = general_settings.get("health_check_details", True) - ### BATCH LINE ITEM CALLBACKS ### _store_batch_line_items: Final = coerce_bool(general_settings.get("store_batch_line_items_in_callbacks")) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # general_settings is untyped upstream if _store_batch_line_items is not None: litellm.store_batch_line_items_in_callbacks = bool(_store_batch_line_items) diff --git a/tests/unit/batches/test_batch_line_item_logging.py b/tests/unit/batches/test_batch_line_item_logging.py index c6cdfdd583e9..a1d51bd3e116 100644 --- a/tests/unit/batches/test_batch_line_item_logging.py +++ b/tests/unit/batches/test_batch_line_item_logging.py @@ -13,9 +13,7 @@ import asyncio import json -import subprocess -import sys -import textwrap +import threading import uuid from datetime import datetime from types import MappingProxyType, SimpleNamespace @@ -32,6 +30,7 @@ drain_batch_line_items, log_batch_line_items, pending_batch_line_item_tasks, + spawn_batch_line_items, ) from litellm.batches.batch_utils import BatchResultFiles from litellm.caching.caching import DualCache @@ -165,7 +164,7 @@ def _parent_logging(custom_llm_provider: str = "openai") -> Logging: messages=[{"role": "user", "content": ""}], stream=False, call_type="aretrieve_batch", - start_time=datetime.now(), + start_time=datetime(2026, 1, 1), litellm_call_id=str(uuid.uuid4()), function_id=str(uuid.uuid4()), ) @@ -1635,52 +1634,35 @@ async def test_a_retrieve_after_the_shutdown_drain_leaves_every_line_to_the_next assert next_retrieve == 2 -_EXIT_FLUSH_SPAWN_SCRIPT: Final = textwrap.dedent( - """ - import asyncio - import atexit - from types import SimpleNamespace - - from litellm.batches.batch_line_item_logging import spawn_batch_line_items - from litellm.types.utils import LiteLLMBatch - - - async def spawn(): - return spawn_batch_line_items( +@pytest.mark.asyncio +async def test_a_retrieve_flushed_at_interpreter_exit_starts_no_line_fan_out(monkeypatch): + """The logging worker's exit flush runs a queued retrieve handler on a loop it + closes as soon as the handler returns, so a fan-out started there would be + destroyed partway with the batch claim held.""" + exited_main_thread: Final = threading.Thread(target=lambda: None) + exited_main_thread.start() + exited_main_thread.join() + with monkeypatch.context() as interpreter_exit: + interpreter_exit.setattr(threading, "main_thread", lambda: exited_main_thread) + spawned: Final = spawn_batch_line_items( batch=LiteLLMBatch( - id="batch_exit", object="batch", endpoint="/v1/chat/completions", input_file_id="in", - status="completed", completion_window="24h", created_at=1, + id="batch_exit", + object="batch", + endpoint="/v1/chat/completions", + input_file_id="in", + status="completed", + completion_window="24h", + created_at=1, ), custom_llm_provider="openai", - parent=SimpleNamespace(litellm_params={}), + parent=_parent_logging(), model_name="gpt-4o", litellm_params=None, model_info=None, result_files=None, ) - - def flush_at_exit(): - loop = asyncio.new_event_loop() - print("spawned:", loop.run_until_complete(spawn()), flush=True) - loop.close() - - - atexit.register(flush_at_exit) - """ -) - - -def test_a_retrieve_flushed_at_interpreter_exit_starts_no_line_fan_out(): - """The logging worker's exit flush runs a queued retrieve handler on a loop it - closes as soon as the handler returns, so a fan-out started there would be - destroyed partway with the batch claim held.""" - result: Final = subprocess.run( - [sys.executable, "-P", "-c", _EXIT_FLUSH_SPAWN_SCRIPT], capture_output=True, text=True, timeout=120 - ) - - assert result.returncode == 0, result.stderr - assert "spawned: None" in result.stdout, result.stdout + result.stderr + assert (spawned, pending_batch_line_item_tasks()) == (None, ()) @pytest.mark.asyncio diff --git a/tests/unit/integrations/clickhouse/test_clickhouse_spend_logger.py b/tests/unit/integrations/clickhouse/test_clickhouse_spend_logger.py index 4cb05f896e1b..5d691d8652bd 100644 --- a/tests/unit/integrations/clickhouse/test_clickhouse_spend_logger.py +++ b/tests/unit/integrations/clickhouse/test_clickhouse_spend_logger.py @@ -518,7 +518,7 @@ async def test_batch_line_item_success_event_does_not_write_spend(monkeypatch: p monkeypatch.setattr(litellm, "store_batch_line_items_in_callbacks", True, raising=False) storage = MagicMock() logger = ClickHouseSpendLogger(storage=storage) - now = datetime.now(timezone.utc) + now = datetime(2026, 1, 1, tzinfo=timezone.utc) await logger.async_log_success_event( { diff --git a/tests/unit/integrations/test_batch_line_item_metering_sinks.py b/tests/unit/integrations/test_batch_line_item_metering_sinks.py index 5430ab969d5a..1121332e8fe3 100644 --- a/tests/unit/integrations/test_batch_line_item_metering_sinks.py +++ b/tests/unit/integrations/test_batch_line_item_metering_sinks.py @@ -18,7 +18,6 @@ import io import json import os -import time import uuid from contextlib import redirect_stdout from datetime import datetime @@ -110,7 +109,7 @@ def _parent_logging() -> Logging: messages=[{"role": "user", "content": ""}], stream=False, call_type="aretrieve_batch", - start_time=datetime.now(), + start_time=datetime(2026, 1, 1), litellm_call_id=str(uuid.uuid4()), function_id=str(uuid.uuid4()), ) @@ -258,7 +257,7 @@ async def test_newrelic_metrics_meters_only_the_aggregate_batch_event(monkeypatc f"New Relic metered per-line costs on top of the aggregate: {costs}; " "line items must be metered only by the aggregate aretrieve_batch event" ) - now = time.time() + now = 1_767_225_600.0 envelopes = build_metric_payload(records=records, window_start=now - 1, now=now) cost_sum = 0.0 for envelope in envelopes: diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index af3c8f0c6c9e..a97cf9cf6c37 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -338,8 +338,8 @@ async def async_post_mcp_tool_call_hook(self, kwargs, response_obj, start_time, returned = await logging_obj.async_post_mcp_tool_call_hook( kwargs={"original_response": result}, response_obj=result, - start_time=datetime.datetime.now(), - end_time=datetime.datetime.now(), + start_time=datetime.datetime(2026, 1, 1), + end_time=datetime.datetime(2026, 1, 1), ) assert returned is result assert result.content == [TextContent(type="text", text="native-safe" if same_content else "final-safe")] @@ -367,8 +367,8 @@ async def async_post_mcp_tool_call_hook(self, kwargs, response_obj, start_time, returned = await logging_obj.async_post_mcp_tool_call_hook( kwargs={"original_response": result}, response_obj=result, - start_time=datetime.datetime.now(), - end_time=datetime.datetime.now(), + start_time=datetime.datetime(2026, 1, 1), + end_time=datetime.datetime(2026, 1, 1), ) assert returned is result assert result.content == [TextContent(type="text", text="[REDACTED]")] @@ -2951,7 +2951,7 @@ async def async_logging_hook( messages=[], stream=stream, call_type="anthropic_messages", - start_time=datetime.datetime.now(), + start_time=datetime.datetime(2026, 1, 1), litellm_call_id="private-snapshot", function_id="private-snapshot", dynamic_async_success_callbacks=[shadow, RecordingLogger(), guardrail], @@ -9724,7 +9724,7 @@ def _logging_obj() -> LitellmLogging: messages=[{"role": "user", "content": "Hey"}], stream=False, call_type="aretrieve_batch", - start_time=time.time(), + start_time=1_767_225_600.0, litellm_call_id="batch-call-reuse", function_id="f", ) diff --git a/tests/unit/proxy/hooks/test_parallel_request_limiter.py b/tests/unit/proxy/hooks/test_parallel_request_limiter.py index 25e3ee101a1d..50d1c8195a76 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter.py @@ -185,8 +185,8 @@ async def test_async_log_failure_event_skips_batch_line_item_events(): "model": "gpt-3.5-turbo", }, response_obj=None, - start_time=datetime.now(), - end_time=datetime.now(), + start_time=datetime(2026, 1, 1), + end_time=datetime(2026, 1, 1), ) assert local_cache.cache_dict == {} diff --git a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py index 84ed0dc5f16f..2467d20ae453 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py @@ -7815,8 +7815,8 @@ async def test_async_log_success_event_skips_batch_line_item_events(): id="x", object="chat.completion", created=1, model="gpt-3.5-turbo", usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15), choices=[], ), - start_time=datetime.now(), - end_time=datetime.now(), + start_time=datetime(2026, 1, 1), + end_time=datetime(2026, 1, 1), ) assert local_cache.in_memory_cache.cache_dict == {} @@ -7844,8 +7844,8 @@ async def test_async_log_failure_event_skips_batch_line_item_events(): "model": "gpt-3.5-turbo", }, response_obj=None, - start_time=datetime.now(), - end_time=datetime.now(), + start_time=datetime(2026, 1, 1), + end_time=datetime(2026, 1, 1), ) assert local_cache.in_memory_cache.cache_dict == {} diff --git a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py index 4e8f32424c0f..df4c7f93efa9 100644 --- a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py @@ -2840,8 +2840,8 @@ async def test_batch_line_item_event_never_updates_spend(): # test-quality-ok: await logger.async_log_success_event( kwargs=kwargs, response_obj=ModelResponse(), - start_time=datetime.now(), - end_time=datetime.now(), + start_time=datetime(2026, 1, 1), + end_time=datetime(2026, 1, 1), ) mock_track.assert_not_awaited() diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 7d793cd8b0ba..442e777c0196 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -1324,7 +1324,7 @@ async def test_deployment_callbacks_skip_batch_line_items(monkeypatch: pytest.Mo } ] ) - now = datetime.now() + now = datetime(2026, 1, 1) line_item_kwargs = { "call_type": "acompletion", "standard_logging_object": {"total_tokens": _BATCH_TOKENS_PER_ROW}, From 56b60e62073838ce148c934dbc8782775d2241fa Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Mon, 12 Oct 2026 03:34:17 -0700 Subject: [PATCH 67/89] Share one batch provider alias across batch_utils and line item logging --- litellm/batches/batch_line_item_logging.py | 18 +++++------- litellm/batches/batch_utils.py | 32 ++++++++++------------ 2 files changed, 22 insertions(+), 28 deletions(-) diff --git a/litellm/batches/batch_line_item_logging.py b/litellm/batches/batch_line_item_logging.py index f508b50ea5f4..845497074999 100644 --- a/litellm/batches/batch_line_item_logging.py +++ b/litellm/batches/batch_line_item_logging.py @@ -13,6 +13,7 @@ from litellm._internal_context import with_service_target from litellm._logging import verbose_logger from litellm.batches.batch_utils import ( + BatchProvider, BatchResultFiles, _batch_response_was_successful, # pyright: ignore[reportPrivateUsage] # batch-internal helper shared with the aggregate cost path by design _fetch_batch_managed_file_content, # pyright: ignore[reportPrivateUsage, reportUnknownVariableType] # same reuse; helper is untyped upstream @@ -35,11 +36,6 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging -_BatchLineProvider: TypeAlias = Literal[ - "openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral" -] - - _SECRET_PARAM_KEYS: Final = frozenset( { "api_key", @@ -55,7 +51,7 @@ ) -def _supported_line_provider(value: str) -> _BatchLineProvider | None: +def _supported_line_provider(value: str) -> BatchProvider | None: match value: case "openai" | "azure" | "vertex_ai" | "hosted_vllm" | "anthropic" | "bedrock" | "mistral": return value @@ -203,7 +199,7 @@ def _line_status_code(entry: Mapping[str, object], custom_llm_provider: str) -> return None -def _line_error_payload(entry: Mapping[str, object], custom_llm_provider: _BatchLineProvider) -> object: +def _line_error_payload(entry: Mapping[str, object], custom_llm_provider: BatchProvider) -> object: if custom_llm_provider == "anthropic": return ( (_as_object_mapping(entry.get("result")) or _EMPTY_BODY).get("error") or entry.get("result") or _EMPTY_BODY @@ -246,7 +242,7 @@ def _line_model( def _line_result( call_type: str, - custom_llm_provider: _BatchLineProvider, + custom_llm_provider: BatchProvider, model: str, response_body: Mapping[str, object], ) -> _BatchLineResult: @@ -273,7 +269,7 @@ def _line_result( def _line_result_or_none( request_call_type: str, - custom_llm_provider: _BatchLineProvider, + custom_llm_provider: BatchProvider, model: str, response_body: Mapping[str, object], custom_id: object, @@ -348,7 +344,7 @@ async def _emit_line_event( entry: Mapping[str, object], requests_by_id: Mapping[str, Mapping[str, object]], batch: LiteLLMBatch, - custom_llm_provider: _BatchLineProvider, + custom_llm_provider: BatchProvider, parent: "Logging", parent_params: Mapping[str, object], model_name: str | None, @@ -442,7 +438,7 @@ async def _emit_line_event( async def _fetch_managed_file_or_empty( file_id: str | None, - custom_llm_provider: _BatchLineProvider, + custom_llm_provider: BatchProvider, fetch_params: dict[str, object] | None, # mutable-ok: batch_utils file fetch takes the shared litellm_params dict ) -> bytes: if file_id is None or file_id in _NO_FILE_IDS: diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index b73a32686ad5..a27b69da0765 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -3,7 +3,7 @@ from dataclasses import dataclass from dataclasses import replace as dataclasses_replace from enum import Enum -from typing import Any, Final, Literal, cast +from typing import Any, Final, Literal, TypeAlias, cast import litellm from litellm._logging import verbose_logger @@ -20,6 +20,8 @@ from litellm.types.workload_identity import ANTHROPIC_WIF_KWARGS_KEYS from litellm.utils import token_counter +BatchProvider: TypeAlias = Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral"] + @dataclass(frozen=True, slots=True) class BatchCostUsageResult: @@ -79,7 +81,7 @@ def batch_cost_is_final(batch: Batch) -> bool: async def calculate_batch_cost_and_usage( file_content_dictionary: list[dict[str, object]], - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral"], + custom_llm_provider: BatchProvider, model_name: str | None = None, model_info: ModelInfo | None = None, ) -> BatchCostUsageResult: @@ -106,7 +108,7 @@ async def calculate_batch_cost_and_usage( async def handle_completed_batch( batch: Batch, - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral"], + custom_llm_provider: BatchProvider, model_name: str | None = None, litellm_params: dict[str, object] | None = None, model_info: ModelInfo | None = None, @@ -137,7 +139,7 @@ async def handle_completed_batch( async def handle_completed_batch_with_files( batch: Batch, - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral"], + custom_llm_provider: BatchProvider, model_name: str | None = None, litellm_params: Mapping[str, object] | None = None, model_info: ModelInfo | None = None, @@ -215,7 +217,7 @@ class _BatchOutputLineStats: def _classify_output_line_stats( entries: Iterable[dict[str, object]], - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral"], + custom_llm_provider: BatchProvider, model_name: str | None, model_info: ModelInfo | None, ) -> Iterator[_BatchOutputLineStats | _LineOutcome]: @@ -234,7 +236,7 @@ def _classify_output_line_stats( def _safe_output_line_stats( entry: Mapping[str, object], - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral"], + custom_llm_provider: BatchProvider, model_name: str | None, model_info: ModelInfo | None, ) -> _BatchOutputLineStats | None: @@ -256,7 +258,7 @@ def _safe_output_line_stats( def _compute_output_line_stats( entry: Mapping[str, object], - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral"], + custom_llm_provider: BatchProvider, model_name: str | None, model_info: ModelInfo | None, ) -> _BatchOutputLineStats: @@ -298,7 +300,7 @@ def _ocr_usage_info_from_response_body(response_body: Mapping[str, object]) -> O def _output_line_cost( response_body: Mapping[str, object], usage: Usage, - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral"], + custom_llm_provider: BatchProvider, model_name: str | None, response_model: str | None, model_info: ModelInfo | None, @@ -327,7 +329,7 @@ def _output_line_cost( def _aggregate_batch_cost_usage_models( entries: Iterable[dict[str, object]], - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral"], + custom_llm_provider: BatchProvider, model_name: str | None = None, model_info: ModelInfo | None = None, ) -> BatchCostUsageResult: @@ -466,9 +468,7 @@ def _provider_output_file_id(output_file_id: str) -> str: async def _fetch_batch_managed_file_content( file_id: str, - custom_llm_provider: Literal[ - "openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral" - ] = "openai", + custom_llm_provider: BatchProvider = "openai", litellm_params: Mapping[str, object] | None = None, ) -> bytes: """ @@ -499,9 +499,7 @@ async def _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", "bedrock", "mistral" - ] = "openai", + custom_llm_provider: BatchProvider = "openai", litellm_params: Mapping[str, object] | None = None, ) -> bytes: """ @@ -523,7 +521,7 @@ async def _fetch_batch_output_file_content( async def fetch_batch_error_file_content( batch: Batch, - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral"], + custom_llm_provider: BatchProvider, litellm_params: Mapping[str, object] | None, ) -> bytes | None: """Fetch the batch's separate error file bytes; None when it has none or the fetch fails.""" @@ -544,7 +542,7 @@ def count_error_file_failed_requests_from_content(content: bytes | None) -> int: async def count_error_file_failed_requests( batch: Batch, - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral"], + custom_llm_provider: BatchProvider, litellm_params: Mapping[str, object] | None, ) -> int: """Count failed requests reported only in the batch's separate error file. From 0470177ab0c192f7b879674b151704c432c37439 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Mon, 12 Oct 2026 03:49:51 -0700 Subject: [PATCH 68/89] Report Bedrock line status codes, skip unknown Bedrock lines without raising, fetch the input file after the native Vertex check --- litellm/batches/batch_line_item_logging.py | 34 ++++++++++++------- .../batches/test_batch_line_item_logging.py | 26 +++++++++++--- 2 files changed, 44 insertions(+), 16 deletions(-) diff --git a/litellm/batches/batch_line_item_logging.py b/litellm/batches/batch_line_item_logging.py index 845497074999..a44c375a5227 100644 --- a/litellm/batches/batch_line_item_logging.py +++ b/litellm/batches/batch_line_item_logging.py @@ -196,6 +196,13 @@ def _line_status_code(entry: Mapping[str, object], custom_llm_provider: str) -> result: Final = _as_object_mapping(entry.get("result")) if result is not None and result.get("type") == "succeeded": return 200 + if custom_llm_provider == "bedrock": + error: Final = _as_object_mapping(entry.get("error")) + error_code: Final = error.get("errorCode") if error is not None else None + if isinstance(error_code, int): + return error_code + if _batch_response_was_successful(entry, custom_llm_provider): + return 200 return None @@ -245,7 +252,7 @@ def _line_result( custom_llm_provider: BatchProvider, model: str, response_body: Mapping[str, object], -) -> _BatchLineResult: +) -> "_BatchLineResult | None": from litellm.types.utils import LlmProviders from litellm.utils import ProviderConfigManager @@ -255,7 +262,7 @@ def _line_result( if transformed is not None: return transformed if custom_llm_provider == "bedrock": - raise ValueError(f"unrecognized bedrock batch output line shape. keys={sorted(response_body)}") + return None if call_type == "aembedding": return EmbeddingResponse(**response_body) # pyright: ignore[reportArgumentType] # provider output bodies are dicts expanded as response ctor kwargs if call_type == "aresponses": @@ -275,14 +282,18 @@ def _line_result_or_none( custom_id: object, ) -> "_BatchLineResult | None": try: - return _line_result(request_call_type, custom_llm_provider, model, response_body) + result: Final = _line_result(request_call_type, custom_llm_provider, model, response_body) except Exception: # noqa: BLE001 # one unparseable line must not drop the rest of the batch's line events - verbose_logger.warning( - "batch output line could not be reconstructed as a %s response, skipping it. custom_id=%s", - request_call_type, - custom_id, - ) - return None + return _skipped_line(request_call_type, custom_id) + return result if result is not None else _skipped_line(request_call_type, custom_id) + + +def _skipped_line(request_call_type: str, custom_id: object) -> None: + verbose_logger.warning( + "batch output line could not be reconstructed as a %s response, skipping it. custom_id=%s", + request_call_type, + custom_id, + ) def _new_child_logging( @@ -586,9 +597,6 @@ async def log_batch_line_items( dict(internal_mapping) if internal_mapping is not None else litellm_params ) - input_file_content: Final = await _fetch_managed_file_or_empty(batch.input_file_id, line_provider, fetch_params) - requests_by_id: Final = _requests_by_custom_id(input_file_content) - output_content: Final = ( result_files.output if result_files is not None and result_files.output is not None @@ -601,6 +609,8 @@ async def log_batch_line_items( batch.id, ) return 0 + input_file_content: Final = await _fetch_managed_file_or_empty(batch.input_file_id, line_provider, fetch_params) + requests_by_id: Final = _requests_by_custom_id(input_file_content) error_content: Final = ( result_files.error if result_files is not None and result_files.error is not None diff --git a/tests/unit/batches/test_batch_line_item_logging.py b/tests/unit/batches/test_batch_line_item_logging.py index a1d51bd3e116..9100ab8399f3 100644 --- a/tests/unit/batches/test_batch_line_item_logging.py +++ b/tests/unit/batches/test_batch_line_item_logging.py @@ -561,6 +561,12 @@ async def test_input_fetch_failure_still_emits_aggregate(recorder): "modelInput": {"messages": [{"role": "user", "content": "hi br-conv"}]}, } ).encode(), + json.dumps( + { + "recordId": "br-err", + "modelInput": {"messages": [{"role": "user", "content": "hi br-err"}]}, + } + ).encode(), ] ) @@ -598,6 +604,13 @@ async def test_input_fetch_failure_still_emits_aggregate(recorder): }, } ).encode(), + json.dumps( + { + "recordId": "br-err", + "modelInput": {"messages": [{"role": "user", "content": "hi br-err"}]}, + "error": {"errorCode": 400, "errorMessage": "bad request"}, + } + ).encode(), ] ) @@ -882,6 +895,11 @@ async def test_line_items_bedrock_shapes(recorder): assert converse_line["prompt_tokens"] == 5 assert converse_line["completion_tokens"] == 6 + assert [_hidden(by_id[line])["batch_line_status_code"] for line in ("br-anth", "br-titan", "br-conv")] == [200] * 3 + (failure,) = recorder.failure_events + assert (_hidden(failure)["batch_custom_id"], _hidden(failure)["batch_line_status_code"]) == ("br-err", 400) + assert "bad request" in _payload(failure)["error_str"] + @pytest.mark.asyncio async def test_line_items_mistral_shape(recorder): @@ -1042,6 +1060,7 @@ async def test_line_items_native_vertex_rows_are_skipped(recorder): assert len(recorder.success_events) == 1 assert len(recorder.failure_events) == 0 assert _hidden(recorder.success_events[0]).get("batch_custom_id") is None + assert "input-vtx" not in [call.kwargs.get("file_id") or call.args[0] for call in file_mock.call_args_list] class _NeverClaimsCache(DualCache): @@ -1494,10 +1513,9 @@ async def test_native_vertex_skip_releases_the_claim(recorder): claim_cache=claim_cache, ) - assert first == 0 - assert second == 0 - assert file_mock.await_count == 4, ( - "a released claim must let the next retrieve reach the fetch step again; already_claimed would return early at 2 fetches" + assert (first, second) == (0, 0) + assert [call.kwargs["file_id"] for call in file_mock.await_args_list] == ["output-vtx-rel", "output-vtx-rel"], ( + "a released claim must let the next retrieve reach the fetch step again; already_claimed would return early at 1 fetch" ) From 8f2c93560523bfb4d28b21ee0a69bfde3087bb42 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Mon, 12 Oct 2026 04:08:45 -0700 Subject: [PATCH 69/89] Read a swallowed Redis claim failure as unavailable, price unpriceable lines at zero, keep the OTel failure docstring --- litellm/batches/batch_line_item_logging.py | 5 ++- litellm/integrations/otel/plumbing/metrics.py | 4 +- .../batches/test_batch_line_item_logging.py | 38 +++++++++++++++++++ 3 files changed, 43 insertions(+), 4 deletions(-) diff --git a/litellm/batches/batch_line_item_logging.py b/litellm/batches/batch_line_item_logging.py index a44c375a5227..91e9aca9cdcc 100644 --- a/litellm/batches/batch_line_item_logging.py +++ b/litellm/batches/batch_line_item_logging.py @@ -104,11 +104,12 @@ async def _claim_line_items(claim_cache: DualCache, claim_key: str, token: str) return "already_claimed" try: ok: Final = await redis_cache.async_set_cache(claim_key, token, ttl=_LINE_ITEM_CLAIM_TTL_SECONDS, nx=True) # pyright: ignore[reportUnknownMemberType] # RedisCache.set is untyped upstream + holder: Final = None if ok else await redis_cache.async_get_cache(claim_key) # pyright: ignore[reportUnknownMemberType] # RedisCache.get is untyped upstream except Exception: # noqa: BLE001 # a redis outage must not emit duplicates; the next retrieve retries return "unavailable" if ok: return "claimed" - return "already_claimed" + return "already_claimed" if holder is not None else "unavailable" @with_service_target("batch_line_items") @@ -426,7 +427,7 @@ async def _emit_line_event( batch, custom_id, status_code, - response_cost=stats.prompt_cost + stats.completion_cost if stats is not None else None, + response_cost=stats.prompt_cost + stats.completion_cost if stats is not None else 0.0, ) if stats is not None and not response_body.get("usage") and isinstance(result, (ModelResponse, EmbeddingResponse)): setattr( # noqa: B010 # ModelResponse.usage is set dynamically by its ctor, so setattr keeps parity for both response types diff --git a/litellm/integrations/otel/plumbing/metrics.py b/litellm/integrations/otel/plumbing/metrics.py index e94f6ab9f25a..c3294c7925d6 100644 --- a/litellm/integrations/otel/plumbing/metrics.py +++ b/litellm/integrations/otel/plumbing/metrics.py @@ -249,8 +249,6 @@ def record_failure( start_time: datetime, end_time: datetime, ) -> None: - if is_batch_line_item_event(kwargs): - return """Record the one metric a failed request can honestly report: the operation's duration, tagged with ``error.type``. @@ -271,6 +269,8 @@ def record_failure( ``gen_ai.token.type``, so an operator's include/exclude list cannot strip the discriminator and silently merge failures back into the success series. """ + if is_batch_line_item_event(kwargs): + return attributes: Final = { **self._filter_attributes(self._bounded_attributes(kwargs)), Error.TYPE: resolve_error_type(kwargs), diff --git a/tests/unit/batches/test_batch_line_item_logging.py b/tests/unit/batches/test_batch_line_item_logging.py index 9100ab8399f3..14b153860f93 100644 --- a/tests/unit/batches/test_batch_line_item_logging.py +++ b/tests/unit/batches/test_batch_line_item_logging.py @@ -667,6 +667,27 @@ def _parent_logging_with_params(litellm_params: dict) -> Logging: return logging_obj +@pytest.mark.asyncio +async def test_a_line_the_batch_could_not_price_reports_the_zero_the_aggregate_billed(recorder): + litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: flag under test; fixture restores it + with ( + patch( + "litellm.files.main.afile_content", new_callable=AsyncMock, side_effect=_file_content + ), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + patch( + "litellm.cost_calculator.batch_cost_calculator", side_effect=ValueError("no batch price") + ), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch + ): + await _log_completed_batch(_parent_logging(), _batch()) + + line_costs: Final = [ + _payload(event)["response_cost"] + for event in recorder.success_events + if _hidden(event).get("batch_custom_id") is not None + ] + assert line_costs == [0.0] + + @pytest.mark.asyncio async def test_line_items_edge_shapes_and_edge_cases(recorder): litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: flag under test; fixture restores it @@ -1251,12 +1272,16 @@ class _FakeRedisCache(RedisCache): def __init__(self) -> None: self._store: dict[str, object] = {} # mutable-ok: a dict-backed fake needs a mutable store self.fail_next_set = False + self.swallow_next_set = False self.swap_owner_on_get: str | None = None async def async_set_cache(self, key, value, nx=False, **_kwargs): if self.fail_next_set: self.fail_next_set = False raise ConnectionError("redis down") + if self.swallow_next_set: + self.swallow_next_set = False + return None if nx and key in self._store: return None self._store[key] = value @@ -1341,6 +1366,19 @@ async def test_line_items_skip_when_redis_claim_backend_is_down(recorder): assert len(recorder.failure_events) == 1 +@pytest.mark.asyncio +async def test_a_redis_set_failure_the_client_swallows_reads_as_unavailable(): + """RedisCache.async_set_cache logs a failed SET and returns None, the same value + SET NX returns when another worker holds the claim.""" + fake: Final = _FakeRedisCache() + claim_cache: Final = DualCache(redis_cache=fake) + fake.swallow_next_set = True + down: Final = await _claim_line_items(claim_cache, "batch_line_items_emitted:batch_down", "first") + claimed: Final = await _claim_line_items(claim_cache, "batch_line_items_emitted:batch_down", "second") + held: Final = await _claim_line_items(claim_cache, "batch_line_items_emitted:batch_down", "third") + assert (down, claimed, held) == ("unavailable", "claimed", "already_claimed") + + @pytest.mark.asyncio async def test_release_line_item_claim_only_deletes_owned_claims(): fake: Final = _FakeRedisCache() From d62c5f2e645b1d13fd9309242b2085af5a5ccb15 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Mon, 12 Oct 2026 13:54:29 -0700 Subject: [PATCH 70/89] Call handle_completed_batch when line items are off, keep a cache test from leaking Redis into the line claim cache, pin an unpriceable line's zero cost on the callback kwargs --- litellm/litellm_core_utils/litellm_logging.py | 34 +++++++++----- .../batches/test_batch_line_item_logging.py | 4 +- tests/unit/batches/test_batch_utils.py | 42 ++++++++---------- .../test_litellm_logging.py | 44 ++++++++----------- .../unit/proxy/test_redis_auth_cache_flag.py | 1 + 5 files changed, 61 insertions(+), 64 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 3d7ef774b7dd..00ebfc631841 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -36,6 +36,7 @@ from litellm.batches.batch_utils import ( BatchResultFiles, batch_cost_is_final, + handle_completed_batch, handle_completed_batch_with_files, ) from litellm.caching.caching import DualCache @@ -3316,6 +3317,9 @@ async def _async_success_handler_body( has_explicit_batch_data: Final = all(x is not None for x in (batch_cost, batch_usage, batch_models)) should_compute_batch_data: Final = not has_explicit_batch_data and batch_cost_is_final(result) result_files: BatchResultFiles | None = None # rebind-ok: one batch-data branch below supplies the bytes + batch_litellm_params: Final = cast( # cast-ok: shared attribute is an untyped dict + "dict[str, object] | None", self.litellm_params + ) if has_explicit_batch_data: set_hidden_param(result, "response_cost", batch_cost) set_hidden_param(result, "batch_models", batch_models) @@ -3342,14 +3346,23 @@ async def _async_success_handler_body( ) elif should_compute_batch_data: - batch_result, fetched_result_files = await handle_completed_batch_with_files( - batch=result, - custom_llm_provider=self.custom_llm_provider, - model_name=self.get_deployment_model_for_cost(), - litellm_params=self.litellm_params, - model_info=self.get_router_deployment_model_info(), - ) - result_files = fetched_result_files # rebind-ok: the files this branch fetched + if litellm.store_batch_line_items_in_callbacks: + batch_result, fetched_result_files = await handle_completed_batch_with_files( + batch=result, + custom_llm_provider=self.custom_llm_provider, + model_name=self.get_deployment_model_for_cost(), + litellm_params=batch_litellm_params, + model_info=self.get_router_deployment_model_info(), + ) + result_files = fetched_result_files # rebind-ok: the files this branch fetched + else: + batch_result = await handle_completed_batch( + batch=result, + custom_llm_provider=self.custom_llm_provider, + model_name=self.get_deployment_model_for_cost(), + litellm_params=batch_litellm_params, + model_info=self.get_router_deployment_model_info(), + ) set_hidden_param(result, "response_cost", batch_result.cost) set_hidden_param(result, "batch_models", batch_result.models) @@ -3366,16 +3379,13 @@ async def _async_success_handler_body( if litellm.store_batch_line_items_in_callbacks and (has_explicit_batch_data or should_compute_batch_data): from litellm.batches.batch_line_item_logging import spawn_batch_line_items - _line_item_litellm_params: Final = cast( # cast-ok: shared attribute is an untyped dict - "dict[str, object] | None", self.litellm_params - ) try: spawn_batch_line_items( batch=result, custom_llm_provider=self.custom_llm_provider, parent=self, model_name=self.get_deployment_model_for_cost(), - litellm_params=_line_item_litellm_params, + litellm_params=batch_litellm_params, model_info=self.get_router_deployment_model_info(), result_files=result_files, ) diff --git a/tests/unit/batches/test_batch_line_item_logging.py b/tests/unit/batches/test_batch_line_item_logging.py index 14b153860f93..8e816f72912b 100644 --- a/tests/unit/batches/test_batch_line_item_logging.py +++ b/tests/unit/batches/test_batch_line_item_logging.py @@ -681,11 +681,11 @@ async def test_a_line_the_batch_could_not_price_reports_the_zero_the_aggregate_b await _log_completed_batch(_parent_logging(), _batch()) line_costs: Final = [ - _payload(event)["response_cost"] + (event["response_cost"], _payload(event)["response_cost"]) for event in recorder.success_events if _hidden(event).get("batch_custom_id") is not None ] - assert line_costs == [0.0] + assert line_costs == [(0.0, 0.0)] @pytest.mark.asyncio diff --git a/tests/unit/batches/test_batch_utils.py b/tests/unit/batches/test_batch_utils.py index 0d00c66550f4..20aba4ea538f 100644 --- a/tests/unit/batches/test_batch_utils.py +++ b/tests/unit/batches/test_batch_utils.py @@ -3016,7 +3016,7 @@ async def test_batch_retrieve_cost_tracking_with_completed_batch_no_explicit_cos logging_obj.custom_llm_provider = "openai" # Mock handle_completed_batch to return cost data - from litellm.batches.batch_utils import BatchCostUsageResult, BatchResultFiles + from litellm.batches.batch_utils import BatchCostUsageResult expected_cost = 0.05 expected_usage = litellm.Usage( @@ -3027,17 +3027,14 @@ async def test_batch_retrieve_cost_tracking_with_completed_batch_no_explicit_cos expected_models = ["gpt-5-mini"] with patch( - "litellm.litellm_core_utils.litellm_logging.handle_completed_batch_with_files", + "litellm.litellm_core_utils.litellm_logging.handle_completed_batch", new=AsyncMock( - return_value=( - BatchCostUsageResult( - cost=expected_cost, - usage=expected_usage, - models=expected_models, - successful_requests=10, - failed_requests=0, - ), - BatchResultFiles(output=None, error=None), + return_value=BatchCostUsageResult( + cost=expected_cost, + usage=expected_usage, + models=expected_models, + successful_requests=10, + failed_requests=0, ) ), ) as mock_handle_batch: @@ -3168,7 +3165,7 @@ async def test_batch_retrieve_cost_tracking_with_explicit_cost_data(): explicit_models = ["gpt-5-mini", "gpt-5.5"] with patch( - "litellm.litellm_core_utils.litellm_logging.handle_completed_batch_with_files", + "litellm.litellm_core_utils.litellm_logging.handle_completed_batch", new=AsyncMock(), ) as mock_handle_batch: # Call async_success_handler with explicit cost data @@ -3301,7 +3298,7 @@ async def test_batch_retrieve_cost_tracking_with_unified_file_id_incomplete_batc logging_obj.custom_llm_provider = "openai" with patch( - "litellm.litellm_core_utils.litellm_logging.handle_completed_batch_with_files", + "litellm.litellm_core_utils.litellm_logging.handle_completed_batch", new=AsyncMock(), ) as mock_handle_batch: # Call async_success_handler with in_progress batch (unified file ID) @@ -3387,20 +3384,17 @@ async def test_batch_retrieve_cost_tracking_with_partial_explicit_data(): ) expected_models = ["gpt-5-mini"] - from litellm.batches.batch_utils import BatchCostUsageResult, BatchResultFiles + from litellm.batches.batch_utils import BatchCostUsageResult with patch( - "litellm.litellm_core_utils.litellm_logging.handle_completed_batch_with_files", + "litellm.litellm_core_utils.litellm_logging.handle_completed_batch", new=AsyncMock( - return_value=( - BatchCostUsageResult( - cost=expected_cost, - usage=expected_usage, - models=expected_models, - successful_requests=8, - failed_requests=0, - ), - BatchResultFiles(output=None, error=None), + return_value=BatchCostUsageResult( + cost=expected_cost, + usage=expected_usage, + models=expected_models, + successful_requests=8, + failed_requests=0, ) ), ) as mock_handle_batch: diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index a97cf9cf6c37..f42798c5616f 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -1287,24 +1287,19 @@ async def test_forwards_deployment_model_and_pricing(self, monkeypatch) -> None: captured: dict[str, object] = {} - from litellm.batches.batch_utils import BatchCostUsageResult, BatchResultFiles + from litellm.batches.batch_utils import BatchCostUsageResult - async def fake_handle_completed_batch( - **kwargs: object, - ) -> tuple[BatchCostUsageResult, BatchResultFiles]: + async def fake_handle_completed_batch(**kwargs: object) -> BatchCostUsageResult: captured.update(kwargs) - return ( - BatchCostUsageResult( - cost=1.25, - usage=Usage(prompt_tokens=1800, completion_tokens=1000, total_tokens=2800), - models=["m"], - successful_requests=1, - failed_requests=0, - ), - BatchResultFiles(output=None, error=None), + return BatchCostUsageResult( + cost=1.25, + usage=Usage(prompt_tokens=1800, completion_tokens=1000, total_tokens=2800), + models=["m"], + successful_requests=1, + failed_requests=0, ) - monkeypatch.setattr(logging_module, "handle_completed_batch_with_files", fake_handle_completed_batch) + monkeypatch.setattr(logging_module, "handle_completed_batch", fake_handle_completed_batch) obj = LitellmLogging( model="bedrock/global.anthropic.claude-sonnet-4-6", @@ -1387,7 +1382,7 @@ async def test_non_final_batch_is_not_priced(self, monkeypatch, status, output_f from litellm.litellm_core_utils import litellm_logging as logging_module handle_completed_batch = AsyncMock() - monkeypatch.setattr(logging_module, "handle_completed_batch_with_files", handle_completed_batch) + monkeypatch.setattr(logging_module, "handle_completed_batch", handle_completed_batch) batch = self._batch(status, output_file_id) await self._logging_obj()._async_success_handler_body(result=batch, start_time=None, end_time=None) @@ -1397,23 +1392,20 @@ async def test_non_final_batch_is_not_priced(self, monkeypatch, status, output_f @pytest.mark.asyncio async def test_completed_batch_with_output_is_priced(self, monkeypatch) -> None: - from litellm.batches.batch_utils import BatchCostUsageResult, BatchResultFiles + from litellm.batches.batch_utils import BatchCostUsageResult from litellm.litellm_core_utils import litellm_logging as logging_module from litellm.types.utils import Usage handle_completed_batch = AsyncMock( - return_value=( - BatchCostUsageResult( - cost=8e-06, - usage=Usage(prompt_tokens=26, completion_tokens=9, total_tokens=35), - models=["gpt-5.6-luna"], - successful_requests=2, - failed_requests=0, - ), - BatchResultFiles(output=None, error=None), + return_value=BatchCostUsageResult( + cost=8e-06, + usage=Usage(prompt_tokens=26, completion_tokens=9, total_tokens=35), + models=["gpt-5.6-luna"], + successful_requests=2, + failed_requests=0, ) ) - monkeypatch.setattr(logging_module, "handle_completed_batch_with_files", handle_completed_batch) + monkeypatch.setattr(logging_module, "handle_completed_batch", handle_completed_batch) batch = self._batch("completed", "file-out") await self._logging_obj()._async_success_handler_body(result=batch, start_time=None, end_time=None) diff --git a/tests/unit/proxy/test_redis_auth_cache_flag.py b/tests/unit/proxy/test_redis_auth_cache_flag.py index e1eb9d13bdcc..398dd1efa970 100644 --- a/tests/unit/proxy/test_redis_auth_cache_flag.py +++ b/tests/unit/proxy/test_redis_auth_cache_flag.py @@ -74,6 +74,7 @@ def _patched_init_cache(litellm_settings: dict, cache_params: dict): patch.object(ps, "spend_counter_cache", fresh_spend_cache), patch.object(ps, "cli_sso_session_cache", fresh_cli_sso_cache), patch.object(ps, "litellm_config_cache", fresh_config_cache), + patch.object(ps, "batch_line_item_claim_cache", DualCache()), patch.object(ps, "llm_router", None), # Cache is locally imported inside _init_cache: patch it at source. patch("litellm.Cache", return_value=mock_litellm_cache), From f6dc4221c904ac6d87cb67180baf437c811c1615 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Mon, 12 Oct 2026 14:45:34 -0700 Subject: [PATCH 71/89] Hand the cost poller's credential snapshot to line item file reads, so a managed Bedrock batch's lines read its input from the deployment bucket --- litellm/batches/batch_line_item_logging.py | 6 +- .../test_batch_output_file_listing.py | 172 +++++++++++++++++- .../batches/test_batch_line_item_logging.py | 42 +++++ 3 files changed, 218 insertions(+), 2 deletions(-) diff --git a/litellm/batches/batch_line_item_logging.py b/litellm/batches/batch_line_item_logging.py index 91e9aca9cdcc..056dfc337156 100644 --- a/litellm/batches/batch_line_item_logging.py +++ b/litellm/batches/batch_line_item_logging.py @@ -594,8 +594,12 @@ async def log_batch_line_items( try: internal_credentials: Final = parent._litellm_internal_model_credentials # pyright: ignore[reportPrivateUsage] # declared transport attribute on Logging internal_mapping: Final = _as_object_mapping(internal_credentials) + # Cloud storage providers resolve their bucket only from the trusted snapshot key, so it + # travels with the flattened credentials, the same shape the cost poller fetches with fetch_params: Final[dict[str, object] | None] = ( # mutable-ok: file fetcher requires a plain dict - dict(internal_mapping) if internal_mapping is not None else litellm_params + {**internal_mapping, "_litellm_internal_model_credentials": MappingProxyType(dict(internal_mapping))} + if internal_mapping is not None + else litellm_params ) output_content: Final = ( diff --git a/tests/integration/management/test_batch_output_file_listing.py b/tests/integration/management/test_batch_output_file_listing.py index 26d1ae8de46e..dbb7d0e14e1f 100644 --- a/tests/integration/management/test_batch_output_file_listing.py +++ b/tests/integration/management/test_batch_output_file_listing.py @@ -16,12 +16,13 @@ from typing import Final import httpx +import yaml from cryptography import x509 from cryptography.hazmat.primitives import hashes, serialization from cryptography.hazmat.primitives.asymmetric import ec from cryptography.x509.oid import NameOID from integration._support.client import JSON_OBJECT, Gateway, Scenario, eventually, object_value, string_value -from integration._support.database import read_rows +from integration._support.database import read_rows, scratch_database from integration._support.process import owned_proxy from integration._support.upstream import ScenarioHandle, delete_scenario, register_scenario from integration._support.wire import Reply, Request, Wire, wire_server @@ -42,6 +43,7 @@ BEDROCK_LAST_MODIFIED: Final = "Thu, 02 Oct 2025 12:00:00 GMT" BEDROCK_OUTPUT_CONTENT: Final = b'{"recordId":"req-1","modelOutput":{}}\n' _BATCH_PROCESSED_SQL: Final = 'SELECT batch_processed FROM "LiteLLM_ManagedObjectTable" WHERE unified_object_id=%s' +_ANY_BATCH_PROCESSED_SQL: Final = 'SELECT 1 FROM "LiteLLM_ManagedObjectTable" WHERE batch_processed' @dataclass(frozen=True, slots=True) @@ -649,3 +651,171 @@ def test_repeated_batch_retrieve_does_not_refetch_saved_file_details(gateway: Ga assert output_after == output_before, output_after additional_hits: Final = _metadata_hit_count(gateway, batch.scenario) assert additional_hits == 0, f"Repeated batch retrieval fetched metadata {additional_hits} more times" + + +@dataclass(frozen=True, slots=True) +class _StoringS3Peer: + objects: dict[str, bytes] # mutable-ok: the peer keeps every uploaded object so later GETs can read it back + output: bytes + + def __call__(self, request: Request) -> Reply: + path: Final = request.target.split("?", 1)[0] + if request.method == "PUT" and path.startswith(f"/{BEDROCK_BUCKET}/"): + self.objects[path] = request.body + return Reply(body=b"") + if request.method == "GET" and path.startswith(f"/{BEDROCK_BUCKET}/{BEDROCK_MANAGED_S3_OUTPUT_PREFIX}"): + if request.headers.get("range") == "bytes=0-0": + return Reply( + status=206, + body=self.output[:1], + headers={ + "Content-Range": f"bytes 0-0/{len(self.output)}", + "Last-Modified": BEDROCK_LAST_MODIFIED, + }, + ) + return Reply(status=200, body=self.output, headers={"Last-Modified": BEDROCK_LAST_MODIFIED}) + stored: Final = self.objects.get(path) + if request.method == "GET" and stored is not None: + return Reply(status=200, body=stored, headers={"Last-Modified": BEDROCK_LAST_MODIFIED}) + return Reply(status=404, body=b'{"message":"not scripted"}') + + +def _bedrock_line_output(marker: str) -> bytes: + succeeded: Final = { + "recordId": "req-1", + "modelInput": {"messages": [{"role": "user", "content": f"{marker} ping"}], "max_tokens": 8}, + "modelOutput": { + "id": "msg_bdrk_line_1", + "type": "message", + "role": "assistant", + "model": BEDROCK_MODEL_ID, + "content": [{"type": "text", "text": "pong"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 8, "output_tokens": 2}, + }, + } + failed: Final = { + "recordId": "req-2", + "modelInput": {"messages": [{"role": "user", "content": f"{marker} bad"}], "max_tokens": 8}, + "error": {"errorCode": 400, "errorMessage": "Malformed input request"}, + } + return (json.dumps(succeeded) + "\n" + json.dumps(failed) + "\n").encode() + + +def _line_items_config(destination: Path) -> Path: + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config["litellm_settings"].update({"callbacks": ["generic_api"], "DEFAULT_FLUSH_INTERVAL_SECONDS": 1}) + config["general_settings"].update({"store_batch_line_items_in_callbacks": True}) + destination.write_text(yaml.safe_dump(config)) + return destination + + +def test_a_bedrock_batch_the_poller_logs_emits_its_lines_from_the_deployment_bucket( + gateway: Gateway, tmp_path: Path +) -> None: + """The cost poller reads a managed Bedrock batch's files with the deployment's S3 bucket, so the + line events it starts must read the input file from that same bucket""" + job_id: Final = f"integration-batch-lines-{uuid.uuid4().hex}" + marker: Final = f"bedrock-poller-lines-{uuid.uuid4().hex[:12]}" + received: Final[list[Request]] = [] # mutable-ok: drain() consumes the queue, later polls keep earlier events + environment: Final = { + "AWS_CA_BUNDLE": str(tmp_path / "bedrock.pem"), + "AWS_EC2_METADATA_DISABLED": "true", + "SSL_VERIFY": "False", + } + with ( + scratch_database() as database_url, + wire_server(lambda _request: Reply()) as sink, + wire_server(_StoringS3Peer({}, _bedrock_line_output(marker))) as s3, + wire_server( + _BedrockControlPlane(SimpleQueue(), job_id, BEDROCK_JOB_ARN_PREFIX + job_id), tls=_tls_context(tmp_path) + ) as bedrock, + _bedrock_tunnel(bedrock) as tunnel, + owned_proxy( + gateway, + tmp_path, + { + **environment, + "HTTPS_PROXY": tunnel.url, + "GENERIC_LOGGER_ENDPOINT": sink.url, + "DATABASE_URL": database_url, + }, + config=_line_items_config(tmp_path / "bedrock_line_items.yaml"), + ) as candidate, + candidate.scenario() as scenario, + ): + model: Final = scenario.model( + model=BEDROCK_MODEL, + api_key=None, + api_base=None, + aws_access_key_id="AKIAIOSFODNN7EXAMPLE", + aws_secret_access_key="wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY", + aws_region_name=BEDROCK_REGION, + s3_bucket_name=BEDROCK_BUCKET, + s3_endpoint_url=s3.url, + aws_batch_role_arn=BEDROCK_ROLE_ARN, + ) + owner_key: Final = scenario.key(user_id=scenario.user(user_role="internal_user"), models=[model]) + input_lines: Final = "".join( + json.dumps( + { + "custom_id": custom_id, + "method": "POST", + "url": "/v1/chat/completions", + "body": { + "model": model, + "messages": [{"role": "user", "content": f"{marker} {text}"}], + "max_tokens": 8, + }, + } + ) + + "\n" + for custom_id, text in (("req-1", "ping"), ("req-2", "bad")) + ) + uploaded: Final = candidate.request_multipart( + "/v1/files", + {"purpose": "batch", "target_model_names": model}, + {"file": ("input.jsonl", input_lines.encode(), "application/jsonl")}, + key=owner_key, + ) + assert uploaded.status_code == 200, uploaded.text + created: Final = candidate.request( + "POST", + "/v1/batches", + { + "input_file_id": string_value(JSON_OBJECT.validate_json(uploaded.content)["id"]), + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + }, + key=owner_key, + ) + assert created.status_code == 200, created.text + + def line_events() -> tuple[dict[str, JsonValue], ...]: + received.extend(sink.drain()) + return tuple( + event + for delivery in received + for event in json.loads(delivery.body) + if object_value(event["hidden_params"]).get("batch_custom_id") is not None + and marker in json.dumps(event["messages"]) + ) + + eventually(lambda: read_rows(_ANY_BATCH_PROCESSED_SQL, (), database_url=database_url), bool, seconds=60) + delivered: Final = eventually( + line_events, + lambda values: ( + {object_value(event["hidden_params"])["batch_custom_id"] for event in values} == {"req-1", "req-2"} + ), + seconds=30, + ) + + assert sorted( + ( + string_value(object_value(event["hidden_params"])["batch_custom_id"]), + event["status"], + object_value(event["hidden_params"])["batch_line_status_code"], + ) + for event in delivered + ) == [("req-1", "success", 200), ("req-2", "failure", 400)], delivered diff --git a/tests/unit/batches/test_batch_line_item_logging.py b/tests/unit/batches/test_batch_line_item_logging.py index 8e816f72912b..2629fa46605b 100644 --- a/tests/unit/batches/test_batch_line_item_logging.py +++ b/tests/unit/batches/test_batch_line_item_logging.py @@ -36,7 +36,9 @@ from litellm.caching.caching import DualCache from litellm.caching.redis_cache import RedisCache from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.get_litellm_params import add_trusted_model_credentials_to_litellm_params from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.llms.bedrock.files.transformation import get_configured_s3_bucket_names from litellm.types.utils import LiteLLMBatch, Usage INPUT_JSONL = b"\n".join( @@ -922,6 +924,46 @@ async def test_line_items_bedrock_shapes(recorder): assert "bad request" in _payload(failure)["error_str"] +@pytest.mark.asyncio +async def test_a_bedrock_batch_the_poller_logs_reads_every_file_from_the_deployment_bucket(recorder): + """Bedrock trusts its S3 bucket only from the server-side credential snapshot, so the line + fan-out the cost poller starts must hand that snapshot to every file read""" + litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: flag under test; fixture restores it + file_mock: Final = AsyncMock(side_effect=_edge_file_content) + parent: Final = _parent_logging(custom_llm_provider="bedrock") + parent._litellm_internal_model_credentials = ( + MappingProxyType( # test-quality-ok: private transport attribute, same channel the batch cost poller sets + {"s3_bucket_name": "line-items-bucket", "aws_region_name": "us-east-1"} + ) + ) + with ( + patch( + "litellm.files.main.afile_content", file_mock + ), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + patch( + "litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02) + ), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch + ): + await _log_completed_batch(parent, _provider_batch("batch_bed_poller", "input-bed", "output-bed")) + + buckets_by_file: Final[dict[str, tuple[str, ...]]] = {} # mutable-ok: collects one entry per fetched file + for call in file_mock.call_args_list: + litellm_params: dict[str, object] = {} # mutable-ok: the dict the real file read builds and passes on + add_trusted_model_credentials_to_litellm_params(litellm_params, call.kwargs) + buckets_by_file[call.kwargs["file_id"]] = get_configured_s3_bucket_names(litellm_params) + assert buckets_by_file == { + "output-bed": ("line-items-bucket",), + "input-bed": ("line-items-bucket",), + } + assert sorted( + _hidden(e)["batch_custom_id"] for e in recorder.success_events if _hidden(e).get("batch_custom_id") + ) == [ + "br-anth", + "br-conv", + "br-titan", + ] + + @pytest.mark.asyncio async def test_line_items_mistral_shape(recorder): litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: flag under test; fixture restores it From 6125f5f320d053c7b7f6521f5a34afc8e218ca8c Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Mon, 12 Oct 2026 15:18:32 -0700 Subject: [PATCH 72/89] Log batch lines through the sync callbacks a direct async call runs, keep failed lines out of router cooldown counts, and emit an error line whose body is text --- litellm/batches/batch_line_item_logging.py | 18 ++- litellm/router.py | 2 + .../spend/test_batch_line_item_callbacks.py | 125 ++++++++++++++- .../batches/test_batch_line_item_logging.py | 143 ++++++++++++++++++ 4 files changed, 278 insertions(+), 10 deletions(-) diff --git a/litellm/batches/batch_line_item_logging.py b/litellm/batches/batch_line_item_logging.py index 056dfc337156..7b2ad261ebd6 100644 --- a/litellm/batches/batch_line_item_logging.py +++ b/litellm/batches/batch_line_item_logging.py @@ -367,7 +367,10 @@ async def _emit_line_event( request_body: Final = _request_body_for_entry(entry, request_line) status_code: Final = _line_status_code(entry, custom_llm_provider) request_call_type: Final = _call_type_for_request(request_line) - response_body: Final = _get_response_from_batch_job_output_file(entry, custom_llm_provider) + # A failed line's body can be the provider's error as text (Mistral writes a JSON string), not an object + response_body: Final = ( + _as_object_mapping(_get_response_from_batch_job_output_file(entry, custom_llm_provider)) or _EMPTY_BODY + ) start_time: Final = parent.start_time model: Final = _line_model(response_body, request_body, parent) @@ -394,6 +397,9 @@ async def _emit_line_event( litellm_params={ **parent_params, "metadata": _metadata_copy(parent_params), + # Flags the line as the async call it stands for, so the logging dispatch runs the + # callbacks a direct async call of that type runs, legacy sync ones included + call_type: True, }, optional_params={}, model=child.model, @@ -415,12 +421,7 @@ async def _emit_line_event( if result is None: exception: Final = _BatchLineFailure(_line_error_payload(entry, custom_llm_provider)) exception._hidden_params = _line_hidden_params(batch, custom_id, status_code) # pyright: ignore[reportPrivateUsage] # _hidden_params is set on the exception instance itself - await child.async_failure_handler( - exception=exception, - traceback_exception="", - start_time=start_time, - end_time=now, - ) + await child.dispatch_failure_handlers(exception, "", prefer_async_handlers=True) return True result._hidden_params = _line_hidden_params( # pyright: ignore[reportPrivateUsage] # same hidden_params channel the aggregate batch event uses @@ -439,11 +440,12 @@ async def _emit_line_event( total_tokens=stats.total_tokens, ), ) - await child.async_success_handler( + await child.dispatch_success_handlers( result=result, start_time=start_time, end_time=now, cache_hit=False, + prefer_async_handlers=True, ) return True diff --git a/litellm/router.py b/litellm/router.py index 02548a6f46c8..66936b11cc32 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -8607,6 +8607,8 @@ def deployment_callback_on_failure( - False if the deployment should not be put in cooldown """ verbose_router_logger.debug("Router: Entering 'deployment_callback_on_failure'") + if is_batch_line_item_event(kwargs): + return False try: exception: Final = kwargs.get("exception", None) diff --git a/tests/integration/spend/test_batch_line_item_callbacks.py b/tests/integration/spend/test_batch_line_item_callbacks.py index 30b80f84bbc6..cec7f0fd3cc5 100644 --- a/tests/integration/spend/test_batch_line_item_callbacks.py +++ b/tests/integration/spend/test_batch_line_item_callbacks.py @@ -2,6 +2,8 @@ import json import uuid +from collections.abc import Callable +from dataclasses import dataclass from datetime import UTC, datetime from hashlib import sha256 from pathlib import Path @@ -77,6 +79,22 @@ def _failed_line(custom_id: str) -> str: ) +def _failed_line_with_text_body(custom_id: str) -> str: + """An error file line whose response body is the provider's error as JSON text, as Mistral writes it.""" + return json.dumps( + { + "id": f"batch_req_{custom_id}", + "custom_id": custom_id, + "response": { + "status_code": 400, + "body": json.dumps({"object": "error", "message": f"synthetic line failure {custom_id}"}), + }, + "error": None, + }, + separators=(",", ":"), + ) + + def _batch(status: str, *, files_ready: bool) -> dict[str, JsonValue]: return { "id": "batch-$REQUEST_ID", @@ -112,9 +130,11 @@ def _input_lines(model_name: str, marker: str) -> tuple[str, ...]: ) -def _provider_routes(input_lines: tuple[str, ...]) -> RoutedResponse: +def _provider_routes( + input_lines: tuple[str, ...], *, error_line: Callable[[str], str] = _failed_line +) -> RoutedResponse: output_lines: Final = (_successful_line("r1"), _failed_line("r2"), _successful_line("r3")) - error_lines: Final = tuple(_failed_line(custom_id) for custom_id in ERROR_FILE_IDS) + error_lines: Final = tuple(error_line(custom_id) for custom_id in ERROR_FILE_IDS) return RoutedResponse( content_type="application/x-routed", routes={ @@ -581,6 +601,107 @@ def test_a_line_body_system_prompt_leads_the_logged_messages(gateway: Gateway, t assert "system" not in object_value(event["metadata"]), event["metadata"] +def test_an_error_file_line_with_a_text_body_emits_a_failure_event(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "batch-line-text-error-" + uuid.uuid4().hex[:12] + + with ( + wire_server(lambda _request: Reply()) as endpoint, + owned_proxy( + gateway, + tmp_path, + {"GENERIC_LOGGER_ENDPOINT": endpoint.url}, + config=_line_item_config(Path("tests/integration/proxy_config.yaml"), tmp_path / "line_items_text.yaml"), + ) as candidate, + candidate.scenario() as scenario, + ): + routed_model: Final = scenario.model( + api_base=f"{gateway.upstream_url}/{marker}", api_key="synthetic-provider-key" + ) + input_lines: Final = _input_lines(routed_model, marker) + handle: Final = register_scenario(marker, _provider_routes(input_lines, error_line=_failed_line_with_text_body)) + scenario.cleanups.callback(delete_scenario, handle) + key: Final = scenario.key(models=[routed_model]) + _completed_batch(candidate, routed_model, input_lines, key) + received: Final[list[Request]] = [] # mutable-ok: drain() consumes the queue, later polls keep earlier batches + + lines: Final = { + string_value(_hidden(event)["batch_custom_id"]): event + for event in _all_lines_delivered(endpoint, received, key) + } + for custom_id in ERROR_FILE_IDS: + assert lines[custom_id]["status"] == "failure", lines[custom_id] + assert _hidden(lines[custom_id])["batch_line_status_code"] == 400, lines[custom_id] + assert f"synthetic line failure {custom_id}" in json.dumps(lines[custom_id]["error_information"]) + + +@dataclass(frozen=True, slots=True) +class _StoringBucket: + objects: dict[str, bytes] # mutable-ok: the bucket keeps every uploaded object for the assertions + + def __call__(self, request: Request) -> Reply: + if request.method == "PUT": + self.objects[request.target.split("?", 1)[0]] = request.body + return Reply(body=b"") + + +def _legacy_s3_config(base: Path, destination: Path, endpoint_url: str) -> Path: + config: Final = yaml.safe_load(base.read_text()) + config["litellm_settings"].update( + { + "success_callback": ["s3"], + "s3_callback_params": { + "s3_bucket_name": "batch-line-items", + "s3_region_name": "us-east-1", + "s3_endpoint_url": endpoint_url, + "s3_aws_access_key_id": "AKIAIOSFODNN7EXAMPLE", + "s3_aws_secret_access_key": "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY", + }, + } + ) + config["general_settings"].update({"store_batch_line_items_in_callbacks": True}) + destination.write_text(yaml.safe_dump(config)) + return destination + + +def test_the_legacy_s3_callback_stores_one_object_per_successful_line(gateway: Gateway, tmp_path: Path) -> None: + marker: Final = "batch-line-s3-" + uuid.uuid4().hex[:12] + bucket: Final = _StoringBucket({}) + + with ( + wire_server(bucket) as s3, + owned_proxy( + gateway, + tmp_path, + {"AWS_EC2_METADATA_DISABLED": "true"}, + config=_legacy_s3_config( + Path("tests/integration/proxy_config.yaml"), tmp_path / "line_items_s3.yaml", s3.url + ), + ) as candidate, + candidate.scenario() as scenario, + ): + routed_model: Final = scenario.model( + api_base=f"{gateway.upstream_url}/{marker}", api_key="synthetic-provider-key" + ) + input_lines: Final = _input_lines(routed_model, marker) + handle: Final = register_scenario(marker, _provider_routes(input_lines)) + scenario.cleanups.callback(delete_scenario, handle) + key: Final = scenario.key(models=[routed_model]) + _completed_batch(candidate, routed_model, input_lines, key) + + def stored_lines() -> list[str]: + payloads = [json.loads(body) for body in tuple(bucket.objects.values())] + return sorted( + string_value(object_value(payload["hidden_params"])["batch_custom_id"]) + for payload in payloads + if object_value(payload["hidden_params"]).get("batch_custom_id") is not None + and marker in json.dumps(payload["messages"]) + ) + + assert eventually(stored_lines, lambda ids: ids == list(OUTPUT_SUCCESS_IDS), seconds=40) == list( + OUTPUT_SUCCESS_IDS + ) + + def test_a_restarted_proxy_does_not_emit_a_completed_batchs_lines_again(gateway: Gateway, tmp_path: Path) -> None: marker: Final = "batch-line-restart-" + uuid.uuid4().hex[:12] config: Final = _line_item_config(Path("tests/integration/proxy_config.yaml"), tmp_path / "line_items_restart.yaml") diff --git a/tests/unit/batches/test_batch_line_item_logging.py b/tests/unit/batches/test_batch_line_item_logging.py index 2629fa46605b..7936c4756c7a 100644 --- a/tests/unit/batches/test_batch_line_item_logging.py +++ b/tests/unit/batches/test_batch_line_item_logging.py @@ -36,9 +36,11 @@ from litellm.caching.caching import DualCache from litellm.caching.redis_cache import RedisCache from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils import litellm_logging, thread_pool_executor from litellm.litellm_core_utils.get_litellm_params import add_trusted_model_credentials_to_litellm_params from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.bedrock.files.transformation import get_configured_s3_bucket_names +from litellm.router_utils.router_callbacks.track_deployment_metrics import get_deployment_failures_for_current_minute from litellm.types.utils import LiteLLMBatch, Usage INPUT_JSONL = b"\n".join( @@ -242,6 +244,92 @@ async def test_line_items_emitted_alongside_aggregate(recorder): assert "bad request boom" in _payload(failure)["error_str"] +class _InlineExecutor: + def submit(self, fn, /, *args, **kwargs): + fn(*args, **kwargs) + + +@pytest.mark.asyncio +async def test_line_events_reach_the_sync_callbacks_an_async_request_runs(recorder, monkeypatch): + litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: flag under test; fixture restores it + s3_uploads: Final[list[dict]] = [] + sync_failures: Final[list[dict]] = [] + + def record_sync_failure(kwargs, response_obj, start_time, end_time): + sync_failures.append(kwargs) + + monkeypatch.setattr(litellm, "success_callback", ["s3"]) # test-quality-ok: legacy string callback under test + monkeypatch.setattr(litellm, "failure_callback", [record_sync_failure]) # test-quality-ok: same, failure side + monkeypatch.setattr( + litellm_logging, "s3Logger", SimpleNamespace(log_event=lambda **kwargs: s3_uploads.append(kwargs["kwargs"])) + ) # test-quality-ok: the S3 upload is the network boundary of the legacy s3 callback + monkeypatch.setattr( + thread_pool_executor, "executor", _InlineExecutor() + ) # test-quality-ok: runs sync handlers inline + monkeypatch.setattr( + litellm_logging, "executor", _InlineExecutor() + ) # test-quality-ok: same executor, module binding + with ( + patch( + "litellm.files.main.afile_content", new_callable=AsyncMock, side_effect=_file_content + ), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + patch( + "litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02) + ), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch + ): + await _log_completed_batch(_parent_logging(), _batch()) + + assert [_hidden(event)["batch_custom_id"] for event in s3_uploads] == ["a"] + assert [ + _hidden(event)["batch_custom_id"] for event in recorder.success_events if _hidden(event).get("batch_custom_id") + ] == ["a"] + assert [ + (_hidden(event)["batch_custom_id"], _hidden(event)["batch_line_status_code"]) for event in sync_failures + ] == [("b", 400)] + assert [_hidden(event)["batch_custom_id"] for event in recorder.failure_events] == ["b"] + + +@pytest.mark.asyncio +async def test_failed_lines_do_not_count_toward_the_deployments_cooldown(recorder, monkeypatch): + litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: flag under test; fixture restores it + monkeypatch.setattr(litellm, "callbacks", []) # test-quality-ok: the router registers into these globals + monkeypatch.setattr(litellm, "success_callback", []) # test-quality-ok: same, restored at teardown + monkeypatch.setattr(litellm, "failure_callback", []) # test-quality-ok: same, restored at teardown + monkeypatch.setattr( + thread_pool_executor, "executor", _InlineExecutor() + ) # test-quality-ok: runs sync handlers inline + monkeypatch.setattr( + litellm_logging, "executor", _InlineExecutor() + ) # test-quality-ok: same executor, module binding + router: Final = litellm.Router( + model_list=[ + { + "model_name": "gpt-4o", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-line-cooldown"}, + "model_info": {"id": "dep-line-cooldown"}, + } + ] + ) + with ( + patch( + "litellm.files.main.afile_content", new_callable=AsyncMock, side_effect=_file_content + ), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + patch( + "litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02) + ), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch + ): + await _log_completed_batch( + _parent_logging_with_params( + {"model_info": {"id": "dep-line-cooldown"}, "metadata": {"model_group": "gpt-4o"}} + ), + _batch(), + ) + + assert [_hidden(event)["batch_custom_id"] for event in recorder.failure_events] == ["b"] + assert router.deployment_callback_on_failure in litellm.failure_callback + assert get_deployment_failures_for_current_minute(router, "dep-line-cooldown") == 0 + + @pytest.mark.asyncio async def test_flag_off_emits_only_aggregate(recorder): assert litellm.store_batch_line_items_in_callbacks is False @@ -1096,6 +1184,61 @@ async def test_line_items_failure_without_error_or_response_body_still_reaches_c assert _payload(failure)["error_str"] == "{}" +MISTRAL_ERR_INPUT_JSONL = json.dumps( + { + "custom_id": "mis-err", + "method": "POST", + "url": "/v1/chat/completions", + "body": {"messages": [{"role": "user", "content": "hi"}], "tool_choice": "any"}, + } +).encode() + +MISTRAL_ERR_OUTPUT_JSONL = json.dumps( + { + "id": "batch-mist-err-1", + "custom_id": "mis-err", + "response": { + "status_code": 400, + "body": json.dumps( + {"object": "error", "message": "No tools provided for tool choice `any`", "type": "invalid_request"} + ), + }, + "error": None, + } +).encode() + + +@pytest.mark.asyncio +async def test_a_mistral_error_line_with_a_text_body_reaches_the_failure_callbacks(recorder): + litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: flag under test; fixture restores it + file_mock: Final = AsyncMock( + side_effect=_scoped_file_content( + { + "input-mist-err": MISTRAL_ERR_INPUT_JSONL, + "output-mist-err": MISTRAL_ERR_OUTPUT_JSONL, + } + ) + ) + with ( + patch( + "litellm.files.main.afile_content", file_mock + ), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + patch( + "litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02) + ), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch + ): + await _log_completed_batch( + _parent_logging(custom_llm_provider="mistral"), + _provider_batch("batch_mist_err", "input-mist-err", "output-mist-err"), + ) + + assert [ + (_hidden(event)["batch_custom_id"], _hidden(event)["batch_line_status_code"]) + for event in recorder.failure_events + ] == [("mis-err", 400)] + assert "No tools provided for tool choice" in _payload(recorder.failure_events[0])["error_str"] + + @pytest.mark.asyncio async def test_line_items_native_vertex_rows_are_skipped(recorder): litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: flag under test; fixture restores it From cbdbbb5514a49e9450d9db63b70f3749606bdb53 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Mon, 12 Oct 2026 22:00:04 -0700 Subject: [PATCH 73/89] fix(batches): log failed Bedrock embedding lines as embeddings with their input A Bedrock batch input record carries no url, so a failed Titan v2 line fell back to acompletion and its inputText never reached the event messages. The failed line now takes aembedding from the deployment model and the record's inputText is the line input --- litellm/batches/batch_line_item_logging.py | 20 +++- .../test_batch_output_file_listing.py | 104 +++++++++++++----- .../batches/test_batch_line_item_logging.py | 62 ++++++++++- 3 files changed, 155 insertions(+), 31 deletions(-) diff --git a/litellm/batches/batch_line_item_logging.py b/litellm/batches/batch_line_item_logging.py index 7b2ad261ebd6..63c39e7fee3c 100644 --- a/litellm/batches/batch_line_item_logging.py +++ b/litellm/batches/batch_line_item_logging.py @@ -220,18 +220,32 @@ def _call_type_for_request(request_line: Mapping[str, object] | None) -> str: return _CALL_TYPE_BY_BATCH_URL.get(url if isinstance(url, str) else "", "acompletion") -def _call_type_for_line(request_line: Mapping[str, object] | None, result: "_BatchLineResult | None") -> str: +def _is_bedrock_embedding_batch_model(model: str) -> bool: + from litellm.llms.bedrock.files.transformation import BedrockFilesConfig + + return BedrockFilesConfig._is_titan_v2_embed_model(model) # pyright: ignore[reportPrivateUsage] # the check that chose the embedding record shape when the batch input was written + + +def _call_type_for_line( + request_line: Mapping[str, object] | None, + result: "_BatchLineResult | None", + custom_llm_provider: BatchProvider, + model: str, +) -> str: if isinstance(result, EmbeddingResponse): return "aembedding" if isinstance(result, ResponsesAPIResponse): return "aresponses" if isinstance(result, ModelResponse): return "acompletion" + # Bedrock input records carry no url, so a failed line takes its type from the deployment model + if custom_llm_provider == "bedrock" and _is_bedrock_embedding_batch_model(model): + return "aembedding" return _call_type_for_request(request_line) def _line_messages(request_body: Mapping[str, object]) -> object: - return request_body.get("messages") or request_body.get("input") or () + return request_body.get("messages") or request_body.get("input") or request_body.get("inputText") or () def _line_model( @@ -384,7 +398,7 @@ async def _emit_line_event( if successful and result is None: return False - call_type: Final = _call_type_for_line(request_line, result) + call_type: Final = _call_type_for_line(request_line, result, custom_llm_provider, model) child: Final = _new_child_logging( parent=parent, model=model, diff --git a/tests/integration/management/test_batch_output_file_listing.py b/tests/integration/management/test_batch_output_file_listing.py index dbb7d0e14e1f..8f093c4e5759 100644 --- a/tests/integration/management/test_batch_output_file_listing.py +++ b/tests/integration/management/test_batch_output_file_listing.py @@ -8,7 +8,7 @@ import ssl import threading import uuid -from collections.abc import Generator +from collections.abc import Callable, Generator from contextlib import contextmanager from dataclasses import dataclass from pathlib import Path @@ -35,6 +35,7 @@ ERROR_BYTES: Final = 512 BEDROCK_MODEL: Final = "bedrock/anthropic.claude-3-haiku-20240307-v1:0" BEDROCK_MODEL_ID: Final = "anthropic.claude-3-haiku-20240307-v1:0" +BEDROCK_TITAN_MODEL_ID: Final = "amazon.titan-embed-text-v2:0" BEDROCK_REGION: Final = "us-east-1" BEDROCK_AUTHORITY: Final = f"bedrock.{BEDROCK_REGION}.amazonaws.com:443" BEDROCK_BUCKET: Final = "integration-batch-listing-bucket" @@ -485,6 +486,7 @@ class _BedrockControlPlane: job_locations: SimpleQueue[tuple[str, str, str]] job_id: str job_arn: str + model_id: str = BEDROCK_MODEL_ID def __call__(self, request: Request) -> Reply: if request.method == "POST" and request.target == "/model-invocation-job": @@ -504,7 +506,7 @@ def __call__(self, request: Request) -> Reply: { "jobArn": self.job_arn, "jobName": retrieved_job_name, - "modelId": BEDROCK_MODEL_ID, + "modelId": self.model_id, "status": "Completed", "submitTime": 1700000000, "lastModifiedTime": 1700000001, @@ -703,6 +705,36 @@ def _bedrock_line_output(marker: str) -> bytes: return (json.dumps(succeeded) + "\n" + json.dumps(failed) + "\n").encode() +def _bedrock_titan_line_output(marker: str) -> bytes: + succeeded: Final = { + "recordId": "req-1", + "modelInput": {"inputText": f"{marker} ping"}, + "modelOutput": {"embedding": [0.1, 0.2], "inputTextTokenCount": 3}, + } + failed: Final = { + "recordId": "req-2", + "modelInput": {"inputText": f"{marker} bad"}, + "error": {"errorCode": 400, "errorMessage": "Malformed input request"}, + } + return (json.dumps(succeeded) + "\n" + json.dumps(failed) + "\n").encode() + + +@dataclass(frozen=True, slots=True) +class _BedrockLineBatch: + model_id: str + endpoint: str + output: Callable[[str], bytes] + + def body(self, content: str) -> dict[str, JsonValue]: + if self.endpoint == "/v1/embeddings": + return {"input": content} + return {"messages": [{"role": "user", "content": content}], "max_tokens": 8} + + +_BEDROCK_CHAT_BATCH: Final = _BedrockLineBatch(BEDROCK_MODEL_ID, "/v1/chat/completions", _bedrock_line_output) +_BEDROCK_TITAN_BATCH: Final = _BedrockLineBatch(BEDROCK_TITAN_MODEL_ID, "/v1/embeddings", _bedrock_titan_line_output) + + def _line_items_config(destination: Path) -> Path: config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) config["litellm_settings"].update({"callbacks": ["generic_api"], "DEFAULT_FLUSH_INTERVAL_SECONDS": 1}) @@ -711,11 +743,9 @@ def _line_items_config(destination: Path) -> Path: return destination -def test_a_bedrock_batch_the_poller_logs_emits_its_lines_from_the_deployment_bucket( - gateway: Gateway, tmp_path: Path -) -> None: - """The cost poller reads a managed Bedrock batch's files with the deployment's S3 bucket, so the - line events it starts must read the input file from that same bucket""" +def _line_events_of_a_polled_bedrock_batch( + gateway: Gateway, tmp_path: Path, batch: _BedrockLineBatch +) -> tuple[dict[str, JsonValue], ...]: job_id: Final = f"integration-batch-lines-{uuid.uuid4().hex}" marker: Final = f"bedrock-poller-lines-{uuid.uuid4().hex[:12]}" received: Final[list[Request]] = [] # mutable-ok: drain() consumes the queue, later polls keep earlier events @@ -727,9 +757,10 @@ def test_a_bedrock_batch_the_poller_logs_emits_its_lines_from_the_deployment_buc with ( scratch_database() as database_url, wire_server(lambda _request: Reply()) as sink, - wire_server(_StoringS3Peer({}, _bedrock_line_output(marker))) as s3, + wire_server(_StoringS3Peer({}, batch.output(marker))) as s3, wire_server( - _BedrockControlPlane(SimpleQueue(), job_id, BEDROCK_JOB_ARN_PREFIX + job_id), tls=_tls_context(tmp_path) + _BedrockControlPlane(SimpleQueue(), job_id, BEDROCK_JOB_ARN_PREFIX + job_id, batch.model_id), + tls=_tls_context(tmp_path), ) as bedrock, _bedrock_tunnel(bedrock) as tunnel, owned_proxy( @@ -746,7 +777,7 @@ def test_a_bedrock_batch_the_poller_logs_emits_its_lines_from_the_deployment_buc candidate.scenario() as scenario, ): model: Final = scenario.model( - model=BEDROCK_MODEL, + model=f"bedrock/{batch.model_id}", api_key=None, api_base=None, aws_access_key_id="AKIAIOSFODNN7EXAMPLE", @@ -762,12 +793,8 @@ def test_a_bedrock_batch_the_poller_logs_emits_its_lines_from_the_deployment_buc { "custom_id": custom_id, "method": "POST", - "url": "/v1/chat/completions", - "body": { - "model": model, - "messages": [{"role": "user", "content": f"{marker} {text}"}], - "max_tokens": 8, - }, + "url": batch.endpoint, + "body": {"model": model, **batch.body(f"{marker} {text}")}, } ) + "\n" @@ -785,7 +812,7 @@ def test_a_bedrock_batch_the_poller_logs_emits_its_lines_from_the_deployment_buc "/v1/batches", { "input_file_id": string_value(JSON_OBJECT.validate_json(uploaded.content)["id"]), - "endpoint": "/v1/chat/completions", + "endpoint": batch.endpoint, "completion_window": "24h", }, key=owner_key, @@ -803,7 +830,7 @@ def line_events() -> tuple[dict[str, JsonValue], ...]: ) eventually(lambda: read_rows(_ANY_BATCH_PROCESSED_SQL, (), database_url=database_url), bool, seconds=60) - delivered: Final = eventually( + return eventually( line_events, lambda values: ( {object_value(event["hidden_params"])["batch_custom_id"] for event in values} == {"req-1", "req-2"} @@ -811,11 +838,36 @@ def line_events() -> tuple[dict[str, JsonValue], ...]: seconds=30, ) - assert sorted( - ( - string_value(object_value(event["hidden_params"])["batch_custom_id"]), - event["status"], - object_value(event["hidden_params"])["batch_line_status_code"], - ) - for event in delivered - ) == [("req-1", "success", 200), ("req-2", "failure", 400)], delivered + +def _line_outcomes(events: tuple[dict[str, JsonValue], ...]) -> list[tuple[str, JsonValue, JsonValue, JsonValue]]: + return sorted( + ( + string_value(object_value(event["hidden_params"])["batch_custom_id"]), + event["status"], + object_value(event["hidden_params"])["batch_line_status_code"], + event["call_type"], + ) + for event in events + ) + + +def test_a_bedrock_batch_the_poller_logs_emits_its_lines_from_the_deployment_bucket( + gateway: Gateway, tmp_path: Path +) -> None: + """The cost poller reads a managed Bedrock batch's files with the deployment's S3 bucket, so the + line events it starts must read the input file from that same bucket""" + delivered: Final = _line_events_of_a_polled_bedrock_batch(gateway, tmp_path, _BEDROCK_CHAT_BATCH) + + assert _line_outcomes(delivered) == [ + ("req-1", "success", 200, "acompletion"), + ("req-2", "failure", 400, "acompletion"), + ], delivered + + +def test_a_failed_line_of_a_bedrock_embedding_batch_is_logged_as_an_embedding(gateway: Gateway, tmp_path: Path) -> None: + delivered: Final = _line_events_of_a_polled_bedrock_batch(gateway, tmp_path, _BEDROCK_TITAN_BATCH) + + assert _line_outcomes(delivered) == [ + ("req-1", "success", 200, "aembedding"), + ("req-2", "failure", 400, "aembedding"), + ], delivered diff --git a/tests/unit/batches/test_batch_line_item_logging.py b/tests/unit/batches/test_batch_line_item_logging.py index 7936c4756c7a..9ada2fae71f8 100644 --- a/tests/unit/batches/test_batch_line_item_logging.py +++ b/tests/unit/batches/test_batch_line_item_logging.py @@ -162,9 +162,9 @@ def recorder(): litellm._async_failure_callback = saved_failure # test-quality-ok: teardown restoring the value set above -def _parent_logging(custom_llm_provider: str = "openai") -> Logging: +def _parent_logging(custom_llm_provider: str = "openai", model: str = "gpt-4o") -> Logging: logging_obj = Logging( - model="gpt-4o", + model=model, messages=[{"role": "user", "content": ""}], stream=False, call_type="aretrieve_batch", @@ -704,6 +704,32 @@ async def test_input_fetch_failure_still_emits_aggregate(recorder): ] ) +BEDROCK_TITAN_INPUT_JSONL = b"\n".join( + [ + json.dumps({"recordId": "titan-ok", "modelInput": {"inputText": "embed me"}}).encode(), + json.dumps({"recordId": "titan-bad", "modelInput": {"inputText": "embed me too", "dimensions": 7}}).encode(), + ] +) + +BEDROCK_TITAN_OUTPUT_JSONL = b"\n".join( + [ + json.dumps( + { + "recordId": "titan-ok", + "modelInput": {"inputText": "embed me"}, + "modelOutput": {"embedding": [0.1, 0.2], "inputTextTokenCount": 3}, + } + ).encode(), + json.dumps( + { + "recordId": "titan-bad", + "modelInput": {"inputText": "embed me too", "dimensions": 7}, + "error": {"errorCode": 400, "errorMessage": "Malformed input request"}, + } + ).encode(), + ] +) + MISTRAL_INPUT_JSONL = json.dumps( { "custom_id": "mis", @@ -741,6 +767,8 @@ def _edge_file_content(file_id: str, **_kwargs): "output-anth": ANTHROPIC_OUTPUT_JSONL, "input-bed": BEDROCK_INPUT_JSONL, "output-bed": BEDROCK_OUTPUT_JSONL, + "input-titan": BEDROCK_TITAN_INPUT_JSONL, + "output-titan": BEDROCK_TITAN_OUTPUT_JSONL, "input-mist": MISTRAL_INPUT_JSONL, "output-mist": MISTRAL_OUTPUT_JSONL, }[file_id] @@ -1009,9 +1037,39 @@ async def test_line_items_bedrock_shapes(recorder): assert [_hidden(by_id[line])["batch_line_status_code"] for line in ("br-anth", "br-titan", "br-conv")] == [200] * 3 (failure,) = recorder.failure_events assert (_hidden(failure)["batch_custom_id"], _hidden(failure)["batch_line_status_code"]) == ("br-err", 400) + assert failure["call_type"] == "acompletion" assert "bad request" in _payload(failure)["error_str"] +@pytest.mark.asyncio +async def test_bedrock_embedding_batch_lines_carry_the_embedding_call_type_and_input(recorder): + litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: flag under test; fixture restores it + with ( + patch( + "litellm.files.main.afile_content", new_callable=AsyncMock, side_effect=_edge_file_content + ), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch + patch( + "litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02) + ), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch + ): + await _log_completed_batch( + _parent_logging(custom_llm_provider="bedrock", model="bedrock/amazon.titan-embed-text-v2:0"), + _provider_batch("batch_titan", "input-titan", "output-titan"), + ) + + (success,) = (e for e in recorder.success_events if _hidden(e).get("batch_custom_id") == "titan-ok") + (failure,) = recorder.failure_events + assert [(_hidden(e)["batch_custom_id"], e["call_type"], _payload(e)["call_type"]) for e in (success, failure)] == [ + ("titan-ok", "aembedding", "aembedding"), + ("titan-bad", "aembedding", "aembedding"), + ] + assert _hidden(failure)["batch_line_status_code"] == 400 + assert [[m["content"] for m in _payload(e)["messages"]] for e in (success, failure)] == [ + ["embed me"], + ["embed me too"], + ] + + @pytest.mark.asyncio async def test_a_bedrock_batch_the_poller_logs_reads_every_file_from_the_deployment_bucket(recorder): """Bedrock trusts its S3 bucket only from the server-side credential snapshot, so the line From 54602427bcd542996db77b2a3a066ccc7fafd16c Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Mon, 12 Oct 2026 23:49:50 -0700 Subject: [PATCH 74/89] fix(batches): keep a Bedrock line's system prompt in its logged messages Bedrock batch input records hold the system prompt as content blocks, which logging only reads off the kwargs as text, so Claude and Nova lines dropped it --- litellm/batches/batch_line_item_logging.py | 7 ++++++- .../test_batch_output_file_listing.py | 18 +++++++++++++++-- .../batches/test_batch_line_item_logging.py | 20 +++++++++++++++---- 3 files changed, 38 insertions(+), 7 deletions(-) diff --git a/litellm/batches/batch_line_item_logging.py b/litellm/batches/batch_line_item_logging.py index 63c39e7fee3c..539660397e8c 100644 --- a/litellm/batches/batch_line_item_logging.py +++ b/litellm/batches/batch_line_item_logging.py @@ -245,7 +245,12 @@ def _call_type_for_line( def _line_messages(request_body: Mapping[str, object]) -> object: - return request_body.get("messages") or request_body.get("input") or request_body.get("inputText") or () + messages: Final = request_body.get("messages") or request_body.get("input") or request_body.get("inputText") or () + system: Final = request_body.get("system") + # Bedrock input records hold the system prompt as content blocks, and logging reads it off the kwargs only as text + if not isinstance(system, list) or not isinstance(messages, list): + return messages + return [{"role": "system", "content": system}, *messages] # pyright: ignore[reportUnknownVariableType] # JSONL values stay unknown, logging takes them as written def _line_model( diff --git a/tests/integration/management/test_batch_output_file_listing.py b/tests/integration/management/test_batch_output_file_listing.py index 8f093c4e5759..7aa10f0d27fa 100644 --- a/tests/integration/management/test_batch_output_file_listing.py +++ b/tests/integration/management/test_batch_output_file_listing.py @@ -21,7 +21,15 @@ from cryptography.hazmat.primitives import hashes, serialization from cryptography.hazmat.primitives.asymmetric import ec from cryptography.x509.oid import NameOID -from integration._support.client import JSON_OBJECT, Gateway, Scenario, eventually, object_value, string_value +from integration._support.client import ( + JSON_OBJECT, + Gateway, + Scenario, + eventually, + list_value, + object_value, + string_value, +) from integration._support.database import read_rows, scratch_database from integration._support.process import owned_proxy from integration._support.upstream import ScenarioHandle, delete_scenario, register_scenario @@ -728,7 +736,10 @@ class _BedrockLineBatch: def body(self, content: str) -> dict[str, JsonValue]: if self.endpoint == "/v1/embeddings": return {"input": content} - return {"messages": [{"role": "user", "content": content}], "max_tokens": 8} + return { + "messages": [{"role": "system", "content": "answer briefly"}, {"role": "user", "content": content}], + "max_tokens": 8, + } _BEDROCK_CHAT_BATCH: Final = _BedrockLineBatch(BEDROCK_MODEL_ID, "/v1/chat/completions", _bedrock_line_output) @@ -862,6 +873,9 @@ def test_a_bedrock_batch_the_poller_logs_emits_its_lines_from_the_deployment_buc ("req-1", "success", 200, "acompletion"), ("req-2", "failure", 400, "acompletion"), ], delivered + first_messages: Final = [object_value(list_value(event["messages"])[0]) for event in delivered] + assert [message["role"] for message in first_messages] == ["system", "system"], delivered + assert all("answer briefly" in json.dumps(message["content"]) for message in first_messages), delivered def test_a_failed_line_of_a_bedrock_embedding_batch_is_logged_as_an_embedding(gateway: Gateway, tmp_path: Path) -> None: diff --git a/tests/unit/batches/test_batch_line_item_logging.py b/tests/unit/batches/test_batch_line_item_logging.py index 9ada2fae71f8..f4eb96eb5ed5 100644 --- a/tests/unit/batches/test_batch_line_item_logging.py +++ b/tests/unit/batches/test_batch_line_item_logging.py @@ -636,7 +636,10 @@ async def test_input_fetch_failure_still_emits_aggregate(recorder): json.dumps( { "recordId": "br-anth", - "modelInput": {"messages": [{"role": "user", "content": "hi br-anth"}]}, + "modelInput": { + "system": [{"type": "text", "text": "answer like claude"}], + "messages": [{"role": "user", "content": "hi br-anth"}], + }, } ).encode(), json.dumps( @@ -648,7 +651,10 @@ async def test_input_fetch_failure_still_emits_aggregate(recorder): json.dumps( { "recordId": "br-conv", - "modelInput": {"messages": [{"role": "user", "content": "hi br-conv"}]}, + "modelInput": { + "system": [{"text": "answer like nova"}], + "messages": [{"role": "user", "content": "hi br-conv"}], + }, } ).encode(), json.dumps( @@ -1013,7 +1019,10 @@ async def test_line_items_bedrock_shapes(recorder): assert anthropic_event["call_type"] == "acompletion" assert _hidden(anthropic_event)["batch_custom_id"] == "br-anth" anthropic_line = _payload(anthropic_event) - assert any(m.get("content") == "hi br-anth" for m in anthropic_line["messages"]) + assert anthropic_line["messages"] == [ + {"role": "system", "content": [{"type": "text", "text": "answer like claude"}]}, + {"role": "user", "content": "hi br-anth"}, + ] assert anthropic_line["response"]["choices"][0]["message"]["content"] == "bedrock claude hi" assert anthropic_line["prompt_tokens"] == 2 assert anthropic_line["completion_tokens"] == 3 @@ -1029,7 +1038,10 @@ async def test_line_items_bedrock_shapes(recorder): converse_event = by_id["br-conv"] assert converse_event["call_type"] == "acompletion" converse_line = _payload(converse_event) - assert any(m.get("content") == "hi br-conv" for m in converse_line["messages"]) + assert converse_line["messages"] == [ + {"role": "system", "content": [{"text": "answer like nova"}]}, + {"role": "user", "content": "hi br-conv"}, + ] assert converse_line["response"]["choices"][0]["message"]["content"] == "converse hi" assert converse_line["prompt_tokens"] == 5 assert converse_line["completion_tokens"] == 6 From 4c120da83d99402b4fc07cdcff0148c74e586356 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Tue, 13 Oct 2026 02:08:29 -0700 Subject: [PATCH 75/89] feat(batches): store per-line batch requests and responses in S3/GCS via batch_line_item_callbacks Replace the per-line logging path with a dedicated hook. When the hourly batch cost poller finishes a managed batch, it pairs every output and error line with its input line, prices each line, and hands one payload to every logger listed in litellm_settings.batch_line_item_callbacks. s3_v2 and gcs_bucket write it as one NDJSON object per batch under batch_line_items//.jsonl. Bodies are dropped and the header kept when the creator key or team, the global setting, or the logger itself turns off message logging, and when the creator can no longer be found. A failed upload is logged and never blocks spend tracking. --- .../pagerduty/pagerduty.py | 3 - .../proxy/common_utils/check_batch_cost.py | 140 +- litellm/__init__.py | 2 +- litellm/batches/batch_line_item_logging.py | 671 ------ litellm/batches/batch_line_items.py | 170 ++ litellm/batches/batch_utils.py | 135 +- .../SlackAlerting/slack_alerting.py | 5 - .../integrations/batch_line_item_objects.py | 17 + .../clickhouse/clickhouse_spend_logger.py | 5 - litellm/integrations/custom_logger.py | 4 + .../datadog/datadog_cost_management.py | 3 - .../integrations/datadog/datadog_metrics.py | 5 - litellm/integrations/gcs_bucket/gcs_bucket.py | 17 +- litellm/integrations/lago.py | 5 - .../integrations/newrelic/newrelic_metrics.py | 5 - litellm/integrations/openmeter.py | 5 - litellm/integrations/opentelemetry.py | 3 - litellm/integrations/otel/plumbing/metrics.py | 5 - litellm/integrations/prometheus.py | 5 - litellm/integrations/s3_v2.py | 25 +- litellm/integrations/shadow_eval_logger.py | 8 +- litellm/litellm_core_utils/core_helpers.py | 12 - litellm/litellm_core_utils/litellm_logging.py | 104 +- litellm/llms/anthropic/chat/transformation.py | 15 - .../llms/base_llm/batches/transformation.py | 11 +- .../llms/bedrock/batches/transformation.py | 48 +- litellm/proxy/_types.py | 4 - .../proxy/hooks/dynamic_rate_limiter_v3.py | 3 - .../hooks/max_budget_per_session_limiter.py | 3 - .../proxy/hooks/model_max_budget_limiter.py | 3 - .../proxy/hooks/parallel_request_limiter.py | 9 +- .../hooks/parallel_request_limiter_v3.py | 5 - .../proxy/hooks/proxy_track_cost_callback.py | 3 - litellm/proxy/proxy_server.py | 27 +- litellm/router.py | 9 +- litellm/router_strategy/budget_limiter.py | 3 - litellm/router_strategy/least_busy.py | 3 +- litellm/router_strategy/lowest_cost.py | 5 +- litellm/router_strategy/lowest_latency.py | 12 +- litellm/router_strategy/lowest_tpm_rpm.py | 5 +- litellm/router_strategy/lowest_tpm_rpm_v2.py | 6 +- .../pre_call_checks/model_rate_limit_check.py | 5 - .../prompt_caching_deployment_check.py | 3 - litellm/types/utils.py | 41 +- .../test_batch_output_file_listing.py | 244 +- .../test_batch_line_items_cold_storage.py | 499 ++++ .../spend/test_batch_line_item_callbacks.py | 817 ------- .../spend/test_batch_parent_id_body_field.py | 254 --- .../batches/test_batch_line_item_logging.py | 2003 ----------------- tests/unit/batches/test_batch_line_items.py | 258 +++ tests/unit/batches/test_batch_utils.py | 82 - .../test_clickhouse_spend_logger.py | 86 - .../gcs_bucket/test_gcs_bucket.py | 38 + .../test_batch_line_item_metering_sinks.py | 484 ---- tests/unit/integrations/test_s3_v2.py | 79 + .../integrations/test_shadow_eval_logger.py | 16 - .../litellm_core_utils/test_core_helpers.py | 11 - .../test_litellm_logging.py | 342 +-- .../bedrock/batches/test_transformation.py | 24 +- .../common_utils/test_check_batch_cost.py | 312 ++- .../hooks/test_model_max_budget_limiter.py | 14 - .../hooks/test_parallel_request_limiter.py | 26 - .../hooks/test_parallel_request_limiter_v3.py | 56 - .../hooks/test_proxy_track_cost_callback.py | 18 - .../unit/proxy/proxy_server/test_lifecycle.py | 32 - ...test_proxy_server_endpoints_and_startup.py | 38 - .../unit/proxy/test_redis_auth_cache_flag.py | 3 - .../test_budget_limiter_hotpath.py | 23 - tests/unit/router_strategy/test_least_busy.py | 25 - .../router_strategy/test_lowest_tpm_rpm.py | 37 - .../test_prompt_caching_deployment_check.py | 35 - .../test_enforce_model_rate_limits.py | 31 - tests/unit/test_router/test_router.py | 122 - ui/litellm-dashboard/src/lib/http/schema.d.ts | 5 - 74 files changed, 1715 insertions(+), 5876 deletions(-) delete mode 100644 litellm/batches/batch_line_item_logging.py create mode 100644 litellm/batches/batch_line_items.py create mode 100644 litellm/integrations/batch_line_item_objects.py create mode 100644 tests/integration/observability/test_batch_line_items_cold_storage.py delete mode 100644 tests/integration/spend/test_batch_line_item_callbacks.py delete mode 100644 tests/integration/spend/test_batch_parent_id_body_field.py delete mode 100644 tests/unit/batches/test_batch_line_item_logging.py create mode 100644 tests/unit/batches/test_batch_line_items.py delete mode 100644 tests/unit/integrations/test_batch_line_item_metering_sinks.py diff --git a/enterprise/litellm_enterprise/enterprise_callbacks/pagerduty/pagerduty.py b/enterprise/litellm_enterprise/enterprise_callbacks/pagerduty/pagerduty.py index c437345a98d7..f920aa7ac13b 100644 --- a/enterprise/litellm_enterprise/enterprise_callbacks/pagerduty/pagerduty.py +++ b/enterprise/litellm_enterprise/enterprise_callbacks/pagerduty/pagerduty.py @@ -17,7 +17,6 @@ from litellm._logging import verbose_logger from litellm.caching import DualCache -from litellm.litellm_core_utils.core_helpers import is_batch_line_item_event from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, @@ -87,8 +86,6 @@ async def async_log_failure_event(self, kwargs, response_obj, start_time, end_ti Record a failure event. Only send an alert to PagerDuty if the configured *failure* threshold is exceeded in the specified window. """ - if is_batch_line_item_event(kwargs): - return now = datetime.now(timezone.utc) standard_logging_payload: Optional[StandardLoggingPayload] = kwargs.get( "standard_logging_object" 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 bec7b6ca1329..7643ac63c027 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -2,16 +2,18 @@ Polls LiteLLM_ManagedObjectTable to check if the batch job is complete, and if the cost has been tracked. """ -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from dataclasses import replace as dataclasses_replace from datetime import datetime, timedelta, timezone from types import MappingProxyType from typing import TYPE_CHECKING, Final, List, Literal, Optional, Protocol, Tuple, cast +import litellm from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid from litellm.constants import ( CLI_SESSION_KEY_PREFIX, + LITELLM_PROXY_MASTER_KEY_ALIAS, MANAGED_OBJECT_STALENESS_CUTOFF_DAYS, MAX_OBJECTS_PER_POLL_CYCLE, ) @@ -23,6 +25,7 @@ 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 @@ -103,6 +106,11 @@ def _team_table(prisma_client: "PrismaClient") -> "TableActions[prisma_models.Li 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, @@ -735,10 +743,9 @@ async def _track_completed_batch_cost( """ from litellm.batches.batch_utils import ( + count_error_file_failed_requests, get_file_content_as_dictionary, calculate_batch_cost_and_usage, - count_error_file_failed_requests_from_content, - fetch_batch_error_file_content, ) from litellm.files.main import afile_content from litellm.files.types import FileContentCallOptions @@ -875,7 +882,7 @@ async def _track_completed_batch_cost( model_name=model_name, model_info=deployment_model_info, ) - error_file_content = await fetch_batch_error_file_content( + error_file_failed_requests: Final = await count_error_file_failed_requests( response, custom_llm_provider=batch_file_provider, litellm_params={ @@ -883,7 +890,6 @@ async def _track_completed_batch_cost( "_litellm_internal_model_credentials": MappingProxyType(dict(credentials)), }, ) - error_file_failed_requests: Final = count_error_file_failed_requests_from_content(error_file_content) batch_result: Final = ( output_file_result if not error_file_failed_requests @@ -923,7 +929,6 @@ async def _track_completed_batch_cost( optional_params={}, custom_llm_provider=str(llm_provider) if llm_provider else None, ) - logging_obj._litellm_internal_model_credentials = MappingProxyType(dict(credentials)) # pyright: ignore[reportPrivateUsage] # trusted credentials transport, consumed by batch_line_item_logging if not await self._claim_job_for_costing(job): verbose_proxy_logger.info( @@ -942,8 +947,6 @@ async def _track_completed_batch_cost( batch_failed_requests=batch_result.failed_requests, batch_prompt_cost=batch_result.prompt_cost, batch_completion_cost=batch_result.completion_cost, - batch_output_file_content=content_bytes, - batch_error_file_content=error_file_content, ) except Exception: await self._release_job_claim(job) @@ -959,8 +962,128 @@ async def _track_completed_batch_cost( api_provider=str(llm_provider) if llm_provider else None, ) + await self._store_batch_line_items(job, response, model_id, batch_id, content_bytes) + return model_name, str(llm_provider) if llm_provider else None + async def _store_batch_line_items( + self, + job: "_ManagedObjectRow", + response: "LiteLLMBatch", + model_id: str, + batch_id: str, + output_content: bytes, + ) -> None: + """Send every line of a finished batch to batch_line_item_callbacks. Failures are logged, never raised, + so the batch stays billed whether or not its lines were stored.""" + if not litellm.batch_line_item_callbacks: + return + from litellm.batches.batch_line_items import ( + BatchFileProvider, + BatchLineItemSource, + build_batch_line_items, + dispatch_batch_line_items, + line_item_metadata, + ) + 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: + 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)), + } + 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 + 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._build_creator_attribution_metadata(job, batch_id)), + ) + payload: Final = await build_batch_line_items( + source, + await self._batch_file_content(response.input_file_id, cost_provider, file_access), + (output_content, error_content), + ) + await dispatch_batch_line_items(payload, redact_bodies=await self._creator_turned_off_message_logging(job)) + 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 key or team that can no longer be read counts as turned off.""" + 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 api_key not in ( + LITELLM_PROXY_MASTER_KEY_ALIAS, + f"{CLI_SESSION_KEY_PREFIX}-{job.created_by}", + ) + 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 api_key is None or (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): """ Check if the batch JOB has been tracked. @@ -1117,6 +1240,7 @@ async def check_batch_cost(self): f"has not appeared yet; leaving job {job.id} for the next poll cycle" ) continue + await self._store_batch_line_items(job, response, model_id, batch_id, b"") await self._finalize_unbilled_terminal_job(job, response) # Record polling run metrics (always, even if nothing was processed) diff --git a/litellm/__init__.py b/litellm/__init__.py index ce6513a7c850..8b2173decceb 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -118,6 +118,7 @@ def _dev_env_hot_reload_enabled() -> bool: 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", @@ -348,7 +349,6 @@ def _dev_env_hot_reload_enabled() -> bool: ) openai_system_messages_first: bool = False disable_vertex_batch_output_transformation: bool = False -store_batch_line_items_in_callbacks: bool = False extra_spend_tag_headers: Optional[List[str]] = None in_memory_llm_clients_cache: "LLMClientCache" safe_memory_mode: bool = False diff --git a/litellm/batches/batch_line_item_logging.py b/litellm/batches/batch_line_item_logging.py deleted file mode 100644 index 539660397e8c..000000000000 --- a/litellm/batches/batch_line_item_logging.py +++ /dev/null @@ -1,671 +0,0 @@ -import asyncio -import json -import threading -import uuid -import weakref -from collections.abc import Iterator, Mapping -from datetime import datetime -from types import MappingProxyType -from typing import TYPE_CHECKING, Final, Literal, TypeAlias - -from typing_extensions import assert_never - -from litellm._internal_context import with_service_target -from litellm._logging import verbose_logger -from litellm.batches.batch_utils import ( - BatchProvider, - BatchResultFiles, - _batch_response_was_successful, # pyright: ignore[reportPrivateUsage] # batch-internal helper shared with the aggregate cost path by design - _fetch_batch_managed_file_content, # pyright: ignore[reportPrivateUsage, reportUnknownVariableType] # same reuse; helper is untyped upstream - _get_response_from_batch_job_output_file, # pyright: ignore[reportPrivateUsage] # same reuse - _iter_batch_output_entries, # pyright: ignore[reportPrivateUsage, reportUnknownVariableType] # same reuse; helper is untyped upstream - _safe_output_line_stats, # pyright: ignore[reportPrivateUsage] # same reuse - _uses_native_vertex_output, # pyright: ignore[reportPrivateUsage] # same reuse -) -from litellm.caching.caching import DualCache -from litellm.caching.in_memory_cache import InMemoryCache -from litellm.types.llms.openai import ResponsesAPIResponse -from litellm.types.utils import ( - EmbeddingResponse, - LiteLLMBatch, - ModelInfo, - ModelResponse, - Usage, -) - -if TYPE_CHECKING: - from litellm.litellm_core_utils.litellm_logging import Logging - -_SECRET_PARAM_KEYS: Final = frozenset( - { - "api_key", - "_litellm_internal_model_credentials", - "azure_ad_token", - "azure_ad_token_provider", - "vertex_credentials", - "aws_access_key_id", - "aws_secret_access_key", - "aws_session_token", - "aws_web_identity_token", - } -) - - -def _supported_line_provider(value: str) -> BatchProvider | None: - match value: - case "openai" | "azure" | "vertex_ai" | "hosted_vllm" | "anthropic" | "bedrock" | "mistral": - return value - case _: - return None - - -_CALL_TYPE_BY_BATCH_URL: Final = MappingProxyType( - { - "/v1/chat/completions": "acompletion", - "/v1/embeddings": "aembedding", - "/v1/responses": "aresponses", - } -) - -_EMPTY_BODY: Final[Mapping[str, object]] = MappingProxyType({}) - -_LINE_ITEM_CLAIM_TTL_SECONDS: Final = 30 * 24 * 60 * 60 - -# Without Redis each worker holds its claims in memory for the claim ttl, so the -# in-memory tier keeps far more than its 200-entry default before evicting a claim -_IN_MEMORY_CLAIM_CAPACITY: Final = 10_000 - -batch_line_item_claim_cache: Final = DualCache( - in_memory_cache=InMemoryCache( - max_size_in_memory=_IN_MEMORY_CLAIM_CAPACITY, default_ttl=_LINE_ITEM_CLAIM_TTL_SECONDS - ) -) - -# asyncio holds tasks weakly, so each spawned line fan-out stays referenced here until it finishes -_LINE_ITEM_TASKS: Final[set["asyncio.Task[int]"]] = set() # mutable-ok: strong refs to running fan-outs -# loops whose shutdown drain has started, where a new fan-out would be cancelled with its claim held -_DRAINED_LOOPS: Final[weakref.WeakSet[asyncio.AbstractEventLoop]] = weakref.WeakSet() - -_SYSTEM_PROMPT_KEYS: Final = ("system", "instructions") - -# Anthropic keeps no input file, so its retrieved batch reports input_file_id as the string "None" -_NO_FILE_IDS: Final = frozenset(("", "None")) - -_ClaimResult: TypeAlias = Literal["claimed", "already_claimed", "unavailable"] - - -@with_service_target("batch_line_items") -async def _claim_line_items(claim_cache: DualCache, claim_key: str, token: str) -> _ClaimResult: - redis_cache: Final = claim_cache.redis_cache - if redis_cache is None: - count: Final = await claim_cache.async_increment_cache(claim_key, 1, ttl=_LINE_ITEM_CLAIM_TTL_SECONDS) # pyright: ignore[reportUnknownMemberType] # DualCache.increment is untyped upstream - if count is None or count == 1: - return "claimed" - return "already_claimed" - try: - ok: Final = await redis_cache.async_set_cache(claim_key, token, ttl=_LINE_ITEM_CLAIM_TTL_SECONDS, nx=True) # pyright: ignore[reportUnknownMemberType] # RedisCache.set is untyped upstream - holder: Final = None if ok else await redis_cache.async_get_cache(claim_key) # pyright: ignore[reportUnknownMemberType] # RedisCache.get is untyped upstream - except Exception: # noqa: BLE001 # a redis outage must not emit duplicates; the next retrieve retries - return "unavailable" - if ok: - return "claimed" - return "already_claimed" if holder is not None else "unavailable" - - -@with_service_target("batch_line_items") -async def _release_line_item_claim(claim_cache: DualCache, claim_key: str, token: str) -> None: - try: - redis_cache: Final = claim_cache.redis_cache - if redis_cache is None: - await claim_cache.async_delete_cache(claim_key) - return - owner: Final = await redis_cache.async_get_cache(claim_key) # pyright: ignore[reportUnknownMemberType] # RedisCache.get is untyped upstream - if owner == token: - await redis_cache.async_delete_cache(claim_key) - except Exception: # noqa: BLE001 # the claim release must never raise; worst case the batch stays claimed - verbose_logger.debug( - "batch line item claim release failed for %s, claim persists until ttl", - claim_key, - ) - - -def _json_fallback(value: object) -> Mapping[str, object] | str: - mapping: Final = _as_object_mapping(value) - if mapping is None: - return str(value) - return dict(mapping) - - -class _BatchLineFailure(Exception): - """A provider-reported per-line batch failure; carries the batch's hidden - params so the failure logging payload can attribute the line.""" - - def __init__(self, error_payload: object) -> None: - super().__init__(json.dumps(error_payload, default=_json_fallback)) - self._hidden_params: dict[str, object] = {} # mutable-ok: plain-dict contract like response _hidden_params - - -def _as_object_mapping(value: object) -> Mapping[str, object] | None: - if isinstance(value, Mapping) and all(isinstance(key, str) for key in value): # pyright: ignore[reportUnknownVariableType] # keys of an unparameterized Mapping are unknown until checked here - return value # pyright: ignore[reportUnknownVariableType, reportReturnType] # every key was verified str above - return None - - -def _output_entries(file_content: bytes) -> Iterator[Mapping[str, object]]: - for entry in _iter_batch_output_entries(file_content): # pyright: ignore[reportUnknownVariableType] # entries are validated into typed mappings below - mapping = _as_object_mapping(entry) # pyright: ignore[reportUnknownArgumentType] # raw entry is unknown until validated here - if mapping is not None: - yield mapping - - -def _line_id(entry: Mapping[str, object]) -> str | None: - line_id: Final = entry.get("custom_id") or entry.get("recordId") - return line_id if isinstance(line_id, str) and line_id else None - - -def _requests_by_custom_id(input_file_content: bytes) -> Mapping[str, Mapping[str, object]]: - """Parse the batch input JSONL into {line id: request line}, keyed by - custom_id or recordId, skipping malformed lines and lines without either.""" - return MappingProxyType( - {line_id: entry for entry in _output_entries(input_file_content) if (line_id := _line_id(entry)) is not None} - ) - - -def _request_body_for_entry( - entry: Mapping[str, object], request_line: Mapping[str, object] | None -) -> Mapping[str, object]: - if request_line is not None: - body: Final = _as_object_mapping(request_line.get("body")) - if body: - return body - params: Final = _as_object_mapping(request_line.get("params")) - if params: - return params - request_model_input: Final = _as_object_mapping(request_line.get("modelInput")) - if request_model_input: - return request_model_input - model_input: Final = _as_object_mapping(entry.get("modelInput")) - return model_input if model_input else _EMPTY_BODY - - -def _line_status_code(entry: Mapping[str, object], custom_llm_provider: str) -> int | None: - response: Final = _as_object_mapping(entry.get("response")) - status: Final = response.get("status_code") if response is not None else None - if isinstance(status, int): - return status - if custom_llm_provider == "anthropic": - result: Final = _as_object_mapping(entry.get("result")) - if result is not None and result.get("type") == "succeeded": - return 200 - if custom_llm_provider == "bedrock": - error: Final = _as_object_mapping(entry.get("error")) - error_code: Final = error.get("errorCode") if error is not None else None - if isinstance(error_code, int): - return error_code - if _batch_response_was_successful(entry, custom_llm_provider): - return 200 - return None - - -def _line_error_payload(entry: Mapping[str, object], custom_llm_provider: BatchProvider) -> object: - if custom_llm_provider == "anthropic": - return ( - (_as_object_mapping(entry.get("result")) or _EMPTY_BODY).get("error") or entry.get("result") or _EMPTY_BODY - ) - return entry.get("error") or entry.get("response") or _EMPTY_BODY - - -def _call_type_for_request(request_line: Mapping[str, object] | None) -> str: - url: Final = request_line.get("url") if request_line is not None else None - return _CALL_TYPE_BY_BATCH_URL.get(url if isinstance(url, str) else "", "acompletion") - - -def _is_bedrock_embedding_batch_model(model: str) -> bool: - from litellm.llms.bedrock.files.transformation import BedrockFilesConfig - - return BedrockFilesConfig._is_titan_v2_embed_model(model) # pyright: ignore[reportPrivateUsage] # the check that chose the embedding record shape when the batch input was written - - -def _call_type_for_line( - request_line: Mapping[str, object] | None, - result: "_BatchLineResult | None", - custom_llm_provider: BatchProvider, - model: str, -) -> str: - if isinstance(result, EmbeddingResponse): - return "aembedding" - if isinstance(result, ResponsesAPIResponse): - return "aresponses" - if isinstance(result, ModelResponse): - return "acompletion" - # Bedrock input records carry no url, so a failed line takes its type from the deployment model - if custom_llm_provider == "bedrock" and _is_bedrock_embedding_batch_model(model): - return "aembedding" - return _call_type_for_request(request_line) - - -def _line_messages(request_body: Mapping[str, object]) -> object: - messages: Final = request_body.get("messages") or request_body.get("input") or request_body.get("inputText") or () - system: Final = request_body.get("system") - # Bedrock input records hold the system prompt as content blocks, and logging reads it off the kwargs only as text - if not isinstance(system, list) or not isinstance(messages, list): - return messages - return [{"role": "system", "content": system}, *messages] # pyright: ignore[reportUnknownVariableType] # JSONL values stay unknown, logging takes them as written - - -def _line_model( - response_body: Mapping[str, object], - request_body: Mapping[str, object], - parent: "Logging", -) -> str: - for candidate in (response_body.get("model"), request_body.get("model"), parent.model): - if isinstance(candidate, str) and candidate: - return candidate - return "" - - -_BatchLineResult: TypeAlias = "ModelResponse | EmbeddingResponse | ResponsesAPIResponse" - - -def _line_result( - call_type: str, - custom_llm_provider: BatchProvider, - model: str, - response_body: Mapping[str, object], -) -> "_BatchLineResult | None": - from litellm.types.utils import LlmProviders - from litellm.utils import ProviderConfigManager - - provider_config: Final = ProviderConfigManager.get_provider_batches_config(model, LlmProviders(custom_llm_provider)) - if provider_config is not None: - transformed: Final = provider_config.transform_batch_output_line(response_body, model) - if transformed is not None: - return transformed - if custom_llm_provider == "bedrock": - return None - if call_type == "aembedding": - return EmbeddingResponse(**response_body) # pyright: ignore[reportArgumentType] # provider output bodies are dicts expanded as response ctor kwargs - if call_type == "aresponses": - return ResponsesAPIResponse(**response_body) # pyright: ignore[reportArgumentType] # same as above - if custom_llm_provider == "anthropic": - from litellm.llms.anthropic.chat.transformation import anthropic_message_to_model_response - - return anthropic_message_to_model_response(response_body, None) - return ModelResponse(**response_body) # pyright: ignore[reportArgumentType] # same as above - - -def _line_result_or_none( - request_call_type: str, - custom_llm_provider: BatchProvider, - model: str, - response_body: Mapping[str, object], - custom_id: object, -) -> "_BatchLineResult | None": - try: - result: Final = _line_result(request_call_type, custom_llm_provider, model, response_body) - except Exception: # noqa: BLE001 # one unparseable line must not drop the rest of the batch's line events - return _skipped_line(request_call_type, custom_id) - return result if result is not None else _skipped_line(request_call_type, custom_id) - - -def _skipped_line(request_call_type: str, custom_id: object) -> None: - verbose_logger.warning( - "batch output line could not be reconstructed as a %s response, skipping it. custom_id=%s", - request_call_type, - custom_id, - ) - - -def _new_child_logging( - parent: "Logging", - model: str, - messages: object, - call_type: str, - start_time: datetime, -) -> "Logging": - from litellm.litellm_core_utils.litellm_logging import Logging - - child: Final = Logging( - model=model, - messages=messages, - stream=False, - call_type=call_type, - start_time=start_time, - litellm_call_id=str(uuid.uuid4()), - function_id=str(uuid.uuid4()), - litellm_trace_id=parent.litellm_trace_id, - dynamic_success_callbacks=parent.dynamic_success_callbacks, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # Logging ctor takes these dynamic callback lists untyped - dynamic_async_success_callbacks=parent.dynamic_async_success_callbacks, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # same as above - dynamic_failure_callbacks=parent.dynamic_failure_callbacks, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # same as above - dynamic_async_failure_callbacks=parent.dynamic_async_failure_callbacks, # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # same as above - kwargs={"litellm_session_id": parent.litellm_session_id}, - ) - # Lines are logged under the retrieve request's key and team callback settings, - # such as turn_off_message_logging and per-team logging credentials - child.standard_callback_dynamic_params = parent.standard_callback_dynamic_params.copy() - return child - - -def _line_hidden_params( - batch: LiteLLMBatch, - custom_id: object, - status_code: int | None, - response_cost: float | None = None, -) -> dict[str, object]: # mutable-ok: response objects declare _hidden_params as a plain dict - return { - "batch_id": batch.id, - "batch_custom_id": custom_id, - "batch_line_status_code": status_code, - "response_cost": response_cost, - } - - -def _optional_params_for_body( - request_body: Mapping[str, object], -) -> dict[str, object]: # mutable-ok: update_environment_variables takes a plain dict - return {key: value for key, value in request_body.items() if key not in ("model", "messages", "input")} - - -def _metadata_copy(params: Mapping[str, object]) -> dict[str, object]: # mutable-ok: dict out for litellm_params - metadata: Final = _as_object_mapping(params.get("metadata")) or _EMPTY_BODY - return {**metadata} - - -async def _emit_line_event( - entry: Mapping[str, object], - requests_by_id: Mapping[str, Mapping[str, object]], - batch: LiteLLMBatch, - custom_llm_provider: BatchProvider, - parent: "Logging", - parent_params: Mapping[str, object], - model_name: str | None, - model_info: ModelInfo | None, -) -> bool: - custom_id: Final = _line_id(entry) - request_line: Final = requests_by_id.get(custom_id or "") - request_body: Final = _request_body_for_entry(entry, request_line) - status_code: Final = _line_status_code(entry, custom_llm_provider) - request_call_type: Final = _call_type_for_request(request_line) - # A failed line's body can be the provider's error as text (Mistral writes a JSON string), not an object - response_body: Final = ( - _as_object_mapping(_get_response_from_batch_job_output_file(entry, custom_llm_provider)) or _EMPTY_BODY - ) - start_time: Final = parent.start_time - model: Final = _line_model(response_body, request_body, parent) - - successful: Final = _batch_response_was_successful(entry, custom_llm_provider) - stats: Final = _safe_output_line_stats(entry, custom_llm_provider, model_name, model_info) if successful else None - result: Final[_BatchLineResult | None] = ( - _line_result_or_none(request_call_type, custom_llm_provider, model, response_body, custom_id) - if successful - else None - ) - if successful and result is None: - return False - - call_type: Final = _call_type_for_line(request_line, result, custom_llm_provider, model) - child: Final = _new_child_logging( - parent=parent, - model=model, - messages=_line_messages(request_body), - call_type=call_type, - start_time=start_time, - ) - line_user: Final = request_body.get("user") - child.update_environment_variables( # pyright: ignore[reportUnknownMemberType] # Logging.update_environment_variables is untyped upstream - litellm_params={ - **parent_params, - "metadata": _metadata_copy(parent_params), - # Flags the line as the async call it stands for, so the logging dispatch runs the - # callbacks a direct async call of that type runs, legacy sync ones included - call_type: True, - }, - optional_params={}, - model=child.model, - user=line_user if isinstance(line_user, str) else None, - custom_llm_provider=custom_llm_provider, - ) - # The line body is caller data, so it is kept as optional_params without its keys - # landing on the callback kwargs, where router hooks read metadata and call_type. - # The system prompt is the one exception: logging reads it off the kwargs, as for a direct call - line_optional_params: Final = _optional_params_for_body(request_body) - child.optional_params = line_optional_params - child.model_call_details["optional_params"] = line_optional_params - child.model_call_details.update({key: request_body[key] for key in _SYSTEM_PROMPT_KEYS if key in request_body}) - child.mark_batch_line_item(batch.id) - for secret_key in _SECRET_PARAM_KEYS: - child.litellm_params.pop(secret_key, None) # pyright: ignore[reportUnknownMemberType] # Logging.litellm_params is untyped upstream - - now: Final = datetime.now() # noqa: DTZ005 # naive to match the logging pipeline start_time - if result is None: - exception: Final = _BatchLineFailure(_line_error_payload(entry, custom_llm_provider)) - exception._hidden_params = _line_hidden_params(batch, custom_id, status_code) # pyright: ignore[reportPrivateUsage] # _hidden_params is set on the exception instance itself - await child.dispatch_failure_handlers(exception, "", prefer_async_handlers=True) - return True - - result._hidden_params = _line_hidden_params( # pyright: ignore[reportPrivateUsage] # same hidden_params channel the aggregate batch event uses - batch, - custom_id, - status_code, - response_cost=stats.prompt_cost + stats.completion_cost if stats is not None else 0.0, - ) - if stats is not None and not response_body.get("usage") and isinstance(result, (ModelResponse, EmbeddingResponse)): - setattr( # noqa: B010 # ModelResponse.usage is set dynamically by its ctor, so setattr keeps parity for both response types - result, - "usage", - Usage( - prompt_tokens=stats.prompt_tokens, - completion_tokens=stats.completion_tokens, - total_tokens=stats.total_tokens, - ), - ) - await child.dispatch_success_handlers( - result=result, - start_time=start_time, - end_time=now, - cache_hit=False, - prefer_async_handlers=True, - ) - return True - - -async def _fetch_managed_file_or_empty( - file_id: str | None, - custom_llm_provider: BatchProvider, - fetch_params: dict[str, object] | None, # mutable-ok: batch_utils file fetch takes the shared litellm_params dict -) -> bytes: - if file_id is None or file_id in _NO_FILE_IDS: - return b"" - return await _fetch_batch_managed_file_content( - file_id, - custom_llm_provider=custom_llm_provider, - litellm_params=fetch_params, # pyright: ignore[reportArgumentType] # batch_utils types this param as an unparameterized dict - ) - - -def spawn_batch_line_items( - batch: LiteLLMBatch, - custom_llm_provider: str, - parent: "Logging", - model_name: str | None, - litellm_params: dict[str, object] | None, # mutable-ok: the logging object's shared litellm_params dict - model_info: ModelInfo | None, - result_files: BatchResultFiles | None, -) -> "asyncio.Task[int] | None": - """Run ``log_batch_line_items`` in its own task, so the fan-out never shares - the logging worker deadline the aggregate aretrieve_batch event runs under. - The parent's litellm_params are copied first, because the aggregate event's - callbacks keep writing request-local markers into its metadata. Once this - loop's shutdown drain has started, or the interpreter is exiting and the - logging worker's exit flush runs the handler on a loop it closes right after, - no fan-out starts or claims the batch, so the next retrieve emits every line.""" - loop: Final = asyncio.get_running_loop() - if loop in _DRAINED_LOOPS or not threading.main_thread().is_alive(): - verbose_logger.warning( - "batch line items for batch_id=%s are skipped during shutdown and will be emitted on the next retrieve", - batch.id, - ) - return None - task: Final = loop.create_task( - log_batch_line_items( - batch=batch, - custom_llm_provider=custom_llm_provider, - parent=parent, - model_name=model_name, - litellm_params=litellm_params, - model_info=model_info, - result_files=result_files, - parent_params=parent_params_snapshot(parent), - ) - ) - _LINE_ITEM_TASKS.add(task) - task.add_done_callback(_LINE_ITEM_TASKS.discard) - task.add_done_callback(_log_line_item_task_failure) - return task - - -def parent_params_snapshot(parent: "Logging") -> Mapping[str, object]: - """The parent's litellm_params with their own copy of the metadata dict.""" - params: Final = _as_object_mapping(parent.litellm_params) or _EMPTY_BODY # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # Logging.litellm_params is untyped upstream - return MappingProxyType({**params, "metadata": _metadata_copy(params)}) - - -def pending_batch_line_item_tasks() -> tuple["asyncio.Task[int]", ...]: - """The spawned line fan-outs that have not finished yet.""" - return tuple(_LINE_ITEM_TASKS) - - -async def drain_batch_line_items(timeout: float) -> int: - """Stop this loop from starting new line fan-outs, then wait up to ``timeout`` - seconds for the running ones, so a shutdown lets them finish before it closes - the clients they write through. Returns how many are still running when the - wait ends.""" - loop: Final = asyncio.get_running_loop() - _DRAINED_LOOPS.add(loop) - running: Final = [task for task in pending_batch_line_item_tasks() if task.get_loop() is loop and not task.done()] - if not running: - return 0 - _, still_running = await asyncio.wait(running, timeout=timeout) - if still_running: - verbose_logger.warning( - "%d batch line item fan-out(s) still running after waiting %ss at shutdown; lines they have not " - "sent yet are skipped until the batch claim expires", - len(still_running), - timeout, - ) - return len(still_running) - - -def _log_line_item_task_failure(task: "asyncio.Task[int]") -> None: - if task.cancelled(): - verbose_logger.warning("batch line item logging was cancelled before it finished") - return - error: Final = task.exception() - if error is not None: - verbose_logger.error("batch line item logging failed; aggregate logging unaffected", exc_info=error) - - -async def log_batch_line_items( - batch: LiteLLMBatch, - custom_llm_provider: str, - parent: "Logging", - model_name: str | None, - litellm_params: dict[str, object] | None, # mutable-ok: the logging object's shared litellm_params dict - model_info: ModelInfo | None, - result_files: BatchResultFiles | None = None, - claim_cache: DualCache = batch_line_item_claim_cache, - parent_params: Mapping[str, object] | None = None, -) -> int: - """Emit one callback event per JSONL line of a completed batch (request - paired with its response/error), behind the opt-in - ``litellm.store_batch_line_items_in_callbacks`` flag. The aggregate - aretrieve_batch event still bills the batch, so per-line events carry - ``batch_parent_id`` and never update spend themselves. Any failure here - is logged and swallowed: aggregate accounting must be unaffected. - ``result_files`` carries the output/error bytes the aggregate path already - fetched, so they are reused instead of refetched. ``parent_params`` is the - parent's litellm_params as of the aggregate event, read now when omitted.""" - line_params: Final = parent_params if parent_params is not None else parent_params_snapshot(parent) - line_provider: Final = _supported_line_provider(custom_llm_provider) - if line_provider is None: - verbose_logger.warning( - "batch line-item callbacks are not supported for provider %s, skipping. batch_id=%s", - custom_llm_provider, - batch.id, - ) - return 0 - claim_key: Final = f"batch_line_items_emitted:{batch.id}" - token: Final = uuid.uuid4().hex - claim: Final = await _claim_line_items(claim_cache, claim_key, token) - match claim: - case "already_claimed": - verbose_logger.debug("batch line items already emitted for batch_id=%s, skipping", batch.id) - return 0 - case "unavailable": - verbose_logger.warning( - "batch line item claim backend unavailable for batch_id=%s, line items will be retried on the next retrieve", - batch.id, - ) - return 0 - case "claimed": - pass - case _: - assert_never(claim) - emitted = 0 # rebind-ok: loop accumulator for emitted line count - try: - internal_credentials: Final = parent._litellm_internal_model_credentials # pyright: ignore[reportPrivateUsage] # declared transport attribute on Logging - internal_mapping: Final = _as_object_mapping(internal_credentials) - # Cloud storage providers resolve their bucket only from the trusted snapshot key, so it - # travels with the flattened credentials, the same shape the cost poller fetches with - fetch_params: Final[dict[str, object] | None] = ( # mutable-ok: file fetcher requires a plain dict - {**internal_mapping, "_litellm_internal_model_credentials": MappingProxyType(dict(internal_mapping))} - if internal_mapping is not None - else litellm_params - ) - - output_content: Final = ( - result_files.output - if result_files is not None and result_files.output is not None - else await _fetch_managed_file_or_empty(batch.output_file_id, line_provider, fetch_params) - ) - first_row: Final = next(_output_entries(output_content), None) - if _uses_native_vertex_output(line_provider, model_name, first_row): - verbose_logger.warning( - "batch line-item callbacks do not support native vertex_ai batch output rows yet, skipping. batch_id=%s", - batch.id, - ) - return 0 - input_file_content: Final = await _fetch_managed_file_or_empty(batch.input_file_id, line_provider, fetch_params) - requests_by_id: Final = _requests_by_custom_id(input_file_content) - error_content: Final = ( - result_files.error - if result_files is not None and result_files.error is not None - else await _fetch_managed_file_or_empty(batch.error_file_id, line_provider, fetch_params) - ) - for content in (output_content, error_content): - for entry in _output_entries(content): - try: - emitted += await _emit_line_event( - entry=entry, - requests_by_id=requests_by_id, - batch=batch, - custom_llm_provider=line_provider, - parent=parent, - parent_params=line_params, - model_name=model_name, - model_info=model_info, - ) - except Exception: # noqa: BLE001 # one bad line must not drop the rest of the batch's line events - verbose_logger.exception( - "batch line item logging failed for entry, continuing with remaining lines. batch_id=%s", - batch.id, - ) - except Exception: # noqa: BLE001 # line-item logging must never break the aggregate aretrieve_batch accounting - verbose_logger.exception( - "batch line item logging failed for batch_id=%s; aggregate logging unaffected", - batch.id, - ) - finally: - if emitted == 0: - await _release_line_item_claim(claim_cache, claim_key, token) - return emitted diff --git a/litellm/batches/batch_line_items.py b/litellm/batches/batch_line_items.py new file mode 100644 index 000000000000..a2ec6b5fd5fe --- /dev/null +++ b/litellm/batches/batch_line_items.py @@ -0,0 +1,170 @@ +"""Per-line records of a finished batch for the loggers listed in batch_line_item_callbacks.""" + +import asyncio +from collections.abc import Iterable, Mapping +from dataclasses import dataclass +from itertools import chain +from typing import Final, Literal, TypeAlias + +import litellm +from litellm._logging import verbose_logger +from litellm.batches.batch_utils import calculate_batch_cost_and_usage, get_file_content_as_dictionary +from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.litellm_logging import init_custom_logger_compatible_class +from litellm.types.utils import ( + ModelInfo, + StandardBatchLineItem, + StandardBatchLineItemMetadata, + StandardBatchLineItemsPayload, + StandardBatchLineItemUsage, +) + +BatchFileProvider: TypeAlias = Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "mistral"] + + +@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=_text(attribution, "user_api_key_hash"), + 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 _line_key(entry: Mapping[str, object]) -> str | None: + key: Final = entry.get("custom_id", entry.get("recordId")) + return key if isinstance(key, str) else None + + +async def build_batch_line_items( + source: BatchLineItemSource, input_content: bytes, result_contents: Iterable[bytes] +) -> StandardBatchLineItemsPayload: + """Pair every output and error line with the input line that shares its custom_id, priced one line at a time.""" + requests: Final = {_line_key(entry): entry for entry in get_file_content_as_dictionary(input_content)} + results: Final = chain.from_iterable(map(get_file_content_as_dictionary, result_contents)) + return StandardBatchLineItemsPayload( + batch_id=source.batch_id, + finished_at=source.finished_at, + lines=tuple([await _line_item(source, requests, entry) for entry in results]), + ) + + +async def _line_item( + source: BatchLineItemSource, + requests: Mapping[str | None, Mapping[str, object]], + entry: dict[str, object], +) -> StandardBatchLineItem: + priced: Final = await calculate_batch_cost_and_usage( + (entry,), source.cost_provider, model_name=source.model, model_info=source.model_info + ) + succeeded: Final = priced.successful_requests > 0 + custom_id: Final = _line_key(entry) + return StandardBatchLineItem( + batch_id=source.batch_id, + custom_id=custom_id, + status="success" if succeeded else "failure", + request=requests.get(custom_id) if custom_id is not None else None, + 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 without_request_and_response(payload: StandardBatchLineItemsPayload) -> StandardBatchLineItemsPayload: + return StandardBatchLineItemsPayload( + batch_id=payload["batch_id"], + finished_at=payload["finished_at"], + 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 the lines to every batch_line_item_callbacks logger. A logger that fails 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, the lines of batch %s were not stored", + callback, + 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 sink failing must not stop the other sinks or the batch cost poller + verbose_logger.exception( + "batch_line_item_callbacks: %r failed to store the %d lines of batch %s: %s", + callback, + len(payload["lines"]), + 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 a27b69da0765..21479bfa419f 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -1,9 +1,9 @@ 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 -from typing import Any, Final, Literal, TypeAlias, cast +from typing import Any, Final, Literal, cast import litellm from litellm._logging import verbose_logger @@ -20,8 +20,6 @@ from litellm.types.workload_identity import ANTHROPIC_WIF_KWARGS_KEYS from litellm.utils import token_counter -BatchProvider: TypeAlias = Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral"] - @dataclass(frozen=True, slots=True) class BatchCostUsageResult: @@ -36,14 +34,6 @@ class BatchCostUsageResult: completion_cost: float = 0.0 -@dataclass(frozen=True, slots=True) -class BatchResultFiles: - """Raw JSONL bytes of a completed batch's result files; None means not fetched.""" - - output: bytes | None - error: bytes | None - - _COMPLETED_BATCH_STATUSES: Final = frozenset({"completed", "complete"}) @@ -80,8 +70,8 @@ def batch_cost_is_final(batch: Batch) -> bool: async def calculate_batch_cost_and_usage( - file_content_dictionary: list[dict[str, object]], - custom_llm_provider: BatchProvider, + 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: @@ -108,7 +98,7 @@ async def calculate_batch_cost_and_usage( async def handle_completed_batch( batch: Batch, - custom_llm_provider: BatchProvider, + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "mistral"], model_name: str | None = None, litellm_params: dict[str, object] | None = None, model_info: ModelInfo | None = None, @@ -126,46 +116,29 @@ async def handle_completed_batch( threaded through so a deployment's configured rates win over the global cost map. """ - return ( - await handle_completed_batch_with_files( - batch, - custom_llm_provider, - model_name=model_name, - litellm_params=litellm_params, - model_info=model_info, - ) - )[0] - - -async def handle_completed_batch_with_files( - batch: Batch, - custom_llm_provider: BatchProvider, - model_name: str | None = None, - litellm_params: Mapping[str, object] | None = None, - model_info: ModelInfo | None = None, -) -> "tuple[BatchCostUsageResult, BatchResultFiles]": - """_handle_completed_batch plus the raw output/error bytes it fetched, so - downstream line-item logging reuses them instead of refetching.""" + # A completed batch whose request lines all failed has no output file - the + # results are written to a separate error_file_id and output_file_id is None. + # There is nothing to price or measure, so report an empty result set instead + # of calling _fetch_batch_output_file_content, which raises on a missing + # output file. Without this guard the logging worker crashes on every + # aretrieve_batch poll and the completed batch's zero-cost accounting is lost. + # The generic retrieval helper keeps raising for callers that explicitly ask + # for a missing output file. if batch.output_file_id is None: - only_error_file_content: Final = await fetch_batch_error_file_content( - batch, custom_llm_provider=custom_llm_provider, litellm_params=litellm_params - ) - return ( - BatchCostUsageResult( - cost=0.0, - usage=Usage(prompt_tokens=0, completion_tokens=0, total_tokens=0), - models=[], - successful_requests=0, - failed_requests=count_error_file_failed_requests_from_content(only_error_file_content), + return BatchCostUsageResult( + cost=0.0, + usage=Usage(prompt_tokens=0, completion_tokens=0, total_tokens=0), + models=[], + successful_requests=0, + failed_requests=await count_error_file_failed_requests( + batch, custom_llm_provider=custom_llm_provider, litellm_params=litellm_params ), - BatchResultFiles(output=None, error=only_error_file_content), ) file_content = await _fetch_batch_output_file_content(batch, custom_llm_provider, litellm_params=litellm_params) - error_file_content: Final = await fetch_batch_error_file_content( + error_file_failed_requests: Final = await count_error_file_failed_requests( batch, custom_llm_provider=custom_llm_provider, litellm_params=litellm_params ) - error_file_failed_requests: Final = count_error_file_failed_requests_from_content(error_file_content) output_file_result: Final = ( calculate_vertex_ai_batch_cost_and_usage( @@ -182,14 +155,11 @@ async def handle_completed_batch_with_files( ) ) - result: Final = ( - output_file_result - if not error_file_failed_requests - else dataclasses_replace( - output_file_result, failed_requests=output_file_result.failed_requests + error_file_failed_requests - ) + if not error_file_failed_requests: + return output_file_result + return dataclasses_replace( + output_file_result, failed_requests=output_file_result.failed_requests + error_file_failed_requests ) - return result, BatchResultFiles(output=file_content, error=error_file_content) _handle_completed_batch = handle_completed_batch @@ -217,7 +187,7 @@ class _BatchOutputLineStats: def _classify_output_line_stats( entries: Iterable[dict[str, object]], - custom_llm_provider: BatchProvider, + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral"], model_name: str | None, model_info: ModelInfo | None, ) -> Iterator[_BatchOutputLineStats | _LineOutcome]: @@ -236,7 +206,7 @@ def _classify_output_line_stats( def _safe_output_line_stats( entry: Mapping[str, object], - custom_llm_provider: BatchProvider, + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral"], model_name: str | None, model_info: ModelInfo | None, ) -> _BatchOutputLineStats | None: @@ -258,7 +228,7 @@ def _safe_output_line_stats( def _compute_output_line_stats( entry: Mapping[str, object], - custom_llm_provider: BatchProvider, + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral"], model_name: str | None, model_info: ModelInfo | None, ) -> _BatchOutputLineStats: @@ -300,7 +270,7 @@ def _ocr_usage_info_from_response_body(response_body: Mapping[str, object]) -> O def _output_line_cost( response_body: Mapping[str, object], usage: Usage, - custom_llm_provider: BatchProvider, + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral"], model_name: str | None, response_model: str | None, model_info: ModelInfo | None, @@ -329,7 +299,7 @@ def _output_line_cost( def _aggregate_batch_cost_usage_models( entries: Iterable[dict[str, object]], - custom_llm_provider: BatchProvider, + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock", "mistral"], model_name: str | None = None, model_info: ModelInfo | None = None, ) -> BatchCostUsageResult: @@ -466,9 +436,9 @@ 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: BatchProvider = "openai", + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "mistral"] = "openai", litellm_params: Mapping[str, object] | None = None, ) -> bytes: """ @@ -497,9 +467,12 @@ 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: BatchProvider = "openai", + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "mistral"] = "openai", litellm_params: Mapping[str, object] | None = None, ) -> bytes: """ @@ -519,30 +492,9 @@ async def _fetch_batch_output_file_content( ) -async def fetch_batch_error_file_content( - batch: Batch, - custom_llm_provider: BatchProvider, - litellm_params: Mapping[str, object] | None, -) -> bytes | None: - """Fetch the batch's separate error file bytes; None when it has none or the fetch fails.""" - if batch.error_file_id is None: - return None - try: - return await _fetch_batch_managed_file_content( - batch.error_file_id, custom_llm_provider=custom_llm_provider, litellm_params=litellm_params - ) - except Exception as e: # noqa: BLE001 # a failed/missing error file must not abort cost tracking for the batch - verbose_logger.debug("Failed to fetch batch error file %s: %s", batch.error_file_id, e) - return None - - -def count_error_file_failed_requests_from_content(content: bytes | None) -> int: - return 0 if content is None else sum(1 for _ in iter_batch_input_lines(content)) - - async def count_error_file_failed_requests( batch: Batch, - custom_llm_provider: BatchProvider, + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "mistral"], litellm_params: Mapping[str, object] | None, ) -> int: """Count failed requests reported only in the batch's separate error file. @@ -552,11 +504,16 @@ async def count_error_file_failed_requests( ``error_file_id`` - they never appear in the output file at all, so counting failures from the output file alone silently undercounts them. """ - return count_error_file_failed_requests_from_content( - await fetch_batch_error_file_content( - batch, custom_llm_provider=custom_llm_provider, litellm_params=litellm_params + if batch.error_file_id is None: + return 0 + try: + error_file_content = await _fetch_batch_managed_file_content( + batch.error_file_id, custom_llm_provider=custom_llm_provider, litellm_params=litellm_params ) - ) + except Exception as e: # noqa: BLE001 # a failed/missing error file must not abort cost tracking for the batch + verbose_logger.debug("Failed to fetch batch error file %s: %s", batch.error_file_id, e) + return 0 + return sum(1 for _ in iter_batch_input_lines(error_file_content)) def extract_file_access_credentials( diff --git a/litellm/integrations/SlackAlerting/slack_alerting.py b/litellm/integrations/SlackAlerting/slack_alerting.py index 43c6a7a77a74..94a3564e5cb9 100644 --- a/litellm/integrations/SlackAlerting/slack_alerting.py +++ b/litellm/integrations/SlackAlerting/slack_alerting.py @@ -29,7 +29,6 @@ from litellm.integrations.SlackAlerting.hanging_request_check import ( AlertingHangingRequestCheck, ) -from litellm.litellm_core_utils.core_helpers import is_batch_line_item_event from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.litellm_core_utils.exception_mapping_utils import ( add_key_name_and_team_to_alert, @@ -1689,8 +1688,6 @@ async def periodic_flush(self): async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): """Log deployment latency""" - if is_batch_line_item_event(kwargs): - return try: if "daily_reports" in self.alert_types: litellm_params: Final = kwargs.get("litellm_params", {}) or {} @@ -1726,8 +1723,6 @@ async def async_log_success_event(self, kwargs, response_obj, start_time, end_ti async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): """Log failure + deployment latency""" - if is_batch_line_item_event(kwargs): - return _litellm_params: Final = kwargs.get("litellm_params", {}) _model_info: Final = _litellm_params.get("model_info", {}) or {} model_id: Final = _model_info.get("id", "") diff --git a/litellm/integrations/batch_line_item_objects.py b/litellm/integrations/batch_line_item_objects.py new file mode 100644 index 000000000000..f9312bdb60d5 --- /dev/null +++ b/litellm/integrations/batch_line_item_objects.py @@ -0,0 +1,17 @@ +"""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 "\n".join(json.dumps(line, default=str, ensure_ascii=False) for line in payload["lines"]) + + +def batch_line_items_object_name(payload: StandardBatchLineItemsPayload) -> str: + """batch_line_items//.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('/', '_')}.jsonl" diff --git a/litellm/integrations/clickhouse/clickhouse_spend_logger.py b/litellm/integrations/clickhouse/clickhouse_spend_logger.py index cb2f4f0bcf8b..585ea6c2afcb 100644 --- a/litellm/integrations/clickhouse/clickhouse_spend_logger.py +++ b/litellm/integrations/clickhouse/clickhouse_spend_logger.py @@ -19,7 +19,6 @@ from litellm.integrations.clickhouse.clickhouse_batch_logger import ClickHouseBatchLogger from litellm.integrations.clickhouse.context import is_lens_analysis from litellm.integrations.clickhouse.schema import SPEND_LOGS_TABLE -from litellm.litellm_core_utils.core_helpers import is_batch_line_item_event from litellm.litellm_core_utils.llm_response_utils.get_headers import get_provider_request_id from litellm.litellm_core_utils.sensitive_data_masker import redact_credentials_in_payload from litellm.tracing.types import SpendLogPayload, SpendLogRecord @@ -206,13 +205,9 @@ class ClickHouseSpendLogger(ClickHouseBatchLogger): table = SPEND_LOGS_TABLE async def async_log_success_event(self, kwargs, response_obj, start_time, end_time) -> None: - if is_batch_line_item_event(kwargs): - return self._log(kwargs) async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time) -> None: - if is_batch_line_item_event(kwargs): - return self._log(kwargs) def _log(self, kwargs: Mapping[str, Any]) -> None: diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py index 5ca9840185b1..4c6410b7e4be 100644 --- a/litellm/integrations/custom_logger.py +++ b/litellm/integrations/custom_logger.py @@ -22,6 +22,7 @@ ModelResponse, ModelResponseStream, StandardAuditLogPayload, + StandardBatchLineItemsPayload, StandardCallbackDynamicParams, StandardLoggingPayload, ) @@ -195,6 +196,9 @@ async def async_log_failure_event(self, kwargs, response_obj, start_time, end_ti 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 finished batch with every line, 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/datadog/datadog_cost_management.py b/litellm/integrations/datadog/datadog_cost_management.py index 1584cfca1c43..538dd95abdd2 100644 --- a/litellm/integrations/datadog/datadog_cost_management.py +++ b/litellm/integrations/datadog/datadog_cost_management.py @@ -14,7 +14,6 @@ get_datadog_service, normalize_datadog_tag_value, ) -from litellm.litellm_core_utils.core_helpers import is_batch_line_item_event from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, @@ -73,8 +72,6 @@ def __init__(self, cost_tag_keys: list[str] | None = None, **kwargs): super().__init__(**kwargs) async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): - if is_batch_line_item_event(kwargs): - return try: standard_logging_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object", None) diff --git a/litellm/integrations/datadog/datadog_metrics.py b/litellm/integrations/datadog/datadog_metrics.py index eae6bda0dc35..5dda336dc944 100644 --- a/litellm/integrations/datadog/datadog_metrics.py +++ b/litellm/integrations/datadog/datadog_metrics.py @@ -14,7 +14,6 @@ get_datadog_service, normalize_datadog_tag_value, ) -from litellm.litellm_core_utils.core_helpers import is_batch_line_item_event from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, @@ -171,8 +170,6 @@ def _add_metrics_from_log( self.log_queue.append(series_count) async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): - if is_batch_line_item_event(kwargs): - return try: standard_logging_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object", None) @@ -188,8 +185,6 @@ async def async_log_success_event(self, kwargs, response_obj, start_time, end_ti verbose_logger.exception("Datadog Metrics: Error in async_log_success_event: %s", e) async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): - if is_batch_line_item_event(kwargs): - return try: standard_logging_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object", None) diff --git a/litellm/integrations/gcs_bucket/gcs_bucket.py b/litellm/integrations/gcs_bucket/gcs_bucket.py index 24a5e940e080..6cad51042792 100644 --- a/litellm/integrations/gcs_bucket/gcs_bucket.py +++ b/litellm/integrations/gcs_bucket/gcs_bucket.py @@ -10,13 +10,14 @@ 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 @@ async def async_log_failure_event(self, kwargs, response_obj, start_time, end_ti 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 the batch's lines as one 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=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/lago.py b/litellm/integrations/lago.py index 5b858f5901da..594427b1e0a0 100644 --- a/litellm/integrations/lago.py +++ b/litellm/integrations/lago.py @@ -11,7 +11,6 @@ from litellm._logging import verbose_logger from litellm._uuid import uuid from litellm.integrations.custom_logger import CustomLogger -from litellm.litellm_core_utils.core_helpers import is_batch_line_item_event from litellm.llms.custom_httpx.http_handler import ( HTTPHandler, get_async_httpx_client, @@ -122,8 +121,6 @@ def _common_logic(self, kwargs: dict, response_obj) -> dict: return returned_val def log_success_event(self, kwargs, response_obj, start_time, end_time): - if is_batch_line_item_event(kwargs): - return _url = os.getenv("LAGO_API_BASE") assert _url is not None and isinstance(_url, str), ( f"LAGO_API_BASE missing or not set correctly. LAGO_API_BASE={_url}" @@ -156,8 +153,6 @@ def log_success_event(self, kwargs, response_obj, start_time, end_time): raise e async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): - if is_batch_line_item_event(kwargs): - return try: verbose_logger.debug("ENTERS LAGO CALLBACK") _url = os.getenv("LAGO_API_BASE") diff --git a/litellm/integrations/newrelic/newrelic_metrics.py b/litellm/integrations/newrelic/newrelic_metrics.py index c8b2c6be3ba8..a2cedeba0f52 100644 --- a/litellm/integrations/newrelic/newrelic_metrics.py +++ b/litellm/integrations/newrelic/newrelic_metrics.py @@ -34,7 +34,6 @@ from litellm._logging import verbose_logger from litellm.integrations.custom_batch_logger import CustomBatchLogger -from litellm.litellm_core_utils.core_helpers import is_batch_line_item_event from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, @@ -302,16 +301,12 @@ async def periodic_flush(self) -> None: await self._final_drain() async def async_log_success_event(self, kwargs, response_obj, start_time, end_time) -> None: - if is_batch_line_item_event(kwargs): - return try: await self._log_async_event(standard_logging_object=kwargs.get("standard_logging_object", None)) except Exception as e: # noqa: BLE001 # logging must never break the request path verbose_logger.exception("New Relic Metrics Layer Error - %s\n%s", e, traceback.format_exc()) async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time) -> None: - if is_batch_line_item_event(kwargs): - return try: await self._log_async_event(standard_logging_object=kwargs.get("standard_logging_object", None)) except Exception as e: # noqa: BLE001 # logging must never break the request path diff --git a/litellm/integrations/openmeter.py b/litellm/integrations/openmeter.py index d0f1dab988ce..db2fe386dec0 100644 --- a/litellm/integrations/openmeter.py +++ b/litellm/integrations/openmeter.py @@ -9,7 +9,6 @@ import litellm from litellm.integrations.custom_logger import CustomLogger -from litellm.litellm_core_utils.core_helpers import is_batch_line_item_event from litellm.llms.custom_httpx.http_handler import ( HTTPHandler, get_async_httpx_client, @@ -98,8 +97,6 @@ def _common_logic(self, kwargs: dict, response_obj): } def log_success_event(self, kwargs, response_obj, start_time, end_time): - if is_batch_line_item_event(kwargs): - return _url = os.getenv("OPENMETER_API_ENDPOINT", "https://openmeter.cloud") if _url.endswith("/"): _url += "api/v1/events" @@ -126,8 +123,6 @@ def log_success_event(self, kwargs, response_obj, start_time, end_time): raise e async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): - if is_batch_line_item_event(kwargs): - return _url = os.getenv("OPENMETER_API_ENDPOINT", "https://openmeter.cloud") if _url.endswith("/"): _url += "api/v1/events" diff --git a/litellm/integrations/opentelemetry.py b/litellm/integrations/opentelemetry.py index 0af657442777..90383ddb60dc 100644 --- a/litellm/integrations/opentelemetry.py +++ b/litellm/integrations/opentelemetry.py @@ -33,7 +33,6 @@ from litellm.integrations.otel.model.semconv import LiteLLM, Metric from litellm.integrations.otel.plumbing.otlp_tls import resolve_otlp_http_tls from litellm.integrations.otel.routing import routing_decision_attributes -from litellm.litellm_core_utils.core_helpers import is_batch_line_item_event from litellm.litellm_core_utils.internal_call_metadata import is_unbilled_non_inference_call_from_params from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.secret_redaction import redact_string @@ -1670,8 +1669,6 @@ def _metric_attribute_allowed(self, key: str) -> bool: return True def _record_metrics(self, kwargs, response_obj, start_time, end_time): - if is_batch_line_item_event(kwargs): - return duration_s: Final = (end_time - start_time).total_seconds() params: Final = kwargs.get("litellm_params") or {} provider: Final = _provider_label(params.get("custom_llm_provider")) diff --git a/litellm/integrations/otel/plumbing/metrics.py b/litellm/integrations/otel/plumbing/metrics.py index c3294c7925d6..7ec3711ee656 100644 --- a/litellm/integrations/otel/plumbing/metrics.py +++ b/litellm/integrations/otel/plumbing/metrics.py @@ -33,7 +33,6 @@ resolve_provider, ) from litellm.integrations.otel.model.utils import to_seconds -from litellm.litellm_core_utils.core_helpers import is_batch_line_item_event from litellm.litellm_core_utils.internal_call_metadata import is_unbilled_non_inference_call_from_params from litellm.litellm_core_utils.safe_json_dumps import safe_dumps @@ -222,8 +221,6 @@ def record( start_time: datetime, end_time: datetime, ) -> None: - if is_batch_line_item_event(kwargs): - return common_attrs: Final = self._filter_attributes(self._bounded_attributes(kwargs)) duration_s: Final = (end_time - start_time).total_seconds() usage_is_replayed: Final = is_unbilled_non_inference_call_from_params( @@ -269,8 +266,6 @@ def record_failure( ``gen_ai.token.type``, so an operator's include/exclude list cannot strip the discriminator and silently merge failures back into the success series. """ - if is_batch_line_item_event(kwargs): - return attributes: Final = { **self._filter_attributes(self._bounded_attributes(kwargs)), Error.TYPE: resolve_error_type(kwargs), diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index b90e571a5c9c..708a75aad996 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -45,7 +45,6 @@ from litellm.litellm_core_utils.core_helpers import ( get_litellm_metadata_from_kwargs, get_metadata_variable_name_from_kwargs, - is_batch_line_item_event, ) from litellm.litellm_core_utils.service_tier_utils import ( get_service_tier_from_standard_logging_payload, @@ -1507,8 +1506,6 @@ def _inc_labeled_counter( self._track_end_user_metric_series(counter, metric_name, _labels) async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): - if is_batch_line_item_event(kwargs): - return # Define prometheus client verbose_logger.debug( "prometheus Logging - Enters success logging function (kwargs keys: %s)", @@ -2605,8 +2602,6 @@ def _set_latency_metrics( ) async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): - if is_batch_line_item_event(kwargs): - return verbose_logger.debug( "prometheus Logging - Enters failure logging function (kwargs keys: %s)", list(kwargs.keys()) if isinstance(kwargs, dict) else type(kwargs).__name__, diff --git a/litellm/integrations/s3_v2.py b/litellm/integrations/s3_v2.py index 47d9b1c7b644..ec5c11b1826b 100644 --- a/litellm/integrations/s3_v2.py +++ b/litellm/integrations/s3_v2.py @@ -32,6 +32,7 @@ S3_PARTITION_GRANULARITY_ENV_VAR, ) 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, @@ -57,7 +58,7 @@ 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 @@ -511,6 +512,28 @@ async def async_log_audit_log_event(self, audit_log: StandardAuditLogPayload) -> 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 the batch's lines as one 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 + delivered: Final = await self.async_upload_data_to_s3( + s3BatchLoggingElement( + payload={}, + body=batch_line_items_ndjson(payload), + content_type="application/x-ndjson", + s3_object_key=s3_object_key, + s3_object_download_filename=object_name.rsplit("/", 1)[-1], + ) + ) + if not delivered: + verbose_logger.error( + "s3 logging: the %d line items of batch %s were not stored at %s", + len(payload["lines"]), + 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) diff --git a/litellm/integrations/shadow_eval_logger.py b/litellm/integrations/shadow_eval_logger.py index d4fbc5d97e8f..962c7707409b 100644 --- a/litellm/integrations/shadow_eval_logger.py +++ b/litellm/integrations/shadow_eval_logger.py @@ -29,11 +29,7 @@ from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.websearch_interception.tools import is_web_search_tool_responses -from litellm.litellm_core_utils.core_helpers import ( - get_litellm_metadata_from_kwargs, - independent_snapshot, - is_batch_line_item_event, -) +from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kwargs, independent_snapshot from litellm.litellm_core_utils.internal_call_metadata import sanitized_forwardable_call_metadata from litellm.litellm_core_utils.llm_judge import ( default_router_provider, @@ -949,8 +945,6 @@ async def async_log_success_event( *, guardrail_snapshot: GuardrailRequestSnapshot | None = None, ) -> None: - if is_batch_line_item_event(kwargs): - return try: payload: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object") # pyright: ignore[reportAssignmentType] # untyped callback kwargs if payload is None: diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index 8ffcc4939dcc..e857461785c3 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -793,18 +793,6 @@ def set_response_cost_in_hidden_params(response: _CarriesHiddenParams, cost: flo hidden_params["additional_headers"] = merged -BATCH_PARENT_ID_KEY: Final = "batch_parent_id" - - -def is_batch_line_item_event(kwargs: Any) -> bool: # noqa: ANN401 # every hook passes its own unannotated callback kwargs - """Whether a callback event is one JSONL line of a completed batch rather than a request""" - match kwargs: - case {"litellm_params": {"batch_parent_id": str() as batch_parent_id}}: - return bool(batch_parent_id) - case _: - return False - - _HIDDEN_PARAMS_ADAPTER: Final = TypeAdapter(Mapping[str, object]) _PROVIDER_HEADERS_ADAPTER: Final = TypeAdapter(Mapping[str, str]) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 37c51836b427..634a1e3f24d3 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -34,12 +34,7 @@ verbose_logger, ) from litellm._uuid import uuid -from litellm.batches.batch_utils import ( - BatchResultFiles, - batch_cost_is_final, - handle_completed_batch, - handle_completed_batch_with_files, -) +from litellm.batches.batch_utils import batch_cost_is_final, handle_completed_batch from litellm.caching.caching import DualCache from litellm.caching.caching_handler import LLMCachingHandler from litellm.caching.redis_batch import flush_post_call_redis_batches @@ -77,7 +72,6 @@ is_classifier_call, ) from litellm.litellm_core_utils.core_helpers import ( - BATCH_PARENT_ID_KEY, get_provider_response_headers_from_hidden_params, is_expected_client_error, proxy_stamped_used_client_oauth_token, @@ -694,7 +688,6 @@ def __init__( self.streaming_chunks: list[object] = [] self.sync_streaming_chunks: list[object] = [] self.log_raw_request_response = log_raw_request_response - self._litellm_internal_model_credentials: Mapping[str, object] | None = None self.raw_request_only = raw_request_only # Initialize dynamic callbacks @@ -967,17 +960,10 @@ def update_environment_variables( if model is not None: self.model = model self.user = user - marked_batch_parent_id: Final[object] = self.litellm_params.get(BATCH_PARENT_ID_KEY) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # Logging.litellm_params is untyped upstream self.litellm_params = { **self.litellm_params, **scrub_sensitive_keys_in_metadata(litellm_params), } - # Every spend, budget and metering hook skips an event carrying batch_parent_id, so only - # mark_batch_line_item may set it; the incoming params can hold caller request fields - if marked_batch_parent_id is None: - self.litellm_params.pop(BATCH_PARENT_ID_KEY, None) # pyright: ignore[reportUnknownMemberType] # Logging.litellm_params is untyped upstream - else: - self.litellm_params[BATCH_PARENT_ID_KEY] = marked_batch_parent_id # pyright: ignore[reportUnknownMemberType] # Logging.litellm_params is untyped upstream self.litellm_request_debug = litellm_params.get("litellm_request_debug", False) self.logger_fn = litellm_params.get("logger_fn", None) if is_debugging_on() or self.litellm_request_debug: @@ -1000,10 +986,6 @@ def update_environment_variables( **additional_params, } ) - # Provider params / additional kwargs are caller-influenced; they must not - # overwrite the internal litellm_params, or a request body carrying - # `litellm_params` could forge logging markers (e.g. batch_parent_id). - self.model_call_details["litellm_params"] = self.litellm_params ## check if stream options is set ## - used by CustomStreamWrapper for easy instrumentation if "stream_options" in additional_params: @@ -1015,10 +997,6 @@ def update_environment_variables( if "custom_llm_provider" in self.model_call_details: self.custom_llm_provider = self.model_call_details["custom_llm_provider"] - def mark_batch_line_item(self, batch_id: str) -> None: - """Mark this logging object as one JSONL line of the completed batch ``batch_id``""" - self.litellm_params[BATCH_PARENT_ID_KEY] = batch_id # pyright: ignore[reportUnknownMemberType] # Logging.litellm_params is untyped upstream - def update_from_kwargs( self, kwargs: dict, @@ -3324,14 +3302,8 @@ async def _async_success_handler_body( batch_models = kwargs.get("batch_models", None) batch_successful_requests: Final = kwargs.get("batch_successful_requests", None) batch_failed_requests: Final = kwargs.get("batch_failed_requests", None) - batch_output_file_content: Final = kwargs.get("batch_output_file_content", None) - batch_error_file_content: Final = kwargs.get("batch_error_file_content", None) has_explicit_batch_data: Final = all(x is not None for x in (batch_cost, batch_usage, batch_models)) should_compute_batch_data: Final = not has_explicit_batch_data and batch_cost_is_final(result) - result_files: BatchResultFiles | None = None # rebind-ok: one batch-data branch below supplies the bytes - batch_litellm_params: Final = cast( # cast-ok: shared attribute is an untyped dict - "dict[str, object] | None", self.litellm_params - ) if has_explicit_batch_data: set_hidden_param(result, "response_cost", batch_cost) set_hidden_param(result, "batch_models", batch_models) @@ -3351,30 +3323,15 @@ async def _async_success_handler_body( total_cost=batch_cost, cost_for_built_in_tools_cost_usd_dollar=0.0, ) - if batch_output_file_content is not None or batch_error_file_content is not None: - result_files = BatchResultFiles( # rebind-ok: the caller supplied the files - output=batch_output_file_content if isinstance(batch_output_file_content, bytes) else None, - error=batch_error_file_content if isinstance(batch_error_file_content, bytes) else None, - ) elif should_compute_batch_data: - if litellm.store_batch_line_items_in_callbacks: - batch_result, fetched_result_files = await handle_completed_batch_with_files( - batch=result, - custom_llm_provider=self.custom_llm_provider, - model_name=self.get_deployment_model_for_cost(), - litellm_params=batch_litellm_params, - model_info=self.get_router_deployment_model_info(), - ) - result_files = fetched_result_files # rebind-ok: the files this branch fetched - else: - batch_result = await handle_completed_batch( - batch=result, - custom_llm_provider=self.custom_llm_provider, - model_name=self.get_deployment_model_for_cost(), - litellm_params=batch_litellm_params, - model_info=self.get_router_deployment_model_info(), - ) + batch_result: Final = await handle_completed_batch( + batch=result, + custom_llm_provider=self.custom_llm_provider, + model_name=self.get_deployment_model_for_cost(), + litellm_params=self.litellm_params, + model_info=self.get_router_deployment_model_info(), + ) set_hidden_param(result, "response_cost", batch_result.cost) set_hidden_param(result, "batch_models", batch_result.models) @@ -3388,25 +3345,6 @@ async def _async_success_handler_body( cost_for_built_in_tools_cost_usd_dollar=0.0, ) - if litellm.store_batch_line_items_in_callbacks and (has_explicit_batch_data or should_compute_batch_data): - from litellm.batches.batch_line_item_logging import spawn_batch_line_items - - try: - spawn_batch_line_items( - batch=result, - custom_llm_provider=self.custom_llm_provider, - parent=self, - model_name=self.get_deployment_model_for_cost(), - litellm_params=batch_litellm_params, - model_info=self.get_router_deployment_model_info(), - result_files=result_files, - ) - except Exception: # noqa: BLE001 # starting the line fan-out must never break the aggregate aretrieve_batch path - verbose_logger.exception( - "batch line item logging failed for batch_id=%s; aggregate logging unaffected", - result.id, - ) - self.truncated_messages_for_logging = await truncate_base64_in_messages_async( StandardLoggingPayloadSetup.append_system_prompt_messages( kwargs=self.model_call_details, messages=self.model_call_details.get("messages") @@ -4363,10 +4301,19 @@ def _handle_anthropic_messages_response_logging(self, result: object) -> ModelRe verbose_logger.warning("LiteLLM: the anthropic_messages stream assembled no response, logging an empty one") return litellm.ModelResponse(model=self.model) else: - from litellm.llms.anthropic.chat.transformation import anthropic_message_to_model_response + from litellm.types.llms.anthropic import AnthropicResponse + + pydantic_result: Final = AnthropicResponse.model_validate(result) + import httpx - result = anthropic_message_to_model_response( - cast(Mapping[str, object], result), # cast-ok: handler result is typed Any upstream + result = litellm.AnthropicConfig().transform_parsed_response( + completion_response=pydantic_result.model_dump(), + raw_response=httpx.Response( + status_code=200, + headers={}, + ), + model_response=litellm.ModelResponse(id=provider_response_id), + json_mode=None, speed=self.optional_params.get("speed") if self.optional_params else None, ) return result @@ -5336,6 +5283,9 @@ def _init_custom_logger_compatible_class( return None +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``. @@ -6615,12 +6565,8 @@ def _extract_response_obj_and_hidden_params( response_obj = {} if original_exception is not None and hidden_params is None: - exception_hidden_params: Final[Mapping[str, object] | None] = getattr( - original_exception, "_hidden_params", None - ) - if isinstance(exception_hidden_params, dict) and exception_hidden_params: - hidden_params = dict(exception_hidden_params) # rebind-ok: the exception carries the hidden params - elif (response_headers := _get_response_headers(original_exception)) is not None: + response_headers: Final = _get_response_headers(original_exception) + if response_headers is not None: hidden_params = dict( StandardLoggingHiddenParams( additional_headers=StandardLoggingPayloadSetup.get_additional_headers(dict(response_headers)), diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index 3943cb823fd0..54a8780ca5b1 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -52,7 +52,6 @@ AnthropicMessagesToolChoice, AnthropicOutputSchema, AnthropicOutputTokensDetails, - AnthropicResponse, AnthropicSystemMessageContent, AnthropicThinkingParam, AnthropicWebSearchTool, @@ -2882,17 +2881,3 @@ def _valid_user_id(user_id: str) -> bool: return False return True - - -def anthropic_message_to_model_response(result: Mapping[str, object], speed: str | None) -> ModelResponse: - pydantic_result: Final = AnthropicResponse.model_validate(result) - return AnthropicConfig().transform_parsed_response( # pyright: ignore[reportUnknownMemberType] # completion_response is a bare dict upstream - completion_response=pydantic_result.model_dump(), - raw_response=httpx.Response( - status_code=200, - headers={}, - ), - model_response=ModelResponse(id=pydantic_result.id or None), - json_mode=None, - speed=speed, - ) diff --git a/litellm/llms/base_llm/batches/transformation.py b/litellm/llms/base_llm/batches/transformation.py index 904eae2f4f27..34c622d4cf6b 100644 --- a/litellm/llms/base_llm/batches/transformation.py +++ b/litellm/llms/base_llm/batches/transformation.py @@ -1,6 +1,5 @@ import types from abc import ABC, abstractmethod -from collections.abc import Mapping from typing import TYPE_CHECKING, Any import httpx @@ -10,7 +9,7 @@ AllMessageValues, CreateBatchRequest, ) -from litellm.types.utils import EmbeddingResponse, LiteLLMBatch, LlmProviders, ModelResponse +from litellm.types.utils import LiteLLMBatch, LlmProviders if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj @@ -40,14 +39,6 @@ def __init__(self): def custom_llm_provider(self) -> LlmProviders: """Return the LLM provider type for this configuration.""" - def transform_batch_output_line( - self, model_output: Mapping[str, object], model: str - ) -> ModelResponse | EmbeddingResponse | None: - """Reconstruct one provider batch output line into a litellm response, or - None when the line is OpenAI-shaped and the caller can use the generic - reconstruction.""" - return None - @classmethod def get_config(cls): """Get configuration dictionary for this class.""" diff --git a/litellm/llms/bedrock/batches/transformation.py b/litellm/llms/bedrock/batches/transformation.py index 136cf8f1b99b..3ad8049ec260 100644 --- a/litellm/llms/bedrock/batches/transformation.py +++ b/litellm/llms/bedrock/batches/transformation.py @@ -27,7 +27,7 @@ AllMessageValues, CreateBatchRequest, ) -from litellm.types.utils import EmbeddingResponse, LiteLLMBatch, LlmProviders, ModelResponse, Usage +from litellm.types.utils import LiteLLMBatch, LlmProviders, Usage from ..base_aws_llm import BaseAWSLLM from ..common_utils import ( @@ -96,47 +96,6 @@ def titan_embedding_usage_from_batch_output(model_output: Mapping[str, object]) ) -def bedrock_batch_line_to_response( - model_output: Mapping[str, object], model: str -) -> ModelResponse | EmbeddingResponse | None: - """Reconstruct a Bedrock batch output line (the ``modelOutput`` object) into - the litellm response type its shape implies, or None when the shape is - unrecognized.""" - if "embedding" in model_output: - embedding: Final = model_output.get("embedding") - return EmbeddingResponse( - model=model, - data=[ - { - "object": "embedding", - "index": 0, - "embedding": embedding if isinstance(embedding, list) else [], - } - ], - usage=titan_embedding_usage_from_batch_output(model_output), - ) - if "output" in model_output: - from ..chat.converse_transformation import AmazonConverseConfig - - return AmazonConverseConfig()._transform_response( # pyright: ignore[reportPrivateUsage, reportUnknownMemberType] # same reconstruction the converse chat path performs on the live response - model=model, - response=Response(200, json=dict(model_output)), - model_response=ModelResponse(), - stream=False, - logging_obj=None, - optional_params={}, - api_key=None, - data="", - messages=[], - encoding=None, - ) - if "content" in model_output: - from litellm.llms.anthropic.chat.transformation import anthropic_message_to_model_response - - return anthropic_message_to_model_response(model_output, None) - return None - - class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): """ Config for Bedrock Batches - handles batch job creation and management for Bedrock @@ -150,11 +109,6 @@ def __init__(self): def custom_llm_provider(self) -> LlmProviders: return LlmProviders.BEDROCK - def transform_batch_output_line( - self, model_output: Mapping[str, object], model: str - ) -> ModelResponse | EmbeddingResponse | None: - return bedrock_batch_line_to_response(model_output, model) - @classmethod def _get_bare_model_name_from_s3_key(cls, object_key: str) -> str | None: if not object_key.startswith(BEDROCK_MANAGED_S3_BATCH_PREFIX): diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index c88c73990a93..2c75a608d0d5 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2786,10 +2786,6 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): None, description="If True, stores request messages and responses in spend logs. Default is False.", ) - store_batch_line_items_in_callbacks: bool | None = Field( - None, - description="If True, a completed batch logged via aretrieve_batch also emits one callback event per JSONL line item (request paired with its response or error). The aggregate batch callback is unchanged. Default is False.", - ) spend_logs_metadata_fields: SpendLogsMetadataFields | None = Field( None, description="Which keys of LiteLLM_SpendLogs.metadata are written to the database. Set exactly one of 'include' (write only these keys) or 'exclude' (drop these keys). 'status' and 'cold_storage_object_key' are always written. Daily spend tables, budgets and logging callbacks still see every key. Unset writes every key", diff --git a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py index 2024d5a1cd5d..d174d9fecb12 100644 --- a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py +++ b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py @@ -15,7 +15,6 @@ from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger -from litellm.litellm_core_utils.core_helpers import is_batch_line_item_event from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.proxy_rate_limit_error import ( ProxyRateLimitError, @@ -701,8 +700,6 @@ async def async_log_success_event(self, kwargs, response_obj, start_time, end_ti - model_saturation_check: Model-wide token tracking - priority_model: Priority-specific token tracking """ - if is_batch_line_item_event(kwargs): - return from litellm.litellm_core_utils.core_helpers import ( get_parent_otel_span_from_kwargs, ) diff --git a/litellm/proxy/hooks/max_budget_per_session_limiter.py b/litellm/proxy/hooks/max_budget_per_session_limiter.py index e74948590cf1..cf3025d8c7bb 100644 --- a/litellm/proxy/hooks/max_budget_per_session_limiter.py +++ b/litellm/proxy/hooks/max_budget_per_session_limiter.py @@ -24,7 +24,6 @@ from litellm.caching.redis_cache import log_redis_failure from litellm.exceptions import RateLimitType from litellm.integrations.custom_logger import CustomLogger -from litellm.litellm_core_utils.core_helpers import is_batch_line_item_event from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit @@ -135,8 +134,6 @@ async def async_log_success_event(self, kwargs, response_obj, start_time, end_ti """ After a successful LLM call, increment the session spend by the response cost. """ - if is_batch_line_item_event(kwargs): # pyright: ignore[reportUnknownArgumentType] # hook kwargs arrive untyped from the logging dispatcher - return try: litellm_params: Final = kwargs.get("litellm_params") or {} metadata: Final = litellm_params.get("metadata") or {} diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index 6ca9c3c5a005..eaa185917e5e 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/litellm/proxy/hooks/model_max_budget_limiter.py @@ -13,7 +13,6 @@ from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import Span -from litellm.litellm_core_utils.core_helpers import is_batch_line_item_event from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.llms.bedrock.common_utils import get_bedrock_base_model from litellm.proxy._types import Litellm_EntityType, UserAPIKeyAuth @@ -537,8 +536,6 @@ async def async_log_success_event(self, kwargs, response_obj, start_time, end_ti Example: key=$LITELLM_API_KEY, model=gpt-4o, max_budget=100, time_period=1d """ - if is_batch_line_item_event(kwargs): - return verbose_proxy_logger.debug("in RouterBudgetLimiting.async_log_success_event") standard_logging_payload: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object", None) if standard_logging_payload is None: diff --git a/litellm/proxy/hooks/parallel_request_limiter.py b/litellm/proxy/hooks/parallel_request_limiter.py index 5c6fa8864a5e..d3ca4597f0da 100644 --- a/litellm/proxy/hooks/parallel_request_limiter.py +++ b/litellm/proxy/hooks/parallel_request_limiter.py @@ -13,10 +13,7 @@ from litellm._logging import verbose_proxy_logger from litellm.exceptions import RateLimitType from litellm.integrations.custom_logger import CustomLogger -from litellm.litellm_core_utils.core_helpers import ( - get_parent_otel_span_from_kwargs, - is_batch_line_item_event, -) +from litellm.litellm_core_utils.core_helpers import get_parent_otel_span_from_kwargs from litellm.proxy._types import CommonProxyErrors, CurrentItemRateLimit, UserAPIKeyAuth from litellm.proxy.auth.auth_utils import ( get_key_model_rpm_limit, @@ -505,8 +502,6 @@ async def async_pre_call_hook( @with_service_target("rate_limits") async def async_log_success_event(self, kwargs, response_obj: object, start_time, end_time) -> None: - if is_batch_line_item_event(kwargs): - return from litellm.proxy.common_utils.callback_utils import ( get_model_group_from_litellm_kwargs, ) @@ -710,8 +705,6 @@ async def async_log_success_event(self, kwargs, response_obj: object, start_time @with_service_target("rate_limits") async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time) -> None: - if is_batch_line_item_event(kwargs): - return try: self.print_verbose("Inside Max Parallel Request Failure Hook") litellm_parent_otel_span: Final[Span | None] = get_parent_otel_span_from_kwargs(kwargs=kwargs) diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 27018b80c142..21025a8a820c 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -43,7 +43,6 @@ from litellm.caching.redis_cache import log_redis_failure from litellm.constants import DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE, INTERNAL_CALL_ORIGIN_METADATA_KEY from litellm.integrations.custom_logger import CustomLogger -from litellm.litellm_core_utils.core_helpers import is_batch_line_item_event from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_str_from_messages, ) @@ -5056,8 +5055,6 @@ async def async_log_success_event(self, kwargs, response_obj, start_time, end_ti """ Update TPM usage on successful API calls by incrementing counters using pipeline """ - if is_batch_line_item_event(kwargs): - return from litellm.litellm_core_utils.core_helpers import ( get_parent_otel_span_from_kwargs, ) @@ -5181,8 +5178,6 @@ async def async_log_failure_event(self, kwargs, response_obj, start_time, end_ti whose partial usage was recovered settles the reservation at that usage instead of refunding it. """ - if is_batch_line_item_event(kwargs): - return from litellm.litellm_core_utils.core_helpers import ( get_parent_otel_span_from_kwargs, ) diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 672a9618f328..5c31bce6f333 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -14,7 +14,6 @@ get_litellm_metadata_from_kwargs, get_metadata_variable_name_from_kwargs, get_parent_otel_span_from_kwargs, - is_batch_line_item_event, ) from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import guardrail_information_cost @@ -135,8 +134,6 @@ def __init__( async def async_log_success_event( self, kwargs: ObjectMapping, response_obj: object, start_time: datetime, end_time: datetime ) -> None: - if is_batch_line_item_event(kwargs): - return if self.spend_event_producer is None or not is_offloadable_success(response_obj): await self._PROXY_track_cost_callback(kwargs, response_obj, start_time, end_time) return diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 58e5f8ae7a8f..c2ca045236c3 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -73,7 +73,6 @@ LITELLM_SETTINGS_SAFE_DB_OVERRIDES, LITELLM_UI_ALLOW_HEADERS, LITELLM_UI_SESSION_DURATION, - MAX_TIME_TO_CLEAR_QUEUE, RUNTIME_UPDATABLE_ROUTER_SETTINGS, ) from litellm.litellm_core_utils.asyncify import asyncify @@ -272,11 +271,6 @@ def generate_feedback_box(): from litellm import Router from litellm._internal_context import service_target, with_service_target from litellm._logging import redact_string, verbose_proxy_logger, verbose_router_logger -from litellm.batches.batch_line_item_logging import ( - _LINE_ITEM_CLAIM_TTL_SECONDS, # pyright: ignore[reportPrivateUsage] # the claim window is defined next to the cache it expires - batch_line_item_claim_cache, - drain_batch_line_items, -) from litellm.caching.caching import DualCache, RedisCache from litellm.caching.dual_cache import DeclaredBatchRead from litellm.caching.redis_batch import ( @@ -1187,10 +1181,6 @@ async def _flush_spend_logs_queue_on_shutdown() -> None: async def proxy_shutdown_event(worker_heartbeat: ProxyWorkerHeartbeat | None = None) -> None: global prisma_client, master_key, user_custom_auth, user_custom_key_generate, user_custom_key_update verbose_proxy_logger.info("Shutting down LiteLLM Proxy Server") - # Batch line fan-outs run outside the logging queue and keep their batch claim - # once started, so let them finish while the claim cache, database and - # callback clients are still open. - await drain_batch_line_items(timeout=MAX_TIME_TO_CLEAR_QUEUE) if worker_heartbeat is not None and prisma_client: await worker_heartbeat.deregister() if prisma_client: @@ -5112,8 +5102,8 @@ def _attach_redis_usage_cache(redis_cache: RedisCache, enable_redis_auth_cache: """ Wires an established coordination Redis into the proxy-level caches that consume it directly: the spend counter cache, the CLI SSO login-session - cache, the batch line-item claim cache, the cluster-wide config cache, - and (only when opted in) the virtual-key auth cache. + cache, the cluster-wide config cache, and (only when opted in) the + virtual-key auth cache. The CLI SSO login-session cache is always backed by Redis when available so that the browser SSO flow behind `lite login` survives landing on different @@ -5127,10 +5117,6 @@ def _attach_redis_usage_cache(redis_cache: RedisCache, enable_redis_auth_cache: redis_cache, default_redis_ttl=CLI_SSO_SESSION_TTL_SECONDS, ) - batch_line_item_claim_cache.attach_redis_cache( - redis_cache, - default_redis_ttl=_LINE_ITEM_CLAIM_TTL_SECONDS, - ) if enable_redis_auth_cache is True: user_api_key_cache.attach_redis_cache( redis_cache, @@ -7067,9 +7053,6 @@ async def load_config(self, router: litellm.Router | None, config_file_path: str health_check_interval = general_settings.get("health_check_interval", DEFAULT_HEALTH_CHECK_INTERVAL) health_check_concurrency = general_settings.get("health_check_concurrency", None) health_check_details = general_settings.get("health_check_details", True) - _store_batch_line_items: Final = coerce_bool(general_settings.get("store_batch_line_items_in_callbacks")) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # general_settings is untyped upstream - if _store_batch_line_items is not None: - litellm.store_batch_line_items_in_callbacks = bool(_store_batch_line_items) # Health-check-driven routing (opt-in, passes through to Router later) _enable_hc_routing = general_settings.get("enable_health_check_routing", False) _hc_staleness = general_settings.get("health_check_staleness_threshold", None) @@ -8095,7 +8078,6 @@ async def _apply_general_settings_side_effects( self._apply_alerting_settings, self._apply_pass_through_settings, self._apply_boolean_settings, - self._apply_batch_line_items_setting, partial(self._apply_cache_size_setting, cache_size_was_db=cache_size_was_db), self._apply_store_model_in_db_setting, partial(self._apply_retention_settings, previous_cleanup_schedule=previous_cleanup_schedule), @@ -8157,11 +8139,6 @@ async def _apply_boolean_settings(self, db_values: Mapping[str, SettingsJsonValu if (value := self.settings.get(key)) is not None: self.settings[key] = coerce_bool(value) - async def _apply_batch_line_items_setting(self, db_values: Mapping[str, SettingsJsonValue]) -> None: - key: Final = "store_batch_line_items_in_callbacks" - value: Final = coerce_bool(self.settings.get(key)) - litellm.store_batch_line_items_in_callbacks = bool(value) if value is not None else False - async def _apply_cache_size_setting( self, db_values: Mapping[str, SettingsJsonValue], diff --git a/litellm/router.py b/litellm/router.py index 66936b11cc32..cdede46ef5c0 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -83,7 +83,6 @@ get_metadata_variable_name_from_kwargs, get_or_create_metadata_bucket, get_parent_otel_span_from_kwargs, - is_batch_line_item_event, ) from litellm.litellm_core_utils.coroutine_checker import coroutine_checker from litellm.litellm_core_utils.credential_accessor import CredentialAccessor @@ -8419,7 +8418,7 @@ async def deployment_callback_on_success( # WS session wrappers fire with result=None; per-turn costs tracked by inner calls. if kwargs.get("call_type") in ("_aresponses_websocket", "_arealtime"): return - if is_batch_retrieve_call_type(kwargs.get("call_type")) or is_batch_line_item_event(kwargs): + if is_batch_retrieve_call_type(kwargs.get("call_type")): return standard_logging_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object", None) if standard_logging_object is None: @@ -8567,7 +8566,7 @@ def sync_deployment_callback_on_success( - key: str - The key used to increment the cache - None: if no key is found """ - if is_batch_retrieve_call_type(kwargs.get("call_type")) or is_batch_line_item_event(kwargs): + if is_batch_retrieve_call_type(kwargs.get("call_type")): return None id = None if kwargs["litellm_params"].get("metadata") is None: @@ -8607,8 +8606,6 @@ def deployment_callback_on_failure( - False if the deployment should not be put in cooldown """ verbose_router_logger.debug("Router: Entering 'deployment_callback_on_failure'") - if is_batch_line_item_event(kwargs): - return False try: exception: Final = kwargs.get("exception", None) @@ -8723,7 +8720,7 @@ async def async_deployment_callback_on_failure( """ Update RPM usage for a deployment """ - if is_batch_retrieve_call_type(kwargs.get("call_type")) or is_batch_line_item_event(kwargs): + if is_batch_retrieve_call_type(kwargs.get("call_type")): return deployment_name: Final = kwargs["litellm_params"]["metadata"].get( "deployment", None diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index 3b632df2178c..8ba1da824c65 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/litellm/router_strategy/budget_limiter.py @@ -35,7 +35,6 @@ from litellm.integrations.custom_logger import CustomLogger, Span from litellm.litellm_core_utils.core_helpers import ( get_metadata_variable_name_from_kwargs, - is_batch_line_item_event, ) from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.router_strategy.tag_based_routing import get_tags_from_request_kwargs @@ -496,8 +495,6 @@ async def _write_queued_increment_operations(self, redis_cache: RedisCache) -> b async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): """Original method now uses helper functions""" verbose_router_logger.debug("in RouterBudgetLimiting.async_log_success_event") - if is_batch_line_item_event(kwargs): - return # WS session wrappers fire with result=None; per-turn costs tracked by inner calls. if kwargs.get("call_type") in ("_aresponses_websocket", "_arealtime"): return diff --git a/litellm/router_strategy/least_busy.py b/litellm/router_strategy/least_busy.py index 7cfada62ecc1..6f3e09366412 100644 --- a/litellm/router_strategy/least_busy.py +++ b/litellm/router_strategy/least_busy.py @@ -10,7 +10,6 @@ from litellm.caching.caching import DualCache from litellm.caching.redis_cache import log_redis_failure from litellm.integrations.custom_logger import CustomLogger -from litellm.litellm_core_utils.core_helpers import is_batch_line_item_event from litellm.router_utils.batch_utils import is_batch_retrieve_call_type IN_FLIGHT_COUNT_TTL_SECONDS: Final = 60 * 60 @@ -51,7 +50,7 @@ def _request_count_key(model_group: str, deployment_id: str) -> str: def _deployment_ref(kwargs: Mapping[str, object]) -> tuple[str, str] | None: - if is_batch_retrieve_call_type(kwargs.get("call_type")) or is_batch_line_item_event(kwargs): + if is_batch_retrieve_call_type(kwargs.get("call_type")): return None try: call: Final = _CALL_KWARGS.validate_python(kwargs) diff --git a/litellm/router_strategy/lowest_cost.py b/litellm/router_strategy/lowest_cost.py index 1bcf248331de..d567b6accccc 100644 --- a/litellm/router_strategy/lowest_cost.py +++ b/litellm/router_strategy/lowest_cost.py @@ -9,7 +9,6 @@ from litellm._logging import verbose_router_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger -from litellm.litellm_core_utils.core_helpers import is_batch_line_item_event from litellm.router_utils.batch_utils import is_batch_retrieve_call_type @@ -23,7 +22,7 @@ def __init__(self, router_cache: DualCache, routing_args: dict = {}): @with_service_target("router_usage") def log_success_event(self, kwargs, response_obj, start_time, end_time): - if is_batch_retrieve_call_type(kwargs.get("call_type")) or is_batch_line_item_event(kwargs): + if is_batch_retrieve_call_type(kwargs.get("call_type")): return try: """ @@ -99,7 +98,7 @@ def log_success_event(self, kwargs, response_obj, start_time, end_time): @with_service_target("router_usage") async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): - if is_batch_retrieve_call_type(kwargs.get("call_type")) or is_batch_line_item_event(kwargs): + if is_batch_retrieve_call_type(kwargs.get("call_type")): return try: """ diff --git a/litellm/router_strategy/lowest_latency.py b/litellm/router_strategy/lowest_latency.py index ca9935a18f8d..83bb22435e67 100644 --- a/litellm/router_strategy/lowest_latency.py +++ b/litellm/router_strategy/lowest_latency.py @@ -13,11 +13,7 @@ from litellm._internal_context import with_service_target from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger -from litellm.litellm_core_utils.core_helpers import ( - get_parent_otel_span_from_kwargs, - is_batch_line_item_event, - safe_divide_seconds, -) +from litellm.litellm_core_utils.core_helpers import get_parent_otel_span_from_kwargs, safe_divide_seconds from litellm.router_utils.batch_utils import is_batch_retrieve_call_type from litellm.types.utils import LiteLLMPydanticObjectBase @@ -65,7 +61,7 @@ def __init__(self, router_cache: DualCache, routing_args: dict = {}): @with_service_target("router_usage") def log_success_event(self, kwargs, response_obj, start_time, end_time): - if is_batch_retrieve_call_type(kwargs.get("call_type")) or is_batch_line_item_event(kwargs): + if is_batch_retrieve_call_type(kwargs.get("call_type")): return try: """ @@ -192,7 +188,7 @@ async def async_log_failure_event(self, kwargs, response_obj, start_time, end_ti """ Check if Timeout Error, if timeout set deployment latency -> 100 """ - if is_batch_retrieve_call_type(kwargs.get("call_type")) or is_batch_line_item_event(kwargs): + if is_batch_retrieve_call_type(kwargs.get("call_type")): return try: metadata_field: Final = self._select_metadata_field(kwargs) @@ -249,7 +245,7 @@ async def async_log_failure_event(self, kwargs, response_obj, start_time, end_ti @with_service_target("router_usage") async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): - if is_batch_retrieve_call_type(kwargs.get("call_type")) or is_batch_line_item_event(kwargs): + if is_batch_retrieve_call_type(kwargs.get("call_type")): return try: """ diff --git a/litellm/router_strategy/lowest_tpm_rpm.py b/litellm/router_strategy/lowest_tpm_rpm.py index 11f28e6c7a54..2d373e0c2661 100644 --- a/litellm/router_strategy/lowest_tpm_rpm.py +++ b/litellm/router_strategy/lowest_tpm_rpm.py @@ -9,7 +9,6 @@ from litellm._logging import verbose_router_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger -from litellm.litellm_core_utils.core_helpers import is_batch_line_item_event from litellm.router_utils.batch_utils import is_batch_retrieve_call_type from litellm.types.utils import LiteLLMPydanticObjectBase from litellm.utils import print_verbose @@ -31,7 +30,7 @@ def __init__(self, router_cache: DualCache, routing_args: dict = {}): @with_service_target("router_usage") def log_success_event(self, kwargs, response_obj, start_time, end_time): - if is_batch_retrieve_call_type(kwargs.get("call_type")) or is_batch_line_item_event(kwargs): + if is_batch_retrieve_call_type(kwargs.get("call_type")): return try: """ @@ -86,7 +85,7 @@ def log_success_event(self, kwargs, response_obj, start_time, end_time): @with_service_target("router_usage") async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): - if is_batch_retrieve_call_type(kwargs.get("call_type")) or is_batch_line_item_event(kwargs): + if is_batch_retrieve_call_type(kwargs.get("call_type")): return try: """ diff --git a/litellm/router_strategy/lowest_tpm_rpm_v2.py b/litellm/router_strategy/lowest_tpm_rpm_v2.py index 225517cb090f..2b37e7dfb404 100644 --- a/litellm/router_strategy/lowest_tpm_rpm_v2.py +++ b/litellm/router_strategy/lowest_tpm_rpm_v2.py @@ -15,7 +15,7 @@ from litellm._logging import verbose_logger, verbose_router_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger -from litellm.litellm_core_utils.core_helpers import get_parent_otel_span_from_kwargs, is_batch_line_item_event +from litellm.litellm_core_utils.core_helpers import get_parent_otel_span_from_kwargs from litellm.router_utils.batch_utils import is_batch_retrieve_call_type from litellm.types.router import RouterErrors from litellm.types.utils import LiteLLMPydanticObjectBase, StandardLoggingPayload @@ -254,7 +254,7 @@ async def async_pre_call_check(self, deployment: dict, parent_otel_span: Span | @with_service_target("router_usage") def log_success_event(self, kwargs, response_obj, start_time, end_time): - if is_batch_retrieve_call_type(kwargs.get("call_type")) or is_batch_line_item_event(kwargs): + if is_batch_retrieve_call_type(kwargs.get("call_type")): return try: """ @@ -297,7 +297,7 @@ def log_success_event(self, kwargs, response_obj, start_time, end_time): @with_service_target("router_usage") async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): - if is_batch_retrieve_call_type(kwargs.get("call_type")) or is_batch_line_item_event(kwargs): + if is_batch_retrieve_call_type(kwargs.get("call_type")): return try: """ diff --git a/litellm/router_utils/pre_call_checks/model_rate_limit_check.py b/litellm/router_utils/pre_call_checks/model_rate_limit_check.py index adeb21855431..17e7ec176ab2 100644 --- a/litellm/router_utils/pre_call_checks/model_rate_limit_check.py +++ b/litellm/router_utils/pre_call_checks/model_rate_limit_check.py @@ -21,7 +21,6 @@ from litellm.caching.dual_cache import DualCache from litellm.caching.redis_cache import RedisCircuitBreakerOpenError from litellm.integrations.custom_logger import CustomLogger -from litellm.litellm_core_utils.core_helpers import is_batch_line_item_event from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import ( ITPM_RESERVED_KEY, async_io_token_pre_call_check, @@ -336,8 +335,6 @@ async def async_log_success_event(self, kwargs, response_obj, start_time, end_ti get_parent_otel_span_from_kwargs, ) - if is_batch_line_item_event(kwargs): - return try: standard_logging_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object") @@ -396,8 +393,6 @@ async def async_log_failure_event(self, kwargs, response_obj, start_time, end_ti get_parent_otel_span_from_kwargs, ) - if is_batch_line_item_event(kwargs): - return # Never fail the primary logging pipeline over an io-token refund error. with contextlib.suppress(Exception): await async_io_token_refund_failure( diff --git a/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py b/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py index 74aaa5c37533..9f1558dedae6 100644 --- a/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py +++ b/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py @@ -15,7 +15,6 @@ AnthropicCacheControlHook, ) from litellm.integrations.custom_logger import CustomLogger, Span -from litellm.litellm_core_utils.core_helpers import is_batch_line_item_event from litellm.litellm_core_utils.token_counter import offload_token_count from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import CallTypes, StandardLoggingPayload @@ -115,8 +114,6 @@ async def async_filter_deployments( return healthy_deployments async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): - if is_batch_line_item_event(kwargs): - return standard_logging_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object", None) if standard_logging_object is None: diff --git a/litellm/types/utils.py b/litellm/types/utils.py index cb0326ec5892..451d4bf7b293 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3313,9 +3313,6 @@ class StandardLoggingHiddenParams(TypedDict): batch_failed_requests: ReadOnly[int | None] litellm_model_name: str | None # the model name sent to the provider by litellm usage_object: dict | None - batch_id: NotRequired[ReadOnly[str | None]] - batch_custom_id: NotRequired[ReadOnly[str | None]] - batch_line_status_code: NotRequired[ReadOnly[int | None]] class StandardLoggingModelInformation(TypedDict): @@ -3650,6 +3647,44 @@ 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): + """Every line of one finished batch, dispatched to batch_line_item_callbacks.""" + + batch_id: ReadOnly[str] + finished_at: 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/management/test_batch_output_file_listing.py b/tests/integration/management/test_batch_output_file_listing.py index 7aa10f0d27fa..26d1ae8de46e 100644 --- a/tests/integration/management/test_batch_output_file_listing.py +++ b/tests/integration/management/test_batch_output_file_listing.py @@ -8,7 +8,7 @@ import ssl import threading import uuid -from collections.abc import Callable, Generator +from collections.abc import Generator from contextlib import contextmanager from dataclasses import dataclass from pathlib import Path @@ -16,21 +16,12 @@ from typing import Final import httpx -import yaml from cryptography import x509 from cryptography.hazmat.primitives import hashes, serialization from cryptography.hazmat.primitives.asymmetric import ec from cryptography.x509.oid import NameOID -from integration._support.client import ( - JSON_OBJECT, - Gateway, - Scenario, - eventually, - list_value, - object_value, - string_value, -) -from integration._support.database import read_rows, scratch_database +from integration._support.client import JSON_OBJECT, Gateway, Scenario, eventually, object_value, string_value +from integration._support.database import read_rows from integration._support.process import owned_proxy from integration._support.upstream import ScenarioHandle, delete_scenario, register_scenario from integration._support.wire import Reply, Request, Wire, wire_server @@ -43,7 +34,6 @@ ERROR_BYTES: Final = 512 BEDROCK_MODEL: Final = "bedrock/anthropic.claude-3-haiku-20240307-v1:0" BEDROCK_MODEL_ID: Final = "anthropic.claude-3-haiku-20240307-v1:0" -BEDROCK_TITAN_MODEL_ID: Final = "amazon.titan-embed-text-v2:0" BEDROCK_REGION: Final = "us-east-1" BEDROCK_AUTHORITY: Final = f"bedrock.{BEDROCK_REGION}.amazonaws.com:443" BEDROCK_BUCKET: Final = "integration-batch-listing-bucket" @@ -52,7 +42,6 @@ BEDROCK_LAST_MODIFIED: Final = "Thu, 02 Oct 2025 12:00:00 GMT" BEDROCK_OUTPUT_CONTENT: Final = b'{"recordId":"req-1","modelOutput":{}}\n' _BATCH_PROCESSED_SQL: Final = 'SELECT batch_processed FROM "LiteLLM_ManagedObjectTable" WHERE unified_object_id=%s' -_ANY_BATCH_PROCESSED_SQL: Final = 'SELECT 1 FROM "LiteLLM_ManagedObjectTable" WHERE batch_processed' @dataclass(frozen=True, slots=True) @@ -494,7 +483,6 @@ class _BedrockControlPlane: job_locations: SimpleQueue[tuple[str, str, str]] job_id: str job_arn: str - model_id: str = BEDROCK_MODEL_ID def __call__(self, request: Request) -> Reply: if request.method == "POST" and request.target == "/model-invocation-job": @@ -514,7 +502,7 @@ def __call__(self, request: Request) -> Reply: { "jobArn": self.job_arn, "jobName": retrieved_job_name, - "modelId": self.model_id, + "modelId": BEDROCK_MODEL_ID, "status": "Completed", "submitTime": 1700000000, "lastModifiedTime": 1700000001, @@ -661,227 +649,3 @@ def test_repeated_batch_retrieve_does_not_refetch_saved_file_details(gateway: Ga assert output_after == output_before, output_after additional_hits: Final = _metadata_hit_count(gateway, batch.scenario) assert additional_hits == 0, f"Repeated batch retrieval fetched metadata {additional_hits} more times" - - -@dataclass(frozen=True, slots=True) -class _StoringS3Peer: - objects: dict[str, bytes] # mutable-ok: the peer keeps every uploaded object so later GETs can read it back - output: bytes - - def __call__(self, request: Request) -> Reply: - path: Final = request.target.split("?", 1)[0] - if request.method == "PUT" and path.startswith(f"/{BEDROCK_BUCKET}/"): - self.objects[path] = request.body - return Reply(body=b"") - if request.method == "GET" and path.startswith(f"/{BEDROCK_BUCKET}/{BEDROCK_MANAGED_S3_OUTPUT_PREFIX}"): - if request.headers.get("range") == "bytes=0-0": - return Reply( - status=206, - body=self.output[:1], - headers={ - "Content-Range": f"bytes 0-0/{len(self.output)}", - "Last-Modified": BEDROCK_LAST_MODIFIED, - }, - ) - return Reply(status=200, body=self.output, headers={"Last-Modified": BEDROCK_LAST_MODIFIED}) - stored: Final = self.objects.get(path) - if request.method == "GET" and stored is not None: - return Reply(status=200, body=stored, headers={"Last-Modified": BEDROCK_LAST_MODIFIED}) - return Reply(status=404, body=b'{"message":"not scripted"}') - - -def _bedrock_line_output(marker: str) -> bytes: - succeeded: Final = { - "recordId": "req-1", - "modelInput": {"messages": [{"role": "user", "content": f"{marker} ping"}], "max_tokens": 8}, - "modelOutput": { - "id": "msg_bdrk_line_1", - "type": "message", - "role": "assistant", - "model": BEDROCK_MODEL_ID, - "content": [{"type": "text", "text": "pong"}], - "stop_reason": "end_turn", - "stop_sequence": None, - "usage": {"input_tokens": 8, "output_tokens": 2}, - }, - } - failed: Final = { - "recordId": "req-2", - "modelInput": {"messages": [{"role": "user", "content": f"{marker} bad"}], "max_tokens": 8}, - "error": {"errorCode": 400, "errorMessage": "Malformed input request"}, - } - return (json.dumps(succeeded) + "\n" + json.dumps(failed) + "\n").encode() - - -def _bedrock_titan_line_output(marker: str) -> bytes: - succeeded: Final = { - "recordId": "req-1", - "modelInput": {"inputText": f"{marker} ping"}, - "modelOutput": {"embedding": [0.1, 0.2], "inputTextTokenCount": 3}, - } - failed: Final = { - "recordId": "req-2", - "modelInput": {"inputText": f"{marker} bad"}, - "error": {"errorCode": 400, "errorMessage": "Malformed input request"}, - } - return (json.dumps(succeeded) + "\n" + json.dumps(failed) + "\n").encode() - - -@dataclass(frozen=True, slots=True) -class _BedrockLineBatch: - model_id: str - endpoint: str - output: Callable[[str], bytes] - - def body(self, content: str) -> dict[str, JsonValue]: - if self.endpoint == "/v1/embeddings": - return {"input": content} - return { - "messages": [{"role": "system", "content": "answer briefly"}, {"role": "user", "content": content}], - "max_tokens": 8, - } - - -_BEDROCK_CHAT_BATCH: Final = _BedrockLineBatch(BEDROCK_MODEL_ID, "/v1/chat/completions", _bedrock_line_output) -_BEDROCK_TITAN_BATCH: Final = _BedrockLineBatch(BEDROCK_TITAN_MODEL_ID, "/v1/embeddings", _bedrock_titan_line_output) - - -def _line_items_config(destination: Path) -> Path: - config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) - config["litellm_settings"].update({"callbacks": ["generic_api"], "DEFAULT_FLUSH_INTERVAL_SECONDS": 1}) - config["general_settings"].update({"store_batch_line_items_in_callbacks": True}) - destination.write_text(yaml.safe_dump(config)) - return destination - - -def _line_events_of_a_polled_bedrock_batch( - gateway: Gateway, tmp_path: Path, batch: _BedrockLineBatch -) -> tuple[dict[str, JsonValue], ...]: - job_id: Final = f"integration-batch-lines-{uuid.uuid4().hex}" - marker: Final = f"bedrock-poller-lines-{uuid.uuid4().hex[:12]}" - received: Final[list[Request]] = [] # mutable-ok: drain() consumes the queue, later polls keep earlier events - environment: Final = { - "AWS_CA_BUNDLE": str(tmp_path / "bedrock.pem"), - "AWS_EC2_METADATA_DISABLED": "true", - "SSL_VERIFY": "False", - } - with ( - scratch_database() as database_url, - wire_server(lambda _request: Reply()) as sink, - wire_server(_StoringS3Peer({}, batch.output(marker))) as s3, - wire_server( - _BedrockControlPlane(SimpleQueue(), job_id, BEDROCK_JOB_ARN_PREFIX + job_id, batch.model_id), - tls=_tls_context(tmp_path), - ) as bedrock, - _bedrock_tunnel(bedrock) as tunnel, - owned_proxy( - gateway, - tmp_path, - { - **environment, - "HTTPS_PROXY": tunnel.url, - "GENERIC_LOGGER_ENDPOINT": sink.url, - "DATABASE_URL": database_url, - }, - config=_line_items_config(tmp_path / "bedrock_line_items.yaml"), - ) as candidate, - candidate.scenario() as scenario, - ): - model: Final = scenario.model( - model=f"bedrock/{batch.model_id}", - api_key=None, - api_base=None, - aws_access_key_id="AKIAIOSFODNN7EXAMPLE", - aws_secret_access_key="wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY", - aws_region_name=BEDROCK_REGION, - s3_bucket_name=BEDROCK_BUCKET, - s3_endpoint_url=s3.url, - aws_batch_role_arn=BEDROCK_ROLE_ARN, - ) - owner_key: Final = scenario.key(user_id=scenario.user(user_role="internal_user"), models=[model]) - input_lines: Final = "".join( - json.dumps( - { - "custom_id": custom_id, - "method": "POST", - "url": batch.endpoint, - "body": {"model": model, **batch.body(f"{marker} {text}")}, - } - ) - + "\n" - for custom_id, text in (("req-1", "ping"), ("req-2", "bad")) - ) - uploaded: Final = candidate.request_multipart( - "/v1/files", - {"purpose": "batch", "target_model_names": model}, - {"file": ("input.jsonl", input_lines.encode(), "application/jsonl")}, - key=owner_key, - ) - assert uploaded.status_code == 200, uploaded.text - created: Final = candidate.request( - "POST", - "/v1/batches", - { - "input_file_id": string_value(JSON_OBJECT.validate_json(uploaded.content)["id"]), - "endpoint": batch.endpoint, - "completion_window": "24h", - }, - key=owner_key, - ) - assert created.status_code == 200, created.text - - def line_events() -> tuple[dict[str, JsonValue], ...]: - received.extend(sink.drain()) - return tuple( - event - for delivery in received - for event in json.loads(delivery.body) - if object_value(event["hidden_params"]).get("batch_custom_id") is not None - and marker in json.dumps(event["messages"]) - ) - - eventually(lambda: read_rows(_ANY_BATCH_PROCESSED_SQL, (), database_url=database_url), bool, seconds=60) - return eventually( - line_events, - lambda values: ( - {object_value(event["hidden_params"])["batch_custom_id"] for event in values} == {"req-1", "req-2"} - ), - seconds=30, - ) - - -def _line_outcomes(events: tuple[dict[str, JsonValue], ...]) -> list[tuple[str, JsonValue, JsonValue, JsonValue]]: - return sorted( - ( - string_value(object_value(event["hidden_params"])["batch_custom_id"]), - event["status"], - object_value(event["hidden_params"])["batch_line_status_code"], - event["call_type"], - ) - for event in events - ) - - -def test_a_bedrock_batch_the_poller_logs_emits_its_lines_from_the_deployment_bucket( - gateway: Gateway, tmp_path: Path -) -> None: - """The cost poller reads a managed Bedrock batch's files with the deployment's S3 bucket, so the - line events it starts must read the input file from that same bucket""" - delivered: Final = _line_events_of_a_polled_bedrock_batch(gateway, tmp_path, _BEDROCK_CHAT_BATCH) - - assert _line_outcomes(delivered) == [ - ("req-1", "success", 200, "acompletion"), - ("req-2", "failure", 400, "acompletion"), - ], delivered - first_messages: Final = [object_value(list_value(event["messages"])[0]) for event in delivered] - assert [message["role"] for message in first_messages] == ["system", "system"], delivered - assert all("answer briefly" in json.dumps(message["content"]) for message in first_messages), delivered - - -def test_a_failed_line_of_a_bedrock_embedding_batch_is_logged_as_an_embedding(gateway: Gateway, tmp_path: Path) -> None: - delivered: Final = _line_events_of_a_polled_bedrock_batch(gateway, tmp_path, _BEDROCK_TITAN_BATCH) - - assert _line_outcomes(delivered) == [ - ("req-1", "success", 200, "aembedding"), - ("req-2", "failure", 400, "aembedding"), - ], delivered 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 000000000000..377c7dc6ff89 --- /dev/null +++ b/tests/integration/observability/test_batch_line_items_cold_storage.py @@ -0,0 +1,499 @@ +"""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 one NDJSON object for the batch holding one record per line. The proxy runs +on its own database so the shared rig's pollers never claim these batches first +""" + +import json +import threading +import uuid +from collections.abc import Callable, Iterator, Mapping +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 pytest +import yaml +from integration._support.client import JSON_OBJECT, Gateway, eventually, gateway_from_environment, 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 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 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 + + +@dataclass(slots=True) +class ObjectSink: + """Stores every accepted object body by name and answers 503 to any upload whose body holds ``reject``.""" + + reject: bytes = b"\x00never" + lock: threading.Lock = field(default_factory=threading.Lock) + 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 + if self.reject in body: + return False + self.stored[name] = body + return True + + 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 + + +def _config(directory: Path, s3_url: str) -> Path: + shipped: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + configuration: Final = { + **shipped, + "litellm_settings": { + **shipped["litellm_settings"], + "callbacks": ["s3_v2"], + "batch_line_item_callbacks": ["s3_v2", "gcs_bucket"], + "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 + + +@pytest.fixture(scope="module") +def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + directory: Final = tmp_path_factory.mktemp("batch-line-items") + certificate, key = write_self_signed_cert(directory, names=("storage.googleapis.com",)) + s3: Final = ObjectSink(reject=b"sink-down-") + gcs: Final = ObjectSink() + with ( + gateway_from_environment() as environment, + scratch_database("batch_line_items") 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", + "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", + }, + config=_config(directory, s3_wire.url), + remove_environment=("DATABASE_URL_READ_REPLICA",), + workers=2, + ) as candidate: + yield Rig(candidate, database_url, s3, gcs) + + +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) -> 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), + "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]) -> SubmittedBatch: + handle: Final = register_scenario(f"batch-line-items-{uuid.uuid4().hex}", _routes(lines)) + model: Final = 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"] + ) + key: Final = 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 _object_name(prefix: str, batch_id: str) -> str: + return f"{prefix}batch_line_items/{LINE_DATE}/{batch_id.replace('/', '_')}.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": { + "logging": [ + { + "callback_name": "s3_v2", + "callback_type": "success", + "callback_vars": {"turn_off_message_logging": "true"}, + } + ] + } + }, + ) + 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_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_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() diff --git a/tests/integration/spend/test_batch_line_item_callbacks.py b/tests/integration/spend/test_batch_line_item_callbacks.py deleted file mode 100644 index cec7f0fd3cc5..000000000000 --- a/tests/integration/spend/test_batch_line_item_callbacks.py +++ /dev/null @@ -1,817 +0,0 @@ -from __future__ import annotations - -import json -import uuid -from collections.abc import Callable -from dataclasses import dataclass -from datetime import UTC, datetime -from hashlib import sha256 -from pathlib import Path -from typing import Final - -import pytest -import yaml -from integration._support.client import JSON_OBJECT, Gateway, eventually, object_value, string_value -from integration._support.database import read_rows, scratch_database -from integration._support.process import owned_proxy -from integration._support.upstream import delete_scenario, register_scenario -from integration._support.wire import Reply, Request, Wire, wire_server -from integration.cost_calculation.cost_tracking_case import JsonResponse, RoutedResponse, TextResponse -from pydantic import JsonValue - -PROMPT_TOKENS: Final = 10 -COMPLETION_TOKENS: Final = 100 -OUTPUT_SUCCESS_IDS: Final = ("r1", "r3") -OUTPUT_FAILED_IDS: Final = ("r2",) -ERROR_FILE_IDS: Final = ("r4", "r5") -ALL_CUSTOM_IDS: Final = ("r1", "r2", "r3", "r4", "r5") -LIVE_TOKENS: Final = 10 -_PROCESSED_BATCH_SQL: Final = ( - 'SELECT unified_object_id FROM "LiteLLM_ManagedObjectTable" WHERE unified_object_id=%s AND batch_processed' -) - - -def _successful_line(custom_id: str) -> str: - return json.dumps( - { - "id": f"batch_req_{custom_id}", - "custom_id": custom_id, - "response": { - "status_code": 200, - "request_id": f"$REQUEST_ID-{custom_id}", - "body": { - "id": f"chatcmpl-$REQUEST_ID-{custom_id}", - "object": "chat.completion", - "model": "gpt-4o-mini", - "choices": [ - { - "index": 0, - "message": {"role": "assistant", "content": f"answer {custom_id}"}, - "finish_reason": "stop", - } - ], - "usage": { - "prompt_tokens": PROMPT_TOKENS, - "completion_tokens": COMPLETION_TOKENS, - "total_tokens": PROMPT_TOKENS + COMPLETION_TOKENS, - }, - }, - }, - "error": None, - }, - separators=(",", ":"), - ) - - -def _failed_line(custom_id: str) -> str: - return json.dumps( - { - "id": f"batch_req_{custom_id}", - "custom_id": custom_id, - "response": { - "status_code": 400, - "request_id": f"$REQUEST_ID-{custom_id}", - "body": {"error": {"message": f"synthetic line failure {custom_id}", "code": "bad_request"}}, - }, - "error": {"code": "bad_request", "message": f"synthetic line failure {custom_id}"}, - }, - separators=(",", ":"), - ) - - -def _failed_line_with_text_body(custom_id: str) -> str: - """An error file line whose response body is the provider's error as JSON text, as Mistral writes it.""" - return json.dumps( - { - "id": f"batch_req_{custom_id}", - "custom_id": custom_id, - "response": { - "status_code": 400, - "body": json.dumps({"object": "error", "message": f"synthetic line failure {custom_id}"}), - }, - "error": None, - }, - separators=(",", ":"), - ) - - -def _batch(status: str, *, files_ready: bool) -> dict[str, JsonValue]: - return { - "id": "batch-$REQUEST_ID", - "object": "batch", - "endpoint": "/v1/chat/completions", - "errors": None, - "input_file_id": "file-in-$REQUEST_ID", - "completion_window": "24h", - "status": status, - "output_file_id": "file-out-$REQUEST_ID" if files_ready else None, - "error_file_id": "file-err-$REQUEST_ID" if files_ready else None, - "created_at": 1, - "in_progress_at": 1, - "completed_at": 1 if files_ready else None, - "expires_at": 1, - "request_counts": {"total": 5, "completed": 2, "failed": 3}, - "metadata": None, - } - - -def _input_lines(model_name: str, marker: str) -> tuple[str, ...]: - return tuple( - json.dumps( - { - "custom_id": custom_id, - "method": "POST", - "url": "/v1/chat/completions", - "body": {"model": model_name, "messages": [{"role": "user", "content": f"{marker} {custom_id}"}]}, - }, - separators=(",", ":"), - ) - for custom_id in ALL_CUSTOM_IDS - ) - - -def _provider_routes( - input_lines: tuple[str, ...], *, error_line: Callable[[str], str] = _failed_line -) -> RoutedResponse: - output_lines: Final = (_successful_line("r1"), _failed_line("r2"), _successful_line("r3")) - error_lines: Final = tuple(error_line(custom_id) for custom_id in ERROR_FILE_IDS) - return RoutedResponse( - content_type="application/x-routed", - routes={ - "POST /chat/completions": JsonResponse( - content_type="application/json", - body={ - "id": "chatcmpl-live-$REQUEST_ID", - "object": "chat.completion", - "created": 1, - "model": "gpt-4o-mini", - "choices": [ - {"index": 0, "message": {"role": "assistant", "content": "live"}, "finish_reason": "stop"} - ], - "usage": {"prompt_tokens": LIVE_TOKENS, "completion_tokens": 0, "total_tokens": LIVE_TOKENS}, - }, - ), - "POST /files": JsonResponse( - content_type="application/json", - body={ - "id": "file-in-$REQUEST_ID", - "object": "file", - "purpose": "batch", - "bytes": 100, - "created_at": 1, - "filename": "in.jsonl", - "status": "processed", - }, - ), - "POST /batches": JsonResponse( - content_type="application/json", body=_batch("validating", files_ready=False) - ), - "GET /batches/batch-$REQUEST_ID": JsonResponse( - content_type="application/json", body=_batch("completed", files_ready=True) - ), - "GET /files/file-in-$REQUEST_ID/content": TextResponse( - content_type="application/jsonl", body="\n".join(input_lines) + "\n" - ), - "GET /files/file-out-$REQUEST_ID/content": TextResponse( - content_type="application/jsonl", body="\n".join(output_lines) + "\n" - ), - "GET /files/file-err-$REQUEST_ID/content": TextResponse( - content_type="application/jsonl", body="\n".join(error_lines) + "\n" - ), - }, - ) - - -def _line_item_config(base: Path, destination: Path, *, enforce_model_rate_limits: bool = False) -> Path: - config: Final = yaml.safe_load(base.read_text()) - config["litellm_settings"].update({"callbacks": ["generic_api"], "DEFAULT_FLUSH_INTERVAL_SECONDS": 1}) - config["general_settings"].update({"store_batch_line_items_in_callbacks": True}) - if enforce_model_rate_limits: - config["router_settings"].update({"optional_pre_call_checks": ["enforce_model_rate_limits"]}) - destination.write_text(yaml.safe_dump(config)) - return destination - - -def _spend_rows(key: str) -> tuple[dict[str, JsonValue], ...]: - return tuple( - read_rows( - 'SELECT call_type, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" WHERE api_key=%s', - (sha256(key.encode()).hexdigest(),), - ) - ) - - -def _hidden(event: dict[str, JsonValue]) -> dict[str, JsonValue]: - return object_value(event["hidden_params"]) - - -@pytest.mark.covers( - "spend.batches.completed_batch_emits_one_callback_event_per_jsonl_line_beside_the_aggregate", - "spend.batches.line_item_callback_events_do_not_bill_spend_twice", -) -def test_completed_batch_emits_paired_request_response_callback_events_per_jsonl_line( - gateway: Gateway, tmp_path: Path -) -> None: - marker: Final = "batch-line-items-" + uuid.uuid4().hex[:12] - sink_secret: Final = "synthetic-sink-secret-" + marker - provider_secret: Final = "synthetic-provider-secret-" + marker - - def sink(request: Request) -> Reply: - assert request.headers["authorization"] == f"Bearer {sink_secret}" - return Reply() - - with ( - wire_server(sink) as endpoint, - owned_proxy( - gateway, - tmp_path, - {"GENERIC_LOGGER_ENDPOINT": endpoint.url, "GENERIC_LOGGER_HEADERS": f"Authorization=Bearer {sink_secret}"}, - config=_line_item_config(Path("tests/integration/proxy_config.yaml"), tmp_path / "line_items.yaml"), - ) as candidate, - candidate.scenario() as scenario, - ): - routed_model: Final = scenario.model(api_base=f"{gateway.upstream_url}/{marker}", api_key=provider_secret) - input_lines: Final = _input_lines(routed_model, marker) - handle: Final = register_scenario(marker, _provider_routes(input_lines)) - scenario.cleanups.callback(delete_scenario, handle) - key: Final = scenario.key(models=[routed_model]) - file_response: Final = candidate.request_multipart( - "/v1/files", - {"purpose": "batch", "model": routed_model}, - {"file": ("in.jsonl", ("\n".join(input_lines) + "\n").encode(), "application/jsonl")}, - key=key, - ) - assert file_response.status_code == 200, file_response.text - input_file_id: Final = string_value(JSON_OBJECT.validate_json(file_response.content)["id"]) - batch_response: Final = candidate.request( - "POST", - "/v1/batches", - { - "input_file_id": input_file_id, - "endpoint": "/v1/chat/completions", - "completion_window": "24h", - "model": routed_model, - }, - key=key, - ) - assert batch_response.status_code == 200, batch_response.text - batch_id: Final = string_value(JSON_OBJECT.validate_json(batch_response.content)["id"]) - retrieval: Final = eventually( - lambda: candidate.request("GET", f"/v1/batches/{batch_id}", key=key), - lambda response: response.status_code == 200 and response.json()["status"] == "completed", - seconds=30, - ) - assert retrieval.status_code == 200, retrieval.text - key_hash: Final = sha256(key.encode()).hexdigest() - batches: Final[ - list[Request] - ] = [] # mutable-ok: drain() consumes the queue, later polls must keep earlier batches - - def delivered() -> tuple[dict[str, JsonValue], ...]: - batches.extend(endpoint.drain()) - return tuple( - event - for batch in batches - for event in json.loads(batch.body) - if object_value(event["metadata"]).get("user_api_key_hash") == key_hash - ) - - events: Final = eventually( - delivered, - lambda values: ( - len([e for e in values if e["call_type"] == "acreate_batch"]) >= 1 - and len([e for e in values if e["call_type"] == "aretrieve_batch"]) >= 1 - and len([e for e in values if _hidden(e).get("batch_custom_id") is not None]) >= len(ALL_CUSTOM_IDS) - ), - seconds=40, - ) - for batch in batches: - for credential in (provider_secret, sink_secret, key): - assert credential.encode() not in batch.body - aggregate_events: Final = tuple(event for event in events if event["call_type"] == "aretrieve_batch") - line_events: Final = tuple(event for event in events if _hidden(event).get("batch_custom_id") is not None) - assert len(aggregate_events) == 1, [event["call_type"] for event in events] - other_events: Final = tuple(event for event in events if event not in aggregate_events + line_events) - assert sorted(event["call_type"] for event in other_events) == ["acreate_batch", "acreate_file"], other_events - assert sorted(string_value(_hidden(event)["batch_custom_id"]) for event in line_events) == list( - ALL_CUSTOM_IDS - ), line_events - by_custom_id: Final = {string_value(_hidden(event)["batch_custom_id"]): event for event in line_events} - for custom_id, line in zip(ALL_CUSTOM_IDS, input_lines, strict=True): - event: Final = by_custom_id[custom_id] - hidden: Final = _hidden(event) - assert hidden["batch_id"] == batch_id, hidden - assert event["call_type"] == "acompletion", event["call_type"] - assert event["messages"] == json.loads(line)["body"]["messages"], (custom_id, event["messages"]) - if custom_id in OUTPUT_SUCCESS_IDS: - assert event["status"] == "success", event - assert hidden["batch_line_status_code"] == 200, hidden - assert event["prompt_tokens"] == PROMPT_TOKENS and event["completion_tokens"] == COMPLETION_TOKENS - assert object_value(event["response"])["id"] == f"chatcmpl-{marker}-{custom_id}", event["response"] - assert object_value(object_value(event["response"])["choices"][0]["message"])["content"] == ( - f"answer {custom_id}" - ) - else: - assert event["status"] == "failure", event - assert hidden["batch_line_status_code"] == 400, hidden - assert f"synthetic line failure {custom_id}" in json.dumps(event["error_information"]), event - aggregate: Final = aggregate_events[0] - assert aggregate["prompt_tokens"] == len(OUTPUT_SUCCESS_IDS) * PROMPT_TOKENS, aggregate - assert aggregate["completion_tokens"] == len(OUTPUT_SUCCESS_IDS) * COMPLETION_TOKENS, aggregate - rows: Final = eventually( - lambda: _spend_rows(key), - lambda values: any(r["call_type"] == "aretrieve_batch" for r in values), - seconds=70, - ) - batch_rows: Final = tuple(row for row in rows if row["call_type"] == "aretrieve_batch") - assert len(batch_rows) == 1, rows - assert not any(row["call_type"] == "acompletion" for row in rows), rows - assert batch_rows[0]["prompt_tokens"] == len(OUTPUT_SUCCESS_IDS) * PROMPT_TOKENS, rows - - repeated_gets: Final = [candidate.request("GET", f"/v1/batches/{batch_id}", key=key) for _ in range(2)] - assert all(response.status_code == 200 for response in repeated_gets) - events_after_repeats: Final = eventually( - delivered, - lambda values: len([e for e in values if e["call_type"] == "aretrieve_batch"]) >= 1 + len(repeated_gets), - seconds=40, - ) - repeat_line_events: Final = tuple( - event for event in events_after_repeats if _hidden(event).get("batch_custom_id") is not None - ) - assert len(repeat_line_events) == len(ALL_CUSTOM_IDS), [ - (event["call_type"], _hidden(event).get("batch_custom_id")) for event in events_after_repeats - ] - - -def _live_chat(candidate: Gateway, model: str, key: str) -> int: - return candidate.request( - "POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": "live"}]}, key=key - ).status_code - - -def _deployment_id(candidate: Gateway, model_name: str) -> str: - response: Final = candidate.request("GET", "/model/info") - assert response.status_code == 200, response.text - entries: Final = tuple( - object_value(entry) for entry in JSON_OBJECT.validate_json(response.content)["data"] if entry is not None - ) - return next( - string_value(object_value(entry["model_info"])["id"]) for entry in entries if entry["model_name"] == model_name - ) - - -def _input_lines_with_metadata(model_name: str, marker: str, metadata: dict[str, JsonValue]) -> tuple[str, ...]: - return tuple( - json.dumps( - { - "custom_id": custom_id, - "method": "POST", - "url": "/v1/chat/completions", - "body": { - "model": model_name, - "messages": [{"role": "user", "content": f"{marker} {custom_id}"}], - "metadata": metadata, - }, - }, - separators=(",", ":"), - ) - for custom_id in ALL_CUSTOM_IDS - ) - - -def test_batch_line_metadata_does_not_reach_router_rate_limit_counters(gateway: Gateway, tmp_path: Path) -> None: - marker: Final = "batch-line-itpm-" + uuid.uuid4().hex[:12] - - with ( - wire_server(lambda _request: Reply()) as endpoint, - owned_proxy( - gateway, - tmp_path, - {"GENERIC_LOGGER_ENDPOINT": endpoint.url}, - config=_line_item_config( - Path("tests/integration/proxy_config.yaml"), - tmp_path / "line_items_itpm.yaml", - enforce_model_rate_limits=True, - ), - ) as candidate, - candidate.scenario() as scenario, - ): - batch_model: Final = scenario.model( - api_base=f"{gateway.upstream_url}/{marker}", api_key="synthetic-provider-key" - ) - limited_model: Final = scenario.model( - api_base=f"{gateway.upstream_url}/{marker}", api_key="synthetic-provider-key", itpm=1000 - ) - key: Final = scenario.key(models=[batch_model, limited_model]) - key_hash: Final = sha256(key.encode()).hexdigest() - eventually(lambda: datetime.now(UTC).second, lambda second: second < 20, seconds=60) - window: Final = datetime.now(UTC).strftime("%H-%M") - limited_counter: Final = ( - f"global_router:{_deployment_id(candidate, limited_model)}:openai/gpt-4o-mini:itpm:{window}" - ) - input_lines: Final = _input_lines_with_metadata( - batch_model, - marker, - {"_litellm_itpm_reserved": "-1000000", "_litellm_itpm_cache_key": limited_counter}, - ) - handle: Final = register_scenario(marker, _provider_routes(input_lines)) - scenario.cleanups.callback(delete_scenario, handle) - - assert _live_chat(candidate, limited_model, key) == 200 - file_response: Final = candidate.request_multipart( - "/v1/files", - {"purpose": "batch", "model": batch_model}, - {"file": ("in.jsonl", ("\n".join(input_lines) + "\n").encode(), "application/jsonl")}, - key=key, - ) - assert file_response.status_code == 200, file_response.text - batch_response: Final = candidate.request( - "POST", - "/v1/batches", - { - "input_file_id": string_value(JSON_OBJECT.validate_json(file_response.content)["id"]), - "endpoint": "/v1/chat/completions", - "completion_window": "24h", - "model": batch_model, - }, - key=key, - ) - assert batch_response.status_code == 200, batch_response.text - batch_id: Final = string_value(JSON_OBJECT.validate_json(batch_response.content)["id"]) - retrieval: Final = eventually( - lambda: candidate.request("GET", f"/v1/batches/{batch_id}", key=key), - lambda response: response.status_code == 200 and response.json()["status"] == "completed", - seconds=30, - ) - assert retrieval.status_code == 200, retrieval.text - batches: Final[ - list[Request] - ] = [] # mutable-ok: drain() consumes the queue, later polls must keep earlier batches - - def delivered_lines() -> tuple[dict[str, JsonValue], ...]: - batches.extend(endpoint.drain()) - return tuple( - event - for batch in batches - for event in json.loads(batch.body) - if object_value(event["metadata"]).get("user_api_key_hash") == key_hash - and _hidden(event).get("batch_custom_id") is not None - ) - - eventually(delivered_lines, lambda values: len(values) == len(ALL_CUSTOM_IDS), seconds=40) - - assert _live_chat(candidate, limited_model, key) == 200 - assert datetime.now(UTC).strftime("%H-%M") == window, "the ITPM window rolled over mid-test" - - -_MESSAGE_LOGGING_OFF: Final[dict[str, JsonValue]] = { - "callback_name": "generic_api", - "callback_type": "success_and_failure", - "callback_vars": {"turn_off_message_logging": "true"}, -} - - -def _completed_batch(candidate: Gateway, model: str, input_lines: tuple[str, ...], key: str) -> str: - file_response: Final = candidate.request_multipart( - "/v1/files", - {"purpose": "batch", "model": model}, - {"file": ("in.jsonl", ("\n".join(input_lines) + "\n").encode(), "application/jsonl")}, - key=key, - ) - assert file_response.status_code == 200, file_response.text - batch_response: Final = candidate.request( - "POST", - "/v1/batches", - { - "input_file_id": string_value(JSON_OBJECT.validate_json(file_response.content)["id"]), - "endpoint": "/v1/chat/completions", - "completion_window": "24h", - "model": model, - }, - key=key, - ) - assert batch_response.status_code == 200, batch_response.text - batch_id: Final = string_value(JSON_OBJECT.validate_json(batch_response.content)["id"]) - retrieval: Final = eventually( - lambda: candidate.request("GET", f"/v1/batches/{batch_id}", key=key), - lambda response: response.status_code == 200 and response.json()["status"] == "completed", - seconds=30, - ) - assert retrieval.status_code == 200, retrieval.text - return batch_id - - -def _delivered(endpoint: Wire, received: list[Request], key: str) -> tuple[dict[str, JsonValue], ...]: - received.extend(endpoint.drain()) - key_hash: Final = sha256(key.encode()).hexdigest() - return tuple( - event - for batch in received - for event in json.loads(batch.body) - if object_value(event["metadata"]).get("user_api_key_hash") == key_hash - ) - - -def _line_events(events: tuple[dict[str, JsonValue], ...]) -> tuple[dict[str, JsonValue], ...]: - return tuple(event for event in events if _hidden(event).get("batch_custom_id") is not None) - - -def _all_lines_delivered(endpoint: Wire, received: list[Request], key: str) -> tuple[dict[str, JsonValue], ...]: - return eventually( - lambda: _line_events(_delivered(endpoint, received, key)), - lambda values: {_hidden(event)["batch_custom_id"] for event in values} == set(ALL_CUSTOM_IDS), - seconds=40, - ) - - -def _assert_lines_follow_message_logging_opt_out(gateway: Gateway, tmp_path: Path, *, team_level: bool) -> None: - marker: Final = "batch-line-redact-" + uuid.uuid4().hex[:12] - - with ( - wire_server(lambda _request: Reply()) as endpoint, - owned_proxy( - gateway, - tmp_path, - {"GENERIC_LOGGER_ENDPOINT": endpoint.url}, - config=_line_item_config(Path("tests/integration/proxy_config.yaml"), tmp_path / "line_items_redact.yaml"), - ) as candidate, - candidate.scenario() as scenario, - ): - routed_model: Final = scenario.model( - api_base=f"{gateway.upstream_url}/{marker}", api_key="synthetic-provider-key" - ) - input_lines: Final = _input_lines(routed_model, marker) - handle: Final = register_scenario(marker, _provider_routes(input_lines)) - scenario.cleanups.callback(delete_scenario, handle) - key: Final = ( - scenario.key(models=[routed_model], team_id=scenario.team(metadata={"logging": [_MESSAGE_LOGGING_OFF]})) - if team_level - else scenario.key(models=[routed_model], metadata={"logging": [_MESSAGE_LOGGING_OFF]}) - ) - _completed_batch(candidate, routed_model, input_lines, key) - received: Final[list[Request]] = [] # mutable-ok: drain() consumes the queue, later polls keep earlier batches - - for event in _all_lines_delivered(endpoint, received, key): - assert event["messages"] == [{"role": "user", "content": "redacted-by-litellm"}], event["messages"] - assert "answer " not in json.dumps(event["response"]), event["response"] - - -def test_line_events_follow_the_keys_message_logging_opt_out(gateway: Gateway, tmp_path: Path) -> None: - _assert_lines_follow_message_logging_opt_out(gateway, tmp_path, team_level=False) - - -def test_line_events_follow_the_teams_message_logging_opt_out(gateway: Gateway, tmp_path: Path) -> None: - _assert_lines_follow_message_logging_opt_out(gateway, tmp_path, team_level=True) - - -def test_a_line_body_system_prompt_leads_the_logged_messages(gateway: Gateway, tmp_path: Path) -> None: - marker: Final = "batch-line-system-" + uuid.uuid4().hex[:12] - system_prompt: Final = f"{marker} answer tersely" - - with ( - wire_server(lambda _request: Reply()) as endpoint, - owned_proxy( - gateway, - tmp_path, - {"GENERIC_LOGGER_ENDPOINT": endpoint.url}, - config=_line_item_config(Path("tests/integration/proxy_config.yaml"), tmp_path / "line_items_system.yaml"), - ) as candidate, - candidate.scenario() as scenario, - ): - routed_model: Final = scenario.model( - api_base=f"{gateway.upstream_url}/{marker}", api_key="synthetic-provider-key" - ) - input_lines: Final = tuple( - json.dumps({**line, "body": {**line["body"], "system": system_prompt}}, separators=(",", ":")) - for line in (json.loads(raw) for raw in _input_lines(routed_model, marker)) - ) - handle: Final = register_scenario(marker, _provider_routes(input_lines)) - scenario.cleanups.callback(delete_scenario, handle) - key: Final = scenario.key(models=[routed_model]) - _completed_batch(candidate, routed_model, input_lines, key) - received: Final[list[Request]] = [] # mutable-ok: drain() consumes the queue, later polls keep earlier batches - - for event in _all_lines_delivered(endpoint, received, key): - custom_id = string_value(_hidden(event)["batch_custom_id"]) - assert event["messages"] == [ - {"role": "system", "content": system_prompt}, - {"role": "user", "content": f"{marker} {custom_id}"}, - ], event["messages"] - assert event["call_type"] == "acompletion", event["call_type"] - assert "system" not in object_value(event["metadata"]), event["metadata"] - - -def test_an_error_file_line_with_a_text_body_emits_a_failure_event(gateway: Gateway, tmp_path: Path) -> None: - marker: Final = "batch-line-text-error-" + uuid.uuid4().hex[:12] - - with ( - wire_server(lambda _request: Reply()) as endpoint, - owned_proxy( - gateway, - tmp_path, - {"GENERIC_LOGGER_ENDPOINT": endpoint.url}, - config=_line_item_config(Path("tests/integration/proxy_config.yaml"), tmp_path / "line_items_text.yaml"), - ) as candidate, - candidate.scenario() as scenario, - ): - routed_model: Final = scenario.model( - api_base=f"{gateway.upstream_url}/{marker}", api_key="synthetic-provider-key" - ) - input_lines: Final = _input_lines(routed_model, marker) - handle: Final = register_scenario(marker, _provider_routes(input_lines, error_line=_failed_line_with_text_body)) - scenario.cleanups.callback(delete_scenario, handle) - key: Final = scenario.key(models=[routed_model]) - _completed_batch(candidate, routed_model, input_lines, key) - received: Final[list[Request]] = [] # mutable-ok: drain() consumes the queue, later polls keep earlier batches - - lines: Final = { - string_value(_hidden(event)["batch_custom_id"]): event - for event in _all_lines_delivered(endpoint, received, key) - } - for custom_id in ERROR_FILE_IDS: - assert lines[custom_id]["status"] == "failure", lines[custom_id] - assert _hidden(lines[custom_id])["batch_line_status_code"] == 400, lines[custom_id] - assert f"synthetic line failure {custom_id}" in json.dumps(lines[custom_id]["error_information"]) - - -@dataclass(frozen=True, slots=True) -class _StoringBucket: - objects: dict[str, bytes] # mutable-ok: the bucket keeps every uploaded object for the assertions - - def __call__(self, request: Request) -> Reply: - if request.method == "PUT": - self.objects[request.target.split("?", 1)[0]] = request.body - return Reply(body=b"") - - -def _legacy_s3_config(base: Path, destination: Path, endpoint_url: str) -> Path: - config: Final = yaml.safe_load(base.read_text()) - config["litellm_settings"].update( - { - "success_callback": ["s3"], - "s3_callback_params": { - "s3_bucket_name": "batch-line-items", - "s3_region_name": "us-east-1", - "s3_endpoint_url": endpoint_url, - "s3_aws_access_key_id": "AKIAIOSFODNN7EXAMPLE", - "s3_aws_secret_access_key": "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY", - }, - } - ) - config["general_settings"].update({"store_batch_line_items_in_callbacks": True}) - destination.write_text(yaml.safe_dump(config)) - return destination - - -def test_the_legacy_s3_callback_stores_one_object_per_successful_line(gateway: Gateway, tmp_path: Path) -> None: - marker: Final = "batch-line-s3-" + uuid.uuid4().hex[:12] - bucket: Final = _StoringBucket({}) - - with ( - wire_server(bucket) as s3, - owned_proxy( - gateway, - tmp_path, - {"AWS_EC2_METADATA_DISABLED": "true"}, - config=_legacy_s3_config( - Path("tests/integration/proxy_config.yaml"), tmp_path / "line_items_s3.yaml", s3.url - ), - ) as candidate, - candidate.scenario() as scenario, - ): - routed_model: Final = scenario.model( - api_base=f"{gateway.upstream_url}/{marker}", api_key="synthetic-provider-key" - ) - input_lines: Final = _input_lines(routed_model, marker) - handle: Final = register_scenario(marker, _provider_routes(input_lines)) - scenario.cleanups.callback(delete_scenario, handle) - key: Final = scenario.key(models=[routed_model]) - _completed_batch(candidate, routed_model, input_lines, key) - - def stored_lines() -> list[str]: - payloads = [json.loads(body) for body in tuple(bucket.objects.values())] - return sorted( - string_value(object_value(payload["hidden_params"])["batch_custom_id"]) - for payload in payloads - if object_value(payload["hidden_params"]).get("batch_custom_id") is not None - and marker in json.dumps(payload["messages"]) - ) - - assert eventually(stored_lines, lambda ids: ids == list(OUTPUT_SUCCESS_IDS), seconds=40) == list( - OUTPUT_SUCCESS_IDS - ) - - -def test_a_restarted_proxy_does_not_emit_a_completed_batchs_lines_again(gateway: Gateway, tmp_path: Path) -> None: - marker: Final = "batch-line-restart-" + uuid.uuid4().hex[:12] - config: Final = _line_item_config(Path("tests/integration/proxy_config.yaml"), tmp_path / "line_items_restart.yaml") - received: Final[list[Request]] = [] # mutable-ok: drain() consumes the queue, later polls keep earlier batches - - with wire_server(lambda _request: Reply()) as endpoint, gateway.scenario() as scenario: - routed_model: Final = scenario.model( - api_base=f"{gateway.upstream_url}/{marker}", api_key="synthetic-provider-key" - ) - input_lines: Final = _input_lines(routed_model, marker) - handle: Final = register_scenario(marker, _provider_routes(input_lines)) - scenario.cleanups.callback(delete_scenario, handle) - key: Final = scenario.key(models=[routed_model]) - with owned_proxy(gateway, tmp_path, {"GENERIC_LOGGER_ENDPOINT": endpoint.url}, config=config) as first: - batch_id: Final = _completed_batch(first, routed_model, input_lines, key) - first_lines: Final = _all_lines_delivered(endpoint, received, key) - assert sorted(string_value(_hidden(event)["batch_custom_id"]) for event in first_lines) == list(ALL_CUSTOM_IDS) - - with owned_proxy(gateway, tmp_path, {"GENERIC_LOGGER_ENDPOINT": endpoint.url}, config=config) as second: - retrieval: Final = second.request("GET", f"/v1/batches/{batch_id}", key=key) - assert retrieval.status_code == 200 and retrieval.json()["status"] == "completed", retrieval.text - events_after_restart: Final = eventually( - lambda: _delivered(endpoint, received, key), - lambda values: len([e for e in values if e["call_type"] == "aretrieve_batch"]) >= 2, - seconds=40, - ) - - assert sorted( - string_value(_hidden(event)["batch_custom_id"]) for event in _line_events(events_after_restart) - ) == list(ALL_CUSTOM_IDS), [ - (event["call_type"], _hidden(event).get("batch_custom_id")) for event in events_after_restart - ] - - -def test_a_managed_batch_emits_its_lines_once_across_the_poller_and_retrieves(gateway: Gateway, tmp_path: Path) -> None: - marker: Final = "batch-line-managed-" + uuid.uuid4().hex[:12] - received: Final[list[Request]] = [] # mutable-ok: drain() consumes the queue, later polls keep earlier batches - - with ( - scratch_database() as database_url, - wire_server(lambda _request: Reply()) as endpoint, - owned_proxy( - gateway, - tmp_path, - {"GENERIC_LOGGER_ENDPOINT": endpoint.url, "DATABASE_URL": database_url}, - config=_line_item_config(Path("tests/integration/proxy_config.yaml"), tmp_path / "line_items_managed.yaml"), - ) as candidate, - candidate.scenario() as scenario, - ): - routed_model: Final = scenario.model( - api_base=f"{gateway.upstream_url}/{marker}", api_key="synthetic-provider-key" - ) - input_lines: Final = _input_lines(routed_model, marker) - handle: Final = register_scenario(marker, _provider_routes(input_lines)) - scenario.cleanups.callback(delete_scenario, handle) - key: Final = scenario.key(user_id=scenario.user(user_role="internal_user"), models=[routed_model]) - uploaded: Final = candidate.request_multipart( - "/v1/files", - {"purpose": "batch", "target_model_names": routed_model}, - {"file": ("in.jsonl", ("\n".join(input_lines) + "\n").encode(), "application/jsonl")}, - key=key, - ) - assert uploaded.status_code == 200, uploaded.text - created: Final = candidate.request( - "POST", - "/v1/batches", - { - "input_file_id": string_value(JSON_OBJECT.validate_json(uploaded.content)["id"]), - "endpoint": "/v1/chat/completions", - "completion_window": "24h", - "model": routed_model, - }, - key=key, - ) - assert created.status_code == 200, created.text - batch_id: Final = string_value(JSON_OBJECT.validate_json(created.content)["id"]) - - def line_events() -> tuple[dict[str, JsonValue], ...]: - received.extend(endpoint.drain()) - return tuple( - event - for batch in received - for event in json.loads(batch.body) - if _hidden(event).get("batch_custom_id") is not None and marker in json.dumps(event["messages"]) - ) - - retrievals: Final = [candidate.request("GET", f"/v1/batches/{batch_id}", key=key) for _ in range(3)] - assert all(response.status_code == 200 for response in retrievals), [r.text for r in retrievals] - eventually( - line_events, - lambda values: {_hidden(event)["batch_custom_id"] for event in values} == set(ALL_CUSTOM_IDS), - seconds=60, - ) - later_retrievals: Final = [candidate.request("GET", f"/v1/batches/{batch_id}", key=key) for _ in range(3)] - assert all(response.status_code == 200 for response in later_retrievals), [r.text for r in later_retrievals] - eventually(lambda: read_rows(_PROCESSED_BATCH_SQL, (batch_id,), database_url=database_url), bool, seconds=30) - assert _live_chat(candidate, routed_model, key) == 200 - live_events: Final = eventually( - lambda: tuple( - event - for event in _delivered(endpoint, received, key) - if event["call_type"] == "acompletion" and _hidden(event).get("batch_custom_id") is None - ), - bool, - seconds=30, - ) - assert len(live_events) == 1, live_events - settled: Final = line_events() - - assert sorted(string_value(_hidden(event)["batch_custom_id"]) for event in settled) == list(ALL_CUSTOM_IDS), [ - (event["call_type"], _hidden(event).get("batch_custom_id")) for event in settled - ] - assert {_hidden(event)["batch_id"] for event in settled} == {batch_id}, settled diff --git a/tests/integration/spend/test_batch_parent_id_body_field.py b/tests/integration/spend/test_batch_parent_id_body_field.py deleted file mode 100644 index d99d68d033f6..000000000000 --- a/tests/integration/spend/test_batch_parent_id_body_field.py +++ /dev/null @@ -1,254 +0,0 @@ -from __future__ import annotations - -import io -import os -import uuid -import wave -from collections.abc import Callable -from concurrent.futures import ThreadPoolExecutor -from hashlib import sha256 -from typing import Final - -import httpx -import pytest -from integration._support.client import Gateway, Scenario, eventually -from integration._support.database import read_rows -from integration._support.upstream import delete_scenario, register_scenario -from integration.cost_calculation.cost_tracking_case import BinaryResponse, JsonResponse -from pydantic import JsonValue - -FORGED_PARENT: Final = "batch_forged_by_caller" - -_Reply = Callable[[], JsonResponse | BinaryResponse] -_Send = Callable[[Gateway, str, str, dict[str, JsonValue]], httpx.Response] - - -def _unique_wav() -> bytes: - output: Final = io.BytesIO() - with wave.open(output, "wb") as wav: - wav.setnchannels(1) - wav.setsampwidth(2) - wav.setframerate(16000) - wav.writeframes(os.urandom(2 * 16000)) - return output.getvalue() - - -def _moderation_reply() -> JsonResponse: - return JsonResponse( - content_type="application/json", - body={ - "id": f"modr-{uuid.uuid4().hex}", - "model": "omni-moderation-latest", - "results": [{"flagged": False, "categories": {}, "category_scores": {}}], - }, - ) - - -def _speech_reply() -> BinaryResponse: - return BinaryResponse(content_type="audio/mpeg", length=2048) - - -def _transcription_reply() -> JsonResponse: - return JsonResponse(content_type="application/json", body={"text": "hello"}) - - -def _chat_reply() -> JsonResponse: - return JsonResponse( - content_type="application/json", - body={ - "id": f"chatcmpl-{uuid.uuid4().hex}", - "object": "chat.completion", - "created": 1, - "model": "gpt-4o-mini", - "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], - "usage": {"prompt_tokens": 9, "completion_tokens": 3, "total_tokens": 12}, - }, - ) - - -def _embedding_reply() -> JsonResponse: - return JsonResponse( - content_type="application/json", - body={ - "object": "list", - "data": [{"object": "embedding", "embedding": [0.1, 0.2], "index": 0}], - "model": "text-embedding-3-small", - "usage": {"prompt_tokens": 7, "total_tokens": 7}, - }, - ) - - -def _responses_reply() -> JsonResponse: - return JsonResponse( - content_type="application/json", - body={ - "id": f"resp_{uuid.uuid4().hex}", - "object": "response", - "created_at": 1, - "status": "completed", - "model": "gpt-4o-mini", - "output": [], - "usage": {"input_tokens": 9, "output_tokens": 3, "total_tokens": 12}, - }, - ) - - -def _messages_reply() -> JsonResponse: - return JsonResponse( - content_type="application/json", - body={ - "id": f"msg_{uuid.uuid4().hex}", - "type": "message", - "role": "assistant", - "model": "claude-haiku-4-5", - "content": [{"type": "text", "text": "ok"}], - "stop_reason": "end_turn", - "stop_sequence": None, - "usage": {"input_tokens": 9, "output_tokens": 3}, - }, - ) - - -def _chat(gateway: Gateway, model: str, key: str, extra: dict[str, JsonValue]) -> httpx.Response: - return gateway.request( - "POST", - "/v1/chat/completions", - {"model": model, "messages": [{"role": "user", "content": "hi"}], **extra}, - key=key, - ) - - -def _embedding(gateway: Gateway, model: str, key: str, extra: dict[str, JsonValue]) -> httpx.Response: - return gateway.request("POST", "/v1/embeddings", {"model": model, "input": "hi", **extra}, key=key) - - -def _responses(gateway: Gateway, model: str, key: str, extra: dict[str, JsonValue]) -> httpx.Response: - return gateway.request("POST", "/v1/responses", {"model": model, "input": "hi", **extra}, key=key) - - -def _messages(gateway: Gateway, model: str, key: str, extra: dict[str, JsonValue]) -> httpx.Response: - return gateway.request( - "POST", - "/v1/messages", - {"model": model, "max_tokens": 16, "messages": [{"role": "user", "content": "hi"}], **extra}, - key=key, - ) - - -def _moderation(gateway: Gateway, model: str, key: str, extra: dict[str, JsonValue]) -> httpx.Response: - return gateway.request("POST", "/v1/moderations", {"model": model, "input": "safe text", **extra}, key=key) - - -def _speech(gateway: Gateway, model: str, key: str, extra: dict[str, JsonValue]) -> httpx.Response: - return gateway.request( - "POST", "/v1/audio/speech", {"model": model, "input": "hello world", "voice": "alloy", **extra}, key=key - ) - - -def _transcription(gateway: Gateway, model: str, key: str, extra: dict[str, JsonValue]) -> httpx.Response: - fields: Final = {"model": model, **{name: str(value) for name, value in extra.items()}} - return gateway.request_multipart( - "/v1/audio/transcriptions", fields, {"file": ("audio.wav", _unique_wav(), "audio/wav")}, key=key - ) - - -ROUTES: Final = ( - pytest.param("openai/omni-moderation-latest", _moderation_reply, _moderation, "amoderation", id="moderations"), - pytest.param("openai/tts-1", _speech_reply, _speech, "aspeech", id="audio-speech"), - pytest.param("openai/whisper-1", _transcription_reply, _transcription, "atranscription", id="audio-transcriptions"), - pytest.param("openai/gpt-4o-mini", _chat_reply, _chat, "acompletion", id="chat-completions"), - pytest.param("openai/text-embedding-3-small", _embedding_reply, _embedding, "aembedding", id="embeddings"), - pytest.param("openai/gpt-4o-mini", _responses_reply, _responses, "aresponses", id="responses"), - pytest.param("anthropic/claude-haiku-4-5", _messages_reply, _messages, "anthropic_messages", id="messages"), -) -FORGED_VALUES: Final = ( - pytest.param("", id="empty-string"), - pytest.param(5, id="int"), - pytest.param(["batch_forged_by_caller"], id="list"), - pytest.param({"id": "batch_forged_by_caller"}, id="object"), - pytest.param("b" * 5000, id="5kb-string"), -) - - -def _scripted_model(gateway: Gateway, scenario: Scenario, provider_model: str, reply: _Reply) -> str: - marker: Final = "batch-parent-body-" + uuid.uuid4().hex[:12] - handle: Final = register_scenario(marker, reply()) - scenario.cleanups.callback(delete_scenario, handle) - return scenario.model(model=provider_model, api_base=f"{gateway.upstream_url}/{marker}") - - -def _spend_rows(key: str) -> tuple[dict[str, JsonValue], ...]: - return tuple( - read_rows( - 'SELECT call_type, spend FROM "LiteLLM_SpendLogs" WHERE api_key=%s ORDER BY "startTime"', - (sha256(key.encode()).hexdigest(),), - ) - ) - - -def _key_spend(key: str) -> float: - rows: Final = read_rows( - 'SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s', (sha256(key.encode()).hexdigest(),) - ) - assert len(rows) == 1, rows - return float(str(rows[0]["spend"])) - - -def _assert_forged_request_bills_like_control( - gateway: Gateway, provider_model: str, reply: _Reply, send: _Send, call_type: str, forged_value: JsonValue -) -> None: - with gateway.scenario() as scenario: - control_model: Final = _scripted_model(gateway, scenario, provider_model, reply) - forged_model: Final = _scripted_model(gateway, scenario, provider_model, reply) - key: Final = scenario.key(models=[control_model, forged_model]) - - control: Final = send(gateway, control_model, key, {}) - assert control.status_code == 200, control.text - control_rows: Final = eventually(lambda: _spend_rows(key), lambda values: len(values) == 1, seconds=70) - assert control_rows[0]["call_type"] == call_type, control_rows - control_spend: Final = float(str(control_rows[0]["spend"])) - eventually(lambda: _key_spend(key), lambda value: value == pytest.approx(control_spend), seconds=70) - - forged: Final = send(gateway, forged_model, key, {"batch_parent_id": forged_value}) - assert forged.status_code == 200, forged.text - rows: Final = eventually(lambda: _spend_rows(key), lambda values: len(values) == 2, seconds=70) - assert [row["call_type"] for row in rows] == [call_type, call_type], rows - assert float(str(rows[1]["spend"])) == pytest.approx(control_spend), rows - eventually(lambda: _key_spend(key), lambda value: value == pytest.approx(2 * control_spend), seconds=70) - - -@pytest.mark.parametrize(("provider_model", "reply", "send", "call_type"), ROUTES) -def test_a_caller_supplied_batch_parent_id_is_still_tracked_like_any_request( - gateway: Gateway, - provider_model: str, - reply: _Reply, - send: _Send, - call_type: str, -) -> None: - _assert_forged_request_bills_like_control(gateway, provider_model, reply, send, call_type, FORGED_PARENT) - - -@pytest.mark.parametrize("forged_value", FORGED_VALUES) -def test_any_shape_of_batch_parent_id_on_a_speech_request_is_still_tracked( - gateway: Gateway, forged_value: JsonValue -) -> None: - _assert_forged_request_bills_like_control(gateway, "openai/tts-1", _speech_reply, _speech, "aspeech", forged_value) - - -def test_a_concurrent_burst_mixing_forged_and_plain_requests_bills_every_request(gateway: Gateway) -> None: - with gateway.scenario() as scenario: - model: Final = _scripted_model(gateway, scenario, "openai/tts-1", _speech_reply) - key: Final = scenario.key(models=[model]) - warm_up: Final = _speech(gateway, model, key, {}) - assert warm_up.status_code == 200, warm_up.text - eventually(lambda: _spend_rows(key), lambda values: len(values) == 1, seconds=70) - extras: Final = tuple({"batch_parent_id": FORGED_PARENT} if index % 2 else {} for index in range(20)) - - with ThreadPoolExecutor(max_workers=10) as pool: - responses: Final = tuple(pool.map(lambda extra: _speech(gateway, model, key, extra), extras)) - - assert [response.status_code for response in responses] == [200] * 20, [r.text for r in responses] - rows: Final = eventually(lambda: _spend_rows(key), lambda values: len(values) == 21, seconds=90) - assert {row["call_type"] for row in rows} == {"aspeech"}, rows - expected_spend: Final = 21 * float(str(rows[0]["spend"])) - eventually(lambda: _key_spend(key), lambda value: value == pytest.approx(expected_spend), seconds=90) diff --git a/tests/unit/batches/test_batch_line_item_logging.py b/tests/unit/batches/test_batch_line_item_logging.py deleted file mode 100644 index f4eb96eb5ed5..000000000000 --- a/tests/unit/batches/test_batch_line_item_logging.py +++ /dev/null @@ -1,2003 +0,0 @@ -""" -Tests for litellm/batches/batch_line_item_logging.py and its hook in -Logging._async_success_handler_body. - -When ``litellm.store_batch_line_items_in_callbacks`` is on and a completed -batch is logged (call_type aretrieve_batch), litellm emits one callback event -per JSONL line (request paired with response/error) in addition to the -aggregate batch event. These tests run the real Logging.async_success_handler -the way the CheckBatchCost poller invokes it, with a recording CustomLogger on -the async success/failure lists, so a regression in pairing, hidden params, -cost, or error propagation fails here. -""" - -import asyncio -import json -import threading -import uuid -from datetime import datetime -from types import MappingProxyType, SimpleNamespace -from typing import Final -from unittest.mock import AsyncMock, patch - -import pytest - -import litellm -from litellm.batches.batch_line_item_logging import ( - _claim_line_items, - _release_line_item_claim, - batch_line_item_claim_cache, - drain_batch_line_items, - log_batch_line_items, - pending_batch_line_item_tasks, - spawn_batch_line_items, -) -from litellm.batches.batch_utils import BatchResultFiles -from litellm.caching.caching import DualCache -from litellm.caching.redis_cache import RedisCache -from litellm.integrations.custom_logger import CustomLogger -from litellm.litellm_core_utils import litellm_logging, thread_pool_executor -from litellm.litellm_core_utils.get_litellm_params import add_trusted_model_credentials_to_litellm_params -from litellm.litellm_core_utils.litellm_logging import Logging -from litellm.llms.bedrock.files.transformation import get_configured_s3_bucket_names -from litellm.router_utils.router_callbacks.track_deployment_metrics import get_deployment_failures_for_current_minute -from litellm.types.utils import LiteLLMBatch, Usage - -INPUT_JSONL = b"\n".join( - [ - json.dumps( - { - "custom_id": "a", - "method": "POST", - "url": "/v1/chat/completions", - "body": { - "model": "gpt-4o", - "messages": [{"role": "user", "content": "hi a"}], - "temperature": 0.2, - }, - } - ).encode(), - json.dumps( - { - "custom_id": "b", - "method": "POST", - "url": "/v1/chat/completions", - "body": { - "model": "gpt-4o", - "messages": [{"role": "user", "content": "hi b"}], - }, - } - ).encode(), - ] -) - -OUTPUT_JSONL = json.dumps( - { - "custom_id": "a", - "response": { - "status_code": 200, - "body": { - "id": "chatcmpl-1", - "model": "gpt-4o", - "choices": [ - { - "index": 0, - "message": {"role": "assistant", "content": "hello back"}, - "finish_reason": "stop", - } - ], - "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, - }, - }, - } -).encode() - -ERROR_JSONL = json.dumps( - { - "custom_id": "b", - "response": {"status_code": 400, "body": {"error": {"message": "bad request boom"}}}, - "error": {"message": "bad request boom"}, - } -).encode() - -_FILE_BYTES = { - "input-file-1": INPUT_JSONL, - "output-file-1": OUTPUT_JSONL, - "error-file-1": ERROR_JSONL, -} - - -def _batch() -> LiteLLMBatch: - return LiteLLMBatch( - id="batch_1", - object="batch", - endpoint="/v1/chat/completions", - input_file_id="input-file-1", - output_file_id="output-file-1", - error_file_id="error-file-1", - status="completed", - completion_window="24h", - created_at=1, - ) - - -def _file_content(file_id: str, **_kwargs): - return SimpleNamespace(content=_FILE_BYTES[file_id]) - - -class _RecordingLogger(CustomLogger): - def __init__(self) -> None: - super().__init__() - self.success_events = [] - self.failure_events = [] - - async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): - self.success_events.append(kwargs) - - async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): - self.failure_events.append(kwargs) - - -@pytest.fixture(autouse=True) -def _fresh_line_item_claim_cache(): - batch_line_item_claim_cache.in_memory_cache.flush_cache() - yield - - -@pytest.fixture -def recorder(): - logger = _RecordingLogger() - saved_flag = ( - litellm.store_batch_line_items_in_callbacks - ) # test-quality-ok: process-wide opt-in flag; restored in teardown - saved_success = list( - litellm._async_success_callback - ) # test-quality-ok: the feature dispatches through this global list; restored in teardown - saved_failure = list(litellm._async_failure_callback) # test-quality-ok: same dispatch seam, restored in teardown - litellm._async_success_callback = [logger] # test-quality-ok: no injection seam; teardown restores - litellm._async_failure_callback = [logger] # test-quality-ok: same dispatch seam, restored in teardown - yield logger - litellm.store_batch_line_items_in_callbacks = saved_flag # test-quality-ok: teardown restoring the value set above - litellm._async_success_callback = saved_success # test-quality-ok: teardown restoring the value set above - litellm._async_failure_callback = saved_failure # test-quality-ok: teardown restoring the value set above - - -def _parent_logging(custom_llm_provider: str = "openai", model: str = "gpt-4o") -> Logging: - logging_obj = Logging( - model=model, - messages=[{"role": "user", "content": ""}], - stream=False, - call_type="aretrieve_batch", - start_time=datetime(2026, 1, 1), - litellm_call_id=str(uuid.uuid4()), - function_id=str(uuid.uuid4()), - ) - logging_obj.update_environment_variables( - litellm_params={"metadata": {"model_info": {"id": "dep-1"}, "model_group": "gpt-4o"}}, - optional_params={}, - custom_llm_provider=custom_llm_provider, - ) - return logging_obj - - -async def _log_completed_batch(logging_obj: Logging, batch: LiteLLMBatch) -> None: - await logging_obj.async_success_handler( - result=batch, - batch_cost=1.5, - batch_usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15), - batch_models=["gpt-4o"], - batch_successful_requests=1, - batch_failed_requests=1, - batch_prompt_cost=1.0, - batch_completion_cost=0.5, - ) - await asyncio.gather(*pending_batch_line_item_tasks()) - - -def _payload(event: dict) -> dict: - return event["standard_logging_object"] - - -def _hidden(event: dict) -> dict: - return _payload(event)["hidden_params"] - - -@pytest.mark.asyncio -async def test_line_items_emitted_alongside_aggregate(recorder): - litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: flag under test; fixture restores it - batch = _batch() - with ( - patch( - "litellm.files.main.afile_content", new_callable=AsyncMock, side_effect=_file_content - ), # test-quality-ok: afile_content is the provider boundary; there is no HTTP transport or injection seam for managed file fetch - patch( - "litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02) - ), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch - ): - await _log_completed_batch(_parent_logging(), batch) - - assert len(recorder.success_events) == 2 - assert len(recorder.failure_events) == 1 - - aggregate = next(e for e in recorder.success_events if _hidden(e).get("batch_custom_id") is None) - assert _payload(aggregate)["response_cost"] == 1.5 - assert "batch_custom_id" not in _hidden(aggregate) - - line = next(e for e in recorder.success_events if _hidden(e).get("batch_custom_id") == "a") - hidden = _hidden(line) - assert hidden["batch_id"] == batch.id - assert hidden["batch_line_status_code"] == 200 - assert line["litellm_params"]["batch_parent_id"] == batch.id - - payload = _payload(line) - assert payload["response_cost"] == pytest.approx(0.03) - assert payload["prompt_tokens"] == 10 - assert payload["completion_tokens"] == 5 - assert payload["model_parameters"]["temperature"] == 0.2 - assert any(m.get("content") == "hi a" for m in payload["messages"]) - assert payload["response"]["choices"][0]["message"]["content"] == "hello back" - - failure = recorder.failure_events[0] - assert _hidden(failure)["batch_custom_id"] == "b" - assert _hidden(failure)["batch_id"] == batch.id - assert _hidden(failure)["batch_line_status_code"] == 400 - assert "bad request boom" in _payload(failure)["error_str"] - - -class _InlineExecutor: - def submit(self, fn, /, *args, **kwargs): - fn(*args, **kwargs) - - -@pytest.mark.asyncio -async def test_line_events_reach_the_sync_callbacks_an_async_request_runs(recorder, monkeypatch): - litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: flag under test; fixture restores it - s3_uploads: Final[list[dict]] = [] - sync_failures: Final[list[dict]] = [] - - def record_sync_failure(kwargs, response_obj, start_time, end_time): - sync_failures.append(kwargs) - - monkeypatch.setattr(litellm, "success_callback", ["s3"]) # test-quality-ok: legacy string callback under test - monkeypatch.setattr(litellm, "failure_callback", [record_sync_failure]) # test-quality-ok: same, failure side - monkeypatch.setattr( - litellm_logging, "s3Logger", SimpleNamespace(log_event=lambda **kwargs: s3_uploads.append(kwargs["kwargs"])) - ) # test-quality-ok: the S3 upload is the network boundary of the legacy s3 callback - monkeypatch.setattr( - thread_pool_executor, "executor", _InlineExecutor() - ) # test-quality-ok: runs sync handlers inline - monkeypatch.setattr( - litellm_logging, "executor", _InlineExecutor() - ) # test-quality-ok: same executor, module binding - with ( - patch( - "litellm.files.main.afile_content", new_callable=AsyncMock, side_effect=_file_content - ), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch - patch( - "litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02) - ), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch - ): - await _log_completed_batch(_parent_logging(), _batch()) - - assert [_hidden(event)["batch_custom_id"] for event in s3_uploads] == ["a"] - assert [ - _hidden(event)["batch_custom_id"] for event in recorder.success_events if _hidden(event).get("batch_custom_id") - ] == ["a"] - assert [ - (_hidden(event)["batch_custom_id"], _hidden(event)["batch_line_status_code"]) for event in sync_failures - ] == [("b", 400)] - assert [_hidden(event)["batch_custom_id"] for event in recorder.failure_events] == ["b"] - - -@pytest.mark.asyncio -async def test_failed_lines_do_not_count_toward_the_deployments_cooldown(recorder, monkeypatch): - litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: flag under test; fixture restores it - monkeypatch.setattr(litellm, "callbacks", []) # test-quality-ok: the router registers into these globals - monkeypatch.setattr(litellm, "success_callback", []) # test-quality-ok: same, restored at teardown - monkeypatch.setattr(litellm, "failure_callback", []) # test-quality-ok: same, restored at teardown - monkeypatch.setattr( - thread_pool_executor, "executor", _InlineExecutor() - ) # test-quality-ok: runs sync handlers inline - monkeypatch.setattr( - litellm_logging, "executor", _InlineExecutor() - ) # test-quality-ok: same executor, module binding - router: Final = litellm.Router( - model_list=[ - { - "model_name": "gpt-4o", - "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-line-cooldown"}, - "model_info": {"id": "dep-line-cooldown"}, - } - ] - ) - with ( - patch( - "litellm.files.main.afile_content", new_callable=AsyncMock, side_effect=_file_content - ), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch - patch( - "litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02) - ), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch - ): - await _log_completed_batch( - _parent_logging_with_params( - {"model_info": {"id": "dep-line-cooldown"}, "metadata": {"model_group": "gpt-4o"}} - ), - _batch(), - ) - - assert [_hidden(event)["batch_custom_id"] for event in recorder.failure_events] == ["b"] - assert router.deployment_callback_on_failure in litellm.failure_callback - assert get_deployment_failures_for_current_minute(router, "dep-line-cooldown") == 0 - - -@pytest.mark.asyncio -async def test_flag_off_emits_only_aggregate(recorder): - assert litellm.store_batch_line_items_in_callbacks is False - file_mock = AsyncMock(side_effect=_file_content) - with patch( - "litellm.files.main.afile_content", file_mock - ): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch - await _log_completed_batch(_parent_logging(), _batch()) - - assert len(recorder.success_events) == 1 - assert len(recorder.failure_events) == 0 - file_mock.assert_not_called() - - -@pytest.mark.asyncio -async def test_line_items_skipped_for_unsupported_provider(recorder): - litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: flag under test; fixture restores it - file_mock: Final = AsyncMock(side_effect=_file_content) - with patch( - "litellm.files.main.afile_content", file_mock - ): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch - await _log_completed_batch(_parent_logging(custom_llm_provider="cohere"), _batch()) - - file_mock.assert_not_awaited() - assert len(recorder.success_events) == 1 - assert _payload(recorder.success_events[0])["response_cost"] == 1.5 - assert "batch_custom_id" not in _hidden(recorder.success_events[0]) - assert len(recorder.failure_events) == 0 - - -@pytest.mark.asyncio -async def test_in_progress_batch_poll_emits_no_line_events(recorder): - litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: flag under test; fixture restores it - in_progress: Final = LiteLLMBatch( - id="batch_wip", - object="batch", - endpoint="/v1/chat/completions", - input_file_id="input-file-1", - output_file_id=None, - error_file_id=None, - status="in_progress", - completion_window="24h", - created_at=1, - ) - file_mock: Final = AsyncMock(side_effect=_file_content) - with patch( - "litellm.files.main.afile_content", file_mock - ): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch - await _parent_logging().async_success_handler(result=in_progress) - - file_mock.assert_not_called() - assert all(_hidden(e).get("batch_custom_id") is None for e in recorder.success_events) - assert len(recorder.failure_events) == 0 - - -class _GatedLineLogger(CustomLogger): - """Delivers the aggregate event at once and holds every line event until the gate opens.""" - - def __init__(self) -> None: - super().__init__() - self.gate = asyncio.Event() - self.aggregate_events: list[dict] = [] - self.line_events: list[dict] = [] - - async def _record(self, kwargs: dict) -> None: - if kwargs["litellm_params"].get("batch_parent_id") is None: - self.aggregate_events.append(kwargs) - return - await self.gate.wait() - self.line_events.append(kwargs) - - async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): - await self._record(kwargs) - - async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): - await self._record(kwargs) - - -@pytest.mark.asyncio -async def test_a_slow_line_fan_out_does_not_hold_back_the_aggregate_event(recorder, monkeypatch): - """The logging worker cancels a callback coroutine after its deadline, so the - aggregate event, which bills the batch, must not wait on the line fan-out.""" - litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: flag under test; fixture restores it - gated: Final = _GatedLineLogger() - monkeypatch.setattr(litellm, "_async_success_callback", [gated]) - monkeypatch.setattr(litellm, "_async_failure_callback", [gated]) - with ( - patch( - "litellm.files.main.afile_content", new_callable=AsyncMock, side_effect=_file_content - ), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch - patch( - "litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02) - ), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch - ): - await _log_aggregate_success() - assert [_payload(e)["response_cost"] for e in gated.aggregate_events] == [1.5] - assert gated.line_events == [] - - gated.gate.set() - await asyncio.gather(*pending_batch_line_item_tasks()) - - assert sorted(_hidden(e)["batch_custom_id"] for e in gated.line_events) == ["a", "b"] - - -async def _log_aggregate_success() -> None: - await asyncio.wait_for( - _parent_logging().async_success_handler( - result=_batch(), - batch_cost=1.5, - batch_usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15), - batch_models=["gpt-4o"], - batch_successful_requests=1, - batch_failed_requests=1, - ), - timeout=5, - ) - - -@pytest.mark.asyncio -async def test_shutdown_drain_lets_a_running_line_fan_out_finish(recorder, monkeypatch): - """A fan-out keeps its batch claim once it starts, so shutdown waits for it - instead of closing the clients its remaining lines are sent through.""" - litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: flag under test; fixture restores it - gated: Final = _GatedLineLogger() - monkeypatch.setattr(litellm, "_async_success_callback", [gated]) - monkeypatch.setattr(litellm, "_async_failure_callback", [gated]) - with ( - patch( - "litellm.files.main.afile_content", new_callable=AsyncMock, side_effect=_file_content - ), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch - patch( - "litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02) - ), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch - ): - await _log_aggregate_success() - asyncio.get_running_loop().call_later(0.05, gated.gate.set) - still_running: Final = await drain_batch_line_items(timeout=5) - - assert still_running == 0 - assert sorted(_hidden(e)["batch_custom_id"] for e in gated.line_events) == ["a", "b"] - - -@pytest.mark.asyncio -async def test_shutdown_drain_stops_waiting_at_its_timeout(recorder, monkeypatch): - """A fan-out that outlasts the shutdown budget does not hold the shutdown open.""" - litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: flag under test; fixture restores it - gated: Final = _GatedLineLogger() - monkeypatch.setattr(litellm, "_async_success_callback", [gated]) - monkeypatch.setattr(litellm, "_async_failure_callback", [gated]) - with ( - patch( - "litellm.files.main.afile_content", new_callable=AsyncMock, side_effect=_file_content - ), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch - patch( - "litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02) - ), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch - ): - await _log_aggregate_success() - still_running: Final = await asyncio.wait_for(drain_batch_line_items(timeout=0.05), timeout=5) - delivered_before_gate: Final = list(gated.line_events) - gated.gate.set() - await asyncio.gather(*pending_batch_line_item_tasks()) - - assert still_running == 1 - assert delivered_before_gate == [] - - -@pytest.mark.asyncio -async def test_input_fetch_failure_still_emits_aggregate(recorder): - litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: flag under test; fixture restores it - with patch( - "litellm.files.main.afile_content", new_callable=AsyncMock, side_effect=ValueError("boom") - ): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch - await _log_completed_batch(_parent_logging(), _batch()) - - assert len(recorder.success_events) == 1 - assert len(recorder.failure_events) == 0 - assert _payload(recorder.success_events[0])["response_cost"] == 1.5 - - -EDGE_INPUT_JSONL = b"\n".join( - [ - json.dumps( - { - "custom_id": "e", - "method": "POST", - "url": "/v1/embeddings", - "body": {"model": "text-embedding-3-small", "input": "embed me", "encoding_format": "float"}, - } - ).encode(), - json.dumps( - { - "custom_id": "r", - "method": "POST", - "url": "/v1/responses", - "body": {"model": "gpt-4o", "input": "respond to me"}, - } - ).encode(), - json.dumps( - { - "custom_id": "badresp", - "method": "POST", - "url": "/v1/responses", - "body": {"model": "gpt-4o", "input": "unreconstructable"}, - } - ).encode(), - json.dumps( - { - "custom_id": "n", - "method": "POST", - "url": "/v1/chat/completions", - "body": {"model": "gpt-4o", "messages": [{"role": "user", "content": "hi n"}]}, - } - ).encode(), - ] -) - -EDGE_OUTPUT_JSONL = b"\n".join( - [ - json.dumps( - { - "custom_id": "e", - "response": { - "status_code": 200, - "body": { - "object": "list", - "data": [{"object": "embedding", "embedding": [0.1], "index": 0}], - "model": "text-embedding-3-small", - "usage": {"prompt_tokens": 3, "total_tokens": 3}, - }, - }, - } - ).encode(), - json.dumps( - { - "custom_id": "r", - "response": { - "status_code": 200, - "body": { - "id": "resp_1", - "object": "response", - "created_at": 1, - "status": "completed", - "output": [], - "model": "gpt-4o", - }, - }, - } - ).encode(), - json.dumps({"custom_id": "badresp", "response": {"status_code": 200, "body": {}}}).encode(), - json.dumps({"custom_id": "n", "response": {"body": {"error": {"message": "no status here"}}}}).encode(), - json.dumps({"custom_id": "boom", "response": "not-a-dict"}).encode(), - json.dumps( - { - "custom_id": "mi", - "modelInput": {"messages": [{"role": "user", "content": "hi mi"}]}, - "response": { - "status_code": 200, - "body": { - "choices": [ - { - "index": 0, - "message": {"role": "assistant", "content": "mi back"}, - "finish_reason": "stop", - } - ], - "usage": {"prompt_tokens": 2, "completion_tokens": 1, "total_tokens": 3}, - }, - }, - } - ).encode(), - ] -) - -ANTHROPIC_INPUT_JSONL = json.dumps( - { - "custom_id": "b2", - "params": {"model": "claude-3", "max_tokens": 5, "messages": [{"role": "user", "content": "hi b2"}]}, - } -).encode() - -ANTHROPIC_OUTPUT_JSONL = json.dumps( - { - "custom_id": "b2", - "result": { - "type": "succeeded", - "message": { - "id": "msg_1", - "type": "message", - "role": "assistant", - "content": [{"type": "text", "text": "hello b2"}], - "model": "claude-3", - "stop_reason": "end_turn", - "stop_sequence": None, - "usage": {"input_tokens": 1, "output_tokens": 2}, - }, - }, - } -).encode() - - -BEDROCK_INPUT_JSONL = b"\n".join( - [ - json.dumps( - { - "recordId": "br-anth", - "modelInput": { - "system": [{"type": "text", "text": "answer like claude"}], - "messages": [{"role": "user", "content": "hi br-anth"}], - }, - } - ).encode(), - json.dumps( - { - "recordId": "br-titan", - "modelInput": {"inputText": "embed me"}, - } - ).encode(), - json.dumps( - { - "recordId": "br-conv", - "modelInput": { - "system": [{"text": "answer like nova"}], - "messages": [{"role": "user", "content": "hi br-conv"}], - }, - } - ).encode(), - json.dumps( - { - "recordId": "br-err", - "modelInput": {"messages": [{"role": "user", "content": "hi br-err"}]}, - } - ).encode(), - ] -) - -BEDROCK_OUTPUT_JSONL = b"\n".join( - [ - json.dumps( - { - "recordId": "br-anth", - "modelOutput": { - "id": "msg_br", - "type": "message", - "role": "assistant", - "content": [{"type": "text", "text": "bedrock claude hi"}], - "model": "anthropic.claude-3", - "stop_reason": "end_turn", - "stop_sequence": None, - "usage": {"input_tokens": 2, "output_tokens": 3}, - }, - } - ).encode(), - json.dumps( - { - "recordId": "br-titan", - "modelOutput": {"embedding": [0.1, 0.2], "inputTextTokenCount": 4}, - } - ).encode(), - json.dumps( - { - "recordId": "br-conv", - "modelOutput": { - "output": {"message": {"role": "assistant", "content": [{"text": "converse hi"}]}}, - "stopReason": "end_turn", - "usage": {"inputTokens": 5, "outputTokens": 6, "totalTokens": 11}, - "metrics": {"latencyMs": 12}, - }, - } - ).encode(), - json.dumps( - { - "recordId": "br-err", - "modelInput": {"messages": [{"role": "user", "content": "hi br-err"}]}, - "error": {"errorCode": 400, "errorMessage": "bad request"}, - } - ).encode(), - ] -) - -BEDROCK_TITAN_INPUT_JSONL = b"\n".join( - [ - json.dumps({"recordId": "titan-ok", "modelInput": {"inputText": "embed me"}}).encode(), - json.dumps({"recordId": "titan-bad", "modelInput": {"inputText": "embed me too", "dimensions": 7}}).encode(), - ] -) - -BEDROCK_TITAN_OUTPUT_JSONL = b"\n".join( - [ - json.dumps( - { - "recordId": "titan-ok", - "modelInput": {"inputText": "embed me"}, - "modelOutput": {"embedding": [0.1, 0.2], "inputTextTokenCount": 3}, - } - ).encode(), - json.dumps( - { - "recordId": "titan-bad", - "modelInput": {"inputText": "embed me too", "dimensions": 7}, - "error": {"errorCode": 400, "errorMessage": "Malformed input request"}, - } - ).encode(), - ] -) - -MISTRAL_INPUT_JSONL = json.dumps( - { - "custom_id": "mis", - "body": {"model": "mistral-small", "messages": [{"role": "user", "content": "hi mis"}]}, - } -).encode() - -MISTRAL_OUTPUT_JSONL = json.dumps( - { - "custom_id": "mis", - "response": { - "status_code": 200, - "body": { - "id": "mis-1", - "choices": [ - { - "index": 0, - "message": {"role": "assistant", "content": "mistral hi"}, - "finish_reason": "stop", - } - ], - "usage": {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5}, - }, - }, - } -).encode() - - -def _edge_file_content(file_id: str, **_kwargs): - return SimpleNamespace( - content={ - "input-2": EDGE_INPUT_JSONL, - "output-2": EDGE_OUTPUT_JSONL, - "input-anth": ANTHROPIC_INPUT_JSONL, - "output-anth": ANTHROPIC_OUTPUT_JSONL, - "input-bed": BEDROCK_INPUT_JSONL, - "output-bed": BEDROCK_OUTPUT_JSONL, - "input-titan": BEDROCK_TITAN_INPUT_JSONL, - "output-titan": BEDROCK_TITAN_OUTPUT_JSONL, - "input-mist": MISTRAL_INPUT_JSONL, - "output-mist": MISTRAL_OUTPUT_JSONL, - }[file_id] - ) - - -def _parent_logging_with_params(litellm_params: dict) -> Logging: - logging_obj = _parent_logging() - logging_obj.update_environment_variables( - litellm_params=litellm_params, - optional_params={}, - custom_llm_provider="openai", - ) - return logging_obj - - -@pytest.mark.asyncio -async def test_a_line_the_batch_could_not_price_reports_the_zero_the_aggregate_billed(recorder): - litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: flag under test; fixture restores it - with ( - patch( - "litellm.files.main.afile_content", new_callable=AsyncMock, side_effect=_file_content - ), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch - patch( - "litellm.cost_calculator.batch_cost_calculator", side_effect=ValueError("no batch price") - ), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch - ): - await _log_completed_batch(_parent_logging(), _batch()) - - line_costs: Final = [ - (event["response_cost"], _payload(event)["response_cost"]) - for event in recorder.success_events - if _hidden(event).get("batch_custom_id") is not None - ] - assert line_costs == [(0.0, 0.0)] - - -@pytest.mark.asyncio -async def test_line_items_edge_shapes_and_edge_cases(recorder): - litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: flag under test; fixture restores it - batch = LiteLLMBatch( - id="batch_edge", - object="batch", - endpoint="/v1/chat/completions", - input_file_id="input-2", - output_file_id="output-2", - error_file_id=None, - status="completed", - completion_window="24h", - created_at=1, - ) - file_mock: Final = AsyncMock(side_effect=_edge_file_content) - parent: Final = _parent_logging_with_params( - { - "metadata": {"model_info": {"id": "dep-1"}, "model_group": "gpt-4o"}, - } - ) - parent._litellm_internal_model_credentials = { - "api_key": "sk-line-items-marker" - } # test-quality-ok: private transport attribute, same channel the batch cost tracker uses - with ( - patch( - "litellm.files.main.afile_content", file_mock - ), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch - patch( - "litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02) - ), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch - ): - await _log_completed_batch(parent, batch) - - by_custom_id = {_hidden(e).get("batch_custom_id"): e for e in recorder.success_events} - aggregate = next(e for e in recorder.success_events if _hidden(e).get("batch_custom_id") is None) - assert _payload(aggregate)["response_cost"] == 1.5 - - assert by_custom_id["e"]["litellm_params"]["batch_parent_id"] == batch.id - assert by_custom_id["e"]["call_type"] == "aembedding" - assert _payload(by_custom_id["e"])["response"]["data"][0]["embedding"] == [0.1] - - assert by_custom_id["r"]["call_type"] == "aresponses" - assert _payload(by_custom_id["r"])["response"]["id"] == "resp_1" - - assert by_custom_id["mi"]["call_type"] == "acompletion" - assert _hidden(by_custom_id["mi"])["batch_line_status_code"] == 200 - - assert "badresp" not in by_custom_id - assert "boom" not in by_custom_id - - failure = recorder.failure_events[0] - assert _hidden(failure)["batch_custom_id"] == "n" - assert _hidden(failure)["batch_line_status_code"] is None - assert "no status here" in _payload(failure)["error_str"] - - assert "sk-line-items-marker" in str(file_mock.call_args_list) - assert "sk-line-items-marker" not in str(by_custom_id["e"]["litellm_params"]) - file_ids_fetched = [call.kwargs.get("file_id") or call.args[0] for call in file_mock.call_args_list] - assert "input-2" in file_ids_fetched and "output-2" in file_ids_fetched - assert not any(file_id is None for file_id in file_ids_fetched) - - -@pytest.mark.asyncio -async def test_line_items_child_params_drop_parent_credentials(recorder): - litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: flag under test; fixture restores it - file_mock: Final = AsyncMock(side_effect=_file_content) - parent: Final = _parent_logging_with_params( - { - "api_key": "sk-parent-secret", - "_litellm_internal_model_credentials": MappingProxyType({"api_key": "sk-parent-secret"}), - "api_base": "https://api.openai.com/v1", - "metadata": {"model_info": {"id": "dep-1"}, "model_group": "gpt-4o"}, - } - ) - batch: Final = _batch() - with ( - patch( - "litellm.files.main.afile_content", file_mock - ), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch - patch( - "litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02) - ), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch - ): - await _log_completed_batch(parent, batch) - - line_events = [ - e for e in [*recorder.success_events, *recorder.failure_events] if _hidden(e).get("batch_custom_id") is not None - ] - assert len(line_events) == 2 - for event in line_events: - params = event["litellm_params"] - assert "api_key" not in params - assert "_litellm_internal_model_credentials" not in params - assert params["batch_parent_id"] == batch.id - assert params["metadata"]["model_group"] == "gpt-4o" - - -@pytest.mark.asyncio -async def test_line_items_anthropic_shapes(recorder): - litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: flag under test; fixture restores it - batch = LiteLLMBatch( - id="batch_anth", - object="batch", - endpoint="/v1/messages", - input_file_id="input-anth", - output_file_id="output-anth", - error_file_id=None, - status="completed", - completion_window="24h", - created_at=1, - ) - file_mock: Final = AsyncMock(side_effect=_edge_file_content) - logging_obj = _parent_logging() - logging_obj.update_environment_variables( - litellm_params={"metadata": {"model_info": {"id": "dep-1"}, "model_group": "claude-3"}}, - optional_params={}, - custom_llm_provider="anthropic", - ) - with ( - patch( - "litellm.files.main.afile_content", file_mock - ), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch - patch( - "litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02) - ), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch - ): - await _log_completed_batch(logging_obj, batch) - - line = next(e for e in recorder.success_events if _hidden(e).get("batch_custom_id") == "b2") - assert _hidden(line)["batch_line_status_code"] == 200 - assert line["litellm_params"]["batch_parent_id"] == batch.id - payload = _payload(line) - assert payload["response"]["choices"][0]["message"]["content"] == "hello b2" - assert payload["prompt_tokens"] == 1 - assert payload["completion_tokens"] == 2 - - -@pytest.mark.asyncio -async def test_anthropic_line_items_arrive_without_an_input_file(recorder): - """Anthropic keeps no input file, so a retrieved Anthropic batch reports input_file_id "None".""" - litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: flag under test; fixture restores it - batch = LiteLLMBatch( - id="batch_anth_no_input", - object="batch", - endpoint="/v1/messages", - input_file_id="None", - output_file_id="output-anth", - error_file_id=None, - status="completed", - completion_window="24h", - created_at=1, - ) - file_mock: Final = AsyncMock(side_effect=_edge_file_content) - logging_obj = _parent_logging(custom_llm_provider="anthropic") - with ( - patch( - "litellm.files.main.afile_content", file_mock - ), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch - patch( - "litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02) - ), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch - ): - await _log_completed_batch(logging_obj, batch) - - assert "None" not in [call.kwargs["file_id"] for call in file_mock.await_args_list] - line = next(e for e in recorder.success_events if _hidden(e).get("batch_custom_id") == "b2") - assert line["litellm_params"]["batch_parent_id"] == batch.id - assert _payload(line)["response"]["choices"][0]["message"]["content"] == "hello b2" - - -def _provider_batch(batch_id: str, input_file_id: str, output_file_id: str) -> LiteLLMBatch: - return LiteLLMBatch( - id=batch_id, - object="batch", - endpoint="/v1/chat/completions", - input_file_id=input_file_id, - output_file_id=output_file_id, - error_file_id=None, - status="completed", - completion_window="24h", - created_at=1, - ) - - -@pytest.mark.asyncio -async def test_line_items_bedrock_shapes(recorder): - litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: flag under test; fixture restores it - file_mock: Final = AsyncMock(side_effect=_edge_file_content) - with ( - patch( - "litellm.files.main.afile_content", file_mock - ), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch - patch( - "litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02) - ), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch - ): - await _log_completed_batch( - _parent_logging(custom_llm_provider="bedrock"), - _provider_batch("batch_bed", "input-bed", "output-bed"), - ) - - by_id = {_hidden(e).get("batch_custom_id"): e for e in recorder.success_events} - - anthropic_event = by_id["br-anth"] - assert anthropic_event["call_type"] == "acompletion" - assert _hidden(anthropic_event)["batch_custom_id"] == "br-anth" - anthropic_line = _payload(anthropic_event) - assert anthropic_line["messages"] == [ - {"role": "system", "content": [{"type": "text", "text": "answer like claude"}]}, - {"role": "user", "content": "hi br-anth"}, - ] - assert anthropic_line["response"]["choices"][0]["message"]["content"] == "bedrock claude hi" - assert anthropic_line["prompt_tokens"] == 2 - assert anthropic_line["completion_tokens"] == 3 - - titan_event = by_id["br-titan"] - assert titan_event["call_type"] == "aembedding" - assert _hidden(titan_event)["batch_custom_id"] == "br-titan" - titan_line = _payload(titan_event) - assert titan_event["optional_params"]["inputText"] == "embed me" - assert titan_line["response"]["data"][0]["embedding"] == [0.1, 0.2] - assert titan_line["prompt_tokens"] == 4 - - converse_event = by_id["br-conv"] - assert converse_event["call_type"] == "acompletion" - converse_line = _payload(converse_event) - assert converse_line["messages"] == [ - {"role": "system", "content": [{"text": "answer like nova"}]}, - {"role": "user", "content": "hi br-conv"}, - ] - assert converse_line["response"]["choices"][0]["message"]["content"] == "converse hi" - assert converse_line["prompt_tokens"] == 5 - assert converse_line["completion_tokens"] == 6 - - assert [_hidden(by_id[line])["batch_line_status_code"] for line in ("br-anth", "br-titan", "br-conv")] == [200] * 3 - (failure,) = recorder.failure_events - assert (_hidden(failure)["batch_custom_id"], _hidden(failure)["batch_line_status_code"]) == ("br-err", 400) - assert failure["call_type"] == "acompletion" - assert "bad request" in _payload(failure)["error_str"] - - -@pytest.mark.asyncio -async def test_bedrock_embedding_batch_lines_carry_the_embedding_call_type_and_input(recorder): - litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: flag under test; fixture restores it - with ( - patch( - "litellm.files.main.afile_content", new_callable=AsyncMock, side_effect=_edge_file_content - ), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch - patch( - "litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02) - ), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch - ): - await _log_completed_batch( - _parent_logging(custom_llm_provider="bedrock", model="bedrock/amazon.titan-embed-text-v2:0"), - _provider_batch("batch_titan", "input-titan", "output-titan"), - ) - - (success,) = (e for e in recorder.success_events if _hidden(e).get("batch_custom_id") == "titan-ok") - (failure,) = recorder.failure_events - assert [(_hidden(e)["batch_custom_id"], e["call_type"], _payload(e)["call_type"]) for e in (success, failure)] == [ - ("titan-ok", "aembedding", "aembedding"), - ("titan-bad", "aembedding", "aembedding"), - ] - assert _hidden(failure)["batch_line_status_code"] == 400 - assert [[m["content"] for m in _payload(e)["messages"]] for e in (success, failure)] == [ - ["embed me"], - ["embed me too"], - ] - - -@pytest.mark.asyncio -async def test_a_bedrock_batch_the_poller_logs_reads_every_file_from_the_deployment_bucket(recorder): - """Bedrock trusts its S3 bucket only from the server-side credential snapshot, so the line - fan-out the cost poller starts must hand that snapshot to every file read""" - litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: flag under test; fixture restores it - file_mock: Final = AsyncMock(side_effect=_edge_file_content) - parent: Final = _parent_logging(custom_llm_provider="bedrock") - parent._litellm_internal_model_credentials = ( - MappingProxyType( # test-quality-ok: private transport attribute, same channel the batch cost poller sets - {"s3_bucket_name": "line-items-bucket", "aws_region_name": "us-east-1"} - ) - ) - with ( - patch( - "litellm.files.main.afile_content", file_mock - ), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch - patch( - "litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02) - ), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch - ): - await _log_completed_batch(parent, _provider_batch("batch_bed_poller", "input-bed", "output-bed")) - - buckets_by_file: Final[dict[str, tuple[str, ...]]] = {} # mutable-ok: collects one entry per fetched file - for call in file_mock.call_args_list: - litellm_params: dict[str, object] = {} # mutable-ok: the dict the real file read builds and passes on - add_trusted_model_credentials_to_litellm_params(litellm_params, call.kwargs) - buckets_by_file[call.kwargs["file_id"]] = get_configured_s3_bucket_names(litellm_params) - assert buckets_by_file == { - "output-bed": ("line-items-bucket",), - "input-bed": ("line-items-bucket",), - } - assert sorted( - _hidden(e)["batch_custom_id"] for e in recorder.success_events if _hidden(e).get("batch_custom_id") - ) == [ - "br-anth", - "br-conv", - "br-titan", - ] - - -@pytest.mark.asyncio -async def test_line_items_mistral_shape(recorder): - litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: flag under test; fixture restores it - file_mock: Final = AsyncMock(side_effect=_edge_file_content) - with ( - patch( - "litellm.files.main.afile_content", file_mock - ), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch - patch( - "litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02) - ), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch - ): - await _log_completed_batch( - _parent_logging(custom_llm_provider="mistral"), - _provider_batch("batch_mist", "input-mist", "output-mist"), - ) - - line = next(e for e in recorder.success_events if _hidden(e).get("batch_custom_id") == "mis") - assert _hidden(line)["batch_line_status_code"] == 200 - assert _payload(line)["response"]["choices"][0]["message"]["content"] == "mistral hi" - - -ANTHROPIC_ERR_INPUT_JSONL = json.dumps( - { - "custom_id": "aerr", - "params": {"model": "claude-3", "max_tokens": 5, "messages": [{"role": "user", "content": "hi aerr"}]}, - } -).encode() - -ANTHROPIC_ERR_OUTPUT_JSONL = json.dumps( - { - "custom_id": "aerr", - "result": { - "type": "errored", - "error": { - "type": "error", - "error": {"type": "invalid_request_error", "message": "max_tokens must be positive"}, - }, - }, - } -).encode() - -VERTEX_NATIVE_OUTPUT_JSONL = b"\n".join( - [ - json.dumps( - { - "request": {"contents": [{"role": "user", "parts": [{"text": "hi v1"}]}]}, - "response": { - "candidates": [{"content": {"role": "model", "parts": [{"text": "hi back"}]}}], - "usageMetadata": {"promptTokenCount": 3, "candidatesTokenCount": 2, "totalTokenCount": 5}, - }, - } - ).encode(), - json.dumps( - { - "request": {"contents": [{"role": "user", "parts": [{"text": "hi v2"}]}]}, - "status": "Internal error", - } - ).encode(), - ] -) - - -def _scoped_file_content(file_map): - return lambda file_id, **kwargs: SimpleNamespace(content=file_map[file_id]) - - -@pytest.mark.asyncio -async def test_line_items_anthropic_failure_keeps_provider_error(recorder): - litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: flag under test; fixture restores it - file_mock: Final = AsyncMock( - side_effect=_scoped_file_content( - { - "input-anth-err": ANTHROPIC_ERR_INPUT_JSONL, - "output-anth-err": ANTHROPIC_ERR_OUTPUT_JSONL, - } - ) - ) - with ( - patch( - "litellm.files.main.afile_content", file_mock - ), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch - patch( - "litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02) - ), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch - ): - await _log_completed_batch( - _parent_logging(custom_llm_provider="anthropic"), - _provider_batch("batch_anth_err", "input-anth-err", "output-anth-err"), - ) - - assert len(recorder.failure_events) == 1 - failure = recorder.failure_events[0] - assert _hidden(failure)["batch_custom_id"] == "aerr" - assert _hidden(failure)["batch_line_status_code"] is None - assert "max_tokens must be positive" in _payload(failure)["error_str"] - - -BODYLESS_ERR_INPUT_JSONL = json.dumps( - {"custom_id": "bare", "url": "/v1/chat/completions", "body": {"model": "gpt-4.1-nano", "messages": []}} -).encode() - -BODYLESS_ERR_OUTPUT_JSONL = json.dumps({"custom_id": "bare", "response": None, "error": None}).encode() - - -@pytest.mark.asyncio -async def test_line_items_failure_without_error_or_response_body_still_reaches_callbacks(recorder): - litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: flag under test; fixture restores it - file_mock: Final = AsyncMock( - side_effect=_scoped_file_content( - { - "input-bare": BODYLESS_ERR_INPUT_JSONL, - "output-bare": BODYLESS_ERR_OUTPUT_JSONL, - } - ) - ) - with ( - patch( - "litellm.files.main.afile_content", file_mock - ), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch - patch( - "litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02) - ), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch - ): - await _log_completed_batch(_parent_logging(), _provider_batch("batch_bare", "input-bare", "output-bare")) - - assert len(recorder.failure_events) == 1 - failure = recorder.failure_events[0] - assert _hidden(failure)["batch_custom_id"] == "bare" - assert _payload(failure)["error_str"] == "{}" - - -MISTRAL_ERR_INPUT_JSONL = json.dumps( - { - "custom_id": "mis-err", - "method": "POST", - "url": "/v1/chat/completions", - "body": {"messages": [{"role": "user", "content": "hi"}], "tool_choice": "any"}, - } -).encode() - -MISTRAL_ERR_OUTPUT_JSONL = json.dumps( - { - "id": "batch-mist-err-1", - "custom_id": "mis-err", - "response": { - "status_code": 400, - "body": json.dumps( - {"object": "error", "message": "No tools provided for tool choice `any`", "type": "invalid_request"} - ), - }, - "error": None, - } -).encode() - - -@pytest.mark.asyncio -async def test_a_mistral_error_line_with_a_text_body_reaches_the_failure_callbacks(recorder): - litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: flag under test; fixture restores it - file_mock: Final = AsyncMock( - side_effect=_scoped_file_content( - { - "input-mist-err": MISTRAL_ERR_INPUT_JSONL, - "output-mist-err": MISTRAL_ERR_OUTPUT_JSONL, - } - ) - ) - with ( - patch( - "litellm.files.main.afile_content", file_mock - ), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch - patch( - "litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02) - ), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch - ): - await _log_completed_batch( - _parent_logging(custom_llm_provider="mistral"), - _provider_batch("batch_mist_err", "input-mist-err", "output-mist-err"), - ) - - assert [ - (_hidden(event)["batch_custom_id"], _hidden(event)["batch_line_status_code"]) - for event in recorder.failure_events - ] == [("mis-err", 400)] - assert "No tools provided for tool choice" in _payload(recorder.failure_events[0])["error_str"] - - -@pytest.mark.asyncio -async def test_line_items_native_vertex_rows_are_skipped(recorder): - litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: flag under test; fixture restores it - file_mock: Final = AsyncMock( - side_effect=_scoped_file_content( - { - "input-vtx": b"", - "output-vtx": VERTEX_NATIVE_OUTPUT_JSONL, - } - ) - ) - with ( - patch( - "litellm.files.main.afile_content", file_mock - ), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch - patch( - "litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02) - ), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch - ): - await _log_completed_batch( - _parent_logging(custom_llm_provider="vertex_ai"), - _provider_batch("batch_vtx", "input-vtx", "output-vtx"), - ) - - assert len(recorder.success_events) == 1 - assert len(recorder.failure_events) == 0 - assert _hidden(recorder.success_events[0]).get("batch_custom_id") is None - assert "input-vtx" not in [call.kwargs.get("file_id") or call.args[0] for call in file_mock.call_args_list] - - -class _NeverClaimsCache(DualCache): - async def async_increment_cache(self, *_args: object, **_kwargs: object) -> None: - return None - - -@pytest.mark.asyncio -async def test_line_items_emit_once_per_batch_id(recorder): - file_mock: Final = AsyncMock(side_effect=_file_content) - claim_cache: Final = DualCache() - parent: Final = _parent_logging() - batch: Final = _batch() - with patch( - "litellm.files.main.afile_content", file_mock - ): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch - first: Final = await log_batch_line_items( - batch=batch, - custom_llm_provider="openai", - parent=parent, - model_name="gpt-4o", - litellm_params=None, - model_info=None, - claim_cache=claim_cache, - ) - second: Final = await log_batch_line_items( - batch=batch, - custom_llm_provider="openai", - parent=parent, - model_name="gpt-4o", - litellm_params=None, - model_info=None, - claim_cache=claim_cache, - ) - - assert first == 2 - assert second == 0 - assert len(recorder.success_events) == 1 - assert len(recorder.failure_events) == 1 - - -@pytest.mark.asyncio -async def test_an_in_memory_claim_survives_a_thousand_newer_batches(): - first: Final = await _claim_line_items(batch_line_item_claim_cache, "batch_line_items_emitted:batch_old", "old") - for index in range(1_000): - await _claim_line_items(batch_line_item_claim_cache, f"batch_line_items_emitted:batch_{index}", str(index)) - again: Final = await _claim_line_items(batch_line_item_claim_cache, "batch_line_items_emitted:batch_old", "again") - - assert (first, again) == ("claimed", "already_claimed") - - -@pytest.mark.asyncio -async def test_line_items_claim_is_scoped_to_the_batch_id(recorder): - file_mock: Final = AsyncMock(side_effect=_file_content) - claim_cache: Final = DualCache() - parent: Final = _parent_logging() - with patch( - "litellm.files.main.afile_content", file_mock - ): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch - first: Final = await log_batch_line_items( - batch=_batch(), - custom_llm_provider="openai", - parent=parent, - model_name="gpt-4o", - litellm_params=None, - model_info=None, - claim_cache=claim_cache, - ) - second: Final = await log_batch_line_items( - batch=_provider_batch("batch_other", "input-file-1", "output-file-1"), - custom_llm_provider="openai", - parent=parent, - model_name="gpt-4o", - litellm_params=None, - model_info=None, - claim_cache=claim_cache, - ) - - assert first == 2 - assert second == 1 - assert len(recorder.success_events) == 2 - assert len(recorder.failure_events) == 1 - - -@pytest.mark.asyncio -async def test_line_items_emit_when_the_claim_backend_returns_nothing(recorder): - file_mock: Final = AsyncMock(side_effect=_file_content) - with patch( - "litellm.files.main.afile_content", file_mock - ): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch - emitted: Final = await log_batch_line_items( - batch=_batch(), - custom_llm_provider="openai", - parent=_parent_logging(), - model_name="gpt-4o", - litellm_params=None, - model_info=None, - claim_cache=_NeverClaimsCache(), - ) - - assert emitted == 2 - assert len(recorder.success_events) == 1 - assert len(recorder.failure_events) == 1 - - -_BROKEN_ERROR_FILE: Final = {**_FILE_BYTES, "error-file-1": "not bytes"} - - -def _broken_error_file_content(file_id: str, **_kwargs): - return SimpleNamespace(content=_BROKEN_ERROR_FILE[file_id]) - - -@pytest.mark.asyncio -async def test_line_items_retry_after_a_failed_fanout(recorder): - claim_cache: Final = DualCache() - batch: Final = _batch() - parent: Final = _parent_logging() - with patch( - "litellm.files.main.afile_content", new_callable=AsyncMock, side_effect=ValueError("boom") - ): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch - first: Final = await log_batch_line_items( - batch=batch, - custom_llm_provider="openai", - parent=parent, - model_name="gpt-4o", - litellm_params=None, - model_info=None, - claim_cache=claim_cache, - ) - file_mock: Final = AsyncMock(side_effect=_file_content) - with patch( - "litellm.files.main.afile_content", file_mock - ): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch - second: Final = await log_batch_line_items( - batch=batch, - custom_llm_provider="openai", - parent=parent, - model_name="gpt-4o", - litellm_params=None, - model_info=None, - claim_cache=claim_cache, - ) - - assert first == 0 - assert second == 2 - assert len(recorder.success_events) == 1 - assert len(recorder.failure_events) == 1 - - -@pytest.mark.asyncio -async def test_line_items_claim_kept_after_partial_emission(recorder): - claim_cache: Final = DualCache() - batch: Final = _batch() - parent: Final = _parent_logging() - file_mock: Final = AsyncMock(side_effect=_broken_error_file_content) - with patch( - "litellm.files.main.afile_content", file_mock - ): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch - first: Final = await log_batch_line_items( - batch=batch, - custom_llm_provider="openai", - parent=parent, - model_name="gpt-4o", - litellm_params=None, - model_info=None, - claim_cache=claim_cache, - ) - second: Final = await log_batch_line_items( - batch=batch, - custom_llm_provider="openai", - parent=parent, - model_name="gpt-4o", - litellm_params=None, - model_info=None, - claim_cache=claim_cache, - ) - - assert first == 1 - assert second == 0 - assert len(recorder.success_events) == 1 - assert len(recorder.failure_events) == 0 - - -class _FakeRedisCache(RedisCache): - """Dict-backed RedisCache stand-in honoring SET NX so the tokenized - line-item claim can be exercised without a Redis process.""" - - def __init__(self) -> None: - self._store: dict[str, object] = {} # mutable-ok: a dict-backed fake needs a mutable store - self.fail_next_set = False - self.swallow_next_set = False - self.swap_owner_on_get: str | None = None - - async def async_set_cache(self, key, value, nx=False, **_kwargs): - if self.fail_next_set: - self.fail_next_set = False - raise ConnectionError("redis down") - if self.swallow_next_set: - self.swallow_next_set = False - return None - if nx and key in self._store: - return None - self._store[key] = value - return True - - async def async_get_cache(self, key, **_kwargs): - if self.swap_owner_on_get is not None and key in self._store: - self._store[key] = self.swap_owner_on_get - return self._store.get(key) - - async def async_delete_cache(self, key): - self._store.pop(key, None) - - -@pytest.mark.asyncio -async def test_line_items_emit_once_per_batch_id_over_redis(recorder): - claim_cache: Final = DualCache(redis_cache=_FakeRedisCache()) - file_mock: Final = AsyncMock(side_effect=_file_content) - parent: Final = _parent_logging() - batch: Final = _batch() - with patch( - "litellm.files.main.afile_content", file_mock - ): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch - first: Final = await log_batch_line_items( - batch=batch, - custom_llm_provider="openai", - parent=parent, - model_name="gpt-4o", - litellm_params=None, - model_info=None, - claim_cache=claim_cache, - ) - second: Final = await log_batch_line_items( - batch=batch, - custom_llm_provider="openai", - parent=parent, - model_name="gpt-4o", - litellm_params=None, - model_info=None, - claim_cache=claim_cache, - ) - - assert first == 2 - assert second == 0 - assert len(recorder.success_events) == 1 - assert len(recorder.failure_events) == 1 - - -@pytest.mark.asyncio -async def test_line_items_skip_when_redis_claim_backend_is_down(recorder): - fake: Final = _FakeRedisCache() - fake.fail_next_set = True - claim_cache: Final = DualCache(redis_cache=fake) - file_mock: Final = AsyncMock(side_effect=_file_content) - parent: Final = _parent_logging() - batch: Final = _batch() - with patch( - "litellm.files.main.afile_content", file_mock - ): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch - first: Final = await log_batch_line_items( - batch=batch, - custom_llm_provider="openai", - parent=parent, - model_name="gpt-4o", - litellm_params=None, - model_info=None, - claim_cache=claim_cache, - ) - second: Final = await log_batch_line_items( - batch=batch, - custom_llm_provider="openai", - parent=parent, - model_name="gpt-4o", - litellm_params=None, - model_info=None, - claim_cache=claim_cache, - ) - - assert first == 0 - assert second == 2 - assert len(recorder.success_events) == 1 - assert len(recorder.failure_events) == 1 - - -@pytest.mark.asyncio -async def test_a_redis_set_failure_the_client_swallows_reads_as_unavailable(): - """RedisCache.async_set_cache logs a failed SET and returns None, the same value - SET NX returns when another worker holds the claim.""" - fake: Final = _FakeRedisCache() - claim_cache: Final = DualCache(redis_cache=fake) - fake.swallow_next_set = True - down: Final = await _claim_line_items(claim_cache, "batch_line_items_emitted:batch_down", "first") - claimed: Final = await _claim_line_items(claim_cache, "batch_line_items_emitted:batch_down", "second") - held: Final = await _claim_line_items(claim_cache, "batch_line_items_emitted:batch_down", "third") - assert (down, claimed, held) == ("unavailable", "claimed", "already_claimed") - - -@pytest.mark.asyncio -async def test_release_line_item_claim_only_deletes_owned_claims(): - fake: Final = _FakeRedisCache() - claim_cache: Final = DualCache(redis_cache=fake) - key: Final = "batch_line_items_emitted:batch_1" - fake._store[key] = ( - "someone-else" # test-quality-ok: seeding the fake's store is the arrangement, like writing Redis directly - ) - - await _release_line_item_claim(claim_cache, key, "my-token") - assert fake._store.get(key) == "someone-else" - - await _release_line_item_claim(claim_cache, key, "someone-else") - assert key not in fake._store - - -@pytest.mark.asyncio -async def test_line_items_failed_fanout_does_not_delete_another_workers_claim(recorder): - fake: Final = _FakeRedisCache() - fake.swap_owner_on_get = "other-worker" - claim_cache: Final = DualCache(redis_cache=fake) - key: Final = "batch_line_items_emitted:batch_1" - batch: Final = _batch() - parent: Final = _parent_logging() - with patch( - "litellm.files.main.afile_content", new_callable=AsyncMock, side_effect=ValueError("boom") - ): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch - first: Final = await log_batch_line_items( - batch=batch, - custom_llm_provider="openai", - parent=parent, - model_name="gpt-4o", - litellm_params=None, - model_info=None, - claim_cache=claim_cache, - ) - file_mock: Final = AsyncMock(side_effect=_file_content) - with patch( - "litellm.files.main.afile_content", file_mock - ): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch - second: Final = await log_batch_line_items( - batch=batch, - custom_llm_provider="openai", - parent=parent, - model_name="gpt-4o", - litellm_params=None, - model_info=None, - claim_cache=claim_cache, - ) - - assert first == 0 - assert second == 0 - assert fake._store.get(key) == "other-worker" - assert len(recorder.success_events) == 0 - assert len(recorder.failure_events) == 0 - - -@pytest.mark.asyncio -async def test_supplied_result_files_skip_the_output_and_error_fetch(recorder): - file_mock: Final = AsyncMock(side_effect=_file_content) - with patch( - "litellm.files.main.afile_content", file_mock - ): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch - emitted: Final = await log_batch_line_items( - batch=_batch(), - custom_llm_provider="openai", - parent=_parent_logging(), - model_name="gpt-4o", - litellm_params=None, - model_info=None, - result_files=BatchResultFiles(output=OUTPUT_JSONL, error=ERROR_JSONL), - claim_cache=DualCache(), - ) - - assert emitted == 2 - assert len(recorder.success_events) == 1 - assert len(recorder.failure_events) == 1 - assert file_mock.await_count == 1 - assert file_mock.await_args.kwargs["file_id"] == "input-file-1" - - -def _error_file_unreadable(file_id: str, **_kwargs): - if file_id == "error-file-1": - raise ValueError("error file gone") - return SimpleNamespace(content=_FILE_BYTES[file_id]) - - -@pytest.mark.asyncio -async def test_an_unreadable_error_file_defers_every_line_to_the_next_retrieve(recorder): - claim_cache: Final = DualCache() - batch: Final = _batch() - parent: Final = _parent_logging() - emitted: list[int] = [] # mutable-ok: collects one count per retrieve - for side_effect in (_error_file_unreadable, _file_content): - with patch( - "litellm.files.main.afile_content", AsyncMock(side_effect=side_effect) - ): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch - emitted.append( - await log_batch_line_items( - batch=batch, - custom_llm_provider="openai", - parent=parent, - model_name="gpt-4o", - litellm_params=None, - model_info=None, - result_files=BatchResultFiles(output=OUTPUT_JSONL, error=None), - claim_cache=claim_cache, - ) - ) - - assert emitted == [0, 2] - assert len(recorder.success_events) == 1 - assert len(recorder.failure_events) == 1 - - -@pytest.mark.asyncio -async def test_empty_result_files_fall_back_to_fetching_everything(recorder): - file_mock: Final = AsyncMock(side_effect=_file_content) - with patch( - "litellm.files.main.afile_content", file_mock - ): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch - emitted: Final = await log_batch_line_items( - batch=_batch(), - custom_llm_provider="openai", - parent=_parent_logging(), - model_name="gpt-4o", - litellm_params=None, - model_info=None, - result_files=BatchResultFiles(output=None, error=None), - claim_cache=DualCache(), - ) - - assert emitted == 2 - fetched: Final = sorted(call.kwargs["file_id"] for call in file_mock.await_args_list) - assert fetched == ["error-file-1", "input-file-1", "output-file-1"] - - -@pytest.mark.asyncio -async def test_native_vertex_skip_releases_the_claim(recorder): - file_mock: Final = AsyncMock( - side_effect=_scoped_file_content( - { - "input-vtx-rel": b"", - "output-vtx-rel": VERTEX_NATIVE_OUTPUT_JSONL, - } - ) - ) - claim_cache: Final = DualCache() - batch: Final = _provider_batch("batch_vtx_rel", "input-vtx-rel", "output-vtx-rel") - parent: Final = _parent_logging(custom_llm_provider="vertex_ai") - with patch( - "litellm.files.main.afile_content", file_mock - ): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch - first: Final = await log_batch_line_items( - batch=batch, - custom_llm_provider="vertex_ai", - parent=parent, - model_name="vertex_ai/gemini-2.5-flash", - litellm_params=None, - model_info=None, - claim_cache=claim_cache, - ) - second: Final = await log_batch_line_items( - batch=batch, - custom_llm_provider="vertex_ai", - parent=parent, - model_name="vertex_ai/gemini-2.5-flash", - litellm_params=None, - model_info=None, - claim_cache=claim_cache, - ) - - assert (first, second) == (0, 0) - assert [call.kwargs["file_id"] for call in file_mock.await_args_list] == ["output-vtx-rel", "output-vtx-rel"], ( - "a released claim must let the next retrieve reach the fetch step again; already_claimed would return early at 1 fetch" - ) - - -@pytest.mark.asyncio -async def test_fanout_emitting_nothing_releases_the_claim(recorder): - file_mock: Final = AsyncMock( - side_effect=_scoped_file_content( - { - "input-empty": b"", - "output-empty": b"", - "error-empty": b"", - } - ) - ) - claim_cache: Final = DualCache() - batch: Final = _provider_batch("batch_empty", "input-empty", "output-empty") - batch = batch.model_copy(update={"error_file_id": "error-empty"}) - parent: Final = _parent_logging() - with patch( - "litellm.files.main.afile_content", file_mock - ): # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch - first: Final = await log_batch_line_items( - batch=batch, - custom_llm_provider="openai", - parent=parent, - model_name="gpt-4o", - litellm_params=None, - model_info=None, - claim_cache=claim_cache, - ) - second: Final = await log_batch_line_items( - batch=batch, - custom_llm_provider="openai", - parent=parent, - model_name="gpt-4o", - litellm_params=None, - model_info=None, - claim_cache=claim_cache, - ) - - assert first == 0 - assert second == 0 - assert file_mock.await_count == 6, ( - "a completed fan-out that emits nothing must release the claim so the next retrieve retries" - ) - - -CALLER_FIELDS_INPUT_JSONL = json.dumps( - { - "custom_id": "a", - "method": "POST", - "url": "/v1/chat/completions", - "body": { - "model": "gpt-4o", - "messages": [{"role": "user", "content": "hi a"}], - "system": "answer briefly", - "temperature": 0.2, - "user": "end-user-a", - "call_type": "embedding", - "metadata": { - "_litellm_itpm_reserved": 5, - "_litellm_itpm_cache_key": "global_router:dep-2:gpt-4o:itpm:00-00", - }, - "standard_callback_dynamic_params": {"langfuse_host": "https://line.example"}, - }, - } -).encode() - - -@pytest.mark.asyncio -async def test_line_body_fields_do_not_replace_the_callback_kwargs(recorder): - """A JSONL line body is the caller's input file, so its fields reach callbacks as - the line's optional_params and never stand in for the kwargs litellm sets, which - router hooks read to settle rate limit reservations. The system prompt is still - logged, and the retrieve request's callback settings still apply.""" - litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: flag under test; fixture restores it - file_mock: Final = AsyncMock( - side_effect=_scoped_file_content({"input-fields": CALLER_FIELDS_INPUT_JSONL, "output-fields": OUTPUT_JSONL}) - ) - parent: Final = _parent_logging() - parent.standard_callback_dynamic_params = {"langfuse_host": "https://team.example"} - with ( - patch( - "litellm.files.main.afile_content", file_mock - ), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch - patch( - "litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02) - ), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch - ): - await _log_completed_batch(parent, _provider_batch("batch_fields", "input-fields", "output-fields")) - - line: Final = next(e for e in recorder.success_events if _hidden(e).get("batch_custom_id") == "a") - assert "metadata" not in line - assert line["call_type"] == "acompletion" - assert line["standard_callback_dynamic_params"] == {"langfuse_host": "https://team.example"} - assert _payload(line)["messages"][0] == {"role": "system", "content": "answer briefly"} - assert line["user"] == "end-user-a" - assert line["optional_params"]["temperature"] == 0.2 - assert line["optional_params"]["metadata"]["_litellm_itpm_reserved"] == 5 - - -@pytest.mark.asyncio -async def test_a_retrieve_after_the_shutdown_drain_leaves_every_line_to_the_next_retrieve(recorder): - """A fan-out started after the shutdown drain would be cancelled with the loop - while holding the batch claim, so it is not started, the batch stays unclaimed, - and the aggregate event still bills the batch.""" - litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: flag under test; fixture restores it - with ( - patch( - "litellm.files.main.afile_content", new_callable=AsyncMock, side_effect=_file_content - ), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch - patch( - "litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02) - ), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch - ): - assert await drain_batch_line_items(timeout=1) == 0 - await _log_aggregate_success() - loop: Final = asyncio.get_running_loop() - started: Final = [task for task in pending_batch_line_item_tasks() if task.get_loop() is loop] - await asyncio.gather(*started) - next_retrieve: Final = await log_batch_line_items( - batch=_batch(), - custom_llm_provider="openai", - parent=_parent_logging(), - model_name="gpt-4o", - litellm_params=None, - model_info=None, - ) - - assert started == [] - assert [ - _payload(e)["response_cost"] for e in recorder.success_events if _hidden(e).get("batch_custom_id") is None - ] == [1.5] - assert next_retrieve == 2 - - -@pytest.mark.asyncio -async def test_a_retrieve_flushed_at_interpreter_exit_starts_no_line_fan_out(monkeypatch): - """The logging worker's exit flush runs a queued retrieve handler on a loop it - closes as soon as the handler returns, so a fan-out started there would be - destroyed partway with the batch claim held.""" - exited_main_thread: Final = threading.Thread(target=lambda: None) - exited_main_thread.start() - exited_main_thread.join() - with monkeypatch.context() as interpreter_exit: - interpreter_exit.setattr(threading, "main_thread", lambda: exited_main_thread) - spawned: Final = spawn_batch_line_items( - batch=LiteLLMBatch( - id="batch_exit", - object="batch", - endpoint="/v1/chat/completions", - input_file_id="in", - status="completed", - completion_window="24h", - created_at=1, - ), - custom_llm_provider="openai", - parent=_parent_logging(), - model_name="gpt-4o", - litellm_params=None, - model_info=None, - result_files=None, - ) - - assert (spawned, pending_batch_line_item_tasks()) == (None, ()) - - -@pytest.mark.asyncio -async def test_line_events_follow_the_retrieve_requests_message_logging_opt_out(recorder): - """A key or team that turned message logging off gets its batch lines logged - without their prompts and responses, the same as its own requests.""" - litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: flag under test; fixture restores it - parent: Final = _parent_logging() - parent.standard_callback_dynamic_params = {"turn_off_message_logging": True} - with ( - patch( - "litellm.files.main.afile_content", new_callable=AsyncMock, side_effect=_file_content - ), # test-quality-ok: afile_content is the provider boundary; no injection seam for managed file fetch - patch( - "litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02) - ), # test-quality-ok: the pricing table boundary, same seam existing batch_utils tests patch - ): - await _log_completed_batch(parent, _batch()) - - line_events: Final = [ - e for e in recorder.success_events + recorder.failure_events if _hidden(e).get("batch_custom_id") - ] - assert sorted(_hidden(e)["batch_custom_id"] for e in line_events) == ["a", "b"] - logged: Final = json.dumps([[_payload(e)["messages"], _payload(e)["response"]] for e in line_events]) - assert "hi a" not in logged - assert "hi b" not in logged - assert "hello back" not in logged 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 000000000000..20211d370d1f --- /dev/null +++ b/tests/unit/batches/test_batch_line_items.py @@ -0,0 +1,258 @@ +import json +import logging + +import pytest + +import litellm +from litellm.batches.batch_line_items import ( + BatchLineItemSource, + build_batch_line_items, + 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.integrations.batch_line_item_objects import batch_line_items_ndjson, batch_line_items_object_name +from litellm.integrations.custom_logger import CustomLogger +from litellm.types.utils import ModelInfo, StandardBatchLineItemsPayload + +INPUT_RATE = 0.001 +OUTPUT_RATE = 0.002 +FINISHED_AT = 1791494130 + + +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": "hash-alice", "user_api_key_team_id": "team-alpha", "tags": ["env:prod"]} + ), + } + return BatchLineItemSource(**{**fields, **overrides}) + + +async def _payload() -> StandardBatchLineItemsPayload: + return await build_batch_line_items( + _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 = await build_batch_line_items( + _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": "hash-alice", + "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 = await build_batch_line_items(_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 = await build_batch_line_items( + _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 = await build_batch_line_items( + _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) + + +@pytest.mark.asyncio +async def test_redaction_drops_only_the_request_and_response(): + payload = await _payload() + + redacted = without_request_and_response(payload) + + assert redacted["batch_id"] == payload["batch_id"] + assert redacted["finished_at"] == payload["finished_at"] + 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 = await _payload() + + assert [json.loads(line) for line in batch_line_items_ndjson(payload).splitlines()] == list(payload["lines"]) + assert batch_line_items_object_name(payload) == "batch_line_items/2026-10-08/bGl0ZWxsbV9wcm94eTtiYXRjaA.jsonl" + assert ( + batch_line_items_object_name({**payload, "batch_id": "projects/p/locations/l/batchPredictionJobs/9"}) + == "batch_line_items/2026-10-08/projects_p_locations_l_batchPredictionJobs_9.jsonl" + ) + + +@pytest.mark.asyncio +async def test_a_logger_gets_the_request_and_response_unless_message_logging_is_off(monkeypatch): + payload = await _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 = await _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 = await _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/batches/test_batch_utils.py b/tests/unit/batches/test_batch_utils.py index 20aba4ea538f..293c38e9f206 100644 --- a/tests/unit/batches/test_batch_utils.py +++ b/tests/unit/batches/test_batch_utils.py @@ -2507,88 +2507,6 @@ async def test_flag_sends_every_vertex_row_down_the_native_path_when_a_model_is_ assert (result.successful_requests, result.failed_requests) == (0, 1) -@pytest.mark.asyncio -async def testhandle_completed_batch_with_files_returns_bytes_and_fetches_error_once(monkeypatch): - rows = [_success_row(model="gpt-4o", usage=_usage(10, 5))] - output_bytes = _vertex_jsonl(rows) - error_bytes = _vertex_jsonl([{"custom_id": "req-bad", "error": {"message": "rejected"}}]) - - fetched = [] - - async def fake_fetch(batch, custom_llm_provider, litellm_params=None): - return output_bytes - - async def fake_afile_content(**kw): - fetched.append(kw["file_id"]) - return type("R", (), {"content": error_bytes})() - - import litellm.cost_calculator as cc - import litellm.files.main as files_main - - monkeypatch.setattr(bu, "_fetch_batch_output_file_content", fake_fetch) - monkeypatch.setattr(files_main, "afile_content", fake_afile_content) - monkeypatch.setattr(cc, "batch_cost_calculator", lambda **kw: (2.0, 1.3)) - - batch = _batch("of").model_copy(update={"error_file_id": "ef"}) - - result, files = await bu.handle_completed_batch_with_files(batch, custom_llm_provider="openai") - - assert result.cost == 3.3 - assert result.failed_requests == 1 - assert files.output == output_bytes - assert files.error == error_bytes - assert fetched == ["ef"] - - -@pytest.mark.asyncio -async def testhandle_completed_batch_with_files_no_output_returns_error_bytes(monkeypatch): - error_bytes = _vertex_jsonl([{"custom_id": "req-bad", "error": {"message": "rejected"}}]) - - async def fake_afile_content(**kw): - return type("R", (), {"content": error_bytes})() - - import litellm.files.main as files_main - - monkeypatch.setattr(files_main, "afile_content", fake_afile_content) - - batch = _batch(None).model_copy(update={"error_file_id": "ef"}) - - result, files = await bu.handle_completed_batch_with_files(batch, custom_llm_provider="openai") - - assert result.failed_requests == 1 - assert result.usage.total_tokens == 0 - assert files.output is None - assert files.error == error_bytes - - -@pytest.mark.asyncio -async def test_handle_completed_batch_with_files_fetches_the_output_file_first(monkeypatch): - """The output file is fetched before the error file, the order - handle_completed_batch used before it returned the fetched bytes.""" - fetched: list[str] = [] - - async def fake_fetch(batch, custom_llm_provider, litellm_params=None): - fetched.append(batch.output_file_id) - return _vertex_jsonl([_success_row(model="gpt-4o", usage=_usage(10, 5))]) - - async def fake_afile_content(**kw): - fetched.append(kw["file_id"]) - return type("R", (), {"content": b""})() - - import litellm.cost_calculator as cc - import litellm.files.main as files_main - - monkeypatch.setattr(bu, "_fetch_batch_output_file_content", fake_fetch) - monkeypatch.setattr(files_main, "afile_content", fake_afile_content) - monkeypatch.setattr(cc, "batch_cost_calculator", lambda **kw: (2.0, 1.3)) - - await bu.handle_completed_batch_with_files( - _batch("of").model_copy(update={"error_file_id": "ef"}), custom_llm_provider="openai" - ) - - assert fetched == ["of", "ef"] - - @pytest.fixture() def _vcr_outcome_gate(request, vcr): install_live_call_probe(request, vcr) diff --git a/tests/unit/integrations/clickhouse/test_clickhouse_spend_logger.py b/tests/unit/integrations/clickhouse/test_clickhouse_spend_logger.py index 5d691d8652bd..563937decb9b 100644 --- a/tests/unit/integrations/clickhouse/test_clickhouse_spend_logger.py +++ b/tests/unit/integrations/clickhouse/test_clickhouse_spend_logger.py @@ -510,92 +510,6 @@ async def test_success_and_failure_events_write_scoped_spend_rows(): assert {key: row[key] for key in original_fields} == original_fields -@pytest.mark.asyncio -async def test_batch_line_item_success_event_does_not_write_spend(monkeypatch: pytest.MonkeyPatch): - # A batch line item carries call_type=acompletion + litellm_params.batch_parent_id. - # The aggregate aretrieve_batch row already bills the batch, so a per-line spend row - # would bill it twice. - monkeypatch.setattr(litellm, "store_batch_line_items_in_callbacks", True, raising=False) - storage = MagicMock() - logger = ClickHouseSpendLogger(storage=storage) - now = datetime(2026, 1, 1, tzinfo=timezone.utc) - - await logger.async_log_success_event( - { - "standard_logging_object": _minimal_payload("chatcmpl-line-1", status="success", cost=0.25), - "call_type": "acompletion", - "litellm_params": {"batch_parent_id": "batch-1"}, - }, - None, - now, - now, - ) - - assert logger.log_queue == [] - - await logger.async_log_success_event( - { - "standard_logging_object": _minimal_payload("chatcmpl-live-1", status="success", cost=0.25), - "call_type": "acompletion", - "litellm_params": {}, - }, - None, - now, - now, - ) - - # The aggregate aretrieve_batch event bills the batch; it must still write one - # row with the batch's full cost even while line items are skipped. - await logger.async_log_success_event( - { - "standard_logging_object": { - **_minimal_payload("batch-1", status="success", cost=0.0001032), - "call_type": "aretrieve_batch", - }, - "call_type": "aretrieve_batch", - "response_cost": 0.0001032, - "litellm_params": {}, - }, - None, - now, - now, - ) - - assert len(logger.log_queue) == 2 - assert logger.log_queue[0]["request_id"] == "chatcmpl-live-1" - assert logger.log_queue[1]["call_type"] == "aretrieve_batch" - assert logger.log_queue[1]["request_id"] == "batch-1" - assert logger.log_queue[1]["spend"] == 0.0001032 - - # A FAILED batch line is still part of the aggregate row's request counts and - # must not add its own spend row; a genuine live failure still writes one. - await logger.async_log_failure_event( - { - "standard_logging_object": _minimal_payload("chatcmpl-line-2", status="failure", cost=0.0), - "call_type": "acompletion", - "litellm_params": {"batch_parent_id": "batch-1"}, - }, - None, - now, - now, - ) - await logger.async_log_failure_event( - { - "standard_logging_object": _minimal_payload("chatcmpl-live-2", status="failure", cost=0.0), - "call_type": "acompletion", - "litellm_params": {}, - }, - None, - now, - now, - ) - - assert len(logger.log_queue) == 3 - assert logger.log_queue[-1]["request_id"] == "chatcmpl-live-2" - assert logger.log_queue[-1]["status"] == "failure" - if logger._flush_task is not None: - logger._flush_task.cancel() - @pytest.mark.parametrize( "status,llm_cost,guardrail_cost,expected", [ diff --git a/tests/unit/integrations/gcs_bucket/test_gcs_bucket.py b/tests/unit/integrations/gcs_bucket/test_gcs_bucket.py index 58b0f26ac3b9..d11ee75d85ab 100644 --- a/tests/unit/integrations/gcs_bucket/test_gcs_bucket.py +++ b/tests/unit/integrations/gcs_bucket/test_gcs_bucket.py @@ -330,3 +330,41 @@ def capture_upload(request: httpx.Request) -> httpx.Response: 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, "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.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_to_log(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, "lines": ({"custom_id": "ok-1"},)} + ) + + assert uploads.call_count == 1 diff --git a/tests/unit/integrations/test_batch_line_item_metering_sinks.py b/tests/unit/integrations/test_batch_line_item_metering_sinks.py deleted file mode 100644 index 1121332e8fe3..000000000000 --- a/tests/unit/integrations/test_batch_line_item_metering_sinks.py +++ /dev/null @@ -1,484 +0,0 @@ -""" -Guards that keep batch line-item callback events out of the built-in metering sinks. - -With ``litellm.store_batch_line_items_in_callbacks`` on, a completed batch emits -one child callback per JSONL line on top of the aggregate ``aretrieve_batch`` -event. Every line carries a real ``response_cost``, so any billing/metering sink -without the ``is_batch_line_item_event`` guard meters aggregate + per-line and -reports roughly 2x the true spend. - -These tests put REAL sink instances on the REAL dispatch lists, drive the REAL -``Logging.async_success_handler`` for a completed 2-line batch (1 success line + -1 error line, aggregate cost $1.50, per-line cost $0.03), and assert each sink -meters exactly the aggregate. OTel spans for line items must still be emitted: -the guard belongs on cost/token metrics, not on tracing. -""" - -import asyncio -import io -import json -import os -import uuid -from contextlib import redirect_stdout -from datetime import datetime -from types import SimpleNamespace -from typing import Any, Final -from unittest.mock import AsyncMock, patch - -import pytest - -import litellm -from litellm.batches.batch_line_item_logging import batch_line_item_claim_cache, pending_batch_line_item_tasks -from litellm.litellm_core_utils.litellm_logging import Logging -from litellm.types.utils import LiteLLMBatch, Usage - -AGGREGATE_COST: Final[float] = 1.5 - -INPUT_JSONL: Final[bytes] = b"\n".join( - [ - json.dumps( - { - "custom_id": "a", - "method": "POST", - "url": "/v1/chat/completions", - "body": {"model": "gpt-4o", "messages": [{"role": "user", "content": "hi a"}]}, - } - ).encode(), - json.dumps( - { - "custom_id": "b", - "method": "POST", - "url": "/v1/chat/completions", - "body": {"model": "gpt-4o", "messages": [{"role": "user", "content": "hi b"}]}, - } - ).encode(), - ] -) - -OUTPUT_JSONL: Final[bytes] = json.dumps( - { - "custom_id": "a", - "response": { - "status_code": 200, - "body": { - "id": "chatcmpl-line-1", - "model": "gpt-4o", - "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], - "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, - }, - }, - } -).encode() - -ERROR_JSONL: Final[bytes] = json.dumps( - { - "custom_id": "b", - "response": {"status_code": 400, "body": {"error": {"message": "boom"}}}, - "error": {"message": "boom"}, - } -).encode() - -_FILE_BYTES: Final[dict[str, bytes]] = { - "input-file-1": INPUT_JSONL, - "output-file-1": OUTPUT_JSONL, - "error-file-1": ERROR_JSONL, -} - - -def _file_content(file_id: str, **_kwargs: Any) -> SimpleNamespace: - return SimpleNamespace(content=_FILE_BYTES[file_id]) - - -def _batch() -> LiteLLMBatch: - return LiteLLMBatch( - id=f"batch_{uuid.uuid4().hex[:8]}", - object="batch", - endpoint="/v1/chat/completions", - input_file_id="input-file-1", - output_file_id="output-file-1", - error_file_id="error-file-1", - status="completed", - completion_window="24h", - created_at=1, - ) - - -def _parent_logging() -> Logging: - logging_obj = Logging( - model="gpt-4o", - messages=[{"role": "user", "content": ""}], - stream=False, - call_type="aretrieve_batch", - start_time=datetime(2026, 1, 1), - litellm_call_id=str(uuid.uuid4()), - function_id=str(uuid.uuid4()), - ) - logging_obj.update_environment_variables( - litellm_params={ - "metadata": { - "model_info": {"id": "dep-1"}, - "model_group": "gpt-4o", - "user_api_key_user_id": "user-77", - "user_api_key_team_id": "team-7", - "user_api_key_team_alias": "team-seven", - "user_api_key_alias": "key-alias-1", - } - }, - optional_params={}, - custom_llm_provider="openai", - ) - return logging_obj - - -class RecordingHTTP: - """Stands in for a sink's HTTP egress object only; the sink logic is real.""" - - def __init__(self) -> None: - self.posts: list[dict[str, Any]] = [] - - async def post(self, url: str, data: Any = None, content: Any = None, **_kw: Any) -> SimpleNamespace: - self.posts.append({"url": url, "body": data if data is not None else content}) - return SimpleNamespace(status_code=200, text="ok", raise_for_status=lambda: None) - - async def put(self, url: str, content: Any = None, **_kw: Any) -> SimpleNamespace: - return SimpleNamespace(status_code=202, text="ok", raise_for_status=lambda: None) - - -async def _log_completed_batch(monkeypatch: pytest.MonkeyPatch, loggers: list, flag: bool = True) -> None: - batch_line_item_claim_cache.in_memory_cache.flush_cache() - monkeypatch.setattr(litellm, "store_batch_line_items_in_callbacks", flag, raising=False) - saved_success = list(litellm._async_success_callback) - saved_failure = list(litellm._async_failure_callback) - litellm._async_success_callback = list(loggers) # test-quality-ok: no injection seam; finally restores - litellm._async_failure_callback = list(loggers) # test-quality-ok: same dispatch seam - buf = io.StringIO() - try: - with ( - patch("litellm.files.main.afile_content", new_callable=AsyncMock, side_effect=_file_content), - patch("litellm.cost_calculator.batch_cost_calculator", return_value=(0.01, 0.02)), - redirect_stdout(buf), - ): - await _parent_logging().async_success_handler( - result=_batch(), - batch_cost=1.5, - batch_usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15), - batch_models=["gpt-4o"], - batch_successful_requests=1, - batch_failed_requests=1, - batch_prompt_cost=1.0, - batch_completion_cost=0.5, - ) - await asyncio.gather(*pending_batch_line_item_tasks()) - finally: - litellm._async_success_callback = saved_success # test-quality-ok: restores the seam - litellm._async_failure_callback = saved_failure # test-quality-ok: restores the seam - - -def _openmeter_sinks() -> tuple[Any, RecordingHTTP]: - from litellm.integrations.openmeter import OpenMeterLogger - - recorder = RecordingHTTP() - logger = OpenMeterLogger() - logger.async_http_handler = recorder - return logger, recorder - - -def _lago_sinks() -> tuple[Any, RecordingHTTP]: - from litellm.integrations.lago import LagoLogger - - recorder = RecordingHTTP() - logger = LagoLogger() - logger.async_http_handler = recorder - return logger, recorder - - -@pytest.mark.asyncio -async def test_openmeter_meters_only_the_aggregate_batch_event(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setenv("OPENMETER_API_KEY", "test-openmeter-key") - logger, recorder = _openmeter_sinks() - - await _log_completed_batch(monkeypatch, [logger]) - - costs = [json.loads(p["body"])["data"]["cost"] for p in recorder.posts] - assert costs == [AGGREGATE_COST], ( - f"OpenMeter metered per-line costs on top of the aggregate: {costs}; " - "line items must be billed only by the aggregate aretrieve_batch event" - ) - - -@pytest.mark.asyncio -async def test_lago_bills_only_the_aggregate_batch_event(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setenv("LAGO_API_KEY", "test-lago-key") - monkeypatch.setenv("LAGO_API_BASE", "http://lago.invalid") - monkeypatch.setenv("LAGO_API_EVENT_CODE", "litellm-usage") - monkeypatch.setenv("LAGO_API_CHARGE_BY", "user_id") - logger, recorder = _lago_sinks() - - await _log_completed_batch(monkeypatch, [logger]) - - costs = [json.loads(p["body"])["event"]["properties"]["response_cost"] for p in recorder.posts] - assert costs == [AGGREGATE_COST], ( - f"Lago billed per-line costs on top of the aggregate: {costs}; " - "line items must be billed only by the aggregate aretrieve_batch event" - ) - - -@pytest.mark.asyncio -async def test_datadog_cost_management_bills_only_the_aggregate_batch_event( - monkeypatch: pytest.MonkeyPatch, -) -> None: - from litellm.integrations.datadog.datadog_cost_management import DatadogCostManagementLogger - - monkeypatch.setenv("DD_API_KEY", "test-dd-key") - monkeypatch.setenv("DD_APP_KEY", "test-dd-app-key") - logger = DatadogCostManagementLogger(cost_tag_keys=[]) - - await _log_completed_batch(monkeypatch, [logger]) - - entries = list(logger.log_queue) - costs = [e.get("response_cost", 0) for e in entries] - assert costs == [AGGREGATE_COST], ( - f"Datadog FOCUS BilledCost queued per-line entries on top of the aggregate: {costs}; " - "line items must be billed only by the aggregate aretrieve_batch event" - ) - - -@pytest.mark.asyncio -async def test_newrelic_metrics_meters_only_the_aggregate_batch_event(monkeypatch: pytest.MonkeyPatch) -> None: - from litellm.integrations.newrelic.newrelic_metrics import NewRelicMetricsLogger, build_metric_payload - - logger = NewRelicMetricsLogger(newrelic_api_key="test-nr-key") - - await _log_completed_batch(monkeypatch, [logger]) - - records = tuple(logger.log_queue) - costs = [r.response_cost for r in records] - assert costs == [AGGREGATE_COST], ( - f"New Relic metered per-line costs on top of the aggregate: {costs}; " - "line items must be metered only by the aggregate aretrieve_batch event" - ) - now = 1_767_225_600.0 - envelopes = build_metric_payload(records=records, window_start=now - 1, now=now) - cost_sum = 0.0 - for envelope in envelopes: - for metric in envelope["metrics"]: - if "cost" in metric["name"]: - value = metric["value"]["sum"] if isinstance(metric["value"], dict) else metric["value"] - cost_sum += value - assert cost_sum == pytest.approx(AGGREGATE_COST) - - -@pytest.mark.asyncio -async def test_prometheus_meters_only_the_aggregate_batch_event(monkeypatch: pytest.MonkeyPatch) -> None: - prometheus_client = pytest.importorskip("prometheus_client") - from litellm.integrations.prometheus import PrometheusLogger - from prometheus_client import REGISTRY - - saved_collectors = list(REGISTRY._collector_to_names.keys()) - for collector in saved_collectors: - REGISTRY.unregister(collector) - logger = PrometheusLogger() - - try: - await _log_completed_batch(monkeypatch, [logger]) - spend_samples = [] - for metric in REGISTRY.collect(): - if metric.name == "litellm_spend_metric": - spend_samples = [sample.value for sample in metric.samples if sample.name.endswith("_total")] - finally: - for collector in list(REGISTRY._collector_to_names.keys()): - REGISTRY.unregister(collector) - for collector in saved_collectors: - try: - REGISTRY.register(collector) - except Exception: # noqa: BLE001 # already re-registered by another holder - pass - - assert spend_samples, "litellm_spend_metric saw no samples at all" - assert sum(spend_samples) == pytest.approx(AGGREGATE_COST), ( - f"Prometheus spend metric double-counted batch line items: {spend_samples}; " - "line items must be metered only by the aggregate aretrieve_batch event" - ) - - -def _otel_v1(monkeypatch: pytest.MonkeyPatch): - otel_sdk = pytest.importorskip("opentelemetry.sdk") - from opentelemetry.sdk.metrics import MeterProvider - from opentelemetry.sdk.metrics.export import InMemoryMetricReader - from opentelemetry.sdk.trace import TracerProvider - from opentelemetry.sdk.trace.export import SimpleSpanProcessor - from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter - - from litellm.integrations.opentelemetry import OpenTelemetry as OTelV1, OpenTelemetryConfig - - reader = InMemoryMetricReader() - span_exporter = InMemorySpanExporter() - tracer_provider = TracerProvider() - tracer_provider.add_span_processor(SimpleSpanProcessor(span_exporter)) - logger = OTelV1( - config=OpenTelemetryConfig(exporter="console", enable_metrics=True), - callback_name="batch_sink_guard_v1", - tracer_provider=tracer_provider, - meter_provider=MeterProvider(metric_readers=[reader]), - ) - return logger, reader, span_exporter - - -def _cost_points(reader: Any) -> list[float]: - data = reader.get_metrics_data() - points: list[float] = [] - if data is None: - return points - for resource_metrics in data.resource_metrics: - for scope_metrics in resource_metrics.scope_metrics: - for metric in scope_metrics.metrics: - if not metric.name.endswith("cost"): - continue - for data_point in metric.data.data_points: - points.append(getattr(data_point, "sum", getattr(data_point, "value", None))) - return points - - -@pytest.mark.asyncio -async def test_otel_v1_meter_cost_skips_line_items_but_spans_stay( - monkeypatch: pytest.MonkeyPatch, -) -> None: - logger, reader, span_exporter = _otel_v1(monkeypatch) - - await _log_completed_batch(monkeypatch, [logger]) - - costs = _cost_points(reader) - assert sum(costs) == pytest.approx(AGGREGATE_COST), ( - f"OTel v1 gen_ai.usage.cost double-counted batch line items: {costs}; " - "line items must be metered only by the aggregate aretrieve_batch event" - ) - finished = span_exporter.get_finished_spans() - assert finished, "OTel v1 emitted no spans at all" - assert any("chatcmpl-line-1" in str(span.attributes) for span in finished), ( - f"OTel v1 dropped the per-line span the feature exists to deliver: {[span.name for span in finished]}" - ) - - -def _reader_snapshot(reader: Any) -> dict[str, list[float]]: - """Every metric's datapoint values, keyed by metric name (sums for histograms).""" - data = reader.get_metrics_data() - out: dict[str, list[float]] = {} - if data is None: - return out - for resource_metrics in data.resource_metrics: - for scope_metrics in resource_metrics.scope_metrics: - for metric in scope_metrics.metrics: - for data_point in metric.data.data_points: - value = getattr(data_point, "sum", None) - if value is None: - value = getattr(data_point, "value", None) - if value is not None: - out.setdefault(metric.name, []).append(value) - return out - - -async def _run_batch_with_otel_v2(monkeypatch: pytest.MonkeyPatch, flag: bool) -> dict[str, list[float]]: - from opentelemetry.sdk.metrics import MeterProvider - from opentelemetry.sdk.metrics.export import InMemoryMetricReader - from opentelemetry.sdk.trace import TracerProvider - - from litellm.integrations.otel.logger import OpenTelemetryV2 - from litellm.integrations.otel.model.config import OpenTelemetryV2Config - - reader = InMemoryMetricReader() - logger = OpenTelemetryV2( - config=OpenTelemetryV2Config(exporter="console", enable_metrics=True), - callback_name="batch_sink_guard_v2", - tracer_provider=TracerProvider(), - meter_provider=MeterProvider(metric_readers=[reader]), - ) - await _log_completed_batch(monkeypatch, [logger], flag=flag) - return _reader_snapshot(reader) - - -@pytest.mark.asyncio -async def test_otel_v2_metrics_are_identical_with_the_flag_on_and_off( - monkeypatch: pytest.MonkeyPatch, -) -> None: - pytest.importorskip("opentelemetry.sdk") - - off = await _run_batch_with_otel_v2(monkeypatch, flag=False) - on = await _run_batch_with_otel_v2(monkeypatch, flag=True) - - for name in ("gen_ai.usage.cost", "gen_ai.client.token.usage"): - assert on.get(name) == off.get(name), ( - f"OTel v2 {name} changed when the line-item flag was turned on; line items " - "must be metered only by the aggregate aretrieve_batch event. " - f"off={off.get(name)} on={on.get(name)}" - ) - for name, values in off.items(): - on_counts = [len(on.get(name, [])), len(values)] - assert on_counts[0] == on_counts[1], ( - f"OTel v2 {name} gained synthetic per-line samples with the flag on: " - f"{on_counts[0]} datapoints on vs {on_counts[1]} off" - ) - assert sum(off.get("gen_ai.usage.cost", [])) == pytest.approx(AGGREGATE_COST) - - -async def _datadog_metric_names(monkeypatch: pytest.MonkeyPatch, flag: bool) -> list[str]: - from litellm.integrations.datadog.datadog_metrics import DatadogMetricsLogger - - monkeypatch.setenv("DD_API_KEY", "test-dd-key") - logger = DatadogMetricsLogger(start_periodic_flush=False) - await _log_completed_batch(monkeypatch, [logger], flag=flag) - return sorted(series["metric"] for series in logger.log_queue) - - -@pytest.mark.asyncio -async def test_datadog_metrics_are_identical_with_the_flag_on_and_off(monkeypatch: pytest.MonkeyPatch) -> None: - off = await _datadog_metric_names(monkeypatch, flag=False) - on = await _datadog_metric_names(monkeypatch, flag=True) - - assert off - assert on == off - - -async def _slack_daily_report_updates(monkeypatch: pytest.MonkeyPatch, flag: bool) -> list[str]: - from litellm.caching.dual_cache import DualCache - from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting - - class _RecordingCache(DualCache): - """The real cache, also listing the daily report keys each event increments.""" - - def __init__(self) -> None: - super().__init__() - self.incremented: list[str] = [] - - async def async_increment_cache(self, key: str, value: float, **kwargs: Any) -> float: - self.incremented.append(key) - return await super().async_increment_cache(key=key, value=value, **kwargs) - - cache = _RecordingCache() - logger = SlackAlerting(internal_usage_cache=cache, alert_types=["daily_reports"]) - await _log_completed_batch(monkeypatch, [logger], flag=flag) - return cache.incremented - - -@pytest.mark.asyncio -async def test_slack_daily_reports_are_identical_with_the_flag_on_and_off(monkeypatch: pytest.MonkeyPatch) -> None: - """Daily reports rank deployments by their live failures and latency, which a - line that ran inside the provider's batch job says nothing about.""" - off = await _slack_daily_report_updates(monkeypatch, flag=False) - on = await _slack_daily_report_updates(monkeypatch, flag=True) - - assert off - assert on == off - - -@pytest.mark.asyncio -async def test_pagerduty_does_not_count_failed_batch_lines(monkeypatch: pytest.MonkeyPatch) -> None: - """A failed line is reported by its batch, so it never counts toward the - failure threshold that pages on-call for live traffic.""" - from litellm_enterprise.enterprise_callbacks.pagerduty.pagerduty import PagerDutyAlerting - - monkeypatch.setenv("PAGERDUTY_API_KEY", "test-pd-key") - logger = PagerDutyAlerting(alerting_args={"failure_threshold": 100}) - - await _log_completed_batch(monkeypatch, [logger]) - - assert logger._failure_events == [] diff --git a/tests/unit/integrations/test_s3_v2.py b/tests/unit/integrations/test_s3_v2.py index 8e9a8982cd97..44aa69904807 100644 --- a/tests/unit/integrations/test_s3_v2.py +++ b/tests/unit/integrations/test_s3_v2.py @@ -5540,3 +5540,82 @@ async def async_log_success_event(self, kwargs, response_obj, start_time, end_ti 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, + "lines": (line, {**line, "custom_id": "ok-2"}), + } + + +def _s3_logger_answering(status: int, s3_path: str | None) -> tuple[S3Logger, list[httpx.Request]]: + requests: list[httpx.Request] = [] + + def respond(request: httpx.Request) -> httpx.Response: + requests.append(request) + return httpx.Response(status, 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.jsonl" + assert request.headers["Content-Type"] == "application/x-ndjson" + 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.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 + ) diff --git a/tests/unit/integrations/test_shadow_eval_logger.py b/tests/unit/integrations/test_shadow_eval_logger.py index cc7171e83e86..55796f7bbcf1 100644 --- a/tests/unit/integrations/test_shadow_eval_logger.py +++ b/tests/unit/integrations/test_shadow_eval_logger.py @@ -1016,22 +1016,6 @@ def test_judge_prompt_is_bounded_however_large_the_inputs(): @pytest.mark.asyncio class TestSuccessHookSkipChain: - async def test_a_batch_line_item_is_not_shadowed(self, monkeypatch: pytest.MonkeyPatch): - """A completed batch's JSONL lines are logged as their own events, but the - batch API served them, so shadowing one would buy a realtime call for it.""" - import litellm as litellm_module - - monkeypatch.setattr(litellm_module, "completion_cost", lambda completion_response: 0.005) - router = _router() - logger = _logger(router=router, prisma=_prisma(), jobs=(_job(),)) - line_kwargs = _success_kwargs() - line_kwargs["litellm_params"]["batch_parent_id"] = "batch_1" - - await logger.async_log_success_event(line_kwargs, RESPONSE, None, None) - await _drain(logger) - - router.acompletion.assert_not_called() - async def test_happy_path_writes_exactly_one_attempt_row(self, monkeypatch: pytest.MonkeyPatch): import litellm as litellm_module diff --git a/tests/unit/litellm_core_utils/test_core_helpers.py b/tests/unit/litellm_core_utils/test_core_helpers.py index 0097b4fc4d70..1f716b41acd3 100644 --- a/tests/unit/litellm_core_utils/test_core_helpers.py +++ b/tests/unit/litellm_core_utils/test_core_helpers.py @@ -16,7 +16,6 @@ get_parent_otel_span_from_kwargs, get_or_create_metadata_bucket, get_provider_response_headers_from_hidden_params, - is_batch_line_item_event, map_finish_reason, normalize_drop_params, process_response_headers, @@ -526,16 +525,6 @@ def test_budget_rejection_decorated_with_provider_is_expected(self): assert is_expected_client_error(vendor_limit) is False -def test_is_batch_line_item_event(): - assert is_batch_line_item_event({"litellm_params": {"batch_parent_id": "batch_1"}}) is True - assert is_batch_line_item_event({"litellm_params": {"batch_parent_id": "batch_1", "metadata": {}}}) is True - assert is_batch_line_item_event({"litellm_params": {"metadata": {}}}) is False - assert is_batch_line_item_event({"litellm_params": {"batch_parent_id": None}}) is False - assert is_batch_line_item_event({}) is False - assert is_batch_line_item_event({"litellm_params": "not-a-mapping"}) is False - assert is_batch_line_item_event({"litellm_params": None}) is False - - class TestProviderResponseHeadersInHiddenParams: def test_records_raw_headers_and_the_processed_additional_headers(self): response = ImageResponse() diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index 553e452c3ada..870378e93f6a 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -32,7 +32,6 @@ from litellm.constants import LOGGING_WORKER_MAX_TIME_PER_COROUTINE, REDACTED_BY_LITELLM, SENTRY_PII_DENYLIST from litellm.cost_calculator import ocr_batch_cost from litellm.integrations.custom_logger import CustomLogger -from litellm.litellm_core_utils.core_helpers import BATCH_PARENT_ID_KEY, is_batch_line_item_event from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging from litellm.litellm_core_utils.litellm_logging import ( Logging, @@ -336,10 +335,8 @@ async def async_post_mcp_tool_call_hook(self, kwargs, response_obj, start_time, ) logging_obj.dynamic_success_callbacks = [NativeReplacement()] returned = await logging_obj.async_post_mcp_tool_call_hook( - kwargs={"original_response": result}, - response_obj=result, - start_time=datetime.datetime(2026, 1, 1), - end_time=datetime.datetime(2026, 1, 1), + kwargs={"original_response": result}, response_obj=result, + start_time=datetime.datetime.now(), end_time=datetime.datetime.now(), ) assert returned is result assert result.content == [TextContent(type="text", text="native-safe" if same_content else "final-safe")] @@ -348,6 +345,7 @@ async def async_post_mcp_tool_call_hook(self, kwargs, response_obj, start_time, assert "SECRET-1234" not in result.model_dump_json() + @pytest.mark.asyncio @pytest.mark.parametrize("mode", ["none", "wrapper", "exception"]) @pytest.mark.parametrize("structured", [False, True]) @@ -365,10 +363,8 @@ async def async_post_mcp_tool_call_hook(self, kwargs, response_obj, start_time, ) logging_obj.dynamic_success_callbacks = [DirectRedactor()] returned = await logging_obj.async_post_mcp_tool_call_hook( - kwargs={"original_response": result}, - response_obj=result, - start_time=datetime.datetime(2026, 1, 1), - end_time=datetime.datetime(2026, 1, 1), + kwargs={"original_response": result}, response_obj=result, + start_time=datetime.datetime.now(), end_time=datetime.datetime.now(), ) assert returned is result assert result.content == [TextContent(type="text", text="[REDACTED]")] @@ -443,8 +439,6 @@ def test_sentry_environment(monkeypatch): set_callbacks(["sentry"]) mock_init.assert_called_once() assert mock_init.call_args[1]["environment"] == environment - - def test_use_custom_pricing_for_model(): from litellm.litellm_core_utils.litellm_logging import use_custom_pricing_for_model @@ -2887,13 +2881,8 @@ def no_prisma() -> None: class RecordingShadowLogger(ShadowEvalLogger): async def async_log_success_event( - self, - kwargs: Mapping[str, object], - response_obj: object, - start_time: object, - end_time: object, - *, - guardrail_snapshot: GuardrailRequestSnapshot | None = None, + self, kwargs: Mapping[str, object], response_obj: object, start_time: object, + end_time: object, *, guardrail_snapshot: GuardrailRequestSnapshot | None = None, ) -> None: shadow_snapshots.append(guardrail_snapshot) await super().async_log_success_event( @@ -2902,20 +2891,13 @@ async def async_log_success_event( class RecordingLogger(CustomLogger): async def async_log_success_event( - self, - kwargs: Mapping[str, object], - response_obj: object, - start_time: object, - end_time: object, + self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object, ) -> None: other_payloads.append(kwargs) class LoggingGuardrail(CustomGuardrail): async def async_logging_hook( - self, - kwargs: dict[str, object], - result: object, - call_type: str, + self, kwargs: dict[str, object], result: object, call_type: str, ) -> tuple[dict[str, object], object]: hook_snapshots.append(logging_obj.shadow_eval_request_snapshot) if hook_mode == "raises": @@ -2927,36 +2909,26 @@ async def async_logging_hook( "user_api_key_hash": "test-key", } snapshot: Final = GuardrailRequestSnapshot.capture( - {"messages": [{"role": "user", "content": "snapshot-only"}]}, - metadata, + {"messages": [{"role": "user", "content": "snapshot-only"}]}, metadata, ) assert snapshot is not None shadow: Final = RecordingShadowLogger(prisma_provider=no_prisma, jobs_cache=InMemoryCache()) guardrail: Final = LoggingGuardrail( - guardrail_name="late-mask", - default_on=True, + guardrail_name="late-mask", default_on=True, event_hook=GuardrailEventHooks.pre_call if hook_mode == "disabled" else GuardrailEventHooks.logging_only, ) monkeypatch.setattr(litellm, "_async_success_callback", []) logging_obj: Final = LitellmLogging( - model="test-model", - messages=[], - stream=stream, - call_type="anthropic_messages", - start_time=datetime.datetime(2026, 1, 1), - litellm_call_id="private-snapshot", - function_id="private-snapshot", + model="test-model", messages=[], stream=stream, call_type="anthropic_messages", + start_time=datetime.datetime.now(), litellm_call_id="private-snapshot", function_id="private-snapshot", dynamic_async_success_callbacks=[shadow, RecordingLogger(), guardrail], ) logging_obj.update_messages([{"role": "user", "content": "logged input"}]) logging_obj.update_environment_variables(litellm_params={"metadata": metadata}, optional_params={}) logging_obj.shadow_eval_request_snapshot = snapshot payload: Final = { - "id": "private-snapshot", - "call_type": "anthropic_messages", - "metadata": metadata, - "model_group": "test-model", - "model_parameters": {}, + "id": "private-snapshot", "call_type": "anthropic_messages", "metadata": metadata, + "model_group": "test-model", "model_parameters": {}, } await logging_obj.async_success_handler(result=ModelResponse(), standard_logging_object=payload) @@ -3411,9 +3383,7 @@ def test_sentry_event_scrubber_initialization(monkeypatch): call_args = mock_init.call_args[1] assert call_args["send_default_pii"] is False assert call_args["event_scrubber"].recursive is True - assert {name.lower() for name in SENTRY_PII_DENYLIST} <= { - name.lower() for name in call_args["event_scrubber"].denylist - } + assert {name.lower() for name in SENTRY_PII_DENYLIST} <= {name.lower() for name in call_args["event_scrubber"].denylist} assert call_args["before_send"] is call_args["before_send_transaction"] @@ -3428,9 +3398,7 @@ def test_sentry_send_default_pii_opt_in(monkeypatch): call_args = mock_init.call_args[1] assert call_args["send_default_pii"] is True - assert not {name.lower() for name in SENTRY_PII_DENYLIST} & { - name.lower() for name in call_args["event_scrubber"].denylist - } + assert not {name.lower() for name in SENTRY_PII_DENYLIST} & {name.lower() for name in call_args["event_scrubber"].denylist} def test_get_masked_values(): @@ -5950,34 +5918,6 @@ def test_handle_anthropic_messages_response_logging_keeps_the_served_response_id assert result.service_tier == "flex" -@pytest.mark.parametrize("as_model", [False, True], ids=["dict", "pydantic"]) -def test_handle_anthropic_messages_response_logging_keeps_the_message_id_of_a_parsed_message(as_model): - from litellm.types.llms.anthropic import AnthropicResponse - - message: Final = { - "id": "msg_served", - "type": "message", - "role": "assistant", - "model": "claude-sonnet-4-5", - "content": [{"type": "text", "text": "hi"}], - "stop_reason": "end_turn", - "stop_sequence": None, - "usage": {"input_tokens": 3, "output_tokens": 1}, - } - logging_obj = _anthropic_messages_logging_obj() - logging_obj.update_environment_variables(litellm_params={}, optional_params={}) - - result = logging_obj._handle_anthropic_messages_response_logging( - result=AnthropicResponse.model_validate(message) if as_model else message - ) - - assert isinstance(result, ModelResponse) - assert result.id == "msg_served" - choice: Final = result.choices[0] - assert isinstance(choice, litellm.Choices) - assert choice.message.content == "hi" - - def test_handle_anthropic_messages_response_logging_passes_model_response_through(): """Anthropic-native path already yields a ModelResponse; it must be returned unchanged.""" logging_obj = _anthropic_messages_logging_obj() @@ -6889,8 +6829,8 @@ def test_restore_correlation_context_safe_to_call_repeatedly(monkeypatch): def test_restore_correlation_context_does_not_resanitize(monkeypatch): - from litellm._logging import _sanitize_correlation_id from litellm.litellm_core_utils.litellm_logging import Logging + from litellm._logging import _sanitize_correlation_id monkeypatch.setattr(litellm, "request_correlation_in_logs", True) trace_id_var.set("outer-trace") @@ -7037,12 +6977,12 @@ def test_read_and_management_calls_cost_nothing(self, call_type): assert cost == 0.0 def test_retrieved_usage_is_not_re_reported_in_standard_logging_payload(self): - from datetime import datetime - from litellm.litellm_core_utils.litellm_logging import ( get_standard_logging_object_payload, ) + from datetime import datetime + logging_obj = self._logging_obj("aget_responses") now = datetime.now() payload = get_standard_logging_object_payload( @@ -7346,9 +7286,7 @@ def test_debugging_log_with_json_logs_tolerates_missing_headers(logging_obj, mon logging_obj.litellm_request_debug = True with patch("litellm.litellm_core_utils.litellm_logging.verbose_logger.warning") as warning: - logging_obj._print_llm_call_debugging_log( - api_base="https://api.openai.com/v1", headers=None, additional_args={} - ) + logging_obj._print_llm_call_debugging_log(api_base="https://api.openai.com/v1", headers=None, additional_args={}) assert "https://api.openai.com/v1" in warning.call_args.kwargs["extra"]["api_base"] @@ -9239,6 +9177,8 @@ def test_get_assembled_streaming_response_bills_a_provider_reported_usage_cost() assert logging_obj.response_cost_calculator(result=assembled) == 0.0042 + + def test_response_cost_calculator_prices_terminal_responses_event_from_its_response(): logging_obj: Final = _responses_stream_logging_obj() inner_response: Final = ResponsesAPIResponse( @@ -9761,230 +9701,20 @@ def test_signoz_dispatch_requires_an_endpoint(monkeypatch): custom_logger_init_args={}, ) assert created is None - assert not [cb for cb in logging_module._in_memory_loggers if getattr(cb, "callback_name", None) == "signoz"] + assert not [ + cb for cb in logging_module._in_memory_loggers if getattr(cb, "callback_name", None) == "signoz" + ] finally: logging_module._in_memory_loggers.clear() monkeypatch.delenv("LITELLM_OTEL_V2", raising=False) is_otel_v2_enabled.cache_clear() -class TestRetrieveBatchReusesFetchedResultFiles: - """The aggregate path and log_batch_line_items must share one fetch of each - result file: output and error are fetched once total (not once each per - consumer), and forwarded bytes keep line-item logging off the wire.""" - - @staticmethod - def _batch_file_bytes() -> dict[str, bytes]: - input_jsonl = json.dumps( - { - "custom_id": "a", - "method": "POST", - "url": "/v1/chat/completions", - "body": {"model": "gpt-4o", "messages": [{"role": "user", "content": "hi"}]}, - } - ).encode() - output_jsonl = json.dumps( - { - "custom_id": "a", - "response": { - "status_code": 200, - "body": { - "id": "chatcmpl-1", - "model": "gpt-4o", - "choices": [ - { - "index": 0, - "message": {"role": "assistant", "content": "hi back"}, - "finish_reason": "stop", - } - ], - "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, - }, - }, - } - ).encode() - error_jsonl = json.dumps({"custom_id": "b", "error": {"message": "rejected"}}).encode() - return {"file-in": input_jsonl, "file-out": output_jsonl, "file-err": error_jsonl} - - @staticmethod - def _logging_obj() -> LitellmLogging: - obj = LitellmLogging( - model="gpt-4o", - messages=[{"role": "user", "content": "Hey"}], - stream=False, - call_type="aretrieve_batch", - start_time=1_767_225_600.0, - litellm_call_id="batch-call-reuse", - function_id="f", - ) - obj.custom_llm_provider = "openai" - obj.update_environment_variables( - litellm_params={"metadata": {}}, - optional_params={}, - custom_llm_provider="openai", - ) - return obj - - @staticmethod - def _batch(): - from litellm.types.utils import LiteLLMBatch - - return LiteLLMBatch( - id="batch_reuse_fetched_files", - completion_window="24h", - created_at=1, - endpoint="/v1/chat/completions", - input_file_id="file-in", - object="batch", - status="completed", - output_file_id="file-out", - error_file_id="file-err", - ) - - @pytest.mark.asyncio - async def test_compute_path_fetches_each_result_file_once(self, monkeypatch) -> None: - from litellm.batches.batch_line_item_logging import batch_line_item_claim_cache, pending_batch_line_item_tasks - - batch_line_item_claim_cache.in_memory_cache.flush_cache() - file_bytes = self._batch_file_bytes() - - async def fake_afile_content(**kwargs): - from types import SimpleNamespace - - return SimpleNamespace(content=file_bytes[kwargs["file_id"]]) - - file_mock = AsyncMock(side_effect=fake_afile_content) - monkeypatch.setattr(litellm, "store_batch_line_items_in_callbacks", True, raising=False) - monkeypatch.setattr("litellm.files.main.afile_content", file_mock) - monkeypatch.setattr("litellm.cost_calculator.batch_cost_calculator", lambda **kw: (0.01, 0.02)) - - await self._logging_obj()._async_success_handler_body(result=self._batch(), start_time=None, end_time=None) - await asyncio.gather(*pending_batch_line_item_tasks()) - - fetched = sorted(call.kwargs["file_id"] for call in file_mock.await_args_list) - assert fetched == ["file-err", "file-in", "file-out"], ( - "each result file must be fetched exactly once across aggregate costing and line-item logging" - ) - - @pytest.mark.asyncio - async def test_explicit_kwargs_path_forwards_result_files(self, monkeypatch) -> None: - from litellm.batches.batch_line_item_logging import batch_line_item_claim_cache, pending_batch_line_item_tasks - - batch_line_item_claim_cache.in_memory_cache.flush_cache() - from litellm.types.utils import Usage - - file_bytes = self._batch_file_bytes() - - async def fake_afile_content(**kwargs): - from types import SimpleNamespace - - return SimpleNamespace(content=file_bytes[kwargs["file_id"]]) - - file_mock = AsyncMock(side_effect=fake_afile_content) - monkeypatch.setattr(litellm, "store_batch_line_items_in_callbacks", True, raising=False) - monkeypatch.setattr("litellm.files.main.afile_content", file_mock) - - await self._logging_obj().async_success_handler( - result=self._batch(), - batch_cost=1.5, - batch_usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15), - batch_models=["gpt-4o"], - batch_successful_requests=1, - batch_failed_requests=1, - batch_output_file_content=file_bytes["file-out"], - batch_error_file_content=file_bytes["file-err"], - ) - await asyncio.gather(*pending_batch_line_item_tasks()) - - assert file_mock.await_count == 1 - assert file_mock.await_args.kwargs["file_id"] == "file-in", ( - "forwarded output/error bytes must leave only the input file to fetch" - ) - - @pytest.mark.asyncio - async def test_line_item_logging_failure_leaves_aggregate_result_intact(self, monkeypatch) -> None: - from litellm.batches.batch_line_item_logging import batch_line_item_claim_cache - - batch_line_item_claim_cache.in_memory_cache.flush_cache() - from litellm.types.utils import Usage - - monkeypatch.setattr(litellm, "store_batch_line_items_in_callbacks", True, raising=False) - monkeypatch.setattr( - "litellm.batches.batch_line_item_logging.log_batch_line_items", - AsyncMock(side_effect=RuntimeError("claim backend exploded")), - ) - - batch = self._batch() - await self._logging_obj().async_success_handler( - result=batch, - batch_cost=1.5, - batch_usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15), - batch_models=["gpt-4o"], - batch_successful_requests=1, - batch_failed_requests=1, - ) - - assert batch._hidden_params["response_cost"] == 1.5 - assert batch.usage.total_tokens == 15 - - -def test_caller_supplied_litellm_params_cannot_forge_logging_markers(logging_obj): - """A request body key `litellm_params` flows into optional_params; the logging - kwargs must keep the internal litellm_params or a caller could forge markers - (batch_parent_id) that spend sinks and router counters use to skip events.""" - from litellm.litellm_core_utils.core_helpers import is_batch_line_item_event - - logging_obj.update_environment_variables( - litellm_params={"metadata": {"model_group": "gpt-4o"}}, - optional_params={"litellm_params": {"batch_parent_id": "fake-batch"}}, - ) - - assert logging_obj.model_call_details["litellm_params"] is logging_obj.litellm_params - assert "batch_parent_id" not in logging_obj.model_call_details["litellm_params"] - assert not is_batch_line_item_event(logging_obj.model_call_details) - - -def test_a_top_level_batch_parent_id_request_field_is_not_a_line_item_marker(logging_obj): - """moderation, transcription and speech build litellm_params from every request kwarg, - so a body field named batch_parent_id lands in the litellm_params argument itself""" - logging_obj.update_environment_variables( - litellm_params={"metadata": {"model_group": "tts-1"}, BATCH_PARENT_ID_KEY: "batch_forged_by_caller"}, - optional_params={}, - ) - - assert BATCH_PARENT_ID_KEY not in logging_obj.model_call_details["litellm_params"] - assert logging_obj.model_call_details["litellm_params"]["metadata"] == {"model_group": "tts-1"} - assert not is_batch_line_item_event(logging_obj.model_call_details) - - -def test_a_marked_batch_line_item_keeps_its_batch_through_later_updates(logging_obj): - logging_obj.update_environment_variables(litellm_params={}, optional_params={}) - logging_obj.mark_batch_line_item("batch_abc") - - logging_obj.update_environment_variables( - litellm_params={BATCH_PARENT_ID_KEY: "batch_forged_by_caller"}, optional_params={} - ) - - assert logging_obj.model_call_details["litellm_params"][BATCH_PARENT_ID_KEY] == "batch_abc" - assert is_batch_line_item_event(logging_obj.model_call_details) - - -def test_hidden_params_carry_batch_line_fields_only_for_batch_line_items(): - plain: Final = StandardLoggingPayloadSetup.get_hidden_params({"model_id": "m"}) - line: Final = StandardLoggingPayloadSetup.get_hidden_params( - {"batch_id": "batch_abc", "batch_custom_id": "r1", "batch_line_status_code": 200} - ) - - assert {"batch_id", "batch_custom_id", "batch_line_status_code"}.isdisjoint(plain) - assert (line["batch_id"], line["batch_custom_id"], line["batch_line_status_code"]) == ("batch_abc", "r1", 200) - - def test_enterprise_alerting_loggers_resolves_real_classes(): from litellm_enterprise.enterprise_callbacks.pagerduty.pagerduty import PagerDutyAlerting from litellm_enterprise.enterprise_callbacks.send_emails.resend_email import ResendEmailLogger from litellm_enterprise.enterprise_callbacks.send_emails.sendgrid_email import SendGridEmailLogger from litellm_enterprise.enterprise_callbacks.send_emails.smtp_email import SMTPEmailLogger - from litellm.litellm_core_utils import litellm_logging as logging_module logging_module._enterprise_alerting_loggers.cache_clear() @@ -10004,7 +9734,6 @@ def test_enterprise_alerting_loggers_falls_back_without_disabling_callback_contr from litellm_enterprise.enterprise_callbacks.callback_controls import ( EnterpriseCallbackControls, ) - from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils import litellm_logging as logging_module @@ -10025,7 +9754,10 @@ def find_spec( ) logging_module._enterprise_alerting_loggers.cache_clear() for module_name in tuple(sys.modules): - if any(module_name == prefix or module_name.startswith(f"{prefix}.") for prefix in module_prefixes): + if any( + module_name == prefix or module_name.startswith(f"{prefix}.") + for prefix in module_prefixes + ): monkeypatch.delitem(sys.modules, module_name, raising=False) monkeypatch.setattr(sys, "meta_path", [BlockPagerDutyFinder(), *sys.meta_path]) @@ -10045,12 +9777,11 @@ def test_init_email_alerting_logger_reuses_instance( integration: Literal["smtp_email", "resend_email", "sendgrid_email"], monkeypatch: pytest.MonkeyPatch, ) -> None: + from litellm.litellm_core_utils import litellm_logging as logging_module from litellm_enterprise.enterprise_callbacks.send_emails.resend_email import ResendEmailLogger from litellm_enterprise.enterprise_callbacks.send_emails.sendgrid_email import SendGridEmailLogger from litellm_enterprise.enterprise_callbacks.send_emails.smtp_email import SMTPEmailLogger - from litellm.litellm_core_utils import litellm_logging as logging_module - logger_class: Final = ( SMTPEmailLogger if integration == "smtp_email" @@ -10084,10 +9815,9 @@ def test_init_email_alerting_logger_reuses_instance( def test_init_pagerduty_logger_reuses_instance(monkeypatch): + from litellm.types.integrations.pagerduty import AlertingConfig from litellm_enterprise.enterprise_callbacks.pagerduty.pagerduty import PagerDutyAlerting - from litellm.litellm_core_utils import litellm_logging as logging_module - from litellm.types.integrations.pagerduty import AlertingConfig monkeypatch.setenv("PAGERDUTY_API_KEY", "test-pagerduty-key") logging_module._in_memory_loggers.clear() @@ -10598,8 +10328,8 @@ def test_get_hidden_params(): assert result["response_cost"] is None assert result["additional_headers"] is None - # assert all required fields in StandardLoggingHiddenParams are present - assert all(field in result for field in StandardLoggingHiddenParams.__required_keys__) + # assert all fields in StandardLoggingHiddenParams are present + assert all(field in result for field in StandardLoggingHiddenParams.__annotations__) # Test with valid params hidden_params = { @@ -10619,8 +10349,8 @@ def test_get_hidden_params(): assert result["response_cost"] == 0.001 assert result["additional_headers"] is not None assert result["additional_headers"]["x_ratelimit_limit_requests"] == 2000 - # assert all required fields in StandardLoggingHiddenParams are present - assert all(field in result for field in StandardLoggingHiddenParams.__required_keys__) + # assert all fields in StandardLoggingHiddenParams are present + assert all(field in result for field in StandardLoggingHiddenParams.__annotations__) @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") def test_get_final_response_obj(): diff --git a/tests/unit/llms/bedrock/batches/test_transformation.py b/tests/unit/llms/bedrock/batches/test_transformation.py index 4a9b175714cf..5e987239c549 100644 --- a/tests/unit/llms/bedrock/batches/test_transformation.py +++ b/tests/unit/llms/bedrock/batches/test_transformation.py @@ -20,7 +20,7 @@ import pytest from litellm.llms.bedrock.batches.transformation import BedrockBatchesConfig -from litellm.types.utils import EmbeddingResponse, LlmProviders, ModelResponse +from litellm.types.utils import LlmProviders # AWS JobStatus -> OpenAI BatchJobStatus, exactly as encoded in transformation.py # (both transform_create_batch_response and transform_retrieve_batch_response). @@ -941,25 +941,3 @@ def test_retrieve_request_accepts_partition_arns(config: BedrockBatchesConfig, a batch_id=arn, optional_params={}, litellm_params={} ) assert result["url"].startswith(expected_prefix) - - -def test_transform_batch_output_line_dispatches_on_shape(config: BedrockBatchesConfig) -> None: - titan = config.transform_batch_output_line( - {"embedding": [0.1, 0.2], "inputTextTokenCount": 4}, model="amazon.titan-embed" - ) - assert isinstance(titan, EmbeddingResponse) - assert titan.data[0]["embedding"] == [0.1, 0.2] - assert titan.usage.prompt_tokens == 4 - - converse = config.transform_batch_output_line( - { - "output": {"message": {"role": "assistant", "content": [{"text": "hi"}]}}, - "stopReason": "end_turn", - "usage": {"inputTokens": 5, "outputTokens": 6, "totalTokens": 11}, - }, - model="amazon.nova-lite", - ) - assert isinstance(converse, ModelResponse) - assert converse.choices[0].message.content == "hi" - - assert config.transform_batch_output_line({"foo": 1}, model="x") is None 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 02ec62ad1c67..2295d86538cf 100644 --- a/tests/unit/proxy/common_utils/test_check_batch_cost.py +++ b/tests/unit/proxy/common_utils/test_check_batch_cost.py @@ -1446,10 +1446,10 @@ async def test_error_file_failures_add_to_failed_request_count( Logging, "async_success_handler", new_callable=AsyncMock ) as success_handler, ): - output_route = provider.get("https://api.openai.com/v1/files/file-output-123/content").mock( + provider.get("https://api.openai.com/v1/files/file-output-123/content").mock( return_value=httpx.Response(200, content=f"{succeeded_line}\n{rejected_line}\n".encode()) ) - error_route = provider.get("https://api.openai.com/v1/files/file-error-456/content").mock( + provider.get("https://api.openai.com/v1/files/file-error-456/content").mock( return_value=httpx.Response(200, content=f"{error_file_lines}\n\n".encode()) ) await check_batch_cost_instance.check_batch_cost() @@ -1458,10 +1458,6 @@ async def test_error_file_failures_add_to_failed_request_count( assert len(spend_log_calls) == 1 handler_kwargs = spend_log_calls[0] assert handler_kwargs["batch_successful_requests"] == 1 - assert output_route.call_count == 1 - assert error_route.call_count == 1 - assert handler_kwargs["batch_output_file_content"] == f"{succeeded_line}\n{rejected_line}\n".encode() - assert handler_kwargs["batch_error_file_content"] == f"{error_file_lines}\n\n".encode() assert handler_kwargs["batch_failed_requests"] == 3, ( "2 error-file lines must add to the output file's 1 rejected request" ) @@ -3225,3 +3221,307 @@ async def test_schema_without_batch_processed_still_bills(self): 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 + after the spend log is written, and a sink failure never unbills or re-polls the 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) + + @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']}") + 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_spend_log_of_the_claimed_batch(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() + + assert journal == [ + f"read {_CLAIM_OUTPUT_FILE_ID}", + "claim", + "bill", + "read file-input-1", + "lines", + "finalize", + ] + (payload,) = logger.payloads + assert payload["batch_id"] == _CLAIM_UNIFIED_BATCH_ID + assert payload["finished_at"] == 2 + (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_a_creator_that_cannot_be_looked_up_gets_lines_without_bodies(self, monkeypatch): + import litellm + + journal: list = [] + _, 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() + + (payload,) = logger.payloads + assert [(line["request"], line["response"]) for line in payload["lines"]] == [(None, None)] + + @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() + + ((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() + + assert [step for step in journal if not step.startswith("read")] == ["claim", "bill", "lines", "finalize"] + 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() + + 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_before_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() + + assert journal == ["read file-err-1", "read file-input-1", "lines", "finalize"] + (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_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() + + assert journal == ["finalize"] + assert logger.payloads == [] + + +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="hash-alice", 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"), + [ + ({"api_key": None}, None, None), + ({}, None, None), + ({"team_id": "team-gone"}, "present", None), + ], + ids=["legacy-row-without-key", "deleted-key", "deleted-team"], + ) + 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 + 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() diff --git a/tests/unit/proxy/hooks/test_model_max_budget_limiter.py b/tests/unit/proxy/hooks/test_model_max_budget_limiter.py index 5d5f30ffbfbc..2678f4205d58 100644 --- a/tests/unit/proxy/hooks/test_model_max_budget_limiter.py +++ b/tests/unit/proxy/hooks/test_model_max_budget_limiter.py @@ -311,17 +311,3 @@ async def test_a_batch_polled_within_every_budget_window_is_never_charged_again( await _poll(limiter, finished, BATCH_COST) assert _local_spend(limiter, KEY_SPEND_KEY) == pytest.approx(BATCH_COST) - - -@pytest.mark.asyncio -async def test_batch_line_item_events_do_not_charge_the_model_budget(): - """Line events carry batch_parent_id; the aggregate aretrieve_batch event is - the one that already bills the batch, so children must not double-charge.""" - limiter: Final = PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache()) - - line_event = _event("acompletion", CHAT_COST) - line_event["litellm_params"]["batch_parent_id"] = "batch_first" - await limiter.async_log_success_event(line_event, response_obj=None, start_time=None, end_time=None) - - assert await _spend(limiter, KEY_SPEND_KEY) == 0.0 - assert await _spend(limiter, USER_SPEND_KEY) == 0.0 diff --git a/tests/unit/proxy/hooks/test_parallel_request_limiter.py b/tests/unit/proxy/hooks/test_parallel_request_limiter.py index 50d1c8195a76..89cb6bd1ef21 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter.py @@ -166,32 +166,6 @@ async def test_async_log_success_event_counts_non_chat_response_tokens(response_ ) -@pytest.mark.asyncio -async def test_async_log_failure_event_skips_batch_line_item_events(): - """Failed line children were never admitted by the limiter, so the failure - hook must not decrement request counters they never incremented.""" - parallel_request_handler = PROXY_MaxParallelRequestsHandler( - internal_usage_cache=InternalUsageCache(DualCache()) - ) - local_cache = parallel_request_handler.internal_usage_cache.dual_cache.in_memory_cache - - await parallel_request_handler.async_log_failure_event( - kwargs={ - "exception": "litellm.APIError: upstream 500", - "litellm_params": { - "batch_parent_id": "batch_1", - "metadata": {"user_api_key": hash_token("sk-line-item")}, - }, - "model": "gpt-3.5-turbo", - }, - response_obj=None, - start_time=datetime(2026, 1, 1), - end_time=datetime(2026, 1, 1), - ) - - assert local_cache.cache_dict == {} - - @pytest.mark.asyncio async def test_a_pre_call_across_a_minute_rollover_lands_in_the_bucket_of_its_first_clock_read(): internal_usage_cache: Final = InternalUsageCache(DualCache()) diff --git a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py index 2467d20ae453..55163c6240b6 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py @@ -7793,59 +7793,3 @@ async def test_a2a_url_target_owns_invocation_fee_and_request_limit( await _rpm_request(limiter, cache, auth, "a2a/cheap") assert denied.value.status_code == 429 assert "expensive" in str(denied.value.detail) - - -@pytest.mark.asyncio -async def test_async_log_success_event_skips_batch_line_item_events(): - """Per-line batch events already ran through the limiter as the aggregate - aretrieve_batch; children must not increment TPM or request counters.""" - local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) - - await handler.async_log_success_event( - kwargs={ - "standard_logging_object": {"metadata": {"user_api_key_hash": hash_token("sk-line-item")}}, - "litellm_params": { - "batch_parent_id": "batch_1", - "metadata": {"user_api_key_hash": hash_token("sk-line-item"), "model_group": "gpt-3.5-turbo"}, - }, - "model": "gpt-3.5-turbo", - }, - response_obj=ModelResponse( - id="x", object="chat.completion", created=1, model="gpt-3.5-turbo", - usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15), choices=[], - ), - start_time=datetime(2026, 1, 1), - end_time=datetime(2026, 1, 1), - ) - - assert local_cache.in_memory_cache.cache_dict == {} - - -@pytest.mark.asyncio -async def test_async_log_failure_event_skips_batch_line_item_events(): - """Failed line children were never admitted by the limiter, so the failure - hook must not release slots or refund TPM they never reserved.""" - local_cache = DualCache() - handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache)) - stash = get_or_create_request_stash() - stash.owner_litellm_call_id = "call-line-item" - stash.reserved_tokens = 50 - stash.reserved_scopes = frozenset({("api_key", hash_token("sk-line-item"))}) - - await handler.async_log_failure_event( - kwargs={ - "litellm_call_id": "call-line-item", - "standard_logging_object": {"metadata": {"user_api_key_hash": hash_token("sk-line-item")}}, - "litellm_params": { - "batch_parent_id": "batch_1", - "metadata": {"user_api_key_hash": hash_token("sk-line-item"), "model_group": "gpt-3.5-turbo"}, - }, - "model": "gpt-3.5-turbo", - }, - response_obj=None, - start_time=datetime(2026, 1, 1), - end_time=datetime(2026, 1, 1), - ) - - assert local_cache.in_memory_cache.cache_dict == {} diff --git a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py index 52a3406d81ac..833af4c03a31 100644 --- a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py @@ -2857,24 +2857,6 @@ async def test_async_post_call_failure_hook_persists_no_raw_model_on_an_unknown_ assert error_information["error_class"] == "ProxyModelNotFoundError" -@pytest.mark.asyncio -async def test_batch_line_item_event_never_updates_spend(): # test-quality-ok: the observable contract is exactly that the spend path is never invoked - logger: Final = ProxyDBLogger() - kwargs: Final = { - "litellm_params": {"batch_parent_id": "batch_1", "metadata": {}}, - "model": "gpt-4o", - "call_type": "acompletion", - } - with patch.object(logger, "_PROXY_track_cost_callback", new_callable=AsyncMock) as mock_track: # test-quality-ok: asserts the callback's own method is skipped; the DB writer is never reached - await logger.async_log_success_event( - kwargs=kwargs, - response_obj=ModelResponse(), - start_time=datetime(2026, 1, 1), - end_time=datetime(2026, 1, 1), - ) - mock_track.assert_not_awaited() - - class _NeverStringifiedMetadataValue: def __repr__(self) -> str: raise AssertionError("a request metadata value was stringified by the cost tracking failure path") diff --git a/tests/unit/proxy/proxy_server/test_lifecycle.py b/tests/unit/proxy/proxy_server/test_lifecycle.py index 046c2b702022..69f9c72c9c82 100644 --- a/tests/unit/proxy/proxy_server/test_lifecycle.py +++ b/tests/unit/proxy/proxy_server/test_lifecycle.py @@ -223,38 +223,6 @@ async def _record_flush(client, accumulator): assert calls == ["flush", "disconnect"] -@pytest.mark.asyncio -async def test_proxy_shutdown_drains_batch_line_items_before_closing_clients(monkeypatch): - """A running batch line fan-out keeps its claim, so shutdown lets it finish - before the database, caches and callback clients it sends through close.""" - from litellm.constants import MAX_TIME_TO_CLEAR_QUEUE - - calls: list = [] # mutable-ok: records call order, which is the assertion - - async def _record_drain(timeout: float) -> int: - calls.append(("drain", timeout)) - return 0 - - monkeypatch.setattr(ps, "drain_batch_line_items", _record_drain, raising=False) - - fake_prisma = MagicMock() - fake_prisma.disconnect = AsyncMock(side_effect=lambda: calls.append("disconnect")) - monkeypatch.setattr(ps, "prisma_client", fake_prisma, raising=False) - monkeypatch.setattr(ps, "flush_gateway_requests", AsyncMock(), raising=False) - monkeypatch.setattr(ps, "flush_request_errors", AsyncMock(), raising=False) - monkeypatch.setattr(ps, "jwt_handler", MagicMock(close=AsyncMock()), raising=False) - monkeypatch.setattr(ps, "db_writer_client", None, raising=False) - - import litellm - - monkeypatch.setattr(litellm, "cache", None, raising=False) - monkeypatch.setattr(litellm, "success_callback", [], raising=False) - - await proxy_shutdown_event() - - assert calls == [("drain", MAX_TIME_TO_CLEAR_QUEUE), "disconnect"] - - @pytest.mark.asyncio async def test_proxy_shutdown_skips_gateway_flush_without_a_database(monkeypatch): """No prisma client means nothing to drain to, and no attempt is made.""" diff --git a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py index 0b0708955bb4..920f9d90bb8c 100644 --- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py +++ b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py @@ -8119,43 +8119,6 @@ async def test_update_general_settings_store_model_in_db_false(): assert ps.general_settings["store_model_in_db"] is False -@pytest.mark.asyncio -async def test_update_general_settings_store_batch_line_items_in_callbacks(): - """ - Verify _update_general_settings sets the litellm module flag when the DB - general_settings carries store_batch_line_items_in_callbacks, and that a - YAML-explicit value wins over the DB value. - """ - from litellm.proxy.proxy_server import ProxyConfig - - proxy_config = ProxyConfig() - saved_flag = litellm.store_batch_line_items_in_callbacks - try: - with patch("litellm.proxy.proxy_server.general_settings", {}): # test-quality-ok: module-global seam - await proxy_config._update_general_settings( - db_general_settings={"store_batch_line_items_in_callbacks": True} - ) - assert litellm.store_batch_line_items_in_callbacks is True - - proxy_config._yaml_general_settings_keys = {"store_batch_line_items_in_callbacks"} - with patch( # test-quality-ok: module-global seam - "litellm.proxy.proxy_server.general_settings", - {"store_batch_line_items_in_callbacks": "false"}, - ): - await proxy_config._update_general_settings( - db_general_settings={"store_batch_line_items_in_callbacks": True} - ) - assert litellm.store_batch_line_items_in_callbacks is False - - proxy_config._yaml_general_settings_keys = set() - litellm.store_batch_line_items_in_callbacks = True # test-quality-ok: seed the prior opt-in so its removal is observable - with patch("litellm.proxy.proxy_server.general_settings", {}): # test-quality-ok: module-global seam - await proxy_config._update_general_settings(db_general_settings={}) - assert litellm.store_batch_line_items_in_callbacks is False - finally: - litellm.store_batch_line_items_in_callbacks = saved_flag # test-quality-ok: restores the prior flag - - @pytest.mark.asyncio async def test_update_general_settings_propagates_apply_user_budget_to_team_keys(): """The Admin UI toggle writes to the DB config, so the flag has to be in the @@ -13332,7 +13295,6 @@ def _patched_coordination_redis_module_state( patch.object(proxy_server_module, "spend_counter_cache", spend_cache), patch.object(proxy_server_module, "user_api_key_cache", DualCache()), patch.object(proxy_server_module, "cli_sso_session_cache", DualCache()), - patch.object(proxy_server_module, "batch_line_item_claim_cache", DualCache()), patch.object(proxy_server_module, "llm_router", None), patch.object(proxy_server_module, "litellm_config_cache", config_cache), patch.object(proxy_server_module, "RedisCache", redis_cache_class), diff --git a/tests/unit/proxy/test_redis_auth_cache_flag.py b/tests/unit/proxy/test_redis_auth_cache_flag.py index 398dd1efa970..cb600bbb5fda 100644 --- a/tests/unit/proxy/test_redis_auth_cache_flag.py +++ b/tests/unit/proxy/test_redis_auth_cache_flag.py @@ -74,7 +74,6 @@ def _patched_init_cache(litellm_settings: dict, cache_params: dict): patch.object(ps, "spend_counter_cache", fresh_spend_cache), patch.object(ps, "cli_sso_session_cache", fresh_cli_sso_cache), patch.object(ps, "litellm_config_cache", fresh_config_cache), - patch.object(ps, "batch_line_item_claim_cache", DualCache()), patch.object(ps, "llm_router", None), # Cache is locally imported inside _init_cache: patch it at source. patch("litellm.Cache", return_value=mock_litellm_cache), @@ -185,12 +184,10 @@ def test_flag_absent_still_shares_the_model_budget_counters_over_redis(self): ps.cli_sso_session_cache, ps.user_api_key_cache, ps.litellm_config_cache, - ps.batch_line_item_claim_cache, ) with ExitStack() as detached: for cache in touched_caches: detached.enter_context(patch.object(cache, "redis_cache", None)) ps._attach_redis_usage_cache(fake_redis, enable_redis_auth_cache=False) assert limiter_cache.redis_cache is fake_redis - assert ps.batch_line_item_claim_cache.redis_cache is fake_redis assert ps.user_api_key_cache.redis_cache is None diff --git a/tests/unit/router_strategy/test_budget_limiter_hotpath.py b/tests/unit/router_strategy/test_budget_limiter_hotpath.py index 80e1862e20f5..4667ac67efa2 100644 --- a/tests/unit/router_strategy/test_budget_limiter_hotpath.py +++ b/tests/unit/router_strategy/test_budget_limiter_hotpath.py @@ -391,29 +391,6 @@ async def test_push_task_failure_is_logged_once_and_not_leaked(disable_budget_sy unretrieved.assert_not_called() -@pytest.mark.asyncio -async def test_batch_line_item_events_do_not_charge_the_provider_budget(disable_budget_sync): - limiter = RouterBudgetLimiting( - dual_cache=DualCache(), - provider_budget_config={"openai": BudgetConfig(max_budget=100.0, budget_duration="1d")}, - ) - await asyncio.gather(*(task for task in asyncio.all_tasks() if task is not asyncio.current_task())) - - await limiter.async_log_success_event( - kwargs={ - "call_type": "acompletion", - "litellm_params": {"custom_llm_provider": "openai", "batch_parent_id": "batch_x"}, - "standard_logging_object": {"response_cost": 1.0, "model_id": "dep-1"}, - }, - response_obj=None, - start_time=None, - end_time=None, - ) - - assert limiter.dual_cache.in_memory_cache.get_cache("provider_spend:openai:1d") in (None, 0) - assert limiter.redis_increment_operation_queue == [] - - _SPEND_KEY = "provider_spend:openai:1d" diff --git a/tests/unit/router_strategy/test_least_busy.py b/tests/unit/router_strategy/test_least_busy.py index c8e936fd57ec..5b3af388ae53 100644 --- a/tests/unit/router_strategy/test_least_busy.py +++ b/tests/unit/router_strategy/test_least_busy.py @@ -219,31 +219,6 @@ async def test_an_open_circuit_breaker_falls_back_without_a_warning_per_request( assert sum("circuit breaker is open" in record.getMessage() for record in caplog.records) == 2 -@pytest.mark.asyncio -async def test_batch_line_item_success_does_not_release_in_flight_slots(monkeypatch: pytest.MonkeyPatch) -> None: - """A batch line item is historical batch traffic (call_type=acompletion + - litellm_params.batch_parent_id): it never started a live request, so its success - callback must not decrement the in-flight count a live request still holds.""" - import litellm - - monkeypatch.setattr(litellm, "store_batch_line_items_in_callbacks", True, raising=False) - shared: Final = SharedRedisCounters() - worker: Final = _worker(shared) - worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) - - line_item_kwargs: Final = { - "call_type": "acompletion", - "litellm_params": { - "batch_parent_id": "batch-1", - "metadata": {"model_group": GROUP}, - "model_info": {"id": "dep-a"}, - }, - } - await worker.async_log_success_event(line_item_kwargs, None, None, None) - - assert shared.count(f"{GROUP}_request_count:dep-a") == 1 - - def test_model_added(): test_cache = DualCache() least_busy_logger = LeastBusyLoggingHandler(router_cache=test_cache) diff --git a/tests/unit/router_strategy/test_lowest_tpm_rpm.py b/tests/unit/router_strategy/test_lowest_tpm_rpm.py index 15d6dd3a9905..f1c56b9d033b 100644 --- a/tests/unit/router_strategy/test_lowest_tpm_rpm.py +++ b/tests/unit/router_strategy/test_lowest_tpm_rpm.py @@ -145,43 +145,6 @@ async def async_get_available_deployments( } -@pytest.mark.asyncio -async def test_usage_based_routing_handlers_skip_batch_line_items(monkeypatch: pytest.MonkeyPatch) -> None: - """ - Batch line-item callbacks carry call_type=acompletion plus - litellm_params.batch_parent_id, so the call-type-only batch guard does not fire. - They are historical batch traffic and must not update the TPM/RPM usage counters. - """ - import litellm - from litellm.router_strategy.lowest_tpm_rpm import LowestTPMLoggingHandler - - monkeypatch.setattr(litellm, "store_batch_line_items_in_callbacks", True, raising=False) - router: Final = Router( - model_list=[_deployment(HIGH_USAGE_DEPLOYMENT_ID)], - routing_strategy="usage-based-routing", - ) - handler: Final = LowestTPMLoggingHandler(router_cache=router.cache, routing_args={}) - line_item_kwargs: Final = { - "call_type": "acompletion", - "litellm_params": { - "batch_parent_id": "batch-1", - "metadata": {"model_group": MODEL_GROUP}, - "model_info": {"id": HIGH_USAGE_DEPLOYMENT_ID}, - }, - } - response_obj: Final = {"usage": {"total_tokens": 600}} - - handler.log_success_event(line_item_kwargs, response_obj, None, None) - await handler.async_log_success_event(line_item_kwargs, response_obj, None, None) - - moved: Final = sorted( - f"{key}={router.cache.in_memory_cache.cache_dict[key]}" - for key in router.cache.in_memory_cache.cache_dict - if ":tpm:" in key or ":rpm:" in key - ) - assert moved == [] - - def _rate_limited_router(num_allowed_send: int) -> tuple[Router, tuple[list[dict[str, str]], ...]]: conversations: Final = tuple( [{"role": "user", "content": f"{index}. Hey, how's it going?"}] for index in range(num_allowed_send) diff --git a/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py b/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py index ebf5db37c32c..c6ab259991e1 100644 --- a/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py @@ -280,41 +280,6 @@ async def test_async_filter_deployments_narrows_for_group_whose_model_minimum_is AUTO_CACHING_MODEL = "anthropic/claude-sonnet-4-5" -@pytest.mark.asyncio -async def test_batch_line_item_event_does_not_pin_a_deployment(): - """A completed batch's JSONL lines are logged as events of their own, but the - batch API served them, not the deployment's realtime endpoint, so a line must - not record a cached prefix that later pins live traffic to that deployment.""" - cache = DualCache() - check = PromptCachingDeploymentCheck(cache=cache) - model = AUTO_CACHING_MODEL - deployments = _deployments(model, model, model) - messages = _messages(word_count=3000) - assert is_prompt_caching_valid_prompt(model=model, messages=messages) is True - - await check.async_log_success_event( - kwargs={ - "litellm_params": {"batch_parent_id": "batch_1"}, - "standard_logging_object": { - "call_type": "acompletion", - "model": model, - "messages": messages, - "model_id": "dep-2", - }, - }, - response_obj=None, - start_time=None, - end_time=None, - ) - filtered = await check.async_filter_deployments( - model=MODEL_GROUP_ALIAS, - healthy_deployments=deployments, - messages=messages, - ) - - assert filtered == deployments - - @pytest.mark.asyncio async def test_replayed_redacted_thinking_block_still_records_and_pins(): """ diff --git a/tests/unit/test_router/test_enforce_model_rate_limits.py b/tests/unit/test_router/test_enforce_model_rate_limits.py index 5e382a2e4a38..7577064b7f91 100644 --- a/tests/unit/test_router/test_enforce_model_rate_limits.py +++ b/tests/unit/test_router/test_enforce_model_rate_limits.py @@ -15,10 +15,6 @@ from litellm import Router from litellm.caching.dual_cache import DualCache from litellm.caching.redis_cache import RedisCircuitBreakerOpenError -from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import ( - ITPM_CACHE_KEY, - ITPM_RESERVED_KEY, -) from litellm.router_utils.pre_call_checks.model_rate_limit_check import ( ModelRateLimitingCheck, ) @@ -385,33 +381,6 @@ async def test_async_log_success_event_increments_cache(self): assert "test-id:gpt-4:tpm:" in kwarg_params["key"] assert kwarg_params["value"] == 50 - @pytest.mark.asyncio - async def test_batch_line_item_events_leave_router_counters_untouched(self): - """A completed batch's JSONL lines are logged as events of their own, and their - metadata comes from the caller's input file, so a line event must not count - tokens against a deployment or settle a reservation the metadata names.""" - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) - itpm_counter = "global_router:other-id:gpt-4:itpm:00-00" - await dual_cache.async_increment_cache(key=itpm_counter, value=8, ttl=60) - minute = litellm.utils.get_utc_datetime().strftime("%H-%M") - line_kwargs = { - "litellm_params": {"batch_parent_id": "batch_1"}, - "metadata": {ITPM_RESERVED_KEY: 5, ITPM_CACHE_KEY: itpm_counter}, - "standard_logging_object": { - "model_id": "test-id", - "total_tokens": 50, - "hidden_params": {"litellm_model_name": "gpt-4"}, - "metadata": {}, - }, - } - - await check.async_log_success_event(line_kwargs, None, None, None) - await check.async_log_failure_event(line_kwargs, None, None, None) - - assert await dual_cache.async_get_cache(key=itpm_counter) == 8 - assert await dual_cache.async_get_cache(key=f"test-id:gpt-4:tpm:{minute}") is None - class TestRouterWithEnforceModelRateLimits: """Test Router integration with enforce_model_rate_limits.""" diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 442e777c0196..4a7a23ff57b5 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -1305,53 +1305,6 @@ def test_sync_deployment_callback_on_success_skips_batch_retrieves( == expected_successes ) - -@pytest.mark.asyncio -async def test_deployment_callbacks_skip_batch_line_items(monkeypatch: pytest.MonkeyPatch): - """ - Batch line-item callbacks carry call_type=acompletion plus - litellm_params.batch_parent_id, so the call-type-only batch guards do not fire. - They describe historical batch traffic already reported by the aggregate - aretrieve_batch event and must not consume live TPM/RPM quota. - """ - monkeypatch.setattr(litellm, "store_batch_line_items_in_callbacks", True, raising=False) - router = litellm.Router( - model_list=[ - { - "model_name": _BATCH_GROUP, - "litellm_params": {"model": _BATCH_DEPLOYMENT_MODEL, "api_base": _BATCH_API_BASE, "api_key": "sk-fake"}, - "model_info": {"id": "batch-dep"}, - } - ] - ) - now = datetime(2026, 1, 1) - line_item_kwargs = { - "call_type": "acompletion", - "standard_logging_object": {"total_tokens": _BATCH_TOKENS_PER_ROW}, - "litellm_params": { - "batch_parent_id": _BATCH_ID, - "metadata": {"model_group": _BATCH_GROUP, "deployment": _BATCH_DEPLOYMENT_MODEL}, - "model_info": {"id": "batch-dep"}, - }, - } - - await router.deployment_callback_on_success( - kwargs=line_item_kwargs, completion_response=None, start_time=now, end_time=now - ) - sync_key = router.sync_deployment_callback_on_success( - kwargs=line_item_kwargs, completion_response=None, start_time=now, end_time=now - ) - await router.async_deployment_callback_on_failure( - kwargs=line_item_kwargs, completion_response=None, start_time=now, end_time=now - ) - - assert sync_key is None - assert ( - get_deployment_successes_for_current_minute(litellm_router_instance=router, deployment_id="batch-dep") == 0 - ) - assert await _moved_routing_counters(router) == [] - assert await _router_usage_keys(router) == [] - _ROUTING_STRATEGY_CACHE_MARKERS = ("_map", "_request_count", ":tpm:", ":rpm:") @@ -1439,81 +1392,6 @@ async def test_arouter_aretrieve_batch_does_not_feed_routing_strategies( assert moved_counters == [] -_BATCH_INPUT_JSONL = "\n".join( - json.dumps( - { - "custom_id": f"row-{row}", - "method": "POST", - "url": "/v1/chat/completions", - "body": {"model": "gpt-4o-mini", "messages": [{"role": "user", "content": "hi"}]}, - } - ) - for row in range(_BATCH_ROWS) -) - - -def _mock_batch_input_file(respx_mock): - respx_mock.get(f"{_BATCH_API_BASE}/files/file-in-1/content").mock( - return_value=httpx.Response(200, text=_BATCH_INPUT_JSONL) - ) - - -async def _line_item_payloads(collector: "_BatchPayloadCollector", minimum: int, timeout: float = 5.0) -> list: - loop = asyncio.get_event_loop() - deadline = loop.time() + timeout - while loop.time() < deadline: - line_items = [p for p in collector.payloads if p and p.get("call_type") == "acompletion"] - if len(line_items) >= minimum: - return line_items - await asyncio.sleep(0.05) - raise AssertionError(f"expected at least {minimum} batch line-item payloads, saw {len(collector.payloads)} total") - - -@pytest.mark.parametrize( - "routing_strategy", - [ - "usage-based-routing", - "usage-based-routing-v2", - "latency-based-routing", - "cost-based-routing", - "least-busy", - ], -) -@pytest.mark.asyncio -async def test_arouter_aretrieve_batch_line_items_do_not_feed_routing_strategies( - monkeypatch: pytest.MonkeyPatch, routing_strategy: str -): - """ - With the line-item flag on, retrieving a completed batch also emits one child - callback per JSONL line (call_type=acompletion + litellm_params.batch_parent_id, - inheriting the parent deployment's model_info). Those children are historical - batch traffic: like the aggregate poll, they must not move the counters that - decide where the next live chat request goes. - """ - from litellm.batches.batch_line_item_logging import batch_line_item_claim_cache - - batch_line_item_claim_cache.in_memory_cache.flush_cache() - monkeypatch.setattr(litellm, "store_batch_line_items_in_callbacks", True, raising=False) - collector = _BatchPayloadCollector() - monkeypatch.setattr(litellm, "callbacks", [collector]) - monkeypatch.setattr(litellm, "input_callback", []) - router = _batch_fan_out_router(routing_strategy) - - with respx.mock(assert_all_called=True) as respx_mock: - _mock_batch_provider(respx_mock) - _mock_batch_input_file(respx_mock) - respx_mock.get(f"{_UNRELATED_BATCH_API_BASE}/batches/{_BATCH_ID}").mock( - return_value=httpx.Response(404, json=_BATCH_NOT_FOUND) - ) - response = await router.aretrieve_batch(batch_id=_BATCH_ID) - line_items = await _line_item_payloads(collector, minimum=_BATCH_ROWS) - moved_counters = await _moved_routing_counters(router) - - assert response.id == _BATCH_ID - assert len(line_items) == _BATCH_ROWS - assert moved_counters == [] - - @pytest.mark.asyncio async def test_arouter_aretrieve_file_content(): """ diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 5cd94d5adacd..7343c13b409b 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -29894,11 +29894,6 @@ export interface components { spend_capture_rate_check?: components["schemas"]["SpendCaptureRateCheckSettings"] | null; /** @description Which keys of LiteLLM_SpendLogs.metadata are written to the database. Set exactly one of 'include' (write only these keys) or 'exclude' (drop these keys). 'status' and 'cold_storage_object_key' are always written. Daily spend tables, budgets and logging callbacks still see every key. Unset writes every key */ spend_logs_metadata_fields?: components["schemas"]["SpendLogsMetadataFields"] | null; - /** - * Store Batch Line Items In Callbacks - * @description If True, a completed batch logged via aretrieve_batch also emits one callback event per JSONL line item (request paired with its response or error). The aggregate batch callback is unchanged. Default is False. - */ - store_batch_line_items_in_callbacks?: boolean | null; /** * Store Model In Db * @description If True, models and config are stored in and loaded from the database. Default is False. From 7ab9a3f6297849c07b85bfd96968879bb553a663 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Tue, 13 Oct 2026 02:35:00 -0700 Subject: [PATCH 76/89] fix(batches): store batch line items after the batch is billed and retired Send a billed batch's lines only after its row is marked complete, so a slow or failing sink never holds a claimed batch open. Skip reading the input file when the bodies are dropped anyway, still send the error lines of a batch whose output file is gone, and end every NDJSON record with a newline. --- .../proxy/common_utils/check_batch_cost.py | 33 ++++++++------ litellm/batches/batch_line_items.py | 4 +- .../integrations/batch_line_item_objects.py | 2 +- tests/unit/batches/test_batch_line_items.py | 5 ++- .../gcs_bucket/test_gcs_bucket.py | 2 +- .../common_utils/test_check_batch_cost.py | 44 ++++++++++++++----- 6 files changed, 60 insertions(+), 30 deletions(-) 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 7643ac63c027..d7cc6550c036 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -732,10 +732,10 @@ async def _track_completed_batch_cost( model_id: str, batch_id: str, prom_logger: Optional["PrometheusLogger"], - ) -> Optional[Tuple[Optional[str], Optional[str]]]: + ) -> tuple[str | None, str | None, bytes] | None: """ Fetch a completed batch's results, compute cost/usage, and emit the - aretrieve_batch spend log. Returns (model_name, llm_provider) on + aretrieve_batch spend log. Returns (model_name, llm_provider, output file bytes) on success, None when the job can't be routed to a deployment or when another pod claimed it. Raises on results-fetch or cost-computation failures so the caller can leave the job unprocessed and retry it on a @@ -962,9 +962,7 @@ async def _track_completed_batch_cost( api_provider=str(llm_provider) if llm_provider else None, ) - await self._store_batch_line_items(job, response, model_id, batch_id, content_bytes) - - return model_name, str(llm_provider) if llm_provider else None + return model_name, str(llm_provider) if llm_provider else None, content_bytes async def _store_batch_line_items( self, @@ -974,8 +972,8 @@ async def _store_batch_line_items( batch_id: str, output_content: bytes, ) -> None: - """Send every line of a finished batch to batch_line_item_callbacks. Failures are logged, never raised, - so the batch stays billed whether or not its lines were stored.""" + """Send every line of a finished batch to batch_line_item_callbacks. Failures are logged, 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 ( @@ -1006,6 +1004,7 @@ async def _store_batch_line_items( 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 @@ -1019,12 +1018,17 @@ async def _store_batch_line_items( model_info=deployment_pricing_model_info(model_id=model_id, deployment_model=litellm_model_name), metadata=line_item_metadata(await self._build_creator_attribution_metadata(job, batch_id)), ) - payload: Final = await build_batch_line_items( - source, - await self._batch_file_content(response.input_file_id, cost_provider, file_access), - (output_content, error_content), + input_content: Final = ( + b"" + if redact_bodies + else await self._batch_file_content(response.input_file_id, cost_provider, file_access) + ) + payload: Final = await build_batch_line_items(source, input_content, (output_content, error_content)) + await dispatch_batch_line_items(payload, redact_bodies=redact_bodies) + verbose_proxy_logger.info( + f"CheckBatchCost: sent the {len(payload['lines'])} line items of batch {batch_id} " + "to batch_line_item_callbacks" ) - await dispatch_batch_line_items(payload, redact_bodies=await self._creator_turned_off_message_logging(job)) 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}" @@ -1198,6 +1202,7 @@ async def check_batch_cost(self): 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._store_batch_line_items(job, response, model_id, batch_id, b"") await self._finalize_unbilled_terminal_job(job, response) continue verbose_proxy_logger.error( @@ -1208,9 +1213,10 @@ async def check_batch_cost(self): continue if tracked is None: continue + tracked_model_name, tracked_provider, output_content = tracked # Track this job for the final metrics summary - processed_models.append(tracked) + processed_models.append((tracked_model_name, tracked_provider)) # mark the job as complete try: @@ -1228,6 +1234,7 @@ async def check_batch_cost(self): verbose_proxy_logger.error( f"CheckBatchCost: failed to mark job {job.id} complete in DB: {db_err}" ) + await self._store_batch_line_items(job, response, model_id, batch_id, output_content) elif response.status in PROVIDER_TERMINAL_BATCH_STATUSES: from litellm.proxy.openai_files_endpoints.common_utils import ( diff --git a/litellm/batches/batch_line_items.py b/litellm/batches/batch_line_items.py index a2ec6b5fd5fe..93ec1daca303 100644 --- a/litellm/batches/batch_line_items.py +++ b/litellm/batches/batch_line_items.py @@ -120,7 +120,7 @@ def without_request_and_response(payload: StandardBatchLineItemsPayload) -> Stan async def dispatch_batch_line_items(payload: StandardBatchLineItemsPayload, redact_bodies: bool) -> None: - """Hand the lines to every batch_line_item_callbacks logger. A logger that fails is logged and never stops the rest.""" + """Hand the 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) ) @@ -143,7 +143,7 @@ async def _deliver( if redact_bodies or _logger_turned_off_message_logging(logger) else payload ) - except Exception as error: # noqa: BLE001 # one sink failing must not stop the other sinks or the batch cost poller + 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 batch %s: %s", callback, diff --git a/litellm/integrations/batch_line_item_objects.py b/litellm/integrations/batch_line_item_objects.py index f9312bdb60d5..23dd2e3200e4 100644 --- a/litellm/integrations/batch_line_item_objects.py +++ b/litellm/integrations/batch_line_item_objects.py @@ -8,7 +8,7 @@ def batch_line_items_ndjson(payload: StandardBatchLineItemsPayload) -> str: - return "\n".join(json.dumps(line, default=str, ensure_ascii=False) for line in payload["lines"]) + return "".join(f"{json.dumps(line, default=str, ensure_ascii=False)}\n" for line in payload["lines"]) def batch_line_items_object_name(payload: StandardBatchLineItemsPayload) -> str: diff --git a/tests/unit/batches/test_batch_line_items.py b/tests/unit/batches/test_batch_line_items.py index 20211d370d1f..02ad67195b72 100644 --- a/tests/unit/batches/test_batch_line_items.py +++ b/tests/unit/batches/test_batch_line_items.py @@ -205,7 +205,10 @@ async def test_redaction_drops_only_the_request_and_response(): async def test_the_object_holds_one_json_line_per_batch_line_under_its_finished_day(): payload = await _payload() - assert [json.loads(line) for line in batch_line_items_ndjson(payload).splitlines()] == list(payload["lines"]) + 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.jsonl" assert ( batch_line_items_object_name({**payload, "batch_id": "projects/p/locations/l/batchPredictionJobs/9"}) diff --git a/tests/unit/integrations/gcs_bucket/test_gcs_bucket.py b/tests/unit/integrations/gcs_bucket/test_gcs_bucket.py index d11ee75d85ab..de2fb6c0f3fb 100644 --- a/tests/unit/integrations/gcs_bucket/test_gcs_bucket.py +++ b/tests/unit/integrations/gcs_bucket/test_gcs_bucket.py @@ -359,7 +359,7 @@ def stored(request: httpx.Request) -> httpx.Response: @pytest.mark.asyncio -async def test_a_rejected_batch_line_items_write_raises_for_the_dispatcher_to_log(monkeypatch: pytest.MonkeyPatch) -> None: +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): 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 2295d86538cf..a862a26c5fd9 100644 --- a/tests/unit/proxy/common_utils/test_check_batch_cost.py +++ b/tests/unit/proxy/common_utils/test_check_batch_cost.py @@ -3225,7 +3225,7 @@ async def test_schema_without_batch_processed_still_bills(self): class TestBatchLineItemsFromThePoller: """When batch_line_item_callbacks is set, the poller hands every line of a finished batch to those loggers - after the spend log is written, and a sink failure never unbills or re-polls the batch.""" + once the batch is billed and retired, so a slow or failing sink never unbills or re-polls the batch.""" _INPUT = b'{"custom_id":"ok-1","body":{"messages":[{"role":"user","content":"hi"}]}}\n' _OUTPUT = ( @@ -3271,6 +3271,12 @@ def _provider(cls, journal: list): 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 @@ -3298,7 +3304,7 @@ def _is_b64(file_id): yield @pytest.mark.asyncio - async def test_lines_reach_the_callbacks_after_the_spend_log_of_the_claimed_batch(self, monkeypatch): + async def test_lines_reach_the_callbacks_after_the_claimed_batch_is_billed_and_retired(self, monkeypatch): import litellm journal: list = [] @@ -3309,14 +3315,7 @@ async def test_lines_reach_the_callbacks_after_the_spend_log_of_the_claimed_batc with self._provider(journal): await instance.check_batch_cost() - assert journal == [ - f"read {_CLAIM_OUTPUT_FILE_ID}", - "claim", - "bill", - "read file-input-1", - "lines", - "finalize", - ] + assert journal == [f"read {_CLAIM_OUTPUT_FILE_ID}", "claim", "bill", "finalize", "lines"] (payload,) = logger.payloads assert payload["batch_id"] == _CLAIM_UNIFIED_BATCH_ID assert payload["finished_at"] == 2 @@ -3340,6 +3339,7 @@ async def test_a_creator_that_cannot_be_looked_up_gets_lines_without_bodies(self (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): @@ -3355,6 +3355,7 @@ async def test_a_creator_with_logging_on_gets_the_request_and_response(self, mon with self._provider(journal): await instance.check_batch_cost() + assert journal[-3:] == ["finalize", "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) @@ -3372,7 +3373,7 @@ async def test_a_failing_sink_leaves_the_batch_billed_and_retired(self, monkeypa with self._provider(journal): await instance.check_batch_cost() - assert [step for step in journal if not step.startswith("read")] == ["claim", "bill", "lines", "finalize"] + 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" @@ -3401,13 +3402,32 @@ async def test_a_batch_where_every_line_failed_sends_its_error_lines_before_it_i with self._provider(journal): await instance.check_batch_cost() - assert journal == ["read file-err-1", "read file-input-1", "lines", "finalize"] + assert journal == ["read file-err-1", "lines", "finalize"] (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_whose_output_file_is_gone_still_sends_its_error_lines_before_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() + + assert journal == ["read file-out-deleted", "read file-err-1", "lines", "finalize"] + (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 From 5cad62022bf7750a0e33ee73b8f36bb60ebe7fff Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Tue, 13 Oct 2026 03:00:02 -0700 Subject: [PATCH 77/89] test(batches): cover team-level and deleted-key redaction of batch line items --- .../test_batch_line_items_cold_storage.py | 88 +++++++++++++++---- 1 file changed, 69 insertions(+), 19 deletions(-) diff --git a/tests/integration/observability/test_batch_line_items_cold_storage.py b/tests/integration/observability/test_batch_line_items_cold_storage.py index 377c7dc6ff89..634e133ab9eb 100644 --- a/tests/integration/observability/test_batch_line_items_cold_storage.py +++ b/tests/integration/observability/test_batch_line_items_cold_storage.py @@ -236,7 +236,7 @@ class Lines: errors: tuple[dict[str, JsonValue], ...] -def _routes(lines: Lines) -> RoutedResponse: +def _routes(lines: Lines, *, finished: bool = True) -> RoutedResponse: has_output: Final = bool(lines.outputs) completed: Final[dict[str, JsonValue]] = { "id": "batch-$REQUEST_ID", @@ -267,7 +267,12 @@ def _routes(lines: Lines) -> RoutedResponse: 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), + "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) ), @@ -290,8 +295,8 @@ class SubmittedBatch: handle: ScenarioHandle -def _submit(rig: Rig, lines: Lines, key_fields: Mapping[str, JsonValue]) -> SubmittedBatch: - handle: Final = register_scenario(f"batch-line-items-{uuid.uuid4().hex}", _routes(lines)) +def _submit(rig: Rig, lines: Lines, key_fields: Mapping[str, JsonValue], *, finished: bool = True) -> SubmittedBatch: + handle: Final = register_scenario(f"batch-line-items-{uuid.uuid4().hex}", _routes(lines, finished=finished)) model: Final = string_value( rig.gateway.post( "/model/new", @@ -330,6 +335,18 @@ def _submit(rig: Rig, lines: Lines, key_fields: Mapping[str, JsonValue]) -> Subm 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) -> str: return f"{prefix}batch_line_items/{LINE_DATE}/{batch_id.replace('/', '_')}.jsonl" @@ -413,21 +430,7 @@ def test_a_key_that_turned_off_message_logging_gets_its_lines_without_request_or lines: Final = Lines( inputs=(_input_line("quiet-1", "a private question"),), outputs=(_success_line("quiet-1"),), errors=() ) - batch: Final = _submit( - rig, - lines, - { - "metadata": { - "logging": [ - { - "callback_name": "s3_v2", - "callback_type": "success", - "callback_vars": {"turn_off_message_logging": "true"}, - } - ] - } - }, - ) + 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) @@ -447,6 +450,53 @@ def test_a_key_that_turned_off_message_logging_gets_its_lines_without_request_or 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_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")), From aa76a2eed8e00882fb085183ce229cba53127841 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Tue, 13 Oct 2026 04:14:27 -0700 Subject: [PATCH 78/89] fix(batches): store batch line items in the background so a hung bucket never holds up billing The poller now retires a batch first and hands its line items to a per-worker background store that runs one batch at a time. A slow, hung or failing bucket no longer delays billing or retiring the other batches in the poll cycle. While a poll cycle's worth of batches is still waiting to store, the next batch's lines are dropped with an error. --- .../proxy/common_utils/check_batch_cost.py | 45 +++++- .../test_batch_line_items_cold_storage.py | 141 ++++++++++++++++-- .../common_utils/test_check_batch_cost.py | 105 ++++++++++++- 3 files changed, 272 insertions(+), 19 deletions(-) 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 d7cc6550c036..fa1338c25829 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -2,6 +2,7 @@ Polls LiteLLM_ManagedObjectTable to check if the batch job is complete, and if the cost has been tracked. """ +import asyncio from collections.abc import Mapping, Sequence from dataclasses import replace as dataclasses_replace from datetime import datetime, timedelta, timezone @@ -130,6 +131,8 @@ def __init__( # 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() @staticmethod def _is_missing_batch_processed_column_error(err: Exception) -> bool: @@ -964,6 +967,42 @@ async def _track_completed_batch_cost( return model_name, str(llm_provider) if llm_provider else None, content_bytes + def _store_batch_line_items_in_background( + self, + job: "_ManagedObjectRow", + response: "LiteLLMBatch", + model_id: str, + batch_id: str, + output_content: bytes, + ) -> None: + """Store a retired batch's lines off the poll loop, one batch at a time, so a slow or hung bucket never 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, output_content) + ) + 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, + output_content: bytes, + ) -> None: + async with self._line_item_store_turn: + await self._store_batch_line_items(job, response, model_id, batch_id, output_content) + async def _store_batch_line_items( self, job: "_ManagedObjectRow", @@ -1202,8 +1241,8 @@ async def check_batch_cost(self): 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._store_batch_line_items(job, response, model_id, batch_id, b"") await self._finalize_unbilled_terminal_job(job, response) + self._store_batch_line_items_in_background(job, response, model_id, batch_id, b"") continue verbose_proxy_logger.error( f"CheckBatchCost: failed to track cost for batch {batch_id} " @@ -1234,7 +1273,7 @@ async def check_batch_cost(self): verbose_proxy_logger.error( f"CheckBatchCost: failed to mark job {job.id} complete in DB: {db_err}" ) - await self._store_batch_line_items(job, response, model_id, batch_id, output_content) + self._store_batch_line_items_in_background(job, response, model_id, batch_id, output_content) elif response.status in PROVIDER_TERMINAL_BATCH_STATUSES: from litellm.proxy.openai_files_endpoints.common_utils import ( @@ -1247,8 +1286,8 @@ async def check_batch_cost(self): f"has not appeared yet; leaving job {job.id} for the next poll cycle" ) continue - await self._store_batch_line_items(job, response, model_id, batch_id, b"") await self._finalize_unbilled_terminal_job(job, response) + self._store_batch_line_items_in_background(job, response, model_id, batch_id, b"") # Record polling run metrics (always, even if nothing was processed) if prom_logger: diff --git a/tests/integration/observability/test_batch_line_items_cold_storage.py b/tests/integration/observability/test_batch_line_items_cold_storage.py index 634e133ab9eb..0fff743572b9 100644 --- a/tests/integration/observability/test_batch_line_items_cold_storage.py +++ b/tests/integration/observability/test_batch_line_items_cold_storage.py @@ -2,13 +2,15 @@ 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 one NDJSON object for the batch holding one record per line. The proxy runs -on its own database so the shared rig's pollers never claim these batches first +on its own database so the shared rig's pollers never claim these batches first. Two more owned proxies list an +unsupported name next to ``s3_v2`` and a null list """ import json import threading 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 @@ -19,7 +21,7 @@ import pytest import yaml -from integration._support.client import JSON_OBJECT, Gateway, eventually, gateway_from_environment, string_value +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 @@ -44,21 +46,39 @@ @dataclass(slots=True) class ObjectSink: - """Stores every accepted object body by name and answers 503 to any upload whose body holds ``reject``.""" + """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)) @@ -103,14 +123,14 @@ class Rig: gcs: ObjectSink -def _config(directory: Path, s3_url: str) -> Path: +def _config(directory: Path, s3_url: str, line_item_callbacks: JsonValue) -> Path: shipped: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) configuration: Final = { **shipped, "litellm_settings": { **shipped["litellm_settings"], "callbacks": ["s3_v2"], - "batch_line_item_callbacks": ["s3_v2", "gcs_bucket"], + "batch_line_item_callbacks": line_item_callbacks, "s3_callback_params": { "s3_bucket_name": S3_BUCKET, "s3_region_name": "us-east-1", @@ -127,15 +147,15 @@ def _config(directory: Path, s3_url: str) -> Path: return path -@pytest.fixture(scope="module") -def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: - directory: Final = tmp_path_factory.mktemp("batch-line-items") +@contextmanager +def _owned_rig(tmp_path_factory: pytest.TempPathFactory, name: str, line_item_callbacks: JsonValue) -> 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-") + s3: Final = ObjectSink(reject=b"sink-down-", hang=b"sink-hang-") gcs: Final = ObjectSink() with ( gateway_from_environment() as environment, - scratch_database("batch_line_items") as database_url, + 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, @@ -157,13 +177,31 @@ def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: "NO_PROXY": "127.0.0.1,localhost", "no_proxy": "127.0.0.1,localhost", }, - config=_config(directory, s3_wire.url), + config=_config(directory, s3_wire.url, line_item_callbacks), remove_environment=("DATABASE_URL_READ_REPLICA",), workers=2, ) as candidate: yield Rig(candidate, database_url, s3, gcs) +@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 + + def _input_line(custom_id: str, text: str) -> dict[str, JsonValue]: return { "custom_id": custom_id, @@ -547,3 +585,84 @@ def test_a_failing_bucket_still_bills_the_batch_and_the_other_sink_keeps_its_lin 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/proxy/common_utils/test_check_batch_cost.py b/tests/unit/proxy/common_utils/test_check_batch_cost.py index a862a26c5fd9..847f9692441f 100644 --- a/tests/unit/proxy/common_utils/test_check_batch_cost.py +++ b/tests/unit/proxy/common_utils/test_check_batch_cost.py @@ -8,6 +8,7 @@ import asyncio import json +import logging from contextlib import contextmanager from typing import TYPE_CHECKING from unittest.mock import AsyncMock, MagicMock, patch @@ -3225,7 +3226,8 @@ async def test_schema_without_batch_processed_still_bills(self): 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, so a slow or failing sink never unbills or re-polls the batch.""" + 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 = ( @@ -3263,6 +3265,11 @@ def _setup(journal: list, *, status="completed", output_file_id=_CLAIM_OUTPUT_FI 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): @@ -3314,6 +3321,7 @@ async def test_lines_reach_the_callbacks_after_the_claimed_batch_is_billed_and_r with self._provider(journal): await instance.check_batch_cost() + await self._stored(instance) assert journal == [f"read {_CLAIM_OUTPUT_FILE_ID}", "claim", "bill", "finalize", "lines"] (payload,) = logger.payloads @@ -3336,6 +3344,7 @@ async def test_a_creator_that_cannot_be_looked_up_gets_lines_without_bodies(self 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)] @@ -3354,6 +3363,7 @@ async def test_a_creator_with_logging_on_gets_the_request_and_response(self, mon with self._provider(journal): await instance.check_batch_cost() + await self._stored(instance) assert journal[-3:] == ["finalize", "read file-input-1", "lines"] ((line,),) = [payload["lines"] for payload in logger.payloads] @@ -3372,6 +3382,7 @@ async def test_a_failing_sink_leaves_the_batch_billed_and_retired(self, monkeypa 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 @@ -3387,11 +3398,12 @@ async def test_without_callbacks_only_the_output_file_is_read(self, monkeypatch) 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_before_it_is_retired(self, monkeypatch): + async def test_a_batch_where_every_line_failed_sends_its_error_lines_after_it_is_retired(self, monkeypatch): import litellm journal: list = [] @@ -3401,8 +3413,9 @@ async def test_a_batch_where_every_line_failed_sends_its_error_lines_before_it_i with self._provider(journal): await instance.check_batch_cost() + await self._stored(instance) - assert journal == ["read file-err-1", "lines", "finalize"] + assert journal == ["finalize", "read file-err-1", "lines"] (payload,) = logger.payloads assert [(line["custom_id"], line["status"], line["response_cost"]) for line in payload["lines"]] == [ ("ok-1", "failure", 0.0) @@ -3410,7 +3423,7 @@ async def test_a_batch_where_every_line_failed_sends_its_error_lines_before_it_i assert row.batch_processed is True @pytest.mark.asyncio - async def test_a_batch_whose_output_file_is_gone_still_sends_its_error_lines_before_it_is_retired( + async def test_a_batch_whose_output_file_is_gone_still_sends_its_error_lines_after_it_is_retired( self, monkeypatch ): import litellm @@ -3422,8 +3435,9 @@ async def test_a_batch_whose_output_file_is_gone_still_sends_its_error_lines_bef with self._provider(journal): await instance.check_batch_cost() + await self._stored(instance) - assert journal == ["read file-out-deleted", "read file-err-1", "lines", "finalize"] + assert journal == ["read file-out-deleted", "finalize", "read file-err-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 @@ -3439,10 +3453,91 @@ async def test_a_terminal_batch_with_no_lines_sends_nothing(self, monkeypatch): 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(),)) + + with self._provider(journal): + await asyncio.wait_for(instance.check_batch_cost(), timeout=5) + while "lines started" not in journal: + await asyncio.sleep(0) + 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", "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, output_content): + 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, b"") + await self._stored(instance) + + assert journal == [f"start batch-{name} with 1 running" for name in "abc"] + class TestCreatorTurnedOffMessageLogging: """Batch line items drop the request and response when the creating key or team turned off message logging, From 15f3d737ebddcfb6b4286b1ef704db7ed2ad61d8 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Tue, 13 Oct 2026 04:35:10 -0700 Subject: [PATCH 79/89] fix(batches): store an unbilled batch's line items only once it is retired When retiring an expired, cancelled or all-failed batch fails, the batch stays on the poll page, so its lines are no longer read and stored again on every poll cycle until it is retired. The poller also releases a billed batch's output bytes before reading the next batch's output, as it did before line items existed. --- .../proxy/common_utils/check_batch_cost.py | 18 +++++++++++------- .../common_utils/test_check_batch_cost.py | 18 ++++++++++++++++++ 2 files changed, 29 insertions(+), 7 deletions(-) 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 fa1338c25829..9cf5bc36c896 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -472,9 +472,9 @@ def _is_output_file_gone_at_provider(error: Exception, output_file_id: Optional[ 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, @@ -502,10 +502,12 @@ async def _finalize_unbilled_terminal_job( 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( @@ -1065,7 +1067,7 @@ async def _store_batch_line_items( payload: Final = await build_batch_line_items(source, input_content, (output_content, error_content)) await dispatch_batch_line_items(payload, redact_bodies=redact_bodies) verbose_proxy_logger.info( - f"CheckBatchCost: sent the {len(payload['lines'])} line items of batch {batch_id} " + f"CheckBatchCost: handed the {len(payload['lines'])} 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 @@ -1241,8 +1243,8 @@ async def check_batch_cost(self): 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) - self._store_batch_line_items_in_background(job, response, model_id, batch_id, b"") + if await self._finalize_unbilled_terminal_job(job, response): + self._store_batch_line_items_in_background(job, response, model_id, batch_id, b"") continue verbose_proxy_logger.error( f"CheckBatchCost: failed to track cost for batch {batch_id} " @@ -1274,6 +1276,8 @@ async def check_batch_cost(self): f"CheckBatchCost: failed to mark job {job.id} complete in DB: {db_err}" ) self._store_batch_line_items_in_background(job, response, model_id, batch_id, output_content) + # release this batch's output before the next batch's output is read + del tracked, output_content elif response.status in PROVIDER_TERMINAL_BATCH_STATUSES: from litellm.proxy.openai_files_endpoints.common_utils import ( @@ -1286,8 +1290,8 @@ async def check_batch_cost(self): f"has not appeared yet; leaving job {job.id} for the next poll cycle" ) continue - await self._finalize_unbilled_terminal_job(job, response) - self._store_batch_line_items_in_background(job, response, model_id, batch_id, b"") + if await self._finalize_unbilled_terminal_job(job, response): + self._store_batch_line_items_in_background(job, response, model_id, batch_id, b"") # Record polling run metrics (always, even if nothing was processed) if prom_logger: 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 847f9692441f..28914d683582 100644 --- a/tests/unit/proxy/common_utils/test_check_batch_cost.py +++ b/tests/unit/proxy/common_utils/test_check_batch_cost.py @@ -3422,6 +3422,24 @@ async def test_a_batch_where_every_line_failed_sends_its_error_lines_after_it_is ] assert row.batch_processed is True + @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 async def test_a_batch_whose_output_file_is_gone_still_sends_its_error_lines_after_it_is_retired( self, monkeypatch From 21a71098c6916f5312ec17c4061322d8b0202478 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Tue, 13 Oct 2026 05:55:24 -0700 Subject: [PATCH 80/89] fix(batches): price and write batch line items off the event loop Pricing every line of a large batch and serializing its object ran on the proxy's event loop, so a 50k-line batch stalled the worker for seconds. Both now run in a worker thread. An object the bucket refuses for good is no longer kept in s3_v2's drop marks, which only a log-queue flush clears. Vertex AI input rows carry the custom_id in their key or request labels, so their lines now get the request too. --- .../proxy/common_utils/check_batch_cost.py | 5 +- litellm/batches/batch_line_items.py | 14 +++-- litellm/batches/batch_utils.py | 10 ++++ litellm/integrations/gcs_bucket/gcs_bucket.py | 2 +- litellm/integrations/s3_v2.py | 17 +++--- .../llms/vertex_ai/files/transformation.py | 9 +++ tests/unit/batches/test_batch_line_items.py | 58 +++++++++++++++---- .../gcs_bucket/test_gcs_bucket.py | 26 +++++++++ tests/unit/integrations/test_s3_v2.py | 40 ++++++++++++- .../common_utils/test_check_batch_cost.py | 26 +++++++++ 10 files changed, 178 insertions(+), 29 deletions(-) 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 9cf5bc36c896..9bc5b0300a2b 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -1064,7 +1064,10 @@ async def _store_batch_line_items( if redact_bodies else await self._batch_file_content(response.input_file_id, cost_provider, file_access) ) - payload: Final = await build_batch_line_items(source, input_content, (output_content, error_content)) + # pricing every line of a large batch is CPU-bound, so it runs off the event loop + payload: Final = await asyncio.to_thread( + build_batch_line_items, source, input_content, (output_content, error_content) + ) await dispatch_batch_line_items(payload, redact_bodies=redact_bodies) verbose_proxy_logger.info( f"CheckBatchCost: handed the {len(payload['lines'])} line items of batch {batch_id} " diff --git a/litellm/batches/batch_line_items.py b/litellm/batches/batch_line_items.py index 93ec1daca303..0dee5a1ab5f4 100644 --- a/litellm/batches/batch_line_items.py +++ b/litellm/batches/batch_line_items.py @@ -8,9 +8,10 @@ import litellm from litellm._logging import verbose_logger -from litellm.batches.batch_utils import calculate_batch_cost_and_usage, get_file_content_as_dictionary +from litellm.batches.batch_utils import batch_cost_and_usage, get_file_content_as_dictionary from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import init_custom_logger_compatible_class +from litellm.llms.vertex_ai.files.transformation import vertex_batch_row_custom_id from litellm.types.utils import ( ModelInfo, StandardBatchLineItem, @@ -50,11 +51,12 @@ def _text(attribution: Mapping[str, object], name: str) -> str | None: 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 None + return key if isinstance(key, str) else vertex_batch_row_custom_id(entry) -async def build_batch_line_items( +def build_batch_line_items( source: BatchLineItemSource, input_content: bytes, result_contents: Iterable[bytes] ) -> StandardBatchLineItemsPayload: """Pair every output and error line with the input line that shares its custom_id, priced one line at a time.""" @@ -63,16 +65,16 @@ async def build_batch_line_items( return StandardBatchLineItemsPayload( batch_id=source.batch_id, finished_at=source.finished_at, - lines=tuple([await _line_item(source, requests, entry) for entry in results]), + lines=tuple([_line_item(source, requests, entry) for entry in results]), ) -async def _line_item( +def _line_item( source: BatchLineItemSource, requests: Mapping[str | None, Mapping[str, object]], entry: dict[str, object], ) -> StandardBatchLineItem: - priced: Final = await calculate_batch_cost_and_usage( + 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 diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 21479bfa419f..424112a53e8a 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -74,6 +74,16 @@ async def calculate_batch_cost_and_usage( 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, ) -> BatchCostUsageResult: """ Calculate the cost and usage of a batch. diff --git a/litellm/integrations/gcs_bucket/gcs_bucket.py b/litellm/integrations/gcs_bucket/gcs_bucket.py index 6cad51042792..c66d3fdee5b7 100644 --- a/litellm/integrations/gcs_bucket/gcs_bucket.py +++ b/litellm/integrations/gcs_bucket/gcs_bucket.py @@ -107,7 +107,7 @@ async def async_log_batch_line_items(self, payload: StandardBatchLineItemsPayloa headers=headers, bucket_name=gcs_logging_config["bucket_name"], object_name=batch_line_items_object_name(payload), - logging_payload=batch_line_items_ndjson(payload), + logging_payload=await asyncio.to_thread(batch_line_items_ndjson, payload), ) def _drain_queue_batch(self) -> list[GCSLogQueueItem]: diff --git a/litellm/integrations/s3_v2.py b/litellm/integrations/s3_v2.py index ec5c11b1826b..f8b44c9981e7 100644 --- a/litellm/integrations/s3_v2.py +++ b/litellm/integrations/s3_v2.py @@ -517,16 +517,17 @@ async def async_log_batch_line_items(self, payload: StandardBatchLineItemsPayloa 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 - delivered: Final = await self.async_upload_data_to_s3( - s3BatchLoggingElement( - payload={}, - body=batch_line_items_ndjson(payload), - content_type="application/x-ndjson", - s3_object_key=s3_object_key, - s3_object_download_filename=object_name.rsplit("/", 1)[-1], - ) + 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=object_name.rsplit("/", 1)[-1], ) + delivered: Final = await self.async_upload_data_to_s3(element) 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 batch %s were not stored at %s", len(payload["lines"]), diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index 09ec2ba6ee42..f84a66d12772 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/tests/unit/batches/test_batch_line_items.py b/tests/unit/batches/test_batch_line_items.py index 02ad67195b72..17b8d6a6c9f4 100644 --- a/tests/unit/batches/test_batch_line_items.py +++ b/tests/unit/batches/test_batch_line_items.py @@ -14,6 +14,7 @@ from litellm.batches.batch_utils import calculate_batch_cost_and_usage, get_file_content_as_dictionary 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 @@ -78,8 +79,8 @@ def _source(**overrides) -> BatchLineItemSource: return BatchLineItemSource(**{**fields, **overrides}) -async def _payload() -> StandardBatchLineItemsPayload: - return await build_batch_line_items( +def _payload() -> StandardBatchLineItemsPayload: + return build_batch_line_items( _source(), _jsonl(_request("ok-1", "first question"), _request("bad-1", "second question")), (_jsonl(_success("ok-1", 100, 50)), _jsonl(_error("bad-1"))), @@ -102,7 +103,7 @@ async def async_log_batch_line_items(self, payload: StandardBatchLineItemsPayloa @pytest.mark.asyncio async def test_each_result_line_is_paired_with_the_request_that_shares_its_custom_id(): - payload = await build_batch_line_items( + payload = build_batch_line_items( _source(), _jsonl(_request("ok-1", "first"), _request("bad-1", "second"), _request("never-ran", "third")), (_jsonl(_success("ok-1", 100, 50)), _jsonl(_error("bad-1"))), @@ -142,7 +143,7 @@ async def test_each_result_line_is_paired_with_the_request_that_shares_its_custo 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 = await build_batch_line_items(_source(), b"", (outputs,)) + payload = build_batch_line_items(_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 ) @@ -164,7 +165,7 @@ async def test_anthropic_results_without_an_input_file_keep_their_status_and_usa } errored = {"custom_id": "a-2", "result": {"type": "errored", "error": {"type": "invalid_request_error"}}} - payload = await build_batch_line_items( + payload = build_batch_line_items( _source(model="claude-sonnet-5-5", custom_llm_provider="anthropic", cost_provider="anthropic"), b"", (_jsonl(succeeded, errored),), @@ -181,7 +182,7 @@ 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 = await build_batch_line_items( + payload = build_batch_line_items( _source(custom_llm_provider="bedrock", cost_provider="bedrock"), _jsonl(request), (_jsonl(result),) ) @@ -189,9 +190,44 @@ async def test_bedrock_lines_pair_by_record_id(): 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 = build_batch_line_items( + _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_input_row_without_a_custom_id_pairs_with_nothing(): + unlabeled = {"request": {"contents": [{"role": "user", "parts": [{"text": "hi"}]}]}} + + payload = build_batch_line_items( + _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.asyncio async def test_redaction_drops_only_the_request_and_response(): - payload = await _payload() + payload = _payload() redacted = without_request_and_response(payload) @@ -203,7 +239,7 @@ async def test_redaction_drops_only_the_request_and_response(): @pytest.mark.asyncio async def test_the_object_holds_one_json_line_per_batch_line_under_its_finished_day(): - payload = await _payload() + payload = _payload() ndjson = batch_line_items_ndjson(payload) @@ -218,7 +254,7 @@ async def test_the_object_holds_one_json_line_per_batch_line_under_its_finished_ @pytest.mark.asyncio async def test_a_logger_gets_the_request_and_response_unless_message_logging_is_off(monkeypatch): - payload = await _payload() + payload = _payload() open_logger, quiet_logger, legacy_logger = ( _RecordingLogger(), _RecordingLogger(turn_off_message_logging=True), @@ -235,7 +271,7 @@ async def test_a_logger_gets_the_request_and_response_unless_message_logging_is_ @pytest.mark.asyncio async def test_redacted_dispatch_never_hands_a_logger_the_bodies(monkeypatch): - payload = await _payload() + payload = _payload() logger = _RecordingLogger() monkeypatch.setattr(litellm, "batch_line_item_callbacks", (logger,)) @@ -246,7 +282,7 @@ async def test_redacted_dispatch_never_hands_a_logger_the_bodies(monkeypatch): @pytest.mark.asyncio async def test_a_failing_or_unknown_sink_is_logged_and_the_other_sinks_still_get_the_lines(monkeypatch, caplog): - payload = await _payload() + payload = _payload() logger = _RecordingLogger() monkeypatch.setattr(litellm, "batch_line_item_callbacks", (_FailingLogger(), "langfuse", logger)) diff --git a/tests/unit/integrations/gcs_bucket/test_gcs_bucket.py b/tests/unit/integrations/gcs_bucket/test_gcs_bucket.py index de2fb6c0f3fb..e7fdf5b3402a 100644 --- a/tests/unit/integrations/gcs_bucket/test_gcs_bucket.py +++ b/tests/unit/integrations/gcs_bucket/test_gcs_bucket.py @@ -368,3 +368,29 @@ async def test_a_rejected_batch_line_items_write_raises_for_the_dispatcher(monke ) 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, "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 44aa69904807..7cd8747c9888 100644 --- a/tests/unit/integrations/test_s3_v2.py +++ b/tests/unit/integrations/test_s3_v2.py @@ -5562,12 +5562,14 @@ def _batch_line_items_payload() -> dict: } -def _s3_logger_answering(status: int, s3_path: str | None) -> tuple[S3Logger, list[httpx.Request]]: +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, request=request) + return httpx.Response(status, content=content, request=request) logger = S3Logger( s3_bucket_name="test-bucket", @@ -5619,3 +5621,37 @@ async def test_a_rejected_batch_line_items_upload_is_logged_with_the_batch_id(ca "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_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 + threads = [] + + def _serialize(payload): + threads.append(threading.current_thread()) + return serialize(payload) + + monkeypatch.setattr(s3_v2_module, "batch_line_items_ndjson", _serialize) + + await logger.async_log_batch_line_items(_batch_line_items_payload()) + + assert len(requests) == 1 + (thread,) = threads + assert thread is not threading.main_thread() 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 28914d683582..9d73c1424b0a 100644 --- a/tests/unit/proxy/common_utils/test_check_batch_cost.py +++ b/tests/unit/proxy/common_utils/test_check_batch_cost.py @@ -3333,6 +3333,32 @@ async def test_lines_reach_the_callbacks_after_the_claimed_batch_is_billed_and_r 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),)) + build = batch_line_items.build_batch_line_items + threads = [] + + def _build(*args): + threads.append(threading.current_thread()) + return build(*args) + + monkeypatch.setattr(batch_line_items, "build_batch_line_items", _build) + 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 From 082a859d5105fba462c8966b691483821607acb7 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Tue, 13 Oct 2026 06:06:48 -0700 Subject: [PATCH 81/89] fix(batches): pair no request for a shared custom_id and wait for an unclaimed batch's retry A Vertex AI embeddings entry whose input is a list is split into one row per item, all with the same custom_id, so its line now gets no request instead of the last row. On a database without batch_processed nothing claims a billed batch, so when marking it complete fails it stays in the poll page and its lines now wait for the next poll instead of being written on every cycle. --- .../proxy/common_utils/check_batch_cost.py | 4 +++ litellm/batches/batch_line_items.py | 11 +++++++- tests/unit/batches/test_batch_line_items.py | 27 +++++++++++++++++++ .../common_utils/test_check_batch_cost.py | 21 +++++++++++++++ 4 files changed, 62 insertions(+), 1 deletion(-) 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 9bc5b0300a2b..8c520897d3e7 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -1278,6 +1278,10 @@ async def check_batch_cost(self): 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 + del tracked, output_content + continue self._store_batch_line_items_in_background(job, response, model_id, batch_id, output_content) # release this batch's output before the next batch's output is read del tracked, output_content diff --git a/litellm/batches/batch_line_items.py b/litellm/batches/batch_line_items.py index 0dee5a1ab5f4..20fe3f677382 100644 --- a/litellm/batches/batch_line_items.py +++ b/litellm/batches/batch_line_items.py @@ -1,6 +1,7 @@ """Per-line records of a finished batch for the loggers listed in batch_line_item_callbacks.""" import asyncio +from collections import Counter from collections.abc import Iterable, Mapping from dataclasses import dataclass from itertools import chain @@ -60,7 +61,7 @@ def build_batch_line_items( source: BatchLineItemSource, input_content: bytes, result_contents: Iterable[bytes] ) -> StandardBatchLineItemsPayload: """Pair every output and error line with the input line that shares its custom_id, priced one line at a time.""" - requests: Final = {_line_key(entry): entry for entry in get_file_content_as_dictionary(input_content)} + requests: Final = _requests_by_custom_id(input_content) results: Final = chain.from_iterable(map(get_file_content_as_dictionary, result_contents)) return StandardBatchLineItemsPayload( batch_id=source.batch_id, @@ -69,6 +70,14 @@ def build_batch_line_items( ) +def _requests_by_custom_id(input_content: bytes) -> Mapping[str | None, Mapping[str, object]]: + """Each input line 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((_line_key(entry), entry) for entry in get_file_content_as_dictionary(input_content)) + counts: Final = Counter(key for key, _ in keyed) + return {key: entry for key, entry in keyed if counts[key] == 1} + + def _line_item( source: BatchLineItemSource, requests: Mapping[str | None, Mapping[str, object]], diff --git a/tests/unit/batches/test_batch_line_items.py b/tests/unit/batches/test_batch_line_items.py index 17b8d6a6c9f4..6dc1c9266526 100644 --- a/tests/unit/batches/test_batch_line_items.py +++ b/tests/unit/batches/test_batch_line_items.py @@ -212,6 +212,33 @@ def test_vertex_lines_pair_with_the_input_row_that_carries_their_custom_id(): ] +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 = build_batch_line_items( + _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"}]}]}} 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 9d73c1424b0a..46f8484e6877 100644 --- a/tests/unit/proxy/common_utils/test_check_batch_cost.py +++ b/tests/unit/proxy/common_utils/test_check_batch_cost.py @@ -3466,6 +3466,27 @@ async def test_a_batch_that_could_not_be_retired_sends_no_lines_until_it_is(self 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 From 459234f5d35683fdece666576afccff55167b19e Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Tue, 13 Oct 2026 06:31:54 -0700 Subject: [PATCH 82/89] fix(batches): store batch line items only once the poll has billed every batch Pricing a large batch's lines on a worker thread still competes with the poll for the GIL, so the next batches in the same poll were billed and retired seconds later. The background stores now wait until the poll cycle that queued them has finished. --- .../proxy/common_utils/check_batch_cost.py | 19 +++++++++--- .../common_utils/test_check_batch_cost.py | 30 +++++++++++++++++++ 2 files changed, 45 insertions(+), 4 deletions(-) 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 8c520897d3e7..e6e3578f4fdf 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -133,6 +133,8 @@ def __init__( 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: @@ -977,9 +979,9 @@ def _store_batch_line_items_in_background( batch_id: str, output_content: bytes, ) -> None: - """Store a retired batch's lines off the poll loop, one batch at a time, so a slow or hung bucket never 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.""" + """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: @@ -1003,6 +1005,7 @@ async def _store_batch_line_items_in_turn( output_content: bytes, ) -> 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, output_content) async def _store_batch_line_items( @@ -1132,7 +1135,15 @@ async def _creator_turned_off_message_logging(self, job: "_ManagedObjectRow") -> 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): + 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 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 46f8484e6877..5f9630efc712 100644 --- a/tests/unit/proxy/common_utils/test_check_batch_cost.py +++ b/tests/unit/proxy/common_utils/test_check_batch_cost.py @@ -3603,6 +3603,36 @@ async def _store(job, response, model_id, batch_id, output_content): 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", b"") + await asyncio.sleep(0.01) + + async def _store(job, response, model_id, batch_id, output_content): + 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", + ] + class TestCreatorTurnedOffMessageLogging: """Batch line items drop the request and response when the creating key or team turned off message logging, From 400a6f9bafde59531aafbe441b6dccee98897a82 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Tue, 13 Oct 2026 07:47:41 -0700 Subject: [PATCH 83/89] fix(batches): hash line-item objects off the event loop and keep lines with a lone surrogate A batch's whole line-items object was hashed for its S3 PUT on the event loop, and a line holding a lone surrogate from a provider's JSON escapes made the object impossible to encode as UTF-8, so the batch's object was never written. The object is now prepared in a worker thread, and only the line that cannot be encoded is written with escaped characters. --- litellm/integrations/batch_line_item_objects.py | 13 ++++++++++++- litellm/integrations/s3_v2.py | 7 +++++-- tests/unit/batches/test_batch_line_items.py | 15 +++++++++++++++ tests/unit/integrations/test_s3_v2.py | 16 +++++++++++----- 4 files changed, 43 insertions(+), 8 deletions(-) diff --git a/litellm/integrations/batch_line_item_objects.py b/litellm/integrations/batch_line_item_objects.py index 23dd2e3200e4..96f7c6d91a6a 100644 --- a/litellm/integrations/batch_line_item_objects.py +++ b/litellm/integrations/batch_line_item_objects.py @@ -8,7 +8,18 @@ def batch_line_items_ndjson(payload: StandardBatchLineItemsPayload) -> str: - return "".join(f"{json.dumps(line, default=str, ensure_ascii=False)}\n" for line in payload["lines"]) + 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: diff --git a/litellm/integrations/s3_v2.py b/litellm/integrations/s3_v2.py index f8b44c9981e7..f36923ba91c3 100644 --- a/litellm/integrations/s3_v2.py +++ b/litellm/integrations/s3_v2.py @@ -524,7 +524,9 @@ async def async_log_batch_line_items(self, payload: StandardBatchLineItemsPayloa s3_object_key=s3_object_key, s3_object_download_filename=object_name.rsplit("/", 1)[-1], ) - delivered: Final = await self.async_upload_data_to_s3(element) + # hashing a whole batch's 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) @@ -579,6 +581,7 @@ def _upload_semaphore(self, value: asyncio.Semaphore | AdaptiveConcurrencyLimite 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 @@ -611,7 +614,7 @@ async def signed_put(prepared: _PreparedPut) -> httpx.Response: return error.response max_retries: Final = 3 - 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) for attempt in range(max_retries): response = await self._recorded_put(partial(signed_put, prepared)) if ( diff --git a/tests/unit/batches/test_batch_line_items.py b/tests/unit/batches/test_batch_line_items.py index 6dc1c9266526..361c9d804a9a 100644 --- a/tests/unit/batches/test_batch_line_items.py +++ b/tests/unit/batches/test_batch_line_items.py @@ -279,6 +279,21 @@ async def test_the_object_holds_one_json_line_per_batch_line_under_its_finished_ ) +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() diff --git a/tests/unit/integrations/test_s3_v2.py b/tests/unit/integrations/test_s3_v2.py index 7cd8747c9888..278480ba605b 100644 --- a/tests/unit/integrations/test_s3_v2.py +++ b/tests/unit/integrations/test_s3_v2.py @@ -5635,23 +5635,29 @@ async def test_a_batch_line_items_object_the_bucket_refuses_for_good_is_not_kept @pytest.mark.asyncio -async def test_batch_line_items_are_serialized_off_the_event_loop(monkeypatch): +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 - threads = [] + prepare = logger._prepare_put + threads = {} def _serialize(payload): - threads.append(threading.current_thread()) + 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 - (thread,) = threads - assert thread is not threading.main_thread() + assert set(threads) == {"serialize", "hash"} + assert threading.main_thread() not in threads.values() From ee61a607c8d2917e035135e88cf0e486853d5705 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Tue, 13 Oct 2026 08:48:51 -0700 Subject: [PATCH 84/89] fix(batches): keep request and response for batches created without a virtual key A batch created with a JWT has no API key on its row, and its line items were stored without request and response as if its creator had been deleted. Such a batch now follows its team's and the proxy's message logging settings, the same as a master key batch. A deleted key or team or a failed lookup still drops the bodies. --- .../proxy/common_utils/check_batch_cost.py | 5 +- .../test_batch_line_items_cold_storage.py | 71 +++++++++++++++++-- .../common_utils/test_check_batch_cost.py | 35 +++++++-- 3 files changed, 97 insertions(+), 14 deletions(-) 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 e6e3578f4fdf..6473981a08bb 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -1100,7 +1100,8 @@ async def _batch_file_content( 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 key or team that can no longer be read counts as turned off.""" + proxy-wide setting. A batch created without a virtual key, such as under JWT 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.proxy._types import UserAPIKeyAuth from litellm.proxy.litellm_pre_call_utils import get_dynamic_logging_metadata from litellm.proxy.proxy_server import proxy_config @@ -1125,7 +1126,7 @@ async def _creator_turned_off_message_logging(self, job: "_ManagedObjectRow") -> f"so its line items are stored without request and response: {error}" ) return True - if api_key is None or (has_key_row and key_row is None) or (team_id and team_row is None): + 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), diff --git a/tests/integration/observability/test_batch_line_items_cold_storage.py b/tests/integration/observability/test_batch_line_items_cold_storage.py index 0fff743572b9..2e864e9fa408 100644 --- a/tests/integration/observability/test_batch_line_items_cold_storage.py +++ b/tests/integration/observability/test_batch_line_items_cold_storage.py @@ -2,12 +2,13 @@ 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 one NDJSON object for the batch holding one record per line. The proxy runs -on its own database so the shared rig's pollers never claim these batches first. Two more owned proxies list an -unsupported name next to ``s3_v2`` and a null list +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. Two more owned proxies list an unsupported name next to ``s3_v2`` and a null list """ import json import threading +import time import uuid from collections.abc import Callable, Iterator, Mapping from contextlib import contextmanager @@ -19,8 +20,10 @@ 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 @@ -30,6 +33,7 @@ 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) @@ -42,6 +46,7 @@ PROMPT_TOKENS: Final = 100 COMPLETION_TOKENS: Final = 50 POLLER_SECONDS: Final = 150 +JWT_KEY_ID: Final = "batch-line-items-signing-key" @dataclass(slots=True) @@ -121,12 +126,38 @@ class Rig: 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 _config(directory: Path, s3_url: str, line_item_callbacks: JsonValue) -> Path: shipped: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) configuration: Final = { **shipped, + "general_settings": { + **shipped["general_settings"], + "enable_jwt_auth": True, + "litellm_jwtauth": {"user_id_jwt_field": "sub", "user_id_upsert": True}, + }, "litellm_settings": { **shipped["litellm_settings"], "callbacks": ["s3_v2"], @@ -153,7 +184,9 @@ def _owned_rig(tmp_path_factory: pytest.TempPathFactory, name: str, line_item_ca 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, gateway_from_environment() as environment, scratch_database(name.replace("-", "_")) as database_url, wire_server(s3_peer(s3)) as s3_wire, @@ -171,6 +204,7 @@ def _owned_rig(tmp_path_factory: pytest.TempPathFactory, name: str, line_item_ca "GCS_BUCKET_NAME": GCS_BUCKET, "GCS_PATH_SERVICE_ACCOUNT": str(service_account), "GCS_FLUSH_INTERVAL": "1", + "JWT_PUBLIC_KEY_URL": jwks_wire.url, "HTTPS_PROXY": tunnel.url, "https_proxy": tunnel.url, "SSL_CERT_FILE": str(certificate), @@ -181,7 +215,7 @@ def _owned_rig(tmp_path_factory: pytest.TempPathFactory, name: str, line_item_ca remove_environment=("DATABASE_URL_READ_REPLICA",), workers=2, ) as candidate: - yield Rig(candidate, database_url, s3, gcs) + yield Rig(candidate, database_url, s3, gcs, signing_key) @pytest.fixture(scope="module") @@ -333,7 +367,9 @@ class SubmittedBatch: handle: ScenarioHandle -def _submit(rig: Rig, lines: Lines, key_fields: Mapping[str, JsonValue], *, finished: bool = True) -> SubmittedBatch: +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)) model: Final = string_value( rig.gateway.post( @@ -350,7 +386,11 @@ def _submit(rig: Rig, lines: Lines, key_fields: Mapping[str, JsonValue], *, fini }, )["model_name"] ) - key: Final = string_value(rig.gateway.post("/key/generate", {"models": [model], **key_fields})["key"]) + 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}, @@ -535,6 +575,27 @@ def test_a_batch_whose_creating_key_was_deleted_gets_its_lines_without_request_o 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_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")), 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 5f9630efc712..96fce13b8913 100644 --- a/tests/unit/proxy/common_utils/test_check_batch_cost.py +++ b/tests/unit/proxy/common_utils/test_check_batch_cost.py @@ -3323,7 +3323,7 @@ async def test_lines_reach_the_callbacks_after_the_claimed_batch_is_billed_and_r await instance.check_batch_cost() await self._stored(instance) - assert journal == [f"read {_CLAIM_OUTPUT_FILE_ID}", "claim", "bill", "finalize", "lines"] + assert journal == [f"read {_CLAIM_OUTPUT_FILE_ID}", "claim", "bill", "finalize", "read file-input-1", "lines"] (payload,) = logger.payloads assert payload["batch_id"] == _CLAIM_UNIFIED_BATCH_ID assert payload["finished_at"] == 2 @@ -3364,7 +3364,8 @@ async def test_a_creator_that_cannot_be_looked_up_gets_lines_without_bodies(self import litellm journal: list = [] - _, instance = self._setup(journal) + row, instance = self._setup(journal) + row.api_key = "hash-of-a-deleted-key" logger = self._recording_logger(journal) monkeypatch.setattr(litellm, "batch_line_item_callbacks", (logger,)) @@ -3441,7 +3442,7 @@ async def test_a_batch_where_every_line_failed_sends_its_error_lines_after_it_is await instance.check_batch_cost() await self._stored(instance) - assert journal == ["finalize", "read file-err-1", "lines"] + 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) @@ -3502,7 +3503,7 @@ async def test_a_batch_whose_output_file_is_gone_still_sends_its_error_lines_aft await instance.check_batch_cost() await self._stored(instance) - assert journal == ["read file-out-deleted", "finalize", "read file-err-1", "lines"] + 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 @@ -3551,7 +3552,7 @@ async def async_log_batch_line_items(self, payload): assert billed_while_hung == ( True, "complete", - [f"read {_CLAIM_OUTPUT_FILE_ID}", "claim", "bill", "finalize", "lines started"], + [f"read {_CLAIM_OUTPUT_FILE_ID}", "claim", "bill", "finalize", "read file-input-1", "lines started"], ) assert journal[-1] == "lines" @@ -3700,11 +3701,11 @@ async def test_a_team_that_turned_off_message_logging_covers_its_keys(self, monk @pytest.mark.parametrize( ("job_fields", "key_row", "team_row"), [ - ({"api_key": None}, None, None), ({}, None, None), ({"team_id": "team-gone"}, "present", None), + ({"api_key": None, "team_id": "team-gone"}, None, None), ], - ids=["legacy-row-without-key", "deleted-key", "deleted-team"], + 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 @@ -3716,6 +3717,26 @@ async def test_a_creator_that_is_gone_counts_as_turned_off(self, monkeypatch, jo assert await instance._creator_turned_off_message_logging(self._job(**job_fields)) is True + @pytest.mark.asyncio + @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_key_follows_its_team_and_the_proxy_setting( + self, monkeypatch, 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=None, team_id="team-alpha" if team_row is not None else None, created_by="jwt-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 From 85169fc62fd683f1db27c73e1d5dc19267c2d5a8 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Tue, 13 Oct 2026 12:02:53 -0700 Subject: [PATCH 85/89] fix(batches): pair line items with the right request on every deployment and skip empty objects A batch retired with nothing billable now reads its input file by the id the provider reported, so a batch that ran on the second deployment of a model group keeps each line's request. A Vertex AI line read as custom_id "unknown" pairs with no input row, and a batch whose files hold no readable line writes no object --- .../proxy/common_utils/check_batch_cost.py | 18 ++++-- litellm/batches/batch_line_items.py | 13 +++- .../test_batch_line_items_cold_storage.py | 59 ++++++++++++++++++- tests/unit/batches/test_batch_line_items.py | 31 ++++++++++ .../common_utils/test_check_batch_cost.py | 45 ++++++++++++++ 5 files changed, 159 insertions(+), 7 deletions(-) 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 6473981a08bb..0ac362f6e4c8 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -971,6 +971,16 @@ async def _track_completed_batch_cost( return model_name, str(llm_provider) if llm_provider else None, content_bytes + 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, b"") + def _store_batch_line_items_in_background( self, job: "_ManagedObjectRow", @@ -1071,6 +1081,8 @@ async def _store_batch_line_items( payload: Final = await asyncio.to_thread( build_batch_line_items, source, input_content, (output_content, error_content) ) + if not payload["lines"]: + return await dispatch_batch_line_items(payload, redact_bodies=redact_bodies) verbose_proxy_logger.info( f"CheckBatchCost: handed the {len(payload['lines'])} line items of batch {batch_id} " @@ -1258,8 +1270,7 @@ async def _poll_batches(self): f"CheckBatchCost: output file {response.output_file_id} of batch {batch_id} " f"does not exist at the provider; retiring job {job.id} unbilled" ) - if await self._finalize_unbilled_terminal_job(job, response): - self._store_batch_line_items_in_background(job, response, model_id, batch_id, b"") + 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} " @@ -1309,8 +1320,7 @@ async def _poll_batches(self): f"has not appeared yet; leaving job {job.id} for the next poll cycle" ) continue - if await self._finalize_unbilled_terminal_job(job, response): - self._store_batch_line_items_in_background(job, response, model_id, batch_id, b"") + 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/batches/batch_line_items.py b/litellm/batches/batch_line_items.py index 20fe3f677382..427560dff37e 100644 --- a/litellm/batches/batch_line_items.py +++ b/litellm/batches/batch_line_items.py @@ -23,6 +23,8 @@ 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 + @dataclass(frozen=True, slots=True) class BatchLineItemSource: @@ -92,7 +94,7 @@ def _line_item( batch_id=source.batch_id, custom_id=custom_id, status="success" if succeeded else "failure", - request=requests.get(custom_id) if custom_id is not None else None, + request=_request_for(source, requests, custom_id), response=entry, usage=StandardBatchLineItemUsage( prompt_tokens=priced.usage.prompt_tokens, @@ -108,6 +110,15 @@ def _line_item( ) +def _request_for( + source: BatchLineItemSource, requests: Mapping[str | None, Mapping[str, object]], custom_id: str | None +) -> Mapping[str, object] | None: + """The input line that shares a result line's custom_id. 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 requests.get(custom_id) + + def without_request_and_response(payload: StandardBatchLineItemsPayload) -> StandardBatchLineItemsPayload: return StandardBatchLineItemsPayload( batch_id=payload["batch_id"], diff --git a/tests/integration/observability/test_batch_line_items_cold_storage.py b/tests/integration/observability/test_batch_line_items_cold_storage.py index 2e864e9fa408..7a0bdde8f226 100644 --- a/tests/integration/observability/test_batch_line_items_cold_storage.py +++ b/tests/integration/observability/test_batch_line_items_cold_storage.py @@ -29,7 +29,7 @@ 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 ScenarioHandle, delete_scenario, register_scenario +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 @@ -47,6 +47,9 @@ 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) @@ -153,6 +156,23 @@ def _config(directory: Path, s3_url: str, line_item_callbacks: JsonValue) -> Pat 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, @@ -371,7 +391,11 @@ 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)) - model: Final = string_value( + 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", { @@ -386,6 +410,11 @@ def _submit( }, )["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 @@ -615,6 +644,32 @@ def test_a_batch_where_every_line_failed_still_writes_its_error_lines(rig: Rig) 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( diff --git a/tests/unit/batches/test_batch_line_items.py b/tests/unit/batches/test_batch_line_items.py index 361c9d804a9a..eae1a2abda94 100644 --- a/tests/unit/batches/test_batch_line_items.py +++ b/tests/unit/batches/test_batch_line_items.py @@ -252,6 +252,37 @@ def test_a_vertex_input_row_without_a_custom_id_pairs_with_nothing(): 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 = build_batch_line_items( + _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() 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 96fce13b8913..77e61d2bcb4e 100644 --- a/tests/unit/proxy/common_utils/test_check_batch_cost.py +++ b/tests/unit/proxy/common_utils/test_check_batch_cost.py @@ -3449,6 +3449,51 @@ async def test_a_batch_where_every_line_failed_sends_its_error_lines_after_it_is ] 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 From 2875f03d5e442d0a1726c7835c2604104635498b Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Tue, 13 Oct 2026 15:19:18 -0700 Subject: [PATCH 86/89] fix(batches): let OAuth2 batches follow team and proxy message logging settings The line-item redaction check looked up every saved creator credential other than the master key or a CLI session as a virtual key. An OAuth2 batch saves the raw token, so it read like a deleted key and lost its request and response It now looks up only a saved value shaped like a virtual key's hash, reusing is_valid_sha256_hash, so OAuth2 batches follow their team's logging settings and the proxy-wide setting while deleted keys still fail closed --- .../proxy/common_utils/check_batch_cost.py | 12 ++- .../test_batch_line_items_cold_storage.py | 82 ++++++++++++++++++- .../common_utils/test_check_batch_cost.py | 12 +-- 3 files changed, 90 insertions(+), 16 deletions(-) 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 0ac362f6e4c8..6a342ec20004 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -14,7 +14,6 @@ from litellm._uuid import uuid from litellm.constants import ( CLI_SESSION_KEY_PREFIX, - LITELLM_PROXY_MASTER_KEY_ALIAS, MANAGED_OBJECT_STALENESS_CUTOFF_DAYS, MAX_OBJECTS_PER_POLL_CYCLE, ) @@ -1112,8 +1111,10 @@ async def _batch_file_content( 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 created without a virtual key, such as under JWT 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.""" + 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 @@ -1121,10 +1122,7 @@ async def _creator_turned_off_message_logging(self, job: "_ManagedObjectRow") -> 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 api_key not in ( - LITELLM_PROXY_MASTER_KEY_ALIAS, - f"{CLI_SESSION_KEY_PREFIX}-{job.created_by}", - ) + 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 diff --git a/tests/integration/observability/test_batch_line_items_cold_storage.py b/tests/integration/observability/test_batch_line_items_cold_storage.py index 7a0bdde8f226..49da6a0619e2 100644 --- a/tests/integration/observability/test_batch_line_items_cold_storage.py +++ b/tests/integration/observability/test_batch_line_items_cold_storage.py @@ -3,7 +3,8 @@ 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 one NDJSON object for the batch holding one record per line. 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. Two more owned proxies list an unsupported name next to ``s3_v2`` and a null list +next to virtual keys. Three more owned proxies list an unsupported name next to ``s3_v2``, list a null list, and +check OAuth2 tokens against a local token info endpoint """ import json @@ -152,7 +153,24 @@ def _signed_token(signing_key: rsa.RSAPrivateKey, subject: str) -> str: ) -def _config(directory: Path, s3_url: str, line_item_callbacks: JsonValue) -> Path: +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, @@ -177,6 +195,7 @@ def _config(directory: Path, s3_url: str, line_item_callbacks: JsonValue) -> Pat **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"], @@ -199,7 +218,9 @@ def _config(directory: Path, s3_url: str, line_item_callbacks: JsonValue) -> Pat @contextmanager -def _owned_rig(tmp_path_factory: pytest.TempPathFactory, name: str, line_item_callbacks: JsonValue) -> Iterator[Rig]: +def _owned_rig( + tmp_path_factory: pytest.TempPathFactory, name: str, line_item_callbacks: JsonValue, *, oauth2: bool = False +) -> 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-") @@ -207,6 +228,7 @@ def _owned_rig(tmp_path_factory: pytest.TempPathFactory, name: str, line_item_ca 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, @@ -225,13 +247,14 @@ def _owned_rig(tmp_path_factory: pytest.TempPathFactory, name: str, line_item_ca "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", }, - config=_config(directory, s3_wire.url, line_item_callbacks), + config=_config(directory, s3_wire.url, line_item_callbacks, oauth2=oauth2), remove_environment=("DATABASE_URL_READ_REPLICA",), workers=2, ) as candidate: @@ -256,6 +279,12 @@ def rig_with_a_null_list(tmp_path_factory: pytest.TempPathFactory) -> Iterator[R 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 + + def _input_line(custom_id: str, text: str) -> dict[str, JsonValue]: return { "custom_id": custom_id, @@ -625,6 +654,51 @@ def test_a_batch_created_with_a_jwt_and_no_key_gets_its_lines_with_request_and_r ] +def test_a_batch_created_with_an_oauth2_token_follows_its_teams_message_logging_setting(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] + + 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")), 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 77e61d2bcb4e..d4faba6e6268 100644 --- a/tests/unit/proxy/common_utils/test_check_batch_cost.py +++ b/tests/unit/proxy/common_utils/test_check_batch_cost.py @@ -10,6 +10,7 @@ import json import logging from contextlib import contextmanager +from hashlib import sha256 from typing import TYPE_CHECKING from unittest.mock import AsyncMock, MagicMock, patch @@ -3365,7 +3366,7 @@ async def test_a_creator_that_cannot_be_looked_up_gets_lines_without_bodies(self journal: list = [] row, instance = self._setup(journal) - row.api_key = "hash-of-a-deleted-key" + row.api_key = sha256(b"sk-deleted-key").hexdigest() logger = self._recording_logger(journal) monkeypatch.setattr(litellm, "batch_line_item_callbacks", (logger,)) @@ -3694,7 +3695,7 @@ def _instance(key_row=None, team_row=None, error: Exception | None = None): return CheckBatchCost(proxy_logging_obj=MagicMock(), prisma_client=prisma, llm_router=MagicMock()) @staticmethod - def _job(api_key="hash-alice", team_id=None, created_by="alice"): + 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) @@ -3763,12 +3764,13 @@ async def test_a_creator_that_is_gone_counts_as_turned_off(self, monkeypatch, jo 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_key_follows_its_team_and_the_proxy_setting( - self, monkeypatch, team_value, proxy_wide, expected + 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 @@ -3777,7 +3779,7 @@ async def test_a_batch_created_without_a_key_follows_its_team_and_the_proxy_sett 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=None, team_id="team-alpha" if team_row is not None else None, created_by="jwt-alice") + 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() From 0ee34c94be0c4bc029fb2f3e0bcc7b63755400b9 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Wed, 14 Oct 2026 14:25:52 -0700 Subject: [PATCH 87/89] fix(batches): hash the creator credential and write batch line items in parts The creator credential reaches line metadata only as a sha256 hash, so an OAuth2 token or the master key alias is never written to the bucket. A queued store keeps only whether the batch was billed and downloads the output file again at its turn, logging an error naming the batch if that fails. Lines are built and uploaded one part at a time, each closing once its raw request and result bytes reach BATCH_LINE_ITEMS_MAX_PART_BYTES (16 MiB by default), with the input file indexed by byte span instead of parsed up front. --- .../proxy/common_utils/check_batch_cost.py | 79 +++++------ litellm/batches/batch_line_items.py | 124 +++++++++++++----- litellm/batches/batch_utils.py | 4 +- litellm/constants.py | 4 + .../integrations/batch_line_item_objects.py | 5 +- litellm/integrations/gcs_bucket/gcs_bucket.py | 2 +- litellm/integrations/s3_v2.py | 9 +- litellm/types/utils.py | 4 +- .../test_batch_line_items_cold_storage.py | 60 +++++++-- tests/unit/batches/test_batch_line_items.py | 94 ++++++++++--- .../gcs_bucket/test_gcs_bucket.py | 8 +- tests/unit/integrations/test_s3_v2.py | 6 +- .../common_utils/test_check_batch_cost.py | 122 +++++++++++++++-- 13 files changed, 393 insertions(+), 128 deletions(-) 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 6a342ec20004..8eb36f1c175d 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -13,6 +13,7 @@ from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid from litellm.constants import ( + BATCH_LINE_ITEMS_MAX_PART_BYTES, CLI_SESSION_KEY_PREFIX, MANAGED_OBJECT_STALENESS_CUTOFF_DAYS, MAX_OBJECTS_PER_POLL_CYCLE, @@ -738,10 +739,10 @@ async def _track_completed_batch_cost( model_id: str, batch_id: str, prom_logger: Optional["PrometheusLogger"], - ) -> tuple[str | None, str | None, bytes] | None: + ) -> Optional[Tuple[Optional[str], Optional[str]]]: """ Fetch a completed batch's results, compute cost/usage, and emit the - aretrieve_batch spend log. Returns (model_name, llm_provider, output file bytes) on + aretrieve_batch spend log. Returns (model_name, llm_provider) on success, None when the job can't be routed to a deployment or when another pod claimed it. Raises on results-fetch or cost-computation failures so the caller can leave the job unprocessed and retry it on a @@ -968,7 +969,7 @@ async def _track_completed_batch_cost( api_provider=str(llm_provider) if llm_provider else None, ) - return model_name, str(llm_provider) if llm_provider else None, content_bytes + return model_name, str(llm_provider) if llm_provider else None async def _retire_unbilled_terminal_job( self, job: "_ManagedObjectRow", response: "LiteLLMBatch", model_id: str, batch_id: str @@ -978,15 +979,10 @@ async def _retire_unbilled_terminal_job( 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, b"") + 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, - output_content: bytes, + 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 @@ -1000,42 +996,35 @@ def _store_batch_line_items_in_background( ) return store: Final = asyncio.create_task( - self._store_batch_line_items_in_turn(job, response, model_id, batch_id, output_content) + 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, - output_content: bytes, + 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, output_content) + 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, - output_content: bytes, + 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. Failures are logged, never - raised, so the batch stays billed and retired whether or not its lines were stored.""" + """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, - build_batch_line_items, + 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 @@ -1054,6 +1043,13 @@ async def _store_batch_line_items( **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 @@ -1076,17 +1072,18 @@ async def _store_batch_line_items( if redact_bodies else await self._batch_file_content(response.input_file_id, cost_provider, file_access) ) - # pricing every line of a large batch is CPU-bound, so it runs off the event loop - payload: Final = await asyncio.to_thread( - build_batch_line_items, source, input_content, (output_content, error_content) - ) - if not payload["lines"]: - return - await dispatch_batch_line_items(payload, redact_bodies=redact_bodies) - verbose_proxy_logger.info( - f"CheckBatchCost: handed the {len(payload['lines'])} line items of batch {batch_id} " - "to batch_line_item_callbacks" + 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}" @@ -1278,10 +1275,9 @@ async def _poll_batches(self): continue if tracked is None: continue - tracked_model_name, tracked_provider, output_content = tracked # Track this job for the final metrics summary - processed_models.append((tracked_model_name, tracked_provider)) + processed_models.append(tracked) # mark the job as complete try: @@ -1301,11 +1297,8 @@ async def _poll_batches(self): ) 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 - del tracked, output_content continue - self._store_batch_line_items_in_background(job, response, model_id, batch_id, output_content) - # release this batch's output before the next batch's output is read - del tracked, output_content + 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 ( diff --git a/litellm/batches/batch_line_items.py b/litellm/batches/batch_line_items.py index 427560dff37e..3d175b22b900 100644 --- a/litellm/batches/batch_line_items.py +++ b/litellm/batches/batch_line_items.py @@ -1,17 +1,19 @@ """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, Mapping +from collections.abc import Iterable, Iterator, Mapping from dataclasses import dataclass -from itertools import chain +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, get_file_content_as_dictionary +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 +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, @@ -24,6 +26,7 @@ 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) @@ -39,7 +42,7 @@ class BatchLineItemSource: def line_item_metadata(attribution: Mapping[str, object]) -> StandardBatchLineItemMetadata: return StandardBatchLineItemMetadata( - user_api_key_hash=_text(attribution, "user_api_key_hash"), + 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"), @@ -53,48 +56,98 @@ def _text(attribution: Mapping[str, object], name: str) -> str | None: 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 build_batch_line_items( +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] -) -> StandardBatchLineItemsPayload: - """Pair every output and error line with the input line that shares its custom_id, priced one line at a time.""" - requests: Final = _requests_by_custom_id(input_content) - results: Final = chain.from_iterable(map(get_file_content_as_dictionary, result_contents)) - return StandardBatchLineItemsPayload( - batch_id=source.batch_id, - finished_at=source.finished_at, - lines=tuple([_line_item(source, requests, entry) for entry in results]), - ) +) -> 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 _requests_by_custom_id(input_content: bytes) -> Mapping[str | None, Mapping[str, object]]: - """Each input line 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((_line_key(entry), entry) for entry in get_file_content_as_dictionary(input_content)) +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: entry for key, entry in keyed if counts[key] == 1} + 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, - requests: Mapping[str | None, Mapping[str, object]], - entry: dict[str, object], + 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 - custom_id: Final = _line_key(entry) return StandardBatchLineItem( batch_id=source.batch_id, custom_id=custom_id, status="success" if succeeded else "failure", - request=_request_for(source, requests, custom_id), + request=request, response=entry, usage=StandardBatchLineItemUsage( prompt_tokens=priced.usage.prompt_tokens, @@ -110,19 +163,21 @@ def _line_item( ) -def _request_for( - source: BatchLineItemSource, requests: Mapping[str | None, Mapping[str, object]], custom_id: str | None -) -> Mapping[str, object] | None: - """The input line that shares a result line's custom_id. A Vertex AI line read as "unknown" pairs with none.""" +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 requests.get(custom_id) + 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"], @@ -142,7 +197,8 @@ def without_request_and_response(payload: StandardBatchLineItemsPayload) -> Stan async def dispatch_batch_line_items(payload: StandardBatchLineItemsPayload, redact_bodies: bool) -> None: - """Hand the lines to every batch_line_item_callbacks logger. A failing logger is logged and never stops the rest.""" + """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) ) @@ -155,8 +211,9 @@ async def _deliver( logger: Final = _resolve_logger(callback) if logger is None: verbose_logger.error( - "batch_line_item_callbacks: %r cannot store batch line items, the lines of batch %s were not stored", + "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 @@ -167,9 +224,10 @@ async def _deliver( ) 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 batch %s: %s", + "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, ) diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 3993b518fd48..555c14ae6c9f 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -617,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 9d18fe5489d0..ac84b7278053 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1880,6 +1880,10 @@ class ServerStreamingClassification(str, Enum): RESET_BUDGET_JOB_LOCK_TTL_SECONDS: Final[int] = 900 PROXY_BATCH_POLLING_INTERVAL: Final = int(os.getenv("PROXY_BATCH_POLLING_INTERVAL", 3600)) MAX_OBJECTS_PER_POLL_CYCLE: Final = max(1, int(os.getenv("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, int(os.getenv("MANAGED_OBJECT_STALENESS_CUTOFF_DAYS", 7))) BATCH_OUTPUT_FILE_LOOKUP_TIMEOUT_SECONDS: Final = 10.0 BATCH_OUTPUT_FILE_FALLBACK_RETRY_AFTER_SECONDS: Final = 60 diff --git a/litellm/integrations/batch_line_item_objects.py b/litellm/integrations/batch_line_item_objects.py index 96f7c6d91a6a..888f16080032 100644 --- a/litellm/integrations/batch_line_item_objects.py +++ b/litellm/integrations/batch_line_item_objects.py @@ -23,6 +23,7 @@ def _ndjson_line(line: object) -> str: def batch_line_items_object_name(payload: StandardBatchLineItemsPayload) -> str: - """batch_line_items//.jsonl, the same name on every retry so a rewrite replaces it.""" + """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('/', '_')}.jsonl" + return f"batch_line_items/{day}/{payload['batch_id'].replace('/', '_')}/part-{payload['part']:05d}.jsonl" diff --git a/litellm/integrations/gcs_bucket/gcs_bucket.py b/litellm/integrations/gcs_bucket/gcs_bucket.py index c66d3fdee5b7..5bc7247e02ac 100644 --- a/litellm/integrations/gcs_bucket/gcs_bucket.py +++ b/litellm/integrations/gcs_bucket/gcs_bucket.py @@ -97,7 +97,7 @@ async def async_log_failure_event(self, kwargs, response_obj, start_time, end_ti verbose_logger.exception("GCS Bucket logging error: %s", e) async def async_log_batch_line_items(self, payload: StandardBatchLineItemsPayload) -> None: - """Write the batch's lines as one NDJSON object under batch_line_items/ in the default bucket.""" + """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"], diff --git a/litellm/integrations/s3_v2.py b/litellm/integrations/s3_v2.py index 9e720a9d1fe2..cfd07c9e2af5 100644 --- a/litellm/integrations/s3_v2.py +++ b/litellm/integrations/s3_v2.py @@ -550,7 +550,7 @@ async def async_log_audit_log_event(self, audit_log: StandardAuditLogPayload) -> verbose_logger.exception("S3 audit log error: %s", e) async def async_log_batch_line_items(self, payload: StandardBatchLineItemsPayload) -> None: - """Upload the batch's lines as one NDJSON object under /batch_line_items/.""" + """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 @@ -559,17 +559,18 @@ async def async_log_batch_line_items(self, payload: StandardBatchLineItemsPayloa 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=object_name.rsplit("/", 1)[-1], + s3_object_download_filename="-".join(object_name.rsplit("/", 2)[-2:]), ) - # hashing a whole batch's object is CPU-bound too, so it is also prepared off the event loop + # 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 batch %s were not stored at %s", + "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, ) diff --git a/litellm/types/utils.py b/litellm/types/utils.py index d999a1bbd59e..2387f447df0a 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3695,10 +3695,12 @@ class StandardBatchLineItem(TypedDict): class StandardBatchLineItemsPayload(TypedDict): - """Every line of one finished batch, dispatched to batch_line_item_callbacks.""" + """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, ...]] diff --git a/tests/integration/observability/test_batch_line_items_cold_storage.py b/tests/integration/observability/test_batch_line_items_cold_storage.py index 49da6a0619e2..9c2a1ecd2c4f 100644 --- a/tests/integration/observability/test_batch_line_items_cold_storage.py +++ b/tests/integration/observability/test_batch_line_items_cold_storage.py @@ -1,10 +1,11 @@ """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 one NDJSON object for the batch holding one record per line. 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. Three more owned proxies list an unsupported name next to ``s3_v2``, list a null list, and -check OAuth2 tokens against a local token info endpoint +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 @@ -219,7 +220,12 @@ def _config(directory: Path, s3_url: str, line_item_callbacks: JsonValue, *, oau @contextmanager def _owned_rig( - tmp_path_factory: pytest.TempPathFactory, name: str, line_item_callbacks: JsonValue, *, oauth2: bool = False + 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",)) @@ -253,6 +259,7 @@ def _owned_rig( "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",), @@ -285,6 +292,12 @@ def rig_with_oauth2(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: 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, @@ -483,8 +496,8 @@ def _quiet_logging() -> dict[str, JsonValue]: } -def _object_name(prefix: str, batch_id: str) -> str: - return f"{prefix}batch_line_items/{LINE_DATE}/{batch_id.replace('/', '_')}.jsonl" +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], ...]: @@ -654,7 +667,9 @@ def test_a_batch_created_with_a_jwt_and_no_key_gets_its_lines_with_request_and_r ] -def test_a_batch_created_with_an_oauth2_token_follows_its_teams_message_logging_setting(rig_with_oauth2: Rig) -> None: +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"] @@ -697,6 +712,35 @@ def test_a_batch_created_with_an_oauth2_token_follows_its_teams_message_logging_ ("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: diff --git a/tests/unit/batches/test_batch_line_items.py b/tests/unit/batches/test_batch_line_items.py index eae1a2abda94..d181a2faa86b 100644 --- a/tests/unit/batches/test_batch_line_items.py +++ b/tests/unit/batches/test_batch_line_items.py @@ -1,17 +1,19 @@ import json import logging +from hashlib import sha256 import pytest import litellm from litellm.batches.batch_line_items import ( BatchLineItemSource, - build_batch_line_items, + 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 @@ -20,6 +22,7 @@ INPUT_RATE = 0.001 OUTPUT_RATE = 0.002 FINISHED_AT = 1791494130 +ALICE_KEY_HASH = sha256(b"sk-alice").hexdigest() def _jsonl(*lines: dict) -> bytes: @@ -73,14 +76,20 @@ def _source(**overrides) -> BatchLineItemSource: "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": "hash-alice", "user_api_key_team_id": "team-alpha", "tags": ["env:prod"]} + {"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 build_batch_line_items( + 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"))), @@ -103,7 +112,7 @@ async def async_log_batch_line_items(self, payload: StandardBatchLineItemsPayloa @pytest.mark.asyncio async def test_each_result_line_is_paired_with_the_request_that_shares_its_custom_id(): - payload = build_batch_line_items( + 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"))), @@ -123,7 +132,7 @@ async def test_each_result_line_is_paired_with_the_request_that_shares_its_custo "model": "gpt-4o-mini", "custom_llm_provider": "openai", "metadata": { - "user_api_key_hash": "hash-alice", + "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, @@ -143,7 +152,7 @@ async def test_each_result_line_is_paired_with_the_request_that_shares_its_custo 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 = build_batch_line_items(_source(), b"", (outputs,)) + 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 ) @@ -165,7 +174,7 @@ async def test_anthropic_results_without_an_input_file_keep_their_status_and_usa } errored = {"custom_id": "a-2", "result": {"type": "errored", "error": {"type": "invalid_request_error"}}} - payload = build_batch_line_items( + payload = _whole_batch( _source(model="claude-sonnet-5-5", custom_llm_provider="anthropic", cost_provider="anthropic"), b"", (_jsonl(succeeded, errored),), @@ -182,7 +191,7 @@ 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 = build_batch_line_items( + payload = _whole_batch( _source(custom_llm_provider="bedrock", cost_provider="bedrock"), _jsonl(request), (_jsonl(result),) ) @@ -200,7 +209,7 @@ def test_vertex_lines_pair_with_the_input_row_that_carries_their_custom_id(): vertex_rows = get_file_content_as_dictionary(vertex_input) assert all("custom_id" not in row for row in vertex_rows) - payload = build_batch_line_items( + 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)),), @@ -227,7 +236,7 @@ def test_a_vertex_embeddings_entry_split_into_several_rows_pairs_with_none_of_th 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 = build_batch_line_items( + 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)),), @@ -242,7 +251,7 @@ def test_a_vertex_embeddings_entry_split_into_several_rows_pairs_with_none_of_th def test_a_vertex_input_row_without_a_custom_id_pairs_with_nothing(): unlabeled = {"request": {"contents": [{"role": "user", "parts": [{"text": "hi"}]}]}} - payload = build_batch_line_items( + 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}),), @@ -274,7 +283,7 @@ def test_a_vertex_input_row_without_a_custom_id_pairs_with_nothing(): ], ) def test_a_vertex_line_read_as_unknown_pairs_with_no_input_row(vertex_input): - payload = build_batch_line_items( + 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)),), @@ -291,6 +300,7 @@ async def test_redaction_drops_only_the_request_and_response(): 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"]) @@ -303,13 +313,67 @@ async def test_the_object_holds_one_json_line_per_batch_line_under_its_finished_ 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.jsonl" assert ( - batch_line_items_object_name({**payload, "batch_id": "projects/p/locations/l/batchPredictionJobs/9"}) - == "batch_line_items/2026-10-08/projects_p_locations_l_batchPredictionJobs_9.jsonl" + 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"] diff --git a/tests/unit/integrations/gcs_bucket/test_gcs_bucket.py b/tests/unit/integrations/gcs_bucket/test_gcs_bucket.py index e7fdf5b3402a..6f143a244f6a 100644 --- a/tests/unit/integrations/gcs_bucket/test_gcs_bucket.py +++ b/tests/unit/integrations/gcs_bucket/test_gcs_bucket.py @@ -347,13 +347,13 @@ def stored(request: httpx.Request) -> httpx.Response: ) await logger.async_log_batch_line_items( - {"batch_id": "bGl0ZWxsbV9wcm94eTtiYXRjaA", "finished_at": 1791494130, "lines": lines} + {"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.jsonl" + 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) @@ -364,7 +364,7 @@ async def test_a_rejected_batch_line_items_write_raises_for_the_dispatcher(monke with pytest.raises(httpx.HTTPStatusError): await logger.async_log_batch_line_items( - {"batch_id": "b", "finished_at": 1791494130, "lines": ({"custom_id": "ok-1"},)} + {"batch_id": "b", "finished_at": 1791494130, "part": 0, "lines": ({"custom_id": "ok-1"},)} ) assert uploads.call_count == 1 @@ -390,7 +390,7 @@ def _serialize(payload: StandardBatchLineItemsPayload) -> str: monkeypatch.setattr(gcs_bucket_module, "batch_line_items_ndjson", _serialize) - await logger.async_log_batch_line_items({"batch_id": "b", "finished_at": 1791494130, "lines": ()}) + 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 90e1d3744836..80f79517fab5 100644 --- a/tests/unit/integrations/test_s3_v2.py +++ b/tests/unit/integrations/test_s3_v2.py @@ -5866,6 +5866,7 @@ def _batch_line_items_payload() -> dict: return { "batch_id": "bGl0ZWxsbV9wcm94eTtiYXRjaA", "finished_at": 1791494130, + "part": 0, "lines": (line, {**line, "custom_id": "ok-2"}), } @@ -5900,8 +5901,9 @@ async def test_batch_line_items_are_one_ndjson_object_under_the_configured_path( (request,) = requests assert request.method == "PUT" - assert request.url.path == "/cold/batch_line_items/2026-10-08/bGl0ZWxsbV9wcm94eTtiYXRjaA.jsonl" + 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 == [] @@ -5913,7 +5915,7 @@ async def test_batch_line_items_without_an_s3_path_sit_at_the_bucket_root(): 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.jsonl" + "/batch_line_items/2026-10-08/bGl0ZWxsbV9wcm94eTtiYXRjaA/part-00000.jsonl" ] 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 d4faba6e6268..75a90be21b86 100644 --- a/tests/unit/proxy/common_utils/test_check_batch_cost.py +++ b/tests/unit/proxy/common_utils/test_check_batch_cost.py @@ -3324,10 +3324,19 @@ async def test_lines_reach_the_callbacks_after_the_claimed_batch_is_billed_and_r await instance.check_batch_cost() await self._stored(instance) - assert journal == [f"read {_CLAIM_OUTPUT_FILE_ID}", "claim", "bill", "finalize", "read file-input-1", "lines"] + 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} @@ -3344,14 +3353,14 @@ async def test_the_lines_are_priced_off_the_event_loop(self, monkeypatch): journal: list = [] row, instance = self._setup(journal) monkeypatch.setattr(litellm, "batch_line_item_callbacks", (self._recording_logger(journal),)) - build = batch_line_items.build_batch_line_items + price = batch_line_items._line_item threads = [] - def _build(*args): + def _price(*args): threads.append(threading.current_thread()) - return build(*args) + return price(*args) - monkeypatch.setattr(batch_line_items, "build_batch_line_items", _build) + monkeypatch.setattr(batch_line_items, "_line_item", _price) with self._provider(journal): await instance.check_batch_cost() await self._stored(instance) @@ -3393,7 +3402,7 @@ async def test_a_creator_with_logging_on_gets_the_request_and_response(self, mon await instance.check_batch_cost() await self._stored(instance) - assert journal[-3:] == ["finalize", "read file-input-1", "lines"] + 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) @@ -3587,10 +3596,13 @@ async def async_log_batch_line_items(self, payload): monkeypatch.setattr(litellm, "batch_line_item_callbacks", (_HungSink(),)) - with self._provider(journal): - await asyncio.wait_for(instance.check_batch_cost(), timeout=5) + 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) @@ -3598,7 +3610,15 @@ async def async_log_batch_line_items(self, payload): assert billed_while_hung == ( True, "complete", - [f"read {_CLAIM_OUTPUT_FILE_ID}", "claim", "bill", "finalize", "read file-input-1", "lines started"], + [ + 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" @@ -3637,7 +3657,7 @@ async def test_batches_store_their_lines_one_at_a_time_in_the_order_they_were_re monkeypatch.setattr(litellm, "batch_line_item_callbacks", (self._recording_logger(journal),)) running: list = [] - async def _store(job, response, model_id, batch_id, output_content): + 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) @@ -3645,7 +3665,7 @@ async def _store(job, response, model_id, batch_id, output_content): 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, b"") + 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"] @@ -3659,10 +3679,12 @@ async def test_lines_are_stored_only_once_the_poll_has_billed_and_retired_every_ 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", b"") + 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, output_content): + 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) @@ -3680,6 +3702,80 @@ async def _store(job, response, model_id, batch_id, output_content): "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_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, From bc925d4e4ea9d66f3aa70316c68f841919f85506 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Wed, 14 Oct 2026 18:28:38 -0700 Subject: [PATCH 88/89] docs(batches): say async_log_batch_line_items is called once per part --- litellm/integrations/custom_logger.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py index 4c6410b7e4be..52577d0c647c 100644 --- a/litellm/integrations/custom_logger.py +++ b/litellm/integrations/custom_logger.py @@ -197,7 +197,7 @@ 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 finished batch with every line, when this logger is listed in batch_line_item_callbacks.""" + """Called once per part of a finished batch's lines, when this logger is listed in batch_line_item_callbacks.""" #### PROMPT MANAGEMENT HOOKS #### From 315d705be711d2872f7c876935762edf9521b148 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Wed, 14 Oct 2026 21:20:14 -0700 Subject: [PATCH 89/89] fix(batches): log an error when a batch's deployment is gone before its line items are stored --- .../proxy/common_utils/check_batch_cost.py | 4 ++++ .../common_utils/test_check_batch_cost.py | 24 +++++++++++++++++++ 2 files changed, 28 insertions(+) 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 8eb36f1c175d..32f06deb3a51 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -1031,6 +1031,10 @@ async def _store_batch_line_items( 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( 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 75a90be21b86..f5c1f8107cb4 100644 --- a/tests/unit/proxy/common_utils/test_check_batch_cost.py +++ b/tests/unit/proxy/common_utils/test_check_batch_cost.py @@ -3731,6 +3731,30 @@ async def _output_gone_after_billing(file_id, **kwargs): 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