From b38b8679189e96c8f20be9d572b1b14318b402d3 Mon Sep 17 00:00:00 2001 From: shivam Date: Thu, 17 Sep 2026 22:48:28 +0000 Subject: [PATCH 01/58] 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 13e9e5093a8..457bcf8c147 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 @@ class CheckBatchCost: } }, **({"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 71857877e53..7aaef9b91d5 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -346,6 +346,7 @@ anthropic_prompt_caching_ttl: Optional[Literal["5m", "1h"]] = ( ) 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 00000000000..077457dd299 --- /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 40621a2f68d..5e8139c1e8d 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -3027,6 +3027,18 @@ class Logging(LiteLLMLoggingBaseClass): 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 680c63393e8..dacbfb8ec81 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 1ae106be390..8ad4de38095 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 @@ class _ProxyDBLogger(CustomLogger): 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 d7d8413d2ce..e47361d1de5 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -6044,6 +6044,13 @@ class ProxyConfig: 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 @@ class ProxyConfig: # 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 aaa16fd2d44..8db91ffcb52 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 00000000000..7fc60ddba77 --- /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 dfc95db3e14..b42dc41f097 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/58] 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 5e8139c1e8d..acba0eb1025 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -3027,7 +3027,7 @@ class Logging(LiteLLMLoggingBaseClass): 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 cb2c70e74c8..20c8258e440 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -34,6 +34,7 @@ if TYPE_CHECKING: _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 7fc60ddba77..78265c2fea3 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 json 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 eba8d912fe0..9b28e42f93b 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/58] 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 dfe84a28d52..14c98048619 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 35eaa20c39e..e7e46cdec1d 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/58] 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 077457dd299..65227c24c8e 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/58] 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 872875cc535..a74b996eaa0 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/58] 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 78265c2fea3..32eeb326c60 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 41c4956dba6..7503a4b30ff 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/58] 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 457bcf8c147..b41b3a6e1ad 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 @@ class CheckBatchCost: } }, **({"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 @@ class CheckBatchCost: 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 65227c24c8e..8a4206d3606 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 32eeb326c60..fe03e348ce0 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/58] 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 b41b3a6e1ad..d3bd830a9c0 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 @@ class CheckBatchCost: 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 8a4206d3606..911a4d478ec 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 acba0eb1025..64cbf8eff4b 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -574,6 +574,7 @@ class Logging(LiteLLMLoggingBaseClass): 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/58] 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 a8d9cf1df8b..828e9d8cab1 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 00000000000..0407f68cc84 --- /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/58] 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 a7b875dc71c..d8b6d0537f1 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -7593,7 +7593,7 @@ class ProxyConfig: 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/58] 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 0407f68cc84..3e1e3ad4e08 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 test_completed_batch_emits_paired_request_response_callback_events_per_jsonl 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/58] 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 d8b6d0537f1..c08a7b8156d 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -6516,13 +6516,6 @@ class ProxyConfig: 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/58] 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 3e1e3ad4e08..e0f32873665 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 test_completed_batch_emits_paired_request_response_callback_events_per_jsonl ) 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 test_completed_batch_emits_paired_request_response_callback_events_per_jsonl 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/58] 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 c08a7b8156d..9eab5478e3f 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -7586,11 +7586,8 @@ class ProxyConfig: 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 7fa63a95264..a52b119e02b 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/58] 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 911a4d478ec..b91ef3dccd3 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 06c8218eea5..11b1d4725f7 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 @@ class Logging(LiteLLMLoggingBaseClass): 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 fe03e348ce0..62c6bd359f2 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 @@ ANTHROPIC_OUTPUT_JSONL = json.dumps( "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/58] 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 11b1d4725f7..2cf1b7e4673 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -4256,7 +4256,7 @@ class Logging(LiteLLMLoggingBaseClass): ) 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/58] 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 2cf1b7e4673..de740fdeca7 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -4256,7 +4256,7 @@ class Logging(LiteLLMLoggingBaseClass): ) 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/58] 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 b91ef3dccd3..a414c419710 100644 --- a/litellm/batches/batch_line_item_logging.py +++ b/litellm/batches/batch_line_item_logging.py @@ -3,7 +3,7 @@ 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 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 @@ if TYPE_CHECKING: _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 62c6bd359f2..1c59eac24bb 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/58] 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 a414c419710..02dd9cceaeb 100644 --- a/litellm/batches/batch_line_item_logging.py +++ b/litellm/batches/batch_line_item_logging.py @@ -25,7 +25,9 @@ from litellm.types.utils import ( 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 @@ _BatchLineResult: TypeAlias = "ModelResponse | EmbeddingResponse | ResponsesAPIR 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 7209ac6a1e7..468a0c63eae 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 de740fdeca7..ea410d6d6ba 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 @@ class Logging(LiteLLMLoggingBaseClass): 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 b7c2ce3c568..7f9294c58df 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -48,6 +48,7 @@ from litellm.types.llms.anthropic import ( 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 e4001566b8c..6d43e9ee7f9 100644 --- a/litellm/llms/bedrock/batches/transformation.py +++ b/litellm/llms/bedrock/batches/transformation.py @@ -27,7 +27,7 @@ from litellm.types.llms.openai import ( 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 1c59eac24bb..0be7f76b594 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 @@ ANTHROPIC_OUTPUT_JSONL = json.dumps( ).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/58] 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 468a0c63eae..240ed048b63 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/58] 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 | 91 +++++++++++++------ .../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, 115 insertions(+), 33 deletions(-) diff --git a/litellm/batches/batch_line_item_logging.py b/litellm/batches/batch_line_item_logging.py index 02dd9cceaeb..3d54625d3fe 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 + from litellm.types.utils import LlmProviders + from litellm.utils import ProviderConfigManager - bedrock_result: Final = bedrock_batch_line_to_response(response_body, model) - if bedrock_result is None: + 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 34c622d4cf6..904eae2f4f2 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 @@ from litellm.types.llms.openai import ( 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 @@ class BaseBatchesConfig(ABC): 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 6d43e9ee7f9..f8d71bab2c7 100644 --- a/litellm/llms/bedrock/batches/transformation.py +++ b/litellm/llms/bedrock/batches/transformation.py @@ -144,6 +144,11 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): 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 0be7f76b594..1f37969bcfb 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 347c459a369..4bce3c37504 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 httpx 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/58] 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 3d54625d3fe..98b6f0a1757 100644 --- a/litellm/batches/batch_line_item_logging.py +++ b/litellm/batches/batch_line_item_logging.py @@ -31,6 +31,20 @@ _BatchLineProvider: TypeAlias = Literal[ _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 d7fbe9f7e09..088a4a41428 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -5,7 +5,7 @@ import logging 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 7f9294c58df..14c28b0ce68 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 f8d71bab2c7..c61cd0dfeef 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 0339cf4dfea..6777e0c88f3 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 import ModelResponse, Router 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 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): - 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 e07b96e5773..a1f60fc4d87 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._logging import verbose_proxy_logger 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 @@ class _PROXY_MaxBudgetPerSessionHandler(CustomLogger): """ 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 cfa54ae01a2..de3d72ddb5b 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/litellm/proxy/hooks/model_max_budget_limiter.py @@ -11,6 +11,7 @@ import litellm 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 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): 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 d41acadc4dd..a6dde585309 100644 --- a/litellm/proxy/hooks/parallel_request_limiter.py +++ b/litellm/proxy/hooks/parallel_request_limiter.py @@ -11,7 +11,10 @@ from litellm import DualCache, EmbeddingResponse, ModelResponse, TextCompletionR 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 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): ) 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 33744906b13..d5b31624064 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._logging import verbose_proxy_logger 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 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): """ 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 66df6578bc8..e8d508feb07 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -13,6 +13,7 @@ from litellm.litellm_core_utils.core_helpers import ( _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 @@ class _ProxyDBLogger(CustomLogger): 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 1f37969bcfb..eb8ae26222d 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 @@ cost, or error propagation fails here. 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 6eeea271127..368248fd956 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 @@ from litellm.litellm_core_utils.core_helpers import ( 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 @@ class TestIsExpectedClientError: 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 ffb60fb4651..b062efb3e81 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 8c0dcd3383c..904b5d67556 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/58] 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 088a4a41428..0bcf222ec6b 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/58] 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 a6dde585309..4ad2062a4e2 100644 --- a/litellm/proxy/hooks/parallel_request_limiter.py +++ b/litellm/proxy/hooks/parallel_request_limiter.py @@ -699,6 +699,8 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): 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 d5b31624064..981c0b81a07 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -4681,6 +4681,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): 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 0e2683dcbfd..cbc16d734a1 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 904b5d67556..ea8d5016463 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/58] 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 e0f32873665..97cad137b51 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 test_completed_batch_emits_paired_request_response_callback_events_per_jsonl ) 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 test_completed_batch_emits_paired_request_response_callback_events_per_jsonl 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 test_completed_batch_emits_paired_request_response_callback_events_per_jsonl 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/58] 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 98b6f0a1757..cf4de486328 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 14c28b0ce68..77b8f906cbe 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 c61cd0dfeef..b4c16f010cc 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/58] 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 3e094df7ac8..2691610a2aa 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/litellm/router_strategy/budget_limiter.py @@ -32,6 +32,7 @@ from litellm.caching.redis_cache import RedisCache, RedisPipelineIncrementOperat 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 @@ class RouterBudgetLimiting(CustomLogger): 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 4cc8fe78811..af64b18388f 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 @@ async def test_get_llm_provider_for_deployment_dict_view_supports_mapping_and_at @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 test_push_returns_before_redis_answers(disable_budget_sync): 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/58] 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 cf4de486328..3be093657ec 100644 --- a/litellm/batches/batch_line_item_logging.py +++ b/litellm/batches/batch_line_item_logging.py @@ -12,6 +12,7 @@ from litellm.batches.batch_utils import ( _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 eb8ae26222d..2d126fc3f29 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 @@ class _RecordingLogger(CustomLogger): @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/58] 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 2d126fc3f29..97442f59a0d 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 @@ class _RecordingLogger(CustomLogger): @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 @@ VERTEX_NATIVE_OUTPUT_JSONL = b"\n".join( ) +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/58] 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/58] 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 3be093657ec..f24ae9fb432 100644 --- a/litellm/batches/batch_line_item_logging.py +++ b/litellm/batches/batch_line_item_logging.py @@ -14,6 +14,7 @@ from litellm.batches.batch_utils import ( _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 @@ _CALL_TYPE_BY_BATCH_URL: Final = MappingProxyType( _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 73c993f57da..bea6a9d2439 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -267,6 +267,10 @@ import litellm 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 97cad137b51..a707a7a4f7a 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 test_completed_batch_emits_paired_request_response_callback_events_per_jsonl 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 test_completed_batch_emits_paired_request_response_callback_events_per_jsonl 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 ffbcd1fba24..e3737b0b761 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 573bfc40c96..c59c4f3422d 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 @@ class TestRedisAuthCacheFlag: 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 97442f59a0d..9a855b602e5 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 @@ from unittest.mock import AsyncMock, patch 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 @@ class _RecordingLogger(CustomLogger): 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/58] 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 f24ae9fb432..21dd23bdb59 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 9a855b602e5..f5514c5814c 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/58] 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 21dd23bdb59..94c2bb0bb95 100644 --- a/litellm/batches/batch_line_item_logging.py +++ b/litellm/batches/batch_line_item_logging.py @@ -5,6 +5,8 @@ from datetime import datetime 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 @@ _LINE_ITEM_CLAIM_TTL_SECONDS: Final = 30 * 24 * 60 * 60 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 f5514c5814c..e78c0a1c307 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 @@ from unittest.mock import AsyncMock, patch 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/58] 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 94c2bb0bb95..171c03d6f00 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 e78c0a1c307..ca99f6229aa 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/58] 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 1365d22d038..83485d1efca 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -5951,6 +5951,9 @@ class StandardLoggingPayloadSetup: 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 ca99f6229aa..877b6b67854 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/58] 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 db28317ae2c..cb41eba32eb 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 @@ class CheckBatchCost: 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 @@ class CheckBatchCost: 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 @@ class CheckBatchCost: "_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 @@ class CheckBatchCost: 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 171c03d6f00..db9f0c9b6b1 100644 --- a/litellm/batches/batch_line_item_logging.py +++ b/litellm/batches/batch_line_item_logging.py @@ -9,6 +9,7 @@ from typing_extensions import assert_never 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 8160618f2c5..3660dfed94b 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 de573d04f98..d5b111a41cf 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -32,7 +32,11 @@ from litellm._logging import ( 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 @@ class Logging(LiteLLMLoggingBaseClass): 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 @@ class Logging(LiteLLMLoggingBaseClass): 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 @@ class Logging(LiteLLMLoggingBaseClass): 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 73adb391481..e8a90b385f2 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 877b6b67854..a8edb920bf7 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 @@ from litellm.batches.batch_line_item_logging import ( 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 b8b922f72a7..367039cc0b4 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 60b7ed32399..9239e47ee7f 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 @@ class TestRetrieveBatchCostPassesModelIdentity: 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 @@ class TestRetrieveBatchPricesOnlyFinalBatches: 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 @@ class TestRetrieveBatchPricesOnlyFinalBatches: @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 4e7effad3bf..41470cc0d2e 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 @@ class TestCheckBatchCost: 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 @@ class TestCheckBatchCost: 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/58] 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 db9f0c9b6b1..580100b736f 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 a8edb920bf7..e338cad5d21 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 9239e47ee7f..1fd8a681873 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 @@ class TestRetrieveBatchReusesFetchedResultFiles: @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 @@ class TestRetrieveBatchReusesFetchedResultFiles: @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 @@ class TestRetrieveBatchReusesFetchedResultFiles: @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/58] 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 cb41eba32eb..d6ef044ef4c 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 @@ class CheckBatchCost: 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 00c940a8116..f2694c6c298 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -3290,7 +3290,7 @@ class Logging(LiteLLMLoggingBaseClass): 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 6777e0c88f3..09cf8af9681 100644 --- a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py +++ b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py @@ -5,7 +5,7 @@ Dynamic rate limiter v3 - Saturation-aware priority-based rate limiting 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 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): 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 cd28a5771d6..cb341356452 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/litellm/proxy/hooks/model_max_budget_limiter.py @@ -4,7 +4,7 @@ import time 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 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): ) -> 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 b04a4823dd5..a628f055f80 100644 --- a/litellm/proxy/hooks/parallel_request_limiter.py +++ b/litellm/proxy/hooks/parallel_request_limiter.py @@ -492,7 +492,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): ) # 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 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): 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 e47dcacd91e..603d3be53c9 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -5008,7 +5008,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): 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 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): 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 e5b2550328a..d2ff21ca3ac 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -6829,7 +6829,7 @@ class ProxyConfig: 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 cb58afef044..7e00793aea0 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/litellm/router_strategy/budget_limiter.py @@ -491,7 +491,7 @@ class RouterBudgetLimiting(CustomLogger): 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/58] 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 f2694c6c298..59a0cf83bb5 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/58] 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 580100b736f..3fb3dd1cb89 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 7af16f712f7..1a0a3c145e4 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -3285,13 +3285,16 @@ class Logging(LiteLLMLoggingBaseClass): 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/58] 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 3fb3dd1cb89..fe6a027da97 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 2e1b9cd5825..591383d2b81 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 2039e3fb954..b9b22fa0da4 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 b4c16f010cc..32fdf7e8997 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/58] 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 c1401e111bb..431512ed33a 100644 --- a/litellm/integrations/clickhouse/clickhouse_spend_logger.py +++ b/litellm/integrations/clickhouse/clickhouse_spend_logger.py @@ -16,6 +16,7 @@ from litellm._logging import verbose_logger 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 591383d2b81..32fd00c4f4d 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -947,6 +947,10 @@ class Logging(LiteLLMLoggingBaseClass): **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 24554e61516..4cc0bdcbe8c 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -78,6 +78,7 @@ from litellm.litellm_core_utils.core_helpers import ( 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 @@ class Router: # 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 @@ class Router: - 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 @@ class Router: """ 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 9ab670e4b95..12afa512e10 100644 --- a/litellm/router_strategy/least_busy.py +++ b/litellm/router_strategy/least_busy.py @@ -9,6 +9,7 @@ from litellm._logging import verbose_router_logger 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 22c321c65fb..d47584ebbfc 100644 --- a/litellm/router_strategy/lowest_cost.py +++ b/litellm/router_strategy/lowest_cost.py @@ -8,6 +8,7 @@ from litellm import ModelResponse, token_counter, verbose_logger 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 @@ class LowestCostLoggingHandler(CustomLogger): 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 @@ class LowestCostLoggingHandler(CustomLogger): ) 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 66c8227195d..a03fcde74e5 100644 --- a/litellm/router_strategy/lowest_latency.py +++ b/litellm/router_strategy/lowest_latency.py @@ -12,7 +12,11 @@ import litellm 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 @@ class LowestLatencyLoggingHandler(CustomLogger): 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 @@ class LowestLatencyLoggingHandler(CustomLogger): """ 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 @@ class LowestLatencyLoggingHandler(CustomLogger): ) 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 d4abf1f8f70..3575db0aa53 100644 --- a/litellm/router_strategy/lowest_tpm_rpm.py +++ b/litellm/router_strategy/lowest_tpm_rpm.py @@ -8,6 +8,7 @@ from litellm import token_counter 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 @@ class LowestTPMLoggingHandler(CustomLogger): 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 @@ class LowestTPMLoggingHandler(CustomLogger): 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 25564a80e0a..02cea320935 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 import token_counter 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 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): 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 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): ) 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 b183bf84ea4..a2eed7ce3e9 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 deeb9c198d2..2690dd06b58 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 @@ class TestRetrieveBatchReusesFetchedResultFiles: 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 c2fa41f4ca8..cd0537dd8f5 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 625f648bec4..55147b316db 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 test_v2_subclass_overriding_async_get_available_deployments_with_the_o 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 96dddf15869..d04495e12d9 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/58] 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 3659248ad3d..4a10e29c702 100644 --- a/litellm/integrations/clickhouse/clickhouse_spend_logger.py +++ b/litellm/integrations/clickhouse/clickhouse_spend_logger.py @@ -205,6 +205,10 @@ class ClickHouseSpendLogger(ClickHouseBatchLogger): 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 0bb7735477a..37716f067e8 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/58] 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 df8034d8764..a53159eb6d5 100644 --- a/litellm/batches/batch_line_item_logging.py +++ b/litellm/batches/batch_line_item_logging.py @@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Final, Literal, TypeAlias, cast, get_args 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 @@ batch_line_item_claim_cache: Final = DualCache() _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/58] 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 538dd95abdd..73017b79f7b 100644 --- a/litellm/integrations/datadog/datadog_cost_management.py +++ b/litellm/integrations/datadog/datadog_cost_management.py @@ -14,6 +14,7 @@ from litellm.integrations.datadog.datadog_handler import ( 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 @@ class DatadogCostManagementLogger(CustomBatchLogger): 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 594427b1e0a..193f809a5fd 100644 --- a/litellm/integrations/lago.py +++ b/litellm/integrations/lago.py @@ -11,6 +11,7 @@ import litellm 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 @@ class LagoLogger(CustomLogger): 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 @@ class LagoLogger(CustomLogger): 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 a2cedeba0f5..8bd77c981f6 100644 --- a/litellm/integrations/newrelic/newrelic_metrics.py +++ b/litellm/integrations/newrelic/newrelic_metrics.py @@ -34,6 +34,7 @@ from httpx import HTTPStatusError, Response 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 @@ class NewRelicMetricsLogger(CustomBatchLogger): 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 db2fe386dec..029c684f969 100644 --- a/litellm/integrations/openmeter.py +++ b/litellm/integrations/openmeter.py @@ -9,6 +9,7 @@ import httpx 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 @@ class OpenMeterLogger(CustomLogger): } 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 @@ class OpenMeterLogger(CustomLogger): 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 8d588896b2f..bd9d1057a9d 100644 --- a/litellm/integrations/opentelemetry.py +++ b/litellm/integrations/opentelemetry.py @@ -28,6 +28,7 @@ from litellm.integrations.otel.model.db_endpoint import db_span_attributes 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 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): 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 e1623f4697f..005658edb45 100644 --- a/litellm/integrations/otel/plumbing/metrics.py +++ b/litellm/integrations/otel/plumbing/metrics.py @@ -33,6 +33,7 @@ from litellm.integrations.otel.model.semconv import ( 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 @@ class GenAIMetricRecorder: 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 6468dc41ea5..47b5cae63fe 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -35,6 +35,7 @@ from litellm.integrations.prometheus_helpers.bounded_prometheus_series_tracker i 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 @@ class PrometheusLogger(CustomLogger): 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 @@ class PrometheusLogger(CustomLogger): ) 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 00000000000..d31c25aa621 --- /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/58] 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 d31c25aa621..a89d8ea9a53 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/58] 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 005658edb45..3af54f7f459 100644 --- a/litellm/integrations/otel/plumbing/metrics.py +++ b/litellm/integrations/otel/plumbing/metrics.py @@ -222,22 +222,24 @@ class GenAIMetricRecorder: 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 @@ class GenAIMetricRecorder: 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 a89d8ea9a53..525f6ae15d0 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 @@ class RecordingHTTP: 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/58] 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 03bca29f47d..86dfb682fe1 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 d7ddd5c3326..354178b6f21 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/58] 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 4e2444b42e1..1a4e7810de7 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -35,8 +35,8 @@ from litellm._logging import ( 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/58] 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 ed564bd552e..b0ace15a84a 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 29d41d7842b..899584520d6 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/58] 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 a53159eb6d5..9565614f78f 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 d10a7a1a46d..0a18fc19cd1 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 62cc69a7518..923059d01b9 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -75,6 +75,7 @@ from litellm.litellm_core_utils.classifier_logging import ( 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 @@ class Logging(LiteLLMLoggingBaseClass): 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 @@ class Logging(LiteLLMLoggingBaseClass): 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 @@ class StandardLoggingPayloadSetup: 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 00000000000..46d36020274 --- /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 e338cad5d21..43da11528b0 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 84aa3e69562..07d1b1024b1 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.caching.caching import DualCache 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/58] 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 167f4dac82c..c885fde6ae4 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 @@ class CheckBatchCost: 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 9565614f78f..854f8c2ef34 100644 --- a/litellm/batches/batch_line_item_logging.py +++ b/litellm/batches/batch_line_item_logging.py @@ -3,7 +3,7 @@ import uuid 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 @@ _BatchLineProvider: TypeAlias = Literal[ "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 @@ _SECRET_PARAM_KEYS: Final = frozenset( 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 26ae92526e0..38736b43224 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 0a18fc19cd1..8ffcc4939dc 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -19,9 +19,9 @@ if TYPE_CHECKING: 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 923059d01b9..26d9035e7a2 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -963,7 +963,7 @@ class Logging(LiteLLMLoggingBaseClass): 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 @@ class Logging(LiteLLMLoggingBaseClass): # 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 @@ class Logging(LiteLLMLoggingBaseClass): 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 @@ class Logging(LiteLLMLoggingBaseClass): 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 @@ class Logging(LiteLLMLoggingBaseClass): 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 0633280ebf4..f65cfcb18fa 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 51608ad3a05..136cf8f1b99 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 707ab48219b..f07656bc51d 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -7087,12 +7087,9 @@ class ProxyConfig: 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 43da11528b0..d383546b7bf 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 @@ BODYLESS_ERR_OUTPUT_JSONL = json.dumps({"custom_id": "bare", "response": None, " @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 525f6ae15d0..9cec64b1273 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 038c991040f..0b0708955bb 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/58] 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 f65cfcb18fa..3943cb823fd 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 07d1b1024b1..9bb6d4ec85d 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/58] 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 854f8c2ef34..bde63efb30f 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 d383546b7bf..aea0ae63ebe 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/58] 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 46d36020274..ea93f27075d 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 @@ 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"), ) @@ -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/58] 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 ea93f27075d..c5502e4957d 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 os 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/58] 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 bde63efb30f..854f8c2ef34 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 aea0ae63ebe..d383546b7bf 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/58] 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 c5502e4957d..d99d68d033f 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 d383546b7bf..f014c6e6fde 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)