From fa5a4f06edec9355d91b864fdc63341b708db8f0 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 29 Jul 2026 00:24:30 +0000 Subject: [PATCH 01/68] fix(headroom): resolve CCR retrieval on streaming /chat/completions Streaming chat completions returned a CustomStreamWrapper, so the agentic loop dispatch in main.py (guarded on ModelResponse) never ran and the headroom_retrieve tool call was streamed straight to a client that never declared the tool. The guardrail now converts a CCR stream request to a non-streaming call in its deployment hook and the loop fake-streams the resolved answer back. --- .../chat_completion_agentic_loop.py | 49 +++++-- .../guardrail_hooks/headroom/headroom.py | 22 ++- litellm/proxy/litellm_pre_call_utils.py | 1 + litellm/types/integrations/custom_logger.py | 8 +- litellm/types/utils.py | 1 + .../guardrail_hooks/test_headroom.py | 136 +++++++++++++++++- 6 files changed, 202 insertions(+), 15 deletions(-) diff --git a/litellm/litellm_core_utils/chat_completion_agentic_loop.py b/litellm/litellm_core_utils/chat_completion_agentic_loop.py index b7262a42324..39d36f25143 100644 --- a/litellm/litellm_core_utils/chat_completion_agentic_loop.py +++ b/litellm/litellm_core_utils/chat_completion_agentic_loop.py @@ -1,12 +1,14 @@ # this is a patch to allow for agentic loops covering llm_http_handler.py and openai sdk based calling flows for the .completion() api import json +from collections.abc import Mapping from typing import cast from litellm._logging import verbose_logger from litellm.integrations.custom_logger import CustomLogger from litellm.types.integrations.custom_logger import ( CHAT_COMPLETION_AGENTIC_SURFACE, + HEADROOM_CONVERTED_STREAM_KEY, NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES, AgenticLoopPlan, AgenticLoopRequestPatch, @@ -46,6 +48,12 @@ def _post_hook_overridden(callback: CustomLogger) -> bool: return getattr(func, "__func__", func) is not getattr(base, "__func__", base) +def _converted_stream_requested(kwargs: Mapping[str, object]) -> bool: + return bool( + kwargs.get("_code_interpreter_interception_converted_stream") or kwargs.get(HEADROOM_CONVERTED_STREAM_KEY) + ) + + def _coerce_int(value: object, default: int) -> int: return int(value) if isinstance(value, (int, str)) else default @@ -80,16 +88,25 @@ def _check_agentic_loop_safety( return fingerprint -def _wrap_response_as_fake_stream(response: object) -> object: - if getattr(response, "object", None) == "chat.completion.chunk": +def _wrap_response_as_fake_stream( + response: object, + *, + model: str, + custom_llm_provider: str, + logging_obj: object, +) -> object: + if isinstance(response, CustomStreamWrapper): return response - if not hasattr(response, "choices"): + if not isinstance(response, ModelResponse): return response - from litellm.llms.base_llm.base_model_iterator import ( - convert_model_response_to_streaming, - ) + from litellm.llms.base_llm.base_model_iterator import MockResponseIterator - return convert_model_response_to_streaming(cast(ModelResponse, response)) + return CustomStreamWrapper( + completion_stream=MockResponseIterator(model_response=response), + model=model, + custom_llm_provider=custom_llm_provider, + logging_obj=logging_obj, + ) def _add_agentic_loop_metadata(kwargs_for_followup: dict[str, object]) -> None: @@ -170,8 +187,13 @@ async def _execute_chat_completion_agentic_plan( model, str(e), ) - if kwargs.get("_code_interpreter_interception_converted_stream") and not depth: - return _wrap_response_as_fake_stream(response_followup) + if _converted_stream_requested(kwargs) and not depth: + return _wrap_response_as_fake_stream( + response_followup, + model=model, + custom_llm_provider=custom_llm_provider, + logging_obj=logging_obj, + ) return response_followup finally: try: @@ -295,9 +317,14 @@ async def maybe_run_chat_completion_agentic_loop( str(e), ) - if kwargs.get("_code_interpreter_interception_converted_stream") and not depth and hasattr(response, "choices"): + if _converted_stream_requested(kwargs) and not depth: return cast( "ModelResponse | CustomStreamWrapper", - _wrap_response_as_fake_stream(response), + _wrap_response_as_fake_stream( + response, + model=model, + custom_llm_provider=custom_llm_provider, + logging_obj=logging_obj, + ), ) return None diff --git a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py index 2735acd7787..1189d9841e4 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py +++ b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py @@ -4,6 +4,7 @@ import json import re import time import uuid +from collections.abc import Mapping from typing import TYPE_CHECKING, Any, ClassVar, List, Literal, Optional import httpx @@ -35,8 +36,12 @@ from litellm.proxy.guardrails.guardrail_hooks.content_text import ( ) from litellm.secret_managers.main import get_secret_str from litellm.types.guardrails import GuardrailEventHooks, Mode -from litellm.types.integrations.custom_logger import AgenticLoopPlan, AgenticLoopRequestPatch -from litellm.types.utils import GenericGuardrailAPIInputs +from litellm.types.integrations.custom_logger import ( + HEADROOM_CONVERTED_STREAM_KEY, + AgenticLoopPlan, + AgenticLoopRequestPatch, +) +from litellm.types.utils import CallTypes, GenericGuardrailAPIInputs if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -606,6 +611,19 @@ class HeadroomGuardrail(CustomGuardrail): return {**inputs, "structured_messages": compressed, "tools": merged_tools} # pyright: ignore[reportReturnType] + async def async_pre_call_deployment_hook( + self, + kwargs: Mapping[str, Any], + call_type: CallTypes | None, + ) -> dict[str, Any] | None: # mutable-ok: overrides CustomLogger hook whose contract is a plain dict + if call_type not in (CallTypes.completion, CallTypes.acompletion): + return None + if not kwargs.get("stream"): + return None + if not has_headroom_retrieve_tool(kwargs.get("tools")): + return None + return {**kwargs, "stream": False, HEADROOM_CONVERTED_STREAM_KEY: True} + async def async_should_run_agentic_loop( self, response: Any, diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index d94fed0ee5b..bd451cb4a18 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -180,6 +180,7 @@ _UNTRUSTED_ROOT_CONTROL_FIELDS = ( "_code_interpreter_interception_converted_stream", "_code_interpreter_interception_sandbox_key", "_code_interpreter_interception_session_scoped", + "_headroom_interception_converted_stream", "max_agentic_loops", ) diff --git a/litellm/types/integrations/custom_logger.py b/litellm/types/integrations/custom_logger.py index 04e490f79ee..8dddeb40e09 100644 --- a/litellm/types/integrations/custom_logger.py +++ b/litellm/types/integrations/custom_logger.py @@ -5,8 +5,14 @@ from pydantic import BaseModel, Field CHAT_COMPLETION_AGENTIC_SURFACE = "chat_completions" RESPONSES_AGENTIC_SURFACE = "responses" CODE_INTERPRETER_INTERCEPTION_PREFIX = "_code_interpreter_interception" +HEADROOM_INTERCEPTION_PREFIX = "_headroom_interception" +HEADROOM_CONVERTED_STREAM_KEY = f"{HEADROOM_INTERCEPTION_PREFIX}_converted_stream" NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES = frozenset( - ("_websearch_interception", "_compression_interception") + ( + "_websearch_interception", + "_compression_interception", + HEADROOM_INTERCEPTION_PREFIX, + ) ) INTERCEPTION_INTERNAL_PREFIXES = frozenset( ( diff --git a/litellm/types/utils.py b/litellm/types/utils.py index e4dfac48141..5edd2806f7d 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3192,6 +3192,7 @@ agentic_loop_internal_litellm_params = [ "_code_interpreter_interception_sandbox_key", "_code_interpreter_interception_session_scoped", "_code_interpreter_interception_converted_stream", + "_headroom_interception_converted_stream", ] all_litellm_params = ( diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py index 248893ed153..9f556a103da 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py @@ -17,10 +17,13 @@ Tests cover: - CCR: headroom_retrieve tool injected when compressed messages contain hashes - CCR: async_should_run_agentic_loop returns True when response has headroom_retrieve tool calls - CCR: async_build_agentic_loop_plan calls retrieve endpoint and builds follow-up messages +- CCR: streaming /chat/completions is converted to a non-streaming call so the agentic + loop resolves the retrieve tool call, then fake-streamed back to the client """ import json import time +from typing import Optional from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -38,7 +41,16 @@ from litellm.proxy.guardrails.guardrail_hooks.headroom.headroom import ( from litellm.proxy.spend_tracking.compression_savings import ( extract_compression_saved_tokens, ) -from litellm.types.utils import GenericGuardrailAPIInputs +from litellm.types.integrations.custom_logger import HEADROOM_CONVERTED_STREAM_KEY +from litellm.types.utils import ( + CallTypes, + ChatCompletionMessageToolCall, + Choices, + Function, + GenericGuardrailAPIInputs, + Message, + ModelResponse, +) FAKE_API_BASE = "https://headroom.example.com" FAKE_API_KEY = "test-key" @@ -1782,3 +1794,125 @@ async def test_fail_open_returns_original_parts_shapes(): messages = result["structured_messages"] assert [m["content"] for m in messages] == [m["content"] for m in PARTS_MESSAGES] + + +CCR_HASH = "b573993006976af767214fac" + + +def _retrieve_tool_definition() -> dict: + return { + "type": "function", + "function": { + "name": HEADROOM_RETRIEVE_TOOL_NAME, + "description": "retrieve compressed content", + "parameters": {"type": "object", "properties": {"hash": {"type": "string"}}}, + }, + } + + +def _model_response_with_retrieve_call() -> ModelResponse: + return ModelResponse( + choices=[ + Choices( + finish_reason="tool_calls", + message=Message( + role="assistant", + content=None, + tool_calls=[ + ChatCompletionMessageToolCall( + id="call_ccr", + type="function", + function=Function( + name=HEADROOM_RETRIEVE_TOOL_NAME, + arguments=json.dumps({"hash": CCR_HASH}), + ), + ) + ], + ), + ) + ] + ) + + +@pytest.mark.parametrize( + "call_type, stream, tools, expect_conversion", + [ + (CallTypes.acompletion, True, [_retrieve_tool_definition()], True), + (CallTypes.completion, True, [_retrieve_tool_definition()], True), + (CallTypes.acompletion, False, [_retrieve_tool_definition()], False), + (CallTypes.acompletion, True, [{"type": "function", "function": {"name": "get_weather"}}], False), + (CallTypes.acompletion, True, None, False), + (CallTypes.aresponses, True, [_retrieve_tool_definition()], False), + (CallTypes.anthropic_messages, True, [_retrieve_tool_definition()], False), + ], +) +@pytest.mark.asyncio +async def test_pre_call_deployment_hook_converts_stream_only_for_ccr_chat_completions( + guardrail: HeadroomGuardrail, + call_type: CallTypes, + stream: bool, + tools: Optional[list], + expect_conversion: bool, +): + kwargs = {"model": "gpt-4o", "stream": stream, "tools": tools} + + result = await guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=call_type) + + if not expect_conversion: + assert result is None + assert kwargs["stream"] is stream + return + + assert result is not None + assert result["stream"] is False + assert result[HEADROOM_CONVERTED_STREAM_KEY] is True + assert kwargs["stream"] is True + + +@pytest.mark.asyncio +async def test_streaming_chat_completion_resolves_ccr_retrieval_end_to_end( + guardrail: HeadroomGuardrail, +): + """Regression test for streaming /chat/completions: the retrieve tool call the + model emits must be resolved by the agentic loop instead of being streamed back + to a client that never declared the tool.""" + original_content = "the full uncompressed document" + final_answer = "the document says hello" + guardrail._issued_hashes_by_call_id["ccr-call-id"] = ( + frozenset({CCR_HASH}), + time.monotonic() + 999, + ) + + real_acompletion = litellm.acompletion + + async def acompletion_with_followup_answer(*args, **kwargs): + if kwargs.get("_agentic_loop_depth"): + kwargs["mock_response"] = final_answer + return await real_acompletion(*args, **kwargs) + + saved_callbacks = list(litellm.callbacks) + litellm.callbacks = [guardrail] + try: + with patch.object( + guardrail.async_handler, + "get", + new_callable=AsyncMock, + return_value=_make_retrieve_response(original_content), + ) as mock_get, patch.object(litellm, "acompletion", new=acompletion_with_followup_answer): + response = await litellm.acompletion( + model="openai/gpt-4o", + messages=[{"role": "user", "content": f"summarize hash={CCR_HASH}"}], + tools=[_retrieve_tool_definition()], + stream=True, + litellm_call_id="ccr-call-id", + mock_response=_model_response_with_retrieve_call(), + ) + chunks = [chunk async for chunk in response] + finally: + litellm.callbacks = saved_callbacks + + streamed_text = "".join(chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices) + assert streamed_text == final_answer + assert not any(chunk.choices and chunk.choices[0].delta.tool_calls for chunk in chunks) + mock_get.assert_called_once() + assert CCR_HASH in (mock_get.call_args.kwargs.get("url") or mock_get.call_args.args[0]) From 8a5135110a0397bc96bf7d2b462fd1783f2bed45 Mon Sep 17 00:00:00 2001 From: Dan Lemon Date: Fri, 24 Jul 2026 00:57:26 +0200 Subject: [PATCH 02/68] fix(soniox): align synthesized SRT/VTT cues to real speech timing The previous cue grouping broke every 15 subword tokens or 5s, which produced uniform ~3s cues that split mid-word and bridged silence gaps, so subtitles did not track the actual speech. Cues are now built from whole words and break on sentence-final punctuation, speaker changes, silence gaps >= 700ms, a 84-char budget, or a 7s duration cap, with timestamps taken directly from token timings. Untimestamped translation tokens are excluded from cues so translated text is never mixed into original-language subtitles. --- litellm/llms/soniox/common_utils.py | 162 +++++++++++------- ...niox_audio_transcription_transformation.py | 98 +++++++++++ 2 files changed, 194 insertions(+), 66 deletions(-) diff --git a/litellm/llms/soniox/common_utils.py b/litellm/llms/soniox/common_utils.py index 90be94b8133..b736d2a5d72 100644 --- a/litellm/llms/soniox/common_utils.py +++ b/litellm/llms/soniox/common_utils.py @@ -113,11 +113,13 @@ def render_soniox_tokens(tokens: list[dict[str, Any]]) -> str: # SRT / VTT subtitle rendering # --------------------------------------------------------------------------- -# Maximum number of tokens to group into a single subtitle cue. -_CUE_MAX_TOKENS: Final[int] = 15 +_CUE_MAX_CHARS: Final[int] = 84 -# Maximum duration (in ms) for a single cue before forcing a break. -_CUE_MAX_DURATION_MS: Final[int] = 5000 +_CUE_MAX_DURATION_MS: Final[int] = 7000 + +_CUE_GAP_MS: Final[int] = 700 + +_SENTENCE_END_CHARS = (".", "!", "?", "。", "！", "？") def _format_timestamp_srt(ms: int) -> str: @@ -144,83 +146,111 @@ def _format_timestamp_vtt(ms: int) -> str: return f"{hours:02d}:{minutes:02d}:{seconds:02d}.{millis:03d}" +def _merge_tokens_into_words(tokens: list[dict[str, Any]]) -> list[dict[str, Any]]: + """ + Merge Soniox subword tokens (e.g. ``"Hel"``, ``"lo"``) into whole words. + + A token starts a new word when its text begins with whitespace, when the + previous token's text ends with whitespace, or when the speaker changes. + Each word carries the first/last available timestamps of its tokens. + + Translation tokens (``translation_status == "translation"``) are excluded: + Soniox does not timestamp them, so they cannot be aligned to the audio and + would otherwise mix translated text into original-language cues. + """ + words: list[dict[str, Any]] = [] + for token in tokens: + text = token.get("text", "") + if not isinstance(text, str) or text == "": + continue + if token.get("translation_status") == "translation": + continue + is_continuation = ( + bool(words) + and not text[0].isspace() + and not words[-1]["text"][-1:].isspace() + and token.get("speaker") == words[-1]["speaker"] + ) + if is_continuation: + last = words[-1] + last["text"] += text + if last["start_ms"] is None: + last["start_ms"] = token.get("start_ms") + if token.get("end_ms") is not None: + last["end_ms"] = token.get("end_ms") + else: + words.append( + { + "text": text, + "start_ms": token.get("start_ms"), + "end_ms": token.get("end_ms"), + "speaker": token.get("speaker"), + } + ) + return words + + def _group_tokens_into_cues( tokens: list[dict[str, Any]], ) -> list[dict[str, Any]]: """ - Group Soniox tokens into subtitle cues. + Group Soniox tokens into subtitle cues aligned to the actual speech. Each cue has: - start_ms: int - end_ms: int - text: str - Grouping heuristics: - - A new cue starts when token count exceeds _CUE_MAX_TOKENS. - - A new cue starts when duration exceeds _CUE_MAX_DURATION_MS. - - A new cue starts when the speaker changes (if diarization is on). - - Tokens without timestamps are appended to the current cue. + Cues only ever break at word boundaries (Soniox tokens are subwords, so + tokens are first merged into words). A new cue starts when: + - the speaker changes (if diarization is on), + - a silence gap of at least _CUE_GAP_MS separates two words, so + subtitles never bridge pauses in speech, + - adding the next word would exceed _CUE_MAX_CHARS (~two subtitle + lines), or + - the cue would span more than _CUE_MAX_DURATION_MS. + A cue also ends after sentence-final punctuation, which keeps cue breaks + at natural seams. Cue timestamps come straight from token timestamps; + words without timestamps stay attached to the surrounding cue. """ - cues: Final[list[dict[str, Any]]] = [] - current_tokens: list[str] = [] - current_start: int | None = None - current_end: int | None = None - current_speaker: Any | None = None + words = _merge_tokens_into_words(tokens) + cues: list[dict[str, Any]] = [] + current: list[dict[str, Any]] = [] + + def _cue_start(ws: list[dict[str, Any]]) -> int | None: + return next((w["start_ms"] for w in ws if w["start_ms"] is not None), None) + + def _cue_end(ws: list[dict[str, Any]]) -> int | None: + return next((w["end_ms"] for w in reversed(ws) if w["end_ms"] is not None), _cue_start(ws)) + + def _cue_text(ws: list[dict[str, Any]]) -> str: + return "".join(w["text"] for w in ws).strip() def _flush() -> None: - if current_tokens and current_start is not None: - text: Final = "".join(current_tokens).strip() - if text: - cues.append( - { - "start_ms": current_start, - "end_ms": (current_end if current_end is not None else current_start), - "text": text, - } - ) + text = _cue_text(current) + start = _cue_start(current) + if text and start is not None: + cues.append({"start_ms": start, "end_ms": _cue_end(current), "text": text}) + current.clear() - for token in tokens: - start_ms = token.get("start_ms") - end_ms = token.get("end_ms") - text = token.get("text", "") - speaker = token.get("speaker") - - # Skip tokens with no timestamp data entirely if we have no cue started - if start_ms is None and current_start is None: - continue - - # Speaker change forces a new cue - if speaker is not None and speaker != current_speaker: + for word in words: + if current: + start_ms = word["start_ms"] + cue_start = _cue_start(current) + cue_end = _cue_end(current) + speaker_changed = word["speaker"] is not None and any( + w["speaker"] is not None and w["speaker"] != word["speaker"] for w in current + ) + gap_exceeded = start_ms is not None and cue_end is not None and (start_ms - cue_end) >= _CUE_GAP_MS + chars_exceeded = len(_cue_text(current)) + len(word["text"]) > _CUE_MAX_CHARS + duration_exceeded = ( + start_ms is not None and cue_start is not None and (start_ms - cue_start) >= _CUE_MAX_DURATION_MS + ) + if speaker_changed or gap_exceeded or chars_exceeded or duration_exceeded: + _flush() + current.append(word) + if word["text"].rstrip().endswith(_SENTENCE_END_CHARS): _flush() - current_tokens = [] - current_start = start_ms - current_end = end_ms - current_speaker = speaker - current_tokens.append(text) - continue - - # Duration or token count exceeded -> flush - should_break = False - if ( - len(current_tokens) >= _CUE_MAX_TOKENS - or current_start is not None - and start_ms is not None - and (start_ms - current_start) >= _CUE_MAX_DURATION_MS - ): - should_break = True - - if should_break: - _flush() - current_tokens = [] - current_start = start_ms - current_end = end_ms - current_tokens.append(text) - else: - if current_start is None: - current_start = start_ms - if end_ms is not None: - current_end = end_ms - current_tokens.append(text) _flush() return cues diff --git a/tests/test_litellm/llms/soniox/audio_transcription/test_soniox_audio_transcription_transformation.py b/tests/test_litellm/llms/soniox/audio_transcription/test_soniox_audio_transcription_transformation.py index 7ee816d5d9e..45626945bd8 100644 --- a/tests/test_litellm/llms/soniox/audio_transcription/test_soniox_audio_transcription_transformation.py +++ b/tests/test_litellm/llms/soniox/audio_transcription/test_soniox_audio_transcription_transformation.py @@ -369,6 +369,104 @@ class TestRenderSonioxTokensAsSrt: assert "01:01:01,000" in result +def _subword_tokens(words, start_ms=0, subword_ms=150, inter_word_gap_ms=50): + tokens = [] + t = start_ms + for word in words: + halves = [word[: len(word) // 2], word[len(word) // 2 :]] if len(word) > 3 else [word] + for i, piece in enumerate(halves): + text = (" " + piece) if i == 0 else piece + tokens.append({"text": text, "start_ms": t, "end_ms": t + subword_ms}) + t += subword_ms + t += inter_word_gap_ms + return tokens, t + + +class TestCueGroupingAlignment: + def test_should_split_cue_on_silence_gap_with_exact_timestamps(self): + from litellm.llms.soniox.common_utils import render_soniox_tokens_as_srt + + before, t = _subword_tokens(["hello", "there"]) + after, _ = _subword_tokens(["welcome", "back"], start_ms=t + 5000) + result = render_soniox_tokens_as_srt(before + after) + cues = result.strip().split("\n\n") + assert len(cues) == 2 + assert "00:00:00,000 --> 00:00:00,650" in cues[0] + assert "hello there" in cues[0] + assert "00:00:05,700 --> 00:00:06,350" in cues[1] + assert "welcome back" in cues[1] + + def test_should_not_bridge_pause_shorter_than_old_duration_cap(self): + from litellm.llms.soniox.common_utils import render_soniox_tokens_as_srt + + before, t = _subword_tokens(["first", "part"]) + after, _ = _subword_tokens(["second", "part"], start_ms=t + 3000) + result = render_soniox_tokens_as_srt(before + after) + cues = result.strip().split("\n\n") + assert len(cues) == 2 + assert "first part" in cues[0] + assert "second part" in cues[1] + + def test_should_never_split_mid_word(self): + from litellm.llms.soniox.common_utils import render_soniox_tokens_as_srt + + tokens, _ = _subword_tokens(["hello"] * 20) + result = render_soniox_tokens_as_srt(tokens) + text_lines = [ + line for line in result.split("\n") if line and "-->" not in line and not line.isdigit() + ] + assert len(text_lines) >= 2 + for line in text_lines: + assert set(line.split()) == {"hello"} + + def test_should_split_after_sentence_final_punctuation(self): + from litellm.llms.soniox.common_utils import render_soniox_tokens_as_srt + + tokens, _ = _subword_tokens(["That", "is", "done.", "Next", "topic"]) + result = render_soniox_tokens_as_srt(tokens) + cues = result.strip().split("\n\n") + assert len(cues) == 2 + assert cues[0].endswith("That is done.") + assert cues[1].endswith("Next topic") + + def test_should_split_on_char_budget_at_word_boundary(self): + from litellm.llms.soniox.common_utils import render_soniox_tokens_as_srt + + tokens, _ = _subword_tokens(["wonderful"] * 12) + result = render_soniox_tokens_as_srt(tokens) + text_lines = [ + line for line in result.split("\n") if line and "-->" not in line and not line.isdigit() + ] + assert len(text_lines) >= 2 + for line in text_lines: + assert len(line) <= 84 + assert set(line.split()) == {"wonderful"} + + def test_should_exclude_untimestamped_translation_tokens_from_cues(self): + from litellm.llms.soniox.common_utils import render_soniox_tokens_as_srt + + tokens = [ + {"text": " Good", "start_ms": 0, "end_ms": 200, "translation_status": "original", "language": "en"}, + {"text": " Guten", "translation_status": "translation", "language": "de", "source_language": "en"}, + {"text": " morning.", "start_ms": 250, "end_ms": 600, "translation_status": "original", "language": "en"}, + ] + result = render_soniox_tokens_as_srt(tokens) + assert "Good morning." in result + assert "Guten" not in result + assert "00:00:00,000 --> 00:00:00,600" in result + + def test_should_keep_untimestamped_word_in_cue(self): + from litellm.llms.soniox.common_utils import render_soniox_tokens_as_srt + + tokens = [ + {"text": " uh", "start_ms": None, "end_ms": None}, + {"text": " hello", "start_ms": 100, "end_ms": 500}, + ] + result = render_soniox_tokens_as_srt(tokens) + assert "uh hello" in result + assert "00:00:00,100 --> 00:00:00,500" in result + + class TestRenderSonioxTokensAsVtt: def test_should_render_basic_vtt_with_header(self): from litellm.llms.soniox.common_utils import render_soniox_tokens_as_vtt From 78fbc57443e59f92b83eaaed01c334a8442dc265 Mon Sep 17 00:00:00 2001 From: Dan Lemon Date: Fri, 24 Jul 2026 14:37:37 +0200 Subject: [PATCH 03/68] fix(soniox): make synthesized subtitle cues work across scripts and languages Cue grouping previously assumed space-separated Latin-style text. Chinese and Japanese audio fused entire utterances into one unbreakable word (and therefore one giant cue) because CJK scripts carry no spaces, and Arabic, Urdu, Hindi and Armenian sentence terminators never triggered a cue break. Words now also split at CJK character boundaries with basic kinsoku handling so punctuation stays attached, the sentence-end set covers script-specific terminators, and the cue length budget counts East Asian wide characters as double width so CJK cues match the same two-line subtitle footprint as Latin text. --- litellm/llms/soniox/common_utils.py | 43 ++++++++-- ...niox_audio_transcription_transformation.py | 85 +++++++++++++++++++ 2 files changed, 123 insertions(+), 5 deletions(-) diff --git a/litellm/llms/soniox/common_utils.py b/litellm/llms/soniox/common_utils.py index b736d2a5d72..2fd8ecf54de 100644 --- a/litellm/llms/soniox/common_utils.py +++ b/litellm/llms/soniox/common_utils.py @@ -2,6 +2,7 @@ Shared utilities for the Soniox provider (https://soniox.com). """ +import unicodedata from typing import Any, Final from litellm.llms.base_llm.chat.transformation import BaseLLMException @@ -119,7 +120,35 @@ _CUE_MAX_DURATION_MS: Final[int] = 7000 _CUE_GAP_MS: Final[int] = 700 -_SENTENCE_END_CHARS = (".", "!", "?", "。", "！", "？") +_SENTENCE_END_CHARS = (".", "!", "?", "。", "！", "？", "؟", "۔", "।", "॥", "։", "።") + +_CJK_RANGES = ( + (0x3400, 0x4DBF), + (0x4E00, 0x9FFF), + (0xF900, 0xFAFF), + (0x3040, 0x309F), + (0x30A0, 0x30FF), + (0x31F0, 0x31FF), +) + +_CJK_NO_BREAK_BEFORE = "、。，．！？：；・ー…」』）〉》】〕" + +_CJK_NO_BREAK_AFTER = "「『（〈《【〔" + + +def _is_cjk(ch: str) -> bool: + cp = ord(ch) + return any(lo <= cp <= hi for lo, hi in _CJK_RANGES) + + +def _is_cjk_word_boundary(prev_ch: str, next_ch: str) -> bool: + if not (_is_cjk(prev_ch) or _is_cjk(next_ch)): + return False + return next_ch not in _CJK_NO_BREAK_BEFORE and prev_ch not in _CJK_NO_BREAK_AFTER + + +def _text_width(text: str) -> int: + return sum(2 if unicodedata.east_asian_width(ch) in ("W", "F") else 1 for ch in text) def _format_timestamp_srt(ms: int) -> str: @@ -151,7 +180,10 @@ def _merge_tokens_into_words(tokens: list[dict[str, Any]]) -> list[dict[str, Any Merge Soniox subword tokens (e.g. ``"Hel"``, ``"lo"``) into whole words. A token starts a new word when its text begins with whitespace, when the - previous token's text ends with whitespace, or when the speaker changes. + previous token's text ends with whitespace, when the speaker changes, or + at a CJK character boundary (CJK scripts carry no spaces, so without this + an entire utterance would fuse into a single unbreakable "word"; CJK + punctuation stays attached to the preceding character per kinsoku rules). Each word carries the first/last available timestamps of its tokens. Translation tokens (``translation_status == "translation"``) are excluded: @@ -170,6 +202,7 @@ def _merge_tokens_into_words(tokens: list[dict[str, Any]]) -> list[dict[str, Any and not text[0].isspace() and not words[-1]["text"][-1:].isspace() and token.get("speaker") == words[-1]["speaker"] + and not _is_cjk_word_boundary(words[-1]["text"][-1:], text[0]) ) if is_continuation: last = words[-1] @@ -206,8 +239,8 @@ def _group_tokens_into_cues( - the speaker changes (if diarization is on), - a silence gap of at least _CUE_GAP_MS separates two words, so subtitles never bridge pauses in speech, - - adding the next word would exceed _CUE_MAX_CHARS (~two subtitle - lines), or + - adding the next word would exceed _CUE_MAX_CHARS of display width + (~two subtitle lines; East-Asian wide characters count double), or - the cue would span more than _CUE_MAX_DURATION_MS. A cue also ends after sentence-final punctuation, which keeps cue breaks at natural seams. Cue timestamps come straight from token timestamps; @@ -242,7 +275,7 @@ def _group_tokens_into_cues( w["speaker"] is not None and w["speaker"] != word["speaker"] for w in current ) gap_exceeded = start_ms is not None and cue_end is not None and (start_ms - cue_end) >= _CUE_GAP_MS - chars_exceeded = len(_cue_text(current)) + len(word["text"]) > _CUE_MAX_CHARS + chars_exceeded = _text_width(_cue_text(current)) + _text_width(word["text"]) > _CUE_MAX_CHARS duration_exceeded = ( start_ms is not None and cue_start is not None and (start_ms - cue_start) >= _CUE_MAX_DURATION_MS ) diff --git a/tests/test_litellm/llms/soniox/audio_transcription/test_soniox_audio_transcription_transformation.py b/tests/test_litellm/llms/soniox/audio_transcription/test_soniox_audio_transcription_transformation.py index 45626945bd8..2917d5663ac 100644 --- a/tests/test_litellm/llms/soniox/audio_transcription/test_soniox_audio_transcription_transformation.py +++ b/tests/test_litellm/llms/soniox/audio_transcription/test_soniox_audio_transcription_transformation.py @@ -467,6 +467,91 @@ class TestCueGroupingAlignment: assert "00:00:00,100 --> 00:00:00,500" in result +def _cue_texts(srt: str) -> list: + return [cue.split("\n", 2)[2] for cue in srt.strip().split("\n\n")] + + +class TestMultilingualCueGrouping: + def test_should_split_spaceless_chinese_on_width_budget(self): + from litellm.llms.soniox.common_utils import _text_width, render_soniox_tokens_as_srt + + tokens = [{"text": "你好", "start_ms": i * 100, "end_ms": i * 100 + 90} for i in range(60)] + result = render_soniox_tokens_as_srt(tokens) + texts = _cue_texts(result) + assert len(texts) >= 3 + for text in texts: + assert _text_width(text) <= 84 + assert set(text) <= {"你", "好"} + + def test_should_split_japanese_after_sentence_end_and_keep_punctuation_attached(self): + from litellm.llms.soniox.common_utils import render_soniox_tokens_as_srt + + tokens = [ + {"text": "今日は", "start_ms": 0, "end_ms": 300}, + {"text": "いい", "start_ms": 300, "end_ms": 500}, + {"text": "天気です", "start_ms": 500, "end_ms": 900}, + {"text": "。", "start_ms": 900, "end_ms": 950}, + {"text": "明日も", "start_ms": 1000, "end_ms": 1300}, + {"text": "晴れ", "start_ms": 1300, "end_ms": 1500}, + ] + texts = _cue_texts(render_soniox_tokens_as_srt(tokens)) + assert texts == ["今日はいい天気です。", "明日も晴れ"] + + def test_should_split_arabic_after_arabic_question_mark(self): + from litellm.llms.soniox.common_utils import render_soniox_tokens_as_srt + + tokens = [ + {"text": " كيف", "start_ms": 0, "end_ms": 300}, + {"text": " حالك؟", "start_ms": 300, "end_ms": 700}, + {"text": " أنا", "start_ms": 800, "end_ms": 1000}, + {"text": " بخير", "start_ms": 1000, "end_ms": 1300}, + ] + texts = _cue_texts(render_soniox_tokens_as_srt(tokens)) + assert texts == ["كيف حالك؟", "أنا بخير"] + + def test_should_split_after_devanagari_and_urdu_terminators(self): + from litellm.llms.soniox.common_utils import render_soniox_tokens_as_srt + + tokens = [ + {"text": " नमस्ते।", "start_ms": 0, "end_ms": 400}, + {"text": " آپ", "start_ms": 500, "end_ms": 700}, + {"text": " ٹھیک۔", "start_ms": 700, "end_ms": 1100}, + {"text": " शुभ", "start_ms": 1200, "end_ms": 1400}, + ] + texts = _cue_texts(render_soniox_tokens_as_srt(tokens)) + assert texts == ["नमस्ते।", "آپ ٹھیک۔", "शुभ"] + + def test_should_split_russian_after_sentence_end(self): + from litellm.llms.soniox.common_utils import render_soniox_tokens_as_srt + + tokens = [ + {"text": " Как", "start_ms": 0, "end_ms": 200}, + {"text": " дела?", "start_ms": 200, "end_ms": 600}, + {"text": " Хорошо.", "start_ms": 700, "end_ms": 1200}, + ] + texts = _cue_texts(render_soniox_tokens_as_srt(tokens)) + assert texts == ["Как дела?", "Хорошо."] + + def test_should_not_split_latin_text_within_width_budget(self): + from litellm.llms.soniox.common_utils import render_soniox_tokens_as_srt + + tokens = [{"text": f" word{i}", "start_ms": i * 100, "end_ms": i * 100 + 90} for i in range(12)] + texts = _cue_texts(render_soniox_tokens_as_srt(tokens)) + assert len(texts) == 1 + + def test_should_merge_subword_tokens_inside_cjk_run(self): + from litellm.llms.soniox.common_utils import _merge_tokens_into_words + + tokens = [ + {"text": "編", "start_ms": 0, "end_ms": 100}, + {"text": "集", "start_ms": 100, "end_ms": 200}, + {"text": "、", "start_ms": 200, "end_ms": 250}, + {"text": "保存", "start_ms": 250, "end_ms": 400}, + ] + words = _merge_tokens_into_words(tokens) + assert [w["text"] for w in words] == ["編", "集、", "保存"] + + class TestRenderSonioxTokensAsVtt: def test_should_render_basic_vtt_with_header(self): from litellm.llms.soniox.common_utils import render_soniox_tokens_as_vtt From 93bfe176b2da58596d54a1a82785306e6fc31e91 Mon Sep 17 00:00:00 2001 From: Dan Lemon Date: Fri, 24 Jul 2026 15:18:33 +0200 Subject: [PATCH 04/68] fix(soniox): enforce cue duration cap using word end timestamp --- litellm/llms/soniox/common_utils.py | 6 ++++-- ...est_soniox_audio_transcription_transformation.py | 13 +++++++++++++ 2 files changed, 17 insertions(+), 2 deletions(-) diff --git a/litellm/llms/soniox/common_utils.py b/litellm/llms/soniox/common_utils.py index 2fd8ecf54de..ee25465e128 100644 --- a/litellm/llms/soniox/common_utils.py +++ b/litellm/llms/soniox/common_utils.py @@ -241,7 +241,8 @@ def _group_tokens_into_cues( subtitles never bridge pauses in speech, - adding the next word would exceed _CUE_MAX_CHARS of display width (~two subtitle lines; East-Asian wide characters count double), or - - the cue would span more than _CUE_MAX_DURATION_MS. + - adding the next word would make the cue span more than + _CUE_MAX_DURATION_MS. A cue also ends after sentence-final punctuation, which keeps cue breaks at natural seams. Cue timestamps come straight from token timestamps; words without timestamps stay attached to the surrounding cue. @@ -276,8 +277,9 @@ def _group_tokens_into_cues( ) gap_exceeded = start_ms is not None and cue_end is not None and (start_ms - cue_end) >= _CUE_GAP_MS chars_exceeded = _text_width(_cue_text(current)) + _text_width(word["text"]) > _CUE_MAX_CHARS + word_end = word["end_ms"] if word["end_ms"] is not None else start_ms duration_exceeded = ( - start_ms is not None and cue_start is not None and (start_ms - cue_start) >= _CUE_MAX_DURATION_MS + word_end is not None and cue_start is not None and (word_end - cue_start) > _CUE_MAX_DURATION_MS ) if speaker_changed or gap_exceeded or chars_exceeded or duration_exceeded: _flush() diff --git a/tests/test_litellm/llms/soniox/audio_transcription/test_soniox_audio_transcription_transformation.py b/tests/test_litellm/llms/soniox/audio_transcription/test_soniox_audio_transcription_transformation.py index 2917d5663ac..cf7b5679dfe 100644 --- a/tests/test_litellm/llms/soniox/audio_transcription/test_soniox_audio_transcription_transformation.py +++ b/tests/test_litellm/llms/soniox/audio_transcription/test_soniox_audio_transcription_transformation.py @@ -455,6 +455,19 @@ class TestCueGroupingAlignment: assert "Guten" not in result assert "00:00:00,000 --> 00:00:00,600" in result + def test_should_split_before_word_whose_end_crosses_duration_cap(self): + from litellm.llms.soniox.common_utils import render_soniox_tokens_as_srt + + tokens = [{"text": " hm", "start_ms": i * 650, "end_ms": i * 650 + 600} for i in range(10)] + [ + {"text": " boom", "start_ms": 6900, "end_ms": 7600} + ] + result = render_soniox_tokens_as_srt(tokens) + cues = result.strip().split("\n\n") + assert len(cues) == 2 + assert "00:00:00,000 --> 00:00:06,450" in cues[0] + assert "00:00:06,900 --> 00:00:07,600" in cues[1] + assert cues[1].endswith("boom") + def test_should_keep_untimestamped_word_in_cue(self): from litellm.llms.soniox.common_utils import render_soniox_tokens_as_srt From d45fe41fe9457708f2d7654ea2a6f43559322eb1 Mon Sep 17 00:00:00 2001 From: Dan Lemon Date: Wed, 12 Aug 2026 09:22:01 +0200 Subject: [PATCH 05/68] refactor(soniox): make cue grouping functional to satisfy type discipline gate --- litellm/llms/soniox/common_utils.py | 237 ++++++++++-------- ...niox_audio_transcription_transformation.py | 2 +- type-discipline-budget.json | 6 +- 3 files changed, 136 insertions(+), 109 deletions(-) diff --git a/litellm/llms/soniox/common_utils.py b/litellm/llms/soniox/common_utils.py index ee25465e128..fc85e2b719c 100644 --- a/litellm/llms/soniox/common_utils.py +++ b/litellm/llms/soniox/common_utils.py @@ -3,6 +3,9 @@ Shared utilities for the Soniox provider (https://soniox.com). """ import unicodedata +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from functools import reduce from typing import Any, Final from litellm.llms.base_llm.chat.transformation import BaseLLMException @@ -120,9 +123,9 @@ _CUE_MAX_DURATION_MS: Final[int] = 7000 _CUE_GAP_MS: Final[int] = 700 -_SENTENCE_END_CHARS = (".", "!", "?", "。", "！", "？", "؟", "۔", "।", "॥", "։", "።") +_SENTENCE_END_CHARS: Final = (".", "!", "?", "。", "！", "？", "؟", "۔", "।", "॥", "։", "።") -_CJK_RANGES = ( +_CJK_RANGES: Final = ( (0x3400, 0x4DBF), (0x4E00, 0x9FFF), (0xF900, 0xFAFF), @@ -131,13 +134,13 @@ _CJK_RANGES = ( (0x31F0, 0x31FF), ) -_CJK_NO_BREAK_BEFORE = "、。，．！？：；・ー…」』）〉》】〕" +_CJK_NO_BREAK_BEFORE: Final = "、。，．！？：；・ー…」』）〉》】〕" -_CJK_NO_BREAK_AFTER = "「『（〈《【〔" +_CJK_NO_BREAK_AFTER: Final = "「『（〈《【〔" def _is_cjk(ch: str) -> bool: - cp = ord(ch) + cp: Final = ord(ch) return any(lo <= cp <= hi for lo, hi in _CJK_RANGES) @@ -175,7 +178,47 @@ def _format_timestamp_vtt(ms: int) -> str: return f"{hours:02d}:{minutes:02d}:{seconds:02d}.{millis:03d}" -def _merge_tokens_into_words(tokens: list[dict[str, Any]]) -> list[dict[str, Any]]: +@dataclass(frozen=True, slots=True) +class _Word: + text: str + start_ms: int | None + end_ms: int | None + speaker: str | int | None + + +@dataclass(frozen=True, slots=True) +class _Cue: + start_ms: int + end_ms: int + text: str + + +def _keeps_token(token: Mapping[str, Any]) -> bool: + text: Final = token.get("text", "") + return isinstance(text, str) and text != "" and token.get("translation_status") != "translation" + + +def _starts_new_word(prev: Mapping[str, Any], token: Mapping[str, Any]) -> bool: + prev_last: Final = prev["text"][-1:] + first: Final = token["text"][0] + return ( + first.isspace() + or prev_last.isspace() + or token.get("speaker") != prev.get("speaker") + or _is_cjk_word_boundary(prev_last, first) + ) + + +def _build_word(group: Sequence[Mapping[str, Any]]) -> _Word: + return _Word( + text="".join(t["text"] for t in group), + start_ms=next((t.get("start_ms") for t in group if t.get("start_ms") is not None), None), + end_ms=next((t.get("end_ms") for t in reversed(group) if t.get("end_ms") is not None), None), + speaker=group[0].get("speaker"), + ) + + +def _merge_tokens_into_words(tokens: Sequence[Mapping[str, Any]]) -> tuple[_Word, ...]: """ Merge Soniox subword tokens (e.g. ``"Hel"``, ``"lo"``) into whole words. @@ -190,50 +233,62 @@ def _merge_tokens_into_words(tokens: list[dict[str, Any]]) -> list[dict[str, Any Soniox does not timestamp them, so they cannot be aligned to the audio and would otherwise mix translated text into original-language cues. """ - words: list[dict[str, Any]] = [] - for token in tokens: - text = token.get("text", "") - if not isinstance(text, str) or text == "": - continue - if token.get("translation_status") == "translation": - continue - is_continuation = ( - bool(words) - and not text[0].isspace() - and not words[-1]["text"][-1:].isspace() - and token.get("speaker") == words[-1]["speaker"] - and not _is_cjk_word_boundary(words[-1]["text"][-1:], text[0]) - ) - if is_continuation: - last = words[-1] - last["text"] += text - if last["start_ms"] is None: - last["start_ms"] = token.get("start_ms") - if token.get("end_ms") is not None: - last["end_ms"] = token.get("end_ms") - else: - words.append( - { - "text": text, - "start_ms": token.get("start_ms"), - "end_ms": token.get("end_ms"), - "speaker": token.get("speaker"), - } - ) - return words + kept: Final = tuple(t for t in tokens if _keeps_token(t)) + starts: Final = tuple(i for i, t in enumerate(kept) if i == 0 or _starts_new_word(kept[i - 1], t)) + return tuple(_build_word(kept[begin:end]) for begin, end in zip(starts, (*starts[1:], len(kept)))) -def _group_tokens_into_cues( - tokens: list[dict[str, Any]], -) -> list[dict[str, Any]]: +def _cue_start(ws: Sequence[_Word]) -> int | None: + return next((w.start_ms for w in ws if w.start_ms is not None), None) + + +def _cue_end(ws: Sequence[_Word]) -> int | None: + return next((w.end_ms for w in reversed(ws) if w.end_ms is not None), _cue_start(ws)) + + +def _cue_text(ws: Sequence[_Word]) -> str: + return "".join(w.text for w in ws).strip() + + +def _should_break(cue: Sequence[_Word], word: _Word) -> bool: + speaker_changed: Final = word.speaker is not None and any( + w.speaker is not None and w.speaker != word.speaker for w in cue + ) + cue_start: Final = _cue_start(cue) + cue_end: Final = _cue_end(cue) + gap_exceeded: Final = word.start_ms is not None and cue_end is not None and (word.start_ms - cue_end) >= _CUE_GAP_MS + chars_exceeded: Final = _text_width(_cue_text(cue)) + _text_width(word.text) > _CUE_MAX_CHARS + word_end: Final = word.end_ms if word.end_ms is not None else word.start_ms + duration_exceeded: Final = ( + word_end is not None and cue_start is not None and (word_end - cue_start) > _CUE_MAX_DURATION_MS + ) + return speaker_changed or gap_exceeded or chars_exceeded or duration_exceeded + + +def _cue_start_indices(words: Sequence[_Word]) -> tuple[int, ...]: + def step(starts: tuple[int, ...], index: int) -> tuple[int, ...]: + if words[index - 1].text.rstrip().endswith(_SENTENCE_END_CHARS): + return (*starts, index) + if _should_break(words[starts[-1] : index], words[index]): + return (*starts, index) + return starts + + return reduce(step, range(1, len(words)), (0,)) if words else () + + +def _build_cue(ws: Sequence[_Word]) -> _Cue | None: + text: Final = _cue_text(ws) + start: Final = _cue_start(ws) + if not text or start is None: + return None + end: Final = _cue_end(ws) + return _Cue(start_ms=start, end_ms=end if end is not None else start, text=text) + + +def _group_tokens_into_cues(tokens: Sequence[Mapping[str, Any]]) -> tuple[_Cue, ...]: """ Group Soniox tokens into subtitle cues aligned to the actual speech. - Each cue has: - - start_ms: int - - end_ms: int - - text: str - Cues only ever break at word boundaries (Soniox tokens are subwords, so tokens are first merged into words). A new cue starts when: - the speaker changes (if diarization is on), @@ -245,50 +300,16 @@ def _group_tokens_into_cues( _CUE_MAX_DURATION_MS. A cue also ends after sentence-final punctuation, which keeps cue breaks at natural seams. Cue timestamps come straight from token timestamps; - words without timestamps stay attached to the surrounding cue. + words without timestamps stay attached to the surrounding cue, and a cue + whose words carry no timestamps at all is dropped. """ - words = _merge_tokens_into_words(tokens) - cues: list[dict[str, Any]] = [] - current: list[dict[str, Any]] = [] - - def _cue_start(ws: list[dict[str, Any]]) -> int | None: - return next((w["start_ms"] for w in ws if w["start_ms"] is not None), None) - - def _cue_end(ws: list[dict[str, Any]]) -> int | None: - return next((w["end_ms"] for w in reversed(ws) if w["end_ms"] is not None), _cue_start(ws)) - - def _cue_text(ws: list[dict[str, Any]]) -> str: - return "".join(w["text"] for w in ws).strip() - - def _flush() -> None: - text = _cue_text(current) - start = _cue_start(current) - if text and start is not None: - cues.append({"start_ms": start, "end_ms": _cue_end(current), "text": text}) - current.clear() - - for word in words: - if current: - start_ms = word["start_ms"] - cue_start = _cue_start(current) - cue_end = _cue_end(current) - speaker_changed = word["speaker"] is not None and any( - w["speaker"] is not None and w["speaker"] != word["speaker"] for w in current - ) - gap_exceeded = start_ms is not None and cue_end is not None and (start_ms - cue_end) >= _CUE_GAP_MS - chars_exceeded = _text_width(_cue_text(current)) + _text_width(word["text"]) > _CUE_MAX_CHARS - word_end = word["end_ms"] if word["end_ms"] is not None else start_ms - duration_exceeded = ( - word_end is not None and cue_start is not None and (word_end - cue_start) > _CUE_MAX_DURATION_MS - ) - if speaker_changed or gap_exceeded or chars_exceeded or duration_exceeded: - _flush() - current.append(word) - if word["text"].rstrip().endswith(_SENTENCE_END_CHARS): - _flush() - - _flush() - return cues + words: Final = _merge_tokens_into_words(tokens) + starts: Final = _cue_start_indices(words) + return tuple( + cue + for begin, end in zip(starts, (*starts[1:], len(words))) + if (cue := _build_cue(words[begin:end])) is not None + ) def render_soniox_tokens_as_srt(tokens: list[dict[str, Any]]) -> str: @@ -301,16 +322,16 @@ def render_soniox_tokens_as_srt(tokens: list[dict[str, Any]]) -> str: if not cues: return "" - lines: Final[list[str]] = [] - for idx, cue in enumerate(cues, start=1): - start = _format_timestamp_srt(cue["start_ms"]) - end = _format_timestamp_srt(cue["end_ms"]) - lines.append(str(idx)) - lines.append(f"{start} --> {end}") - lines.append(cue["text"]) - lines.append("") # blank line between cues - - return "\n".join(lines) + return "\n".join( + line + for idx, cue in enumerate(cues, start=1) + for line in ( + str(idx), + f"{_format_timestamp_srt(cue.start_ms)} --> {_format_timestamp_srt(cue.end_ms)}", + cue.text, + "", + ) + ) def render_soniox_tokens_as_vtt(tokens: list[dict[str, Any]]) -> str: @@ -321,12 +342,18 @@ def render_soniox_tokens_as_vtt(tokens: list[dict[str, Any]]) -> str: """ cues: Final = _group_tokens_into_cues(tokens) - lines: Final[list[str]] = ["WEBVTT", ""] - for cue in cues: - start = _format_timestamp_vtt(cue["start_ms"]) - end = _format_timestamp_vtt(cue["end_ms"]) - lines.append(f"{start} --> {end}") - lines.append(cue["text"]) - lines.append("") # blank line between cues + lines: Final = ( + "WEBVTT", + "", + *( + line + for cue in cues + for line in ( + f"{_format_timestamp_vtt(cue.start_ms)} --> {_format_timestamp_vtt(cue.end_ms)}", + cue.text, + "", + ) + ), + ) return "\n".join(lines) diff --git a/tests/test_litellm/llms/soniox/audio_transcription/test_soniox_audio_transcription_transformation.py b/tests/test_litellm/llms/soniox/audio_transcription/test_soniox_audio_transcription_transformation.py index cf7b5679dfe..cb053ac5f1e 100644 --- a/tests/test_litellm/llms/soniox/audio_transcription/test_soniox_audio_transcription_transformation.py +++ b/tests/test_litellm/llms/soniox/audio_transcription/test_soniox_audio_transcription_transformation.py @@ -562,7 +562,7 @@ class TestMultilingualCueGrouping: {"text": "保存", "start_ms": 250, "end_ms": 400}, ] words = _merge_tokens_into_words(tokens) - assert [w["text"] for w in words] == ["編", "集、", "保存"] + assert [w.text for w in words] == ["編", "集、", "保存"] class TestRenderSonioxTokensAsVtt: diff --git a/type-discipline-budget.json b/type-discipline-budget.json index fdacf375844..daa3e89fa0a 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 23003 + "limit": 22994 }, "LIT002": { - "limit": 27146 + "limit": 27139 }, "LIT003": { "limit": 269 @@ -27,7 +27,7 @@ "limit": 0 }, "LIT010": { - "limit": 16731 + "limit": 16727 }, "LIT011": { "limit": 5596 From 0a582975de53eb74f60fb5a8dc10ea5cbb4098c7 Mon Sep 17 00:00:00 2001 From: Dan Lemon Date: Wed, 12 Aug 2026 12:10:17 +0200 Subject: [PATCH 06/68] fix(soniox): accumulate cue start indices in linear time --- litellm/llms/soniox/common_utils.py | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/litellm/llms/soniox/common_utils.py b/litellm/llms/soniox/common_utils.py index fc85e2b719c..b4597ef2a50 100644 --- a/litellm/llms/soniox/common_utils.py +++ b/litellm/llms/soniox/common_utils.py @@ -5,7 +5,7 @@ Shared utilities for the Soniox provider (https://soniox.com). import unicodedata from collections.abc import Mapping, Sequence from dataclasses import dataclass -from functools import reduce +from itertools import accumulate, groupby from typing import Any, Final from litellm.llms.base_llm.chat.transformation import BaseLLMException @@ -266,14 +266,16 @@ def _should_break(cue: Sequence[_Word], word: _Word) -> bool: def _cue_start_indices(words: Sequence[_Word]) -> tuple[int, ...]: - def step(starts: tuple[int, ...], index: int) -> tuple[int, ...]: + def next_start(start: int, index: int) -> int: if words[index - 1].text.rstrip().endswith(_SENTENCE_END_CHARS): - return (*starts, index) - if _should_break(words[starts[-1] : index], words[index]): - return (*starts, index) - return starts + return index + if _should_break(words[start:index], words[index]): + return index + return start - return reduce(step, range(1, len(words)), (0,)) if words else () + if not words: + return () + return tuple(start for start, _ in groupby(accumulate(range(1, len(words)), next_start, initial=0))) def _build_cue(ws: Sequence[_Word]) -> _Cue | None: From c72ffe4bc8c6008c7da98b1f65ab72a22f5df5bc Mon Sep 17 00:00:00 2001 From: Dan Lemon Date: Wed, 12 Aug 2026 12:45:48 +0200 Subject: [PATCH 07/68] chore: restore type-discipline-budget.json to base --- type-discipline-budget.json | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/type-discipline-budget.json b/type-discipline-budget.json index daa3e89fa0a..fdacf375844 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 22994 + "limit": 23003 }, "LIT002": { - "limit": 27139 + "limit": 27146 }, "LIT003": { "limit": 269 @@ -27,7 +27,7 @@ "limit": 0 }, "LIT010": { - "limit": 16727 + "limit": 16731 }, "LIT011": { "limit": 5596 From dfc22d31b428c9ba6fb61357cd6cf5438918e3dc Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Thu, 18 Jun 2026 15:24:35 -0500 Subject: [PATCH 08/68] feat(vertex-ai): add veo 3.1 lite model metadata --- ...odel_prices_and_context_window_backup.json | 15 ++++ model_prices_and_context_window.json | 15 ++++ .../test_vertex_video_transformation.py | 81 ++++++++++++++++++- 3 files changed, 109 insertions(+), 2 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index b12e4a9fea3..33580506977 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -39343,6 +39343,21 @@ "video" ] }, + "vertex_ai/veo-3.1-lite-generate-001": { + "litellm_provider": "vertex_ai-video-models", + "max_input_tokens": 1024, + "max_tokens": 1024, + "mode": "video_generation", + "output_cost_per_second": 0.05, + "output_cost_per_second_1080p": 0.08, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing#veo", + "supported_modalities": [ + "text" + ], + "supported_output_modalities": [ + "video" + ] + }, "voyage/rerank-2": { "input_cost_per_token": 5e-08, "litellm_provider": "voyage", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index b12e4a9fea3..33580506977 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -39343,6 +39343,21 @@ "video" ] }, + "vertex_ai/veo-3.1-lite-generate-001": { + "litellm_provider": "vertex_ai-video-models", + "max_input_tokens": 1024, + "max_tokens": 1024, + "mode": "video_generation", + "output_cost_per_second": 0.05, + "output_cost_per_second_1080p": 0.08, + "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing#veo", + "supported_modalities": [ + "text" + ], + "supported_output_modalities": [ + "video" + ] + }, "voyage/rerank-2": { "input_cost_per_token": 5e-08, "litellm_provider": "voyage", diff --git a/tests/test_litellm/llms/vertex_ai/videos/test_vertex_video_transformation.py b/tests/test_litellm/llms/vertex_ai/videos/test_vertex_video_transformation.py index 55197d3165c..a0fe57a6eba 100644 --- a/tests/test_litellm/llms/vertex_ai/videos/test_vertex_video_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/videos/test_vertex_video_transformation.py @@ -4,12 +4,16 @@ Tests for Vertex AI (Veo) video generation transformation. import base64 import json -import os -from unittest.mock import MagicMock, Mock, patch +from pathlib import Path +from typing import Mapping, cast +from unittest.mock import Mock, patch import httpx import pytest +import litellm +from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider +from litellm.llms.openai.cost_calculation import video_generation_cost from litellm.llms.vertex_ai.videos.transformation import ( VertexAIVideoConfig, _convert_image_to_vertex_format, @@ -17,6 +21,21 @@ from litellm.llms.vertex_ai.videos.transformation import ( from litellm.types.router import GenericLiteLLMParams from litellm.types.videos.main import VideoObject +VEO_31_LITE_VERTEX_MODEL = "vertex_ai/veo-3.1-lite-generate-001" +ROOT_MODEL_COST_PATH = ( + Path(__file__).parents[5] / "model_prices_and_context_window.json" +) +BACKUP_MODEL_COST_PATH = ( + Path(__file__).parents[5] + / "litellm" + / "model_prices_and_context_window_backup.json" +) +ModelCostMap = Mapping[str, Mapping[str, object]] + + +def _load_model_cost_map(path: Path) -> ModelCostMap: + return cast(ModelCostMap, json.loads(path.read_text())) + class TestVertexAIVideoConfig: """Test VertexAIVideoConfig transformation class.""" @@ -123,6 +142,64 @@ class TestVertexAIVideoConfig: # Should NOT include endpoint assert not url.endswith(":predictLongRunning") + def test_veo_31_lite_model_cost_entries_match_pricing(self): + for path in (ROOT_MODEL_COST_PATH, BACKUP_MODEL_COST_PATH): + model_cost = _load_model_cost_map(path) + info = model_cost.get(VEO_31_LITE_VERTEX_MODEL) + + assert info is not None, f"{VEO_31_LITE_VERTEX_MODEL} missing from {path}" + assert info["litellm_provider"] == "vertex_ai-video-models" + assert info["mode"] == "video_generation" + assert info["max_input_tokens"] == 1024 + assert info["output_cost_per_second"] == 0.05 + assert info["output_cost_per_second_1080p"] == 0.08 + + def test_veo_31_lite_provider_routing_from_local_model_map(self): + original_model_cost = litellm.model_cost + original_vertex_video_models = set(litellm.vertex_ai_video_models) + + try: + model_cost = _load_model_cost_map(BACKUP_MODEL_COST_PATH) + litellm.model_cost = dict(model_cost) + litellm.vertex_ai_video_models.clear() + litellm.add_known_models(model_cost_map=litellm.model_cost) + + model, custom_llm_provider, _, _ = get_llm_provider( + model="veo-3.1-lite-generate-001" + ) + + assert model == "veo-3.1-lite-generate-001" + assert custom_llm_provider == "vertex_ai" + finally: + litellm.model_cost = original_model_cost + litellm.vertex_ai_video_models.clear() + litellm.vertex_ai_video_models.update(original_vertex_video_models) + + def test_veo_31_lite_cost_uses_resolution_tiers(self): + model_cost = _load_model_cost_map(BACKUP_MODEL_COST_PATH) + model_info = model_cost[VEO_31_LITE_VERTEX_MODEL] + + assert ( + video_generation_cost( + model=VEO_31_LITE_VERTEX_MODEL, + duration_seconds=10.0, + custom_llm_provider="vertex_ai", + model_info=dict(model_info), + video_resolution="720p", + ) + == 0.5 + ) + assert ( + video_generation_cost( + model=VEO_31_LITE_VERTEX_MODEL, + duration_seconds=10.0, + custom_llm_provider="vertex_ai", + model_info=dict(model_info), + video_resolution="1080p", + ) + == 0.8 + ) + def test_transform_video_create_request(self): """Test transformation of video creation request.""" prompt = "A cat playing with a ball of yarn" From cc8e128d53e4c25d2da8c1b035ce435a735777dc Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Thu, 18 Jun 2026 16:26:12 -0500 Subject: [PATCH 09/68] test(vertex-ai): tighten veo lite review fixes --- .../test_vertex_video_transformation.py | 67 ++++++++----------- 1 file changed, 29 insertions(+), 38 deletions(-) diff --git a/tests/test_litellm/llms/vertex_ai/videos/test_vertex_video_transformation.py b/tests/test_litellm/llms/vertex_ai/videos/test_vertex_video_transformation.py index a0fe57a6eba..927c7dd94b2 100644 --- a/tests/test_litellm/llms/vertex_ai/videos/test_vertex_video_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/videos/test_vertex_video_transformation.py @@ -154,51 +154,42 @@ class TestVertexAIVideoConfig: assert info["output_cost_per_second"] == 0.05 assert info["output_cost_per_second_1080p"] == 0.08 - def test_veo_31_lite_provider_routing_from_local_model_map(self): - original_model_cost = litellm.model_cost - original_vertex_video_models = set(litellm.vertex_ai_video_models) + def test_veo_31_lite_provider_routing_from_local_model_map( + self, monkeypatch: pytest.MonkeyPatch + ): + model_cost = _load_model_cost_map(BACKUP_MODEL_COST_PATH) + vertex_video_models = { + model_name.removeprefix("vertex_ai/") + for model_name, info in model_cost.items() + if info.get("litellm_provider") == "vertex_ai-video-models" + } + monkeypatch.setattr(litellm, "vertex_ai_video_models", vertex_video_models) - try: - model_cost = _load_model_cost_map(BACKUP_MODEL_COST_PATH) - litellm.model_cost = dict(model_cost) - litellm.vertex_ai_video_models.clear() - litellm.add_known_models(model_cost_map=litellm.model_cost) + model, custom_llm_provider, _, _ = get_llm_provider( + model="veo-3.1-lite-generate-001" + ) - model, custom_llm_provider, _, _ = get_llm_provider( - model="veo-3.1-lite-generate-001" - ) - - assert model == "veo-3.1-lite-generate-001" - assert custom_llm_provider == "vertex_ai" - finally: - litellm.model_cost = original_model_cost - litellm.vertex_ai_video_models.clear() - litellm.vertex_ai_video_models.update(original_vertex_video_models) + assert model == "veo-3.1-lite-generate-001" + assert custom_llm_provider == "vertex_ai" def test_veo_31_lite_cost_uses_resolution_tiers(self): model_cost = _load_model_cost_map(BACKUP_MODEL_COST_PATH) model_info = model_cost[VEO_31_LITE_VERTEX_MODEL] - assert ( - video_generation_cost( - model=VEO_31_LITE_VERTEX_MODEL, - duration_seconds=10.0, - custom_llm_provider="vertex_ai", - model_info=dict(model_info), - video_resolution="720p", - ) - == 0.5 - ) - assert ( - video_generation_cost( - model=VEO_31_LITE_VERTEX_MODEL, - duration_seconds=10.0, - custom_llm_provider="vertex_ai", - model_info=dict(model_info), - video_resolution="1080p", - ) - == 0.8 - ) + assert video_generation_cost( + model=VEO_31_LITE_VERTEX_MODEL, + duration_seconds=10.0, + custom_llm_provider="vertex_ai", + model_info=dict(model_info), + video_resolution="720p", + ) == pytest.approx(0.5) + assert video_generation_cost( + model=VEO_31_LITE_VERTEX_MODEL, + duration_seconds=10.0, + custom_llm_provider="vertex_ai", + model_info=dict(model_info), + video_resolution="1080p", + ) == pytest.approx(0.8) def test_transform_video_create_request(self): """Test transformation of video creation request.""" From 1e4b30e4ed198fa532af923f39f38134af1f8a1d Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Thu, 18 Jun 2026 17:35:28 -0500 Subject: [PATCH 10/68] fix(vertex-ai): map veo video size to resolution --- .../llms/vertex_ai/videos/transformation.py | 42 +++++++++++++++---- ...odel_prices_and_context_window_backup.json | 3 +- model_prices_and_context_window.json | 3 +- .../test_vertex_video_transformation.py | 39 +++++++++++++++++ 4 files changed, 78 insertions(+), 9 deletions(-) diff --git a/litellm/llms/vertex_ai/videos/transformation.py b/litellm/llms/vertex_ai/videos/transformation.py index d28f5b5b120..3e7859c49a8 100644 --- a/litellm/llms/vertex_ai/videos/transformation.py +++ b/litellm/llms/vertex_ai/videos/transformation.py @@ -91,6 +91,13 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase): 3. Extract video data (base64) from response """ + _OPENAI_VIDEO_SIZE_TO_ASPECT_RATIO: dict[str, str] = { + "1280x720": "16:9", + "1920x1080": "16:9", + "720x1280": "9:16", + "1080x1920": "9:16", + } + def __init__(self): BaseVideoConfig.__init__(self) VertexBase.__init__(self) @@ -133,6 +140,8 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase): - prompt → prompt (in instances) - input_reference → image (in instances) - size → aspectRatio (e.g., "1280x720" → "16:9") + - size → resolution when inferable ("1280x720"/"720x1280" → "720p", + "1920x1080"/"1080x1920" → "1080p"); skipped if ``resolution`` is already set - seconds → durationSeconds (defaults to 4 seconds if not provided) """ mapped_params: Final[dict[str, Any]] = {} @@ -147,6 +156,9 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase): if "parameters" in video_create_optional_params: mapped_params["parameters"] = video_create_optional_params["parameters"] + if "resolution" in video_create_optional_params: + mapped_params["resolution"] = video_create_optional_params["resolution"] + # Map size to aspectRatio if "size" in video_create_optional_params: size: Final = video_create_optional_params["size"] @@ -154,6 +166,15 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase): aspect_ratio: Final = self._convert_size_to_aspect_ratio(size) if aspect_ratio: mapped_params["aspectRatio"] = aspect_ratio + nested_params: Final = video_create_optional_params.get("parameters") + has_resolution = "resolution" in mapped_params or ( + isinstance(nested_params, dict) + and nested_params.get("resolution") is not None + ) + if not has_resolution: + inferred_resolution = self._convert_size_to_resolution(size) + if inferred_resolution is not None: + mapped_params["resolution"] = inferred_resolution # Map seconds to durationSeconds, default to 4 seconds (matching OpenAI) if "seconds" in video_create_optional_params: @@ -177,14 +198,21 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase): if not size: return None - aspect_ratio_map: Final = { - "1280x720": "16:9", - "1920x1080": "16:9", - "720x1280": "9:16", - "1080x1920": "9:16", - } + return self._OPENAI_VIDEO_SIZE_TO_ASPECT_RATIO.get(size, "16:9") - return aspect_ratio_map.get(size, "16:9") + def _convert_size_to_resolution(self, size: str) -> str | None: + if not size or size not in self._OPENAI_VIDEO_SIZE_TO_ASPECT_RATIO: + return None + try: + width, height = size.split("x", 1) + smaller_edge = min(int(width), int(height)) + except (ValueError, TypeError): + return None + if smaller_edge == 720: + return "720p" + if smaller_edge == 1080: + return "1080p" + return None def validate_environment( self, diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 33580506977..80cb98d91f5 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -39352,7 +39352,8 @@ "output_cost_per_second_1080p": 0.08, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing#veo", "supported_modalities": [ - "text" + "text", + "image" ], "supported_output_modalities": [ "video" diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 33580506977..80cb98d91f5 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -39352,7 +39352,8 @@ "output_cost_per_second_1080p": 0.08, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing#veo", "supported_modalities": [ - "text" + "text", + "image" ], "supported_output_modalities": [ "video" diff --git a/tests/test_litellm/llms/vertex_ai/videos/test_vertex_video_transformation.py b/tests/test_litellm/llms/vertex_ai/videos/test_vertex_video_transformation.py index 927c7dd94b2..ec0aa187ff0 100644 --- a/tests/test_litellm/llms/vertex_ai/videos/test_vertex_video_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/videos/test_vertex_video_transformation.py @@ -153,6 +153,7 @@ class TestVertexAIVideoConfig: assert info["max_input_tokens"] == 1024 assert info["output_cost_per_second"] == 0.05 assert info["output_cost_per_second_1080p"] == 0.08 + assert info["supported_modalities"] == ["text", "image"] def test_veo_31_lite_provider_routing_from_local_model_map( self, monkeypatch: pytest.MonkeyPatch @@ -284,6 +285,44 @@ class TestVertexAIVideoConfig: assert mapped["durationSeconds"] == 8 assert mapped["aspectRatio"] == "16:9" + assert mapped["resolution"] == "720p" + + def test_map_openai_size_to_1080p_resolution(self): + mapped = self.config.map_openai_params( + video_create_optional_params={"size": "1920x1080"}, + model=VEO_31_LITE_VERTEX_MODEL, + drop_params=False, + ) + + assert mapped["aspectRatio"] == "16:9" + assert mapped["resolution"] == "1080p" + + def test_map_openai_size_does_not_override_provider_resolution(self): + mapped = self.config.map_openai_params( + video_create_optional_params={ + "size": "1920x1080", + "parameters": {"resolution": "720p"}, + }, + model=VEO_31_LITE_VERTEX_MODEL, + drop_params=False, + ) + + assert mapped["aspectRatio"] == "16:9" + assert "resolution" not in mapped + assert mapped["parameters"] == {"resolution": "720p"} + + def test_map_openai_size_does_not_override_direct_resolution(self): + mapped = self.config.map_openai_params( + video_create_optional_params={ + "size": "1920x1080", + "resolution": "720p", + }, + model=VEO_31_LITE_VERTEX_MODEL, + drop_params=False, + ) + + assert mapped["aspectRatio"] == "16:9" + assert mapped["resolution"] == "720p" def test_map_openai_params_default_duration(self): """Test that durationSeconds is omitted when not provided.""" From 9012842c2ffd3d93f0ff2a372de87c787280b3fb Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Wed, 15 Jul 2026 21:04:48 -0500 Subject: [PATCH 11/68] fix(vertex-ai): gate veo resolution inference --- .../llms/vertex_ai/videos/transformation.py | 30 ++++++++----------- litellm/types/videos/main.py | 1 + .../test_vertex_video_transformation.py | 24 ++++++++++++--- 3 files changed, 34 insertions(+), 21 deletions(-) diff --git a/litellm/llms/vertex_ai/videos/transformation.py b/litellm/llms/vertex_ai/videos/transformation.py index 3e7859c49a8..9389a5f419d 100644 --- a/litellm/llms/vertex_ai/videos/transformation.py +++ b/litellm/llms/vertex_ai/videos/transformation.py @@ -97,6 +97,12 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase): "720x1280": "9:16", "1080x1920": "9:16", } + _OPENAI_VIDEO_SIZE_TO_RESOLUTION: dict[str, str] = { + "1280x720": "720p", + "1920x1080": "1080p", + "720x1280": "720p", + "1080x1920": "1080p", + } def __init__(self): BaseVideoConfig.__init__(self) @@ -140,8 +146,9 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase): - prompt → prompt (in instances) - input_reference → image (in instances) - size → aspectRatio (e.g., "1280x720" → "16:9") - - size → resolution when inferable ("1280x720"/"720x1280" → "720p", - "1920x1080"/"1080x1920" → "1080p"); skipped if ``resolution`` is already set + - size → resolution for Veo 3 models when inferable + ("1280x720"/"720x1280" → "720p", "1920x1080"/"1080x1920" → "1080p"); + skipped if ``resolution`` is already set - seconds → durationSeconds (defaults to 4 seconds if not provided) """ mapped_params: Final[dict[str, Any]] = {} @@ -168,10 +175,10 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase): mapped_params["aspectRatio"] = aspect_ratio nested_params: Final = video_create_optional_params.get("parameters") has_resolution = "resolution" in mapped_params or ( - isinstance(nested_params, dict) - and nested_params.get("resolution") is not None + isinstance(nested_params, dict) and nested_params.get("resolution") is not None ) - if not has_resolution: + supports_resolution = model.removeprefix("vertex_ai/").startswith("veo-3.") + if supports_resolution and not has_resolution: inferred_resolution = self._convert_size_to_resolution(size) if inferred_resolution is not None: mapped_params["resolution"] = inferred_resolution @@ -201,18 +208,7 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase): return self._OPENAI_VIDEO_SIZE_TO_ASPECT_RATIO.get(size, "16:9") def _convert_size_to_resolution(self, size: str) -> str | None: - if not size or size not in self._OPENAI_VIDEO_SIZE_TO_ASPECT_RATIO: - return None - try: - width, height = size.split("x", 1) - smaller_edge = min(int(width), int(height)) - except (ValueError, TypeError): - return None - if smaller_edge == 720: - return "720p" - if smaller_edge == 1080: - return "1080p" - return None + return self._OPENAI_VIDEO_SIZE_TO_RESOLUTION.get(size) def validate_environment( self, diff --git a/litellm/types/videos/main.py b/litellm/types/videos/main.py index 3677cec3c8f..4c57efe05b0 100644 --- a/litellm/types/videos/main.py +++ b/litellm/types/videos/main.py @@ -76,6 +76,7 @@ class VideoCreateOptionalRequestParams(TypedDict, total=False): image: Any | None # Image for image-to-video; dict with gcsUri/bytesBase64Encoded, or file-like object parameters: dict[str, Any] | None # Provider-specific parameters block passed directly to the API model: str | None + resolution: str | None seconds: str | None size: str | None characters: list[dict[str, str]] | None diff --git a/tests/test_litellm/llms/vertex_ai/videos/test_vertex_video_transformation.py b/tests/test_litellm/llms/vertex_ai/videos/test_vertex_video_transformation.py index ec0aa187ff0..1c2952cbb0b 100644 --- a/tests/test_litellm/llms/vertex_ai/videos/test_vertex_video_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/videos/test_vertex_video_transformation.py @@ -285,17 +285,33 @@ class TestVertexAIVideoConfig: assert mapped["durationSeconds"] == 8 assert mapped["aspectRatio"] == "16:9" - assert mapped["resolution"] == "720p" + assert "resolution" not in mapped - def test_map_openai_size_to_1080p_resolution(self): + @pytest.mark.parametrize( + ("size", "expected_resolution"), + (("1280x720", "720p"), ("1920x1080", "1080p")), + ) + def test_map_openai_size_to_resolution_for_veo_3( + self, size: str, expected_resolution: str + ): mapped = self.config.map_openai_params( - video_create_optional_params={"size": "1920x1080"}, + video_create_optional_params={"size": size}, model=VEO_31_LITE_VERTEX_MODEL, drop_params=False, ) assert mapped["aspectRatio"] == "16:9" - assert mapped["resolution"] == "1080p" + assert mapped["resolution"] == expected_resolution + + def test_map_openai_size_does_not_infer_resolution_for_veo_2(self): + mapped = self.config.map_openai_params( + video_create_optional_params={"size": "1920x1080"}, + model="vertex_ai/veo-2.0-generate-001", + drop_params=False, + ) + + assert mapped["aspectRatio"] == "16:9" + assert "resolution" not in mapped def test_map_openai_size_does_not_override_provider_resolution(self): mapped = self.config.map_openai_params( From d6e859cc97cc99d8fa249627f4459c23fe75d8c3 Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Wed, 15 Jul 2026 21:22:39 -0500 Subject: [PATCH 12/68] fix(vertex-ai): gate resolution inference by pricing metadata --- .../llms/vertex_ai/videos/transformation.py | 11 ++++- .../test_vertex_video_transformation.py | 45 ++++++++++++++++--- 2 files changed, 49 insertions(+), 7 deletions(-) diff --git a/litellm/llms/vertex_ai/videos/transformation.py b/litellm/llms/vertex_ai/videos/transformation.py index 9389a5f419d..bac3d31501b 100644 --- a/litellm/llms/vertex_ai/videos/transformation.py +++ b/litellm/llms/vertex_ai/videos/transformation.py @@ -12,6 +12,7 @@ from typing import TYPE_CHECKING, Any, Final, cast import httpx from httpx._types import RequestFiles +import litellm from litellm.constants import DEFAULT_GOOGLE_VIDEO_DURATION_SECONDS from litellm.images.utils import ImageEditRequestUtils from litellm.llms.base_llm.videos.transformation import BaseVideoConfig @@ -146,7 +147,7 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase): - prompt → prompt (in instances) - input_reference → image (in instances) - size → aspectRatio (e.g., "1280x720" → "16:9") - - size → resolution for Veo 3 models when inferable + - size → resolution for models with resolution-tier pricing when inferable ("1280x720"/"720x1280" → "720p", "1920x1080"/"1080x1920" → "1080p"); skipped if ``resolution`` is already set - seconds → durationSeconds (defaults to 4 seconds if not provided) @@ -177,7 +178,7 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase): has_resolution = "resolution" in mapped_params or ( isinstance(nested_params, dict) and nested_params.get("resolution") is not None ) - supports_resolution = model.removeprefix("vertex_ai/").startswith("veo-3.") + supports_resolution = self._supports_resolution_inference(model) if supports_resolution and not has_resolution: inferred_resolution = self._convert_size_to_resolution(size) if inferred_resolution is not None: @@ -210,6 +211,12 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase): def _convert_size_to_resolution(self, size: str) -> str | None: return self._OPENAI_VIDEO_SIZE_TO_RESOLUTION.get(size) + @staticmethod + def _supports_resolution_inference(model: str) -> bool: + model_key = model if model.startswith("vertex_ai/") else f"vertex_ai/{model}" + model_info = litellm.model_cost.get(model_key, {}) + return model_info.get("output_cost_per_second_1080p") is not None + def validate_environment( self, headers: dict, diff --git a/tests/test_litellm/llms/vertex_ai/videos/test_vertex_video_transformation.py b/tests/test_litellm/llms/vertex_ai/videos/test_vertex_video_transformation.py index 1c2952cbb0b..11cea7c3d23 100644 --- a/tests/test_litellm/llms/vertex_ai/videos/test_vertex_video_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/videos/test_vertex_video_transformation.py @@ -288,15 +288,33 @@ class TestVertexAIVideoConfig: assert "resolution" not in mapped @pytest.mark.parametrize( - ("size", "expected_resolution"), - (("1280x720", "720p"), ("1920x1080", "1080p")), + ("model", "size", "expected_resolution"), + ( + (VEO_31_LITE_VERTEX_MODEL, "1280x720", "720p"), + ( + VEO_31_LITE_VERTEX_MODEL.removeprefix("vertex_ai/"), + "1920x1080", + "1080p", + ), + ), ) - def test_map_openai_size_to_resolution_for_veo_3( - self, size: str, expected_resolution: str + def test_map_openai_size_to_resolution_for_resolution_tier_model( + self, + model: str, + size: str, + expected_resolution: str, + monkeypatch: pytest.MonkeyPatch, ): + model_cost = _load_model_cost_map(BACKUP_MODEL_COST_PATH) + monkeypatch.setitem( + litellm.model_cost, + VEO_31_LITE_VERTEX_MODEL, + dict(model_cost[VEO_31_LITE_VERTEX_MODEL]), + ) + mapped = self.config.map_openai_params( video_create_optional_params={"size": size}, - model=VEO_31_LITE_VERTEX_MODEL, + model=model, drop_params=False, ) @@ -313,6 +331,23 @@ class TestVertexAIVideoConfig: assert mapped["aspectRatio"] == "16:9" assert "resolution" not in mapped + def test_map_openai_size_does_not_infer_resolution_for_existing_veo_3( + self, monkeypatch: pytest.MonkeyPatch + ): + model = "veo-3.1-generate-001" + model_key = f"vertex_ai/{model}" + model_cost = _load_model_cost_map(BACKUP_MODEL_COST_PATH) + monkeypatch.setitem(litellm.model_cost, model_key, dict(model_cost[model_key])) + + mapped = self.config.map_openai_params( + video_create_optional_params={"size": "1920x1080"}, + model=model, + drop_params=False, + ) + + assert mapped["aspectRatio"] == "16:9" + assert "resolution" not in mapped + def test_map_openai_size_does_not_override_provider_resolution(self): mapped = self.config.map_openai_params( video_create_optional_params={ From d40bf05b119a01599c2254e527521eb2da6d1035 Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Wed, 12 Aug 2026 12:15:30 -0500 Subject: [PATCH 13/68] style(vertex-ai): update video test imports --- .../llms/vertex_ai/videos/test_vertex_video_transformation.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/test_litellm/llms/vertex_ai/videos/test_vertex_video_transformation.py b/tests/test_litellm/llms/vertex_ai/videos/test_vertex_video_transformation.py index 11cea7c3d23..f4d3c95a2b1 100644 --- a/tests/test_litellm/llms/vertex_ai/videos/test_vertex_video_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/videos/test_vertex_video_transformation.py @@ -4,8 +4,9 @@ Tests for Vertex AI (Veo) video generation transformation. import base64 import json +from collections.abc import Mapping from pathlib import Path -from typing import Mapping, cast +from typing import cast from unittest.mock import Mock, patch import httpx From 6f3a7c80ca58bf611a087ad533502906deccf85c Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Wed, 12 Aug 2026 12:56:51 -0500 Subject: [PATCH 14/68] style(vertex-ai): annotate Veo class mappings --- litellm/llms/vertex_ai/videos/transformation.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/litellm/llms/vertex_ai/videos/transformation.py b/litellm/llms/vertex_ai/videos/transformation.py index bac3d31501b..eb368245eb7 100644 --- a/litellm/llms/vertex_ai/videos/transformation.py +++ b/litellm/llms/vertex_ai/videos/transformation.py @@ -7,7 +7,7 @@ Based on: https://docs.cloud.google.com/vertex-ai/generative-ai/docs/model-refer import base64 import time -from typing import TYPE_CHECKING, Any, Final, cast +from typing import TYPE_CHECKING, Any, ClassVar, Final, cast import httpx from httpx._types import RequestFiles @@ -92,13 +92,13 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase): 3. Extract video data (base64) from response """ - _OPENAI_VIDEO_SIZE_TO_ASPECT_RATIO: dict[str, str] = { + _OPENAI_VIDEO_SIZE_TO_ASPECT_RATIO: ClassVar[dict[str, str]] = { "1280x720": "16:9", "1920x1080": "16:9", "720x1280": "9:16", "1080x1920": "9:16", } - _OPENAI_VIDEO_SIZE_TO_RESOLUTION: dict[str, str] = { + _OPENAI_VIDEO_SIZE_TO_RESOLUTION: ClassVar[dict[str, str]] = { "1280x720": "720p", "1920x1080": "1080p", "720x1280": "720p", From bbe9fc78e9bb3492831bdcaa1f16a6b432e7d448 Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Wed, 12 Aug 2026 13:39:06 -0500 Subject: [PATCH 15/68] style(vertex-ai): satisfy Veo quality gates --- .../llms/vertex_ai/videos/transformation.py | 36 +++++++++++-------- litellm/types/videos/main.py | 4 +-- 2 files changed, 23 insertions(+), 17 deletions(-) diff --git a/litellm/llms/vertex_ai/videos/transformation.py b/litellm/llms/vertex_ai/videos/transformation.py index eb368245eb7..b925f82c886 100644 --- a/litellm/llms/vertex_ai/videos/transformation.py +++ b/litellm/llms/vertex_ai/videos/transformation.py @@ -7,6 +7,8 @@ Based on: https://docs.cloud.google.com/vertex-ai/generative-ai/docs/model-refer import base64 import time +from collections.abc import Mapping +from types import MappingProxyType from typing import TYPE_CHECKING, Any, ClassVar, Final, cast import httpx @@ -92,18 +94,22 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase): 3. Extract video data (base64) from response """ - _OPENAI_VIDEO_SIZE_TO_ASPECT_RATIO: ClassVar[dict[str, str]] = { - "1280x720": "16:9", - "1920x1080": "16:9", - "720x1280": "9:16", - "1080x1920": "9:16", - } - _OPENAI_VIDEO_SIZE_TO_RESOLUTION: ClassVar[dict[str, str]] = { - "1280x720": "720p", - "1920x1080": "1080p", - "720x1280": "720p", - "1080x1920": "1080p", - } + _OPENAI_VIDEO_SIZE_TO_ASPECT_RATIO: ClassVar[Mapping[str, str]] = MappingProxyType( + { + "1280x720": "16:9", + "1920x1080": "16:9", + "720x1280": "9:16", + "1080x1920": "9:16", + } + ) + _OPENAI_VIDEO_SIZE_TO_RESOLUTION: ClassVar[Mapping[str, str]] = MappingProxyType( + { + "1280x720": "720p", + "1920x1080": "1080p", + "720x1280": "720p", + "1080x1920": "1080p", + } + ) def __init__(self): BaseVideoConfig.__init__(self) @@ -213,9 +219,9 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase): @staticmethod def _supports_resolution_inference(model: str) -> bool: - model_key = model if model.startswith("vertex_ai/") else f"vertex_ai/{model}" - model_info = litellm.model_cost.get(model_key, {}) - return model_info.get("output_cost_per_second_1080p") is not None + model_key: Final = model if model.startswith("vertex_ai/") else f"vertex_ai/{model}" + model_info: Final = litellm.model_cost.get(model_key) + return model_info is not None and model_info.get("output_cost_per_second_1080p") is not None def validate_environment( self, diff --git a/litellm/types/videos/main.py b/litellm/types/videos/main.py index 4c57efe05b0..99b08f6caf6 100644 --- a/litellm/types/videos/main.py +++ b/litellm/types/videos/main.py @@ -2,7 +2,7 @@ from typing import Any, Literal from openai.types.audio.transcription_create_params import FileTypes from pydantic import BaseModel -from typing_extensions import TypedDict +from typing_extensions import ReadOnly, TypedDict class VideoObject(BaseModel): @@ -76,7 +76,7 @@ class VideoCreateOptionalRequestParams(TypedDict, total=False): image: Any | None # Image for image-to-video; dict with gcsUri/bytesBase64Encoded, or file-like object parameters: dict[str, Any] | None # Provider-specific parameters block passed directly to the API model: str | None - resolution: str | None + resolution: ReadOnly[str | None] seconds: str | None size: str | None characters: list[dict[str, str]] | None From 141ada1118bc64c011b67ef9fd7ee2f15854865c Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Mon, 17 Aug 2026 11:25:55 -0400 Subject: [PATCH 16/68] feat(batches): aggregate reasoning tokens and per-line pass/fail counts Batch retrieval already computed cost/usage on completion, but silently dropped reasoning tokens and never counted per-line success/failure. Adds BatchCostUsageResult (replacing bare cost/usage/models tuples) with successful_requests/failed_requests, and threads reasoning_tokens through the aggregated Usage. Both surface on SpendLogs the same way batch_models already does. --- .../proxy/common_utils/check_batch_cost.py | 22 +- litellm/batches/batch_utils.py | 92 +++++-- litellm/litellm_core_utils/litellm_logging.py | 24 +- litellm/proxy/_types.py | 2 + .../spend_tracking/spend_tracking_utils.py | 16 ++ litellm/types/utils.py | 4 +- .../test_batch_custom_pricing.py | 16 +- tests/batches_tests/test_batch_rate_limits.py | 14 +- .../test_batches_logging_unit_tests.py | 58 ++-- .../proxy_unit_tests/test_check_batch_cost.py | 32 ++- .../test_litellm/batches/test_batch_utils.py | 259 ++++++++++++------ .../test_vertex_ai_batch_passthrough.py | 78 +++--- 12 files changed, 406 insertions(+), 211 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 25b00597355..85b9bd77bc3 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -651,16 +651,14 @@ class CheckBatchCost: # Pass deployment model_info so custom batch pricing # (input_cost_per_token_batches etc.) is used for cost calc deployment_model_info = deployment_info.model_info.model_dump() if deployment_info.model_info else {} - batch_cost, batch_usage, batch_models = ( - await calculate_batch_cost_and_usage( - file_content_dictionary=file_content_as_dict, - custom_llm_provider=llm_provider, # type: ignore - model_name=model_name, - model_info=deployment_model_info, # type: ignore[arg-type] - ) + batch_result = await calculate_batch_cost_and_usage( + file_content_dictionary=file_content_as_dict, + custom_llm_provider=llm_provider, # type: ignore + model_name=model_name, + model_info=deployment_model_info, # type: ignore[arg-type] ) logging_obj = LiteLLMLogging( - model=batch_models[0], + model=batch_result.models[0], messages=[{"role": "user", "content": ""}], stream=False, call_type="aretrieve_batch", @@ -684,9 +682,11 @@ class CheckBatchCost: await logging_obj.async_success_handler( result=response, - batch_cost=batch_cost, - batch_usage=batch_usage, - batch_models=batch_models, + batch_cost=batch_result.cost, + batch_usage=batch_result.usage, + batch_models=batch_result.models, + batch_successful_requests=batch_result.successful_requests, + batch_failed_requests=batch_result.failed_requests, ) # Record batch duration (completed_at - created_at) diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 9681d64f656..20aa3c755bd 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -12,12 +12,23 @@ from litellm.types.utils import CallTypes, ModelInfo, Usage from litellm.utils import token_counter +@dataclass(frozen=True, slots=True) +class BatchCostUsageResult: + """Aggregate cost, usage, and per-line pass/fail counts for a completed batch.""" + + cost: float + usage: Usage + models: list[str] + successful_requests: int + failed_requests: int + + async def calculate_batch_cost_and_usage( file_content_dictionary: list[dict], custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"], model_name: str | None = None, model_info: ModelInfo | None = None, -) -> tuple[float, Usage, list[str]]: +) -> BatchCostUsageResult: """ Calculate the cost and usage of a batch. @@ -32,8 +43,7 @@ async def calculate_batch_cost_and_usage( and model_name and getattr(litellm, "disable_vertex_batch_output_transformation", False) ): - batch_cost, batch_usage = calculate_vertex_ai_batch_cost_and_usage(file_content_dictionary, model_name) - return batch_cost, batch_usage, [model_name] + return calculate_vertex_ai_batch_cost_and_usage(file_content_dictionary, model_name) return _aggregate_batch_cost_usage_models( entries=file_content_dictionary, @@ -48,7 +58,7 @@ async def _handle_completed_batch( custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"], model_name: str | None = None, litellm_params: dict | None = None, -) -> tuple[float, Usage, list[str]]: +) -> BatchCostUsageResult: """Fetch a completed batch's output file and aggregate its cost, usage, and models in a single pass over the JSONL lines, so the parsed file content is never materialized in memory. @@ -66,10 +76,7 @@ async def _handle_completed_batch( and model_name and getattr(litellm, "disable_vertex_batch_output_transformation", False) ): - batch_cost, batch_usage = calculate_vertex_ai_batch_cost_and_usage( - _get_file_content_as_dictionary(file_content), model_name - ) - return batch_cost, batch_usage, [model_name] + return calculate_vertex_ai_batch_cost_and_usage(_get_file_content_as_dictionary(file_content), model_name) return _aggregate_batch_cost_usage_models( entries=_iter_batch_input_entries(file_content), @@ -86,19 +93,24 @@ class _BatchOutputLineStats: total_tokens: int cache_read_tokens: int cache_creation_tokens: int + reasoning_tokens: int model: str | None -def _iter_successful_output_line_stats( +def _classify_output_line_stats( entries: Iterable[dict], custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"], model_name: str | None, model_info: ModelInfo | None, -) -> Iterator[_BatchOutputLineStats]: +) -> Iterator[_BatchOutputLineStats | None]: + """Classify every output line in a single pass: yields stats for a + successful line, ``None`` for a failed one (per ``_batch_response_was_successful``). + Counting failures this way avoids a second pass over a potentially huge output file.""" from litellm.cost_calculator import batch_cost_calculator for entry in entries: if not _batch_response_was_successful(entry, custom_llm_provider): + yield None continue response_body = _get_response_from_batch_job_output_file(entry, custom_llm_provider) usage = _get_batch_job_usage_from_response_body(response_body, custom_llm_provider) @@ -123,6 +135,7 @@ def _iter_successful_output_line_stats( custom_llm_provider=custom_llm_provider, call_type=CallTypes.aretrieve_batch.value, ) + reasoning_tokens = usage.completion_tokens_details.reasoning_tokens if usage.completion_tokens_details else None yield _BatchOutputLineStats( cost=line_cost, prompt_tokens=usage.prompt_tokens, @@ -130,6 +143,7 @@ def _iter_successful_output_line_stats( total_tokens=usage.total_tokens, cache_read_tokens=prompt_details["cache_hit_tokens"], cache_creation_tokens=prompt_details["cache_creation_tokens"], + reasoning_tokens=reasoning_tokens or 0, model=response_model, ) @@ -139,10 +153,14 @@ def _aggregate_batch_cost_usage_models( custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"], model_name: str | None = None, model_info: ModelInfo | None = None, -) -> tuple[float, Usage, list[str]]: - """Aggregate cost, usage, and models from batch output entries in a single - pass, holding one small stats record per line instead of the parsed file.""" - line_stats: Final = tuple(_iter_successful_output_line_stats(entries, custom_llm_provider, model_name, model_info)) +) -> BatchCostUsageResult: + """Aggregate cost, usage, models, and pass/fail counts from batch output + entries in a single pass, holding one small stats record per line instead + of the parsed file.""" + all_results: Final = tuple(_classify_output_line_stats(entries, custom_llm_provider, model_name, model_info)) + line_stats: Final = tuple(stats for stats in all_results if stats is not None) + successful_requests: Final = len(line_stats) + failed_requests: Final = len(all_results) - successful_requests cache_token_params: Final = { key: tokens @@ -156,18 +174,32 @@ def _aggregate_batch_cost_usage_models( total_tokens=sum(stats.total_tokens for stats in line_stats), prompt_tokens=sum(stats.prompt_tokens for stats in line_stats), completion_tokens=sum(stats.completion_tokens for stats in line_stats), + reasoning_tokens=sum(stats.reasoning_tokens for stats in line_stats), **cache_token_params, ) batch_models: Final = [model_name] if model_name else [stats.model for stats in line_stats if stats.model] total_cost: Final = sum((stats.cost for stats in line_stats), 0.0) - verbose_logger.debug("batch output aggregate: cost=%s usage=%s models=%s", total_cost, batch_usage, batch_models) - return total_cost, batch_usage, batch_models + verbose_logger.debug( + "batch output aggregate: cost=%s usage=%s models=%s successful=%d failed=%d", + total_cost, + batch_usage, + batch_models, + successful_requests, + failed_requests, + ) + return BatchCostUsageResult( + cost=total_cost, + usage=batch_usage, + models=batch_models, + successful_requests=successful_requests, + failed_requests=failed_requests, + ) def calculate_vertex_ai_batch_cost_and_usage( vertex_ai_batch_responses: list[dict], model_name: str | None = None, -) -> tuple[float, Usage]: +) -> BatchCostUsageResult: """ Calculate both cost and usage from raw Vertex AI batch responses. @@ -178,6 +210,10 @@ def calculate_vertex_ai_batch_cost_and_usage( {"request": ..., "response": {"candidates": [...], "usageMetadata": {...}}} usageMetadata contains promptTokenCount, candidatesTokenCount, totalTokenCount. + + A row with no ``response`` is counted as failed - the same signal already + used to skip it from cost/usage aggregation, since Vertex batch prediction + output doesn't establish a distinct error shape in this (non-default) path. """ from litellm.cost_calculator import batch_cost_calculator @@ -185,12 +221,16 @@ def calculate_vertex_ai_batch_cost_and_usage( total_tokens = 0 prompt_tokens = 0 completion_tokens = 0 + successful_requests = 0 + failed_requests = 0 actual_model_name: Final = model_name or "gemini-2.0-flash-001" for response in vertex_ai_batch_responses: response_body = response.get("response") if response_body is None: + failed_requests += 1 continue + successful_requests += 1 usage_metadata = response_body.get("usageMetadata", {}) _prompt = usage_metadata.get("promptTokenCount", 0) or 0 @@ -218,17 +258,25 @@ def calculate_vertex_ai_batch_cost_and_usage( total_tokens += _total verbose_logger.info( - "vertex_ai batch cost: cost=%s, prompt=%d, completion=%d, total=%d", + "vertex_ai batch cost: cost=%s, prompt=%d, completion=%d, total=%d, successful=%d, failed=%d", total_cost, prompt_tokens, completion_tokens, total_tokens, + successful_requests, + failed_requests, ) - return total_cost, Usage( - total_tokens=total_tokens, - prompt_tokens=prompt_tokens, - completion_tokens=completion_tokens, + return BatchCostUsageResult( + cost=total_cost, + usage=Usage( + total_tokens=total_tokens, + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + ), + models=[actual_model_name], + successful_requests=successful_requests, + failed_requests=failed_requests, ) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index a72d46e3fe8..9d035afbde2 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -2574,6 +2574,8 @@ class Logging(LiteLLMLoggingBaseClass): batch_cost: Final = kwargs.get("batch_cost", None) batch_usage = kwargs.get("batch_usage", None) batch_models = kwargs.get("batch_models", None) + batch_successful_requests = kwargs.get("batch_successful_requests", None) + batch_failed_requests = kwargs.get("batch_failed_requests", 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 = ( @@ -2582,22 +2584,22 @@ class Logging(LiteLLMLoggingBaseClass): if has_explicit_batch_data: result._hidden_params["response_cost"] = batch_cost result._hidden_params["batch_models"] = batch_models + result._hidden_params["batch_successful_requests"] = batch_successful_requests + result._hidden_params["batch_failed_requests"] = batch_failed_requests result.usage = batch_usage elif should_compute_batch_data: - ( - response_cost, - batch_usage, - batch_models, - ) = await _handle_completed_batch( + batch_result = await _handle_completed_batch( batch=result, custom_llm_provider=self.custom_llm_provider, litellm_params=self.litellm_params, ) - result._hidden_params["response_cost"] = response_cost - result._hidden_params["batch_models"] = batch_models - result.usage = batch_usage + result._hidden_params["response_cost"] = batch_result.cost + result._hidden_params["batch_models"] = batch_result.models + result._hidden_params["batch_successful_requests"] = batch_result.successful_requests + result._hidden_params["batch_failed_requests"] = batch_result.failed_requests + result.usage = batch_result.usage start_time, end_time, result = self._success_handler_helper_fn( start_time=start_time, @@ -5062,6 +5064,8 @@ class StandardLoggingPayloadSetup: additional_headers=None, litellm_overhead_time_ms=None, batch_models=None, + batch_successful_requests=None, + batch_failed_requests=None, litellm_model_name=None, usage_object=None, ) @@ -5451,6 +5455,8 @@ def _extract_response_obj_and_hidden_params( response_cost=None, litellm_overhead_time_ms=None, batch_models=None, + batch_successful_requests=None, + batch_failed_requests=None, litellm_model_name=None, usage_object=None, ) @@ -5819,6 +5825,8 @@ def create_dummy_standard_logging_payload() -> StandardLoggingPayload: additional_headers=None, litellm_overhead_time_ms=None, batch_models=None, + batch_successful_requests=None, + batch_failed_requests=None, litellm_model_name=None, usage_object=None, ) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index a566d491597..297fc1bd201 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3462,6 +3462,8 @@ class SpendLogsMetadata(TypedDict): status: StandardLoggingPayloadStatus proxy_server_request: str | None batch_models: list[str] | None + batch_successful_requests: int | None + batch_failed_requests: int | None error_information: StandardLoggingPayloadErrorInformation | None usage_object: dict | None model_map_information: StandardLoggingModelInformation | None diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 3146d8bccfb..f08236adab2 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -74,6 +74,8 @@ def _get_spend_logs_metadata( metadata: dict | None, applied_guardrails: list[str] | None = None, batch_models: list[str] | None = None, + batch_successful_requests: int | None = None, + batch_failed_requests: int | None = None, mcp_tool_call_metadata: StandardLoggingMCPToolCall | None = None, vector_store_request_metadata: list[StandardLoggingVectorStoreRequest] | None = None, guardrail_information: list[StandardLoggingGuardrailInformation] | None = None, @@ -102,6 +104,8 @@ def _get_spend_logs_metadata( error_information=None, proxy_server_request=None, batch_models=None, + batch_successful_requests=None, + batch_failed_requests=None, mcp_tool_call_metadata=None, vector_store_request_metadata=None, model_map_information=None, @@ -128,6 +132,8 @@ def _get_spend_logs_metadata( clean_metadata["user_api_key"] = _hash_api_key_for_spend_log(raw_user_api_key) clean_metadata["applied_guardrails"] = applied_guardrails clean_metadata["batch_models"] = batch_models + clean_metadata["batch_successful_requests"] = batch_successful_requests + clean_metadata["batch_failed_requests"] = batch_failed_requests clean_metadata["mcp_tool_call_metadata"] = mcp_tool_call_metadata clean_metadata["vector_store_request_metadata"] = _get_vector_store_request_for_spend_logs_payload( vector_store_request_metadata @@ -310,6 +316,16 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs if standard_logging_payload is not None else None ), + batch_successful_requests=( + standard_logging_payload.get("hidden_params", {}).get("batch_successful_requests", None) + if standard_logging_payload is not None + else None + ), + batch_failed_requests=( + standard_logging_payload.get("hidden_params", {}).get("batch_failed_requests", None) + if standard_logging_payload is not None + else None + ), mcp_tool_call_metadata=( standard_logging_payload["metadata"].get("mcp_tool_call_metadata", None) if standard_logging_payload is not None diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 272fbabf807..dc8adf9b88e 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -39,7 +39,7 @@ from pydantic import ( field_serializer, field_validator, ) -from typing_extensions import Required, TypedDict +from typing_extensions import ReadOnly, Required, TypedDict from litellm._logging import verbose_logger from litellm._uuid import uuid @@ -2880,6 +2880,8 @@ class StandardLoggingHiddenParams(TypedDict): litellm_overhead_time_ms: float | None additional_headers: StandardLoggingAdditionalHeaders | None batch_models: list[str] | None + batch_successful_requests: ReadOnly[int | None] + batch_failed_requests: ReadOnly[int | None] litellm_model_name: str | None # the model name sent to the provider by litellm usage_object: dict | None diff --git a/tests/batches_tests/test_batch_custom_pricing.py b/tests/batches_tests/test_batch_custom_pricing.py index c2159b564a8..b76b865862a 100644 --- a/tests/batches_tests/test_batch_custom_pricing.py +++ b/tests/batches_tests/test_batch_custom_pricing.py @@ -116,16 +116,16 @@ def test_aggregate_batch_cost_uses_custom_model_info(): """_aggregate_batch_cost_usage_models should thread model_info to batch_cost_calculator.""" file_content = [_make_batch_output_line(prompt_tokens=10, completion_tokens=5)] - cost, _, _ = _aggregate_batch_cost_usage_models( + result = _aggregate_batch_cost_usage_models( entries=file_content, custom_llm_provider="openai", model_info=CUSTOM_MODEL_INFO, ) expected = (10 * 0.00125) + (5 * 0.005) - assert cost == pytest.approx( + assert result.cost == pytest.approx( expected - ), f"Expected total cost {expected}, got {cost}" + ), f"Expected total cost {expected}, got {result.cost}" @pytest.mark.parametrize("data_residency", ["eu", "us"]) @@ -164,15 +164,15 @@ async def test_calculate_batch_cost_and_usage_uses_custom_model_info(): """calculate_batch_cost_and_usage should thread model_info.""" file_content = [_make_batch_output_line(prompt_tokens=10, completion_tokens=5)] - batch_cost, batch_usage, batch_models = await calculate_batch_cost_and_usage( + result = await calculate_batch_cost_and_usage( file_content_dictionary=file_content, custom_llm_provider="openai", model_info=CUSTOM_MODEL_INFO, ) expected = (10 * 0.00125) + (5 * 0.005) - assert batch_cost == pytest.approx( + assert result.cost == pytest.approx( expected - ), f"Expected total cost {expected}, got {batch_cost}" - assert batch_usage.prompt_tokens == 10 - assert batch_usage.completion_tokens == 5 + ), f"Expected total cost {expected}, got {result.cost}" + assert result.usage.prompt_tokens == 10 + assert result.usage.completion_tokens == 5 diff --git a/tests/batches_tests/test_batch_rate_limits.py b/tests/batches_tests/test_batch_rate_limits.py index 2c804d21ace..0baff6c17be 100644 --- a/tests/batches_tests/test_batch_rate_limits.py +++ b/tests/batches_tests/test_batch_rate_limits.py @@ -1027,7 +1027,7 @@ async def test_batch_logging_azure_credentials_regression(): with patch( "litellm.files.main.afile_content", side_effect=mock_afile_content_tracker ): - cost, usage, models = await _handle_completed_batch( + result = await _handle_completed_batch( batch=mock_batch, custom_llm_provider="azure", litellm_params=azure_credentials, @@ -1039,13 +1039,13 @@ async def test_batch_logging_azure_credentials_regression(): ], "REGRESSION: Credentials not passed through _handle_completed_batch" # Verify cost and usage were calculated - assert cost > 0, "Cost should be calculated" - assert usage.total_tokens == 40, "Usage should be calculated correctly" + assert result.cost > 0, "Cost should be calculated" + assert result.usage.total_tokens == 40, "Usage should be calculated correctly" print(" ✓ Credentials passed through full flow") - print(f" ✓ Cost: {cost}") - print(f" ✓ Usage: {usage.total_tokens} tokens") - print(f" ✓ Models: {models}") + print(f" ✓ Cost: {result.cost}") + print(f" ✓ Usage: {result.usage.total_tokens} tokens") + print(f" ✓ Models: {result.models}") # Test 4: Verify error prevention print("\n4. Testing 'Missing credentials' error prevention...") @@ -1064,7 +1064,7 @@ async def test_batch_logging_azure_credentials_regression(): "litellm.files.main.afile_content", side_effect=mock_afile_content_tracker ): try: - cost, usage, models = await _handle_completed_batch( + result = await _handle_completed_batch( batch=mock_batch, custom_llm_provider="azure", litellm_params=azure_credentials, diff --git a/tests/batches_tests/test_batches_logging_unit_tests.py b/tests/batches_tests/test_batches_logging_unit_tests.py index 62b6f5b08e4..9c26514f872 100644 --- a/tests/batches_tests/test_batches_logging_unit_tests.py +++ b/tests/batches_tests/test_batches_logging_unit_tests.py @@ -138,12 +138,12 @@ def test_get_file_content_as_dictionary(sample_file_content): def test_get_batch_job_total_usage_from_file_content(sample_file_content_dict): with patch("litellm.completion_cost", return_value=0.0): - _, usage, _ = _aggregate_batch_cost_usage_models( + result = _aggregate_batch_cost_usage_models( entries=sample_file_content_dict, custom_llm_provider="openai" ) - assert usage.total_tokens == 62 # 30 + 32 - assert usage.prompt_tokens == 42 # 20 + 22 - assert usage.completion_tokens == 20 # 10 + 10 + assert result.usage.total_tokens == 62 # 30 + 32 + assert result.usage.prompt_tokens == 42 # 20 + 22 + assert result.usage.completion_tokens == 20 # 10 + 10 @pytest.mark.asyncio @@ -156,11 +156,11 @@ async def test_batch_cost_calculator(sample_file_content_dict): so we expect the cost to be 0.5 * 2 = 1.0 """ with patch("litellm.completion_cost", return_value=0.5): - cost, _, _ = _aggregate_batch_cost_usage_models( + result = _aggregate_batch_cost_usage_models( entries=sample_file_content_dict, custom_llm_provider="openai", ) - assert cost == 1.0 # 0.5 * 2 successful responses + assert result.cost == 1.0 # 0.5 * 2 successful responses def test_get_response_from_batch_job_output_file(sample_file_content_dict): @@ -226,6 +226,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 + expected_cost = 0.05 expected_usage = litellm.Usage( prompt_tokens=100, @@ -236,7 +238,15 @@ async def test_batch_retrieve_cost_tracking_with_completed_batch_no_explicit_cos with patch( "litellm.litellm_core_utils.litellm_logging._handle_completed_batch", - new=AsyncMock(return_value=(expected_cost, expected_usage, expected_models)), + new=AsyncMock( + return_value=BatchCostUsageResult( + cost=expected_cost, + usage=expected_usage, + models=expected_models, + successful_requests=10, + failed_requests=0, + ) + ), ) as mock_handle_batch: # Call async_success_handler await logging_obj.async_success_handler( @@ -251,6 +261,8 @@ async def test_batch_retrieve_cost_tracking_with_completed_batch_no_explicit_cos # Verify cost and usage were set on the batch result assert mock_batch._hidden_params["response_cost"] == expected_cost assert mock_batch._hidden_params["batch_models"] == expected_models + assert mock_batch._hidden_params["batch_successful_requests"] == 10 + assert mock_batch._hidden_params["batch_failed_requests"] == 0 assert mock_batch.usage == expected_usage @@ -284,7 +296,7 @@ async def test_handle_completed_batch_computes_real_cost_from_output_file( "litellm.batches.batch_utils._fetch_batch_output_file_content", new=AsyncMock(return_value=sample_file_content_bytes), ): - cost, usage, models = await _handle_completed_batch( + result = await _handle_completed_batch( batch=batch, custom_llm_provider="openai" ) @@ -294,16 +306,18 @@ async def test_handle_completed_batch_computes_real_cost_from_output_file( + 20 * pricing["output_cost_per_token_batches"] ) - assert cost == pytest.approx(expected_cost) - assert cost > 0 + assert result.cost == pytest.approx(expected_cost) + assert result.cost > 0 assert ( - cost + result.cost < 42 * pricing["input_cost_per_token"] + 20 * pricing["output_cost_per_token"] ) - assert usage.prompt_tokens == 42 - assert usage.completion_tokens == 20 - assert usage.total_tokens == 62 - assert models == ["gpt-4o-mini-2024-07-18", "gpt-4o-mini-2024-07-18"] + assert result.usage.prompt_tokens == 42 + assert result.usage.completion_tokens == 20 + assert result.usage.total_tokens == 62 + assert result.models == ["gpt-4o-mini-2024-07-18", "gpt-4o-mini-2024-07-18"] + assert result.successful_requests == 2 + assert result.failed_requests == 0 @pytest.mark.asyncio @@ -542,9 +556,19 @@ async def test_batch_retrieve_cost_tracking_with_partial_explicit_data(): ) expected_models = ["gpt-5-mini"] + from litellm.batches.batch_utils import BatchCostUsageResult + with patch( "litellm.litellm_core_utils.litellm_logging._handle_completed_batch", - new=AsyncMock(return_value=(expected_cost, expected_usage, expected_models)), + new=AsyncMock( + return_value=BatchCostUsageResult( + cost=expected_cost, + usage=expected_usage, + models=expected_models, + successful_requests=8, + failed_requests=0, + ) + ), ) as mock_handle_batch: # Call async_success_handler with partial explicit data await logging_obj.async_success_handler( @@ -560,4 +584,6 @@ async def test_batch_retrieve_cost_tracking_with_partial_explicit_data(): # Verify computed cost data was used (not partial explicit data) assert mock_batch._hidden_params["response_cost"] == expected_cost assert mock_batch._hidden_params["batch_models"] == expected_models + assert mock_batch._hidden_params["batch_successful_requests"] == 8 + assert mock_batch._hidden_params["batch_failed_requests"] == 0 assert mock_batch.usage == expected_usage diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index 72f8b87dd16..79739669ff4 100644 --- a/tests/proxy_unit_tests/test_check_batch_cost.py +++ b/tests/proxy_unit_tests/test_check_batch_cost.py @@ -13,6 +13,20 @@ import pytest _IS_B64 = "litellm.proxy.openai_files_endpoints.common_utils._is_base64_encoded_unified_file_id" +def _batch_cost_result(cost, usage, models, successful_requests=1, failed_requests=0): + """Build the BatchCostUsageResult calculate_batch_cost_and_usage now returns, + for mocking it in tests that only care about cost/usage/models.""" + from litellm.batches.batch_utils import BatchCostUsageResult + + return BatchCostUsageResult( + cost=cost, + usage=usage, + models=models, + successful_requests=successful_requests, + failed_requests=failed_requests, + ) + + def _unmanaged_vertex_file_object( input_file_id="gs://bucket/litellm-vertex-files/publishers/google/models/gemini-2.5-flash/abc.jsonl", status="validating", @@ -321,7 +335,7 @@ class TestCheckBatchCost: patch( "litellm.batches.batch_utils.calculate_batch_cost_and_usage", new_callable=AsyncMock, - return_value=( + return_value=_batch_cost_result( 0.01, {"prompt_tokens": 10, "completion_tokens": 5}, ["gpt-4"], @@ -426,7 +440,7 @@ class TestCheckBatchCost: patch( "litellm.batches.batch_utils.calculate_batch_cost_and_usage", new_callable=AsyncMock, - return_value=(0.01, {"prompt_tokens": 10, "completion_tokens": 5}, ["claude-haiku-4-5"]), + return_value=_batch_cost_result(0.01, {"prompt_tokens": 10, "completion_tokens": 5}, ["claude-haiku-4-5"]), ), patch( "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", @@ -526,7 +540,7 @@ class TestCheckBatchCost: patch( "litellm.batches.batch_utils.calculate_batch_cost_and_usage", new_callable=AsyncMock, - return_value=( + return_value=_batch_cost_result( 0.01, {"prompt_tokens": 10, "completion_tokens": 5}, ["gpt-4"], @@ -656,7 +670,7 @@ class TestCheckBatchCost: patch( "litellm.batches.batch_utils.calculate_batch_cost_and_usage", new_callable=AsyncMock, - return_value=( + return_value=_batch_cost_result( 0.01, {"prompt_tokens": 10, "completion_tokens": 5}, ["gpt-4"], @@ -1142,7 +1156,7 @@ class TestCheckBatchCost: patch( "litellm.batches.batch_utils.calculate_batch_cost_and_usage", new_callable=AsyncMock, - return_value=( + return_value=_batch_cost_result( 0.01, {"prompt_tokens": 10, "completion_tokens": 5}, ["gpt-4"], @@ -1271,7 +1285,7 @@ class TestCheckBatchCost: patch( "litellm.batches.batch_utils.calculate_batch_cost_and_usage", new_callable=AsyncMock, - return_value=( + return_value=_batch_cost_result( 0.01, {"prompt_tokens": 10, "completion_tokens": 5}, ["gpt-4"], @@ -1529,7 +1543,7 @@ class TestUnmanagedVertexRouting: patch( "litellm.batches.batch_utils.calculate_batch_cost_and_usage", new_callable=AsyncMock, - return_value=( + return_value=_batch_cost_result( 0.01, {"prompt_tokens": 10, "completion_tokens": 5}, ["gemini-2.5-flash"], @@ -1759,7 +1773,7 @@ class TestUnmanagedBedrockRouting: patch( "litellm.batches.batch_utils.calculate_batch_cost_and_usage", new_callable=AsyncMock, - return_value=( + return_value=_batch_cost_result( 0.02, {"prompt_tokens": 10, "completion_tokens": 5}, ["claude-sonnet-4"], @@ -1951,7 +1965,7 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: patch( "litellm.batches.batch_utils.calculate_batch_cost_and_usage", new_callable=AsyncMock, - return_value=(0.01, {"prompt_tokens": 10}, ["gpt-5.5"]), + return_value=_batch_cost_result(0.01, {"prompt_tokens": 10}, ["gpt-5.5"]), ), patch("litellm.litellm_core_utils.litellm_logging.Logging") as logging_cls, ): diff --git a/tests/test_litellm/batches/test_batch_utils.py b/tests/test_litellm/batches/test_batch_utils.py index d2074853f2b..46969dfc033 100644 --- a/tests/test_litellm/batches/test_batch_utils.py +++ b/tests/test_litellm/batches/test_batch_utils.py @@ -210,10 +210,10 @@ def test_estimate_tokens_never_zero_for_short_rows(): def test_output_models_uses_model_name_override(monkeypatch): monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 0.0) - _, _, models = bu._aggregate_batch_cost_usage_models( + result = bu._aggregate_batch_cost_usage_models( entries=[_success_row(model="ignored")], custom_llm_provider="openai", model_name="forced-model" ) - assert models == ["forced-model"] + assert result.models == ["forced-model"] def test_output_models_collects_from_successful_only(monkeypatch): @@ -223,15 +223,15 @@ def test_output_models_collects_from_successful_only(monkeypatch): _failed_row(model="should-be-skipped"), _success_row(model="claude-3"), ] - _, _, models = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai") - assert models == ["gpt-4o", "claude-3"] + result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai") + assert result.models == ["gpt-4o", "claude-3"] def test_output_models_skips_successful_without_model(monkeypatch): monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 0.0) rows = [{"response": {"status_code": 200, "body": {}}}] - _, _, models = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai") - assert models == [] + result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai") + assert result.models == [] # =========================================================================== # @@ -398,8 +398,8 @@ def test_total_usage_sums_successful_only(monkeypatch): _failed_row(), # excluded _success_row(usage=_usage(20, 10)), # 30 ] - _, usage, _ = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai") - assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == ( + result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai") + assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == ( 30, 15, 45, @@ -417,7 +417,7 @@ def test_total_usage_and_cost_normalize_mixed_responses_and_chat(): ) chat_row = _success_row(usage=_usage(10, 5)) - cost, usage, _ = bu._aggregate_batch_cost_usage_models( + result = bu._aggregate_batch_cost_usage_models( entries=[responses_row, chat_row], custom_llm_provider="openai", model_info={ @@ -426,22 +426,79 @@ def test_total_usage_and_cost_normalize_mixed_responses_and_chat(): }, ) - assert usage.prompt_tokens == 30 - assert usage.completion_tokens == 12 - assert usage.total_tokens == 42 - assert usage.cache_read_input_tokens == 3 - assert cost == pytest.approx((30 * 0.00125) + (12 * 0.005)) + assert result.usage.prompt_tokens == 30 + assert result.usage.completion_tokens == 12 + assert result.usage.total_tokens == 42 + assert result.usage.cache_read_input_tokens == 3 + assert result.cost == pytest.approx((30 * 0.00125) + (12 * 0.005)) def test_total_usage_empty_is_zero(): - cost, usage, models = bu._aggregate_batch_cost_usage_models(entries=[], custom_llm_provider="openai") - assert cost == 0.0 - assert models == [] - assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == ( + result = bu._aggregate_batch_cost_usage_models(entries=[], custom_llm_provider="openai") + assert result.cost == 0.0 + assert result.models == [] + assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == ( 0, 0, 0, ) + assert result.successful_requests == 0 + assert result.failed_requests == 0 + + +def test_total_usage_includes_reasoning_tokens(monkeypatch): + monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 0.0) + rows = [ + _success_row( + usage={ + "prompt_tokens": 10, + "completion_tokens": 50, + "total_tokens": 60, + "completion_tokens_details": {"reasoning_tokens": 30}, + } + ), + _success_row( + usage={ + "prompt_tokens": 5, + "completion_tokens": 20, + "total_tokens": 25, + "completion_tokens_details": {"reasoning_tokens": 8}, + } + ), + _failed_row(), # excluded, must not contribute reasoning tokens either + ] + result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai") + assert result.usage.completion_tokens_details is not None + assert result.usage.completion_tokens_details.reasoning_tokens == 38 + + +def test_aggregate_counts_successful_and_failed_requests(monkeypatch): + monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 0.0) + rows = [ + _success_row(usage=_usage(10, 5)), + _failed_row(), + _success_row(usage=_usage(20, 10)), + _failed_row(), + _failed_row(), + ] + result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai") + assert result.successful_requests == 2 + assert result.failed_requests == 3 + assert result.successful_requests + result.failed_requests == len(rows) + + +def test_aggregate_returns_batch_cost_usage_result_dataclass(monkeypatch): + monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 1.0) + result = bu._aggregate_batch_cost_usage_models( + entries=[_success_row(usage=_usage(10, 5))], custom_llm_provider="openai" + ) + assert isinstance(result, bu.BatchCostUsageResult) + assert (result.cost, result.models, result.successful_requests, result.failed_requests) == ( + 1.0, + ["gpt-4o"], + 1, + 0, + ) # =========================================================================== # @@ -464,10 +521,12 @@ def test_cost_from_content_completion_cost_path(monkeypatch): _success_row(usage=_usage(20, 10)), ] - total, _, _ = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai") + result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai") - assert total == 1.0 # 2 successful * 0.5 + assert result.cost == 1.0 # 2 successful * 0.5 assert len(calls) == 2 # failed row not costed + assert result.successful_requests == 2 + assert result.failed_requests == 1 def test_cost_from_content_model_info_path(monkeypatch): @@ -480,13 +539,13 @@ def test_cost_from_content_model_info_path(monkeypatch): _success_row(usage=_usage(20, 10)), ] - total, _, _ = bu._aggregate_batch_cost_usage_models( + result = bu._aggregate_batch_cost_usage_models( entries=rows, custom_llm_provider="openai", model_info={"input_cost_per_token": 0.0}, # type: ignore[arg-type] # truthy -> model_info path ) - assert total == pytest.approx(0.6) # 2 * (0.1 + 0.2) + assert result.cost == pytest.approx(0.6) # 2 * (0.1 + 0.2) def test_aggregate_consumes_entries_in_a_single_pass(monkeypatch): @@ -496,11 +555,13 @@ def test_aggregate_consumes_entries_in_a_single_pass(monkeypatch): monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 0.5) one_shot = (row for row in [_success_row(usage=_usage(10, 5)), _failed_row(), _success_row(usage=_usage(20, 10))]) - cost, usage, models = bu._aggregate_batch_cost_usage_models(entries=one_shot, custom_llm_provider="openai") + result = bu._aggregate_batch_cost_usage_models(entries=one_shot, custom_llm_provider="openai") - assert cost == 1.0 - assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (30, 15, 45) - assert models == ["gpt-4o", "gpt-4o"] + assert result.cost == 1.0 + assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (30, 15, 45) + assert result.models == ["gpt-4o", "gpt-4o"] + assert result.successful_requests == 2 + assert result.failed_requests == 1 # =========================================================================== # @@ -514,7 +575,13 @@ async def test_calculate_vertex_disable_transform_path(monkeypatch): monkeypatch.setattr( bu, "calculate_vertex_ai_batch_cost_and_usage", - lambda content, model: (9.9, Usage(prompt_tokens=1, completion_tokens=2, total_tokens=3)), + lambda content, model: bu.BatchCostUsageResult( + cost=9.9, + usage=Usage(prompt_tokens=1, completion_tokens=2, total_tokens=3), + models=["gemini-2.0-flash-001"], + successful_requests=1, + failed_requests=0, + ), ) # generic path must NOT be taken monkeypatch.setattr( @@ -523,12 +590,12 @@ async def test_calculate_vertex_disable_transform_path(monkeypatch): lambda **kw: pytest.fail("generic path should not run"), ) - cost, usage, models = await bu.calculate_batch_cost_and_usage( + result = await bu.calculate_batch_cost_and_usage( file_content_dictionary=[], custom_llm_provider="vertex_ai", model_name="gemini-2.0-flash-001" ) - assert cost == 9.9 - assert usage.total_tokens == 3 - assert models == ["gemini-2.0-flash-001"] + assert result.cost == 9.9 + assert result.usage.total_tokens == 3 + assert result.models == ["gemini-2.0-flash-001"] @pytest.mark.asyncio @@ -542,12 +609,12 @@ async def test_calculate_vertex_disable_transform_needs_model_name(monkeypatch): lambda content, model: pytest.fail("raw vertex path should not run"), ) - cost, usage, models = await bu.calculate_batch_cost_and_usage( + result = await bu.calculate_batch_cost_and_usage( file_content_dictionary=[], custom_llm_provider="vertex_ai" ) - assert cost == 0.0 - assert usage.total_tokens == 0 - assert models == [] + assert result.cost == 0.0 + assert result.usage.total_tokens == 0 + assert result.models == [] # =========================================================================== # @@ -580,14 +647,16 @@ def test_vertex_cost_and_usage_aggregation(monkeypatch): }, ] - cost, usage = bu.calculate_vertex_ai_batch_cost_and_usage(responses, "gemini-x") + result = bu.calculate_vertex_ai_batch_cost_and_usage(responses, "gemini-x") - assert cost == pytest.approx(0.6) # 2 * (0.1 + 0.2) - assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == ( + assert result.cost == pytest.approx(0.6) # 2 * (0.1 + 0.2) + assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == ( 30, 15, 45, ) + assert result.successful_requests == 2 + assert result.failed_requests == 0 def test_vertex_cost_skips_none_response_body(monkeypatch): @@ -607,10 +676,12 @@ def test_vertex_cost_skips_none_response_body(monkeypatch): }, ] - cost, usage = bu.calculate_vertex_ai_batch_cost_and_usage(responses, "gemini-x") + result = bu.calculate_vertex_ai_batch_cost_and_usage(responses, "gemini-x") - assert cost == pytest.approx(1.0) # only one line costed - assert usage.total_tokens == 10 + assert result.cost == pytest.approx(1.0) # only one line costed + assert result.usage.total_tokens == 10 + assert result.successful_requests == 1 + assert result.failed_requests == 1 def test_vertex_usage_total_token_fallback(monkeypatch): @@ -620,8 +691,8 @@ def test_vertex_usage_total_token_fallback(monkeypatch): monkeypatch.setattr(cc, "batch_cost_calculator", lambda **kw: (0.0, 0.0)) responses = [{"response": {"usageMetadata": {"promptTokenCount": 8, "candidatesTokenCount": 4}}}] - _, usage = bu.calculate_vertex_ai_batch_cost_and_usage(responses, "gemini-x") - assert usage.total_tokens == 12 + result = bu.calculate_vertex_ai_batch_cost_and_usage(responses, "gemini-x") + assert result.usage.total_tokens == 12 def test_vertex_cost_error_in_line_is_swallowed(monkeypatch): @@ -644,9 +715,9 @@ def test_vertex_cost_error_in_line_is_swallowed(monkeypatch): } ] - cost, usage = bu.calculate_vertex_ai_batch_cost_and_usage(responses, "gemini-x") - assert cost == 0.0 - assert usage.total_tokens == 10 + result = bu.calculate_vertex_ai_batch_cost_and_usage(responses, "gemini-x") + assert result.cost == 0.0 + assert result.usage.total_tokens == 10 # =========================================================================== # @@ -659,13 +730,11 @@ async def test_calculate_batch_cost_and_usage_orchestration(monkeypatch): rows = [_success_row(model="gpt-4o", usage=_usage(10, 5))] monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 2.5) - cost, usage, models = await bu.calculate_batch_cost_and_usage( - file_content_dictionary=rows, custom_llm_provider="openai" - ) + result = await bu.calculate_batch_cost_and_usage(file_content_dictionary=rows, custom_llm_provider="openai") - assert cost == 2.5 - assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (10, 5, 15) - assert models == ["gpt-4o"] + assert result.cost == 2.5 + assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (10, 5, 15) + assert result.models == ["gpt-4o"] # =========================================================================== # @@ -883,16 +952,18 @@ async def test_handle_completed_vertex_batch_computes_cost_usage_and_models(monk monkeypatch.setattr(files_main, "afile_content", fake_afile_content) - cost, usage, models = await bu._handle_completed_batch( + result = await bu._handle_completed_batch( _batch("gs://litellm-bucket/output/predictions.jsonl"), custom_llm_provider="vertex_ai", litellm_params={"vertex_project": "proj-1", "vertex_location": "us-central1"}, ) - assert cost > 0 - assert cost == pytest.approx(30 * 7.5e-07 + 15 * 3.75e-06) - assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (30, 15, 45) - assert models == ["gemini-3.6-flash", "gemini-3.6-flash"] + assert result.cost > 0 + assert result.cost == pytest.approx(30 * 7.5e-07 + 15 * 3.75e-06) + assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (30, 15, 45) + assert result.models == ["gemini-3.6-flash", "gemini-3.6-flash"] + assert result.successful_requests == 2 + assert result.failed_requests == 0 @pytest.mark.asyncio @@ -970,11 +1041,11 @@ async def test_handle_completed_batch_orchestration(monkeypatch): monkeypatch.setattr(bu, "_fetch_batch_output_file_content", fake_fetch) monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 3.3) - cost, usage, models = await bu._handle_completed_batch(_batch("of"), custom_llm_provider="openai") + result = await bu._handle_completed_batch(_batch("of"), custom_llm_provider="openai") - assert cost == 3.3 - assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (10, 5, 15) - assert models == ["gpt-4o"] + assert result.cost == 3.3 + assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (10, 5, 15) + assert result.models == ["gpt-4o"] @pytest.mark.asyncio @@ -991,19 +1062,25 @@ async def test_handle_completed_batch_vertex_disable_transform_path(monkeypatch) def fake_vertex_calc(content, model): seen["content"] = content seen["model"] = model - return 7.7, Usage(prompt_tokens=1, completion_tokens=2, total_tokens=3) + return bu.BatchCostUsageResult( + cost=7.7, + usage=Usage(prompt_tokens=1, completion_tokens=2, total_tokens=3), + models=["gemini-x"], + successful_requests=1, + failed_requests=0, + ) monkeypatch.setattr(bu, "calculate_vertex_ai_batch_cost_and_usage", fake_vertex_calc) - cost, usage, models = await bu._handle_completed_batch( + result = await bu._handle_completed_batch( _batch("gs://litellm-bucket/output/predictions.jsonl"), custom_llm_provider="vertex_ai", model_name="gemini-x", ) - assert cost == 7.7 - assert usage.total_tokens == 3 - assert models == ["gemini-x"] + assert result.cost == 7.7 + assert result.usage.total_tokens == 3 + assert result.models == ["gemini-x"] assert seen["content"] == raw_rows assert seen["model"] == "gemini-x" @@ -1105,14 +1182,14 @@ def test_bedrock_cost_uses_deployment_model_name(): "recordId": "1", "modelOutput": {"model": "claude-sonnet-4-6", "usage": {"input_tokens": 13, "output_tokens": 5}}, } - cost, _, models = bu._aggregate_batch_cost_usage_models( + result = bu._aggregate_batch_cost_usage_models( entries=[row], custom_llm_provider="bedrock", model_name="us.anthropic.claude-sonnet-4-6", model_info={}, ) - assert cost > 0 - assert models == ["us.anthropic.claude-sonnet-4-6"] + assert result.cost > 0 + assert result.models == ["us.anthropic.claude-sonnet-4-6"] def test_anthropic_total_usage_sums_succeeded_only(monkeypatch): @@ -1124,8 +1201,10 @@ def test_anthropic_total_usage_sums_succeeded_only(monkeypatch): _anthropic_errored_row(), _anthropic_succeeded_row(usage=_anthropic_usage(20, 10, cache_read=100)), ] - _, usage, _ = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="anthropic") - assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (130, 15, 145) + result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="anthropic") + assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (130, 15, 145) + assert result.successful_requests == 2 + assert result.failed_requests == 1 def test_anthropic_total_usage_aggregates_cache_token_details(monkeypatch): @@ -1137,11 +1216,11 @@ def test_anthropic_total_usage_aggregates_cache_token_details(monkeypatch): _anthropic_errored_row(), _anthropic_succeeded_row(usage=_anthropic_usage(50, 20, cache_creation=300, cache_read=700)), ] - _, usage, _ = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="anthropic") - assert usage.prompt_tokens_details.cached_tokens == 8700 - assert usage.prompt_tokens_details.cache_creation_tokens == 2300 - assert usage.cache_read_input_tokens == 8700 - assert usage.cache_creation_input_tokens == 2300 + result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="anthropic") + assert result.usage.prompt_tokens_details.cached_tokens == 8700 + assert result.usage.prompt_tokens_details.cache_creation_tokens == 2300 + assert result.usage.cache_read_input_tokens == 8700 + assert result.usage.cache_creation_input_tokens == 2300 def test_total_usage_without_cache_tokens_has_no_prompt_details(monkeypatch): @@ -1152,9 +1231,9 @@ def test_total_usage_without_cache_tokens_has_no_prompt_details(monkeypatch): "response": {"status_code": 200, "body": {"model": "gpt-5.2", "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}}}, } ] - _, usage, _ = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai") - assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (10, 5, 15) - assert usage.prompt_tokens_details is None + result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai") + assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (10, 5, 15) + assert result.usage.prompt_tokens_details is None def test_anthropic_cost_applies_batch_discount_and_cache_pricing(): @@ -1165,14 +1244,14 @@ def test_anthropic_cost_applies_batch_discount_and_cache_pricing(): _anthropic_errored_row(), ] - total, _, _ = bu._aggregate_batch_cost_usage_models( + result = bu._aggregate_batch_cost_usage_models( entries=rows, custom_llm_provider="anthropic", model_info=_ANTHROPIC_MODEL_INFO, # type: ignore[arg-type] ) expected_half_price = (1000 * 3e-6 + 8000 * 3e-7 + 2000 * 3.75e-6 + 200 * 15e-6) / 2 - assert total == pytest.approx(expected_half_price) + assert result.cost == pytest.approx(expected_half_price) def test_anthropic_cost_without_model_info_uses_batch_cost_calculator(monkeypatch): @@ -1191,11 +1270,9 @@ def test_anthropic_cost_without_model_info_uses_batch_cost_calculator(monkeypatc lambda **kw: pytest.fail("anthropic rows must not go through completion_cost"), ) - total, _, _ = bu._aggregate_batch_cost_usage_models( - entries=[_anthropic_succeeded_row()], custom_llm_provider="anthropic" - ) + result = bu._aggregate_batch_cost_usage_models(entries=[_anthropic_succeeded_row()], custom_llm_provider="anthropic") - assert total == pytest.approx(0.3) + assert result.cost == pytest.approx(0.3) assert seen[0]["model"] == "claude-sonnet-4-5-20250929" assert seen[0]["custom_llm_provider"] == "anthropic" assert seen[0]["usage"].prompt_tokens == 10 @@ -1209,8 +1286,8 @@ def test_anthropic_batch_models_collected_from_succeeded_rows(monkeypatch): _anthropic_succeeded_row(model="claude-sonnet-4-5-20250929"), _anthropic_errored_row(), ] - _, _, models = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="anthropic") - assert models == ["claude-sonnet-4-5-20250929"] + result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="anthropic") + assert result.models == ["claude-sonnet-4-5-20250929"] @pytest.mark.asyncio @@ -1220,16 +1297,16 @@ async def test_calculate_batch_cost_and_usage_anthropic_end_to_end(): _anthropic_errored_row(), ] - cost, usage, models = await bu.calculate_batch_cost_and_usage( + result = await bu.calculate_batch_cost_and_usage( file_content_dictionary=rows, custom_llm_provider="anthropic", model_name="claude-sonnet-4-5", model_info=_ANTHROPIC_MODEL_INFO, # type: ignore[arg-type] ) - assert cost == pytest.approx(1000 * 3e-6 / 2 + 8000 * 3e-7 / 2 + 2000 * 3.75e-6 / 2 + 200 * 15e-6 / 2) - assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (11000, 200, 11200) - assert models == ["claude-sonnet-4-5"] + assert result.cost == pytest.approx(1000 * 3e-6 / 2 + 8000 * 3e-7 / 2 + 2000 * 3.75e-6 / 2 + 200 * 15e-6 / 2) + assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (11000, 200, 11200) + assert result.models == ["claude-sonnet-4-5"] def test_extract_credentials_forwards_the_trusted_model_credential_snapshot(): diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_ai_batch_passthrough.py b/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_ai_batch_passthrough.py index ac79c183ca3..1d2d7d4d5c3 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_ai_batch_passthrough.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_ai_batch_passthrough.py @@ -478,14 +478,14 @@ class TestVertexAIBatchPassthroughHandler: } ] - total_cost, usage = calculate_vertex_ai_batch_cost_and_usage( + result = calculate_vertex_ai_batch_cost_and_usage( vertex_ai_batch_responses, model_name="gemini-2.0-flash-001" ) - assert usage.total_tokens == 15 - assert usage.prompt_tokens == 10 - assert usage.completion_tokens == 5 - assert total_cost > 0, "batch_cost_calculator should return a non-zero cost" + assert result.usage.total_tokens == 15 + assert result.usage.prompt_tokens == 10 + assert result.usage.completion_tokens == 5 + assert result.cost > 0, "batch_cost_calculator should return a non-zero cost" def test_batch_response_transformation(self): """Test transformation of Vertex AI batch responses to OpenAI format""" @@ -664,14 +664,14 @@ class TestVertexAIBatchCostCalculation: }, ] - total_cost, usage = calculate_vertex_ai_batch_cost_and_usage( + result = calculate_vertex_ai_batch_cost_and_usage( responses, model_name="gemini-2.0-flash-001" ) - assert usage.prompt_tokens == 18 - assert usage.completion_tokens == 8 - assert usage.total_tokens == 26 - assert total_cost > 0, "batch_cost_calculator should return a non-zero cost" + assert result.usage.prompt_tokens == 18 + assert result.usage.completion_tokens == 8 + assert result.usage.total_tokens == 26 + assert result.cost > 0, "batch_cost_calculator should return a non-zero cost" def test_should_skip_responses_with_null_response_body(self): """Failed lines (response: None) are skipped without error.""" @@ -699,27 +699,29 @@ class TestVertexAIBatchCostCalculation: }, ] - total_cost, usage = calculate_vertex_ai_batch_cost_and_usage( + result = calculate_vertex_ai_batch_cost_and_usage( responses, model_name="gemini-2.0-flash-001" ) - assert usage.prompt_tokens == 18 - assert usage.completion_tokens == 8 - assert usage.total_tokens == 26 - assert total_cost > 0 + assert result.usage.prompt_tokens == 18 + assert result.usage.completion_tokens == 8 + assert result.usage.total_tokens == 26 + assert result.cost > 0 + assert result.successful_requests == 2 + assert result.failed_requests == 1 def test_should_return_zeros_for_empty_response_list(self): """Empty input → zero cost and zero usage.""" from litellm.batches.batch_utils import calculate_vertex_ai_batch_cost_and_usage - total_cost, usage = calculate_vertex_ai_batch_cost_and_usage( + result = calculate_vertex_ai_batch_cost_and_usage( [], model_name="gemini-2.0-flash-001" ) - assert total_cost == 0.0 - assert usage.total_tokens == 0 - assert usage.prompt_tokens == 0 - assert usage.completion_tokens == 0 + assert result.cost == 0.0 + assert result.usage.total_tokens == 0 + assert result.usage.prompt_tokens == 0 + assert result.usage.completion_tokens == 0 def test_should_handle_missing_usage_metadata_gracefully(self): """Response without usageMetadata → 0 tokens, 0 cost for that line.""" @@ -729,13 +731,13 @@ class TestVertexAIBatchCostCalculation: {"response": {"candidates": [{"content": {"parts": [{"text": "hi"}]}}]}}, ] - total_cost, usage = calculate_vertex_ai_batch_cost_and_usage( + result = calculate_vertex_ai_batch_cost_and_usage( responses, model_name="gemini-2.0-flash-001" ) - assert usage.prompt_tokens == 0 - assert usage.completion_tokens == 0 - assert usage.total_tokens == 0 + assert result.usage.prompt_tokens == 0 + assert result.usage.completion_tokens == 0 + assert result.usage.total_tokens == 0 @pytest.mark.asyncio async def test_openai_shaped_output_records_nonzero_cost_and_usage(self): @@ -813,7 +815,7 @@ class TestVertexAIBatchCostCalculation: try: litellm.disable_vertex_batch_output_transformation = False - cost, usage, _ = await calculate_batch_cost_and_usage( + result = await calculate_batch_cost_and_usage( file_content_dictionary=openai_shaped_responses, custom_llm_provider="vertex_ai", model_name="gemini-2.0-flash-001", @@ -822,17 +824,17 @@ class TestVertexAIBatchCostCalculation: litellm.disable_vertex_batch_output_transformation = original_flag assert ( - usage.prompt_tokens == 18 - ), f"expected 18 prompt tokens, got {usage.prompt_tokens}" + result.usage.prompt_tokens == 18 + ), f"expected 18 prompt tokens, got {result.usage.prompt_tokens}" assert ( - usage.completion_tokens == 8 - ), f"expected 8 completion tokens, got {usage.completion_tokens}" + result.usage.completion_tokens == 8 + ), f"expected 8 completion tokens, got {result.usage.completion_tokens}" assert ( - usage.total_tokens == 26 - ), f"expected 26 total tokens, got {usage.total_tokens}" + result.usage.total_tokens == 26 + ), f"expected 26 total tokens, got {result.usage.total_tokens}" assert ( - cost > 0 - ), f"expected non-zero cost for completed Vertex batch, got {cost}" + result.cost > 0 + ), f"expected non-zero cost for completed Vertex batch, got {result.cost}" @pytest.mark.asyncio async def test_raw_vertex_output_still_works_when_transformation_disabled(self): @@ -865,7 +867,7 @@ class TestVertexAIBatchCostCalculation: try: litellm.disable_vertex_batch_output_transformation = True - cost, usage, _ = await calculate_batch_cost_and_usage( + result = await calculate_batch_cost_and_usage( file_content_dictionary=raw_vertex_responses, custom_llm_provider="vertex_ai", model_name="gemini-2.0-flash-001", @@ -873,7 +875,7 @@ class TestVertexAIBatchCostCalculation: finally: litellm.disable_vertex_batch_output_transformation = original_flag - assert usage.prompt_tokens == 10 - assert usage.completion_tokens == 5 - assert usage.total_tokens == 15 - assert cost > 0, "raw Vertex shape should also produce non-zero cost" + assert result.usage.prompt_tokens == 10 + assert result.usage.completion_tokens == 5 + assert result.usage.total_tokens == 15 + assert result.cost > 0, "raw Vertex shape should also produce non-zero cost" From 2bfa1613b4837db3bc01e48297588a123c0c08e9 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Mon, 17 Aug 2026 11:48:04 -0400 Subject: [PATCH 17/68] fix(batches): satisfy LIT010/LIT011/reportPrivateUsage budgets for new fields Suppress the mutation/private-access lints the new batch_successful_requests/batch_failed_requests plumbing triggers, matching the existing suppressed pattern already used for response_cost/batch_models on the same lines. --- litellm/batches/batch_utils.py | 4 ++-- litellm/litellm_core_utils/litellm_logging.py | 14 +++++++------- litellm/proxy/_types.py | 6 +++--- 3 files changed, 12 insertions(+), 12 deletions(-) diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 20aa3c755bd..6f6de18b04e 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -221,8 +221,8 @@ def calculate_vertex_ai_batch_cost_and_usage( total_tokens = 0 prompt_tokens = 0 completion_tokens = 0 - successful_requests = 0 - failed_requests = 0 + successful_requests = 0 # rebind-ok: loop accumulator, matches total_cost/total_tokens above + failed_requests = 0 # rebind-ok: loop accumulator, matches total_cost/total_tokens above actual_model_name: Final = model_name or "gemini-2.0-flash-001" for response in vertex_ai_batch_responses: diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 9d035afbde2..c795b181064 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -2574,8 +2574,8 @@ class Logging(LiteLLMLoggingBaseClass): batch_cost: Final = kwargs.get("batch_cost", None) batch_usage = kwargs.get("batch_usage", None) batch_models = kwargs.get("batch_models", None) - batch_successful_requests = kwargs.get("batch_successful_requests", None) - batch_failed_requests = kwargs.get("batch_failed_requests", None) + batch_successful_requests: Final = kwargs.get("batch_successful_requests", None) + batch_failed_requests: Final = kwargs.get("batch_failed_requests", 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 = ( @@ -2584,12 +2584,12 @@ class Logging(LiteLLMLoggingBaseClass): if has_explicit_batch_data: result._hidden_params["response_cost"] = batch_cost result._hidden_params["batch_models"] = batch_models - result._hidden_params["batch_successful_requests"] = batch_successful_requests - result._hidden_params["batch_failed_requests"] = batch_failed_requests + result._hidden_params["batch_successful_requests"] = batch_successful_requests # pyright: ignore[reportPrivateUsage] # rebind-ok: same result._hidden_params pattern as response_cost/batch_models above + result._hidden_params["batch_failed_requests"] = batch_failed_requests # pyright: ignore[reportPrivateUsage] # rebind-ok: same pattern as above result.usage = batch_usage elif should_compute_batch_data: - batch_result = await _handle_completed_batch( + batch_result: Final = await _handle_completed_batch( batch=result, custom_llm_provider=self.custom_llm_provider, litellm_params=self.litellm_params, @@ -2597,8 +2597,8 @@ class Logging(LiteLLMLoggingBaseClass): result._hidden_params["response_cost"] = batch_result.cost result._hidden_params["batch_models"] = batch_result.models - result._hidden_params["batch_successful_requests"] = batch_result.successful_requests - result._hidden_params["batch_failed_requests"] = batch_result.failed_requests + result._hidden_params["batch_successful_requests"] = batch_result.successful_requests # pyright: ignore[reportPrivateUsage] # rebind-ok: same pattern as above + result._hidden_params["batch_failed_requests"] = batch_result.failed_requests # pyright: ignore[reportPrivateUsage] # rebind-ok: same pattern as above result.usage = batch_result.usage start_time, end_time, result = self._success_handler_helper_fn( diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 297fc1bd201..cbc8c003868 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -15,7 +15,7 @@ from pydantic import ( field_validator, model_validator, ) -from typing_extensions import NotRequired, Required, TypedDict +from typing_extensions import NotRequired, ReadOnly, Required, TypedDict from litellm._uuid import uuid from litellm.constants import DEFAULT_STAGGER_WINDOW_SECONDS, MCP_STDIO_ALLOWED_COMMANDS @@ -3462,8 +3462,8 @@ class SpendLogsMetadata(TypedDict): status: StandardLoggingPayloadStatus proxy_server_request: str | None batch_models: list[str] | None - batch_successful_requests: int | None - batch_failed_requests: int | None + batch_successful_requests: ReadOnly[int | None] + batch_failed_requests: ReadOnly[int | None] error_information: StandardLoggingPayloadErrorInformation | None usage_object: dict | None model_map_information: StandardLoggingModelInformation | None From 02cf319b483d8c6961b5464f0a72875c0ee6ae4a Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Mon, 17 Aug 2026 14:58:14 -0400 Subject: [PATCH 18/68] fix(batches): count failures reported only in the batch error file Live verification against a real OpenAI batch showed per-request failures (e.g. a rejected param) land in error_file_id, never in the output file, so failed_requests silently undercounted them (0 instead of the real 1). _handle_completed_batch now also fetches error_file_id when present and folds its line count into failed_requests. --- litellm/batches/batch_utils.py | 124 +++++++++++++----- .../test_litellm/batches/test_batch_utils.py | 63 +++++++++ 2 files changed, 152 insertions(+), 35 deletions(-) diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 6f6de18b04e..4889be3c0f6 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -1,6 +1,7 @@ import json from collections.abc import Iterable, Iterator from dataclasses import dataclass +from dataclasses import replace as dataclasses_replace from typing import Any, Final, Literal import litellm @@ -70,18 +71,28 @@ async def _handle_completed_batch( litellm_params: Optional litellm parameters containing credentials (api_key, api_base, etc.) """ 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 + ) - if ( - custom_llm_provider == "vertex_ai" - and model_name - and getattr(litellm, "disable_vertex_batch_output_transformation", False) - ): - return calculate_vertex_ai_batch_cost_and_usage(_get_file_content_as_dictionary(file_content), model_name) + output_file_result: Final = ( + calculate_vertex_ai_batch_cost_and_usage(_get_file_content_as_dictionary(file_content), model_name) + if ( + custom_llm_provider == "vertex_ai" + and model_name + and getattr(litellm, "disable_vertex_batch_output_transformation", False) + ) + else _aggregate_batch_cost_usage_models( + entries=_iter_batch_input_entries(file_content), + custom_llm_provider=custom_llm_provider, + model_name=model_name, + ) + ) - return _aggregate_batch_cost_usage_models( - entries=_iter_batch_input_entries(file_content), - custom_llm_provider=custom_llm_provider, - model_name=model_name, + 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 ) @@ -280,6 +291,50 @@ def calculate_vertex_ai_batch_cost_and_usage( ) +async def _fetch_batch_managed_file_content( + file_id: str, + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"] = "openai", + litellm_params: dict | None = None, +) -> bytes: + """ + Fetch a batch's output or error file and return its raw JSONL bytes. + + Args: + file_id: The provider or unified (litellm-managed) file id to fetch + custom_llm_provider: The LLM provider + litellm_params: Optional litellm parameters containing credentials (api_key, api_base, etc.) + Required for Azure and other providers that need authentication + """ + from litellm.files.main import afile_content + from litellm.proxy.openai_files_endpoints.common_utils import ( + _is_base64_encoded_unified_file_id, + ) + + resolved_file_id = file_id + is_base64_unified_file_id: Final = _is_base64_encoded_unified_file_id(file_id) + if is_base64_unified_file_id: + try: + resolved_file_id = is_base64_unified_file_id.split("llm_output_file_id,")[1].split(";")[0] + verbose_logger.debug("Extracted LLM output file ID from unified file ID: %s", resolved_file_id) + except (IndexError, AttributeError) as e: + verbose_logger.error( + "Failed to extract LLM output file ID from unified file ID: %s, error: %s", file_id, e + ) + + # Build kwargs for afile_content with credentials from litellm_params + file_content_kwargs: Final = { + "file_id": resolved_file_id, + "custom_llm_provider": custom_llm_provider, + } + + # Extract and add credentials for file access + credentials: Final = _extract_file_access_credentials(litellm_params) + file_content_kwargs.update(credentials) + + _file_content: Final = await afile_content(**file_content_kwargs) + return _file_content.content + + async def _fetch_batch_output_file_content( batch: Batch, custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"] = "openai", @@ -294,37 +349,36 @@ async def _fetch_batch_output_file_content( litellm_params: Optional litellm parameters containing credentials (api_key, api_base, etc.) Required for Azure and other providers that need authentication """ - from litellm.files.main import afile_content - from litellm.proxy.openai_files_endpoints.common_utils import ( - _is_base64_encoded_unified_file_id, - ) - if batch.output_file_id is None: raise ValueError("Output file id is None cannot retrieve file content") - file_id = batch.output_file_id - is_base64_unified_file_id: Final = _is_base64_encoded_unified_file_id(file_id) - if is_base64_unified_file_id: - try: - file_id = is_base64_unified_file_id.split("llm_output_file_id,")[1].split(";")[0] - verbose_logger.debug("Extracted LLM output file ID from unified file ID: %s", file_id) - except (IndexError, AttributeError) as e: - verbose_logger.error( - "Failed to extract LLM output file ID from unified file ID: %s, error: %s", batch.output_file_id, e - ) + return await _fetch_batch_managed_file_content( + batch.output_file_id, custom_llm_provider=custom_llm_provider, litellm_params=litellm_params + ) - # Build kwargs for afile_content with credentials from litellm_params - file_content_kwargs: Final = { - "file_id": file_id, - "custom_llm_provider": custom_llm_provider, - } - # Extract and add credentials for file access - credentials: Final = _extract_file_access_credentials(litellm_params) - file_content_kwargs.update(credentials) +async def _count_error_file_failed_requests( + batch: Batch, + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"], + litellm_params: dict | None, +) -> int: + """Count failed requests reported only in the batch's separate error file. - _file_content: Final = await afile_content(**file_content_kwargs) - return _file_content.content + OpenAI-shaped batch providers write successful lines to ``output_file_id`` + and per-request failures (e.g. a rejected param) to a distinct + ``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 + ) + 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/tests/test_litellm/batches/test_batch_utils.py b/tests/test_litellm/batches/test_batch_utils.py index 46969dfc033..06f42ce7b51 100644 --- a/tests/test_litellm/batches/test_batch_utils.py +++ b/tests/test_litellm/batches/test_batch_utils.py @@ -1048,6 +1048,69 @@ async def test_handle_completed_batch_orchestration(monkeypatch): assert result.models == ["gpt-4o"] +@pytest.mark.asyncio +async def test_handle_completed_batch_counts_error_file_failures(monkeypatch): + """Regression test: OpenAI writes per-request failures (e.g. a rejected param) + to a separate error_file_id, never into the output file - so failed_requests + must include them or it silently undercounts real batch failures.""" + from litellm.types.llms.openai import Batch + + rows = [_success_row(model="gpt-5-mini", usage=_usage(24, 107))] + error_rows = [ + { + "id": "batch_req_err1", + "custom_id": "req-2-bad", + "response": {"status_code": 400, "body": {"error": {"message": "Invalid 'temperature'"}}}, + "error": None, + } + ] + + async def fake_fetch(batch, custom_llm_provider, litellm_params=None): + return _vertex_jsonl(rows) + + async def fake_afile_content(**kw): + return type("R", (), {"content": _vertex_jsonl(error_rows)})() + + 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(litellm, "completion_cost", lambda **kw: 0.0) + + batch = Batch( + id="b", + completion_window="24h", + created_at=1, + endpoint="/v1/chat/completions", + input_file_id="f", + object="batch", + status="completed", + output_file_id="of", + error_file_id="ef", + ) + + result = await bu._handle_completed_batch(batch, custom_llm_provider="openai") + + assert result.successful_requests == 1 + assert result.failed_requests == 1 + + +@pytest.mark.asyncio +async def test_handle_completed_batch_no_error_file_id_reports_zero_error_failures(monkeypatch): + rows = [_success_row(model="gpt-4o", usage=_usage(10, 5))] + + async def fake_fetch(batch, custom_llm_provider, litellm_params=None): + return _vertex_jsonl(rows) + + monkeypatch.setattr(bu, "_fetch_batch_output_file_content", fake_fetch) + monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 0.0) + + result = await bu._handle_completed_batch(_batch("of"), custom_llm_provider="openai") + + assert result.successful_requests == 1 + assert result.failed_requests == 0 + + @pytest.mark.asyncio async def test_handle_completed_batch_vertex_disable_transform_path(monkeypatch): raw_rows = [{"response": {"usageMetadata": {"promptTokenCount": 1, "candidatesTokenCount": 2}}}] From 682032dd1b23ba49aab8ba486b44716805619ec1 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Tue, 18 Aug 2026 19:29:25 -0400 Subject: [PATCH 19/68] fix(batches): fix CI failures from lint format and merged upstream guard - ruff format litellm/batches/batch_utils.py - reconcile the upstream output_file_id=None guard (merged in from litellm_internal_staging) with BatchCostUsageResult, and count that batch's error_file_id failures instead of always reporting 0 - fix test_handle_completed_batch_no_output_file_is_zero's tuple unpacking, which predated the BatchCostUsageResult refactor - commit the batch_successful_requests/batch_failed_requests fixture fix to test_spend_management_endpoints.py that was left uncommitted --- litellm/batches/batch_utils.py | 14 ++++++++++---- tests/test_litellm/batches/test_batch_utils.py | 10 ++++++---- .../test_spend_management_endpoints.py | 6 +++--- 3 files changed, 19 insertions(+), 11 deletions(-) diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 8a42efd6dc5..4eb9c7a5dfa 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -79,7 +79,15 @@ async def _handle_completed_batch( # The generic retrieval helper keeps raising for callers that explicitly ask # for a missing output file. if batch.output_file_id is None: - return 0.0, Usage(prompt_tokens=0, completion_tokens=0, total_tokens=0), [] + return BatchCostUsageResult( + cost=0.0, + usage=Usage(prompt_tokens=0, completion_tokens=0, total_tokens=0), + models=[], + successful_requests=0, + failed_requests=await _count_error_file_failed_requests( + batch, custom_llm_provider=custom_llm_provider, litellm_params=litellm_params + ), + ) 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( @@ -328,9 +336,7 @@ async def _fetch_batch_managed_file_content( resolved_file_id = is_base64_unified_file_id.split("llm_output_file_id,")[1].split(";")[0] verbose_logger.debug("Extracted LLM output file ID from unified file ID: %s", resolved_file_id) except (IndexError, AttributeError) as e: - verbose_logger.error( - "Failed to extract LLM output file ID from unified file ID: %s, error: %s", file_id, e - ) + verbose_logger.error("Failed to extract LLM output file ID from unified file ID: %s, error: %s", file_id, e) # Build kwargs for afile_content with credentials from litellm_params file_content_kwargs: Final = { diff --git a/tests/test_litellm/batches/test_batch_utils.py b/tests/test_litellm/batches/test_batch_utils.py index daf3b91a26b..3a69911dfe6 100644 --- a/tests/test_litellm/batches/test_batch_utils.py +++ b/tests/test_litellm/batches/test_batch_utils.py @@ -1126,11 +1126,13 @@ async def test_handle_completed_batch_no_output_file_is_zero(monkeypatch): monkeypatch.setattr(bu, "_fetch_batch_output_file_content", _must_not_fetch) - cost, usage, models = await bu._handle_completed_batch(_batch(None), custom_llm_provider="openai") + result = await bu._handle_completed_batch(_batch(None), custom_llm_provider="openai") - assert cost == 0.0 - assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (0, 0, 0) - assert models == [] + assert result.cost == 0.0 + assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (0, 0, 0) + assert result.models == [] + assert result.successful_requests == 0 + assert result.failed_requests == 0 @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py index 7052e050806..b8ed2b04b9d 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py @@ -2633,7 +2633,7 @@ class TestSpendLogsPayload: "model": "gpt-4o", "user": "", "team_id": "", - "metadata": '{"applied_guardrails": [], "batch_models": null, "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "routing_decision": null, "internal_call_origin": null, "guardrail_information": null, "compression_savings": null, "usage_object": {"completion_tokens": 20, "prompt_tokens": 10, "total_tokens": 30, "completion_tokens_details": null, "prompt_tokens_details": null}, "model_map_information": {"model_map_key": "gpt-4o", "model_map_value": {"key": "gpt-4o", "max_tokens": 16384, "max_input_tokens": 128000, "max_output_tokens": 16384, "input_cost_per_token": 2.5e-06, "cache_creation_input_token_cost": null, "cache_read_input_token_cost": 1.25e-06, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": 1.25e-06, "output_cost_per_token_batches": 5e-06, "output_cost_per_token": 1e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_reasoning_token": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "openai", "mode": "chat", "supports_system_messages": true, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": false, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": false, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": true, "supports_reasoning": false, "search_context_cost_per_query": {"search_context_size_low": 0.03, "search_context_size_medium": 0.035, "search_context_size_high": 0.05}, "tpm": null, "rpm": null, "supported_openai_params": ["frequency_penalty", "logit_bias", "logprobs", "top_logprobs", "max_tokens", "max_completion_tokens", "modalities", "prediction", "n", "presence_penalty", "seed", "stop", "stream", "stream_options", "temperature", "top_p", "tools", "tool_choice", "function_call", "functions", "max_retries", "extra_headers", "parallel_tool_calls", "audio", "response_format", "user"]}}, "additional_usage_values": {"completion_tokens_details": null, "prompt_tokens_details": null}}', + "metadata": '{"applied_guardrails": [], "batch_models": null, "batch_successful_requests": null, "batch_failed_requests": null, "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "routing_decision": null, "internal_call_origin": null, "guardrail_information": null, "compression_savings": null, "usage_object": {"completion_tokens": 20, "prompt_tokens": 10, "total_tokens": 30, "completion_tokens_details": null, "prompt_tokens_details": null}, "model_map_information": {"model_map_key": "gpt-4o", "model_map_value": {"key": "gpt-4o", "max_tokens": 16384, "max_input_tokens": 128000, "max_output_tokens": 16384, "input_cost_per_token": 2.5e-06, "cache_creation_input_token_cost": null, "cache_read_input_token_cost": 1.25e-06, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": 1.25e-06, "output_cost_per_token_batches": 5e-06, "output_cost_per_token": 1e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_reasoning_token": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "openai", "mode": "chat", "supports_system_messages": true, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": false, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": false, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": true, "supports_reasoning": false, "search_context_cost_per_query": {"search_context_size_low": 0.03, "search_context_size_medium": 0.035, "search_context_size_high": 0.05}, "tpm": null, "rpm": null, "supported_openai_params": ["frequency_penalty", "logit_bias", "logprobs", "top_logprobs", "max_tokens", "max_completion_tokens", "modalities", "prediction", "n", "presence_penalty", "seed", "stop", "stream", "stream_options", "temperature", "top_p", "tools", "tool_choice", "function_call", "functions", "max_retries", "extra_headers", "parallel_tool_calls", "audio", "response_format", "user"]}}, "additional_usage_values": {"completion_tokens_details": null, "prompt_tokens_details": null}}', "cache_key": "Cache OFF", "spend": 0.00022500000000000002, "total_tokens": 30, @@ -2729,7 +2729,7 @@ class TestSpendLogsPayload: "model": "claude-4-sonnet-20250514", "user": "", "team_id": "", - "metadata": '{"applied_guardrails": [], "batch_models": null, "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "routing_decision": null, "internal_call_origin": null, "guardrail_information": null, "compression_savings": null, "usage_object": {"completion_tokens": 503, "prompt_tokens": 2095, "total_tokens": 2598, "completion_tokens_details": null, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}, "model_map_information": {"model_map_key": "claude-4-sonnet-20250514", "model_map_value": {"key": "claude-4-sonnet-20250514", "max_tokens": 128000, "max_input_tokens": 200000, "max_output_tokens": 128000, "input_cost_per_token": 3e-06, "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": null, "output_cost_per_token_batches": null, "output_cost_per_token": 1.5e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "anthropic", "mode": "chat", "supports_system_messages": null, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": true, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": true, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": false, "supports_reasoning": true, "search_context_cost_per_query": null, "tpm": null, "rpm": null, "supported_openai_params": ["stream", "stop", "temperature", "top_p", "max_tokens", "max_completion_tokens", "tools", "tool_choice", "extra_headers", "parallel_tool_calls", "response_format", "user", "reasoning_effort", "thinking"]}}, "additional_usage_values": {"completion_tokens_details": {"accepted_prediction_tokens": null, "audio_tokens": null, "reasoning_tokens": null, "rejected_prediction_tokens": null, "text_tokens": 503, "image_tokens": null}, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0, "text_tokens": null, "image_tokens": null}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}}', + "metadata": '{"applied_guardrails": [], "batch_models": null, "batch_successful_requests": null, "batch_failed_requests": null, "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "routing_decision": null, "internal_call_origin": null, "guardrail_information": null, "compression_savings": null, "usage_object": {"completion_tokens": 503, "prompt_tokens": 2095, "total_tokens": 2598, "completion_tokens_details": null, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}, "model_map_information": {"model_map_key": "claude-4-sonnet-20250514", "model_map_value": {"key": "claude-4-sonnet-20250514", "max_tokens": 128000, "max_input_tokens": 200000, "max_output_tokens": 128000, "input_cost_per_token": 3e-06, "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": null, "output_cost_per_token_batches": null, "output_cost_per_token": 1.5e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "anthropic", "mode": "chat", "supports_system_messages": null, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": true, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": true, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": false, "supports_reasoning": true, "search_context_cost_per_query": null, "tpm": null, "rpm": null, "supported_openai_params": ["stream", "stop", "temperature", "top_p", "max_tokens", "max_completion_tokens", "tools", "tool_choice", "extra_headers", "parallel_tool_calls", "response_format", "user", "reasoning_effort", "thinking"]}}, "additional_usage_values": {"completion_tokens_details": {"accepted_prediction_tokens": null, "audio_tokens": null, "reasoning_tokens": null, "rejected_prediction_tokens": null, "text_tokens": 503, "image_tokens": null}, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0, "text_tokens": null, "image_tokens": null}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}}', "cache_key": "Cache OFF", "spend": 0.01383, "total_tokens": 2598, @@ -2823,7 +2823,7 @@ class TestSpendLogsPayload: "model": "claude-4-sonnet-20250514", "user": "", "team_id": "", - "metadata": '{"applied_guardrails": [], "batch_models": null, "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "routing_decision": null, "internal_call_origin": null, "guardrail_information": null, "compression_savings": null, "usage_object": {"completion_tokens": 503, "prompt_tokens": 2095, "total_tokens": 2598, "completion_tokens_details": null, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}, "model_map_information": {"model_map_key": "claude-4-sonnet-20250514", "model_map_value": {"key": "claude-4-sonnet-20250514", "max_tokens": 128000, "max_input_tokens": 200000, "max_output_tokens": 128000, "input_cost_per_token": 3e-06, "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": null, "output_cost_per_token_batches": null, "output_cost_per_token": 1.5e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "anthropic", "mode": "chat", "supports_system_messages": null, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": true, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": true, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": false, "supports_reasoning": true, "search_context_cost_per_query": null, "tpm": null, "rpm": null, "supported_openai_params": ["stream", "stop", "temperature", "top_p", "max_tokens", "max_completion_tokens", "tools", "tool_choice", "extra_headers", "parallel_tool_calls", "response_format", "user", "reasoning_effort", "thinking"]}}, "additional_usage_values": {"completion_tokens_details": {"accepted_prediction_tokens": null, "audio_tokens": null, "reasoning_tokens": null, "rejected_prediction_tokens": null, "text_tokens": 503, "image_tokens": null}, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0, "text_tokens": null, "image_tokens": null}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}}', + "metadata": '{"applied_guardrails": [], "batch_models": null, "batch_successful_requests": null, "batch_failed_requests": null, "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "routing_decision": null, "internal_call_origin": null, "guardrail_information": null, "compression_savings": null, "usage_object": {"completion_tokens": 503, "prompt_tokens": 2095, "total_tokens": 2598, "completion_tokens_details": null, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}, "model_map_information": {"model_map_key": "claude-4-sonnet-20250514", "model_map_value": {"key": "claude-4-sonnet-20250514", "max_tokens": 128000, "max_input_tokens": 200000, "max_output_tokens": 128000, "input_cost_per_token": 3e-06, "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": null, "output_cost_per_token_batches": null, "output_cost_per_token": 1.5e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "anthropic", "mode": "chat", "supports_system_messages": null, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": true, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": true, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": false, "supports_reasoning": true, "search_context_cost_per_query": null, "tpm": null, "rpm": null, "supported_openai_params": ["stream", "stop", "temperature", "top_p", "max_tokens", "max_completion_tokens", "tools", "tool_choice", "extra_headers", "parallel_tool_calls", "response_format", "user", "reasoning_effort", "thinking"]}}, "additional_usage_values": {"completion_tokens_details": {"accepted_prediction_tokens": null, "audio_tokens": null, "reasoning_tokens": null, "rejected_prediction_tokens": null, "text_tokens": 503, "image_tokens": null}, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0, "text_tokens": null, "image_tokens": null}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}}', "cache_key": "Cache OFF", "spend": 0.01383, "total_tokens": 2598, From ff833a5872cc167a7166a58fdf78aeb2d175b570 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Tue, 18 Aug 2026 19:37:07 -0400 Subject: [PATCH 20/68] fix(batches): address Greptile type-discipline feedback Add Final to the reasoning_tokens local var and type the _batch_cost_result test helper's parameters, per review feedback on PR #37208. --- litellm/batches/batch_utils.py | 4 +- .../proxy_unit_tests/test_check_batch_cost.py | 519 ++++++------------ 2 files changed, 167 insertions(+), 356 deletions(-) diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 4eb9c7a5dfa..057c9978879 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -165,7 +165,9 @@ def _classify_output_line_stats( custom_llm_provider=custom_llm_provider, call_type=CallTypes.aretrieve_batch.value, ) - reasoning_tokens = usage.completion_tokens_details.reasoning_tokens if usage.completion_tokens_details else None + reasoning_tokens: Final = ( + usage.completion_tokens_details.reasoning_tokens if usage.completion_tokens_details else None + ) yield _BatchOutputLineStats( cost=line_cost, prompt_tokens=usage.prompt_tokens, diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index 139bf583e6a..ff2dce498f0 100644 --- a/tests/proxy_unit_tests/test_check_batch_cost.py +++ b/tests/proxy_unit_tests/test_check_batch_cost.py @@ -6,14 +6,24 @@ Vertex (raw gs:// input_file_id) and Bedrock (raw s3:// input_file_id, ARN unified_object_id) batches with no managed unified id. """ +from typing import TYPE_CHECKING from unittest.mock import AsyncMock, MagicMock, patch import pytest +if TYPE_CHECKING: + from litellm.batches.batch_utils import BatchCostUsageResult + _IS_B64 = "litellm.proxy.openai_files_endpoints.common_utils._is_base64_encoded_unified_file_id" -def _batch_cost_result(cost, usage, models, successful_requests=1, failed_requests=0): +def _batch_cost_result( + cost: float, + usage: dict, + models: list[str], + successful_requests: int = 1, + failed_requests: int = 0, +) -> "BatchCostUsageResult": """Build the BatchCostUsageResult calculate_batch_cost_and_usage now returns, for mocking it in tests that only care about cost/usage/models.""" from litellm.batches.batch_utils import BatchCostUsageResult @@ -90,9 +100,7 @@ class TestCheckBatchCost: return MagicMock() @pytest.fixture - def check_batch_cost_instance( - self, mock_proxy_logging_obj, mock_prisma_client, mock_llm_router - ): + def check_batch_cost_instance(self, mock_proxy_logging_obj, mock_prisma_client, mock_llm_router): from litellm_enterprise.proxy.common_utils.check_batch_cost import ( CheckBatchCost, ) @@ -104,23 +112,15 @@ class TestCheckBatchCost: ) @pytest.mark.asyncio - async def test_cleanup_scoped_to_batch_file_purpose( - self, check_batch_cost_instance, mock_prisma_client - ): + async def test_cleanup_scoped_to_batch_file_purpose(self, check_batch_cost_instance, mock_prisma_client): """_cleanup_stale_managed_objects scopes its update to file_purpose='batch' only.""" - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 - ) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) # Return empty so the main poll loop exits immediately - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( - return_value=[] - ) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[]) await check_batch_cost_instance.check_batch_cost() - calls = ( - mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list - ) + calls = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list stale_call = calls[0] assert stale_call[1]["data"] == {"status": "stale_expired"} where = stale_call[1]["where"] @@ -129,9 +129,7 @@ class TestCheckBatchCost: assert "created_at" in where @pytest.mark.asyncio - async def test_startup_probe_confirms_batch_processed_support( - self, check_batch_cost_instance, mock_prisma_client - ): + async def test_startup_probe_confirms_batch_processed_support(self, check_batch_cost_instance, mock_prisma_client): mock_prisma_client.db.litellm_managedobjecttable.find_first = AsyncMock(return_value=None) await check_batch_cost_instance.confirm_batch_processed_support() @@ -142,9 +140,7 @@ class TestCheckBatchCost: assert check_batch_cost_instance._has_batch_processed_column is True @pytest.mark.asyncio - async def test_startup_probe_marks_column_absent( - self, check_batch_cost_instance, mock_prisma_client - ): + async def test_startup_probe_marks_column_absent(self, check_batch_cost_instance, mock_prisma_client): mock_prisma_client.db.litellm_managedobjecttable.find_first = AsyncMock( side_effect=Exception("column batch_processed does not exist") ) @@ -168,18 +164,12 @@ class TestCheckBatchCost: assert check_batch_cost_instance._has_batch_processed_column is True @pytest.mark.asyncio - async def test_find_many_uses_pagination_and_excludes_stale( - self, check_batch_cost_instance, mock_prisma_client - ): + async def test_find_many_uses_pagination_and_excludes_stale(self, check_batch_cost_instance, mock_prisma_client): """find_many is called with take, order, and all terminal statuses excluded.""" from litellm.constants import MAX_OBJECTS_PER_POLL_CYCLE - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 - ) - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( - return_value=[] - ) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[]) await check_batch_cost_instance.check_batch_cost() @@ -205,9 +195,7 @@ class TestCheckBatchCost: """Falls back to query without batch_processed when primary query raises.""" from litellm.constants import MAX_OBJECTS_PER_POLL_CYCLE - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 - ) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) # First find_many (primary query) raises with a schema error; second (fallback) returns empty mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( side_effect=[Exception("column batch_processed does not exist"), []] @@ -215,9 +203,7 @@ class TestCheckBatchCost: await check_batch_cost_instance.check_batch_cost() - calls = ( - mock_prisma_client.db.litellm_managedobjecttable.find_many.call_args_list - ) + calls = mock_prisma_client.db.litellm_managedobjecttable.find_many.call_args_list assert len(calls) == 2 fallback_where = calls[1][1]["where"] assert "batch_processed" not in fallback_where @@ -228,32 +214,20 @@ class TestCheckBatchCost: assert check_batch_cost_instance.batch_processed_support_confirmed is False @pytest.mark.asyncio - async def test_column_absence_cached_across_cycles( - self, check_batch_cost_instance, mock_prisma_client - ): + async def test_column_absence_cached_across_cycles(self, check_batch_cost_instance, mock_prisma_client): """After column absence is discovered, subsequent cycles skip the primary query entirely.""" from litellm.constants import MAX_OBJECTS_PER_POLL_CYCLE - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 - ) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) # Simulate column already known absent from a previous cycle check_batch_cost_instance._has_batch_processed_column = False - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( - return_value=[] - ) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[]) await check_batch_cost_instance.check_batch_cost() # Only one find_many call — the fallback directly, no primary query attempt - assert ( - mock_prisma_client.db.litellm_managedobjecttable.find_many.call_count == 1 - ) - fallback_where = ( - mock_prisma_client.db.litellm_managedobjecttable.find_many.call_args[1][ - "where" - ] - ) + assert mock_prisma_client.db.litellm_managedobjecttable.find_many.call_count == 1 + fallback_where = mock_prisma_client.db.litellm_managedobjecttable.find_many.call_args[1]["where"] assert "batch_processed" not in fallback_where @pytest.mark.asyncio @@ -267,13 +241,9 @@ class TestCheckBatchCost: """ from unittest.mock import patch - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 - ) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( - return_value=None - ) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) mock_job = MagicMock() mock_job.id = "job-fallback-1" @@ -282,22 +252,16 @@ class TestCheckBatchCost: # Simulate column already known absent (e.g. discovered on a previous cycle) check_batch_cost_instance._has_batch_processed_column = False - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( - return_value=[mock_job] - ) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) # Build a fake batch response whose status triggers the completion branch mock_response = MagicMock() mock_response.status = "completed" mock_response.output_file_id = "file-output-123" - mock_response.model_dump_json.return_value = ( - '{"id":"batch-1","status":"completed"}' - ) + mock_response.model_dump_json.return_value = '{"id":"batch-1","status":"completed"}' mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) - mock_llm_router.get_deployment_credentials_with_provider = MagicMock( - return_value={"api_key": "sk-test"} - ) + mock_llm_router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"}) mock_deployment = MagicMock() mock_deployment.litellm_params.custom_llm_provider = "openai" @@ -345,9 +309,7 @@ class TestCheckBatchCost: "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", return_value=("gpt-4", "openai", None, None), ), - patch( - "litellm.litellm_core_utils.litellm_logging.Logging" - ) as mock_logging_cls, + patch("litellm.litellm_core_utils.litellm_logging.Logging") as mock_logging_cls, ): mock_logging_obj = MagicMock() mock_logging_obj.async_success_handler = AsyncMock() @@ -356,15 +318,11 @@ class TestCheckBatchCost: await check_batch_cost_instance.check_batch_cost() # The update must have been called — this is the core assertion. - assert ( - mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1 - ), "Expected update() to be called exactly once for the completed job" - update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[ - 1 - ]["data"] - assert ( - "batch_processed" not in update_data - ), "update() must NOT include batch_processed when column is absent" + assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1, ( + "Expected update() to be called exactly once for the completed job" + ) + update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[1]["data"] + assert "batch_processed" not in update_data, "update() must NOT include batch_processed when column is absent" assert update_data["status"] == "complete" @pytest.mark.asyncio @@ -440,7 +398,9 @@ class TestCheckBatchCost: patch( "litellm.batches.batch_utils.calculate_batch_cost_and_usage", new_callable=AsyncMock, - return_value=_batch_cost_result(0.01, {"prompt_tokens": 10, "completion_tokens": 5}, ["claude-haiku-4-5"]), + return_value=_batch_cost_result( + 0.01, {"prompt_tokens": 10, "completion_tokens": 5}, ["claude-haiku-4-5"] + ), ), patch( "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", @@ -458,9 +418,9 @@ class TestCheckBatchCost: passed_kwargs = mock_afile_content.await_args[1] snapshot = passed_kwargs.get("_litellm_internal_model_credentials") assert snapshot is not None, "cost poller must pass the trusted credential snapshot" - assert isinstance( - snapshot, MappingProxyType - ), "snapshot must be a MappingProxyType; a plain dict is rejected by get_configured_s3_bucket_name" + assert isinstance(snapshot, MappingProxyType), ( + "snapshot must be a MappingProxyType; a plain dict is rejected by get_configured_s3_bucket_name" + ) assert snapshot["s3_bucket_name"] == "configured-batch-bucket" @pytest.mark.asyncio @@ -474,13 +434,9 @@ class TestCheckBatchCost: """ from unittest.mock import patch - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 - ) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( - return_value=None - ) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) mock_job = MagicMock() mock_job.id = "job-primary-1" @@ -488,21 +444,15 @@ class TestCheckBatchCost: mock_job.created_by = "user-1" assert check_batch_cost_instance._has_batch_processed_column is True - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( - return_value=[mock_job] - ) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) mock_response = MagicMock() mock_response.status = "completed" mock_response.output_file_id = "file-output-123" - mock_response.model_dump_json.return_value = ( - '{"id":"batch-1","status":"completed"}' - ) + mock_response.model_dump_json.return_value = '{"id":"batch-1","status":"completed"}' mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) - mock_llm_router.get_deployment_credentials_with_provider = MagicMock( - return_value={"api_key": "sk-test"} - ) + mock_llm_router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"}) mock_deployment = MagicMock() mock_deployment.litellm_params.custom_llm_provider = "openai" @@ -550,9 +500,7 @@ class TestCheckBatchCost: "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", return_value=("gpt-4", "openai", None, None), ), - patch( - "litellm.litellm_core_utils.litellm_logging.Logging" - ) as mock_logging_cls, + patch("litellm.litellm_core_utils.litellm_logging.Logging") as mock_logging_cls, ): mock_logging_obj = MagicMock() mock_logging_obj.async_success_handler = AsyncMock() @@ -560,15 +508,13 @@ class TestCheckBatchCost: await check_batch_cost_instance.check_batch_cost() - assert ( - mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1 - ), "Expected update() to be called exactly once for the completed job" - update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[ - 1 - ]["data"] - assert ( - update_data["batch_processed"] is True - ), "update() must include batch_processed=True when column is present" + assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1, ( + "Expected update() to be called exactly once for the completed job" + ) + update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[1]["data"] + assert update_data["batch_processed"] is True, ( + "update() must include batch_processed=True when column is present" + ) assert update_data["status"] == "complete" @pytest.mark.asyncio @@ -712,22 +658,16 @@ class TestCheckBatchCost: """ from unittest.mock import patch - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 - ) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( - return_value=None - ) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) mock_job = MagicMock() mock_job.id = "job-anthropic-1" mock_job.unified_object_id = "dW5pZmllZF9iYXRjaF9pZA==" mock_job.created_by = "user-1" - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( - return_value=[mock_job] - ) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) mock_response = MagicMock() mock_response.status = "completed" @@ -761,9 +701,9 @@ class TestCheckBatchCost: ): await check_batch_cost_instance.check_batch_cost() - assert ( - mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 0 - ), "a failed cost tracking attempt must not mark the job processed" + assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 0, ( + "a failed cost tracking attempt must not mark the job processed" + ) @pytest.mark.asyncio @pytest.mark.parametrize("terminal_status", ["failed", "expired", "cancelled"]) @@ -780,13 +720,9 @@ class TestCheckBatchCost: """ import base64 - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 - ) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( - return_value=None - ) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) mock_job = MagicMock() mock_job.id = "job-terminal-1" @@ -796,31 +732,25 @@ class TestCheckBatchCost: mock_job.created_by = "user-1" assert check_batch_cost_instance._has_batch_processed_column is True - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( - return_value=[mock_job] - ) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) mock_response = MagicMock() mock_response.status = terminal_status mock_response.output_file_id = None - mock_response.model_dump_json.return_value = ( - f'{{"id":"batch-1","status":"{terminal_status}"}}' - ) + mock_response.model_dump_json.return_value = f'{{"id":"batch-1","status":"{terminal_status}"}}' mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) await check_batch_cost_instance.check_batch_cost() - assert ( - mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1 - ), f"Expected update() to be called exactly once for a {terminal_status} job" - update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[ - 1 - ]["data"] + assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1, ( + f"Expected update() to be called exactly once for a {terminal_status} job" + ) + update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[1]["data"] assert update_data["status"] == terminal_status - assert ( - update_data["batch_processed"] is True - ), "terminal-status update() must set batch_processed=True so polling stops" + assert update_data["batch_processed"] is True, ( + "terminal-status update() must set batch_processed=True so polling stops" + ) @pytest.mark.asyncio @pytest.mark.parametrize("terminal_status", ["failed", "cancelled"]) @@ -855,13 +785,9 @@ class TestCheckBatchCost: f"litellm_proxy:application/octet-stream;unified_id,u-2;llm_output_file_id,{raw_error_file_id}".encode() ).decode() - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 - ) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( - return_value=None - ) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) input_file_row = MagicMock() input_file_row.unified_file_id = unified_input_file_id @@ -871,9 +797,7 @@ class TestCheckBatchCost: return input_file_row return None - mock_prisma_client.db.litellm_managedfiletable.find_first = AsyncMock( - side_effect=find_managed_file - ) + mock_prisma_client.db.litellm_managedfiletable.find_first = AsyncMock(side_effect=find_managed_file) mock_job = MagicMock() mock_job.id = "job-terminal-mint-1" @@ -882,9 +806,7 @@ class TestCheckBatchCost: mock_job.team_id = "team-1" check_batch_cost_instance._has_batch_processed_column = True - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( - return_value=[mock_job] - ) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) response = LiteLLMBatch( id="batch-456", @@ -902,9 +824,7 @@ class TestCheckBatchCost: mock_hook = MagicMock() mock_hook.get_unified_output_file_id.side_effect = [unified_error_file_id] mock_hook.store_unified_file_id = AsyncMock() - check_batch_cost_instance.proxy_logging_obj.get_proxy_hook.return_value = ( - mock_hook - ) + check_batch_cost_instance.proxy_logging_obj.get_proxy_hook.return_value = mock_hook await check_batch_cost_instance.check_batch_cost() @@ -955,13 +875,9 @@ class TestCheckBatchCost: import base64 from unittest.mock import patch - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 - ) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( - return_value=None - ) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) mock_job = MagicMock() mock_job.id = "job-completed-no-output-1" @@ -971,24 +887,18 @@ class TestCheckBatchCost: mock_job.created_by = "user-1" assert check_batch_cost_instance._has_batch_processed_column is True - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( - return_value=[mock_job] - ) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) mock_response = MagicMock() mock_response.status = completed_status mock_response.output_file_id = None mock_response.error_file_id = "file-error-123" - mock_response.model_dump_json.return_value = ( - f'{{"id":"batch-1","status":"{completed_status}"}}' - ) + mock_response.model_dump_json.return_value = f'{{"id":"batch-1","status":"{completed_status}"}}' mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) # Billing reads credentials off the router; if it is touched we billed a batch # that has no output, which is the behaviour this test guards against. - mock_llm_router.get_deployment_credentials_with_provider = MagicMock( - return_value={"api_key": "sk-test"} - ) + mock_llm_router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"}) with patch( "litellm.files.main.afile_content", @@ -996,22 +906,18 @@ class TestCheckBatchCost: ) as mock_afile_content: await check_batch_cost_instance.check_batch_cost() - assert ( - mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1 - ), "a completed batch with no output file must be marked processed exactly once" - update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[ - 1 - ]["data"] + assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1, ( + "a completed batch with no output file must be marked processed exactly once" + ) + update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[1]["data"] assert update_data["status"] == completed_status - assert ( - update_data["batch_processed"] is True - ), "completed-without-output update() must set batch_processed=True so polling stops" - assert ( - mock_afile_content.await_count == 0 - ), "a batch with no output file must not be billed" - assert ( - mock_llm_router.get_deployment_credentials_with_provider.call_count == 0 - ), "a batch with no output file must not enter the cost-tracking path" + assert update_data["batch_processed"] is True, ( + "completed-without-output update() must set batch_processed=True so polling stops" + ) + assert mock_afile_content.await_count == 0, "a batch with no output file must not be billed" + assert mock_llm_router.get_deployment_credentials_with_provider.call_count == 0, ( + "a batch with no output file must not enter the cost-tracking path" + ) @pytest.mark.asyncio async def test_non_terminal_status_left_unprocessed( @@ -1022,9 +928,7 @@ class TestCheckBatchCost: """ from unittest.mock import patch - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 - ) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() mock_job = MagicMock() @@ -1032,9 +936,7 @@ class TestCheckBatchCost: mock_job.unified_object_id = "dW5pZmllZF9iYXRjaF9pZA==" mock_job.created_by = "user-1" - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( - return_value=[mock_job] - ) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) mock_response = MagicMock() mock_response.status = "in_progress" @@ -1060,9 +962,9 @@ class TestCheckBatchCost: ): await check_batch_cost_instance.check_batch_cost() - assert ( - mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 0 - ), "a non-terminal batch must not be written back (would stop polling prematurely)" + assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 0, ( + "a non-terminal batch must not be written back (would stop polling prematurely)" + ) @pytest.mark.asyncio @pytest.mark.parametrize("terminal_status", ["expired", "cancelled", "failed"]) @@ -1079,13 +981,9 @@ class TestCheckBatchCost: """ from unittest.mock import patch - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 - ) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( - return_value=None - ) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) mock_job = MagicMock() mock_job.id = "job-terminal-with-output-1" @@ -1093,21 +991,15 @@ class TestCheckBatchCost: mock_job.created_by = "user-1" assert check_batch_cost_instance._has_batch_processed_column is True - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( - return_value=[mock_job] - ) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) mock_response = MagicMock() mock_response.status = terminal_status mock_response.output_file_id = "file-output-123" - mock_response.model_dump_json.return_value = ( - f'{{"id":"batch-1","status":"{terminal_status}"}}' - ) + mock_response.model_dump_json.return_value = f'{{"id":"batch-1","status":"{terminal_status}"}}' mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) - mock_llm_router.get_deployment_credentials_with_provider = MagicMock( - return_value={"api_key": "sk-test"} - ) + mock_llm_router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"}) mock_deployment = MagicMock() mock_deployment.litellm_params.custom_llm_provider = "openai" @@ -1155,9 +1047,7 @@ class TestCheckBatchCost: "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", return_value=("gpt-4", "openai", None, None), ), - patch( - "litellm.litellm_core_utils.litellm_logging.Logging" - ) as mock_logging_cls, + patch("litellm.litellm_core_utils.litellm_logging.Logging") as mock_logging_cls, ): mock_logging_obj = MagicMock() mock_logging_obj.async_success_handler = AsyncMock() @@ -1165,20 +1055,16 @@ class TestCheckBatchCost: await check_batch_cost_instance.check_batch_cost() - assert ( - mock_afile_content.await_count == 1 - ), f"{terminal_status} batch with an output file must fetch results and be billed" - mock_logging_obj.async_success_handler.assert_awaited_once() - assert ( - mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1 + assert mock_afile_content.await_count == 1, ( + f"{terminal_status} batch with an output file must fetch results and be billed" ) - update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[ - 1 - ]["data"] + mock_logging_obj.async_success_handler.assert_awaited_once() + assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1 + update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[1]["data"] assert update_data["batch_processed"] is True - assert ( - update_data["status"] == terminal_status - ), f"billed {terminal_status} batch must keep its real terminal status in the DB" + assert update_data["status"] == terminal_status, ( + f"billed {terminal_status} batch must keep its real terminal status in the DB" + ) @pytest.mark.asyncio async def test_terminal_batch_with_missing_output_file_is_retired_unbilled( @@ -1195,13 +1081,9 @@ class TestCheckBatchCost: from litellm.exceptions import NotFoundError - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 - ) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( - return_value=None - ) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) mock_job = MagicMock() mock_job.id = "job-output-gone-1" @@ -1211,23 +1093,17 @@ class TestCheckBatchCost: mock_job.created_by = "user-1" assert check_batch_cost_instance._has_batch_processed_column is True - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( - return_value=[mock_job] - ) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) missing_output_file_id = "gs://batch-out/job-1/predictions.jsonl" mock_response = MagicMock() mock_response.status = "failed" mock_response.output_file_id = missing_output_file_id mock_response.error_file_id = None - mock_response.model_dump_json.return_value = ( - '{"id":"batch-1","status":"failed"}' - ) + mock_response.model_dump_json.return_value = '{"id":"batch-1","status":"failed"}' mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) - mock_llm_router.get_deployment_credentials_with_provider = MagicMock( - return_value={"api_key": "sk-test"} - ) + mock_llm_router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"}) with ( patch( @@ -1248,12 +1124,10 @@ class TestCheckBatchCost: assert mock_afile_content.await_count == 1 mock_calculate.assert_not_awaited() - assert ( - mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1 - ), "a terminal batch with a 404ing output file must be retired, not retried forever" - update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[ - 1 - ]["data"] + assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1, ( + "a terminal batch with a 404ing output file must be retired, not retried forever" + ) + update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[1]["data"] assert update_data["status"] == "failed" assert update_data["batch_processed"] is True @@ -1266,13 +1140,9 @@ class TestCheckBatchCost: Without this, GET /batches/{id} returns a raw file ID that cannot be routed through the proxy, causing API_KEY errors when clients call GET /files/{id}/content. """ - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=0 - ) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( - return_value=None - ) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) mock_job = MagicMock() mock_job.id = "job-raw-file-1" @@ -1281,9 +1151,7 @@ class TestCheckBatchCost: mock_job.team_id = None check_batch_cost_instance._has_batch_processed_column = True - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( - return_value=[mock_job] - ) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) raw_output_file_id = "file-batch-output-abc123" raw_error_file_id = "file-batch-error-xyz456" @@ -1294,14 +1162,10 @@ class TestCheckBatchCost: mock_response.status = "completed" mock_response.output_file_id = raw_output_file_id mock_response.error_file_id = raw_error_file_id - mock_response.model_dump_json.return_value = ( - '{"id":"batch-1","status":"completed"}' - ) + mock_response.model_dump_json.return_value = '{"id":"batch-1","status":"completed"}' mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) - mock_llm_router.get_deployment_credentials_with_provider = MagicMock( - return_value={"api_key": "sk-test"} - ) + mock_llm_router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"}) mock_deployment = MagicMock() mock_deployment.litellm_params.custom_llm_provider = "azure" @@ -1316,9 +1180,7 @@ class TestCheckBatchCost: fake_managed_error_id, ] mock_hook.store_unified_file_id = AsyncMock() - check_batch_cost_instance.proxy_logging_obj.get_proxy_hook.return_value = ( - mock_hook - ) + check_batch_cost_instance.proxy_logging_obj.get_proxy_hook.return_value = mock_hook mock_file_content = MagicMock() mock_file_content.content = b'{"id":"req-1"}' @@ -1361,9 +1223,7 @@ class TestCheckBatchCost: "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", return_value=("gpt-5-mini", "azure", None, None), ), - patch( - "litellm.litellm_core_utils.litellm_logging.Logging" - ) as mock_logging_cls, + patch("litellm.litellm_core_utils.litellm_logging.Logging") as mock_logging_cls, ): mock_logging_obj = MagicMock() mock_logging_obj.async_success_handler = AsyncMock() @@ -1414,9 +1274,7 @@ class TestUnmanagedVertexRouting: def _job(self, file_object=None): job = MagicMock() job.unified_object_id = "8823717160934178816" - job.file_object = ( - file_object if file_object is not None else _unmanaged_vertex_file_object() - ) + job.file_object = file_object if file_object is not None else _unmanaged_vertex_file_object() return job def test_flag_off_skips_unmanaged_id_unchanged(self): @@ -1454,9 +1312,7 @@ class TestUnmanagedVertexRouting: assert result == ("deploy-1", "8823717160934178816") # bare model name (trailing GCS segment), not the full publishers/.. path - router.resolve_model_name_from_model_id.assert_called_once_with( - "gemini-2.5-flash" - ) + router.resolve_model_name_from_model_id.assert_called_once_with("gemini-2.5-flash") router.get_model_ids.assert_called_once_with(model_name="gemini-2.5-flash") def test_flag_on_skips_non_vertex_deployment_sharing_model_group(self): @@ -1476,9 +1332,7 @@ class TestUnmanagedVertexRouting: result = instance._resolve_job_routing(self._job(), prom) assert result is None - prom.record_check_batch_cost_error.assert_called_once_with( - "unmanaged_no_matching_deployment" - ) + prom.record_check_batch_cost_error.assert_called_once_with("unmanaged_no_matching_deployment") def test_flag_on_uses_later_vertex_deployment_with_matching_suffix(self): router = MagicMock() @@ -1526,9 +1380,7 @@ class TestUnmanagedVertexRouting: result = instance._resolve_job_routing(self._job(), prom) assert result is None - prom.record_check_batch_cost_error.assert_called_once_with( - "unmanaged_no_matching_deployment" - ) + prom.record_check_batch_cost_error.assert_called_once_with("unmanaged_no_matching_deployment") def test_flag_on_non_gcs_input_is_not_unmanaged_vertex(self): """Flag on, but input_file_id is not a gs:// publishers path: treat as unroutable, @@ -1536,9 +1388,7 @@ class TestUnmanagedVertexRouting: router = MagicMock() instance = self._instance(track_unmanaged=True, router=router) prom = MagicMock() - job = self._job( - file_object=_unmanaged_vertex_file_object(input_file_id="file-abc-123") - ) + job = self._job(file_object=_unmanaged_vertex_file_object(input_file_id="file-abc-123")) with patch(_IS_B64, return_value=False): result = instance._resolve_job_routing(job, prom) @@ -1562,9 +1412,7 @@ class TestUnmanagedVertexRouting: mock_response.error_file_id = None mock_response.completed_at = None mock_response.created_at = None - mock_response.model_dump_json.return_value = ( - '{"id":"8823717160934178816","status":"completed"}' - ) + mock_response.model_dump_json.return_value = '{"id":"8823717160934178816","status":"completed"}' router.aretrieve_batch = AsyncMock(return_value=mock_response) router.get_deployment_credentials_with_provider = MagicMock( return_value={"vertex_project": "p", "vertex_location": "us-central1"} @@ -1586,9 +1434,7 @@ class TestUnmanagedVertexRouting: prisma.db.litellm_managedobjecttable = MagicMock() prisma.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) prisma.db.litellm_managedobjecttable.update = AsyncMock() - prisma.db.litellm_managedobjecttable.find_many = AsyncMock( - return_value=[self._job()] - ) + prisma.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[self._job()]) prisma.db.litellm_usertable = MagicMock() prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=None) @@ -1619,9 +1465,7 @@ class TestUnmanagedVertexRouting: "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", return_value=("gemini-2.5-flash", "vertex_ai", None, None), ), - patch( - "litellm.litellm_core_utils.litellm_logging.Logging" - ) as mock_logging_cls, + patch("litellm.litellm_core_utils.litellm_logging.Logging") as mock_logging_cls, ): mock_logging_obj = MagicMock() mock_logging_obj.async_success_handler = AsyncMock() @@ -1662,9 +1506,7 @@ class TestUnmanagedBedrockRouting: def _job(self, file_object=None): job = MagicMock() job.unified_object_id = self._ARN - job.file_object = ( - file_object if file_object is not None else _unmanaged_bedrock_file_object() - ) + job.file_object = file_object if file_object is not None else _unmanaged_bedrock_file_object() return job def _bedrock_deployment(self): @@ -1719,9 +1561,7 @@ class TestUnmanagedBedrockRouting: result = instance._resolve_job_routing(self._job(), prom) assert result is None - prom.record_check_batch_cost_error.assert_called_once_with( - "unmanaged_no_matching_deployment" - ) + prom.record_check_batch_cost_error.assert_called_once_with("unmanaged_no_matching_deployment") def test_flag_on_matches_deployment_despite_colon_dash_mismatch(self): """The S3 object key has ':' replaced with '-' (e.g. 'v1-0'), but the configured @@ -1759,9 +1599,7 @@ class TestUnmanagedBedrockRouting: result = instance._resolve_job_routing(self._job(), prom) assert result is None - prom.record_check_batch_cost_error.assert_called_once_with( - "unmanaged_no_matching_deployment" - ) + prom.record_check_batch_cost_error.assert_called_once_with("unmanaged_no_matching_deployment") def test_flag_on_non_s3_input_is_not_unmanaged_bedrock(self): """Flag on, but input_file_id is not a litellm-bedrock-files- s3:// key: treat as @@ -1769,9 +1607,7 @@ class TestUnmanagedBedrockRouting: router = MagicMock() instance = self._instance(track_unmanaged=True, router=router) prom = MagicMock() - job = self._job( - file_object=_unmanaged_bedrock_file_object(input_file_id="file-abc-123") - ) + job = self._job(file_object=_unmanaged_bedrock_file_object(input_file_id="file-abc-123")) with patch(_IS_B64, return_value=False): result = instance._resolve_job_routing(job, prom) @@ -1794,13 +1630,9 @@ class TestUnmanagedBedrockRouting: mock_response.error_file_id = None mock_response.completed_at = None mock_response.created_at = None - mock_response.model_dump_json.return_value = ( - f'{{"id":"{self._ARN}","status":"completed"}}' - ) + mock_response.model_dump_json.return_value = f'{{"id":"{self._ARN}","status":"completed"}}' router.aretrieve_batch = AsyncMock(return_value=mock_response) - router.get_deployment_credentials_with_provider = MagicMock( - return_value={"aws_region_name": "us-east-1"} - ) + router.get_deployment_credentials_with_provider = MagicMock(return_value={"aws_region_name": "us-east-1"}) deployment = self._bedrock_deployment() deployment.model_name = "claude-sonnet-4" @@ -1816,9 +1648,7 @@ class TestUnmanagedBedrockRouting: prisma.db.litellm_managedobjecttable = MagicMock() prisma.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) prisma.db.litellm_managedobjecttable.update = AsyncMock() - prisma.db.litellm_managedobjecttable.find_many = AsyncMock( - return_value=[self._job()] - ) + prisma.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[self._job()]) prisma.db.litellm_usertable = MagicMock() prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=None) @@ -1849,9 +1679,7 @@ class TestUnmanagedBedrockRouting: "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", return_value=("claude-sonnet-4", "bedrock", None, None), ), - patch( - "litellm.litellm_core_utils.litellm_logging.Logging" - ) as mock_logging_cls, + patch("litellm.litellm_core_utils.litellm_logging.Logging") as mock_logging_cls, ): mock_logging_obj = MagicMock() mock_logging_obj.async_success_handler = AsyncMock() @@ -1975,9 +1803,7 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: ) router = MagicMock() - router.get_deployment_credentials_with_provider = MagicMock( - return_value={"api_key": "sk-test"} - ) + router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"}) deployment = MagicMock() deployment.litellm_params.custom_llm_provider = "azure" deployment.litellm_params.model = "azure/gpt-5.5" @@ -1986,8 +1812,8 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: router.get_deployment = MagicMock(return_value=deployment) hook = MagicMock() - hook.get_unified_output_file_id = ( - lambda output_file_id, model_id, model_name: _PROXY_LiteLLMManagedFiles.get_unified_output_file_id( + hook.get_unified_output_file_id = lambda output_file_id, model_id, model_name: ( + _PROXY_LiteLLMManagedFiles.get_unified_output_file_id( None, output_file_id=output_file_id, model_id=model_id, model_name=model_name ) ) @@ -2056,9 +1882,7 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: get_models_from_unified_file_id, ) - output_file_id = await self._run( - self._job(self._managed_input_file_id(self._PUBLIC_MODEL_GROUP)) - ) + output_file_id = await self._run(self._job(self._managed_input_file_id(self._PUBLIC_MODEL_GROUP))) decoded = _is_base64_encoded_unified_file_id(output_file_id) assert get_models_from_unified_file_id(decoded) == [self._PUBLIC_MODEL_GROUP] @@ -2072,9 +1896,7 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: _extract_models_from_managed_resource_id, ) - output_file_id = await self._run( - self._job(self._managed_input_file_id(self._PUBLIC_MODEL_GROUP)) - ) + output_file_id = await self._run(self._job(self._managed_input_file_id(self._PUBLIC_MODEL_GROUP))) models = _extract_models_from_managed_resource_id(output_file_id, "file_id", None) assert models == [self._PUBLIC_MODEL_GROUP] @@ -2082,9 +1904,7 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: await can_key_call_model( model=models[0], llm_model_list=None, - valid_token=UserAPIKeyAuth( - api_key="sk-test", models=[self._PUBLIC_MODEL_GROUP] - ), + valid_token=UserAPIKeyAuth(api_key="sk-test", models=[self._PUBLIC_MODEL_GROUP]), llm_router=None, ) is True @@ -2101,6 +1921,8 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: decoded = _is_base64_encoded_unified_file_id(output_file_id) assert get_models_from_unified_file_id(decoded) == [self._PUBLIC_MODEL_GROUP] + + class TestBatchCostAttribution: """CheckBatchCost rebuilds the creator's spend metadata from the managed-object row so the batch-cost log is attributed like a non-batch request.""" @@ -2196,9 +2018,7 @@ class TestBatchCostAttribution: """An alias lookup failure must not lose the spend row; the key hash and team still attribute it.""" instance = self._instance() - instance.prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( - side_effect=Exception("db down") - ) + instance.prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(side_effect=Exception("db down")) metadata = await instance._build_creator_attribution_metadata(self._job(), "batch-1") @@ -2294,9 +2114,7 @@ class TestPollPageStarvation: async def test_unified_id_without_model_id_is_retired(self): """A unified id that decodes but carries no model_id is unroutable no matter what the config says, so it must leave the poll page instead of being retried forever.""" - prisma = self._prisma( - [self._job("job-no-model", self._encode("litellm_proxy;llm_batch_id:poison-no-model"))] - ) + prisma = self._prisma([self._job("job-no-model", self._encode("litellm_proxy;llm_batch_id:poison-no-model"))]) llm_router = MagicMock() llm_router.aretrieve_batch = AsyncMock() @@ -2334,9 +2152,7 @@ class TestPollPageStarvation: await self._instance(prisma, llm_router).check_batch_cost() prisma.db.litellm_managedobjecttable.update.assert_awaited_once() - assert prisma.db.litellm_managedobjecttable.update.call_args[1]["data"] == { - "batch_processed": True - } + assert prisma.db.litellm_managedobjecttable.update.call_args[1]["data"] == {"batch_processed": True} @pytest.mark.asyncio async def test_provider_404_with_deployment_gone_keeps_job(self): @@ -2389,17 +2205,13 @@ class TestPollPageStarvation: async def test_retirement_falls_back_to_status_without_batch_processed_column(self): """Older schemas have no batch_processed column, so the only way to stop selecting the row is the status filter the poll query already applies.""" - prisma = self._prisma( - [self._job("job-legacy", self._encode("litellm_proxy;llm_batch_id:poison-no-model"))] - ) + prisma = self._prisma([self._job("job-legacy", self._encode("litellm_proxy;llm_batch_id:poison-no-model"))]) instance = self._instance(prisma, MagicMock()) instance._has_batch_processed_column = False await instance.check_batch_cost() - assert prisma.db.litellm_managedobjecttable.update.call_args[1]["data"] == { - "status": "stale_expired" - } + assert prisma.db.litellm_managedobjecttable.update.call_args[1]["data"] == {"status": "stale_expired"} @pytest.mark.asyncio async def test_stale_cleanup_gives_up_on_never_costed_completed_rows(self): @@ -2454,14 +2266,11 @@ class TestPollPageStarvation: await self._instance(prisma, llm_router).check_batch_cost() - retired = [ - call[1]["where"]["id"] - for call in prisma.db.litellm_managedobjecttable.update.call_args_list - ] + retired = [call[1]["where"]["id"] for call in prisma.db.litellm_managedobjecttable.update.call_args_list] assert retired == ["job-no-model", "job-gone"] - assert ( - llm_router.aretrieve_batch.await_args_list[-1][1]["batch_id"] == "batch_live" - ), "the newer healthy batch must still be polled in the same cycle" + assert llm_router.aretrieve_batch.await_args_list[-1][1]["batch_id"] == "batch_live", ( + "the newer healthy batch must still be polled in the same cycle" + ) @pytest.mark.asyncio async def test_404_that_does_not_name_the_batch_keeps_job_for_retry(self): From 3bfeaa78ddd1bb7cae9ab6576e21a699a466b642 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Tue, 18 Aug 2026 19:39:28 -0400 Subject: [PATCH 21/68] revert: undo accidental whole-file ruff-format of test_check_batch_cost.py The prior commit ran ruff format on the whole file to type the _batch_cost_result helper, reflowing hundreds of unrelated pre-existing lines that were never ruff-format-clean to begin with (confirmed at the PR's base commit, before any of these changes). CI's ruff-format gate only checks litellm/**/*.py, not tests/, so this reformatting served no CI purpose and only bloated the diff. Restores everything except the intended TYPE_CHECKING import and _batch_cost_result annotations. --- .../proxy_unit_tests/test_check_batch_cost.py | 507 ++++++++++++------ 1 file changed, 354 insertions(+), 153 deletions(-) diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index ff2dce498f0..cd7f28007af 100644 --- a/tests/proxy_unit_tests/test_check_batch_cost.py +++ b/tests/proxy_unit_tests/test_check_batch_cost.py @@ -100,7 +100,9 @@ class TestCheckBatchCost: return MagicMock() @pytest.fixture - def check_batch_cost_instance(self, mock_proxy_logging_obj, mock_prisma_client, mock_llm_router): + def check_batch_cost_instance( + self, mock_proxy_logging_obj, mock_prisma_client, mock_llm_router + ): from litellm_enterprise.proxy.common_utils.check_batch_cost import ( CheckBatchCost, ) @@ -112,15 +114,23 @@ class TestCheckBatchCost: ) @pytest.mark.asyncio - async def test_cleanup_scoped_to_batch_file_purpose(self, check_batch_cost_instance, mock_prisma_client): + async def test_cleanup_scoped_to_batch_file_purpose( + self, check_batch_cost_instance, mock_prisma_client + ): """_cleanup_stale_managed_objects scopes its update to file_purpose='batch' only.""" - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( + return_value=0 + ) # Return empty so the main poll loop exits immediately - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[]) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[] + ) await check_batch_cost_instance.check_batch_cost() - calls = mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list + calls = ( + mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args_list + ) stale_call = calls[0] assert stale_call[1]["data"] == {"status": "stale_expired"} where = stale_call[1]["where"] @@ -129,7 +139,9 @@ class TestCheckBatchCost: assert "created_at" in where @pytest.mark.asyncio - async def test_startup_probe_confirms_batch_processed_support(self, check_batch_cost_instance, mock_prisma_client): + async def test_startup_probe_confirms_batch_processed_support( + self, check_batch_cost_instance, mock_prisma_client + ): mock_prisma_client.db.litellm_managedobjecttable.find_first = AsyncMock(return_value=None) await check_batch_cost_instance.confirm_batch_processed_support() @@ -140,7 +152,9 @@ class TestCheckBatchCost: assert check_batch_cost_instance._has_batch_processed_column is True @pytest.mark.asyncio - async def test_startup_probe_marks_column_absent(self, check_batch_cost_instance, mock_prisma_client): + async def test_startup_probe_marks_column_absent( + self, check_batch_cost_instance, mock_prisma_client + ): mock_prisma_client.db.litellm_managedobjecttable.find_first = AsyncMock( side_effect=Exception("column batch_processed does not exist") ) @@ -164,12 +178,18 @@ class TestCheckBatchCost: assert check_batch_cost_instance._has_batch_processed_column is True @pytest.mark.asyncio - async def test_find_many_uses_pagination_and_excludes_stale(self, check_batch_cost_instance, mock_prisma_client): + async def test_find_many_uses_pagination_and_excludes_stale( + self, check_batch_cost_instance, mock_prisma_client + ): """find_many is called with take, order, and all terminal statuses excluded.""" from litellm.constants import MAX_OBJECTS_PER_POLL_CYCLE - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[]) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( + return_value=0 + ) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[] + ) await check_batch_cost_instance.check_batch_cost() @@ -195,7 +215,9 @@ class TestCheckBatchCost: """Falls back to query without batch_processed when primary query raises.""" from litellm.constants import MAX_OBJECTS_PER_POLL_CYCLE - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( + return_value=0 + ) # First find_many (primary query) raises with a schema error; second (fallback) returns empty mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( side_effect=[Exception("column batch_processed does not exist"), []] @@ -203,7 +225,9 @@ class TestCheckBatchCost: await check_batch_cost_instance.check_batch_cost() - calls = mock_prisma_client.db.litellm_managedobjecttable.find_many.call_args_list + calls = ( + mock_prisma_client.db.litellm_managedobjecttable.find_many.call_args_list + ) assert len(calls) == 2 fallback_where = calls[1][1]["where"] assert "batch_processed" not in fallback_where @@ -214,20 +238,32 @@ class TestCheckBatchCost: assert check_batch_cost_instance.batch_processed_support_confirmed is False @pytest.mark.asyncio - async def test_column_absence_cached_across_cycles(self, check_batch_cost_instance, mock_prisma_client): + async def test_column_absence_cached_across_cycles( + self, check_batch_cost_instance, mock_prisma_client + ): """After column absence is discovered, subsequent cycles skip the primary query entirely.""" from litellm.constants import MAX_OBJECTS_PER_POLL_CYCLE - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( + return_value=0 + ) # Simulate column already known absent from a previous cycle check_batch_cost_instance._has_batch_processed_column = False - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[]) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[] + ) await check_batch_cost_instance.check_batch_cost() # Only one find_many call — the fallback directly, no primary query attempt - assert mock_prisma_client.db.litellm_managedobjecttable.find_many.call_count == 1 - fallback_where = mock_prisma_client.db.litellm_managedobjecttable.find_many.call_args[1]["where"] + assert ( + mock_prisma_client.db.litellm_managedobjecttable.find_many.call_count == 1 + ) + fallback_where = ( + mock_prisma_client.db.litellm_managedobjecttable.find_many.call_args[1][ + "where" + ] + ) assert "batch_processed" not in fallback_where @pytest.mark.asyncio @@ -241,9 +277,13 @@ class TestCheckBatchCost: """ from unittest.mock import patch - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( + return_value=0 + ) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=None + ) mock_job = MagicMock() mock_job.id = "job-fallback-1" @@ -252,16 +292,22 @@ class TestCheckBatchCost: # Simulate column already known absent (e.g. discovered on a previous cycle) check_batch_cost_instance._has_batch_processed_column = False - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[mock_job] + ) # Build a fake batch response whose status triggers the completion branch mock_response = MagicMock() mock_response.status = "completed" mock_response.output_file_id = "file-output-123" - mock_response.model_dump_json.return_value = '{"id":"batch-1","status":"completed"}' + mock_response.model_dump_json.return_value = ( + '{"id":"batch-1","status":"completed"}' + ) mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) - mock_llm_router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"}) + mock_llm_router.get_deployment_credentials_with_provider = MagicMock( + return_value={"api_key": "sk-test"} + ) mock_deployment = MagicMock() mock_deployment.litellm_params.custom_llm_provider = "openai" @@ -309,7 +355,9 @@ class TestCheckBatchCost: "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", return_value=("gpt-4", "openai", None, None), ), - patch("litellm.litellm_core_utils.litellm_logging.Logging") as mock_logging_cls, + patch( + "litellm.litellm_core_utils.litellm_logging.Logging" + ) as mock_logging_cls, ): mock_logging_obj = MagicMock() mock_logging_obj.async_success_handler = AsyncMock() @@ -318,11 +366,15 @@ class TestCheckBatchCost: await check_batch_cost_instance.check_batch_cost() # The update must have been called — this is the core assertion. - assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1, ( - "Expected update() to be called exactly once for the completed job" - ) - update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[1]["data"] - assert "batch_processed" not in update_data, "update() must NOT include batch_processed when column is absent" + assert ( + mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1 + ), "Expected update() to be called exactly once for the completed job" + update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[ + 1 + ]["data"] + assert ( + "batch_processed" not in update_data + ), "update() must NOT include batch_processed when column is absent" assert update_data["status"] == "complete" @pytest.mark.asyncio @@ -398,9 +450,7 @@ class TestCheckBatchCost: patch( "litellm.batches.batch_utils.calculate_batch_cost_and_usage", new_callable=AsyncMock, - return_value=_batch_cost_result( - 0.01, {"prompt_tokens": 10, "completion_tokens": 5}, ["claude-haiku-4-5"] - ), + return_value=_batch_cost_result(0.01, {"prompt_tokens": 10, "completion_tokens": 5}, ["claude-haiku-4-5"]), ), patch( "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", @@ -418,9 +468,9 @@ class TestCheckBatchCost: passed_kwargs = mock_afile_content.await_args[1] snapshot = passed_kwargs.get("_litellm_internal_model_credentials") assert snapshot is not None, "cost poller must pass the trusted credential snapshot" - assert isinstance(snapshot, MappingProxyType), ( - "snapshot must be a MappingProxyType; a plain dict is rejected by get_configured_s3_bucket_name" - ) + assert isinstance( + snapshot, MappingProxyType + ), "snapshot must be a MappingProxyType; a plain dict is rejected by get_configured_s3_bucket_name" assert snapshot["s3_bucket_name"] == "configured-batch-bucket" @pytest.mark.asyncio @@ -434,9 +484,13 @@ class TestCheckBatchCost: """ from unittest.mock import patch - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( + return_value=0 + ) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=None + ) mock_job = MagicMock() mock_job.id = "job-primary-1" @@ -444,15 +498,21 @@ class TestCheckBatchCost: mock_job.created_by = "user-1" assert check_batch_cost_instance._has_batch_processed_column is True - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[mock_job] + ) mock_response = MagicMock() mock_response.status = "completed" mock_response.output_file_id = "file-output-123" - mock_response.model_dump_json.return_value = '{"id":"batch-1","status":"completed"}' + mock_response.model_dump_json.return_value = ( + '{"id":"batch-1","status":"completed"}' + ) mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) - mock_llm_router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"}) + mock_llm_router.get_deployment_credentials_with_provider = MagicMock( + return_value={"api_key": "sk-test"} + ) mock_deployment = MagicMock() mock_deployment.litellm_params.custom_llm_provider = "openai" @@ -500,7 +560,9 @@ class TestCheckBatchCost: "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", return_value=("gpt-4", "openai", None, None), ), - patch("litellm.litellm_core_utils.litellm_logging.Logging") as mock_logging_cls, + patch( + "litellm.litellm_core_utils.litellm_logging.Logging" + ) as mock_logging_cls, ): mock_logging_obj = MagicMock() mock_logging_obj.async_success_handler = AsyncMock() @@ -508,13 +570,15 @@ class TestCheckBatchCost: await check_batch_cost_instance.check_batch_cost() - assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1, ( - "Expected update() to be called exactly once for the completed job" - ) - update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[1]["data"] - assert update_data["batch_processed"] is True, ( - "update() must include batch_processed=True when column is present" - ) + assert ( + mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1 + ), "Expected update() to be called exactly once for the completed job" + update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[ + 1 + ]["data"] + assert ( + update_data["batch_processed"] is True + ), "update() must include batch_processed=True when column is present" assert update_data["status"] == "complete" @pytest.mark.asyncio @@ -658,16 +722,22 @@ class TestCheckBatchCost: """ from unittest.mock import patch - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( + return_value=0 + ) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=None + ) mock_job = MagicMock() mock_job.id = "job-anthropic-1" mock_job.unified_object_id = "dW5pZmllZF9iYXRjaF9pZA==" mock_job.created_by = "user-1" - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[mock_job] + ) mock_response = MagicMock() mock_response.status = "completed" @@ -701,9 +771,9 @@ class TestCheckBatchCost: ): await check_batch_cost_instance.check_batch_cost() - assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 0, ( - "a failed cost tracking attempt must not mark the job processed" - ) + assert ( + mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 0 + ), "a failed cost tracking attempt must not mark the job processed" @pytest.mark.asyncio @pytest.mark.parametrize("terminal_status", ["failed", "expired", "cancelled"]) @@ -720,9 +790,13 @@ class TestCheckBatchCost: """ import base64 - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( + return_value=0 + ) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=None + ) mock_job = MagicMock() mock_job.id = "job-terminal-1" @@ -732,25 +806,31 @@ class TestCheckBatchCost: mock_job.created_by = "user-1" assert check_batch_cost_instance._has_batch_processed_column is True - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[mock_job] + ) mock_response = MagicMock() mock_response.status = terminal_status mock_response.output_file_id = None - mock_response.model_dump_json.return_value = f'{{"id":"batch-1","status":"{terminal_status}"}}' + mock_response.model_dump_json.return_value = ( + f'{{"id":"batch-1","status":"{terminal_status}"}}' + ) mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) await check_batch_cost_instance.check_batch_cost() - assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1, ( - f"Expected update() to be called exactly once for a {terminal_status} job" - ) - update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[1]["data"] + assert ( + mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1 + ), f"Expected update() to be called exactly once for a {terminal_status} job" + update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[ + 1 + ]["data"] assert update_data["status"] == terminal_status - assert update_data["batch_processed"] is True, ( - "terminal-status update() must set batch_processed=True so polling stops" - ) + assert ( + update_data["batch_processed"] is True + ), "terminal-status update() must set batch_processed=True so polling stops" @pytest.mark.asyncio @pytest.mark.parametrize("terminal_status", ["failed", "cancelled"]) @@ -785,9 +865,13 @@ class TestCheckBatchCost: f"litellm_proxy:application/octet-stream;unified_id,u-2;llm_output_file_id,{raw_error_file_id}".encode() ).decode() - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( + return_value=0 + ) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=None + ) input_file_row = MagicMock() input_file_row.unified_file_id = unified_input_file_id @@ -797,7 +881,9 @@ class TestCheckBatchCost: return input_file_row return None - mock_prisma_client.db.litellm_managedfiletable.find_first = AsyncMock(side_effect=find_managed_file) + mock_prisma_client.db.litellm_managedfiletable.find_first = AsyncMock( + side_effect=find_managed_file + ) mock_job = MagicMock() mock_job.id = "job-terminal-mint-1" @@ -806,7 +892,9 @@ class TestCheckBatchCost: mock_job.team_id = "team-1" check_batch_cost_instance._has_batch_processed_column = True - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[mock_job] + ) response = LiteLLMBatch( id="batch-456", @@ -824,7 +912,9 @@ class TestCheckBatchCost: mock_hook = MagicMock() mock_hook.get_unified_output_file_id.side_effect = [unified_error_file_id] mock_hook.store_unified_file_id = AsyncMock() - check_batch_cost_instance.proxy_logging_obj.get_proxy_hook.return_value = mock_hook + check_batch_cost_instance.proxy_logging_obj.get_proxy_hook.return_value = ( + mock_hook + ) await check_batch_cost_instance.check_batch_cost() @@ -875,9 +965,13 @@ class TestCheckBatchCost: import base64 from unittest.mock import patch - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( + return_value=0 + ) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=None + ) mock_job = MagicMock() mock_job.id = "job-completed-no-output-1" @@ -887,18 +981,24 @@ class TestCheckBatchCost: mock_job.created_by = "user-1" assert check_batch_cost_instance._has_batch_processed_column is True - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[mock_job] + ) mock_response = MagicMock() mock_response.status = completed_status mock_response.output_file_id = None mock_response.error_file_id = "file-error-123" - mock_response.model_dump_json.return_value = f'{{"id":"batch-1","status":"{completed_status}"}}' + mock_response.model_dump_json.return_value = ( + f'{{"id":"batch-1","status":"{completed_status}"}}' + ) mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) # Billing reads credentials off the router; if it is touched we billed a batch # that has no output, which is the behaviour this test guards against. - mock_llm_router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"}) + mock_llm_router.get_deployment_credentials_with_provider = MagicMock( + return_value={"api_key": "sk-test"} + ) with patch( "litellm.files.main.afile_content", @@ -906,18 +1006,22 @@ class TestCheckBatchCost: ) as mock_afile_content: await check_batch_cost_instance.check_batch_cost() - assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1, ( - "a completed batch with no output file must be marked processed exactly once" - ) - update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[1]["data"] + assert ( + mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1 + ), "a completed batch with no output file must be marked processed exactly once" + update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[ + 1 + ]["data"] assert update_data["status"] == completed_status - assert update_data["batch_processed"] is True, ( - "completed-without-output update() must set batch_processed=True so polling stops" - ) - assert mock_afile_content.await_count == 0, "a batch with no output file must not be billed" - assert mock_llm_router.get_deployment_credentials_with_provider.call_count == 0, ( - "a batch with no output file must not enter the cost-tracking path" - ) + assert ( + update_data["batch_processed"] is True + ), "completed-without-output update() must set batch_processed=True so polling stops" + assert ( + mock_afile_content.await_count == 0 + ), "a batch with no output file must not be billed" + assert ( + mock_llm_router.get_deployment_credentials_with_provider.call_count == 0 + ), "a batch with no output file must not enter the cost-tracking path" @pytest.mark.asyncio async def test_non_terminal_status_left_unprocessed( @@ -928,7 +1032,9 @@ class TestCheckBatchCost: """ from unittest.mock import patch - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( + return_value=0 + ) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() mock_job = MagicMock() @@ -936,7 +1042,9 @@ class TestCheckBatchCost: mock_job.unified_object_id = "dW5pZmllZF9iYXRjaF9pZA==" mock_job.created_by = "user-1" - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[mock_job] + ) mock_response = MagicMock() mock_response.status = "in_progress" @@ -962,9 +1070,9 @@ class TestCheckBatchCost: ): await check_batch_cost_instance.check_batch_cost() - assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 0, ( - "a non-terminal batch must not be written back (would stop polling prematurely)" - ) + assert ( + mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 0 + ), "a non-terminal batch must not be written back (would stop polling prematurely)" @pytest.mark.asyncio @pytest.mark.parametrize("terminal_status", ["expired", "cancelled", "failed"]) @@ -981,9 +1089,13 @@ class TestCheckBatchCost: """ from unittest.mock import patch - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( + return_value=0 + ) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=None + ) mock_job = MagicMock() mock_job.id = "job-terminal-with-output-1" @@ -991,15 +1103,21 @@ class TestCheckBatchCost: mock_job.created_by = "user-1" assert check_batch_cost_instance._has_batch_processed_column is True - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[mock_job] + ) mock_response = MagicMock() mock_response.status = terminal_status mock_response.output_file_id = "file-output-123" - mock_response.model_dump_json.return_value = f'{{"id":"batch-1","status":"{terminal_status}"}}' + mock_response.model_dump_json.return_value = ( + f'{{"id":"batch-1","status":"{terminal_status}"}}' + ) mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) - mock_llm_router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"}) + mock_llm_router.get_deployment_credentials_with_provider = MagicMock( + return_value={"api_key": "sk-test"} + ) mock_deployment = MagicMock() mock_deployment.litellm_params.custom_llm_provider = "openai" @@ -1047,7 +1165,9 @@ class TestCheckBatchCost: "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", return_value=("gpt-4", "openai", None, None), ), - patch("litellm.litellm_core_utils.litellm_logging.Logging") as mock_logging_cls, + patch( + "litellm.litellm_core_utils.litellm_logging.Logging" + ) as mock_logging_cls, ): mock_logging_obj = MagicMock() mock_logging_obj.async_success_handler = AsyncMock() @@ -1055,16 +1175,20 @@ class TestCheckBatchCost: await check_batch_cost_instance.check_batch_cost() - assert mock_afile_content.await_count == 1, ( - f"{terminal_status} batch with an output file must fetch results and be billed" - ) + assert ( + mock_afile_content.await_count == 1 + ), f"{terminal_status} batch with an output file must fetch results and be billed" mock_logging_obj.async_success_handler.assert_awaited_once() - assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1 - update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[1]["data"] - assert update_data["batch_processed"] is True - assert update_data["status"] == terminal_status, ( - f"billed {terminal_status} batch must keep its real terminal status in the DB" + assert ( + mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1 ) + update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[ + 1 + ]["data"] + assert update_data["batch_processed"] is True + assert ( + update_data["status"] == terminal_status + ), f"billed {terminal_status} batch must keep its real terminal status in the DB" @pytest.mark.asyncio async def test_terminal_batch_with_missing_output_file_is_retired_unbilled( @@ -1081,9 +1205,13 @@ class TestCheckBatchCost: from litellm.exceptions import NotFoundError - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( + return_value=0 + ) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=None + ) mock_job = MagicMock() mock_job.id = "job-output-gone-1" @@ -1093,17 +1221,23 @@ class TestCheckBatchCost: mock_job.created_by = "user-1" assert check_batch_cost_instance._has_batch_processed_column is True - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[mock_job] + ) missing_output_file_id = "gs://batch-out/job-1/predictions.jsonl" mock_response = MagicMock() mock_response.status = "failed" mock_response.output_file_id = missing_output_file_id mock_response.error_file_id = None - mock_response.model_dump_json.return_value = '{"id":"batch-1","status":"failed"}' + mock_response.model_dump_json.return_value = ( + '{"id":"batch-1","status":"failed"}' + ) mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) - mock_llm_router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"}) + mock_llm_router.get_deployment_credentials_with_provider = MagicMock( + return_value={"api_key": "sk-test"} + ) with ( patch( @@ -1124,10 +1258,12 @@ class TestCheckBatchCost: assert mock_afile_content.await_count == 1 mock_calculate.assert_not_awaited() - assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1, ( - "a terminal batch with a 404ing output file must be retired, not retried forever" - ) - update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[1]["data"] + assert ( + mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1 + ), "a terminal batch with a 404ing output file must be retired, not retried forever" + update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[ + 1 + ]["data"] assert update_data["status"] == "failed" assert update_data["batch_processed"] is True @@ -1140,9 +1276,13 @@ class TestCheckBatchCost: Without this, GET /batches/{id} returns a raw file ID that cannot be routed through the proxy, causing API_KEY errors when clients call GET /files/{id}/content. """ - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( + return_value=0 + ) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=None + ) mock_job = MagicMock() mock_job.id = "job-raw-file-1" @@ -1151,7 +1291,9 @@ class TestCheckBatchCost: mock_job.team_id = None check_batch_cost_instance._has_batch_processed_column = True - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[mock_job] + ) raw_output_file_id = "file-batch-output-abc123" raw_error_file_id = "file-batch-error-xyz456" @@ -1162,10 +1304,14 @@ class TestCheckBatchCost: mock_response.status = "completed" mock_response.output_file_id = raw_output_file_id mock_response.error_file_id = raw_error_file_id - mock_response.model_dump_json.return_value = '{"id":"batch-1","status":"completed"}' + mock_response.model_dump_json.return_value = ( + '{"id":"batch-1","status":"completed"}' + ) mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) - mock_llm_router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"}) + mock_llm_router.get_deployment_credentials_with_provider = MagicMock( + return_value={"api_key": "sk-test"} + ) mock_deployment = MagicMock() mock_deployment.litellm_params.custom_llm_provider = "azure" @@ -1180,7 +1326,9 @@ class TestCheckBatchCost: fake_managed_error_id, ] mock_hook.store_unified_file_id = AsyncMock() - check_batch_cost_instance.proxy_logging_obj.get_proxy_hook.return_value = mock_hook + check_batch_cost_instance.proxy_logging_obj.get_proxy_hook.return_value = ( + mock_hook + ) mock_file_content = MagicMock() mock_file_content.content = b'{"id":"req-1"}' @@ -1223,7 +1371,9 @@ class TestCheckBatchCost: "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", return_value=("gpt-5-mini", "azure", None, None), ), - patch("litellm.litellm_core_utils.litellm_logging.Logging") as mock_logging_cls, + patch( + "litellm.litellm_core_utils.litellm_logging.Logging" + ) as mock_logging_cls, ): mock_logging_obj = MagicMock() mock_logging_obj.async_success_handler = AsyncMock() @@ -1274,7 +1424,9 @@ class TestUnmanagedVertexRouting: def _job(self, file_object=None): job = MagicMock() job.unified_object_id = "8823717160934178816" - job.file_object = file_object if file_object is not None else _unmanaged_vertex_file_object() + job.file_object = ( + file_object if file_object is not None else _unmanaged_vertex_file_object() + ) return job def test_flag_off_skips_unmanaged_id_unchanged(self): @@ -1312,7 +1464,9 @@ class TestUnmanagedVertexRouting: assert result == ("deploy-1", "8823717160934178816") # bare model name (trailing GCS segment), not the full publishers/.. path - router.resolve_model_name_from_model_id.assert_called_once_with("gemini-2.5-flash") + router.resolve_model_name_from_model_id.assert_called_once_with( + "gemini-2.5-flash" + ) router.get_model_ids.assert_called_once_with(model_name="gemini-2.5-flash") def test_flag_on_skips_non_vertex_deployment_sharing_model_group(self): @@ -1332,7 +1486,9 @@ class TestUnmanagedVertexRouting: result = instance._resolve_job_routing(self._job(), prom) assert result is None - prom.record_check_batch_cost_error.assert_called_once_with("unmanaged_no_matching_deployment") + prom.record_check_batch_cost_error.assert_called_once_with( + "unmanaged_no_matching_deployment" + ) def test_flag_on_uses_later_vertex_deployment_with_matching_suffix(self): router = MagicMock() @@ -1380,7 +1536,9 @@ class TestUnmanagedVertexRouting: result = instance._resolve_job_routing(self._job(), prom) assert result is None - prom.record_check_batch_cost_error.assert_called_once_with("unmanaged_no_matching_deployment") + prom.record_check_batch_cost_error.assert_called_once_with( + "unmanaged_no_matching_deployment" + ) def test_flag_on_non_gcs_input_is_not_unmanaged_vertex(self): """Flag on, but input_file_id is not a gs:// publishers path: treat as unroutable, @@ -1388,7 +1546,9 @@ class TestUnmanagedVertexRouting: router = MagicMock() instance = self._instance(track_unmanaged=True, router=router) prom = MagicMock() - job = self._job(file_object=_unmanaged_vertex_file_object(input_file_id="file-abc-123")) + job = self._job( + file_object=_unmanaged_vertex_file_object(input_file_id="file-abc-123") + ) with patch(_IS_B64, return_value=False): result = instance._resolve_job_routing(job, prom) @@ -1412,7 +1572,9 @@ class TestUnmanagedVertexRouting: mock_response.error_file_id = None mock_response.completed_at = None mock_response.created_at = None - mock_response.model_dump_json.return_value = '{"id":"8823717160934178816","status":"completed"}' + mock_response.model_dump_json.return_value = ( + '{"id":"8823717160934178816","status":"completed"}' + ) router.aretrieve_batch = AsyncMock(return_value=mock_response) router.get_deployment_credentials_with_provider = MagicMock( return_value={"vertex_project": "p", "vertex_location": "us-central1"} @@ -1434,7 +1596,9 @@ class TestUnmanagedVertexRouting: prisma.db.litellm_managedobjecttable = MagicMock() prisma.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) prisma.db.litellm_managedobjecttable.update = AsyncMock() - prisma.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[self._job()]) + prisma.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[self._job()] + ) prisma.db.litellm_usertable = MagicMock() prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=None) @@ -1465,7 +1629,9 @@ class TestUnmanagedVertexRouting: "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", return_value=("gemini-2.5-flash", "vertex_ai", None, None), ), - patch("litellm.litellm_core_utils.litellm_logging.Logging") as mock_logging_cls, + patch( + "litellm.litellm_core_utils.litellm_logging.Logging" + ) as mock_logging_cls, ): mock_logging_obj = MagicMock() mock_logging_obj.async_success_handler = AsyncMock() @@ -1506,7 +1672,9 @@ class TestUnmanagedBedrockRouting: def _job(self, file_object=None): job = MagicMock() job.unified_object_id = self._ARN - job.file_object = file_object if file_object is not None else _unmanaged_bedrock_file_object() + job.file_object = ( + file_object if file_object is not None else _unmanaged_bedrock_file_object() + ) return job def _bedrock_deployment(self): @@ -1561,7 +1729,9 @@ class TestUnmanagedBedrockRouting: result = instance._resolve_job_routing(self._job(), prom) assert result is None - prom.record_check_batch_cost_error.assert_called_once_with("unmanaged_no_matching_deployment") + prom.record_check_batch_cost_error.assert_called_once_with( + "unmanaged_no_matching_deployment" + ) def test_flag_on_matches_deployment_despite_colon_dash_mismatch(self): """The S3 object key has ':' replaced with '-' (e.g. 'v1-0'), but the configured @@ -1599,7 +1769,9 @@ class TestUnmanagedBedrockRouting: result = instance._resolve_job_routing(self._job(), prom) assert result is None - prom.record_check_batch_cost_error.assert_called_once_with("unmanaged_no_matching_deployment") + prom.record_check_batch_cost_error.assert_called_once_with( + "unmanaged_no_matching_deployment" + ) def test_flag_on_non_s3_input_is_not_unmanaged_bedrock(self): """Flag on, but input_file_id is not a litellm-bedrock-files- s3:// key: treat as @@ -1607,7 +1779,9 @@ class TestUnmanagedBedrockRouting: router = MagicMock() instance = self._instance(track_unmanaged=True, router=router) prom = MagicMock() - job = self._job(file_object=_unmanaged_bedrock_file_object(input_file_id="file-abc-123")) + job = self._job( + file_object=_unmanaged_bedrock_file_object(input_file_id="file-abc-123") + ) with patch(_IS_B64, return_value=False): result = instance._resolve_job_routing(job, prom) @@ -1630,9 +1804,13 @@ class TestUnmanagedBedrockRouting: mock_response.error_file_id = None mock_response.completed_at = None mock_response.created_at = None - mock_response.model_dump_json.return_value = f'{{"id":"{self._ARN}","status":"completed"}}' + mock_response.model_dump_json.return_value = ( + f'{{"id":"{self._ARN}","status":"completed"}}' + ) router.aretrieve_batch = AsyncMock(return_value=mock_response) - router.get_deployment_credentials_with_provider = MagicMock(return_value={"aws_region_name": "us-east-1"}) + router.get_deployment_credentials_with_provider = MagicMock( + return_value={"aws_region_name": "us-east-1"} + ) deployment = self._bedrock_deployment() deployment.model_name = "claude-sonnet-4" @@ -1648,7 +1826,9 @@ class TestUnmanagedBedrockRouting: prisma.db.litellm_managedobjecttable = MagicMock() prisma.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0) prisma.db.litellm_managedobjecttable.update = AsyncMock() - prisma.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[self._job()]) + prisma.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[self._job()] + ) prisma.db.litellm_usertable = MagicMock() prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=None) @@ -1679,7 +1859,9 @@ class TestUnmanagedBedrockRouting: "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", return_value=("claude-sonnet-4", "bedrock", None, None), ), - patch("litellm.litellm_core_utils.litellm_logging.Logging") as mock_logging_cls, + patch( + "litellm.litellm_core_utils.litellm_logging.Logging" + ) as mock_logging_cls, ): mock_logging_obj = MagicMock() mock_logging_obj.async_success_handler = AsyncMock() @@ -1803,7 +1985,9 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: ) router = MagicMock() - router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"}) + router.get_deployment_credentials_with_provider = MagicMock( + return_value={"api_key": "sk-test"} + ) deployment = MagicMock() deployment.litellm_params.custom_llm_provider = "azure" deployment.litellm_params.model = "azure/gpt-5.5" @@ -1812,8 +1996,8 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: router.get_deployment = MagicMock(return_value=deployment) hook = MagicMock() - hook.get_unified_output_file_id = lambda output_file_id, model_id, model_name: ( - _PROXY_LiteLLMManagedFiles.get_unified_output_file_id( + hook.get_unified_output_file_id = ( + lambda output_file_id, model_id, model_name: _PROXY_LiteLLMManagedFiles.get_unified_output_file_id( None, output_file_id=output_file_id, model_id=model_id, model_name=model_name ) ) @@ -1882,7 +2066,9 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: get_models_from_unified_file_id, ) - output_file_id = await self._run(self._job(self._managed_input_file_id(self._PUBLIC_MODEL_GROUP))) + output_file_id = await self._run( + self._job(self._managed_input_file_id(self._PUBLIC_MODEL_GROUP)) + ) decoded = _is_base64_encoded_unified_file_id(output_file_id) assert get_models_from_unified_file_id(decoded) == [self._PUBLIC_MODEL_GROUP] @@ -1896,7 +2082,9 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: _extract_models_from_managed_resource_id, ) - output_file_id = await self._run(self._job(self._managed_input_file_id(self._PUBLIC_MODEL_GROUP))) + output_file_id = await self._run( + self._job(self._managed_input_file_id(self._PUBLIC_MODEL_GROUP)) + ) models = _extract_models_from_managed_resource_id(output_file_id, "file_id", None) assert models == [self._PUBLIC_MODEL_GROUP] @@ -1904,7 +2092,9 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: await can_key_call_model( model=models[0], llm_model_list=None, - valid_token=UserAPIKeyAuth(api_key="sk-test", models=[self._PUBLIC_MODEL_GROUP]), + valid_token=UserAPIKeyAuth( + api_key="sk-test", models=[self._PUBLIC_MODEL_GROUP] + ), llm_router=None, ) is True @@ -1921,8 +2111,6 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: decoded = _is_base64_encoded_unified_file_id(output_file_id) assert get_models_from_unified_file_id(decoded) == [self._PUBLIC_MODEL_GROUP] - - class TestBatchCostAttribution: """CheckBatchCost rebuilds the creator's spend metadata from the managed-object row so the batch-cost log is attributed like a non-batch request.""" @@ -2018,7 +2206,9 @@ class TestBatchCostAttribution: """An alias lookup failure must not lose the spend row; the key hash and team still attribute it.""" instance = self._instance() - instance.prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(side_effect=Exception("db down")) + instance.prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( + side_effect=Exception("db down") + ) metadata = await instance._build_creator_attribution_metadata(self._job(), "batch-1") @@ -2114,7 +2304,9 @@ class TestPollPageStarvation: async def test_unified_id_without_model_id_is_retired(self): """A unified id that decodes but carries no model_id is unroutable no matter what the config says, so it must leave the poll page instead of being retried forever.""" - prisma = self._prisma([self._job("job-no-model", self._encode("litellm_proxy;llm_batch_id:poison-no-model"))]) + prisma = self._prisma( + [self._job("job-no-model", self._encode("litellm_proxy;llm_batch_id:poison-no-model"))] + ) llm_router = MagicMock() llm_router.aretrieve_batch = AsyncMock() @@ -2152,7 +2344,9 @@ class TestPollPageStarvation: await self._instance(prisma, llm_router).check_batch_cost() prisma.db.litellm_managedobjecttable.update.assert_awaited_once() - assert prisma.db.litellm_managedobjecttable.update.call_args[1]["data"] == {"batch_processed": True} + assert prisma.db.litellm_managedobjecttable.update.call_args[1]["data"] == { + "batch_processed": True + } @pytest.mark.asyncio async def test_provider_404_with_deployment_gone_keeps_job(self): @@ -2205,13 +2399,17 @@ class TestPollPageStarvation: async def test_retirement_falls_back_to_status_without_batch_processed_column(self): """Older schemas have no batch_processed column, so the only way to stop selecting the row is the status filter the poll query already applies.""" - prisma = self._prisma([self._job("job-legacy", self._encode("litellm_proxy;llm_batch_id:poison-no-model"))]) + prisma = self._prisma( + [self._job("job-legacy", self._encode("litellm_proxy;llm_batch_id:poison-no-model"))] + ) instance = self._instance(prisma, MagicMock()) instance._has_batch_processed_column = False await instance.check_batch_cost() - assert prisma.db.litellm_managedobjecttable.update.call_args[1]["data"] == {"status": "stale_expired"} + assert prisma.db.litellm_managedobjecttable.update.call_args[1]["data"] == { + "status": "stale_expired" + } @pytest.mark.asyncio async def test_stale_cleanup_gives_up_on_never_costed_completed_rows(self): @@ -2266,11 +2464,14 @@ class TestPollPageStarvation: await self._instance(prisma, llm_router).check_batch_cost() - retired = [call[1]["where"]["id"] for call in prisma.db.litellm_managedobjecttable.update.call_args_list] + retired = [ + call[1]["where"]["id"] + for call in prisma.db.litellm_managedobjecttable.update.call_args_list + ] assert retired == ["job-no-model", "job-gone"] - assert llm_router.aretrieve_batch.await_args_list[-1][1]["batch_id"] == "batch_live", ( - "the newer healthy batch must still be polled in the same cycle" - ) + assert ( + llm_router.aretrieve_batch.await_args_list[-1][1]["batch_id"] == "batch_live" + ), "the newer healthy batch must still be polled in the same cycle" @pytest.mark.asyncio async def test_404_that_does_not_name_the_batch_keeps_job_for_retry(self): From 98447e00436443f26ac5538bb4ad646cc0b32f75 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Tue, 18 Aug 2026 19:58:48 -0400 Subject: [PATCH 22/68] fix(batches): reconcile merge with upstream batch pricing changes - Resolved merge conflicts in _handle_completed_batch and CheckBatchCost._track_completed_batch_cost, keeping BatchCostUsageResult while adopting upstream's model_info threading and improved deployment-pricing lookup - Fixed two upstream tests and one mock that still expected the old tuple(cost, usage, models) return shape - Suppressed the one new LIT002 violation from the empty-output-file BatchCostUsageResult literal --- litellm/batches/batch_utils.py | 2 +- .../test_litellm/batches/test_batch_utils.py | 20 +++++++++---------- .../test_litellm_logging.py | 12 +++++++++-- 3 files changed, 21 insertions(+), 13 deletions(-) diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index bae8e198223..d0874b5aff5 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -86,7 +86,7 @@ async def _handle_completed_batch( return BatchCostUsageResult( cost=0.0, usage=Usage(prompt_tokens=0, completion_tokens=0, total_tokens=0), - models=[], + 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 diff --git a/tests/test_litellm/batches/test_batch_utils.py b/tests/test_litellm/batches/test_batch_utils.py index 0adebf2506b..64b3df4180b 100644 --- a/tests/test_litellm/batches/test_batch_utils.py +++ b/tests/test_litellm/batches/test_batch_utils.py @@ -1501,24 +1501,24 @@ async def test_handle_completed_bedrock_batch_prices_from_deployment_model(monke monkeypatch.setattr(bu, "_fetch_batch_output_file_content", fake_fetch) - cost, usage, _ = await bu._handle_completed_batch( + result = await bu._handle_completed_batch( _batch("of"), custom_llm_provider="bedrock", model_name="bedrock/global.anthropic.claude-sonnet-4-6", ) - assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (1800, 1000, 2800) + assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (1800, 1000, 2800) # 3e-06 / 1.5e-05 on-demand, halved for batch. - assert cost == pytest.approx(1800 * 3e-06 / 2 + 1000 * 1.5e-05 / 2) + assert result.cost == pytest.approx(1800 * 3e-06 / 2 + 1000 * 1.5e-05 / 2) # The response model alone cannot price a bedrock batch: this is the $0 bug. - zero_cost, zero_usage, _ = await bu._handle_completed_batch( + zero_result = await bu._handle_completed_batch( _batch("of"), custom_llm_provider="bedrock", model_name=None, ) - assert zero_cost == 0.0 - assert zero_usage.total_tokens == 2800 + assert zero_result.cost == 0.0 + assert zero_result.usage.total_tokens == 2800 @pytest.mark.asyncio @@ -1531,7 +1531,7 @@ async def test_handle_completed_batch_honors_deployment_pricing(monkeypatch) -> monkeypatch.setattr(bu, "_fetch_batch_output_file_content", fake_fetch) - free_cost, _, _ = await bu._handle_completed_batch( + free_result = await bu._handle_completed_batch( _batch("of"), custom_llm_provider="vertex_ai", model_name="vertex_ai/gemini-2.5-flash", @@ -1542,15 +1542,15 @@ async def test_handle_completed_batch_honors_deployment_pricing(monkeypatch) -> "output_cost_per_token_batches": 0.0, }, ) - assert free_cost == 0.0 + assert free_result.cost == 0.0 - billed_cost, _, _ = await bu._handle_completed_batch( + billed_result = await bu._handle_completed_batch( _batch("of"), custom_llm_provider="vertex_ai", model_name="vertex_ai/gemini-2.5-flash", model_info=None, ) - assert billed_cost > 0.0 + assert billed_result.cost > 0.0 # =========================================================================== # diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 0d54680fa81..e46057fe11a 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -586,9 +586,17 @@ class TestRetrieveBatchCostPassesModelIdentity: captured: dict[str, object] = {} - async def fake_handle_completed_batch(**kwargs: object) -> tuple[float, Usage, list[str]]: + from litellm.batches.batch_utils import BatchCostUsageResult + + async def fake_handle_completed_batch(**kwargs: object) -> BatchCostUsageResult: captured.update(kwargs) - return 1.25, Usage(prompt_tokens=1800, completion_tokens=1000, total_tokens=2800), ["m"] + return BatchCostUsageResult( + cost=1.25, + usage=Usage(prompt_tokens=1800, completion_tokens=1000, total_tokens=2800), + models=["m"], + successful_requests=1, + failed_requests=0, + ) monkeypatch.setattr(logging_module, "_handle_completed_batch", fake_handle_completed_batch) From d989f172e9d78ff63c5867f45850d5065b1ad7c8 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Tue, 18 Aug 2026 20:32:11 -0400 Subject: [PATCH 23/68] fix(batches): revert invalid Final on a loop-scoped local basedpyright rejects a Final variable assigned inside a loop body (reportGeneralTypeIssues); the type-discipline checker doesn't flag this line without Final either, so the annotation only bought a basedpyright budget regression. --- litellm/batches/batch_utils.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index d0874b5aff5..4e62c4a0130 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -170,9 +170,7 @@ def _classify_output_line_stats( custom_llm_provider=custom_llm_provider, call_type=CallTypes.aretrieve_batch.value, ) - reasoning_tokens: Final = ( - usage.completion_tokens_details.reasoning_tokens if usage.completion_tokens_details else None - ) + reasoning_tokens = usage.completion_tokens_details.reasoning_tokens if usage.completion_tokens_details else None yield _BatchOutputLineStats( cost=line_cost, prompt_tokens=usage.prompt_tokens, From 5a7edc9c77836f8f47634d8d47719bb44059a1fc Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Mon, 24 Aug 2026 19:26:52 -0400 Subject: [PATCH 24/68] fix(spend-tracking): make the new batch count keys writable SpendLogsMetadata is built by assigning each key in turn, so ReadOnly on the two new ones breached the basedpyright reportTypedDictNotRequiredAccess ceiling. Every sibling key in this TypedDict is writable for the same reason. --- litellm/proxy/_types.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index bf0dbc97482..377eb7dce0e 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3528,8 +3528,8 @@ class SpendLogsMetadata(TypedDict): status: StandardLoggingPayloadStatus proxy_server_request: str | None batch_models: list[str] | None - batch_successful_requests: ReadOnly[int | None] - batch_failed_requests: ReadOnly[int | None] + batch_successful_requests: int | None # writable-ok: built by assignment like every sibling key in this TypedDict + batch_failed_requests: int | None # writable-ok: built by assignment like every sibling key in this TypedDict error_information: StandardLoggingPayloadErrorInformation | None usage_object: dict | None model_map_information: StandardLoggingModelInformation | None From 4c2f0f3632c39ae8015f136b02eb002ee24ad248 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Mon, 24 Aug 2026 19:34:16 -0400 Subject: [PATCH 25/68] test(batches): cover decoding a model-encoded error file id The error file now resolves through _provider_output_file_id like the output file does. Sending the encoded id straight to the provider 404s, and the swallowed fetch failure would silently report zero failures. --- .../test_litellm/batches/test_batch_utils.py | 47 +++++++++++++++++++ 1 file changed, 47 insertions(+) diff --git a/tests/test_litellm/batches/test_batch_utils.py b/tests/test_litellm/batches/test_batch_utils.py index 1540a587349..c86c7c4df03 100644 --- a/tests/test_litellm/batches/test_batch_utils.py +++ b/tests/test_litellm/batches/test_batch_utils.py @@ -1164,6 +1164,53 @@ async def test_handle_completed_batch_counts_error_file_failures(monkeypatch): assert result.failed_requests == 1 +@pytest.mark.asyncio +async def test_handle_completed_batch_decodes_model_encoded_error_file_id(monkeypatch): + """A model-encoded error file id must be decoded to the raw provider id before + the fetch, exactly like the output file id. Sending the encoded id straight to + the provider 404s, and the swallowed fetch failure silently reports 0 failures.""" + import base64 + + from litellm.types.llms.openai import Batch + + provider_error_file_id = "file-real-error-id" + encoded_error_file_id = "file-" + base64.urlsafe_b64encode( + f"litellm:{provider_error_file_id};model,model-abc".encode() + ).decode().rstrip("=") + + requested_file_ids = [] + + async def fake_fetch(batch, custom_llm_provider, litellm_params=None): + return _vertex_jsonl([_success_row(model="gpt-4o", usage=_usage(10, 5))]) + + async def fake_afile_content(**kw): + requested_file_ids.append(kw["file_id"]) + return type("R", (), {"content": _vertex_jsonl([{"custom_id": "bad-1"}])})() + + 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(litellm, "completion_cost", lambda **kw: 0.0) + + batch = Batch( + id="b", + completion_window="24h", + created_at=1, + endpoint="/v1/chat/completions", + input_file_id="f", + object="batch", + status="completed", + output_file_id="of", + error_file_id=encoded_error_file_id, + ) + + result = await bu._handle_completed_batch(batch, custom_llm_provider="openai") + + assert requested_file_ids == [provider_error_file_id] + assert result.failed_requests == 1 + + @pytest.mark.asyncio async def test_handle_completed_batch_no_error_file_id_reports_zero_error_failures(monkeypatch): rows = [_success_row(model="gpt-4o", usage=_usage(10, 5))] From 3ebf09464a9fd81c4aa0f0fa8edde0016ed54104 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Tue, 25 Aug 2026 10:12:33 -0700 Subject: [PATCH 26/68] fix(proxy): give every `requests` call a timeout so a silent server cannot hang the caller `requests` has no default timeout, so a host that accepts the connection and never answers blocks the calling thread forever. The one on the request path is the HiddenLayer guardrail's `_get_jwt`. It runs synchronously inside `_call_hiddenlayer` whenever the hour-long JWT expires and the API answers 401, so a stalled auth host parked the worker's whole event loop, not just the guarded request. The other eight are the teams and users CLI clients, which pin the operator's terminal instead. `TeamsManagementClient` and `UsersManagementClient` now take the same `timeout: int = 30` their `HTTPClient` sibling already had, and `Client` threads its own timeout down to teams. `_poll_for_ready_data` already passed a timeout through a TypedDict that ruff could not see into; passing the argument directly retires both the TypedDict and the suppression it would have needed. Graduate S113 into ruff.toml so the next `requests` call without a timeout fails the lint step. --- litellm/proxy/client/cli/commands/auth.py | 10 +--- litellm/proxy/client/client.py | 2 +- litellm/proxy/client/teams.py | 10 ++-- litellm/proxy/client/users.py | 15 +++--- .../hiddenlayer/hiddenlayer.py | 7 ++- ruff.toml | 3 +- tests/test_litellm/proxy/client/conftest.py | 38 +++++++++++++++ tests/test_litellm/proxy/client/test_teams.py | 20 ++++++++ tests/test_litellm/proxy/client/test_users.py | 16 +++++++ .../guardrail_hooks/test_hiddenlayer.py | 48 +++++++++++++++++++ type-discipline-budget.json | 4 +- 11 files changed, 148 insertions(+), 25 deletions(-) create mode 100644 tests/test_litellm/proxy/client/conftest.py create mode 100644 tests/test_litellm/proxy/client/test_teams.py diff --git a/litellm/proxy/client/cli/commands/auth.py b/litellm/proxy/client/cli/commands/auth.py index 550b11311f5..2fad9f933c1 100644 --- a/litellm/proxy/client/cli/commands/auth.py +++ b/litellm/proxy/client/cli/commands/auth.py @@ -99,11 +99,6 @@ class CliPollData(TypedDict, total=False): team_id: str -class CliPollRequestKwargs(TypedDict, total=False): - timeout: int - headers: dict[str, str] - - class CliSsoStartData(TypedDict): login_id: str poll_secret: str @@ -518,10 +513,7 @@ def _poll_for_ready_data( ) -> CliPollData | None: for attempt in range(total_timeout // poll_interval): try: - request_kwargs: CliPollRequestKwargs = {"timeout": request_timeout} - if headers is not None: - request_kwargs["headers"] = headers - response = requests.get(url, **request_kwargs) + response = requests.get(url, headers=headers, timeout=request_timeout) if response.status_code == 200: data: CliPollData = response.json() status = data.get("status") diff --git a/litellm/proxy/client/client.py b/litellm/proxy/client/client.py index d71802e06c8..560523db189 100644 --- a/litellm/proxy/client/client.py +++ b/litellm/proxy/client/client.py @@ -38,4 +38,4 @@ class Client: self.chat = ChatClient(base_url=self._base_url, api_key=self._api_key) self.keys = KeysManagementClient(base_url=self._base_url, api_key=self._api_key) self.credentials = CredentialsManagementClient(base_url=self._base_url, api_key=self._api_key) - self.teams = TeamsManagementClient(base_url=self._base_url, api_key=self._api_key) + self.teams = TeamsManagementClient(base_url=self._base_url, api_key=self._api_key, timeout=timeout) diff --git a/litellm/proxy/client/teams.py b/litellm/proxy/client/teams.py index ef2ac53f9c4..105060e5ca9 100644 --- a/litellm/proxy/client/teams.py +++ b/litellm/proxy/client/teams.py @@ -11,16 +11,18 @@ from .exceptions import UnauthorizedError class TeamsManagementClient: """Client for managing teams in LiteLLM proxy.""" - def __init__(self, base_url: str, api_key: str | None = None): + def __init__(self, base_url: str, api_key: str | None = None, timeout: int = 30): """ Initialize the TeamsManagementClient. Args: base_url (str): The base URL of the LiteLLM proxy server (e.g., "http://localhost:4000") api_key (Optional[str]): API key for authentication. If provided, it will be sent as a Bearer token. + timeout (int): Request timeout in seconds (default: 30) """ self._base_url = base_url.rstrip("/") # Remove trailing slash if present self._api_key = api_key + self._timeout = timeout def _get_headers(self) -> dict[str, str]: """ @@ -60,7 +62,7 @@ class TeamsManagementClient: if organization_id: params["organization_id"] = organization_id - response: Final = requests.get(url, headers=self._get_headers(), params=params) + response: Final = requests.get(url, headers=self._get_headers(), params=params, timeout=self._timeout) if response.status_code == 401: raise UnauthorizedError("Authentication failed. Check your API key.") @@ -117,7 +119,7 @@ class TeamsManagementClient: if sort_by: params["sort_by"] = sort_by - response: Final = requests.get(url, headers=self._get_headers(), params=params) + response: Final = requests.get(url, headers=self._get_headers(), params=params, timeout=self._timeout) if response.status_code == 401: raise UnauthorizedError("Authentication failed. Check your API key.") @@ -138,7 +140,7 @@ class TeamsManagementClient: """ url: Final = f"{self._base_url}/team/available" - response: Final = requests.get(url, headers=self._get_headers()) + response: Final = requests.get(url, headers=self._get_headers(), timeout=self._timeout) if response.status_code == 401: raise UnauthorizedError("Authentication failed. Check your API key.") diff --git a/litellm/proxy/client/users.py b/litellm/proxy/client/users.py index df5f9aad23e..3f11fe94043 100644 --- a/litellm/proxy/client/users.py +++ b/litellm/proxy/client/users.py @@ -6,9 +6,10 @@ from .exceptions import NotFoundError, UnauthorizedError class UsersManagementClient: - def __init__(self, base_url: str, api_key: str | None = None): + def __init__(self, base_url: str, api_key: str | None = None, timeout: int = 30): self.base_url = base_url.rstrip("/") self.api_key = api_key + self.timeout = timeout def _get_headers(self) -> dict[str, str]: headers: Final = {"Content-Type": "application/json"} @@ -19,7 +20,7 @@ class UsersManagementClient: def list_users(self, params: dict[str, Any] | None = None) -> list[dict[str, Any]]: """List users (GET /user/list)""" url: Final = f"{self.base_url}/user/list" - response: Final = requests.get(url, headers=self._get_headers(), params=params) + response: Final = requests.get(url, headers=self._get_headers(), params=params, timeout=self.timeout) if response.status_code == 401: raise UnauthorizedError(response.text) response.raise_for_status() @@ -29,7 +30,7 @@ class UsersManagementClient: """Get user info (GET /user/info)""" url: Final = f"{self.base_url}/user/info" params: Final = {"user_id": user_id} if user_id else {} - response: Final = requests.get(url, headers=self._get_headers(), params=params) + response: Final = requests.get(url, headers=self._get_headers(), params=params, timeout=self.timeout) if response.status_code == 401: raise UnauthorizedError(response.text) if response.status_code == 404: @@ -41,7 +42,7 @@ class UsersManagementClient: """Get user info v2 - lightweight, returns only user object (GET /v2/user/info)""" url: Final = f"{self.base_url}/v2/user/info" params: Final = {"user_id": user_id} if user_id else {} - response: Final = requests.get(url, headers=self._get_headers(), params=params) + response: Final = requests.get(url, headers=self._get_headers(), params=params, timeout=self.timeout) if response.status_code == 401: raise UnauthorizedError(response.text) if response.status_code == 404: @@ -52,7 +53,7 @@ class UsersManagementClient: def create_user(self, user_data: dict[str, Any]) -> dict[str, Any]: """Create a new user (POST /user/new)""" url: Final = f"{self.base_url}/user/new" - response: Final = requests.post(url, headers=self._get_headers(), json=user_data) + response: Final = requests.post(url, headers=self._get_headers(), json=user_data, timeout=self.timeout) if response.status_code == 401: raise UnauthorizedError(response.text) response.raise_for_status() @@ -61,7 +62,9 @@ class UsersManagementClient: def delete_user(self, user_ids: list[str]) -> dict[str, Any]: """Delete users (POST /user/delete)""" url: Final = f"{self.base_url}/user/delete" - response: Final = requests.post(url, headers=self._get_headers(), json={"user_ids": user_ids}) + response: Final = requests.post( + url, headers=self._get_headers(), json={"user_ids": user_ids}, timeout=self.timeout + ) if response.status_code == 401: raise UnauthorizedError(response.text) response.raise_for_status() diff --git a/litellm/proxy/guardrails/guardrail_hooks/hiddenlayer/hiddenlayer.py b/litellm/proxy/guardrails/guardrail_hooks/hiddenlayer/hiddenlayer.py index 507dd645953..f15b8ec1e74 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/hiddenlayer/hiddenlayer.py +++ b/litellm/proxy/guardrails/guardrail_hooks/hiddenlayer/hiddenlayer.py @@ -36,6 +36,9 @@ if TYPE_CHECKING: from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel +_AUTH_TIMEOUT_SECONDS: Final[float] = 30.0 + + class _HiddenlayerEvaluation(TypedDict, total=False): action: str threat_level: str @@ -117,10 +120,10 @@ def is_saas(host: str) -> bool: return False -def _get_jwt(auth_url, api_id, api_key) -> str: +def _get_jwt(auth_url, api_id, api_key, timeout: float = _AUTH_TIMEOUT_SECONDS) -> str: token_url: Final = f"{auth_url}/oauth2/token?grant_type=client_credentials" - resp: Final = requests.post(token_url, auth=HTTPBasicAuth(api_id, api_key)) + resp: Final = requests.post(token_url, auth=HTTPBasicAuth(api_id, api_key), timeout=timeout) if not resp.ok: raise RuntimeError( diff --git a/ruff.toml b/ruff.toml index 44bdf9d8125..3ac4c1fc94d 100644 --- a/ruff.toml +++ b/ruff.toml @@ -6,7 +6,8 @@ lint.extend-select = [ "T20", "PGH004", "RUF008", "RUF009", "RUF100", "B033", "FURB136", "FURB168", "FURB188", "I001", "PERF402", "PIE790", "PIE800", "PLC0208", "PLR0402", "PLR1711", "PLR1730", "PLR2044", "PLW0133", "PYI030", "PYI041", "PYI064", "RET501", - "RUF010", "RUF022", "RUF023", "RUF051", "SIM114", "SIM118", "TC005", "UP006", "UP007", "UP008", + "RUF010", "RUF022", "RUF023", "RUF051", "S113", "SIM114", "SIM118", "TC005", "UP006", "UP007", + "UP008", "UP012", "UP018", "UP024", "UP032", "UP034", "UP035", "UP037", "UP045", ] # RUF100 (unused-noqa) only knows the rules enabled in THIS config, so it would strip diff --git a/tests/test_litellm/proxy/client/conftest.py b/tests/test_litellm/proxy/client/conftest.py new file mode 100644 index 00000000000..c8b7951e284 --- /dev/null +++ b/tests/test_litellm/proxy/client/conftest.py @@ -0,0 +1,38 @@ +import threading + +import pytest + + +@pytest.fixture +def hanging_server(): + """A server that accepts the connection and never answers, so only a timeout ends the call.""" + from http.server import BaseHTTPRequestHandler, HTTPServer + from socketserver import ThreadingMixIn + + stop: threading.Event = threading.Event() + + class SilentRequestHandler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def _hang(self): + stop.wait(timeout=30) + + do_GET = _hang + do_POST = _hang + + def log_message(self, format, *args): + pass + + class ThreadedServer(ThreadingMixIn, HTTPServer): + daemon_threads = True + + server = ThreadedServer(("127.0.0.1", 0), SilentRequestHandler) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + yield f"http://127.0.0.1:{server.server_port}" + finally: + stop.set() + server.shutdown() + server.server_close() + thread.join(timeout=5) diff --git a/tests/test_litellm/proxy/client/test_teams.py b/tests/test_litellm/proxy/client/test_teams.py new file mode 100644 index 00000000000..b61091ca44b --- /dev/null +++ b/tests/test_litellm/proxy/client/test_teams.py @@ -0,0 +1,20 @@ +import time + +import pytest +import requests + +from litellm.proxy.client.teams import TeamsManagementClient + + +def test_list_gives_up_at_the_timeout_instead_of_hanging(hanging_server): + """ + A proxy that accepts the connection but never answers used to pin the caller's + process forever, since the request carried no timeout at all. + """ + client = TeamsManagementClient(base_url=hanging_server, api_key="sk-test", timeout=1) + + started = time.monotonic() + with pytest.raises(requests.exceptions.Timeout): + client.list() + + assert time.monotonic() - started < 10 diff --git a/tests/test_litellm/proxy/client/test_users.py b/tests/test_litellm/proxy/client/test_users.py index 87b8392e402..5b4d89420ab 100644 --- a/tests/test_litellm/proxy/client/test_users.py +++ b/tests/test_litellm/proxy/client/test_users.py @@ -1,6 +1,8 @@ +import time from unittest.mock import MagicMock, patch import pytest +import requests @@ -82,3 +84,17 @@ def test_delete_user_unauthorized(mock_post, client): mock_post.return_value.text = "unauthorized" with pytest.raises(UnauthorizedError): client.delete_user(["u1"]) + + +def test_delete_user_gives_up_at_the_timeout_instead_of_hanging(hanging_server): + """ + A proxy that accepts the connection but never answers used to pin the caller's + process forever, since the request carried no timeout at all. + """ + client = UsersManagementClient(base_url=hanging_server, api_key="sk-test", timeout=1) + + started = time.monotonic() + with pytest.raises(requests.exceptions.Timeout): + client.delete_user(["u1"]) + + assert time.monotonic() - started < 10 diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_hiddenlayer.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_hiddenlayer.py index 1b2108c837d..b140082a3bf 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_hiddenlayer.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_hiddenlayer.py @@ -1,4 +1,6 @@ import os +import threading +import time import uuid from typing import List, cast from unittest.mock import AsyncMock, MagicMock, patch @@ -6,6 +8,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi import HTTPException from httpx import Request, Response +import requests import litellm @@ -14,6 +17,7 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.proxy.guardrails.guardrail_hooks.hiddenlayer.hiddenlayer import ( HiddenlayerGuardrail, HiddenlayerGuardrailV2, + _get_jwt, ) from litellm.proxy.guardrails.init_guardrails import init_guardrails_v2 from litellm.types.utils import ( @@ -1088,3 +1092,47 @@ class TestHiddenlayerGuardrailV2: config_model = HiddenlayerGuardrailV2.get_config_model() assert config_model is not None assert config_model.__name__ == "HiddenlayerGuardrailConfigModel" + + +@pytest.fixture +def hanging_auth_server(): + """A server that accepts the connection and never answers, so only a timeout ends the call.""" + from http.server import BaseHTTPRequestHandler, HTTPServer + from socketserver import ThreadingMixIn + + stop: threading.Event = threading.Event() + + class SilentRequestHandler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def do_POST(self): + stop.wait(timeout=30) + + def log_message(self, format, *args): + pass + + class ThreadedServer(ThreadingMixIn, HTTPServer): + daemon_threads = True + + server = ThreadedServer(("127.0.0.1", 0), SilentRequestHandler) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + yield f"http://127.0.0.1:{server.server_port}" + finally: + stop.set() + server.shutdown() + server.server_close() + thread.join(timeout=5) + + +def test_get_jwt_gives_up_at_the_timeout_instead_of_blocking_the_event_loop(hanging_auth_server): + """ + `_get_jwt` runs synchronously inside `_call_hiddenlayer`, so an auth host that + accepts and never answers used to park the whole worker's event loop. + """ + started = time.monotonic() + with pytest.raises(requests.exceptions.Timeout): + _get_jwt(auth_url=hanging_auth_server, api_id="id", api_key="secret", timeout=1) + + assert time.monotonic() - started < 10 diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 627811a7f1d..6dcabe076c9 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,6 +1,6 @@ { "LIT001": { - "limit": 22805 + "limit": 22804 }, "LIT002": { "limit": 26873 @@ -33,6 +33,6 @@ "limit": 5588 }, "LIT012": { - "limit": 4510 + "limit": 4508 } } From cbc931da5465e1ea08fa9d5cf702c2cc38f1ab35 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Tue, 25 Aug 2026 11:00:29 -0700 Subject: [PATCH 27/68] test: assert the poll call shape after the timeout refactor --- tests/test_litellm/proxy/auth/test_cli_auth.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/tests/test_litellm/proxy/auth/test_cli_auth.py b/tests/test_litellm/proxy/auth/test_cli_auth.py index c9b31a1d776..5cde5522376 100644 --- a/tests/test_litellm/proxy/auth/test_cli_auth.py +++ b/tests/test_litellm/proxy/auth/test_cli_auth.py @@ -82,7 +82,7 @@ async def test_poll_for_ready_404(sleep_mock, request_mock): _poll_for_ready_data( "https://litellm.com", poll_interval=1, total_timeout=1, request_timeout=42 ) - request_mock.assert_called_once_with("https://litellm.com", timeout=42) + request_mock.assert_called_once_with("https://litellm.com", headers=None, timeout=42) @pytest.mark.asyncio @@ -103,7 +103,7 @@ async def test_poll_for_ready_200_ready(sleep_mock, click_mock, request_mock): ) assert actual == {"status": "ready", "json": "data"} click_mock.assert_not_called() - request_mock.assert_called_once_with("https://litellm.com", timeout=42) + request_mock.assert_called_once_with("https://litellm.com", headers=None, timeout=42) sleep_mock.assert_not_called() @@ -131,8 +131,8 @@ async def test_poll_for_ready_single_pending(sleep_mock, click_mock, request_moc click_mock.assert_not_called() request_mock.assert_has_calls( [ - call("https://litellm.com", timeout=42), - call("https://litellm.com", timeout=42), + call("https://litellm.com", headers=None, timeout=42), + call("https://litellm.com", headers=None, timeout=42), ] ) sleep_mock.assert_called_once_with(1) @@ -168,8 +168,8 @@ async def test_poll_for_ready_pending(sleep_mock, click_mock, request_mock): click_mock.assert_has_calls([call("Pending message"), call("Pending message")]) request_mock.assert_has_calls( [ - call("https://litellm.com", timeout=42), - call("https://litellm.com", timeout=42), + call("https://litellm.com", headers=None, timeout=42), + call("https://litellm.com", headers=None, timeout=42), ] ) sleep_mock.assert_has_calls([call(1), call(1)]) @@ -194,7 +194,7 @@ async def test_poll_for_ready_connection_failure(sleep_mock, click_mock, request click_mock.assert_called_once_with("Connection error (will retry): ERROR") request_mock.assert_has_calls( [ - call("https://litellm.com", timeout=42), + call("https://litellm.com", headers=None, timeout=42), ] ) sleep_mock.assert_has_calls([call(1), call(1)]) From 80fd9970c26ec1c81955dcba55e293ed73d4d69e Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 26 Aug 2026 01:25:26 +0000 Subject: [PATCH 28/68] fix(azure): route image generation and edits through /openai/v1 for v1 api versions Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/azure/azure.py | 7 +++ .../llms/azure/image_edit/transformation.py | 11 +++- .../test_azure_image_edit_transformation.py | 53 +++++++++++++++++++ .../test_azure_image_generation_init.py | 32 +++++++++++ 4 files changed, 101 insertions(+), 2 deletions(-) diff --git a/litellm/llms/azure/azure.py b/litellm/llms/azure/azure.py index 980b27cda55..3a30444f06d 100644 --- a/litellm/llms/azure/azure.py +++ b/litellm/llms/azure/azure.py @@ -1113,6 +1113,13 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): api_version=api_version, ) + if BaseAzureLLM._is_azure_v1_api_version(api_version): + return BaseAzureLLM._get_base_azure_url( + api_base=api_base, + litellm_params={"api_version": api_version}, + route="/openai/images/generations", + ) + if "/openai/deployments/" in api_base: base_url_with_deployment = api_base else: diff --git a/litellm/llms/azure/image_edit/transformation.py b/litellm/llms/azure/image_edit/transformation.py index 15592968bad..eb1a51ca3e8 100644 --- a/litellm/llms/azure/image_edit/transformation.py +++ b/litellm/llms/azure/image_edit/transformation.py @@ -93,8 +93,6 @@ class AzureImageEditConfig(OpenAIImageEditConfig): raise ValueError( f"api_base is required for Azure AI Studio. Please set the api_base parameter. Passed `api_base={api_base}`" ) - original_url: Final = httpx.URL(api_base) - # Resolve api_version: litellm_params > litellm.api_version > AZURE_API_VERSION env > default. # Mirrors the fallback chain used by the Azure chat path in common_utils.py, # so callers that set a global / env api_version don't get an unversioned URL. @@ -105,6 +103,15 @@ class AzureImageEditConfig(OpenAIImageEditConfig): or litellm.AZURE_DEFAULT_API_VERSION ) + if BaseAzureLLM._is_azure_v1_api_version(api_version): + return BaseAzureLLM._get_base_azure_url( + api_base=api_base, + litellm_params={"api_version": api_version}, + route="/openai/images/edits", + ) + + original_url: Final = httpx.URL(api_base) + # Create a new dictionary with existing params query_params: Final = dict(original_url.params) diff --git a/tests/test_litellm/llms/azure/image_edit/test_azure_image_edit_transformation.py b/tests/test_litellm/llms/azure/image_edit/test_azure_image_edit_transformation.py index 59472d1a49d..251b98aa8d4 100644 --- a/tests/test_litellm/llms/azure/image_edit/test_azure_image_edit_transformation.py +++ b/tests/test_litellm/llms/azure/image_edit/test_azure_image_edit_transformation.py @@ -233,3 +233,56 @@ def test_api_version_in_api_base_query_is_preserved(monkeypatch): ) assert _query_params(url) == {"api-version": "2024-05-01-preview"} + + +# --------------------------------------------------------------------------- +# Azure v1 API surface (api_version in {"v1", "preview", "latest"}) +# +# The v1 surface exposes /openai/v1/images/edits and routes by ``model`` in the +# multipart form. Building the deployment-scoped path instead makes Azure 404. +# --------------------------------------------------------------------------- + + +def test_v1_api_version_uses_v1_route_and_keeps_model(monkeypatch): + monkeypatch.setattr(litellm, "api_version", None, raising=False) + monkeypatch.delenv("AZURE_API_VERSION", raising=False) + config = AzureImageEditConfig() + + for api_version in ("v1", "preview", "latest"): + url = config.get_complete_url( + model=_FALLBACK_MODEL, + api_base=_FALLBACK_API_BASE, + litellm_params={"api_version": api_version}, + ) + assert urllib.parse.urlparse(url).path == "/openai/v1/images/edits" + assert _query_params(url) == {"api-version": api_version} + assert config.finalize_image_edit_request_data({"model": _FALLBACK_MODEL, "prompt": "x"}, url) == { + "model": _FALLBACK_MODEL, + "prompt": "x", + } + + +def test_v1_api_version_from_global_uses_v1_route(monkeypatch): + monkeypatch.setattr(litellm, "api_version", "preview", raising=False) + monkeypatch.delenv("AZURE_API_VERSION", raising=False) + + url = AzureImageEditConfig().get_complete_url( + model=_FALLBACK_MODEL, + api_base=_FALLBACK_API_BASE, + litellm_params={}, + ) + + assert urllib.parse.urlparse(url).path == "/openai/v1/images/edits" + + +def test_dated_api_version_still_uses_deployment_route(monkeypatch): + monkeypatch.setattr(litellm, "api_version", None, raising=False) + monkeypatch.delenv("AZURE_API_VERSION", raising=False) + + url = AzureImageEditConfig().get_complete_url( + model=_FALLBACK_MODEL, + api_base=_FALLBACK_API_BASE, + litellm_params={"api_version": "2024-10-21"}, + ) + + assert urllib.parse.urlparse(url).path == f"/openai/deployments/{_FALLBACK_MODEL}/images/edits" diff --git a/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py b/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py index 560fee17328..ebb6c672f96 100644 --- a/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py +++ b/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py @@ -433,3 +433,35 @@ async def test_azure_aimage_generation_base_model_vs_deployment_name(): wire_json = post_kwargs.get("json") or {} assert "model" not in wire_json assert data.get("model") == base_model + + +@pytest.mark.parametrize("api_version", ["v1", "preview", "latest"]) +def test_azure_image_generation_v1_api_version_uses_v1_route(api_version): + """The v1 Azure surface exposes /openai/v1/images/generations and routes by body ``model``.""" + url = AzureChatCompletion().create_azure_base_url( + azure_client_params={ + "azure_endpoint": "https://my-resource.openai.azure.com", + "api_version": api_version, + }, + model="gpt-image-1", + base_model=None, + ) + assert url == f"https://my-resource.openai.azure.com/openai/v1/images/generations?api-version={api_version}" + data = {"model": "gpt-image-1", "prompt": "x"} + assert azure_deployment_image_generation_json_body(url, data) == data + + +def test_azure_image_generation_dated_api_version_uses_deployment_route(): + url = AzureChatCompletion().create_azure_base_url( + azure_client_params={ + "azure_endpoint": "https://my-resource.openai.azure.com", + "api_version": "2024-10-21", + }, + model="gpt-image-1", + base_model=None, + ) + assert ( + url + == "https://my-resource.openai.azure.com/openai/deployments/gpt-image-1/images/generations?api-version=2024-10-21" + ) + assert "model" not in azure_deployment_image_generation_json_body(url, {"model": "gpt-image-1", "prompt": "x"}) From a6e68b1b0ea9baceaf695eadc14a62163d2de2d6 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 26 Aug 2026 01:29:57 +0000 Subject: [PATCH 29/68] fix(azure): wrap v1 image route litellm_params in MappingProxyType Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/azure/azure.py | 3 ++- litellm/llms/azure/image_edit/transformation.py | 3 ++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/litellm/llms/azure/azure.py b/litellm/llms/azure/azure.py index 3a30444f06d..c23a4dd04fa 100644 --- a/litellm/llms/azure/azure.py +++ b/litellm/llms/azure/azure.py @@ -2,6 +2,7 @@ import asyncio import json import time from collections.abc import Callable, Coroutine +from types import MappingProxyType from typing import Any, Final import httpx @@ -1116,7 +1117,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): if BaseAzureLLM._is_azure_v1_api_version(api_version): return BaseAzureLLM._get_base_azure_url( api_base=api_base, - litellm_params={"api_version": api_version}, + litellm_params=MappingProxyType({"api_version": api_version}), route="/openai/images/generations", ) diff --git a/litellm/llms/azure/image_edit/transformation.py b/litellm/llms/azure/image_edit/transformation.py index eb1a51ca3e8..274bb49f400 100644 --- a/litellm/llms/azure/image_edit/transformation.py +++ b/litellm/llms/azure/image_edit/transformation.py @@ -1,3 +1,4 @@ +from types import MappingProxyType from typing import Final, cast import httpx @@ -106,7 +107,7 @@ class AzureImageEditConfig(OpenAIImageEditConfig): if BaseAzureLLM._is_azure_v1_api_version(api_version): return BaseAzureLLM._get_base_azure_url( api_base=api_base, - litellm_params={"api_version": api_version}, + litellm_params=MappingProxyType({"api_version": api_version}), route="/openai/images/edits", ) From 5802cf0d8d8772cfe401b3eab8a5680110de22a3 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 26 Aug 2026 01:40:59 +0000 Subject: [PATCH 30/68] fix(azure): drop deployment path from api_base when building v1 image routes Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/azure/azure.py | 14 +++++------ litellm/llms/azure/common_utils.py | 23 +++++++++++++++++++ .../llms/azure/image_edit/transformation.py | 14 +++++------ .../test_azure_image_edit_transformation.py | 22 +++++++++++------- .../test_azure_image_generation_init.py | 12 ++++++++++ 5 files changed, 63 insertions(+), 22 deletions(-) diff --git a/litellm/llms/azure/azure.py b/litellm/llms/azure/azure.py index c23a4dd04fa..ce65d62ca45 100644 --- a/litellm/llms/azure/azure.py +++ b/litellm/llms/azure/azure.py @@ -2,7 +2,6 @@ import asyncio import json import time from collections.abc import Callable, Coroutine -from types import MappingProxyType from typing import Any, Final import httpx @@ -1114,12 +1113,13 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): api_version=api_version, ) - if BaseAzureLLM._is_azure_v1_api_version(api_version): - return BaseAzureLLM._get_base_azure_url( - api_base=api_base, - litellm_params=MappingProxyType({"api_version": api_version}), - route="/openai/images/generations", - ) + v1_url: Final = BaseAzureLLM.get_azure_v1_image_url( + api_base=api_base, + api_version=api_version, + route="/openai/images/generations", + ) + if v1_url is not None: + return v1_url if "/openai/deployments/" in api_base: base_url_with_deployment = api_base diff --git a/litellm/llms/azure/common_utils.py b/litellm/llms/azure/common_utils.py index 2c34851d275..ae9a43a540c 100644 --- a/litellm/llms/azure/common_utils.py +++ b/litellm/llms/azure/common_utils.py @@ -4,6 +4,7 @@ import json import os from collections.abc import Callable, Mapping from functools import lru_cache +from types import MappingProxyType from typing import Any, Final, Literal, NamedTuple, cast import httpx @@ -789,6 +790,28 @@ class BaseAzureLLM(BaseOpenAILLM): return str(final_url) + @staticmethod + def get_azure_v1_image_url(api_base: str, api_version: str | None, route: str) -> str | None: + """ + Azure's v1 surface serves images at ``/openai/v1/images/{generations,edits}`` and routes by + ``model`` in the request body, so any deployment path in ``api_base`` has to be dropped. + + Returns None when ``api_version`` is a dated one, which still uses the deployment route. + """ + if not BaseAzureLLM._is_azure_v1_api_version(api_version): + return None + + base_url: Final = httpx.URL(api_base) + openai_path_start: Final = base_url.path.find("/openai") + resource_base: Final = ( + api_base if openai_path_start == -1 else str(base_url.copy_with(path=base_url.path[:openai_path_start])) + ) + return BaseAzureLLM._get_base_azure_url( + api_base=resource_base, + litellm_params=MappingProxyType({"api_version": api_version}), + route=route, + ) + @staticmethod def _is_azure_v1_api_version(api_version: str | None) -> bool: if api_version is None: diff --git a/litellm/llms/azure/image_edit/transformation.py b/litellm/llms/azure/image_edit/transformation.py index 274bb49f400..e4716289a34 100644 --- a/litellm/llms/azure/image_edit/transformation.py +++ b/litellm/llms/azure/image_edit/transformation.py @@ -1,4 +1,3 @@ -from types import MappingProxyType from typing import Final, cast import httpx @@ -104,12 +103,13 @@ class AzureImageEditConfig(OpenAIImageEditConfig): or litellm.AZURE_DEFAULT_API_VERSION ) - if BaseAzureLLM._is_azure_v1_api_version(api_version): - return BaseAzureLLM._get_base_azure_url( - api_base=api_base, - litellm_params=MappingProxyType({"api_version": api_version}), - route="/openai/images/edits", - ) + v1_url: Final = BaseAzureLLM.get_azure_v1_image_url( + api_base=api_base, + api_version=api_version, + route="/openai/images/edits", + ) + if v1_url is not None: + return v1_url original_url: Final = httpx.URL(api_base) diff --git a/tests/test_litellm/llms/azure/image_edit/test_azure_image_edit_transformation.py b/tests/test_litellm/llms/azure/image_edit/test_azure_image_edit_transformation.py index 251b98aa8d4..48fa7221dff 100644 --- a/tests/test_litellm/llms/azure/image_edit/test_azure_image_edit_transformation.py +++ b/tests/test_litellm/llms/azure/image_edit/test_azure_image_edit_transformation.py @@ -235,14 +235,6 @@ def test_api_version_in_api_base_query_is_preserved(monkeypatch): assert _query_params(url) == {"api-version": "2024-05-01-preview"} -# --------------------------------------------------------------------------- -# Azure v1 API surface (api_version in {"v1", "preview", "latest"}) -# -# The v1 surface exposes /openai/v1/images/edits and routes by ``model`` in the -# multipart form. Building the deployment-scoped path instead makes Azure 404. -# --------------------------------------------------------------------------- - - def test_v1_api_version_uses_v1_route_and_keeps_model(monkeypatch): monkeypatch.setattr(litellm, "api_version", None, raising=False) monkeypatch.delenv("AZURE_API_VERSION", raising=False) @@ -286,3 +278,17 @@ def test_dated_api_version_still_uses_deployment_route(monkeypatch): ) assert urllib.parse.urlparse(url).path == f"/openai/deployments/{_FALLBACK_MODEL}/images/edits" + + +def test_v1_api_version_replaces_deployment_scoped_api_base(monkeypatch): + monkeypatch.setattr(litellm, "api_version", None, raising=False) + monkeypatch.delenv("AZURE_API_VERSION", raising=False) + + url = AzureImageEditConfig().get_complete_url( + model=_FALLBACK_MODEL, + api_base=f"{_FALLBACK_API_BASE}/openai/deployments/{_FALLBACK_MODEL}/images/edits", + litellm_params={"api_version": "preview"}, + ) + + assert urllib.parse.urlparse(url).path == "/openai/v1/images/edits" + assert _query_params(url) == {"api-version": "preview"} diff --git a/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py b/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py index ebb6c672f96..cfec4eea2f4 100644 --- a/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py +++ b/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py @@ -465,3 +465,15 @@ def test_azure_image_generation_dated_api_version_uses_deployment_route(): == "https://my-resource.openai.azure.com/openai/deployments/gpt-image-1/images/generations?api-version=2024-10-21" ) assert "model" not in azure_deployment_image_generation_json_body(url, {"model": "gpt-image-1", "prompt": "x"}) + + +def test_azure_image_generation_v1_api_version_replaces_deployment_scoped_api_base(): + url = AzureChatCompletion().create_azure_base_url( + azure_client_params={ + "azure_endpoint": "https://my-resource.openai.azure.com/openai/deployments/gpt-image-1/images/generations", + "api_version": "preview", + }, + model="gpt-image-1", + base_model=None, + ) + assert url == "https://my-resource.openai.azure.com/openai/v1/images/generations?api-version=preview" From 1cf683c9806246fa9af96d30194a0b68b5c83f2d Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 26 Aug 2026 01:49:48 +0000 Subject: [PATCH 31/68] fix(azure): honor base_url client param and drop stale api-version on v1 image routes Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/azure/azure.py | 7 ++++--- litellm/llms/azure/common_utils.py | 10 +++++++--- .../test_azure_image_edit_transformation.py | 2 +- .../test_azure_image_generation_init.py | 12 ++++++++++++ 4 files changed, 24 insertions(+), 7 deletions(-) diff --git a/litellm/llms/azure/azure.py b/litellm/llms/azure/azure.py index ce65d62ca45..c2d4d8b9306 100644 --- a/litellm/llms/azure/azure.py +++ b/litellm/llms/azure/azure.py @@ -1091,9 +1091,10 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): AzureFoundryMAIImageGenerationConfig, ) - api_base: str = azure_client_params.get("azure_endpoint", "") # "https://example-endpoint.openai.azure.com" - if api_base.endswith("/"): - api_base = api_base.rstrip("/") + # deployment-scoped endpoints are moved to "base_url" by select_azure_base_url_or_endpoint + api_base: str = (azure_client_params.get("azure_endpoint") or azure_client_params.get("base_url") or "").rstrip( + "/" + ) api_version: Final[str] = azure_client_params.get("api_version", "") if model is None: model = "" diff --git a/litellm/llms/azure/common_utils.py b/litellm/llms/azure/common_utils.py index ae9a43a540c..8d5c7687ec8 100644 --- a/litellm/llms/azure/common_utils.py +++ b/litellm/llms/azure/common_utils.py @@ -794,7 +794,8 @@ class BaseAzureLLM(BaseOpenAILLM): def get_azure_v1_image_url(api_base: str, api_version: str | None, route: str) -> str | None: """ Azure's v1 surface serves images at ``/openai/v1/images/{generations,edits}`` and routes by - ``model`` in the request body, so any deployment path in ``api_base`` has to be dropped. + ``model`` in the request body, so any deployment path and stale ``api-version`` in + ``api_base`` have to be dropped. Returns None when ``api_version`` is a dated one, which still uses the deployment route. """ @@ -803,8 +804,11 @@ class BaseAzureLLM(BaseOpenAILLM): base_url: Final = httpx.URL(api_base) openai_path_start: Final = base_url.path.find("/openai") - resource_base: Final = ( - api_base if openai_path_start == -1 else str(base_url.copy_with(path=base_url.path[:openai_path_start])) + resource_base: Final = str( + base_url.copy_with( + path=base_url.path if openai_path_start == -1 else base_url.path[:openai_path_start], + params=httpx.QueryParams({k: v for k, v in base_url.params.items() if k != "api-version"}), + ) ) return BaseAzureLLM._get_base_azure_url( api_base=resource_base, diff --git a/tests/test_litellm/llms/azure/image_edit/test_azure_image_edit_transformation.py b/tests/test_litellm/llms/azure/image_edit/test_azure_image_edit_transformation.py index 48fa7221dff..5c7b249ae72 100644 --- a/tests/test_litellm/llms/azure/image_edit/test_azure_image_edit_transformation.py +++ b/tests/test_litellm/llms/azure/image_edit/test_azure_image_edit_transformation.py @@ -286,7 +286,7 @@ def test_v1_api_version_replaces_deployment_scoped_api_base(monkeypatch): url = AzureImageEditConfig().get_complete_url( model=_FALLBACK_MODEL, - api_base=f"{_FALLBACK_API_BASE}/openai/deployments/{_FALLBACK_MODEL}/images/edits", + api_base=f"{_FALLBACK_API_BASE}/openai/deployments/{_FALLBACK_MODEL}/images/edits?api-version=2024-10-21", litellm_params={"api_version": "preview"}, ) diff --git a/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py b/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py index cfec4eea2f4..65beec2b707 100644 --- a/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py +++ b/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py @@ -477,3 +477,15 @@ def test_azure_image_generation_v1_api_version_replaces_deployment_scoped_api_ba base_model=None, ) assert url == "https://my-resource.openai.azure.com/openai/v1/images/generations?api-version=preview" + + +def test_azure_image_generation_v1_api_version_uses_base_url_client_param(): + url = AzureChatCompletion().create_azure_base_url( + azure_client_params={ + "base_url": "https://my-resource.openai.azure.com/openai/deployments/gpt-image-1?api-version=2024-10-21", + "api_version": "preview", + }, + model="gpt-image-1", + base_model=None, + ) + assert url == "https://my-resource.openai.azure.com/openai/v1/images/generations?api-version=preview" From ed9575520b289b70986a3e9a3382c7638d1723a8 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 26 Aug 2026 01:55:57 +0000 Subject: [PATCH 32/68] fix(azure): build v1 image query params without a mutable dict Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/azure/common_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/llms/azure/common_utils.py b/litellm/llms/azure/common_utils.py index 8d5c7687ec8..6cb7d09cec4 100644 --- a/litellm/llms/azure/common_utils.py +++ b/litellm/llms/azure/common_utils.py @@ -807,7 +807,7 @@ class BaseAzureLLM(BaseOpenAILLM): resource_base: Final = str( base_url.copy_with( path=base_url.path if openai_path_start == -1 else base_url.path[:openai_path_start], - params=httpx.QueryParams({k: v for k, v in base_url.params.items() if k != "api-version"}), + params=httpx.QueryParams(tuple((k, v) for k, v in base_url.params.multi_items() if k != "api-version")), ) ) return BaseAzureLLM._get_base_azure_url( From bff46fd782c709dc5ccc37d10082d52ae59e9441 Mon Sep 17 00:00:00 2001 From: shivam Date: Thu, 27 Aug 2026 22:15:41 +0000 Subject: [PATCH 33/68] chore: reformat headroom stream constant and regenerate stale openapi snapshot Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/_lazy_openapi_snapshot.json | 44 +++++++++++++++++++ litellm/types/integrations/custom_logger.py | 4 +- ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 + 3 files changed, 47 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index e30750ef565..1963c7799a2 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -15038,6 +15038,17 @@ } ], "title": "Upstream Resource" + }, + "upstream_token_header": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Upstream Token Header" } }, "title": "MCPCredentials", @@ -17518,6 +17529,17 @@ } ], "title": "Upstream Resource" + }, + "upstream_token_header": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Upstream Token Header" } }, "title": "MCPCredentials", @@ -20352,6 +20374,17 @@ } ], "title": "Upstream Resource" + }, + "upstream_token_header": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Upstream Token Header" } }, "title": "MCPCredentials", @@ -23699,6 +23732,17 @@ } ], "title": "Upstream Resource" + }, + "upstream_token_header": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Upstream Token Header" } }, "title": "MCPCredentials", diff --git a/litellm/types/integrations/custom_logger.py b/litellm/types/integrations/custom_logger.py index fd4648bff14..9a714e1724e 100644 --- a/litellm/types/integrations/custom_logger.py +++ b/litellm/types/integrations/custom_logger.py @@ -6,9 +6,7 @@ CHAT_COMPLETION_AGENTIC_SURFACE: Final = "chat_completions" RESPONSES_AGENTIC_SURFACE: Final = "responses" CODE_INTERPRETER_INTERCEPTION_PREFIX: Final = "_code_interpreter_interception" HEADROOM_INTERCEPTION_PREFIX: Final = "_headroom_interception" -HEADROOM_CONVERTED_STREAM_KEY: Final = ( - f"{HEADROOM_INTERCEPTION_PREFIX}_converted_stream" -) +HEADROOM_CONVERTED_STREAM_KEY: Final = f"{HEADROOM_INTERCEPTION_PREFIX}_converted_stream" NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES: Final = frozenset( ( "_websearch_interception", diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index c124cc2e9c8..03a4d1142f8 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -30208,6 +30208,8 @@ export interface components { token_exchange_profile?: string | null; /** Upstream Resource */ upstream_resource?: string | null; + /** Upstream Token Header */ + upstream_token_header?: string | null; }; /** * MCPEnvVar From 1cf1a97b731f99b9276d299bf5c7fc0487802b52 Mon Sep 17 00:00:00 2001 From: shivam Date: Thu, 27 Aug 2026 22:37:08 +0000 Subject: [PATCH 34/68] chore: annotate headroom pre-call hook dict for the type-discipline budget Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/guardrails/guardrail_hooks/headroom/headroom.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py index e4fe392b97f..2b34c0eeccb 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py +++ b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py @@ -728,7 +728,11 @@ class HeadroomGuardrail(CustomGuardrail): return None if not has_headroom_retrieve_tool(kwargs.get("tools")): return None - return {**kwargs, "stream": False, HEADROOM_CONVERTED_STREAM_KEY: True} + return { # mutable-ok: the hook contract is a plain dict the router merges into the request kwargs + **kwargs, + "stream": False, + HEADROOM_CONVERTED_STREAM_KEY: True, + } async def async_should_run_agentic_loop( self, From d556a042866921af4f0ec583167159be660d551b Mon Sep 17 00:00:00 2001 From: shivam Date: Thu, 27 Aug 2026 22:49:23 +0000 Subject: [PATCH 35/68] test(headroom): fake the HTTP boundary in the streaming CCR regression test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/_lazy_openapi_snapshot.json | 35 ++++++ .../guardrail_hooks/test_headroom.py | 116 ++++++++++-------- ui/litellm-dashboard/src/lib/http/schema.d.ts | 5 + 3 files changed, 104 insertions(+), 52 deletions(-) diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 1963c7799a2..040d258f97a 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -10238,6 +10238,18 @@ "description": "AWS Bedrock runtime endpoint URL", "title": "Aws Bedrock Runtime Endpoint" }, + "aws_external_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "description": "External ID required by the target role's trust policy on sts:AssumeRole", + "title": "Aws External Id" + }, "aws_profile_name": { "anyOf": [ { @@ -25237,6 +25249,9 @@ }, { "$ref": "#/components/schemas/ChatCompletionImageObject" + }, + { + "$ref": "#/components/schemas/ChatCompletionToolReferenceObject" } ] }, @@ -25324,6 +25339,26 @@ "title": "ChatCompletionToolParamFunctionChunk", "type": "object" }, + "ChatCompletionToolReferenceObject": { + "description": "Anthropic tool-search result block, carried through untouched so it survives a round trip.", + "properties": { + "tool_name": { + "title": "Tool Name", + "type": "string" + }, + "type": { + "const": "tool_reference", + "title": "Type", + "type": "string" + } + }, + "required": [ + "type", + "tool_name" + ], + "title": "ChatCompletionToolReferenceObject", + "type": "object" + }, "ChatCompletionUserMessage": { "properties": { "cache_control": { diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py index 420b91c983b..52df875e714 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py @@ -28,6 +28,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest +import respx from fastapi import HTTPException import litellm @@ -44,12 +45,7 @@ from litellm.proxy.spend_tracking.compression_savings import ( from litellm.types.integrations.custom_logger import HEADROOM_CONVERTED_STREAM_KEY from litellm.types.utils import ( CallTypes, - ChatCompletionMessageToolCall, - Choices, - Function, GenericGuardrailAPIInputs, - Message, - ModelResponse, ) FAKE_API_BASE = "https://headroom.example.com" @@ -1919,30 +1915,41 @@ def _retrieve_tool_definition() -> dict: } -def _model_response_with_retrieve_call() -> ModelResponse: - return ModelResponse( - choices=[ - Choices( - finish_reason="tool_calls", - message=Message( - role="assistant", - content=None, - tool_calls=[ - ChatCompletionMessageToolCall( - id="call_ccr", - type="function", - function=Function( - name=HEADROOM_RETRIEVE_TOOL_NAME, - arguments=json.dumps({"hash": CCR_HASH}), - ), - ) - ], - ), - ) - ] +def _openai_completion_payload(message: dict, finish_reason: str) -> dict: + return { + "id": "chatcmpl-ccr", + "object": "chat.completion", + "created": 1700000000, + "model": "gpt-4o", + "choices": [{"index": 0, "message": message, "finish_reason": finish_reason}], + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + } + + +def _openai_tool_call_payload() -> dict: + return _openai_completion_payload( + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_ccr", + "type": "function", + "function": { + "name": HEADROOM_RETRIEVE_TOOL_NAME, + "arguments": json.dumps({"hash": CCR_HASH}), + }, + } + ], + }, + "tool_calls", ) +def _openai_text_payload(content: str) -> dict: + return _openai_completion_payload({"role": "assistant", "content": content}, "stop") + + @pytest.mark.parametrize( "call_type, stream, tools, expect_conversion", [ @@ -1981,6 +1988,8 @@ async def test_pre_call_deployment_hook_converts_stream_only_for_ccr_chat_comple @pytest.mark.asyncio async def test_streaming_chat_completion_resolves_ccr_retrieval_end_to_end( guardrail: HeadroomGuardrail, + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, ): """Regression test for streaming /chat/completions: the retrieve tool call the model emits must be resolved by the agentic loop instead of being streamed back @@ -1992,33 +2001,30 @@ async def test_streaming_chat_completion_resolves_ccr_retrieval_end_to_end( time.monotonic() + 999, ) - real_acompletion = litellm.acompletion + monkeypatch.setenv("OPENAI_API_KEY", "sk-test") + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + upstream = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( + side_effect=[ + httpx.Response(200, json=_openai_tool_call_payload()), + httpx.Response(200, json=_openai_text_payload(final_answer)), + ] + ) - async def acompletion_with_followup_answer(*args, **kwargs): - if kwargs.get("_agentic_loop_depth"): - kwargs["mock_response"] = final_answer - return await real_acompletion(*args, **kwargs) - - saved_callbacks = list(litellm.callbacks) - litellm.callbacks = [guardrail] - try: - with patch.object( - guardrail.async_handler, - "get", - new_callable=AsyncMock, - return_value=_make_retrieve_response(original_content), - ) as mock_get, patch.object(litellm, "acompletion", new=acompletion_with_followup_answer): - response = await litellm.acompletion( - model="openai/gpt-4o", - messages=[{"role": "user", "content": f"summarize hash={CCR_HASH}"}], - tools=[_retrieve_tool_definition()], - stream=True, - litellm_call_id="ccr-call-id", - mock_response=_model_response_with_retrieve_call(), - ) - chunks = [chunk async for chunk in response] - finally: - litellm.callbacks = saved_callbacks + with patch.object( + guardrail.async_handler, + "get", + new_callable=AsyncMock, + return_value=_make_retrieve_response(original_content), + ) as mock_get: + response = await litellm.acompletion( + model="openai/gpt-4o", + messages=[{"role": "user", "content": f"summarize hash={CCR_HASH}"}], + tools=[_retrieve_tool_definition()], + stream=True, + litellm_call_id="ccr-call-id", + ) + chunks = [chunk async for chunk in response] streamed_text = "".join(chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices) assert streamed_text == final_answer @@ -2026,6 +2032,12 @@ async def test_streaming_chat_completion_resolves_ccr_retrieval_end_to_end( mock_get.assert_called_once() assert CCR_HASH in (mock_get.call_args.kwargs.get("url") or mock_get.call_args.args[0]) + assert len(upstream.calls) == 2 + followup_body = json.loads(upstream.calls[1].request.content) + assert not followup_body.get("stream") + assert original_content in json.dumps(followup_body["messages"]) + assert not any(key.startswith("_headroom_interception") for key in followup_body) + # --------------------------------------------------------------------------- # LIT-5018: the turn the model is being asked to act on is never compressed. diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 25fbd53018a..8945710cba2 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -29466,6 +29466,11 @@ export interface components { * @description AWS Bedrock runtime endpoint URL */ aws_bedrock_runtime_endpoint?: string | null; + /** + * Aws External Id + * @description External ID required by the target role's trust policy on sts:AssumeRole + */ + aws_external_id?: string | null; /** * Aws Profile Name * @description AWS profile name for credential retrieval From a14daf18d6ebbe399ec271207a9ea7f9124155b8 Mon Sep 17 00:00:00 2001 From: shivam Date: Thu, 27 Aug 2026 23:19:17 +0000 Subject: [PATCH 36/68] refactor(agentic-loop): narrow logging_obj before building the fake stream wrapper Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/litellm_core_utils/chat_completion_agentic_loop.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/litellm/litellm_core_utils/chat_completion_agentic_loop.py b/litellm/litellm_core_utils/chat_completion_agentic_loop.py index a5f5c6ac68f..c8e9e2583ba 100644 --- a/litellm/litellm_core_utils/chat_completion_agentic_loop.py +++ b/litellm/litellm_core_utils/chat_completion_agentic_loop.py @@ -10,6 +10,8 @@ from litellm.litellm_core_utils.agentic_loop_settings import ( DEFAULT_MAX_AGENTIC_LOOPS, validated_max_agentic_loops, ) +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObject +from litellm.llms.base_llm.base_model_iterator import MockResponseIterator from litellm.types.integrations.custom_logger import ( CHAT_COMPLETION_AGENTIC_SURFACE, HEADROOM_CONVERTED_STREAM_KEY, @@ -104,9 +106,8 @@ def _wrap_response_as_fake_stream( ) -> object: if isinstance(response, CustomStreamWrapper): return response - if not isinstance(response, ModelResponse): + if not isinstance(response, ModelResponse) or not isinstance(logging_obj, LiteLLMLoggingObject): return response - from litellm.llms.base_llm.base_model_iterator import MockResponseIterator return CustomStreamWrapper( completion_stream=MockResponseIterator(model_response=response), From 7f3ff3b47f350df5a394e23d7e89fb0592b79de9 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 28 Aug 2026 15:31:11 +0000 Subject: [PATCH 37/68] fix(bedrock): route all cohere.embed models to BedrockCohereEmbeddingConfig Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/utils.py | 2 +- tests/test_litellm/test_utils.py | 31 +++++++++++++++++++++++++++++++ 2 files changed, 32 insertions(+), 1 deletion(-) diff --git a/litellm/utils.py b/litellm/utils.py index 520c40f67c0..0cceda85a09 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -3536,7 +3536,7 @@ def get_optional_params_embeddings( object = litellm.AmazonTitanMultimodalEmbeddingG1Config() elif "amazon.titan-embed-text-v2:0" in model: object = litellm.AmazonTitanV2Config() - elif "cohere.embed-multilingual-v3" in model or "cohere.embed-v4" in model: + elif "cohere.embed" in model: object = litellm.BedrockCohereEmbeddingConfig() elif "twelvelabs" in model or "marengo" in model: object = litellm.TwelveLabsMarengoEmbeddingConfig() diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 92033251b13..6381e1a1274 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -4412,6 +4412,37 @@ class TestVertexEmbeddingEncodingFormat: assert optional_params.get("outputDimensionality") == 256 +class TestBedrockCohereEmbeddingDispatch: + """All bedrock cohere.embed models must route to BedrockCohereEmbeddingConfig, + not just multilingual-v3/v4: english-v3 was falling into the unmapped + else-branch and rejecting encoding_format. Issue #38659.""" + + @pytest.mark.parametrize( + "model", + [ + "cohere.embed-english-v3", + "cohere.embed-multilingual-v3", + "cohere.embed-v4:0", + ], + ) + def test_cohere_embed_models_accept_encoding_format(self, model): + optional_params = litellm.utils.get_optional_params_embeddings( + model=model, + encoding_format="float", + custom_llm_provider="bedrock", + ) + assert optional_params.get("embedding_types") == ["float"] + + def test_cohere_embed_english_v3_maps_dimensions(self): + optional_params = litellm.utils.get_optional_params_embeddings( + model="cohere.embed-english-v3", + encoding_format="float", + dimensions=512, + custom_llm_provider="bedrock", + ) + assert optional_params.get("output_dimension") == 512 + + @pytest.mark.parametrize( "model", [ From db1b1e219551bc53bec9ed6fb60cb90502dd87ce Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 29 Aug 2026 09:45:52 +0000 Subject: [PATCH 38/68] fix: bound Hugging Face config fetch and keep embedding tests off the network Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/constants.py | 1 + litellm/utils.py | 5 +++-- .../embedding/test_huggingface_embedding_handler.py | 13 ++++++++++++- 3 files changed, 16 insertions(+), 3 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index fc88086805f..aa7ab1d81ff 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -35,6 +35,7 @@ DEFAULT_COOLDOWN_TIME_SECONDS: Final = int(os.getenv("DEFAULT_COOLDOWN_TIME_SECO DEFAULT_REPLICATE_POLLING_RETRIES: Final = int(os.getenv("DEFAULT_REPLICATE_POLLING_RETRIES", 5)) DEFAULT_REPLICATE_POLLING_DELAY_SECONDS: Final = int(os.getenv("DEFAULT_REPLICATE_POLLING_DELAY_SECONDS", 1)) DEFAULT_IMAGE_TOKEN_COUNT: Final = int(os.getenv("DEFAULT_IMAGE_TOKEN_COUNT", 250)) +HF_CONFIG_FETCH_TIMEOUT_SECONDS: Final = float(os.getenv("HF_CONFIG_FETCH_TIMEOUT_SECONDS", 10)) # Maximum wall-clock seconds a streaming response is allowed to run. # Streams exceeding this duration are terminated with a Timeout error. diff --git a/litellm/utils.py b/litellm/utils.py index 5cd9bfc5f32..1f4674fedf0 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -69,6 +69,7 @@ from litellm.constants import ( DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET, DEFAULT_TRIM_RATIO, FUNCTION_DEFINITION_TOKEN_COUNT, + HF_CONFIG_FETCH_TIMEOUT_SECONDS, INITIAL_RETRY_DELAY, JITTER, MAX_RETRY_DELAY, @@ -5168,7 +5169,7 @@ def get_max_tokens(model: str) -> int | None: config_url: Final = f"https://huggingface.co/{model_name}/raw/main/config.json" try: # Make the HTTP request to get the raw JSON file - response: Final = litellm.module_level_client.get(config_url) + response: Final = litellm.module_level_client.get(config_url, timeout=HF_CONFIG_FETCH_TIMEOUT_SECONDS) response.raise_for_status() # Raise an exception for bad responses (4xx or 5xx) # Parse the JSON response @@ -5522,7 +5523,7 @@ def _get_max_position_embeddings(model_name: str) -> int | None: try: # Make the HTTP request to get the raw JSON file - response: Final = litellm.module_level_client.get(config_url) + response: Final = litellm.module_level_client.get(config_url, timeout=HF_CONFIG_FETCH_TIMEOUT_SECONDS) response.raise_for_status() # Raise an exception for bad responses (4xx or 5xx) # Parse the JSON response diff --git a/tests/test_litellm/llms/huggingface/embedding/test_huggingface_embedding_handler.py b/tests/test_litellm/llms/huggingface/embedding/test_huggingface_embedding_handler.py index f1226311b5e..0384fb796d9 100644 --- a/tests/test_litellm/llms/huggingface/embedding/test_huggingface_embedding_handler.py +++ b/tests/test_litellm/llms/huggingface/embedding/test_huggingface_embedding_handler.py @@ -4,6 +4,7 @@ from unittest.mock import patch, MagicMock, AsyncMock import litellm import pytest +import respx MOCK_EMBEDDING_RESPONSE = [[0.1, 0.2, 0.3, 0.4, 0.5]] @@ -21,6 +22,16 @@ def mock_embedding_http_handler(): yield mock_post +@pytest.fixture +def mock_hf_config_fetch(): + """Serve the Hugging Face config.json fetched during cost calculation, so no test leaves the process""" + with respx.mock(assert_all_called=False) as respx_mock: + respx_mock.get(url__regex=r"https://huggingface\.co/.*/config\.json").respond( + json={"max_position_embeddings": 512} + ) + yield respx_mock + + @pytest.fixture def mock_embedding_async_http_handler(): """Fixture to mock the async HTTP handler for embedding tests""" @@ -39,7 +50,7 @@ def mock_embedding_async_http_handler(): class TestHuggingFaceEmbedding: @pytest.fixture(autouse=True) - def setup(self, mock_embedding_http_handler, mock_embedding_async_http_handler): + def setup(self, mock_embedding_http_handler, mock_embedding_async_http_handler, mock_hf_config_fetch): self.mock_get_task_patcher = patch( "litellm.llms.huggingface.embedding.handler.get_hf_task_embedding_for_model" ) From 1687c65823948fbf0a2eebc084583144f5f52165 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 29 Aug 2026 09:53:43 +0000 Subject: [PATCH 39/68] fix: hardcode HF config fetch timeout instead of reading an env var Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/constants.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/constants.py b/litellm/constants.py index aa7ab1d81ff..bce54e8f080 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -35,7 +35,7 @@ DEFAULT_COOLDOWN_TIME_SECONDS: Final = int(os.getenv("DEFAULT_COOLDOWN_TIME_SECO DEFAULT_REPLICATE_POLLING_RETRIES: Final = int(os.getenv("DEFAULT_REPLICATE_POLLING_RETRIES", 5)) DEFAULT_REPLICATE_POLLING_DELAY_SECONDS: Final = int(os.getenv("DEFAULT_REPLICATE_POLLING_DELAY_SECONDS", 1)) DEFAULT_IMAGE_TOKEN_COUNT: Final = int(os.getenv("DEFAULT_IMAGE_TOKEN_COUNT", 250)) -HF_CONFIG_FETCH_TIMEOUT_SECONDS: Final = float(os.getenv("HF_CONFIG_FETCH_TIMEOUT_SECONDS", 10)) +HF_CONFIG_FETCH_TIMEOUT_SECONDS: Final = 10.0 # Maximum wall-clock seconds a streaming response is allowed to run. # Streams exceeding this duration are terminated with a Timeout error. From e535724923439c809640534e4cfebf552ee69eb0 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 29 Aug 2026 10:10:03 +0000 Subject: [PATCH 40/68] test: cover the bounded Hugging Face config fetch Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm/test_utils.py | 29 +++++++++++++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 1ff50bd0116..d93970bf5de 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -6,6 +6,7 @@ from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest +import respx from jsonschema import validate @@ -5689,3 +5690,31 @@ class TestDefaultReasoningEffortHydration: model_info = dict(_get_model_info_helper(model="gpt-5.6-terra", custom_llm_provider="openai")) assert model_info.get("default_reasoning_effort") is None + + +class TestHuggingFaceConfigFetch: + """The Hugging Face config.json fetch runs on background logging threads during cost + calculation, so an unbounded request can hang a whole test job; the timeout is the fix.""" + + @pytest.fixture + def hf_config_route(self): + with respx.mock(assert_all_called=True) as respx_mock: + yield respx_mock.get(url__regex=r"https://huggingface\.co/.*/config\.json").respond( + json={"max_position_embeddings": 512} + ) + + def test_get_max_tokens_reads_hf_config_with_a_bounded_timeout(self, hf_config_route): + from litellm.constants import HF_CONFIG_FETCH_TIMEOUT_SECONDS + from litellm.utils import get_max_tokens + + assert get_max_tokens("huggingface/some-org/some-model") == 512 + request_timeout = hf_config_route.calls.last.request.extensions["timeout"] + assert request_timeout["read"] == HF_CONFIG_FETCH_TIMEOUT_SECONDS + + def test_get_max_position_embeddings_reads_hf_config_with_a_bounded_timeout(self, hf_config_route): + from litellm.constants import HF_CONFIG_FETCH_TIMEOUT_SECONDS + from litellm.utils import _get_max_position_embeddings + + assert _get_max_position_embeddings("some-org/some-model") == 512 + request_timeout = hf_config_route.calls.last.request.extensions["timeout"] + assert request_timeout["read"] == HF_CONFIG_FETCH_TIMEOUT_SECONDS From 8278a92a065506e592010bdc07c18e70a2100267 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sat, 29 Aug 2026 10:42:48 -0700 Subject: [PATCH 41/68] fix(ui): let the log drawer's trace sidebar expand again once collapsed The sidebar toggle sits absolutely positioned over the drawer's flex row. With the sidebar expanded it lands on the sidebar header, but once collapsed it lands on the drawer header, which is sticky at z-chrome (10). The named-z-scale refactor moved the toggle from z-20 to z-raised (1), so from then on the header painted over it and swallowed the click: collapse the trace list and there was no way to bring it back. Moves the toggle to z-floating (30) and folds the two mirrored buttons into one, since they only ever differed by icon, label and handler. Covered by a Playwright spec, which is the tier that can see the layering: the button stays visible and enabled either way, so the pre-fix failure is a click interception that jsdom cannot reproduce. --- tests/e2e/ui/tests/logs/logs.spec.ts | 26 ++++++++++++++++ .../LogDetailsDrawer/LogDetailsDrawer.tsx | 30 ++++++------------- 2 files changed, 35 insertions(+), 21 deletions(-) diff --git a/tests/e2e/ui/tests/logs/logs.spec.ts b/tests/e2e/ui/tests/logs/logs.spec.ts index b29a72cbf81..56a3d0f0109 100644 --- a/tests/e2e/ui/tests/logs/logs.spec.ts +++ b/tests/e2e/ui/tests/logs/logs.spec.ts @@ -148,6 +148,32 @@ test.describe("Logs page", () => { }); }); + test("the trace sidebar collapses and expands again", async ({ page, request }) => { + const prompt = `logs-sidebar-prompt-${uniqueSuffix()}`; + const requestId = await sendChatCompletion(request, { + model: CHAT_MODEL_A, + prompt, + }); + await waitForSpendLog(request, requestId); + + const row = await openLogsForRequest(page, requestId); + await row.click(); + + const drawer = page.getByRole("dialog").first(); + await expect(drawer.getByText("Request & Response")).toBeVisible({ timeout: 20_000 }); + + const toggle = drawer.getByLabel("Collapse trace sidebar"); + await expect(toggle).toBeVisible({ timeout: 10_000 }); + await toggle.click(); + + const expandToggle = drawer.getByLabel("Expand trace sidebar"); + await expect(expandToggle).toBeVisible({ timeout: 10_000 }); + await expandToggle.click({ timeout: 10_000 }); + + await expect(drawer.getByLabel("Collapse trace sidebar")).toBeVisible({ timeout: 10_000 }); + await expect(drawer.getByText(prompt, { exact: false })).toBeVisible({ timeout: 10_000 }); + }); + test("the JSON view exposes Request and Response tabs", async ({ page, request }) => { const prompt = `logs-json-prompt-${uniqueSuffix()}`; const requestId = await sendChatCompletion(request, { diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailsDrawer.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailsDrawer.tsx index f677a66cc54..c1e01d4d0ba 100644 --- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailsDrawer.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailsDrawer.tsx @@ -313,27 +313,15 @@ export function LogDetailsDrawer({ {logEntry?.request_id ? `Request ${logEntry.request_id} details` : "Request details"}
- {!isSidebarCollapsed ? ( - - ) : ( - - )} + {!isSidebarCollapsed && (
From fcd6ea46ce5fb9fb098a615514a7b0d012d41f17 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sat, 29 Aug 2026 11:00:47 -0700 Subject: [PATCH 42/68] refactor(ui): keep the log drawer's trace sidebar in flow as a collapsed rail The toggle was absolutely positioned over the drawer's flex row, owned by neither column. That forced two coupled workarounds: a stacking level so it could beat whatever it landed on, and pl-12 on the sidebar header to reserve space for a button that was not its child. The sidebar column now always renders, at 224px expanded and a 40px rail collapsed, and the toggle is a normal in-flow child of the column it controls. No absolute, no z-index, no reserved padding, and nothing that can paint over the button. It also stops the toggle from clipping the provider logo, which it did in the collapsed state even before the z-index scale landed. Costs 40px of drawer width while collapsed. --- .../LogDetailsDrawer/LogDetailsDrawer.tsx | 181 +++++++++--------- 1 file changed, 94 insertions(+), 87 deletions(-) diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailsDrawer.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailsDrawer.tsx index c1e01d4d0ba..511e1a18253 100644 --- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailsDrawer.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailsDrawer.tsx @@ -30,6 +30,7 @@ export interface LogDetailsDrawerProps { } const SIDEBAR_WIDTH_PX = 224; +const SIDEBAR_RAIL_WIDTH_PX = 40; // Session logs are fetched page-by-page from the paginated backend and // accumulated so the drawer can show the whole session. page_size is the @@ -312,98 +313,104 @@ export function LogDetailsDrawer({ {logEntry?.request_id ? `Request ${logEntry.request_id} details` : "Request details"} -
- - {!isSidebarCollapsed && ( -
-
-
-
-
- {isSessionMode ? "Session" : "Trace"} -
-
- {leftPanelDisplayId} - -
+
+ + {!isSidebarCollapsed && ( +
+
+ {isSessionMode ? "Session" : "Trace"}
-
-
- {logsForList.length} req - {[ - isSessionMode - ? llmCount - : logsForList.filter( - (row) => !MCP_CALL_TYPES.includes(row.call_type) && !AGENT_CALL_TYPES.includes(row.call_type), - ).length, - isSessionMode - ? agentCount - : logsForList.filter((row) => AGENT_CALL_TYPES.includes(row.call_type)).length, - isSessionMode - ? mcpCount - : logsForList.filter((row) => MCP_CALL_TYPES.includes(row.call_type)).length, - ].map((count, i) => { - const label = [" LLM", " Agent", " MCP"][i]; - return count > 0 ? ( - +
+ {leftPanelDisplayId} + +
+
+ {logsForList.length} req + {[ + isSessionMode + ? llmCount + : logsForList.filter( + (row) => + !MCP_CALL_TYPES.includes(row.call_type) && !AGENT_CALL_TYPES.includes(row.call_type), + ).length, + isSessionMode + ? agentCount + : logsForList.filter((row) => AGENT_CALL_TYPES.includes(row.call_type)).length, + isSessionMode + ? mcpCount + : logsForList.filter((row) => MCP_CALL_TYPES.includes(row.call_type)).length, + ].map((count, i) => { + const label = [" LLM", " Agent", " MCP"][i]; + return count > 0 ? ( + + · + {count} + {label} + + ) : null; + })} + · + {isSessionMode ? getSpendString(totalSessionCost) : getSpendString(currentLog.spend || 0)} + {isSessionMode && ( + <> · - {count} - {label} - - ) : null; - })} - · - {isSessionMode ? getSpendString(totalSessionCost) : getSpendString(currentLog.spend || 0)} + {sessionDurationSeconds}s + + )} +
{isSessionMode && ( - <> - · - {sessionDurationSeconds}s - +
+ {cacheHitCount}/{logsForList.length} cached +
+ )} + {isSessionMode && sessionTruncated && ( +
+ Showing most recent {logsForList.length} of {sessionTotalCount} +
+ )} + {isSessionMode && ( + setSessionSortMode(value as SessionLogSortMode)} + > + + + Duration + + + Start time + + + )}
- {isSessionMode && ( -
- {cacheHitCount}/{logsForList.length} cached -
- )} - {isSessionMode && sessionTruncated && ( -
- Showing most recent {logsForList.length} of {sessionTotalCount} -
- )} - {isSessionMode && ( - setSessionSortMode(value as SessionLogSortMode)} - > - - - Duration - - - Start time - - - - )} -
+ )} +
+ {!isSidebarCollapsed && (
{normalizeGuardrailEntries(metadata?.guardrail_information).length > 0 && (
@@ -447,8 +454,8 @@ export function LogDetailsDrawer({
)}
-
- )} + )} +
Date: Sat, 29 Aug 2026 11:26:39 -0700 Subject: [PATCH 43/68] test: add batch request count keys to gcs pub sub spend logs fixture --- .../gcs_pub_sub_body/spend_logs_payload.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/logging_callback_tests/gcs_pub_sub_body/spend_logs_payload.json b/tests/logging_callback_tests/gcs_pub_sub_body/spend_logs_payload.json index 5789f19aa55..1838fb16e91 100644 --- a/tests/logging_callback_tests/gcs_pub_sub_body/spend_logs_payload.json +++ b/tests/logging_callback_tests/gcs_pub_sub_body/spend_logs_payload.json @@ -11,7 +11,7 @@ "user": "", "team_id": "", "organization_id": "", - "metadata": "{\"applied_guardrails\": [], \"attempted_fallbacks\": null, \"original_model_group\": null, \"batch_models\": null, \"mcp_tool_call_metadata\": null, \"vector_store_request_metadata\": null, \"routing_decision\": null, \"internal_call_origin\": null, \"guardrail_information\": null, \"compression_savings\": null, \"litellm_gateway_injected_cache\": null, \"usage_object\": {\"completion_tokens\": 20, \"prompt_tokens\": 10, \"total_tokens\": 30, \"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"model_map_information\": {\"model_map_key\": \"gpt-4o\", \"model_map_value\": {\"key\": \"gpt-4o\", \"max_tokens\": 16384, \"max_input_tokens\": 128000, \"max_output_tokens\": 16384, \"input_cost_per_token\": 2.5e-06, \"cache_creation_input_token_cost\": null, \"cache_read_input_token_cost\": 1.25e-06, \"input_cost_per_character\": null, \"input_cost_per_token_above_128k_tokens\": null, \"input_cost_per_token_above_200k_tokens\": null, \"input_cost_per_query\": null, \"input_cost_per_second\": null, \"input_cost_per_audio_token\": null, \"input_cost_per_token_batches\": 1.25e-06, \"output_cost_per_token_batches\": 5e-06, \"output_cost_per_token\": 1e-05, \"output_cost_per_audio_token\": null, \"output_cost_per_character\": null, \"output_cost_per_token_above_128k_tokens\": null, \"output_cost_per_character_above_128k_tokens\": null, \"output_cost_per_token_above_200k_tokens\": null, \"output_cost_per_second\": null, \"output_cost_per_image\": null, \"output_vector_size\": null, \"litellm_provider\": \"openai\", \"mode\": \"chat\", \"supports_system_messages\": true, \"supports_response_schema\": true, \"supports_vision\": true, \"supports_function_calling\": true, \"supports_tool_choice\": true, \"supports_assistant_prefill\": false, \"supports_prompt_caching\": true, \"supports_audio_input\": false, \"supports_audio_output\": false, \"supports_pdf_input\": false, \"supports_embedding_image_input\": false, \"supports_native_streaming\": null, \"supports_web_search\": true, \"supports_reasoning\": false, \"search_context_cost_per_query\": {\"search_context_size_low\": 0.03, \"search_context_size_medium\": 0.035, \"search_context_size_high\": 0.05}, \"tpm\": null, \"rpm\": null, \"supported_openai_params\": [\"frequency_penalty\", \"logit_bias\", \"logprobs\", \"top_logprobs\", \"max_tokens\", \"max_completion_tokens\", \"modalities\", \"prediction\", \"n\", \"presence_penalty\", \"seed\", \"stop\", \"stream\", \"stream_options\", \"temperature\", \"top_p\", \"tools\", \"tool_choice\", \"function_call\", \"functions\", \"max_retries\", \"extra_headers\", \"parallel_tool_calls\", \"audio\", \"response_format\", \"user\"]}}, \"additional_usage_values\": {\"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"user_api_key\": null, \"user_api_key_alias\": null, \"user_api_key_team_id\": null, \"user_api_key_project_id\": null, \"user_api_key_project_alias\": null, \"user_api_key_org_id\": null, \"user_api_key_user_id\": null, \"user_api_key_team_alias\": null, \"spend_logs_metadata\": null, \"requester_ip_address\": null, \"status\": null, \"proxy_server_request\": null, \"error_information\": null, \"attempted_retries\": null, \"max_retries\": null}", + "metadata": "{\"applied_guardrails\": [], \"attempted_fallbacks\": null, \"original_model_group\": null, \"batch_models\": null, \"batch_successful_requests\": null, \"batch_failed_requests\": null, \"mcp_tool_call_metadata\": null, \"vector_store_request_metadata\": null, \"routing_decision\": null, \"internal_call_origin\": null, \"guardrail_information\": null, \"compression_savings\": null, \"litellm_gateway_injected_cache\": null, \"usage_object\": {\"completion_tokens\": 20, \"prompt_tokens\": 10, \"total_tokens\": 30, \"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"model_map_information\": {\"model_map_key\": \"gpt-4o\", \"model_map_value\": {\"key\": \"gpt-4o\", \"max_tokens\": 16384, \"max_input_tokens\": 128000, \"max_output_tokens\": 16384, \"input_cost_per_token\": 2.5e-06, \"cache_creation_input_token_cost\": null, \"cache_read_input_token_cost\": 1.25e-06, \"input_cost_per_character\": null, \"input_cost_per_token_above_128k_tokens\": null, \"input_cost_per_token_above_200k_tokens\": null, \"input_cost_per_query\": null, \"input_cost_per_second\": null, \"input_cost_per_audio_token\": null, \"input_cost_per_token_batches\": 1.25e-06, \"output_cost_per_token_batches\": 5e-06, \"output_cost_per_token\": 1e-05, \"output_cost_per_audio_token\": null, \"output_cost_per_character\": null, \"output_cost_per_token_above_128k_tokens\": null, \"output_cost_per_character_above_128k_tokens\": null, \"output_cost_per_token_above_200k_tokens\": null, \"output_cost_per_second\": null, \"output_cost_per_image\": null, \"output_vector_size\": null, \"litellm_provider\": \"openai\", \"mode\": \"chat\", \"supports_system_messages\": true, \"supports_response_schema\": true, \"supports_vision\": true, \"supports_function_calling\": true, \"supports_tool_choice\": true, \"supports_assistant_prefill\": false, \"supports_prompt_caching\": true, \"supports_audio_input\": false, \"supports_audio_output\": false, \"supports_pdf_input\": false, \"supports_embedding_image_input\": false, \"supports_native_streaming\": null, \"supports_web_search\": true, \"supports_reasoning\": false, \"search_context_cost_per_query\": {\"search_context_size_low\": 0.03, \"search_context_size_medium\": 0.035, \"search_context_size_high\": 0.05}, \"tpm\": null, \"rpm\": null, \"supported_openai_params\": [\"frequency_penalty\", \"logit_bias\", \"logprobs\", \"top_logprobs\", \"max_tokens\", \"max_completion_tokens\", \"modalities\", \"prediction\", \"n\", \"presence_penalty\", \"seed\", \"stop\", \"stream\", \"stream_options\", \"temperature\", \"top_p\", \"tools\", \"tool_choice\", \"function_call\", \"functions\", \"max_retries\", \"extra_headers\", \"parallel_tool_calls\", \"audio\", \"response_format\", \"user\"]}}, \"additional_usage_values\": {\"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"user_api_key\": null, \"user_api_key_alias\": null, \"user_api_key_team_id\": null, \"user_api_key_project_id\": null, \"user_api_key_project_alias\": null, \"user_api_key_org_id\": null, \"user_api_key_user_id\": null, \"user_api_key_team_alias\": null, \"spend_logs_metadata\": null, \"requester_ip_address\": null, \"status\": null, \"proxy_server_request\": null, \"error_information\": null, \"attempted_retries\": null, \"max_retries\": null}", "cache_key": "Cache OFF", "spend": 0.00022500000000000002, "total_tokens": 30, From 9e01bd1441e80c18d5d64a3fc22970ba00589ba7 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 11:41:58 -0700 Subject: [PATCH 44/68] fix(azure): send the deployment name as the body model on v1 image routes --- litellm/llms/azure/azure.py | 8 +- .../llms/azure/image_generation/http_utils.py | 21 +++- .../test_azure_image_generation_init.py | 98 +++++++++++++++++++ 3 files changed, 121 insertions(+), 6 deletions(-) diff --git a/litellm/llms/azure/azure.py b/litellm/llms/azure/azure.py index c2d4d8b9306..2bcc830851a 100644 --- a/litellm/llms/azure/azure.py +++ b/litellm/llms/azure/azure.py @@ -846,6 +846,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): api_key: str, data: dict, headers: dict, + deployment_name: str | None = None, ) -> httpx.Response: """ Implemented for azure dall-e-2 image gen calls @@ -957,7 +958,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): content=json.dumps(result).encode("utf-8"), request=httpx.Request(method="POST", url="https://api.openai.com/v1"), ) - request_json: Final = azure_deployment_image_generation_json_body(api_base, data) + request_json: Final = azure_deployment_image_generation_json_body(api_base, data, deployment_name) return await async_handler.post( url=api_base, json=request_json, @@ -973,6 +974,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): api_key: str, data: dict, headers: dict, + deployment_name: str | None = None, ) -> httpx.Response: """ Implemented for azure dall-e-2 image gen calls @@ -1073,7 +1075,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): content=json.dumps(result).encode("utf-8"), request=httpx.Request(method="POST", url="https://api.openai.com/v1"), ) - request_json: Final = azure_deployment_image_generation_json_body(api_base, data) + request_json: Final = azure_deployment_image_generation_json_body(api_base, data, deployment_name) return sync_handler.post( url=api_base, json=request_json, @@ -1176,6 +1178,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): api_key=api_key, data=data, headers=headers, + deployment_name=model, ) provider_config: Final = get_azure_image_generation_config(data.get("model", "dall-e-2")) @@ -1311,6 +1314,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): api_key=api_key or "", data=data, headers=headers, + deployment_name=model, ) provider_config: Final = get_azure_image_generation_config(data.get("model", "dall-e-2")) if isinstance(provider_config, AzureFoundryMAIImageGenerationConfig): diff --git a/litellm/llms/azure/image_generation/http_utils.py b/litellm/llms/azure/image_generation/http_utils.py index 03c425eeffc..1aa5757ca95 100644 --- a/litellm/llms/azure/image_generation/http_utils.py +++ b/litellm/llms/azure/image_generation/http_utils.py @@ -1,7 +1,9 @@ """HTTP helpers for Azure OpenAI image generation (REST, not SDK).""" +from typing import Final -def azure_deployment_image_generation_json_body(api_base: str, data: dict) -> dict: + +def azure_deployment_image_generation_json_body(api_base: str, data: dict, deployment_name: str | None = None) -> dict: """ Build the JSON body for Azure OpenAI image generation POSTs. @@ -9,9 +11,20 @@ def azure_deployment_image_generation_json_body(api_base: str, data: dict) -> di deployment in the URL only; sending ``model`` in the body (especially the deployment name) breaks some models (e.g. gpt-image-2). See LiteLLM #26316. + For the v1 surface (``.../openai/v1/images/...``), Azure routes by the deployment + name in the body ``model`` field, so the deployment name must replace any base + model name there or Azure answers 404 DeploymentNotFound. + Provider-style URLs (e.g. ``/providers/...`` for FLUX on Azure AI) keep all keys so non–OpenAI-deployment payloads still work. """ - if "images/generations" in api_base and "/openai/deployments/" in api_base: - return {k: v for k, v in data.items() if k != "model"} - return data + drop_model: Final = "images/generations" in api_base and "/openai/deployments/" in api_base + v1_route: Final = "/openai/v1/images/" in api_base and bool(deployment_name) + if not drop_model and not v1_route: + return data + entries: Final = ( + tuple((k, v) for k, v in data.items() if k != "model") + if drop_model + else (*data.items(), ("model", deployment_name)) + ) + return {k: v for k, v in entries} diff --git a/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py b/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py index 65beec2b707..70b5eab5c37 100644 --- a/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py +++ b/tests/test_litellm/llms/azure/image_generation/test_azure_image_generation_init.py @@ -3,9 +3,12 @@ import traceback from typing import Callable, Optional from unittest.mock import AsyncMock, MagicMock, Mock, patch +import httpx import pytest +import respx import litellm +from litellm.caching.llm_caching_handler import LLMClientCache from litellm.llms.azure.azure import AzureChatCompletion from litellm.llms.azure.image_generation.http_utils import ( azure_deployment_image_generation_json_body, @@ -489,3 +492,98 @@ def test_azure_image_generation_v1_api_version_uses_base_url_client_param(): base_model=None, ) assert url == "https://my-resource.openai.azure.com/openai/v1/images/generations?api-version=preview" + + +def test_azure_v1_image_generation_json_body_sends_deployment_name(): + """The v1 route ignores the URL and routes by body ``model``, which must be the deployment name.""" + url = "https://my-resource.openai.azure.com/openai/v1/images/generations?api-version=preview" + data = {"model": "gpt-image-2", "prompt": "x", "n": 1} + out = azure_deployment_image_generation_json_body(url, data, deployment_name="img-dep") + assert out["model"] == "img-dep" + assert out["prompt"] == "x" + assert data["model"] == "gpt-image-2" + assert azure_deployment_image_generation_json_body(url, data) == data + + +@pytest.mark.asyncio +async def test_azure_aimage_generation_v1_route_sends_deployment_name_in_body( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "in_memory_llm_clients_cache", LLMClientCache()) + azure_chat_completion = AzureChatCompletion() + model = "img-dep" + base_model = "gpt-image-2" + data = {"model": base_model, "prompt": "A beautiful image of a cat", "n": 1} + azure_client_params = { + "azure_endpoint": "https://my-resource.openai.azure.com", + "api_version": "preview", + } + + route = respx_mock.post("https://my-resource.openai.azure.com/openai/v1/images/generations").mock( + return_value=httpx.Response(200, json={"created": 1234567890, "data": [{"b64_json": "aaaa"}]}) + ) + + logging_obj = MagicMock() + logging_obj.pre_call = MagicMock() + logging_obj.post_call = MagicMock() + + await azure_chat_completion.aimage_generation( + data=data, + model_response=None, + azure_client_params=azure_client_params, + api_key="test-api-key", + input=[], + logging_obj=logging_obj, + headers={}, + model=model, + timeout=60.0, + ) + + request = route.calls.last.request + assert str(request.url) == ("https://my-resource.openai.azure.com/openai/v1/images/generations?api-version=preview") + sent_body = json.loads(request.content) + assert sent_body["model"] == model + assert sent_body["prompt"] == data["prompt"] + + +def test_azure_image_generation_v1_route_base_model_vs_deployment_name(respx_mock: respx.MockRouter): + """On the v1 surface the body ``model`` must be the deployment name, never base_model.""" + azure_chat_completion = AzureChatCompletion() + prompt = "A beautiful image of a cat" + model = "img-dep" + base_model = "gpt-image-2" + api_base = "https://my-resource.openai.azure.com" + api_version = "v1" + litellm_params = { + "base_model": base_model, + "api_base": api_base, + "api_version": api_version, + } + + route = respx_mock.post(f"{api_base}/openai/v1/images/generations").mock( + return_value=httpx.Response(200, json={"created": 1234567890, "data": [{"b64_json": "aaaa"}]}) + ) + + logging_obj = MagicMock() + logging_obj.pre_call = MagicMock() + logging_obj.post_call = MagicMock() + + azure_chat_completion.image_generation( + prompt=prompt, + timeout=60.0, + optional_params={"n": 1, "size": "1024x1024"}, + logging_obj=logging_obj, + headers={}, + model=model, + api_key="test-api-key", + api_base=api_base, + api_version=api_version, + litellm_params=litellm_params, + ) + + request = route.calls.last.request + assert str(request.url) == f"{api_base}/openai/v1/images/generations?api-version={api_version}" + sent_body = json.loads(request.content) + assert sent_body["model"] == model + assert sent_body["prompt"] == prompt From b27a1a13a28422074f8fde1c6aaf062408802589 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sat, 29 Aug 2026 11:42:45 -0700 Subject: [PATCH 45/68] refactor(ui): move the trace sidebar toggle into the log drawer header The collapsed rail kept the toggle in flow but left a 40px stub of empty sidebar on screen. The toggle now leads the drawer header's first row, ahead of the provider logo and the model name, so it reads as part of the header and the sidebar goes back to unmounting when collapsed. Still no absolute positioning and no stacking level: the button is a normal in-flow child of the header row it sits in. DrawerHeader takes the collapsed state and the toggle handler as props rather than reaching for the drawer's state. --- .../LogDetailsDrawer/DrawerHeader.tsx | 33 ++++++++++++----- .../LogDetailsDrawer/LogDetailsDrawer.tsx | 36 ++++++------------- 2 files changed, 34 insertions(+), 35 deletions(-) diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/DrawerHeader.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/DrawerHeader.tsx index 751667b917b..d6eb33d0b0a 100644 --- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/DrawerHeader.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/DrawerHeader.tsx @@ -1,5 +1,5 @@ import { useState } from "react"; -import { Check, ChevronDown, ChevronUp, Copy, X } from "lucide-react"; +import { Check, ChevronDown, ChevronLeft, ChevronRight, ChevronUp, Copy, X } from "lucide-react"; import moment from "moment"; import { Badge } from "@/components/ui/badge"; import { Button } from "@/components/ui/button"; @@ -26,6 +26,8 @@ interface DrawerHeaderProps { statusLabel: string; statusColor: "error" | "success"; environment: string; + isSidebarCollapsed: boolean; + onToggleSidebar: () => void; } /** @@ -40,6 +42,8 @@ export function DrawerHeader({ statusLabel, statusColor, environment, + isSidebarCollapsed, + onToggleSidebar, }: DrawerHeaderProps) { const provider = log.custom_llm_provider || ""; const providerInfo = provider ? getProviderLogoAndName(provider) : null; @@ -56,13 +60,24 @@ export function DrawerHeader({ }} > {/* Row 0: Model + Provider with Logo */} - +
+ + +
{/* Row 1: Request ID + Actions */}
+
{providerLogo && (
-
-
- - {!isSidebarCollapsed && ( -
+ {!isSidebarCollapsed && ( +
+
+
{isSessionMode ? "Session" : "Trace"}
@@ -407,10 +391,8 @@ export function LogDetailsDrawer({ )}
- )} -
+
- {!isSidebarCollapsed && (
{normalizeGuardrailEntries(metadata?.guardrail_information).length > 0 && (
@@ -454,13 +436,15 @@ export function LogDetailsDrawer({
)}
- )} -
+
+ )}
setIsSidebarCollapsed((collapsed) => !collapsed)} onPrevious={selectPreviousLog} onNext={selectNextLog} statusLabel={statusLabel} From f2f988bd577bf4bcaffe0ab0b0c8de3ce254927e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 11:46:23 -0700 Subject: [PATCH 46/68] test: cover forged _headroom_interception_converted_stream strip at the proxy boundary --- tests/test_litellm/proxy/test_litellm_pre_call_utils.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index 9fb31d2a6db..649bcbb6511 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -1073,6 +1073,7 @@ async def test_key_metadata_enable_prompt_caching_promoted_to_request_root(key_v "_code_interpreter_interception_active", "_code_interpreter_interception_converted_stream", "_code_interpreter_interception_sandbox_key", + "_headroom_interception_converted_stream", "max_agentic_loops", ], ) @@ -1107,6 +1108,7 @@ async def test_add_litellm_data_to_request_strips_callback_control_fields( "_code_interpreter_interception_active": True, "_code_interpreter_interception_converted_stream": True, "_code_interpreter_interception_sandbox_key": "forged-key", + "_headroom_interception_converted_stream": True, "max_agentic_loops": 9999, } sample_value = sample_values[control_field] From 20e6d6457a9a0ef5fe15bf516bb71e55710500ce Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sat, 29 Aug 2026 11:48:04 -0700 Subject: [PATCH 47/68] refactor(ui): show the log drawer's sidebar toggle only where it has a row Putting the toggle in the drawer header unconditionally stranded it on its own line: the model row renders empty for a log that names no model or provider, so the chevron sat alone above the request id. The sidebar keeps the toggle whenever it is open, in its own header. Collapsed, the toggle moves into the drawer header and joins the model row, or the request id row when there is no model to join. Shared between both through SidebarToggle so the two call sites cannot drift. --- .../LogDetailsDrawer/DrawerHeader.test.tsx | 67 +++++++++++++++++++ .../LogDetailsDrawer/DrawerHeader.tsx | 24 ++++--- .../LogDetailsDrawer/LogDetailsDrawer.tsx | 9 ++- .../LogDetailsDrawer/SidebarToggle.tsx | 21 ++++++ 4 files changed, 108 insertions(+), 13 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/DrawerHeader.test.tsx create mode 100644 ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/SidebarToggle.tsx diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/DrawerHeader.test.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/DrawerHeader.test.tsx new file mode 100644 index 00000000000..a8e27504019 --- /dev/null +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/DrawerHeader.test.tsx @@ -0,0 +1,67 @@ +import { screen, within } from "@testing-library/react"; +import { describe, expect, it, vi } from "vitest"; + +import { render } from "../../../../tests/test-utils"; +import type { LogEntry } from "../columns"; +import { DrawerHeader } from "./DrawerHeader"; + +const logEntry = (overrides: Partial): LogEntry => + ({ + request_id: "170d64ea-69f0-431a-be72-332f8f78c18a", + api_key: "key-1", + team_id: "team-1", + model: "gpt-4o", + model_id: "model-1", + custom_llm_provider: "openai", + call_type: "acompletion", + spend: 0.01, + total_tokens: 10, + prompt_tokens: 5, + completion_tokens: 5, + startTime: "2026-07-07T09:50:13Z", + endTime: "2026-07-07T09:50:14Z", + cache_hit: "false", + messages: [], + response: {}, + ...overrides, + }) as LogEntry; + +const renderHeader = (log: LogEntry, isSidebarCollapsed: boolean) => + render( + , + ); + +const expandToggle = () => screen.getByLabelText("Expand trace sidebar"); + +describe("DrawerHeader sidebar toggle", () => { + it("stays out of the header while the sidebar owns it", () => { + renderHeader(logEntry({}), false); + + expect(screen.queryByLabelText("Expand trace sidebar")).not.toBeInTheDocument(); + expect(screen.queryByLabelText("Collapse trace sidebar")).not.toBeInTheDocument(); + }); + + it("shares the model row once the sidebar is collapsed", () => { + renderHeader(logEntry({}), true); + + const row = expandToggle().parentElement as HTMLElement; + expect(within(row).getByText("gpt-4o")).toBeInTheDocument(); + }); + + it("falls back to the request id row when the log names no model", () => { + renderHeader(logEntry({ model: "", custom_llm_provider: "" }), true); + + const row = expandToggle().parentElement as HTMLElement; + expect(within(row).getByText("170d64ea-69f0-431a-be72-332f8f78c18a")).toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/DrawerHeader.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/DrawerHeader.tsx index d6eb33d0b0a..65b5801602c 100644 --- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/DrawerHeader.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/DrawerHeader.tsx @@ -1,5 +1,5 @@ import { useState } from "react"; -import { Check, ChevronDown, ChevronLeft, ChevronRight, ChevronUp, Copy, X } from "lucide-react"; +import { Check, ChevronDown, ChevronUp, Copy, X } from "lucide-react"; import moment from "moment"; import { Badge } from "@/components/ui/badge"; import { Button } from "@/components/ui/button"; @@ -7,6 +7,7 @@ import { Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from "@/comp import { LogEntry } from "../columns"; import { AutoRouterTag } from "@/components/shared/table_cells"; import { ClassifyTag } from "./ClassifyTag"; +import { SidebarToggle } from "./SidebarToggle"; import { getProviderLogoAndName } from "../../provider_info_helpers"; import { DRAWER_HEADER_PADDING, @@ -47,6 +48,8 @@ export function DrawerHeader({ }: DrawerHeaderProps) { const provider = log.custom_llm_provider || ""; const providerInfo = provider ? getProviderLogoAndName(provider) : null; + const showToggleWithProvider = isSidebarCollapsed && Boolean(providerInfo || log.model); + const showToggleWithRequestId = isSidebarCollapsed && !showToggleWithProvider; return (
{/* Row 0: Model + Provider with Logo */}
- + {showToggleWithProvider && } + {showToggleWithRequestId && }
diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailsDrawer.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailsDrawer.tsx index e25c1a1b8b2..9849c95de71 100644 --- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailsDrawer.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailsDrawer.tsx @@ -8,6 +8,7 @@ import { AGENT_CALL_TYPES, MCP_CALL_TYPES } from "../constants"; import { getEventDisplayName } from "../utils"; import { ClassifyTag } from "./ClassifyTag"; import { DrawerHeader } from "./DrawerHeader"; +import { SidebarToggle } from "./SidebarToggle"; import { useKeyboardNavigation } from "./useKeyboardNavigation"; import { LogDetailContent, GuardrailJumpLink } from "./LogDetailContent"; import { sessionSpendLogsCall } from "../../networking"; @@ -314,8 +315,12 @@ export function LogDetailsDrawer({
{!isSidebarCollapsed && (
-
-
+
+ setIsSidebarCollapsed((collapsed) => !collapsed)} + /> +
{isSessionMode ? "Session" : "Trace"}
diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/SidebarToggle.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/SidebarToggle.tsx new file mode 100644 index 00000000000..b4ee191f4cd --- /dev/null +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/SidebarToggle.tsx @@ -0,0 +1,21 @@ +import { ChevronLeft, ChevronRight } from "lucide-react"; +import { Button } from "@/components/ui/button"; + +export interface SidebarToggleProps { + isCollapsed: boolean; + onToggle: () => void; +} + +export function SidebarToggle({ isCollapsed, onToggle }: SidebarToggleProps) { + return ( + + ); +} From 4a3dcd5e8ebc579e5e3dc5f2bcf00c702899e57d Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sat, 29 Aug 2026 11:53:14 -0700 Subject: [PATCH 48/68] refactor(ui): give the collapsed log drawer its own sidebar toggle Open, the trace sidebar is byte-identical to before: the toggle sits over its header exactly where it did, and the header keeps the padding that makes room for it. Collapsed, that button has nowhere to live, so the drawer header shows one instead, on the model row or the request id row when the log names no model. Both come from SidebarToggle, so they cannot drift in design. The chevrons now point the way the sidebar will move: right while it is open, left while it is collapsed. --- .../LogDetailsDrawer/LogDetailsDrawer.tsx | 160 +++++++++--------- .../LogDetailsDrawer/SidebarToggle.tsx | 10 +- 2 files changed, 88 insertions(+), 82 deletions(-) diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailsDrawer.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailsDrawer.tsx index 9849c95de71..ddd0a650c04 100644 --- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailsDrawer.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailsDrawer.tsx @@ -312,90 +312,94 @@ export function LogDetailsDrawer({ {logEntry?.request_id ? `Request ${logEntry.request_id} details` : "Request details"} -
+
{!isSidebarCollapsed && ( -
-
- setIsSidebarCollapsed((collapsed) => !collapsed)} - /> -
-
- {isSessionMode ? "Session" : "Trace"} + setIsSidebarCollapsed(true)} + className="absolute top-2 left-2 z-raised" + /> + )} + {!isSidebarCollapsed && ( +
+
+
+
+
+ {isSessionMode ? "Session" : "Trace"} +
+
+ {leftPanelDisplayId} + +
-
- {leftPanelDisplayId} - -
-
- {logsForList.length} req - {[ - isSessionMode - ? llmCount - : logsForList.filter( - (row) => - !MCP_CALL_TYPES.includes(row.call_type) && !AGENT_CALL_TYPES.includes(row.call_type), - ).length, - isSessionMode - ? agentCount - : logsForList.filter((row) => AGENT_CALL_TYPES.includes(row.call_type)).length, - isSessionMode - ? mcpCount - : logsForList.filter((row) => MCP_CALL_TYPES.includes(row.call_type)).length, - ].map((count, i) => { - const label = [" LLM", " Agent", " MCP"][i]; - return count > 0 ? ( - - · - {count} - {label} - - ) : null; - })} - · - {isSessionMode ? getSpendString(totalSessionCost) : getSpendString(currentLog.spend || 0)} - {isSessionMode && ( - <> +
+
+ {logsForList.length} req + {[ + isSessionMode + ? llmCount + : logsForList.filter( + (row) => !MCP_CALL_TYPES.includes(row.call_type) && !AGENT_CALL_TYPES.includes(row.call_type), + ).length, + isSessionMode + ? agentCount + : logsForList.filter((row) => AGENT_CALL_TYPES.includes(row.call_type)).length, + isSessionMode + ? mcpCount + : logsForList.filter((row) => MCP_CALL_TYPES.includes(row.call_type)).length, + ].map((count, i) => { + const label = [" LLM", " Agent", " MCP"][i]; + return count > 0 ? ( + · - {sessionDurationSeconds}s - - )} -
+ {count} + {label} + + ) : null; + })} + · + {isSessionMode ? getSpendString(totalSessionCost) : getSpendString(currentLog.spend || 0)} {isSessionMode && ( -
- {cacheHitCount}/{logsForList.length} cached -
- )} - {isSessionMode && sessionTruncated && ( -
- Showing most recent {logsForList.length} of {sessionTotalCount} -
- )} - {isSessionMode && ( - setSessionSortMode(value as SessionLogSortMode)} - > - - - Duration - - - Start time - - - + <> + · + {sessionDurationSeconds}s + )}
+ {isSessionMode && ( +
+ {cacheHitCount}/{logsForList.length} cached +
+ )} + {isSessionMode && sessionTruncated && ( +
+ Showing most recent {logsForList.length} of {sessionTotalCount} +
+ )} + {isSessionMode && ( + setSessionSortMode(value as SessionLogSortMode)} + > + + + Duration + + + Start time + + + + )}
diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/SidebarToggle.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/SidebarToggle.tsx index b4ee191f4cd..a1ba7ff2491 100644 --- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/SidebarToggle.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/SidebarToggle.tsx @@ -1,21 +1,23 @@ import { ChevronLeft, ChevronRight } from "lucide-react"; import { Button } from "@/components/ui/button"; +import { cn } from "@/lib/cva.config"; export interface SidebarToggleProps { isCollapsed: boolean; onToggle: () => void; + className?: string; } -export function SidebarToggle({ isCollapsed, onToggle }: SidebarToggleProps) { +export function SidebarToggle({ isCollapsed, onToggle, className }: SidebarToggleProps) { return ( ); } From 4d4cf403347059c3c1307d817179b809c2a5f74a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 12:06:38 -0700 Subject: [PATCH 49/68] fix(headroom): delegate to the parent deployment hook so deployment-level configs still compress --- .../guardrail_hooks/headroom/headroom.py | 16 +++++----- .../guardrail_hooks/test_headroom.py | 31 ++++++++++++++++++- 2 files changed, 39 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py index 6ae16281983..e2d2fffb2df 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py +++ b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py @@ -719,17 +719,19 @@ class HeadroomGuardrail(CustomGuardrail): async def async_pre_call_deployment_hook( self, - kwargs: Mapping[str, Any], + kwargs: dict[str, Any], call_type: CallTypes | None, ) -> dict[str, Any] | None: # mutable-ok: overrides CustomLogger hook whose contract is a plain dict + base_result: Final = await super().async_pre_call_deployment_hook(kwargs, call_type) + effective: Final = base_result if base_result is not None else kwargs if call_type not in (CallTypes.completion, CallTypes.acompletion): - return None - if not kwargs.get("stream"): - return None - if not has_headroom_retrieve_tool(kwargs.get("tools")): - return None + return base_result + if not effective.get("stream"): + return base_result + if not has_headroom_retrieve_tool(effective.get("tools")): + return base_result return { # mutable-ok: the hook contract is a plain dict the router merges into the request kwargs - **kwargs, + **effective, "stream": False, HEADROOM_CONVERTED_STREAM_KEY: True, } diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py index 52df875e714..22031abad8a 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py @@ -1975,7 +1975,8 @@ async def test_pre_call_deployment_hook_converts_stream_only_for_ccr_chat_comple result = await guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=call_type) if not expect_conversion: - assert result is None + assert result is kwargs + assert HEADROOM_CONVERTED_STREAM_KEY not in kwargs assert kwargs["stream"] is stream return @@ -1985,6 +1986,34 @@ async def test_pre_call_deployment_hook_converts_stream_only_for_ccr_chat_comple assert kwargs["stream"] is True +@pytest.mark.asyncio +async def test_pre_call_deployment_hook_still_compresses_for_deployment_level_configs( + guardrail: HeadroomGuardrail, +): + """Regression for the stream-conversion override swallowing the parent hook: + when the guardrail is attached at the deployment level and proxy pre_call never + ran, the deployment hook is the only place compression executes, so the + override must delegate to CustomGuardrail.async_pre_call_deployment_hook.""" + kwargs = { + "model": "gpt-4o", + "messages": [dict(m) for m in ORIGINAL_MESSAGES], + "stream": False, + "guardrails": ["headroom"], + "metadata": {}, + } + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + return_value=_make_compress_response(COMPRESSED_MESSAGES), + ): + result = await guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.acompletion) + + assert result is not None + assert result["messages"] == EXPECTED_MESSAGES + + @pytest.mark.asyncio async def test_streaming_chat_completion_resolves_ccr_retrieval_end_to_end( guardrail: HeadroomGuardrail, From 4ef012627ddbbc08d87fc85f01367bc66029b7ae Mon Sep 17 00:00:00 2001 From: Mateo Wang Date: Sat, 29 Aug 2026 12:06:42 -0700 Subject: [PATCH 50/68] fix: count error-file failures in the batch cost poller path --- .../proxy/common_utils/check_batch_cost.py | 27 +++- .../proxy_unit_tests/test_check_batch_cost.py | 125 ++++++++++++++++++ 2 files changed, 149 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 163b1592dd9..d34635fa253 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -2,9 +2,10 @@ Polls LiteLLM_ManagedObjectTable to check if the batch job is complete, and if the cost has been tracked. """ +from dataclasses import replace as dataclasses_replace from datetime import datetime, timedelta, timezone from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Dict, Final, List, Optional, Tuple +from typing import TYPE_CHECKING, Any, Dict, Final, List, Literal, Optional, Tuple, cast from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid @@ -626,6 +627,7 @@ 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, ) @@ -761,12 +763,31 @@ class CheckBatchCost: model_id=model_id, deployment_model=litellm_model_name, ) - batch_result = await calculate_batch_cost_and_usage( + batch_file_provider: Final = cast( + Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"], llm_provider + ) + output_file_result: Final = await calculate_batch_cost_and_usage( file_content_dictionary=file_content_as_dict, - custom_llm_provider=llm_provider, # type: ignore + custom_llm_provider=batch_file_provider, model_name=model_name, model_info=deployment_model_info, ) + error_file_failed_requests: Final = await _count_error_file_failed_requests( + response, + custom_llm_provider=batch_file_provider, + litellm_params={ + **credentials, + "_litellm_internal_model_credentials": MappingProxyType(dict(credentials)), + }, + ) + batch_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, + ) + ) logging_obj = LiteLLMLogging( model=batch_result.models[0], messages=[{"role": "user", "content": ""}], diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index 7f4fac4e6e9..bf5add54139 100644 --- a/tests/proxy_unit_tests/test_check_batch_cost.py +++ b/tests/proxy_unit_tests/test_check_batch_cost.py @@ -1373,6 +1373,131 @@ class TestCheckBatchCost: update_data["status"] == terminal_status ), f"billed {terminal_status} batch must keep its real terminal status in the DB" + @pytest.mark.asyncio + async def test_error_file_failures_add_to_failed_request_count( + self, check_batch_cost_instance, mock_prisma_client, mock_llm_router + ): + """OpenAI-shaped providers report per-request failures only in a separate + error file. The poller prices from the output file, so without also counting + the error file's lines, batch_failed_requests on the spend log undercounts: + regression test for the poller path merging error-file failures. + """ + from unittest.mock import patch + + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( + return_value=1 + ) + mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=None + ) + + mock_job = MagicMock() + mock_job.id = "job-error-file-1" + mock_job.unified_object_id = "dW5pZmllZF9iYXRjaF9pZA==" + mock_job.created_by = "user-1" + + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[mock_job] + ) + + mock_response = MagicMock() + mock_response.status = "completed" + mock_response.output_file_id = "file-output-123" + mock_response.error_file_id = "file-error-456" + mock_response.model_dump_json.return_value = ( + '{"id":"batch-1","status":"completed"}' + ) + + mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) + mock_llm_router.get_deployment_credentials_with_provider = MagicMock( + return_value={"api_key": "sk-test"} + ) + + mock_deployment = MagicMock() + mock_deployment.litellm_params.custom_llm_provider = "openai" + mock_deployment.litellm_params.model = "gpt-4" + mock_deployment.model_info.model_dump.return_value = {} + mock_llm_router.get_deployment = MagicMock(return_value=mock_deployment) + + output_file_content = MagicMock() + output_file_content.content = b'{"id":"req-1"}' + error_file_content = MagicMock() + error_file_content.content = ( + b'{"id":"err-1","error":{"message":"rejected"}}\n' + b'{"id":"err-2","error":{"message":"rejected"}}\n\n' + ) + + def _file_content_for(**kwargs): + if kwargs["file_id"] == "file-error-456": + return error_file_content + return output_file_content + + decoded_id = "llm_model_id,model-123;llm_batch_id,batch-456;" + + with ( + patch( + "litellm.proxy.openai_files_endpoints.common_utils._is_base64_encoded_unified_file_id", + side_effect=[decoded_id, None, None], + ), + patch( + "litellm.proxy.openai_files_endpoints.common_utils.get_model_id_from_unified_batch_id", + return_value="model-123", + ), + patch( + "litellm.proxy.openai_files_endpoints.common_utils.get_batch_id_from_unified_batch_id", + return_value="batch-456", + ), + patch( + "litellm.files.main.afile_content", + new_callable=AsyncMock, + side_effect=_file_content_for, + ) as mock_afile_content, + patch( + "litellm.batches.batch_utils._get_file_content_as_dictionary", + return_value=[{"id": "req-1"}], + ), + patch( + "litellm.batches.batch_utils.calculate_batch_cost_and_usage", + new_callable=AsyncMock, + return_value=_batch_cost_result( + 0.01, + {"prompt_tokens": 10, "completion_tokens": 5}, + ["gpt-4"], + successful_requests=3, + failed_requests=1, + ), + ), + patch( + "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", + return_value=("gpt-4", "openai", None, None), + ), + patch( + "litellm.litellm_core_utils.litellm_logging.Logging" + ) as mock_logging_cls, + ): + mock_logging_obj = MagicMock() + mock_logging_obj.async_success_handler = AsyncMock() + mock_logging_cls.return_value = mock_logging_obj + + await check_batch_cost_instance.check_batch_cost() + + assert mock_afile_content.await_count == 2, ( + "the poller must fetch the error file in addition to the output file" + ) + fetched_file_ids = { + call.kwargs["file_id"] for call in mock_afile_content.await_args_list + } + assert fetched_file_ids == {"file-output-123", "file-error-456"} + + mock_logging_obj.async_success_handler.assert_awaited_once() + handler_kwargs = mock_logging_obj.async_success_handler.await_args.kwargs + assert handler_kwargs["batch_successful_requests"] == 3 + assert handler_kwargs["batch_failed_requests"] == 3, ( + "2 error-file lines must add to the output file's 1 failed request" + ) + assert handler_kwargs["batch_cost"] == 0.01 + @pytest.mark.asyncio async def test_terminal_batch_with_missing_output_file_is_retired_unbilled( self, check_batch_cost_instance, mock_prisma_client, mock_llm_router From a3eac3f7716f91ec728848b8200d8384d5ee1f83 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 12:13:09 -0700 Subject: [PATCH 51/68] fix(bedrock): normalize encoding_format base64 to float for cohere embed models --- .../llms/bedrock/embed/cohere_transformation.py | 4 +++- tests/test_litellm/test_utils.py | 16 ++++++++++++++++ 2 files changed, 19 insertions(+), 1 deletion(-) diff --git a/litellm/llms/bedrock/embed/cohere_transformation.py b/litellm/llms/bedrock/embed/cohere_transformation.py index d1c9ceb99d1..8a17bb9d595 100644 --- a/litellm/llms/bedrock/embed/cohere_transformation.py +++ b/litellm/llms/bedrock/embed/cohere_transformation.py @@ -20,7 +20,9 @@ class BedrockCohereEmbeddingConfig: def map_openai_params(self, non_default_params: dict, optional_params: dict) -> dict: for k, v in non_default_params.items(): if k == "encoding_format": - optional_params["embedding_types"] = v if isinstance(v, list) else [v] + optional_params["embedding_types"] = [ + "float" if fmt == "base64" else fmt for fmt in (tuple(v) if isinstance(v, list) else (v,)) + ] elif k == "dimensions": optional_params["output_dimension"] = v return optional_params diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 6381e1a1274..b34d8360183 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -4433,6 +4433,22 @@ class TestBedrockCohereEmbeddingDispatch: ) assert optional_params.get("embedding_types") == ["float"] + @pytest.mark.parametrize( + "model", + [ + "cohere.embed-english-v3", + "cohere.embed-multilingual-v3", + "cohere.embed-v4:0", + ], + ) + def test_cohere_embed_models_map_base64_to_float(self, model): + optional_params = litellm.utils.get_optional_params_embeddings( + model=model, + encoding_format="base64", + custom_llm_provider="bedrock", + ) + assert optional_params.get("embedding_types") == ["float"] + def test_cohere_embed_english_v3_maps_dimensions(self): optional_params = litellm.utils.get_optional_params_embeddings( model="cohere.embed-english-v3", From ae1039f811ad78c2f7228a1e9305bb33ab22546f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 12:29:31 -0700 Subject: [PATCH 52/68] ci: rerun checks on an identical tree after a windows job infra failure From 886d39c3a226869d913a6d82976952ec607b3766 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 12:44:01 -0700 Subject: [PATCH 53/68] test(bedrock): expect cohere embed base64 encoding_format to normalize to float --- tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py b/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py index 114e473be98..9dff24b3977 100644 --- a/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py +++ b/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py @@ -945,7 +945,7 @@ def test_titan_image_embedding_cost_uses_per_image_rate(): "encoding_format,expected_embedding_types", [ ("float", ["float"]), - ("base64", ["base64"]), + ("base64", ["float"]), (["float", "int8"], ["float", "int8"]), ], ) From 2affd800eca0f6c4bdc418cb9154c0b9c5893252 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 12:45:56 -0700 Subject: [PATCH 54/68] test(headroom): cover stream conversion after deployment-level compression --- .../guardrail_hooks/test_headroom.py | 26 +++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py index 22031abad8a..1fbc975e40a 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py @@ -2014,6 +2014,32 @@ async def test_pre_call_deployment_hook_still_compresses_for_deployment_level_co assert result["messages"] == EXPECTED_MESSAGES +@pytest.mark.asyncio +async def test_pre_call_deployment_hook_converts_stream_after_deployment_level_compression( + guardrail: HeadroomGuardrail, +): + kwargs = { + "model": "gpt-4o", + "messages": [dict(m) for m in ORIGINAL_MESSAGES], + "stream": True, + "guardrails": ["headroom"], + "metadata": {}, + } + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + return_value=_make_compress_response(COMPRESSED_MESSAGES_WITH_HASH), + ): + result = await guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.acompletion) + + assert result is not None + assert has_headroom_retrieve_tool(result["tools"]) + assert result["stream"] is False + assert result[HEADROOM_CONVERTED_STREAM_KEY] is True + + @pytest.mark.asyncio async def test_streaming_chat_completion_resolves_ccr_retrieval_end_to_end( guardrail: HeadroomGuardrail, From 99884f0eaa42cd5d8dfd98ee8e411d1e203b88eb Mon Sep 17 00:00:00 2001 From: Mateo Wang Date: Sat, 29 Aug 2026 12:53:52 -0700 Subject: [PATCH 55/68] test: fake the provider file boundary in the poller error-file regression test --- .../proxy_unit_tests/test_check_batch_cost.py | 338 +++++++----------- 1 file changed, 122 insertions(+), 216 deletions(-) diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index bf5add54139..4f9fcf1952b 100644 --- a/tests/proxy_unit_tests/test_check_batch_cost.py +++ b/tests/proxy_unit_tests/test_check_batch_cost.py @@ -1382,37 +1382,31 @@ class TestCheckBatchCost: the error file's lines, batch_failed_requests on the spend log undercounts: regression test for the poller path merging error-file failures. """ + import base64 from unittest.mock import patch - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=1 - ) + import httpx + import respx + + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=1) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( - return_value=None - ) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) mock_job = MagicMock() mock_job.id = "job-error-file-1" - mock_job.unified_object_id = "dW5pZmllZF9iYXRjaF9pZA==" + mock_job.unified_object_id = base64.urlsafe_b64encode( + b"litellm_proxy;model_id:model-123;llm_batch_id:batch-456" + ).decode() mock_job.created_by = "user-1" - - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( - return_value=[mock_job] - ) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) mock_response = MagicMock() mock_response.status = "completed" mock_response.output_file_id = "file-output-123" mock_response.error_file_id = "file-error-456" - mock_response.model_dump_json.return_value = ( - '{"id":"batch-1","status":"completed"}' - ) - + mock_response.model_dump_json.return_value = '{"id":"batch-1","status":"completed"}' mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) - mock_llm_router.get_deployment_credentials_with_provider = MagicMock( - return_value={"api_key": "sk-test"} - ) + mock_llm_router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"}) mock_deployment = MagicMock() mock_deployment.litellm_params.custom_llm_provider = "openai" @@ -1420,83 +1414,73 @@ class TestCheckBatchCost: mock_deployment.model_info.model_dump.return_value = {} mock_llm_router.get_deployment = MagicMock(return_value=mock_deployment) - output_file_content = MagicMock() - output_file_content.content = b'{"id":"req-1"}' - error_file_content = MagicMock() - error_file_content.content = ( - b'{"id":"err-1","error":{"message":"rejected"}}\n' - b'{"id":"err-2","error":{"message":"rejected"}}\n\n' + succeeded_line = json.dumps( + { + "custom_id": "req-1", + "response": { + "status_code": 200, + "body": { + "id": "chatcmpl-1", + "object": "chat.completion", + "model": "gpt-4", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "hi"}, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 10, + "completion_tokens": 5, + "total_tokens": 15, + }, + }, + }, + "error": None, + } + ) + rejected_line = json.dumps( + { + "custom_id": "req-2", + "response": { + "status_code": 400, + "body": {"error": {"message": "bad request"}}, + }, + "error": None, + } + ) + error_file_lines = "\n".join( + json.dumps({"custom_id": custom_id, "error": {"message": "rejected"}}) for custom_id in ("req-3", "req-4") ) - def _file_content_for(**kwargs): - if kwargs["file_id"] == "file-error-456": - return error_file_content - return output_file_content - - decoded_id = "llm_model_id,model-123;llm_batch_id,batch-456;" - with ( - patch( - "litellm.proxy.openai_files_endpoints.common_utils._is_base64_encoded_unified_file_id", - side_effect=[decoded_id, None, None], - ), - patch( - "litellm.proxy.openai_files_endpoints.common_utils.get_model_id_from_unified_batch_id", - return_value="model-123", - ), - patch( - "litellm.proxy.openai_files_endpoints.common_utils.get_batch_id_from_unified_batch_id", - return_value="batch-456", - ), - patch( - "litellm.files.main.afile_content", - new_callable=AsyncMock, - side_effect=_file_content_for, - ) as mock_afile_content, - patch( - "litellm.batches.batch_utils._get_file_content_as_dictionary", - return_value=[{"id": "req-1"}], - ), - patch( - "litellm.batches.batch_utils.calculate_batch_cost_and_usage", - new_callable=AsyncMock, - return_value=_batch_cost_result( - 0.01, - {"prompt_tokens": 10, "completion_tokens": 5}, - ["gpt-4"], - successful_requests=3, - failed_requests=1, - ), - ), - patch( - "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", - return_value=("gpt-4", "openai", None, None), - ), - patch( + respx.mock(assert_all_called=True) as provider, + patch( # test-quality-ok: the poller builds Logging inline, the only seam to its handler kwargs "litellm.litellm_core_utils.litellm_logging.Logging" ) as mock_logging_cls, ): + 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( + return_value=httpx.Response(200, content=f"{error_file_lines}\n\n".encode()) + ) mock_logging_obj = MagicMock() mock_logging_obj.async_success_handler = AsyncMock() mock_logging_cls.return_value = mock_logging_obj await check_batch_cost_instance.check_batch_cost() - assert mock_afile_content.await_count == 2, ( - "the poller must fetch the error file in addition to the output file" - ) - fetched_file_ids = { - call.kwargs["file_id"] for call in mock_afile_content.await_args_list - } - assert fetched_file_ids == {"file-output-123", "file-error-456"} - mock_logging_obj.async_success_handler.assert_awaited_once() handler_kwargs = mock_logging_obj.async_success_handler.await_args.kwargs - assert handler_kwargs["batch_successful_requests"] == 3 + assert handler_kwargs["batch_successful_requests"] == 1 assert handler_kwargs["batch_failed_requests"] == 3, ( - "2 error-file lines must add to the output file's 1 failed request" + "2 error-file lines must add to the output file's 1 rejected request" ) - assert handler_kwargs["batch_cost"] == 0.01 + assert handler_kwargs["batch_models"] == ["gpt-4"] + assert handler_kwargs["batch_usage"].total_tokens == 15 + assert handler_kwargs["batch_cost"] > 0 @pytest.mark.asyncio async def test_terminal_batch_with_missing_output_file_is_retired_unbilled( @@ -1513,13 +1497,9 @@ class TestCheckBatchCost: from litellm.exceptions import NotFoundError - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=1 - ) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=1) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( - return_value=None - ) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) mock_job = MagicMock() mock_job.id = "job-output-gone-1" @@ -1529,23 +1509,17 @@ class TestCheckBatchCost: mock_job.created_by = "user-1" assert check_batch_cost_instance._has_batch_processed_column is True - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( - return_value=[mock_job] - ) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) missing_output_file_id = "gs://batch-out/job-1/predictions.jsonl" mock_response = MagicMock() mock_response.status = "failed" mock_response.output_file_id = missing_output_file_id mock_response.error_file_id = None - mock_response.model_dump_json.return_value = ( - '{"id":"batch-1","status":"failed"}' - ) + mock_response.model_dump_json.return_value = '{"id":"batch-1","status":"failed"}' mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) - mock_llm_router.get_deployment_credentials_with_provider = MagicMock( - return_value={"api_key": "sk-test"} - ) + mock_llm_router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"}) with ( patch( @@ -1566,12 +1540,10 @@ class TestCheckBatchCost: assert mock_afile_content.await_count == 1 mock_calculate.assert_not_awaited() - assert ( - mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1 - ), "a terminal batch with a 404ing output file must be retired, not retried forever" - update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[ - 1 - ]["data"] + assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1, ( + "a terminal batch with a 404ing output file must be retired, not retried forever" + ) + update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[1]["data"] assert update_data["status"] == "failed" assert update_data["batch_processed"] is True @@ -1584,13 +1556,9 @@ class TestCheckBatchCost: Without this, GET /batches/{id} returns a raw file ID that cannot be routed through the proxy, causing API_KEY errors when clients call GET /files/{id}/content. """ - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( - return_value=1 - ) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=1) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( - return_value=None - ) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) mock_job = MagicMock() mock_job.id = "job-raw-file-1" @@ -1599,9 +1567,7 @@ class TestCheckBatchCost: mock_job.team_id = None check_batch_cost_instance._has_batch_processed_column = True - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( - return_value=[mock_job] - ) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) raw_output_file_id = "file-batch-output-abc123" raw_error_file_id = "file-batch-error-xyz456" @@ -1612,14 +1578,10 @@ class TestCheckBatchCost: mock_response.status = "completed" mock_response.output_file_id = raw_output_file_id mock_response.error_file_id = raw_error_file_id - mock_response.model_dump_json.return_value = ( - '{"id":"batch-1","status":"completed"}' - ) + mock_response.model_dump_json.return_value = '{"id":"batch-1","status":"completed"}' mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) - mock_llm_router.get_deployment_credentials_with_provider = MagicMock( - return_value={"api_key": "sk-test"} - ) + mock_llm_router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"}) mock_deployment = MagicMock() mock_deployment.litellm_params.custom_llm_provider = "azure" @@ -1634,9 +1596,7 @@ class TestCheckBatchCost: fake_managed_error_id, ] mock_hook.store_unified_file_id = AsyncMock() - check_batch_cost_instance.proxy_logging_obj.get_proxy_hook.return_value = ( - mock_hook - ) + check_batch_cost_instance.proxy_logging_obj.get_proxy_hook.return_value = mock_hook mock_file_content = MagicMock() mock_file_content.content = b'{"id":"req-1"}' @@ -1679,9 +1639,7 @@ class TestCheckBatchCost: "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", return_value=("gpt-5-mini", "azure", None, None), ), - patch( - "litellm.litellm_core_utils.litellm_logging.Logging" - ) as mock_logging_cls, + patch("litellm.litellm_core_utils.litellm_logging.Logging") as mock_logging_cls, ): mock_logging_obj = MagicMock() mock_logging_obj.async_success_handler = AsyncMock() @@ -1732,9 +1690,7 @@ class TestUnmanagedVertexRouting: def _job(self, file_object=None): job = MagicMock() job.unified_object_id = "8823717160934178816" - job.file_object = ( - file_object if file_object is not None else _unmanaged_vertex_file_object() - ) + job.file_object = file_object if file_object is not None else _unmanaged_vertex_file_object() return job def test_flag_off_skips_unmanaged_id_unchanged(self): @@ -1772,9 +1728,7 @@ class TestUnmanagedVertexRouting: assert result == ("deploy-1", "8823717160934178816") # bare model name (trailing GCS segment), not the full publishers/.. path - router.resolve_model_name_from_model_id.assert_called_once_with( - "gemini-2.5-flash" - ) + router.resolve_model_name_from_model_id.assert_called_once_with("gemini-2.5-flash") router.get_model_ids.assert_called_once_with(model_name="gemini-2.5-flash") def test_flag_on_skips_non_vertex_deployment_sharing_model_group(self): @@ -1794,9 +1748,7 @@ class TestUnmanagedVertexRouting: result = instance._resolve_job_routing(self._job(), prom) assert result is None - prom.record_check_batch_cost_error.assert_called_once_with( - "unmanaged_no_matching_deployment" - ) + prom.record_check_batch_cost_error.assert_called_once_with("unmanaged_no_matching_deployment") def test_flag_on_uses_later_vertex_deployment_with_matching_suffix(self): router = MagicMock() @@ -1844,9 +1796,7 @@ class TestUnmanagedVertexRouting: result = instance._resolve_job_routing(self._job(), prom) assert result is None - prom.record_check_batch_cost_error.assert_called_once_with( - "unmanaged_no_matching_deployment" - ) + prom.record_check_batch_cost_error.assert_called_once_with("unmanaged_no_matching_deployment") def test_flag_on_non_gcs_input_is_not_unmanaged_vertex(self): """Flag on, but input_file_id is not a gs:// publishers path: treat as unroutable, @@ -1854,9 +1804,7 @@ class TestUnmanagedVertexRouting: router = MagicMock() instance = self._instance(track_unmanaged=True, router=router) prom = MagicMock() - job = self._job( - file_object=_unmanaged_vertex_file_object(input_file_id="file-abc-123") - ) + job = self._job(file_object=_unmanaged_vertex_file_object(input_file_id="file-abc-123")) with patch(_IS_B64, return_value=False): result = instance._resolve_job_routing(job, prom) @@ -1880,9 +1828,7 @@ class TestUnmanagedVertexRouting: mock_response.error_file_id = None mock_response.completed_at = None mock_response.created_at = None - mock_response.model_dump_json.return_value = ( - '{"id":"8823717160934178816","status":"completed"}' - ) + mock_response.model_dump_json.return_value = '{"id":"8823717160934178816","status":"completed"}' router.aretrieve_batch = AsyncMock(return_value=mock_response) router.get_deployment_credentials_with_provider = MagicMock( return_value={"vertex_project": "p", "vertex_location": "us-central1"} @@ -1904,9 +1850,7 @@ class TestUnmanagedVertexRouting: prisma.db.litellm_managedobjecttable = MagicMock() prisma.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=1) prisma.db.litellm_managedobjecttable.update = AsyncMock() - prisma.db.litellm_managedobjecttable.find_many = AsyncMock( - return_value=[self._job()] - ) + prisma.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[self._job()]) prisma.db.litellm_usertable = MagicMock() prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=None) @@ -1937,9 +1881,7 @@ class TestUnmanagedVertexRouting: "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", return_value=("gemini-2.5-flash", "vertex_ai", None, None), ), - patch( - "litellm.litellm_core_utils.litellm_logging.Logging" - ) as mock_logging_cls, + patch("litellm.litellm_core_utils.litellm_logging.Logging") as mock_logging_cls, ): mock_logging_obj = MagicMock() mock_logging_obj.async_success_handler = AsyncMock() @@ -1980,9 +1922,7 @@ class TestUnmanagedBedrockRouting: def _job(self, file_object=None): job = MagicMock() job.unified_object_id = self._ARN - job.file_object = ( - file_object if file_object is not None else _unmanaged_bedrock_file_object() - ) + job.file_object = file_object if file_object is not None else _unmanaged_bedrock_file_object() return job def _bedrock_deployment(self): @@ -2037,9 +1977,7 @@ class TestUnmanagedBedrockRouting: result = instance._resolve_job_routing(self._job(), prom) assert result is None - prom.record_check_batch_cost_error.assert_called_once_with( - "unmanaged_no_matching_deployment" - ) + prom.record_check_batch_cost_error.assert_called_once_with("unmanaged_no_matching_deployment") def test_flag_on_matches_deployment_despite_colon_dash_mismatch(self): """The S3 object key has ':' replaced with '-' (e.g. 'v1-0'), but the configured @@ -2077,9 +2015,7 @@ class TestUnmanagedBedrockRouting: result = instance._resolve_job_routing(self._job(), prom) assert result is None - prom.record_check_batch_cost_error.assert_called_once_with( - "unmanaged_no_matching_deployment" - ) + prom.record_check_batch_cost_error.assert_called_once_with("unmanaged_no_matching_deployment") def test_flag_on_non_s3_input_is_not_unmanaged_bedrock(self): """Flag on, but input_file_id is not a litellm-bedrock-files- s3:// key: treat as @@ -2087,9 +2023,7 @@ class TestUnmanagedBedrockRouting: router = MagicMock() instance = self._instance(track_unmanaged=True, router=router) prom = MagicMock() - job = self._job( - file_object=_unmanaged_bedrock_file_object(input_file_id="file-abc-123") - ) + job = self._job(file_object=_unmanaged_bedrock_file_object(input_file_id="file-abc-123")) with patch(_IS_B64, return_value=False): result = instance._resolve_job_routing(job, prom) @@ -2112,13 +2046,9 @@ class TestUnmanagedBedrockRouting: mock_response.error_file_id = None mock_response.completed_at = None mock_response.created_at = None - mock_response.model_dump_json.return_value = ( - f'{{"id":"{self._ARN}","status":"completed"}}' - ) + mock_response.model_dump_json.return_value = f'{{"id":"{self._ARN}","status":"completed"}}' router.aretrieve_batch = AsyncMock(return_value=mock_response) - router.get_deployment_credentials_with_provider = MagicMock( - return_value={"aws_region_name": "us-east-1"} - ) + router.get_deployment_credentials_with_provider = MagicMock(return_value={"aws_region_name": "us-east-1"}) deployment = self._bedrock_deployment() deployment.model_name = "claude-sonnet-4" @@ -2134,9 +2064,7 @@ class TestUnmanagedBedrockRouting: prisma.db.litellm_managedobjecttable = MagicMock() prisma.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=1) prisma.db.litellm_managedobjecttable.update = AsyncMock() - prisma.db.litellm_managedobjecttable.find_many = AsyncMock( - return_value=[self._job()] - ) + prisma.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[self._job()]) prisma.db.litellm_usertable = MagicMock() prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=None) @@ -2167,9 +2095,7 @@ class TestUnmanagedBedrockRouting: "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", return_value=("claude-sonnet-4", "bedrock", None, None), ), - patch( - "litellm.litellm_core_utils.litellm_logging.Logging" - ) as mock_logging_cls, + patch("litellm.litellm_core_utils.litellm_logging.Logging") as mock_logging_cls, ): mock_logging_obj = MagicMock() mock_logging_obj.async_success_handler = AsyncMock() @@ -2293,9 +2219,7 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: ) router = MagicMock() - router.get_deployment_credentials_with_provider = MagicMock( - return_value={"api_key": "sk-test"} - ) + router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"}) deployment = MagicMock() deployment.litellm_params.custom_llm_provider = "azure" deployment.litellm_params.model = "azure/gpt-5.5" @@ -2304,8 +2228,8 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: router.get_deployment = MagicMock(return_value=deployment) hook = MagicMock() - hook.get_unified_output_file_id = ( - lambda output_file_id, model_id, model_name: _PROXY_LiteLLMManagedFiles.get_unified_output_file_id( + hook.get_unified_output_file_id = lambda output_file_id, model_id, model_name: ( + _PROXY_LiteLLMManagedFiles.get_unified_output_file_id( None, output_file_id=output_file_id, model_id=model_id, model_name=model_name ) ) @@ -2374,9 +2298,7 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: get_models_from_unified_file_id, ) - output_file_id = await self._run( - self._job(self._managed_input_file_id(self._PUBLIC_MODEL_GROUP)) - ) + output_file_id = await self._run(self._job(self._managed_input_file_id(self._PUBLIC_MODEL_GROUP))) decoded = _is_base64_encoded_unified_file_id(output_file_id) assert get_models_from_unified_file_id(decoded) == [self._PUBLIC_MODEL_GROUP] @@ -2390,9 +2312,7 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: _extract_models_from_managed_resource_id, ) - output_file_id = await self._run( - self._job(self._managed_input_file_id(self._PUBLIC_MODEL_GROUP)) - ) + output_file_id = await self._run(self._job(self._managed_input_file_id(self._PUBLIC_MODEL_GROUP))) models = _extract_models_from_managed_resource_id(output_file_id, "file_id", None) assert models == [self._PUBLIC_MODEL_GROUP] @@ -2400,9 +2320,7 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: await can_key_call_model( model=models[0], llm_model_list=None, - valid_token=UserAPIKeyAuth( - api_key="sk-test", models=[self._PUBLIC_MODEL_GROUP] - ), + valid_token=UserAPIKeyAuth(api_key="sk-test", models=[self._PUBLIC_MODEL_GROUP]), llm_router=None, ) is True @@ -2419,6 +2337,8 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: decoded = _is_base64_encoded_unified_file_id(output_file_id) assert get_models_from_unified_file_id(decoded) == [self._PUBLIC_MODEL_GROUP] + + class TestBatchCostAttribution: """CheckBatchCost rebuilds the creator's spend metadata from the managed-object row so the batch-cost log is attributed like a non-batch request.""" @@ -2514,9 +2434,7 @@ class TestBatchCostAttribution: """An alias lookup failure must not lose the spend row; the key hash and team still attribute it.""" instance = self._instance() - instance.prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( - side_effect=Exception("db down") - ) + instance.prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(side_effect=Exception("db down")) metadata = await instance._build_creator_attribution_metadata(self._job(), "batch-1") @@ -2612,9 +2530,7 @@ class TestPollPageStarvation: async def test_unified_id_without_model_id_is_retired(self): """A unified id that decodes but carries no model_id is unroutable no matter what the config says, so it must leave the poll page instead of being retried forever.""" - prisma = self._prisma( - [self._job("job-no-model", self._encode("litellm_proxy;llm_batch_id:poison-no-model"))] - ) + prisma = self._prisma([self._job("job-no-model", self._encode("litellm_proxy;llm_batch_id:poison-no-model"))]) llm_router = MagicMock() llm_router.aretrieve_batch = AsyncMock() @@ -2652,9 +2568,7 @@ class TestPollPageStarvation: await self._instance(prisma, llm_router).check_batch_cost() prisma.db.litellm_managedobjecttable.update.assert_awaited_once() - assert prisma.db.litellm_managedobjecttable.update.call_args[1]["data"] == { - "batch_processed": True - } + assert prisma.db.litellm_managedobjecttable.update.call_args[1]["data"] == {"batch_processed": True} @pytest.mark.asyncio async def test_provider_404_with_deployment_gone_keeps_job(self): @@ -2707,17 +2621,13 @@ class TestPollPageStarvation: async def test_retirement_falls_back_to_status_without_batch_processed_column(self): """Older schemas have no batch_processed column, so the only way to stop selecting the row is the status filter the poll query already applies.""" - prisma = self._prisma( - [self._job("job-legacy", self._encode("litellm_proxy;llm_batch_id:poison-no-model"))] - ) + prisma = self._prisma([self._job("job-legacy", self._encode("litellm_proxy;llm_batch_id:poison-no-model"))]) instance = self._instance(prisma, MagicMock()) instance._has_batch_processed_column = False await instance.check_batch_cost() - assert prisma.db.litellm_managedobjecttable.update.call_args[1]["data"] == { - "status": "stale_expired" - } + assert prisma.db.litellm_managedobjecttable.update.call_args[1]["data"] == {"status": "stale_expired"} @pytest.mark.asyncio async def test_stale_cleanup_gives_up_on_never_costed_completed_rows(self): @@ -2772,14 +2682,11 @@ class TestPollPageStarvation: await self._instance(prisma, llm_router).check_batch_cost() - retired = [ - call[1]["where"]["id"] - for call in prisma.db.litellm_managedobjecttable.update.call_args_list - ] + retired = [call[1]["where"]["id"] for call in prisma.db.litellm_managedobjecttable.update.call_args_list] assert retired == ["job-no-model", "job-gone"] - assert ( - llm_router.aretrieve_batch.await_args_list[-1][1]["batch_id"] == "batch_live" - ), "the newer healthy batch must still be polled in the same cycle" + assert llm_router.aretrieve_batch.await_args_list[-1][1]["batch_id"] == "batch_live", ( + "the newer healthy batch must still be polled in the same cycle" + ) @pytest.mark.asyncio async def test_404_that_does_not_name_the_batch_keeps_job_for_retry(self): @@ -2808,6 +2715,7 @@ class TestPollPageStarvation: prisma.db.litellm_managedobjecttable.update.assert_not_awaited() + class _FakeManagedObjectRow: """One managed batch row the provider has finished but nothing has costed yet.""" @@ -2824,8 +2732,12 @@ class _FakeManagedObjectRow: self.request_tags = None self.created_at = 1700000000 self.file_object = json.dumps( - {"id": "batch-456", "status": "in_progress", "input_file_id": "file-input-1", - "output_file_id": _CLAIM_OUTPUT_FILE_ID} + { + "id": "batch-456", + "status": "in_progress", + "input_file_id": "file-input-1", + "output_file_id": _CLAIM_OUTPUT_FILE_ID, + } ) @@ -2925,9 +2837,7 @@ class TestMultiPodBatchCostClaim: router = MagicMock() router.aretrieve_batch = AsyncMock(return_value=response) - router.get_deployment_credentials_with_provider = MagicMock( - return_value={"api_key": "sk-test"} - ) + router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"}) router.get_deployment = MagicMock(return_value=deployment) return router @@ -3087,9 +2997,7 @@ class TestMultiPodBatchCostClaim: await asyncio.Event().wait() with self._billing_patches(journal, during_fetch=_never_returns) as logging_obj: - interrupted = asyncio.create_task( - self._instance(prisma, self._router()).check_batch_cost() - ) + interrupted = asyncio.create_task(self._instance(prisma, self._router()).check_batch_cost()) await asyncio.wait_for(reached_fetch.wait(), timeout=5) assert row.batch_processed is False, "an in-flight costing must not mark the row processed" interrupted.cancel() @@ -3124,9 +3032,7 @@ class TestMultiPodBatchCostClaim: await finish_fetch.wait() with self._billing_patches(journal, during_fetch=_wait_for_the_delete_attempt): - costing = asyncio.create_task( - self._instance(prisma, self._router()).check_batch_cost() - ) + costing = asyncio.create_task(self._instance(prisma, self._router()).check_batch_cost()) await asyncio.wait_for(reached_fetch.wait(), timeout=5) with pytest.raises(HTTPException) as blocked: From f0340fef16e04f097d960d20422244ff5e2537e7 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 29 Aug 2026 12:56:04 -0700 Subject: [PATCH 56/68] feat(mcp_gateway): add RFC 7662 introspection for gateway session tokens (#38726) * feat(mcp_gateway): add RFC 7662 introspection for gateway session tokens Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(mcp_gateway): omit Bearer token_type for refresh introspection and allow mcp-scoped keys Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(mcp_gateway): cover introspection of RS256-signed session tokens Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(mcp_gateway): load the discoverable router on a cold /introspect request Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * chore(openapi): regenerate lazy snapshot and schema.d.ts for /introspect Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../mcp_server/discoverable_endpoints.py | 25 +++- .../mcp_server/gateway_dcr_flow.py | 103 +++++++++++++ .../outbound_credentials/session_token.py | 12 +- litellm/proxy/_lazy_features.py | 1 + litellm/proxy/_lazy_openapi_snapshot.json | 58 +++++++ litellm/proxy/_types.py | 1 + .../mcp_server/test_discoverable_endpoints.py | 79 ++++++++++ .../mcp_server/test_gateway_dcr_flow.py | 141 +++++++++++++++++- ui/litellm-dashboard/src/lib/http/schema.d.ts | 62 ++++++++ 9 files changed, 475 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 93b85edd88d..eff467072a8 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -9,7 +9,7 @@ from typing import TYPE_CHECKING, Any, Final, Literal, Optional from urllib.parse import parse_qsl, urlencode, urlparse, urlunparse import httpx -from fastapi import APIRouter, Form, HTTPException, Request +from fastapi import APIRouter, Depends, Form, HTTPException, Request from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse, Response from pydantic import BaseModel, ConfigDict, Field, SecretStr, ValidationError @@ -46,6 +46,7 @@ from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import ( aggregate_authorize, aggregate_token, complete_connect_flow, + introspect_gateway_token, is_gateway_dcr_client_id, is_proxy_api_resource, native_client_auth_contract, @@ -67,6 +68,7 @@ from litellm.proxy._experimental.mcp_server.proxy_api_credentials import ( mint_proxy_credential, ) from litellm.proxy.auth.ip_address_utils import IPAddressUtils +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.encrypt_decrypt_utils import ( decrypt_value_helper, encrypt_value_helper, @@ -1951,6 +1953,26 @@ async def revoke_endpoint(request: Request, token: str = Form(...), client_id: s return await revoke_refresh_token(token=token, client_id=client_id, master_key=master_key, cache=user_api_key_cache) +@router.post("/introspect", dependencies=[Depends(user_api_key_auth)]) +async def introspect_endpoint(token: str = Form(...)) -> Response: + """RFC 7662 introspection for gateway-issued session tokens (``llm_session_`` / + ``llm_srefresh_``), so an external gateway can validate them without the signing + secret. The caller authenticates with a LiteLLM virtual key (section 2.1, enforced by + the route dependency); any token the gateway cannot vouch for answers + ``{"active": false}`` with no further detail.""" + from litellm.proxy.proxy_server import ( # noqa: PLC0415 # circular import at module load + master_key, + user_api_key_cache, + ) + + return await introspect_gateway_token( + token=token, + master_key=master_key, + reload_user=_reload_active_user_by_id, + cache=user_api_key_cache, + ) + + @router.get("/.well-known/litellm-cli-auth") async def native_client_auth_discovery(request: Request) -> JSONResponse: """The versioned contract a native client (``lite login --pkce``, or a CLI in any other @@ -2456,6 +2478,7 @@ def _build_aggregate_authorization_server_response(request: Request) -> dict: "issuer": f"{request_base_url}/mcp", "authorization_endpoint": f"{request_base_url}/authorize", "token_endpoint": f"{request_base_url}/token", + "introspection_endpoint": f"{request_base_url}/introspect", "registration_endpoint": f"{request_base_url}/register", "response_types_supported": ["code"], "scopes_supported": [], diff --git a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py index 853a07972c1..a43e762a456 100644 --- a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py +++ b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py @@ -70,13 +70,19 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credent open_session_refresh_bearer, ) from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token import ( + SESSION_ISSUER, SESSION_REFRESH_TTL_SECONDS, MintedSessionToken, + OpenedSessionToken, SessionAudience, SessionPrincipal, SessionSigningKeys, + is_session_refresh_token, + is_session_token, mint_session_refresh_token, mint_session_token, + open_session_refresh_token, + open_session_token, ) from litellm.proxy.common_utils.encrypt_decrypt_utils import ( decrypt_value_helper, @@ -885,6 +891,23 @@ class _SingleUseGuard: count = await self._cache.async_increment_cache(key, 1, ttl=ttl_seconds, local_only=True) return "first" if count == 1 else "replayed" + async def peek(self, key: str) -> Literal["unclaimed", "claimed", "unavailable"]: + """Read-only view of a single-use marker, resolved against the same shared authority as + :meth:`claim` so introspection observes exactly the record redemption and revocation wrote. + A backend fault is ``"unavailable"`` (fail closed) rather than a guess either way.""" + from litellm.proxy.proxy_server import redis_usage_cache # noqa: PLC0415 # circular import at module load + + redis_cache: Final = redis_usage_cache or getattr(self._cache, "redis_cache", None) + if redis_cache is not None: + try: + value = await redis_cache.async_get_cache(key) + except Exception as e: # noqa: BLE001 # ANY Redis fault fails the read closed + verbose_logger.warning("mcp gateway single-use peek: shared cache backend unavailable: %s", e) + return "unavailable" + return "unclaimed" if value is None else "claimed" + local: Final = await self._cache.async_get_cache(key, local_only=True) + return "unclaimed" if local is None else "claimed" + def _session_token_pair(principal: SessionPrincipal, keys: SessionSigningKeys, now: datetime) -> Response: access: Final = mint_session_token(principal, keys, now) @@ -1199,3 +1222,83 @@ async def revoke_refresh_token(token: str, client_id: str, master_key: str | Non if burned == "unavailable": return _oauth_error(503, "temporarily_unavailable", _CLAIM_UNAVAILABLE_DESCRIPTION) return Response(content="{}", media_type="application/json", headers=TOKEN_NO_CACHE_HEADERS) + + +def _inactive_introspection_response() -> Response: + """RFC 7662 section 2.2: any token the gateway cannot vouch for, whatever the reason + (wrong family, bad signature, expired, revoked, or a deactivated user), answers 200 + with ``active: false`` and nothing else, so introspection is not a token oracle.""" + return JSONResponse(status_code=200, content={"active": False}, headers=TOKEN_NO_CACHE_HEADERS) + + +def _active_introspection_response(opened: OpenedSessionToken) -> Response: + principal: Final = opened.principal + optional_claims: Final = { + key: value + for key, value in ( + ("token_type", "Bearer" if opened.kind == "session" else None), + ("team_id", principal.team_id), + ("resource_server_id", principal.resource_server_id), + ("audience", principal.audience), + ) + if value is not None + } + return JSONResponse( + status_code=200, + content={ + "active": True, + "iss": SESSION_ISSUER, + "sub": principal.user_id, + "client_id": principal.client_id, + "jti": opened.jti, + "iat": opened.iat, + "exp": opened.exp, + "kind": opened.kind, + **optional_claims, + }, + headers=TOKEN_NO_CACHE_HEADERS, + ) + + +async def introspect_gateway_token( + token: str, + master_key: str | None, + reload_user: ReloadUser, + cache: DualCache, +) -> Response: + """RFC 7662 introspection for the gateway's session tokens, so an external gateway + (Kong, an API management layer) can validate a LiteLLM-issued MCP session credential + without holding the signing secret. The caller is already authenticated by the route + (section 2.1). Active means everything admission itself would require: valid signature + under the configured session signing keys, unexpired, not a revoked or rotated refresh + token, and a litellm user that is still live, so a deactivated user's outstanding + tokens introspect as inactive immediately. A shared-backend or DB outage answers 503 + rather than guessing in either direction.""" + if master_key is None: + verbose_logger.error("mcp_gateway_dcr introspect rejected: no master_key configured") + return _oauth_error(500, "server_error", "the gateway has no master key configured") + keys: Final = active_session_signing_keys(master_key) + if isinstance(keys, SessionSigningConfigError): + verbose_logger.error("mcp_gateway_dcr introspect rejected: %s", keys.detail) + return _oauth_error(500, "server_error", keys.detail) + now: Final = datetime.now(timezone.utc) + if is_session_token(token): + opened = open_session_token(token, keys, now) + elif is_session_refresh_token(token): + opened = open_session_refresh_token(token, keys, now) + else: + return _inactive_introspection_response() + if not isinstance(opened, OpenedSessionToken): + return _inactive_introspection_response() + if opened.kind == "session_refresh": + peeked: Final = await _SingleUseGuard(cache).peek(f"{_USED_REFRESH_CACHE_PREFIX}{opened.jti}") + if peeked == "unavailable": + return _oauth_error(503, "temporarily_unavailable", _CLAIM_UNAVAILABLE_DESCRIPTION) + if peeked == "claimed": + return _inactive_introspection_response() + failure: Final = await reload_user(opened.principal.user_id) + if failure == "unavailable": + return _oauth_error(503, "temporarily_unavailable", "the gateway database is unavailable; retry") + if failure is not None: + return _inactive_introspection_response() + return _active_introspection_response(opened) diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_token.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_token.py index 6824f96f927..0fa750a4c4a 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_token.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_token.py @@ -221,12 +221,17 @@ class MintedSessionToken(BaseModel): class OpenedSessionToken(BaseModel): - """A validated session token of either kind: the principal it was minted for, plus the - ``jti`` so the token endpoint can enforce single-use rotation on a refresh token.""" + """A validated session token of either kind: the principal it was minted for, the + ``jti`` so the token endpoint can enforce single-use rotation on a refresh token, and + the signed ``kind``/``iat``/``exp`` so an introspection response can report the + token's metadata without re-decoding.""" model_config = ConfigDict(frozen=True) principal: SessionPrincipal jti: str + kind: SessionTokenKind + iat: int + exp: int class SessionTokenTooLarge(BaseModel): @@ -458,6 +463,9 @@ def _open( team_id=claims.team_id, ), jti=claims.jti, + kind=claims.kind, + iat=claims.iat, + exp=claims.exp, ) diff --git a/litellm/proxy/_lazy_features.py b/litellm/proxy/_lazy_features.py index c435234cbbc..3f90e6c0a7a 100644 --- a/litellm/proxy/_lazy_features.py +++ b/litellm/proxy/_lazy_features.py @@ -161,6 +161,7 @@ LAZY_FEATURES: Final[tuple[LazyFeature, ...]] = ( "/callback", "/register", "/revoke", + "/introspect", ), # Catches the /{mcp_server_name}/authorize|token|register variants. path_suffixes=("/authorize", "/token", "/register"), diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index c1e89f8aa75..b5f4fa01a7b 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -16555,6 +16555,19 @@ "title": "Body_authorize_complete_authorize_complete_post", "type": "object" }, + "Body_introspect_endpoint_introspect_post": { + "properties": { + "token": { + "title": "Token", + "type": "string" + } + }, + "required": [ + "token" + ], + "title": "Body_introspect_endpoint_introspect_post", + "type": "object" + }, "Body_revoke_endpoint_revoke_post": { "properties": { "client_id": { @@ -19134,6 +19147,51 @@ ] } }, + "/introspect": { + "post": { + "description": "RFC 7662 introspection for gateway-issued session tokens (``llm_session_`` /\n``llm_srefresh_``), so an external gateway can validate them without the signing\nsecret. The caller authenticates with a LiteLLM virtual key (section 2.1, enforced by\nthe route dependency); any token the gateway cannot vouch for answers\n``{\"active\": false}`` with no further detail.", + "operationId": "introspect_endpoint_introspect_post", + "requestBody": { + "content": { + "application/x-www-form-urlencoded": { + "schema": { + "$ref": "#/components/schemas/Body_introspect_endpoint_introspect_post" + } + } + }, + "required": true + }, + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Introspect Endpoint", + "tags": [ + "mcp_discoverable" + ] + } + }, "/register": { "post": { "operationId": "register_client_register_post", diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index f40b0632398..21b3a210877 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -504,6 +504,7 @@ class LiteLLMRoutes(enum.Enum): "/mcp-rest/tools/list", "/mcp-rest/tools/call", "/v1/mcp/tools", + "/introspect", ] # MCP server CRUD routes — control-plane. Gated by DISABLE_ADMIN_ENDPOINTS. diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index 0c809940b84..3279c59acd4 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -10290,3 +10290,82 @@ def test_native_client_authorize_without_the_proxy_resource_keeps_the_mcp_flow(m assert 'name="decision"' not in response.text assert "team-b" not in response.text assert minted == [] + + +def test_introspect_route_requires_virtual_key_auth_and_is_advertised(): + """RFC 7662 section 2.1: introspection must not be anonymous. Pins the route-level + user_api_key_auth dependency (structure, so removing it fails here without a proxy), + and that the aggregate AS metadata advertises the endpoint for discovery.""" + from fastapi import FastAPI + from fastapi.routing import APIRoute + from fastapi.testclient import TestClient + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import router + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + route = next(r for r in router.routes if isinstance(r, APIRoute) and r.path == "/introspect") + assert route.methods == {"POST"} + assert any(dependency.call is user_api_key_auth for dependency in route.dependant.dependencies) + + from litellm.proxy._types import LiteLLMRoutes + + assert "/introspect" in LiteLLMRoutes.mcp_routes.value + + from litellm.proxy._lazy_features import LAZY_FEATURES + + discoverable = next(feature for feature in LAZY_FEATURES if feature.name == "mcp_discoverable") + assert "/introspect" in discoverable.path_prefixes + + app = FastAPI() + app.include_router(router) + client = TestClient(app) + asm = client.get("/.well-known/oauth-authorization-server/mcp") + assert asm.json()["introspection_endpoint"] == "http://testserver/introspect" + + +def test_introspect_route_answers_for_authenticated_caller(monkeypatch): + """End-to-end over the real route with the auth dependency satisfied: a garbage token + is active false, a freshly minted session access token is active true with its claims.""" + from datetime import datetime, timezone + + from fastapi import FastAPI + from fastapi.testclient import TestClient + + from litellm.proxy._experimental.mcp_server.discoverable_endpoints import router + from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credentials import ( + session_keys_from_master_key, + ) + from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token import ( + SessionPrincipal, + mint_session_token, + ) + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + introspect_master_key = "sk-introspect-route-test" + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", introspect_master_key, raising=False) + + async def fake_reload(user_id: str): + return None + + monkeypatch.setattr( + "litellm.proxy._experimental.mcp_server.discoverable_endpoints._reload_active_user_by_id", fake_reload + ) + app = FastAPI() + app.include_router(router) + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth() + client = TestClient(app) + + garbage = client.post("/introspect", data={"token": "llm_session_garbage"}) + assert garbage.status_code == 200 + assert garbage.json() == {"active": False} + + minted = mint_session_token( + SessionPrincipal(user_id="u1", client_id="llm_dcrc_client"), + session_keys_from_master_key(introspect_master_key), + datetime.now(timezone.utc), + ) + active = client.post("/introspect", data={"token": minted.token.get_secret_value()}) + assert active.status_code == 200 + assert active.json()["active"] is True + assert active.json()["sub"] == "u1" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py index 761f823076b..32a3f70c357 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py @@ -26,6 +26,7 @@ from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import ( aggregate_authorize, aggregate_token, complete_connect_flow, + introspect_gateway_token, is_gateway_dcr_client_id, is_proxy_api_resource, native_client_auth_contract, @@ -41,7 +42,13 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credent resolve_session_bearer, session_keys_from_master_key, ) -from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token import SESSION_REFRESH_PREFIX +from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token import ( + SESSION_ISSUER, + SESSION_REFRESH_PREFIX, + SessionPrincipal, + mint_session_refresh_token, + mint_session_token, +) MASTER_KEY = "sk-gateway-dcr-flow-tests" REDIRECT_URI = "https://claude.ai/api/mcp/auth_callback" @@ -1598,17 +1605,24 @@ async def test_refresh_answers_503_without_burning_the_token_while_redis_is_down ) redis_down = await _refresh_native( - payload["refresh_token"], client_id, _Minter(), _redis_that(AsyncMock(side_effect=ConnectionError("redis down"))) + payload["refresh_token"], + client_id, + _Minter(), + _redis_that(AsyncMock(side_effect=ConnectionError("redis down"))), ) assert redis_down.status_code == 503 assert json.loads(redis_down.body)["error"] == "temporarily_unavailable" assert "refresh_token" not in json.loads(redis_down.body) - redis_back = await _refresh_native(payload["refresh_token"], client_id, _Minter(), _redis_that(AsyncMock(return_value=1))) + redis_back = await _refresh_native( + payload["refresh_token"], client_id, _Minter(), _redis_that(AsyncMock(return_value=1)) + ) assert redis_back.status_code == 200 assert json.loads(redis_back.body)["refresh_token"] != payload["refresh_token"] - replayed = await _refresh_native(payload["refresh_token"], client_id, _Minter(), _redis_that(AsyncMock(return_value=2))) + replayed = await _refresh_native( + payload["refresh_token"], client_id, _Minter(), _redis_that(AsyncMock(return_value=2)) + ) assert replayed.status_code == 400 assert json.loads(replayed.body)["error"] == "invalid_grant" @@ -1674,3 +1688,122 @@ def test_native_client_auth_contract_points_every_endpoint_at_this_proxy(): ) def test_is_proxy_api_resource_matches_only_this_proxy(resource, expected): assert is_proxy_api_resource(_request(), resource) is expected + + +def _introspection_fixtures(): + keys = session_keys_from_master_key(MASTER_KEY) + now = datetime.now(timezone.utc) + principal = SessionPrincipal(user_id="u1", client_id="llm_dcrc_client", team_id="t1") + return keys, now, principal + + +async def _introspect(token, cache=None, reload_user=_reload_user_active, master_key=MASTER_KEY): + response = await introspect_gateway_token( + token=token, master_key=master_key, reload_user=reload_user, cache=cache or DualCache() + ) + return response.status_code, json.loads(response.body) + + +@pytest.mark.asyncio +async def test_introspect_active_access_token_reports_rfc7662_claims(): + keys, now, principal = _introspection_fixtures() + minted = mint_session_token(principal, keys, now) + status, body = await _introspect(minted.token.get_secret_value()) + assert status == 200 + assert body["active"] is True + assert body["token_type"] == "Bearer" + assert body["iss"] == SESSION_ISSUER + assert body["sub"] == "u1" + assert body["client_id"] == "llm_dcrc_client" + assert body["kind"] == "session" + assert body["team_id"] == "t1" + assert body["exp"] - body["iat"] == 3600 + assert body["jti"] + + +@pytest.mark.asyncio +async def test_introspect_invalid_tokens_answer_active_false(): + keys, now, principal = _introspection_fixtures() + wrong_key = mint_session_token(principal, session_keys_from_master_key("sk-a-rotated-master-key"), now) + expired = mint_session_token(principal, keys, now - timedelta(seconds=7200)) + for candidate in ( + "sk-not-a-session-token", + "llm_session_malformed", + wrong_key.token.get_secret_value(), + expired.token.get_secret_value(), + ): + status, body = await _introspect(candidate) + assert (status, body) == (200, {"active": False}) + + +@pytest.mark.asyncio +async def test_introspect_refresh_token_goes_inactive_once_rotated(): + keys, now, _ = _introspection_fixtures() + client_id = (await _register([REDIRECT_URI]))["client_id"] + minted = mint_session_refresh_token(SessionPrincipal(user_id="u1", client_id=client_id), keys, now) + cache = DualCache() + status, body = await _introspect(minted.token.get_secret_value(), cache=cache) + assert (status, body["active"], body["kind"]) == (200, True, "session_refresh") + assert "token_type" not in body + + revoked = await revoke_refresh_token( + token=minted.token.get_secret_value(), client_id=client_id, master_key=MASTER_KEY, cache=cache + ) + assert revoked.status_code == 200 + status, body = await _introspect(minted.token.get_secret_value(), cache=cache) + assert (status, body) == (200, {"active": False}) + + +@pytest.mark.asyncio +async def test_introspect_accepts_rs256_signed_tokens_under_configured_signing(monkeypatch): + from cryptography.hazmat.primitives import serialization + from cryptography.hazmat.primitives.asymmetric import rsa + from pydantic import SecretStr + + from litellm.proxy import proxy_server + from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token import AsymmetricSessionKeys + + private_pem = ( + rsa.generate_private_key(public_exponent=65537, key_size=2048) + .private_bytes( + serialization.Encoding.PEM, + serialization.PrivateFormat.PKCS8, + serialization.NoEncryption(), + ) + .decode() + ) + monkeypatch.setitem( + proxy_server.general_settings, + "mcp_session_token_signing", + {"algorithm": "RS256", "kid": "k1", "private_key": private_pem}, + ) + _, now, principal = _introspection_fixtures() + rs_keys = AsymmetricSessionKeys(private_key_pem=SecretStr(private_pem), kid="k1") + minted = mint_session_token(principal, rs_keys, now) + status, body = await _introspect(minted.token.get_secret_value()) + assert (status, body["active"], body["kind"]) == (200, True, "session") + + hs_signed = mint_session_token(principal, session_keys_from_master_key(MASTER_KEY), now) + status, body = await _introspect(hs_signed.token.get_secret_value()) + assert (status, body) == (200, {"active": False}) + + +@pytest.mark.asyncio +async def test_introspect_fails_closed_on_dead_user_and_503s_on_outage(): + keys, now, principal = _introspection_fixtures() + minted = mint_session_token(principal, keys, now) + + async def _reload_user_gone(user_id: str): + return "unresolvable" + + status, body = await _introspect(minted.token.get_secret_value(), reload_user=_reload_user_gone) + assert (status, body) == (200, {"active": False}) + + async def _reload_user_outage(user_id: str): + return "unavailable" + + status, body = await _introspect(minted.token.get_secret_value(), reload_user=_reload_user_outage) + assert (status, body["error"]) == (503, "temporarily_unavailable") + + status, body = await _introspect(minted.token.get_secret_value(), master_key=None) + assert (status, body["error"]) == (500, "server_error") diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index eaa05ddc005..525642702ca 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -6974,6 +6974,30 @@ export interface paths { patch?: never; trace?: never; }; + "/introspect": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + get?: never; + put?: never; + /** + * Introspect Endpoint + * @description RFC 7662 introspection for gateway-issued session tokens (``llm_session_`` / + * ``llm_srefresh_``), so an external gateway can validate them without the signing + * secret. The caller authenticates with a LiteLLM virtual key (section 2.1, enforced by + * the route dependency); any token the gateway cannot vouch for answers + * ``{"active": false}`` with no further detail. + */ + post: operations["introspect_endpoint_introspect_post"]; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/invitation/delete": { parameters: { query?: never; @@ -23474,6 +23498,11 @@ export interface components { /** Mask[] */ "mask[]"?: string[] | null; }; + /** Body_introspect_endpoint_introspect_post */ + Body_introspect_endpoint_introspect_post: { + /** Token */ + token: string; + }; /** Body_revoke_endpoint_revoke_post */ Body_revoke_endpoint_revoke_post: { /** Client Id */ @@ -47651,6 +47680,39 @@ export interface operations { }; }; }; + introspect_endpoint_introspect_post: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody: { + content: { + "application/x-www-form-urlencoded": components["schemas"]["Body_introspect_endpoint_introspect_post"]; + }; + }; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; invitation_delete_invitation_delete_post: { parameters: { query?: never; From 38145c20826697c13e3b1955e41d53061376eb94 Mon Sep 17 00:00:00 2001 From: Mateo Wang Date: Sat, 29 Aug 2026 13:17:43 -0700 Subject: [PATCH 57/68] test: undo the drive-by reformat below the poller error-file regression test --- .../proxy_unit_tests/test_check_batch_cost.py | 196 ++++++++++++------ 1 file changed, 137 insertions(+), 59 deletions(-) diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index 4f9fcf1952b..757762eac87 100644 --- a/tests/proxy_unit_tests/test_check_batch_cost.py +++ b/tests/proxy_unit_tests/test_check_batch_cost.py @@ -1497,9 +1497,13 @@ class TestCheckBatchCost: from litellm.exceptions import NotFoundError - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=1) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( + return_value=1 + ) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=None + ) mock_job = MagicMock() mock_job.id = "job-output-gone-1" @@ -1509,17 +1513,23 @@ class TestCheckBatchCost: mock_job.created_by = "user-1" assert check_batch_cost_instance._has_batch_processed_column is True - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[mock_job] + ) missing_output_file_id = "gs://batch-out/job-1/predictions.jsonl" mock_response = MagicMock() mock_response.status = "failed" mock_response.output_file_id = missing_output_file_id mock_response.error_file_id = None - mock_response.model_dump_json.return_value = '{"id":"batch-1","status":"failed"}' + mock_response.model_dump_json.return_value = ( + '{"id":"batch-1","status":"failed"}' + ) mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) - mock_llm_router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"}) + mock_llm_router.get_deployment_credentials_with_provider = MagicMock( + return_value={"api_key": "sk-test"} + ) with ( patch( @@ -1540,10 +1550,12 @@ class TestCheckBatchCost: assert mock_afile_content.await_count == 1 mock_calculate.assert_not_awaited() - assert mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1, ( - "a terminal batch with a 404ing output file must be retired, not retried forever" - ) - update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[1]["data"] + assert ( + mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1 + ), "a terminal batch with a 404ing output file must be retired, not retried forever" + update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[ + 1 + ]["data"] assert update_data["status"] == "failed" assert update_data["batch_processed"] is True @@ -1556,9 +1568,13 @@ class TestCheckBatchCost: Without this, GET /batches/{id} returns a raw file ID that cannot be routed through the proxy, causing API_KEY errors when clients call GET /files/{id}/content. """ - mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=1) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock( + return_value=1 + ) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() - mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock( + return_value=None + ) mock_job = MagicMock() mock_job.id = "job-raw-file-1" @@ -1567,7 +1583,9 @@ class TestCheckBatchCost: mock_job.team_id = None check_batch_cost_instance._has_batch_processed_column = True - mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[mock_job] + ) raw_output_file_id = "file-batch-output-abc123" raw_error_file_id = "file-batch-error-xyz456" @@ -1578,10 +1596,14 @@ class TestCheckBatchCost: mock_response.status = "completed" mock_response.output_file_id = raw_output_file_id mock_response.error_file_id = raw_error_file_id - mock_response.model_dump_json.return_value = '{"id":"batch-1","status":"completed"}' + mock_response.model_dump_json.return_value = ( + '{"id":"batch-1","status":"completed"}' + ) mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) - mock_llm_router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"}) + mock_llm_router.get_deployment_credentials_with_provider = MagicMock( + return_value={"api_key": "sk-test"} + ) mock_deployment = MagicMock() mock_deployment.litellm_params.custom_llm_provider = "azure" @@ -1596,7 +1618,9 @@ class TestCheckBatchCost: fake_managed_error_id, ] mock_hook.store_unified_file_id = AsyncMock() - check_batch_cost_instance.proxy_logging_obj.get_proxy_hook.return_value = mock_hook + check_batch_cost_instance.proxy_logging_obj.get_proxy_hook.return_value = ( + mock_hook + ) mock_file_content = MagicMock() mock_file_content.content = b'{"id":"req-1"}' @@ -1639,7 +1663,9 @@ class TestCheckBatchCost: "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", return_value=("gpt-5-mini", "azure", None, None), ), - patch("litellm.litellm_core_utils.litellm_logging.Logging") as mock_logging_cls, + patch( + "litellm.litellm_core_utils.litellm_logging.Logging" + ) as mock_logging_cls, ): mock_logging_obj = MagicMock() mock_logging_obj.async_success_handler = AsyncMock() @@ -1690,7 +1716,9 @@ class TestUnmanagedVertexRouting: def _job(self, file_object=None): job = MagicMock() job.unified_object_id = "8823717160934178816" - job.file_object = file_object if file_object is not None else _unmanaged_vertex_file_object() + job.file_object = ( + file_object if file_object is not None else _unmanaged_vertex_file_object() + ) return job def test_flag_off_skips_unmanaged_id_unchanged(self): @@ -1728,7 +1756,9 @@ class TestUnmanagedVertexRouting: assert result == ("deploy-1", "8823717160934178816") # bare model name (trailing GCS segment), not the full publishers/.. path - router.resolve_model_name_from_model_id.assert_called_once_with("gemini-2.5-flash") + router.resolve_model_name_from_model_id.assert_called_once_with( + "gemini-2.5-flash" + ) router.get_model_ids.assert_called_once_with(model_name="gemini-2.5-flash") def test_flag_on_skips_non_vertex_deployment_sharing_model_group(self): @@ -1748,7 +1778,9 @@ class TestUnmanagedVertexRouting: result = instance._resolve_job_routing(self._job(), prom) assert result is None - prom.record_check_batch_cost_error.assert_called_once_with("unmanaged_no_matching_deployment") + prom.record_check_batch_cost_error.assert_called_once_with( + "unmanaged_no_matching_deployment" + ) def test_flag_on_uses_later_vertex_deployment_with_matching_suffix(self): router = MagicMock() @@ -1796,7 +1828,9 @@ class TestUnmanagedVertexRouting: result = instance._resolve_job_routing(self._job(), prom) assert result is None - prom.record_check_batch_cost_error.assert_called_once_with("unmanaged_no_matching_deployment") + prom.record_check_batch_cost_error.assert_called_once_with( + "unmanaged_no_matching_deployment" + ) def test_flag_on_non_gcs_input_is_not_unmanaged_vertex(self): """Flag on, but input_file_id is not a gs:// publishers path: treat as unroutable, @@ -1804,7 +1838,9 @@ class TestUnmanagedVertexRouting: router = MagicMock() instance = self._instance(track_unmanaged=True, router=router) prom = MagicMock() - job = self._job(file_object=_unmanaged_vertex_file_object(input_file_id="file-abc-123")) + job = self._job( + file_object=_unmanaged_vertex_file_object(input_file_id="file-abc-123") + ) with patch(_IS_B64, return_value=False): result = instance._resolve_job_routing(job, prom) @@ -1828,7 +1864,9 @@ class TestUnmanagedVertexRouting: mock_response.error_file_id = None mock_response.completed_at = None mock_response.created_at = None - mock_response.model_dump_json.return_value = '{"id":"8823717160934178816","status":"completed"}' + mock_response.model_dump_json.return_value = ( + '{"id":"8823717160934178816","status":"completed"}' + ) router.aretrieve_batch = AsyncMock(return_value=mock_response) router.get_deployment_credentials_with_provider = MagicMock( return_value={"vertex_project": "p", "vertex_location": "us-central1"} @@ -1850,7 +1888,9 @@ class TestUnmanagedVertexRouting: prisma.db.litellm_managedobjecttable = MagicMock() prisma.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=1) prisma.db.litellm_managedobjecttable.update = AsyncMock() - prisma.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[self._job()]) + prisma.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[self._job()] + ) prisma.db.litellm_usertable = MagicMock() prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=None) @@ -1881,7 +1921,9 @@ class TestUnmanagedVertexRouting: "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", return_value=("gemini-2.5-flash", "vertex_ai", None, None), ), - patch("litellm.litellm_core_utils.litellm_logging.Logging") as mock_logging_cls, + patch( + "litellm.litellm_core_utils.litellm_logging.Logging" + ) as mock_logging_cls, ): mock_logging_obj = MagicMock() mock_logging_obj.async_success_handler = AsyncMock() @@ -1922,7 +1964,9 @@ class TestUnmanagedBedrockRouting: def _job(self, file_object=None): job = MagicMock() job.unified_object_id = self._ARN - job.file_object = file_object if file_object is not None else _unmanaged_bedrock_file_object() + job.file_object = ( + file_object if file_object is not None else _unmanaged_bedrock_file_object() + ) return job def _bedrock_deployment(self): @@ -1977,7 +2021,9 @@ class TestUnmanagedBedrockRouting: result = instance._resolve_job_routing(self._job(), prom) assert result is None - prom.record_check_batch_cost_error.assert_called_once_with("unmanaged_no_matching_deployment") + prom.record_check_batch_cost_error.assert_called_once_with( + "unmanaged_no_matching_deployment" + ) def test_flag_on_matches_deployment_despite_colon_dash_mismatch(self): """The S3 object key has ':' replaced with '-' (e.g. 'v1-0'), but the configured @@ -2015,7 +2061,9 @@ class TestUnmanagedBedrockRouting: result = instance._resolve_job_routing(self._job(), prom) assert result is None - prom.record_check_batch_cost_error.assert_called_once_with("unmanaged_no_matching_deployment") + prom.record_check_batch_cost_error.assert_called_once_with( + "unmanaged_no_matching_deployment" + ) def test_flag_on_non_s3_input_is_not_unmanaged_bedrock(self): """Flag on, but input_file_id is not a litellm-bedrock-files- s3:// key: treat as @@ -2023,7 +2071,9 @@ class TestUnmanagedBedrockRouting: router = MagicMock() instance = self._instance(track_unmanaged=True, router=router) prom = MagicMock() - job = self._job(file_object=_unmanaged_bedrock_file_object(input_file_id="file-abc-123")) + job = self._job( + file_object=_unmanaged_bedrock_file_object(input_file_id="file-abc-123") + ) with patch(_IS_B64, return_value=False): result = instance._resolve_job_routing(job, prom) @@ -2046,9 +2096,13 @@ class TestUnmanagedBedrockRouting: mock_response.error_file_id = None mock_response.completed_at = None mock_response.created_at = None - mock_response.model_dump_json.return_value = f'{{"id":"{self._ARN}","status":"completed"}}' + mock_response.model_dump_json.return_value = ( + f'{{"id":"{self._ARN}","status":"completed"}}' + ) router.aretrieve_batch = AsyncMock(return_value=mock_response) - router.get_deployment_credentials_with_provider = MagicMock(return_value={"aws_region_name": "us-east-1"}) + router.get_deployment_credentials_with_provider = MagicMock( + return_value={"aws_region_name": "us-east-1"} + ) deployment = self._bedrock_deployment() deployment.model_name = "claude-sonnet-4" @@ -2064,7 +2118,9 @@ class TestUnmanagedBedrockRouting: prisma.db.litellm_managedobjecttable = MagicMock() prisma.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=1) prisma.db.litellm_managedobjecttable.update = AsyncMock() - prisma.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[self._job()]) + prisma.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[self._job()] + ) prisma.db.litellm_usertable = MagicMock() prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=None) @@ -2095,7 +2151,9 @@ class TestUnmanagedBedrockRouting: "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", return_value=("claude-sonnet-4", "bedrock", None, None), ), - patch("litellm.litellm_core_utils.litellm_logging.Logging") as mock_logging_cls, + patch( + "litellm.litellm_core_utils.litellm_logging.Logging" + ) as mock_logging_cls, ): mock_logging_obj = MagicMock() mock_logging_obj.async_success_handler = AsyncMock() @@ -2219,7 +2277,9 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: ) router = MagicMock() - router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"}) + router.get_deployment_credentials_with_provider = MagicMock( + return_value={"api_key": "sk-test"} + ) deployment = MagicMock() deployment.litellm_params.custom_llm_provider = "azure" deployment.litellm_params.model = "azure/gpt-5.5" @@ -2228,8 +2288,8 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: router.get_deployment = MagicMock(return_value=deployment) hook = MagicMock() - hook.get_unified_output_file_id = lambda output_file_id, model_id, model_name: ( - _PROXY_LiteLLMManagedFiles.get_unified_output_file_id( + hook.get_unified_output_file_id = ( + lambda output_file_id, model_id, model_name: _PROXY_LiteLLMManagedFiles.get_unified_output_file_id( None, output_file_id=output_file_id, model_id=model_id, model_name=model_name ) ) @@ -2298,7 +2358,9 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: get_models_from_unified_file_id, ) - output_file_id = await self._run(self._job(self._managed_input_file_id(self._PUBLIC_MODEL_GROUP))) + output_file_id = await self._run( + self._job(self._managed_input_file_id(self._PUBLIC_MODEL_GROUP)) + ) decoded = _is_base64_encoded_unified_file_id(output_file_id) assert get_models_from_unified_file_id(decoded) == [self._PUBLIC_MODEL_GROUP] @@ -2312,7 +2374,9 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: _extract_models_from_managed_resource_id, ) - output_file_id = await self._run(self._job(self._managed_input_file_id(self._PUBLIC_MODEL_GROUP))) + output_file_id = await self._run( + self._job(self._managed_input_file_id(self._PUBLIC_MODEL_GROUP)) + ) models = _extract_models_from_managed_resource_id(output_file_id, "file_id", None) assert models == [self._PUBLIC_MODEL_GROUP] @@ -2320,7 +2384,9 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: await can_key_call_model( model=models[0], llm_model_list=None, - valid_token=UserAPIKeyAuth(api_key="sk-test", models=[self._PUBLIC_MODEL_GROUP]), + valid_token=UserAPIKeyAuth( + api_key="sk-test", models=[self._PUBLIC_MODEL_GROUP] + ), llm_router=None, ) is True @@ -2337,8 +2403,6 @@ class TestManagedOutputFileIdEncodesPublicModelGroup: decoded = _is_base64_encoded_unified_file_id(output_file_id) assert get_models_from_unified_file_id(decoded) == [self._PUBLIC_MODEL_GROUP] - - class TestBatchCostAttribution: """CheckBatchCost rebuilds the creator's spend metadata from the managed-object row so the batch-cost log is attributed like a non-batch request.""" @@ -2434,7 +2498,9 @@ class TestBatchCostAttribution: """An alias lookup failure must not lose the spend row; the key hash and team still attribute it.""" instance = self._instance() - instance.prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(side_effect=Exception("db down")) + instance.prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( + side_effect=Exception("db down") + ) metadata = await instance._build_creator_attribution_metadata(self._job(), "batch-1") @@ -2530,7 +2596,9 @@ class TestPollPageStarvation: async def test_unified_id_without_model_id_is_retired(self): """A unified id that decodes but carries no model_id is unroutable no matter what the config says, so it must leave the poll page instead of being retried forever.""" - prisma = self._prisma([self._job("job-no-model", self._encode("litellm_proxy;llm_batch_id:poison-no-model"))]) + prisma = self._prisma( + [self._job("job-no-model", self._encode("litellm_proxy;llm_batch_id:poison-no-model"))] + ) llm_router = MagicMock() llm_router.aretrieve_batch = AsyncMock() @@ -2568,7 +2636,9 @@ class TestPollPageStarvation: await self._instance(prisma, llm_router).check_batch_cost() prisma.db.litellm_managedobjecttable.update.assert_awaited_once() - assert prisma.db.litellm_managedobjecttable.update.call_args[1]["data"] == {"batch_processed": True} + assert prisma.db.litellm_managedobjecttable.update.call_args[1]["data"] == { + "batch_processed": True + } @pytest.mark.asyncio async def test_provider_404_with_deployment_gone_keeps_job(self): @@ -2621,13 +2691,17 @@ class TestPollPageStarvation: async def test_retirement_falls_back_to_status_without_batch_processed_column(self): """Older schemas have no batch_processed column, so the only way to stop selecting the row is the status filter the poll query already applies.""" - prisma = self._prisma([self._job("job-legacy", self._encode("litellm_proxy;llm_batch_id:poison-no-model"))]) + prisma = self._prisma( + [self._job("job-legacy", self._encode("litellm_proxy;llm_batch_id:poison-no-model"))] + ) instance = self._instance(prisma, MagicMock()) instance._has_batch_processed_column = False await instance.check_batch_cost() - assert prisma.db.litellm_managedobjecttable.update.call_args[1]["data"] == {"status": "stale_expired"} + assert prisma.db.litellm_managedobjecttable.update.call_args[1]["data"] == { + "status": "stale_expired" + } @pytest.mark.asyncio async def test_stale_cleanup_gives_up_on_never_costed_completed_rows(self): @@ -2682,11 +2756,14 @@ class TestPollPageStarvation: await self._instance(prisma, llm_router).check_batch_cost() - retired = [call[1]["where"]["id"] for call in prisma.db.litellm_managedobjecttable.update.call_args_list] + retired = [ + call[1]["where"]["id"] + for call in prisma.db.litellm_managedobjecttable.update.call_args_list + ] assert retired == ["job-no-model", "job-gone"] - assert llm_router.aretrieve_batch.await_args_list[-1][1]["batch_id"] == "batch_live", ( - "the newer healthy batch must still be polled in the same cycle" - ) + assert ( + llm_router.aretrieve_batch.await_args_list[-1][1]["batch_id"] == "batch_live" + ), "the newer healthy batch must still be polled in the same cycle" @pytest.mark.asyncio async def test_404_that_does_not_name_the_batch_keeps_job_for_retry(self): @@ -2715,7 +2792,6 @@ class TestPollPageStarvation: prisma.db.litellm_managedobjecttable.update.assert_not_awaited() - class _FakeManagedObjectRow: """One managed batch row the provider has finished but nothing has costed yet.""" @@ -2732,12 +2808,8 @@ class _FakeManagedObjectRow: self.request_tags = None self.created_at = 1700000000 self.file_object = json.dumps( - { - "id": "batch-456", - "status": "in_progress", - "input_file_id": "file-input-1", - "output_file_id": _CLAIM_OUTPUT_FILE_ID, - } + {"id": "batch-456", "status": "in_progress", "input_file_id": "file-input-1", + "output_file_id": _CLAIM_OUTPUT_FILE_ID} ) @@ -2837,7 +2909,9 @@ class TestMultiPodBatchCostClaim: router = MagicMock() router.aretrieve_batch = AsyncMock(return_value=response) - router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"}) + router.get_deployment_credentials_with_provider = MagicMock( + return_value={"api_key": "sk-test"} + ) router.get_deployment = MagicMock(return_value=deployment) return router @@ -2997,7 +3071,9 @@ class TestMultiPodBatchCostClaim: await asyncio.Event().wait() with self._billing_patches(journal, during_fetch=_never_returns) as logging_obj: - interrupted = asyncio.create_task(self._instance(prisma, self._router()).check_batch_cost()) + interrupted = asyncio.create_task( + self._instance(prisma, self._router()).check_batch_cost() + ) await asyncio.wait_for(reached_fetch.wait(), timeout=5) assert row.batch_processed is False, "an in-flight costing must not mark the row processed" interrupted.cancel() @@ -3032,7 +3108,9 @@ class TestMultiPodBatchCostClaim: await finish_fetch.wait() with self._billing_patches(journal, during_fetch=_wait_for_the_delete_attempt): - costing = asyncio.create_task(self._instance(prisma, self._router()).check_batch_cost()) + costing = asyncio.create_task( + self._instance(prisma, self._router()).check_batch_cost() + ) await asyncio.wait_for(reached_fetch.wait(), timeout=5) with pytest.raises(HTTPException) as blocked: From 4d5205c355113d19721899aad55463735f3637b3 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 13:32:39 -0700 Subject: [PATCH 58/68] fix(proxy): give the remaining CLI clients a request timeout The keys, credentials, models, model groups, and chat clients still sent requests with no timeout, so a proxy that accepts the connection and never answers pinned the caller forever. They now default to the same 30 seconds as their teams and users siblings, with chat on the OpenAI SDK's 600 second default, and Client wires its timeout through to all of them. S113 cannot see Session methods, so each client gets a hanging-server regression test instead. --- litellm/proxy/client/chat.py | 11 +++++-- litellm/proxy/client/client.py | 11 +++---- litellm/proxy/client/credentials.py | 12 ++++---- litellm/proxy/client/keys.py | 14 +++++---- litellm/proxy/client/model_groups.py | 6 ++-- litellm/proxy/client/models.py | 14 +++++---- tests/test_litellm/proxy/client/test_chat.py | 29 +++++++++++++++++++ .../test_litellm/proxy/client/test_client.py | 8 +++++ .../proxy/client/test_credentials.py | 15 ++++++++++ tests/test_litellm/proxy/client/test_keys.py | 15 ++++++++++ .../proxy/client/test_model_groups.py | 15 ++++++++++ .../test_litellm/proxy/client/test_models.py | 15 ++++++++++ type-discipline-budget.json | 4 +-- 13 files changed, 140 insertions(+), 29 deletions(-) diff --git a/litellm/proxy/client/chat.py b/litellm/proxy/client/chat.py index 2953ed7f683..bd4d0df3ed0 100644 --- a/litellm/proxy/client/chat.py +++ b/litellm/proxy/client/chat.py @@ -8,16 +8,19 @@ from .exceptions import UnauthorizedError class ChatClient: - def __init__(self, base_url: str, api_key: str | None = None): + def __init__(self, base_url: str, api_key: str | None = None, timeout: int = 600): """ Initialize the ChatClient. Args: base_url (str): The base URL of the LiteLLM proxy server (e.g., "http://localhost:8000") api_key (Optional[str]): API key for authentication. If provided, it will be sent as a Bearer token. + timeout (int): Request timeout in seconds (default: 600, the OpenAI SDK default, since a completion + can legitimately take minutes) """ self._base_url = base_url.rstrip("/") # Remove trailing slash if present self._api_key = api_key + self._timeout = timeout def _get_headers(self) -> dict[str, str]: """ @@ -96,7 +99,7 @@ class ChatClient: # Prepare and send the request session: Final = requests.Session() try: - response: Final = session.send(request.prepare()) + response: Final = session.send(request.prepare(), timeout=self._timeout) response.raise_for_status() return response.json() except requests.exceptions.HTTPError as e: @@ -161,7 +164,9 @@ class ChatClient: # Make streaming request session: Final = requests.Session() try: - response: Final = session.post(url, headers=self._get_headers(), json=data, stream=True) + response: Final = session.post( + url, headers=self._get_headers(), json=data, stream=True, timeout=self._timeout + ) response.raise_for_status() # Parse SSE stream diff --git a/litellm/proxy/client/client.py b/litellm/proxy/client/client.py index 560523db189..de1e45b91be 100644 --- a/litellm/proxy/client/client.py +++ b/litellm/proxy/client/client.py @@ -24,7 +24,8 @@ class Client: Args: base_url (str): The base URL of the LiteLLM proxy server (e.g., "http://localhost:4000") api_key (Optional[str]): API key for authentication. If provided, it will be sent as a Bearer token. - timeout: Request timeout in seconds (default: 30) + timeout: Request timeout in seconds for management calls (default: 30). Chat completions keep + ChatClient's own 600 second default, since a completion can legitimately take minutes """ self._base_url = base_url.rstrip("/") # Only use the stored CLI key when it was issued for this server. @@ -33,9 +34,9 @@ class Client: # Initialize resource clients self.http = HTTPClient(base_url=base_url, api_key=self._api_key, timeout=timeout) - self.models = ModelsManagementClient(base_url=self._base_url, api_key=self._api_key) - self.model_groups = ModelGroupsManagementClient(base_url=self._base_url, api_key=self._api_key) + self.models = ModelsManagementClient(base_url=self._base_url, api_key=self._api_key, timeout=timeout) + self.model_groups = ModelGroupsManagementClient(base_url=self._base_url, api_key=self._api_key, timeout=timeout) self.chat = ChatClient(base_url=self._base_url, api_key=self._api_key) - self.keys = KeysManagementClient(base_url=self._base_url, api_key=self._api_key) - self.credentials = CredentialsManagementClient(base_url=self._base_url, api_key=self._api_key) + self.keys = KeysManagementClient(base_url=self._base_url, api_key=self._api_key, timeout=timeout) + self.credentials = CredentialsManagementClient(base_url=self._base_url, api_key=self._api_key, timeout=timeout) self.teams = TeamsManagementClient(base_url=self._base_url, api_key=self._api_key, timeout=timeout) diff --git a/litellm/proxy/client/credentials.py b/litellm/proxy/client/credentials.py index 136bdf3f293..a9bff67b1c5 100644 --- a/litellm/proxy/client/credentials.py +++ b/litellm/proxy/client/credentials.py @@ -6,16 +6,18 @@ from .exceptions import UnauthorizedError class CredentialsManagementClient: - def __init__(self, base_url: str, api_key: str | None = None): + def __init__(self, base_url: str, api_key: str | None = None, timeout: int = 30): """ Initialize the CredentialsManagementClient. Args: base_url (str): The base URL of the LiteLLM proxy server (e.g., "http://localhost:8000") api_key (Optional[str]): API key for authentication. If provided, it will be sent as a Bearer token. + timeout (int): Request timeout in seconds (default: 30) """ self._base_url = base_url.rstrip("/") # Remove trailing slash if present self._api_key = api_key + self._timeout = timeout def _get_headers(self) -> dict[str, str]: """ @@ -56,7 +58,7 @@ class CredentialsManagementClient: session: Final = requests.Session() try: - response: Final = session.send(request.prepare()) + response: Final = session.send(request.prepare(), timeout=self._timeout) response.raise_for_status() return response.json() except requests.exceptions.HTTPError as e: @@ -103,7 +105,7 @@ class CredentialsManagementClient: session: Final = requests.Session() try: - response: Final = session.send(request.prepare()) + response: Final = session.send(request.prepare(), timeout=self._timeout) response.raise_for_status() return response.json() except requests.exceptions.HTTPError as e: @@ -140,7 +142,7 @@ class CredentialsManagementClient: session: Final = requests.Session() try: - response: Final = session.send(request.prepare()) + response: Final = session.send(request.prepare(), timeout=self._timeout) response.raise_for_status() return response.json() except requests.exceptions.HTTPError as e: @@ -177,7 +179,7 @@ class CredentialsManagementClient: session: Final = requests.Session() try: - response: Final = session.send(request.prepare()) + response: Final = session.send(request.prepare(), timeout=self._timeout) response.raise_for_status() return response.json() except requests.exceptions.HTTPError as e: diff --git a/litellm/proxy/client/keys.py b/litellm/proxy/client/keys.py index 5b66567363d..fe100c5f676 100644 --- a/litellm/proxy/client/keys.py +++ b/litellm/proxy/client/keys.py @@ -9,16 +9,18 @@ from .exceptions import UnauthorizedError class KeysManagementClient: - def __init__(self, base_url: str, api_key: str | None = None): + def __init__(self, base_url: str, api_key: str | None = None, timeout: int = 30): """ Initialize the KeysManagementClient. Args: base_url (str): The base URL of the LiteLLM proxy server (e.g., "http://localhost:8000") api_key (Optional[str]): API key for authentication. If provided, it will be sent as a Bearer token. + timeout (int): Request timeout in seconds (default: 30) """ self._base_url = base_url.rstrip("/") # Remove trailing slash if present self._api_key = api_key + self._timeout = timeout def _get_headers(self) -> dict[str, str]: """ @@ -99,7 +101,7 @@ class KeysManagementClient: session: Final = requests.Session() try: - response: Final = session.send(request.prepare()) + response: Final = session.send(request.prepare(), timeout=self._timeout) response.raise_for_status() return response.json() except requests.exceptions.HTTPError as e: @@ -174,7 +176,7 @@ class KeysManagementClient: session: Final = requests.Session() try: - response: Final = session.send(request.prepare()) + response: Final = session.send(request.prepare(), timeout=self._timeout) response.raise_for_status() return response.json() except requests.exceptions.HTTPError as e: @@ -218,7 +220,7 @@ class KeysManagementClient: session: Final = requests.Session() try: - response: Final = session.send(request.prepare()) + response: Final = session.send(request.prepare(), timeout=self._timeout) response.raise_for_status() return response.json() except requests.exceptions.HTTPError as e: @@ -279,7 +281,7 @@ class KeysManagementClient: session: Final = requests.Session() response_text: str | None = None try: - response: Final = session.send(request.prepare()) + response: Final = session.send(request.prepare(), timeout=self._timeout) response_text = response.text response.raise_for_status() return response.json() @@ -309,7 +311,7 @@ class KeysManagementClient: session: Final = requests.Session() try: - response: Final = session.send(request.prepare()) + response: Final = session.send(request.prepare(), timeout=self._timeout) response.raise_for_status() return response.json() except requests.exceptions.HTTPError as e: diff --git a/litellm/proxy/client/model_groups.py b/litellm/proxy/client/model_groups.py index 9c7c38dc67c..fef307600c4 100644 --- a/litellm/proxy/client/model_groups.py +++ b/litellm/proxy/client/model_groups.py @@ -6,16 +6,18 @@ from .exceptions import UnauthorizedError class ModelGroupsManagementClient: - def __init__(self, base_url: str, api_key: str | None = None): + def __init__(self, base_url: str, api_key: str | None = None, timeout: int = 30): """ Initialize the ModelGroupsManagementClient. Args: base_url (str): The base URL of the LiteLLM proxy server (e.g., "http://localhost:8000") api_key (Optional[str]): API key for authentication. If provided, it will be sent as a Bearer token. + timeout (int): Request timeout in seconds (default: 30) """ self._base_url = base_url.rstrip("/") # Remove trailing slash if present self._api_key = api_key + self._timeout = timeout def _get_headers(self) -> dict[str, str]: """ @@ -53,7 +55,7 @@ class ModelGroupsManagementClient: # Prepare and send the request session: Final = requests.Session() try: - response: Final = session.send(request.prepare()) + response: Final = session.send(request.prepare(), timeout=self._timeout) response.raise_for_status() return response.json()["data"] except requests.exceptions.HTTPError as e: diff --git a/litellm/proxy/client/models.py b/litellm/proxy/client/models.py index 0f1dd2b5bab..4b16087e15b 100644 --- a/litellm/proxy/client/models.py +++ b/litellm/proxy/client/models.py @@ -7,16 +7,18 @@ from .exceptions import NotFoundError, UnauthorizedError class ModelsManagementClient: - def __init__(self, base_url: str, api_key: str | None = None): + def __init__(self, base_url: str, api_key: str | None = None, timeout: int = 30): """ Initialize the ModelsManagementClient. Args: base_url (str): The base URL of the LiteLLM proxy server (e.g., "http://localhost:8000") api_key (Optional[str]): API key for authentication. If provided, it will be sent as a Bearer token. + timeout (int): Request timeout in seconds (default: 30) """ self._base_url = base_url.rstrip("/") # Remove trailing slash if present self._api_key = api_key + self._timeout = timeout def _get_headers(self) -> dict[str, str]: """ @@ -55,7 +57,7 @@ class ModelsManagementClient: # Prepare and send the request session: Final = requests.Session() try: - response: Final = session.send(request.prepare()) + response: Final = session.send(request.prepare(), timeout=self._timeout) response.raise_for_status() return response.json()["data"] except requests.exceptions.HTTPError as e: @@ -104,7 +106,7 @@ class ModelsManagementClient: # Prepare and send the request session: Final = requests.Session() try: - response: Final = session.send(request.prepare()) + response: Final = session.send(request.prepare(), timeout=self._timeout) response.raise_for_status() return response.json() except requests.exceptions.HTTPError as e: @@ -140,7 +142,7 @@ class ModelsManagementClient: # Prepare and send the request session: Final = requests.Session() try: - response: Final = session.send(request.prepare()) + response: Final = session.send(request.prepare(), timeout=self._timeout) response.raise_for_status() return response.json() except requests.exceptions.HTTPError as e: @@ -232,7 +234,7 @@ class ModelsManagementClient: # Prepare and send the request session: Final = requests.Session() try: - response: Final = session.send(request.prepare()) + response: Final = session.send(request.prepare(), timeout=self._timeout) response.raise_for_status() return response.json()["data"] except requests.exceptions.HTTPError as e: @@ -282,7 +284,7 @@ class ModelsManagementClient: # Prepare and send the request session: Final = requests.Session() try: - response: Final = session.send(request.prepare()) + response: Final = session.send(request.prepare(), timeout=self._timeout) response.raise_for_status() return response.json() except requests.exceptions.HTTPError as e: diff --git a/tests/test_litellm/proxy/client/test_chat.py b/tests/test_litellm/proxy/client/test_chat.py index b8e55c45502..67b6ee833f2 100644 --- a/tests/test_litellm/proxy/client/test_chat.py +++ b/tests/test_litellm/proxy/client/test_chat.py @@ -1,6 +1,7 @@ import importlib import importlib.util from importlib.machinery import PathFinder +import time import site import sys @@ -227,3 +228,31 @@ def test_completions_other_errors(client, sample_messages): with pytest.raises(requests.exceptions.HTTPError) as exc_info: client.completions(model="gpt-4", messages=sample_messages) assert exc_info.value.response.status_code == 500 + + +def test_completions_gives_up_at_the_timeout_instead_of_hanging(hanging_server): + """ + A proxy that accepts the connection but never answers used to pin the caller's + process forever, since the request carried no timeout at all. + """ + client = ChatClient(base_url=hanging_server, api_key="sk-test", timeout=1) + + started = time.monotonic() + with pytest.raises(requests.exceptions.Timeout): + client.completions(model="gpt-5.4", messages=[{"role": "user", "content": "hi"}]) + + assert time.monotonic() - started < 10 + + +def test_completions_stream_gives_up_at_the_timeout_instead_of_hanging(hanging_server): + """ + The streaming call opens the response before reading chunks, so a proxy that never + sends its headers used to hang here forever too. + """ + client = ChatClient(base_url=hanging_server, api_key="sk-test", timeout=1) + + started = time.monotonic() + with pytest.raises(requests.exceptions.Timeout): + next(client.completions_stream(model="gpt-5.4", messages=[{"role": "user", "content": "hi"}])) + + assert time.monotonic() - started < 10 diff --git a/tests/test_litellm/proxy/client/test_client.py b/tests/test_litellm/proxy/client/test_client.py index fe3e2c52ce5..87eb3400b8c 100644 --- a/tests/test_litellm/proxy/client/test_client.py +++ b/tests/test_litellm/proxy/client/test_client.py @@ -82,6 +82,12 @@ def test_client_initialization(): assert client.http._base_url == "http://localhost:4000" assert client.http._api_key == "test-key" assert client.http._timeout == 60 + assert client.teams._timeout == 60 + assert client.keys._timeout == 60 + assert client.credentials._timeout == 60 + assert client.models._timeout == 60 + assert client.model_groups._timeout == 60 + assert client.chat._timeout == 600 def test_client_default_timeout(): @@ -92,6 +98,8 @@ def test_client_default_timeout(): ) assert client.http._timeout == 30 + assert client.keys._timeout == 30 + assert client.chat._timeout == 600 def test_client_without_api_key(): diff --git a/tests/test_litellm/proxy/client/test_credentials.py b/tests/test_litellm/proxy/client/test_credentials.py index 41886e3b292..666c5dac2b0 100644 --- a/tests/test_litellm/proxy/client/test_credentials.py +++ b/tests/test_litellm/proxy/client/test_credentials.py @@ -1,4 +1,5 @@ +import time import pytest import requests @@ -276,3 +277,17 @@ def test_encrypt_credential_values_does_not_mutate_original(monkeypatch): assert encrypted.credential_values["api_key"] != "sk-123" assert credential.credential_values["api_key"] == "sk-123" assert encrypted.credential_name == credential.credential_name + + +def test_list_gives_up_at_the_timeout_instead_of_hanging(hanging_server): + """ + A proxy that accepts the connection but never answers used to pin the caller's + process forever, since the request carried no timeout at all. + """ + client = CredentialsManagementClient(base_url=hanging_server, api_key="sk-test", timeout=1) + + started = time.monotonic() + with pytest.raises(requests.exceptions.Timeout): + client.list() + + assert time.monotonic() - started < 10 diff --git a/tests/test_litellm/proxy/client/test_keys.py b/tests/test_litellm/proxy/client/test_keys.py index 282b97b1c09..b9b07bddf1f 100644 --- a/tests/test_litellm/proxy/client/test_keys.py +++ b/tests/test_litellm/proxy/client/test_keys.py @@ -1,3 +1,4 @@ +import time import traceback import pytest @@ -509,3 +510,17 @@ def test_not_found_error_redacts_wrapped_key(): assert "REDACTED" in str(wrapped) assert LEAKY_KEY not in str(wrapped.orig_exception) assert wrapped.orig_exception.response.status_code == 404 + + +def test_list_gives_up_at_the_timeout_instead_of_hanging(hanging_server): + """ + A proxy that accepts the connection but never answers used to pin the caller's + process forever, since the request carried no timeout at all. + """ + client = KeysManagementClient(base_url=hanging_server, api_key="sk-test", timeout=1) + + started = time.monotonic() + with pytest.raises(requests.exceptions.Timeout): + client.list() + + assert time.monotonic() - started < 10 diff --git a/tests/test_litellm/proxy/client/test_model_groups.py b/tests/test_litellm/proxy/client/test_model_groups.py index 9ea8e94ff95..4a513a127b8 100644 --- a/tests/test_litellm/proxy/client/test_model_groups.py +++ b/tests/test_litellm/proxy/client/test_model_groups.py @@ -1,4 +1,5 @@ +import time import pytest import requests @@ -172,3 +173,17 @@ def test_client_initialization_without_api_key(base_url): assert client._api_key is None assert client.model_groups._api_key is None + + +def test_info_gives_up_at_the_timeout_instead_of_hanging(hanging_server): + """ + A proxy that accepts the connection but never answers used to pin the caller's + process forever, since the request carried no timeout at all. + """ + client = ModelGroupsManagementClient(base_url=hanging_server, api_key="sk-test", timeout=1) + + started = time.monotonic() + with pytest.raises(requests.exceptions.Timeout): + client.info() + + assert time.monotonic() - started < 10 diff --git a/tests/test_litellm/proxy/client/test_models.py b/tests/test_litellm/proxy/client/test_models.py index fe053ffd683..9aa5a6cf0b3 100644 --- a/tests/test_litellm/proxy/client/test_models.py +++ b/tests/test_litellm/proxy/client/test_models.py @@ -1,4 +1,5 @@ +import time import pytest import requests @@ -732,3 +733,17 @@ def test_update_other_errors(client): with pytest.raises(requests.exceptions.HTTPError) as exc_info: client.update(model_id=model_id, model_params=model_params) assert exc_info.value.response.status_code == 500 + + +def test_list_gives_up_at_the_timeout_instead_of_hanging(hanging_server): + """ + A proxy that accepts the connection but never answers used to pin the caller's + process forever, since the request carried no timeout at all. + """ + client = ModelsManagementClient(base_url=hanging_server, api_key="sk-test", timeout=1) + + started = time.monotonic() + with pytest.raises(requests.exceptions.Timeout): + client.list() + + assert time.monotonic() - started < 10 diff --git a/type-discipline-budget.json b/type-discipline-budget.json index f3f1a7defe7..7365cec4fdd 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,6 +1,6 @@ { "LIT001": { - "limit": 22705 + "limit": 22704 }, "LIT002": { "limit": 26854 @@ -33,6 +33,6 @@ "limit": 5577 }, "LIT012": { - "limit": 4508 + "limit": 4506 } } From 645792955dd2b03c7df3fa105f272fafe01dee1e Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 29 Aug 2026 13:36:08 -0700 Subject: [PATCH 59/68] feat(proxy): cyberark conjur secret manager configuration via Admin UI (#38445) * feat(proxy): CyberArk Conjur secret manager configuration via Admin UI Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(ui): mock networking base-url helpers in AdminPanel test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): restore deployment CyberArk env config on delete and roll back on persist failure Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): reinit env-configured hashicorp vault manager after cyberark persist rollback Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- basedpyright-code-budget.json | 2 +- litellm/proxy/_lazy_openapi_snapshot.json | 289 +++++++++++++ .../config_override_endpoints.py | 408 +++++++++++++++++- litellm/proxy/proxy_server.py | 61 +++ .../management_endpoints/config_overrides.py | 37 ++ ruff-strict-budget.json | 4 +- .../test_config_override_endpoints.py | 391 +++++++++++++++++ .../_components/AdminPanel.test.tsx | 2 + .../admin-panel/_components/AdminPanel.tsx | 6 + .../hooks/configOverrides/cyberArkApi.ts | 38 ++ .../configOverrides/useCyberArkConfig.ts | 24 ++ .../useDeleteCyberArkConfig.ts | 19 + .../useUpdateCyberArkConfig.ts | 19 + .../AdminSettings/CyberArk/CyberArk.test.tsx | 78 ++++ .../AdminSettings/CyberArk/CyberArk.tsx | 232 ++++++++++ .../CyberArk/CyberArkEmptyPlaceholder.tsx | 24 ++ .../CyberArk/EditCyberArkModal.test.tsx | 176 ++++++++ .../CyberArk/EditCyberArkModal.tsx | 176 ++++++++ .../AdminSettings/CyberArk/constants.ts | 12 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 213 +++++++++ 20 files changed, 2195 insertions(+), 16 deletions(-) create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/configOverrides/cyberArkApi.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/configOverrides/useCyberArkConfig.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/configOverrides/useDeleteCyberArkConfig.ts create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/hooks/configOverrides/useUpdateCyberArkConfig.ts create mode 100644 ui/litellm-dashboard/src/components/Settings/AdminSettings/CyberArk/CyberArk.test.tsx create mode 100644 ui/litellm-dashboard/src/components/Settings/AdminSettings/CyberArk/CyberArk.tsx create mode 100644 ui/litellm-dashboard/src/components/Settings/AdminSettings/CyberArk/CyberArkEmptyPlaceholder.tsx create mode 100644 ui/litellm-dashboard/src/components/Settings/AdminSettings/CyberArk/EditCyberArkModal.test.tsx create mode 100644 ui/litellm-dashboard/src/components/Settings/AdminSettings/CyberArk/EditCyberArkModal.tsx create mode 100644 ui/litellm-dashboard/src/components/Settings/AdminSettings/CyberArk/constants.ts diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index ef88ae574fb..95e69bf7f1d 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -24,7 +24,7 @@ "limit": 19 }, "reportExplicitAny": { - "limit": 5486 + "limit": 5485 }, "reportFunctionMemberAccess": { "limit": 7 diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index b5f4fa01a7b..08a35b41126 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -6459,6 +6459,109 @@ "title": "ConfigOverrideSettingsResponse", "type": "object" }, + "CyberArkConfig": { + "description": "Configuration for CyberArk Conjur secret manager integration.", + "properties": { + "client_cert": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "description": "Path to the client TLS certificate for certificate-based authentication", + "title": "Client Cert" + }, + "client_key": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "description": "Path to the client TLS private key for certificate-based authentication", + "title": "Client Key" + }, + "cyberark_account": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "description": "The Conjur organization account name", + "title": "Cyberark Account" + }, + "cyberark_api_base": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "description": "The address of the CyberArk Conjur server (e.g., https://conjur.example.com)", + "title": "Cyberark Api Base" + }, + "cyberark_api_key": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "description": "API key for Conjur API-key authentication", + "title": "Cyberark Api Key" + }, + "cyberark_username": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "description": "The Conjur username (login) to authenticate as", + "title": "Cyberark Username" + }, + "refresh_interval": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "description": "Auth token cache TTL in seconds (default: 300)", + "title": "Refresh Interval" + }, + "ssl_verify": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "description": "Set to false to disable SSL verification (e.g., for self-signed certificates)", + "title": "Ssl Verify" + } + }, + "title": "CyberArkConfig", + "type": "object" + }, "HTTPValidationError": { "properties": { "detail": { @@ -6654,6 +6757,192 @@ } }, "paths": { + "/config_overrides/cyberark": { + "delete": { + "description": "Delete CyberArk Conjur configuration. Idempotent.", + "operationId": "delete_cyberark_config_config_overrides_cyberark_delete", + "parameters": [ + { + "description": "The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability", + "in": "header", + "name": "litellm-changed-by", + "required": false, + "schema": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "description": "The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability", + "title": "Litellm-Changed-By" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "additionalProperties": { + "type": "string" + }, + "title": "Response Delete Cyberark Config Config Overrides Cyberark Delete", + "type": "object" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Delete Cyberark Config", + "tags": [ + "config_overrides" + ] + }, + "get": { + "description": "Get current CyberArk Conjur configuration.\nReturns decrypted values from DB, or falls back to current env vars.\nSensitive fields are masked before leaving the server.", + "operationId": "get_cyberark_config_config_overrides_cyberark_get", + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/ConfigOverrideSettingsResponse" + } + } + }, + "description": "Successful Response" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Get Cyberark Config", + "tags": [ + "config_overrides" + ] + }, + "post": { + "description": "Update CyberArk Conjur secret manager configuration.\nSets environment variables, encrypts sensitive fields, and stores in DB.\nReinitializes the secret manager on this pod.", + "operationId": "update_cyberark_config_config_overrides_cyberark_post", + "parameters": [ + { + "description": "The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability", + "in": "header", + "name": "litellm-changed-by", + "required": false, + "schema": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "description": "The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability", + "title": "Litellm-Changed-By" + } + } + ], + "requestBody": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/CyberArkConfig" + } + } + }, + "required": true + }, + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "additionalProperties": { + "type": "string" + }, + "title": "Response Update Cyberark Config Config Overrides Cyberark Post", + "type": "object" + } + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Update Cyberark Config", + "tags": [ + "config_overrides" + ] + } + }, + "/config_overrides/cyberark/test_connection": { + "post": { + "description": "Test the connection to the currently configured CyberArk Conjur server.\nUses the already-initialized secret manager client. Does not modify any state.", + "operationId": "test_cyberark_connection_config_overrides_cyberark_test_connection_post", + "responses": { + "200": { + "content": { + "application/json": { + "schema": { + "additionalProperties": { + "type": "string" + }, + "title": "Response Test Cyberark Connection Config Overrides Cyberark Test Connection Post", + "type": "object" + } + } + }, + "description": "Successful Response" + } + }, + "security": [ + { + "APIKeyHeader": [] + } + ], + "summary": "Test Cyberark Connection", + "tags": [ + "config_overrides" + ] + } + }, "/config_overrides/hashicorp_vault": { "delete": { "description": "Delete Hashicorp Vault configuration. Idempotent.", diff --git a/litellm/proxy/management_endpoints/config_override_endpoints.py b/litellm/proxy/management_endpoints/config_override_endpoints.py index 7f4faddf178..84593460704 100644 --- a/litellm/proxy/management_endpoints/config_override_endpoints.py +++ b/litellm/proxy/management_endpoints/config_override_endpoints.py @@ -3,7 +3,7 @@ import json import os from collections.abc import Mapping, Sequence from datetime import datetime, timezone -from typing import TYPE_CHECKING, Any, Final, Protocol +from typing import TYPE_CHECKING, Final, Protocol from fastapi import APIRouter, Depends, Header, HTTPException from pydantic import BaseModel, TypeAdapter @@ -36,10 +36,12 @@ from litellm.repositories.table_repositories import ConfigOverridesRepository from litellm.types.llms.custom_http import httpxSpecialProvider from litellm.types.proxy.management_endpoints.config_overrides import ( ConfigOverrideSettingsResponse, + CyberArkConfig, HashicorpVaultConfig, ) if TYPE_CHECKING: + from litellm.proxy.proxy_server import ProxyConfig from litellm.proxy.utils import PrismaClient router: Final = APIRouter() @@ -83,18 +85,19 @@ def _log_audit_task_exception(task: "asyncio.Task[None]") -> None: return exc: Final = task.exception() if exc is not None: - verbose_proxy_logger.warning("Failed to write hashicorp-vault config audit log: %s", exc) + verbose_proxy_logger.warning("Failed to write config override audit log: %s", exc) -async def _emit_hashicorp_vault_audit_log( +async def _emit_config_override_audit_log( *, + object_id: str, action: AUDIT_ACTIONS, before_config: Mapping[str, object] | None, after_config: Mapping[str, object] | None, user_api_key_dict: UserAPIKeyAuth, litellm_changed_by: str | None, ) -> None: - """Emit an audit-log row for a /config_overrides/hashicorp_vault mutation. + """Emit an audit-log row for a /config_overrides/{object_id} mutation. Mirrors the ``store_audit_logs``-gated pattern from ``team_callback_endpoints.py``. Captured under @@ -118,7 +121,7 @@ async def _emit_hashicorp_vault_audit_log( changed_by=litellm_changed_by or user_api_key_dict.user_id or litellm_proxy_admin_name, changed_by_api_key=user_api_key_dict.api_key, table_name=LitellmTableNames.CONFIG_OVERRIDES_TABLE_NAME, - object_id="hashicorp_vault", + object_id=object_id, action=action, updated_values=json.dumps({"config": _redact_config(after_config)}, default=str), before_value=json.dumps({"config": _redact_config(before_config)}, default=str), @@ -150,6 +153,24 @@ HASHICORP_SENSITIVE_FIELDS: Final[set[str]] = { "client_key", } +# --- CyberArk Conjur constants --- + +CYBERARK_ENV_VAR_MAPPING: Final[dict[str, str]] = { # mutable-ok: module-level env mapping + "cyberark_api_base": "CYBERARK_API_BASE", + "cyberark_account": "CYBERARK_ACCOUNT", + "cyberark_username": "CYBERARK_USERNAME", + "cyberark_api_key": "CYBERARK_API_KEY", + "client_cert": "CYBERARK_CLIENT_CERT", + "client_key": "CYBERARK_CLIENT_KEY", + "ssl_verify": "CYBERARK_SSL_VERIFY", + "refresh_interval": "CYBERARK_REFRESH_INTERVAL", +} + +CYBERARK_SENSITIVE_FIELDS: Final[set[str]] = { # mutable-ok: module-level constant, mirrors HASHICORP_SENSITIVE_FIELDS + "cyberark_api_key", + "client_key", +} + _sensitive_masker: Final = SensitiveDataMasker() @@ -215,9 +236,12 @@ def _parse_config_value(raw: str | Mapping[str, object]) -> dict[str, object]: return dict(raw) -def _set_env_vars(config_data: Mapping[str, object]) -> None: - """Set HCP_VAULT_* env vars from config data. Unsets vars for missing/None/empty fields.""" - for field_name, env_var_name in HASHICORP_ENV_VAR_MAPPING.items(): +def _set_env_vars( + config_data: Mapping[str, object], + env_var_mapping: Mapping[str, str] = HASHICORP_ENV_VAR_MAPPING, +) -> None: + """Set mapped env vars from config data. Unsets vars for missing/None/empty fields.""" + for field_name, env_var_name in env_var_mapping.items(): value = config_data.get(field_name) if value is not None and value != "": os.environ[env_var_name] = str(value) @@ -225,13 +249,74 @@ def _set_env_vars(config_data: Mapping[str, object]) -> None: os.environ.pop(env_var_name, None) -def _clear_hashicorp_vault_state(proxy_config: Any) -> None: +def _clear_hashicorp_vault_state(proxy_config: "ProxyConfig") -> None: """Clear all Hashicorp Vault state: env vars, secret manager, and change-detection cache.""" _set_env_vars({}) if litellm._key_management_system == KeyManagementSystem.HASHICORP_VAULT: litellm.secret_manager_client = None litellm._key_management_system = None - proxy_config._last_hashicorp_vault_config = None + proxy_config._last_hashicorp_vault_config = None # pyright: ignore[reportPrivateUsage] # proxy-internal change-detection cache + + +def _snapshot_cyberark_boot_env(proxy_config: "ProxyConfig") -> None: + """Capture deployment-provided CYBERARK_* env vars once, before the first DB-driven overwrite.""" + if proxy_config._cyberark_boot_env is None: # pyright: ignore[reportPrivateUsage] # proxy-internal boot snapshot + proxy_config._cyberark_boot_env = _get_current_env_values(CYBERARK_ENV_VAR_MAPPING) # pyright: ignore[reportPrivateUsage] # proxy-internal boot snapshot + + +def _restore_cyberark_runtime(proxy_config: "ProxyConfig", env_values: Mapping[str, str | None]) -> None: + """Restore CYBERARK_* env vars and reinitialize (or drop) the secret manager to match them.""" + _set_env_vars(env_values, CYBERARK_ENV_VAR_MAPPING) + if env_values.get("cyberark_api_base"): + try: + proxy_config.initialize_secret_manager(key_management_system="cyberark") + except Exception: # noqa: BLE001 # restore is best-effort; fall through to dropping the manager + verbose_proxy_logger.exception("Failed to restore previous CyberArk configuration") + else: + return + if litellm._key_management_system != KeyManagementSystem.CYBERARK: # pyright: ignore[reportPrivateUsage] # proxy-internal helper, mirrors hashicorp endpoint usage + return + litellm.secret_manager_client = None + litellm._key_management_system = None # pyright: ignore[reportPrivateUsage] # proxy-internal helper, mirrors hashicorp endpoint usage + # Force the vault reload to re-init from its own row so no manager is stranded inactive + proxy_config._last_hashicorp_vault_config = None # pyright: ignore[reportPrivateUsage] # proxy-internal change-detection cache + if os.environ.get("HCP_VAULT_ADDR"): + try: + proxy_config.initialize_secret_manager(key_management_system="hashicorp_vault") + except Exception: # noqa: BLE001 # restore is best-effort; the vault reload loop retries from its own row + verbose_proxy_logger.exception("Failed to reinitialize Hashicorp Vault after CyberArk rollback") + + +def _clear_cyberark_state(proxy_config: "ProxyConfig") -> None: + """Drop DB-driven CyberArk state, restoring deployment-provided env vars if any.""" + boot_env: Final[Mapping[str, str | None]] = ( + proxy_config._cyberark_boot_env or {} # pyright: ignore[reportPrivateUsage] # proxy-internal boot snapshot + ) + _restore_cyberark_runtime(proxy_config, boot_env) + proxy_config._last_cyberark_config = None # pyright: ignore[reportPrivateUsage] # proxy-internal helper, mirrors hashicorp endpoint usage + + +async def _persist_cyberark_config( + prisma_client: "PrismaClient", + proxy_config: "ProxyConfig", + config_data: Mapping[str, object], +) -> dict[str, object]: + """Encrypt and upsert the CyberArk config row; returns the stored (encrypted) payload.""" + encrypted_data: Final = proxy_config._encrypt_env_variables(dict(config_data)) # pyright: ignore[reportPrivateUsage] # proxy-internal helper, mirrors hashicorp endpoint usage + config_value: Final = safe_dumps(encrypted_data) + await _config_overrides_table(prisma_client).upsert( + where={"config_type": "cyberark"}, # mutable-ok: prisma upsert payload + data={ # mutable-ok: prisma upsert payload + "create": { # mutable-ok: prisma upsert payload + "config_type": "cyberark", + "config_value": config_value, + }, + "update": { # mutable-ok: prisma upsert payload + "config_value": config_value, + }, + }, + ) + return safe_json_loads(config_value) # --- Hashicorp Vault endpoints --- @@ -358,7 +443,8 @@ async def update_hashicorp_vault_config( # row was absent or its ``config_value`` was NULL. before_config: Final = existing_decrypted if existing_decrypted is not None else env_values action: Final[AUDIT_ACTIONS] = "updated" if existing_record is not None else "created" - await _emit_hashicorp_vault_audit_log( + await _emit_config_override_audit_log( + object_id="hashicorp_vault", action=action, before_config=before_config, after_config=config_data, @@ -484,7 +570,8 @@ async def delete_hashicorp_vault_config( # Only emit audit log if a row was actually removed; an idempotent # delete on a non-existent row produces no security-relevant change. if deleted: - await _emit_hashicorp_vault_audit_log( + await _emit_config_override_audit_log( + object_id="hashicorp_vault", action="deleted", before_config=before_config, after_config=None, @@ -529,7 +616,7 @@ async def test_hashicorp_vault_connection( # Step 1: Authenticate (exercises AppRole login, TLS cert login, or direct token) try: - headers: Final[dict[str, str]] = await asyncio.to_thread(client._get_request_headers) + headers: Final[Mapping[str, str]] = await asyncio.to_thread(client._get_request_headers) except Exception as e: raise HTTPException( status_code=502, @@ -554,3 +641,298 @@ async def test_hashicorp_vault_connection( "status": "success", "message": f"Successfully connected to Vault at {client.vault_addr}", } + + +# --- CyberArk Conjur endpoints --- + + +@router.post( + "/config_overrides/cyberark", + tags=["Config Overrides"], # mutable-ok: FastAPI route decorator metadata + dependencies=[Depends(user_api_key_auth)], # mutable-ok: FastAPI route decorator metadata +) +async def update_cyberark_config( + config: CyberArkConfig, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), # noqa: B008 # FastAPI dependency injection + litellm_changed_by: str | None = Header( + None, + description="The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability", + ), +) -> dict[str, str]: + """ + Update CyberArk Conjur secret manager configuration. + Sets environment variables, encrypts sensitive fields, and stores in DB. + Reinitializes the secret manager on this pod. + """ + from litellm.proxy.proxy_server import prisma_client, proxy_config + + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + raise HTTPException( + status_code=403, + detail="Only admin users can update config overrides", + ) + + if prisma_client is None: + raise HTTPException( + status_code=500, + detail=CommonProxyErrors.db_not_connected_error.value, + ) + + config_data: dict[str, object] = config.model_dump(exclude_none=True) # mutable-ok: merged # rebind-ok: stripped + + # Merge ALL fields the user didn't send: try DB first, fall back to env vars. + # Omitted field = keep existing; empty string = clear/remove the field. + existing_record: Final = await _config_overrides_table(prisma_client).find_unique( + where={"config_type": "cyberark"} # mutable-ok: prisma where clause + ) + existing_decrypted: dict[str, object] | None = None # mutable-ok: DB payload # rebind-ok: set when record exists + env_values: dict[str, str | None] = {} # mutable-ok: env snapshot # rebind-ok: populated when no DB record exists + if existing_record is not None and existing_record.config_value is not None: + existing_data: Final = _parse_config_value(existing_record.config_value) + existing_decrypted = proxy_config._decrypt_db_variables(existing_data) # pyright: ignore[reportPrivateUsage] # rebind-ok: populated when a prior record decrypts + for field in CYBERARK_ENV_VAR_MAPPING: + if field not in config_data and existing_decrypted.get(field): + config_data[field] = existing_decrypted[field] + else: + env_values = _get_current_env_values(CYBERARK_ENV_VAR_MAPPING) # rebind-ok: populated when no DB record exists + for field in CYBERARK_ENV_VAR_MAPPING: + if field not in config_data and env_values.get(field): + config_data[field] = env_values[field] + + config_data = {k: v for k, v in config_data.items() if v != ""} # mutable-ok: dict # rebind-ok: "" means clear + + has_api_base: Final = bool(config_data.get("cyberark_api_base")) + has_api_key_auth: Final = bool(config_data.get("cyberark_api_key")) + has_tls_cert_auth: Final = bool(config_data.get("client_cert") and config_data.get("client_key")) + + if not has_api_base: + raise HTTPException( + status_code=400, + detail="CyberArk API Base is required", + ) + + if not has_api_key_auth and not has_tls_cert_auth: + raise HTTPException( + status_code=400, + detail="At least one authentication method is required: " + "provide an API Key, or both Client Certificate and Client Key", + ) + + _snapshot_cyberark_boot_env(proxy_config) + previous_env: Final = _get_current_env_values(CYBERARK_ENV_VAR_MAPPING) + _set_env_vars(config_data, CYBERARK_ENV_VAR_MAPPING) + + try: + proxy_config.initialize_secret_manager(key_management_system="cyberark") + except Exception as e: # noqa: BLE001 # any init failure must roll back env vars + _set_env_vars(previous_env, CYBERARK_ENV_VAR_MAPPING) + verbose_proxy_logger.exception("Error reinitializing CyberArk secret manager: %s", str(e)) + raise HTTPException( + status_code=500, + detail=f"Failed to initialize secret manager: {e}", + ) + + try: + proxy_config._last_cyberark_config = await _persist_cyberark_config( # pyright: ignore[reportPrivateUsage] # proxy-internal helper, mirrors hashicorp endpoint usage + prisma_client, proxy_config, config_data + ) + except Exception as e: # noqa: BLE001 # persistence failure must roll back the runtime state set above + _restore_cyberark_runtime(proxy_config, previous_env) + verbose_proxy_logger.exception("Error persisting CyberArk configuration: %s", str(e)) + raise HTTPException( + status_code=500, + detail=f"Failed to persist CyberArk configuration: {e}", + ) + + before_config: Final = existing_decrypted if existing_decrypted is not None else env_values + action: Final[AUDIT_ACTIONS] = "updated" if existing_record is not None else "created" + await _emit_config_override_audit_log( + object_id="cyberark", + action=action, + before_config=before_config, + after_config=config_data, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=litellm_changed_by, + ) + + return { # mutable-ok: JSON response payload + "message": "CyberArk configuration updated successfully", + "status": "success", + } + + +@router.get( + "/config_overrides/cyberark", + tags=["Config Overrides"], # mutable-ok: FastAPI route decorator metadata + dependencies=[Depends(user_api_key_auth)], # mutable-ok: FastAPI route decorator metadata + response_model=ConfigOverrideSettingsResponse, +) +async def get_cyberark_config( + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), # noqa: B008 # FastAPI dependency injection +) -> ConfigOverrideSettingsResponse: + """ + Get current CyberArk Conjur configuration. + Returns decrypted values from DB, or falls back to current env vars. + Sensitive fields are masked before leaving the server. + """ + from litellm.proxy.management_endpoints.common_utils import ( + _user_has_admin_view, # pyright: ignore[reportPrivateUsage] # proxy-internal helper, mirrors hashicorp endpoint usage + ) + from litellm.proxy.proxy_server import prisma_client, proxy_config + + if not _user_has_admin_view(user_api_key_dict): + raise HTTPException( + status_code=403, + detail="Only admin users can view config overrides", + ) + + if prisma_client is None: + raise HTTPException( + status_code=500, + detail=CommonProxyErrors.db_not_connected_error.value, + ) + + field_schema: Final = _build_field_schema(CyberArkConfig) + + db_record: Final = await _config_overrides_table(prisma_client).find_unique( + where={"config_type": "cyberark"} + ) # mutable-ok: prisma where clause + + if db_record is not None and db_record.config_value is not None: + config_data: Final = _parse_config_value(db_record.config_value) + decrypted_data: Final[Mapping[str, object]] = proxy_config._decrypt_db_variables(config_data) # pyright: ignore[reportPrivateUsage] # proxy-internal helper, mirrors hashicorp endpoint usage + masked_data: Final = _mask_sensitive_fields(decrypted_data, CYBERARK_SENSITIVE_FIELDS) + + return ConfigOverrideSettingsResponse( + config_type="cyberark", + values=masked_data, + field_schema=field_schema, + ) + + env_values: Final = _get_current_env_values(CYBERARK_ENV_VAR_MAPPING) + masked_env_values: Final = _mask_sensitive_fields(env_values, CYBERARK_SENSITIVE_FIELDS) + + return ConfigOverrideSettingsResponse( + config_type="cyberark", + values=masked_env_values, + field_schema=field_schema, + ) + + +@router.delete( + "/config_overrides/cyberark", + tags=["Config Overrides"], # mutable-ok: FastAPI route decorator metadata + dependencies=[Depends(user_api_key_auth)], # mutable-ok: FastAPI route decorator metadata +) +async def delete_cyberark_config( + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), # noqa: B008 # FastAPI dependency injection + litellm_changed_by: str | None = Header( + None, + description="The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability", + ), +) -> dict[str, str]: + """Delete CyberArk Conjur configuration. Idempotent.""" + from litellm.proxy.proxy_server import prisma_client, proxy_config + + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + raise HTTPException( + status_code=403, + detail="Only admin users can delete config overrides", + ) + + if prisma_client is None: + raise HTTPException( + status_code=500, + detail=CommonProxyErrors.db_not_connected_error.value, + ) + + existing_record: Final = await _config_overrides_table(prisma_client).find_unique( + where={"config_type": "cyberark"} # mutable-ok: prisma where clause + ) + before_config: dict[str, object] | None = None # mutable-ok: audit snapshot # rebind-ok: set when decrypts + if existing_record is not None and existing_record.config_value is not None: + try: + before_config = proxy_config._decrypt_db_variables(_parse_config_value(existing_record.config_value)) # pyright: ignore[reportPrivateUsage] # rebind-ok: populated when the prior record decrypts + except Exception: # noqa: BLE001 # undecryptable prior config must not block deletion + before_config = None # rebind-ok: reset when decryption fails + + deleted = False # rebind-ok: set true once the DB row is removed + try: + await _config_overrides_table(prisma_client).delete( + where={"config_type": "cyberark"} + ) # mutable-ok: prisma where clause + deleted = True # rebind-ok: set true once the DB row is removed + except RecordNotFoundError: + verbose_proxy_logger.debug("No existing CyberArk config record to delete") + + _clear_cyberark_state(proxy_config) + + if deleted: + await _emit_config_override_audit_log( + object_id="cyberark", + action="deleted", + before_config=before_config, + after_config=None, + user_api_key_dict=user_api_key_dict, + litellm_changed_by=litellm_changed_by, + ) + + return { # mutable-ok: JSON response payload + "message": "CyberArk configuration deleted successfully", + "status": "success", + } + + +@router.post( + "/config_overrides/cyberark/test_connection", + tags=["Config Overrides"], # mutable-ok: FastAPI route decorator metadata + dependencies=[Depends(user_api_key_auth)], # mutable-ok: FastAPI route decorator metadata +) +async def test_cyberark_connection( + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), # noqa: B008 # FastAPI dependency injection +) -> dict[str, str]: + """ + Test the connection to the currently configured CyberArk Conjur server. + Uses the already-initialized secret manager client. Does not modify any state. + """ + from litellm.secret_managers.cyberark_secret_manager import CyberArkSecretManager + + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + raise HTTPException( + status_code=403, + detail="Only admin users can test CyberArk connection", + ) + + client: Final = litellm.secret_manager_client + if not isinstance(client, CyberArkSecretManager): + raise HTTPException( + status_code=400, + detail="CyberArk is not configured. Save a configuration first.", + ) + + try: + headers: Final[Mapping[str, str]] = await asyncio.to_thread(client._get_request_headers) # pyright: ignore[reportPrivateUsage] # proxy-internal helper, mirrors hashicorp endpoint usage + except Exception as e: # noqa: BLE001 # surface any auth failure as a 502 with detail + raise HTTPException( + status_code=502, + detail=f"CyberArk authentication failed: {e}", + ) + + try: + async_client: Final = get_async_httpx_client( + llm_provider=httpxSpecialProvider.SecretManager, + params={"ssl_verify": client.ssl_verify}, # mutable-ok: httpx client params + ) + whoami_url: Final = f"{client.conjur_addr}/whoami" + response: Final = await async_client.get(whoami_url, headers=headers) + response.raise_for_status() + except Exception as e: # noqa: BLE001 # surface any connectivity/TLS failure as a 502 with detail + raise HTTPException( + status_code=502, + detail=f"CyberArk token validation failed: {e}", + ) + + return { # mutable-ok: JSON response payload + "status": "success", + "message": f"Successfully connected to CyberArk Conjur at {client.conjur_addr}", + } diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index af0aa9743bc..53a99065f2d 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -4281,6 +4281,8 @@ class ProxyConfig: self.config: dict[str, Any] = {} self._last_semantic_filter_config: dict[str, object] | None = None self._last_hashicorp_vault_config: dict[str, object] | None = None + self._last_cyberark_config: dict[str, object] | None = None # mutable-ok: change-detection cache + self._cyberark_boot_env: dict[str, str | None] | None = None # mutable-ok: deployment env snapshot, set once self.worker_registry: list[WorkerRegistryEntry] = [] self.config_sync_subscriber: ConfigSyncSubscriber | None = None self.auth_cache_invalidation_subscriber: AuthCacheInvalidationSubscriber | None = None @@ -6977,6 +6979,7 @@ class ProxyConfig: if self._should_load_db_object(object_type="config_overrides"): await self._init_hashicorp_vault_config_override(prisma_client=prisma_client) + await self._init_cyberark_config_override(prisma_client=prisma_client) await self._apply_safe_litellm_settings_overrides_from_db(prisma_client=prisma_client) @@ -7141,6 +7144,64 @@ class ProxyConfig: str(e), ) + async def _init_cyberark_config_override(self, prisma_client: PrismaClient) -> None: + """ + Load CyberArk Conjur config override from DB. + Decrypts sensitive fields, sets CYBERARK_* env vars, and reinitializes the secret manager. + Called periodically via _init_non_llm_objects_in_db to sync config across pods. + """ + from litellm.proxy.management_endpoints.config_override_endpoints import ( + CYBERARK_ENV_VAR_MAPPING, + _clear_cyberark_state, # pyright: ignore[reportPrivateUsage] # module-internal helper shared with the endpoint module + _get_current_env_values, # pyright: ignore[reportPrivateUsage] # module-internal helper shared with the endpoint module + _parse_config_value, # pyright: ignore[reportPrivateUsage] # module-internal helper shared with the endpoint module + _set_env_vars, # pyright: ignore[reportPrivateUsage] # module-internal helper shared with the endpoint module + _snapshot_cyberark_boot_env, # pyright: ignore[reportPrivateUsage] # module-internal helper shared with the endpoint module + ) + + try: + db_record: Final[_ConfigOverridesRow | None] = cast( # cast-ok: prisma Json stub is `str`, runtime dict + "_ConfigOverridesRow | None", + await call_with_db_reconnect_retry( + prisma_client, + lambda: ConfigOverridesRepository(prisma_client).table.find_unique( + where={"config_type": "cyberark"} # mutable-ok: prisma where clause + ), + reason="init_cyberark_config_override_lookup_failure", + ), + ) + + if db_record is None or db_record.config_value is None: + if self._last_cyberark_config is not None: + _clear_cyberark_state(self) + return + + config_data: Final = _parse_config_value(db_record.config_value) + + # Skip reinit if config hasn't changed since last poll + if self._last_cyberark_config == config_data: + return + + decrypted_data: Final = self._decrypt_db_variables(config_data) + + _snapshot_cyberark_boot_env(self) + previous_env: Final = _get_current_env_values(CYBERARK_ENV_VAR_MAPPING) + _set_env_vars(decrypted_data, CYBERARK_ENV_VAR_MAPPING) + + try: + self.initialize_secret_manager(key_management_system="cyberark") + except Exception: + _set_env_vars(previous_env, CYBERARK_ENV_VAR_MAPPING) + raise + + self._last_cyberark_config = config_data.copy() + verbose_proxy_logger.debug("CyberArk config override loaded from DB") + except Exception as e: # noqa: BLE001 # any DB/decrypt/init failure must not break proxy boot + verbose_proxy_logger.exception( + "Error loading CyberArk config override from DB: %s", + str(e), + ) + async def check_periodic_reloads(self, prisma_client: PrismaClient): """ Run the admin-configured periodic model cost map reload. diff --git a/litellm/types/proxy/management_endpoints/config_overrides.py b/litellm/types/proxy/management_endpoints/config_overrides.py index 9e1ea23ac46..f9cba6983db 100644 --- a/litellm/types/proxy/management_endpoints/config_overrides.py +++ b/litellm/types/proxy/management_endpoints/config_overrides.py @@ -52,6 +52,43 @@ class HashicorpVaultConfig(BaseModel): ) +class CyberArkConfig(BaseModel): + """Configuration for CyberArk Conjur secret manager integration.""" + + cyberark_api_base: str | None = Field( + default=None, + description="The address of the CyberArk Conjur server (e.g., https://conjur.example.com)", + ) + cyberark_account: str | None = Field( + default=None, + description="The Conjur organization account name", + ) + cyberark_username: str | None = Field( + default=None, + description="The Conjur username (login) to authenticate as", + ) + cyberark_api_key: str | None = Field( + default=None, + description="API key for Conjur API-key authentication", + ) + client_cert: str | None = Field( + default=None, + description="Path to the client TLS certificate for certificate-based authentication", + ) + client_key: str | None = Field( + default=None, + description="Path to the client TLS private key for certificate-based authentication", + ) + ssl_verify: str | None = Field( + default=None, + description="Set to false to disable SSL verification (e.g., for self-signed certificates)", + ) + refresh_interval: str | None = Field( + default=None, + description="Auth token cache TTL in seconds (default: 300)", + ) + + class ConfigOverrideSettingsResponse(BaseModel): """Response model for config override settings GET endpoints.""" diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index c60988eccc0..b4107582000 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -24,7 +24,7 @@ "limit": 133 }, "ANN401": { - "limit": 655 + "limit": 654 }, "ASYNC230": { "limit": 11 @@ -231,7 +231,7 @@ "limit": 5 }, "TID251": { - "limit": 1117 + "limit": 1116 }, "TRY002": { "limit": 524 diff --git a/tests/test_litellm/proxy/management_endpoints/test_config_override_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_config_override_endpoints.py index d90d589c504..03f94fbe94c 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_config_override_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_config_override_endpoints.py @@ -11,16 +11,19 @@ import litellm import litellm.proxy.proxy_server as ps from litellm.proxy._types import KeyManagementSystem, LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.management_endpoints.config_override_endpoints import ( + CYBERARK_ENV_VAR_MAPPING, HASHICORP_ENV_VAR_MAPPING, _build_field_schema, _set_env_vars, ) from litellm.proxy.proxy_server import app from litellm.types.proxy.management_endpoints.config_overrides import ( + CyberArkConfig, HashicorpVaultConfig, ) VAULT_URL = "/config_overrides/hashicorp_vault" +CYBERARK_URL = "/config_overrides/cyberark" @pytest.fixture @@ -42,6 +45,7 @@ def _make_mock_proxy_config(): cfg = MagicMock() cfg.initialize_secret_manager = MagicMock() cfg._last_hashicorp_vault_config = None + cfg._cyberark_boot_env = None cfg._encrypt_env_variables = MagicMock( side_effect=lambda d: {k: f"enc_{v}" for k, v in d.items()} ) @@ -67,6 +71,8 @@ def _cleanup(): app.dependency_overrides.pop(ps.user_api_key_auth, None) for env_var in HASHICORP_ENV_VAR_MAPPING.values(): os.environ.pop(env_var, None) + for env_var in CYBERARK_ENV_VAR_MAPPING.values(): + os.environ.pop(env_var, None) def _set_admin(): @@ -275,6 +281,391 @@ async def test_hashicorp_vault_validation_errors_and_access_control( _cleanup() +@pytest.mark.asyncio +async def test_cyberark_crud_lifecycle(client, monkeypatch): + """Create → read (masked) → partial update (merge from DB) → clear field → + delete → idempotent delete → env fallback → merge from env → schema.""" + mock_prisma, mock_db = _make_mock_db() + mock_cfg = _make_mock_proxy_config() + mock_cfg._last_cyberark_config = None + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + monkeypatch.setattr(ps, "proxy_config", mock_cfg) + old_client, old_kms = litellm.secret_manager_client, litellm._key_management_system + _set_admin() + + try: + # 1. POST: create with API-key auth + r = client.post( + CYBERARK_URL, + json={ + "cyberark_api_base": "https://conjur.example.com", + "cyberark_account": "myorg", + "cyberark_username": "litellm-user", + "cyberark_api_key": "my-secret-api-key", + }, + ) + assert r.status_code == 200 + assert os.environ["CYBERARK_API_BASE"] == "https://conjur.example.com" + assert os.environ["CYBERARK_API_KEY"] == "my-secret-api-key" + data = _upserted_data(mock_db) + assert data["cyberark_api_key"] == "enc_my-secret-api-key" + mock_cfg.initialize_secret_manager.assert_called_with( + key_management_system="cyberark" + ) + assert mock_cfg._last_cyberark_config is not None + + # 2. GET: sensitive fields masked + mock_db.find_unique = AsyncMock(return_value=_db_record(data)) + r = client.get(CYBERARK_URL) + assert r.status_code == 200 + vals = r.json()["values"] + assert vals["cyberark_api_base"] == "https://conjur.example.com" + assert "*" in vals["cyberark_api_key"] + assert "properties" in r.json()["field_schema"] + + # 3. POST partial: omitted fields merge from DB + r = client.post(CYBERARK_URL, json={"cyberark_api_base": "https://conjur.new.com"}) + assert r.status_code == 200 + data = _upserted_data(mock_db) + assert data["cyberark_api_base"] == "enc_https://conjur.new.com" + assert data["cyberark_api_key"] == "enc_my-secret-api-key" + assert data["cyberark_account"] == "enc_myorg" + + # 4. POST empty string: clears field, switches to cert auth + step3 = { + **data, + "client_cert": "enc_/certs/client.pem", + "client_key": "enc_/certs/client.key", + } + mock_db.find_unique = AsyncMock(return_value=_db_record(step3)) + mock_db.upsert = AsyncMock(return_value=None) + r = client.post(CYBERARK_URL, json={"cyberark_api_key": ""}) + assert r.status_code == 200 + data = _upserted_data(mock_db) + assert "cyberark_api_key" not in data + assert data["client_cert"] == "enc_/certs/client.pem" + + # 5. DELETE: clears everything + litellm.secret_manager_client = MagicMock() # test-quality-ok: endpoint hot-reloads litellm globals; test must set and restore them + litellm._key_management_system = KeyManagementSystem.CYBERARK # test-quality-ok: endpoint hot-reloads litellm globals; test must set and restore them + r = client.delete(CYBERARK_URL) + assert r.status_code == 200 + assert os.environ.get("CYBERARK_API_BASE") is None + assert litellm.secret_manager_client is None + assert mock_cfg._last_cyberark_config is None + + # 6. DELETE idempotent + mock_db.delete = AsyncMock( + side_effect=RecordNotFoundError( + data={"clientVersion": "0.0.0"}, message="Not found" + ) + ) + assert client.delete(CYBERARK_URL).status_code == 200 + + # 7. GET: env var fallback with masking + mock_db.find_unique = AsyncMock(return_value=None) + monkeypatch.setenv("CYBERARK_API_BASE", "https://conjur.env.com") + monkeypatch.setenv("CYBERARK_API_KEY", "env-api-key") + r = client.get(CYBERARK_URL) + vals = r.json()["values"] + assert vals["cyberark_api_base"] == "https://conjur.env.com" + assert "*" in vals["cyberark_api_key"] + + # 8. POST: merge from env vars + mock_cfg.initialize_secret_manager = MagicMock() + mock_db.upsert = AsyncMock(return_value=None) + r = client.post(CYBERARK_URL, json={"cyberark_api_base": "https://conjur.merged.com"}) + assert r.status_code == 200 + data = _upserted_data(mock_db) + assert data["cyberark_api_key"] == "enc_env-api-key" + + # 9. _build_field_schema + schema = _build_field_schema(CyberArkConfig) + assert "cyberark_api_base" in schema["properties"] + assert len(schema["properties"]["cyberark_api_base"]["description"]) > 0 + + finally: + litellm.secret_manager_client = old_client # test-quality-ok: endpoint hot-reloads litellm globals; test must set and restore them + litellm._key_management_system = old_kms # test-quality-ok: endpoint hot-reloads litellm globals; test must set and restore them + _cleanup() + + +@pytest.mark.asyncio +async def test_cyberark_validation_errors_and_access_control(client, monkeypatch): + """Validation (missing api base, missing auth, init failure rollback), + DELETE preserves non-CyberArk secret managers, non-admin 403.""" + mock_prisma, mock_db = _make_mock_db() + mock_cfg = MagicMock() + mock_cfg._last_cyberark_config = {"cyberark_api_base": "old"} + mock_cfg._cyberark_boot_env = None + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + monkeypatch.setattr(ps, "proxy_config", mock_cfg) + old_client, old_kms = litellm.secret_manager_client, litellm._key_management_system + _set_admin() + + try: + # 1. Missing cyberark_api_base → 400 + r = client.post(CYBERARK_URL, json={"cyberark_api_key": "key"}) + assert r.status_code == 400 + assert "API Base" in r.json()["detail"] + + # 2. Missing auth → 400 (cert without key is not valid auth) + r = client.post( + CYBERARK_URL, + json={"cyberark_api_base": "https://c.com", "client_cert": "/c.pem"}, + ) + assert r.status_code == 400 + assert "authentication" in r.json()["detail"].lower() + + # 3. Init failure → 500, env vars restored, nothing persisted + mock_cfg.initialize_secret_manager = MagicMock(side_effect=Exception("fail")) + monkeypatch.setenv("CYBERARK_API_BASE", "https://conjur.old.com") + monkeypatch.setenv("CYBERARK_API_KEY", "old-key") + r = client.post( + CYBERARK_URL, + json={"cyberark_api_base": "https://bad.com", "cyberark_api_key": "bad"}, + ) + assert r.status_code == 500 + assert os.environ["CYBERARK_API_BASE"] == "https://conjur.old.com" + mock_db.upsert.assert_not_awaited() + + # 4. DELETE preserves non-CyberArk secret manager + aws = MagicMock() + litellm.secret_manager_client = aws # test-quality-ok: endpoint hot-reloads litellm globals; test must set and restore them + litellm._key_management_system = KeyManagementSystem.AWS_SECRET_MANAGER # test-quality-ok: endpoint hot-reloads litellm globals; test must set and restore them + assert client.delete(CYBERARK_URL).status_code == 200 + assert litellm.secret_manager_client is aws + + # 5. Non-admin → 403 + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, user_id="user" + ) + assert client.get(CYBERARK_URL).status_code == 403 + assert ( + client.post( + CYBERARK_URL, json={"cyberark_api_base": "https://c.com"} + ).status_code + == 403 + ) + assert client.delete(CYBERARK_URL).status_code == 403 + assert client.post(CYBERARK_URL + "/test_connection").status_code == 403 + + finally: + litellm.secret_manager_client = old_client # test-quality-ok: endpoint hot-reloads litellm globals; test must set and restore them + litellm._key_management_system = old_kms # test-quality-ok: endpoint hot-reloads litellm globals; test must set and restore them + _cleanup() + + +@pytest.mark.asyncio +async def test_cyberark_delete_restores_deployment_env_config(client, monkeypatch): + """Deleting the DB override must restore env vars the deployment started with, + and reinitialize the manager from them, instead of wiping CyberArk entirely.""" + mock_prisma, mock_db = _make_mock_db() + mock_cfg = _make_mock_proxy_config() + mock_cfg._last_cyberark_config = None + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + monkeypatch.setattr(ps, "proxy_config", mock_cfg) + old_client, old_kms = litellm.secret_manager_client, litellm._key_management_system + _set_admin() + + try: + monkeypatch.setenv("CYBERARK_API_BASE", "https://conjur.boot.com") + monkeypatch.setenv("CYBERARK_API_KEY", "boot-key") + + r = client.post( + CYBERARK_URL, + json={"cyberark_api_base": "https://conjur.db.com", "cyberark_api_key": "db-key"}, + ) + assert r.status_code == 200 + assert os.environ["CYBERARK_API_BASE"] == "https://conjur.db.com" + + mock_cfg.initialize_secret_manager.reset_mock() + r = client.delete(CYBERARK_URL) + assert r.status_code == 200 + assert os.environ["CYBERARK_API_BASE"] == "https://conjur.boot.com" + assert os.environ["CYBERARK_API_KEY"] == "boot-key" + mock_cfg.initialize_secret_manager.assert_called_with(key_management_system="cyberark") + assert mock_cfg._last_cyberark_config is None + finally: + litellm.secret_manager_client = old_client # test-quality-ok: endpoint hot-reloads litellm globals; test must set and restore them + litellm._key_management_system = old_kms # test-quality-ok: endpoint hot-reloads litellm globals; test must set and restore them + _cleanup() + + +@pytest.mark.asyncio +async def test_cyberark_persist_failure_rolls_back_runtime_state(client, monkeypatch): + """If the DB upsert fails after the manager was reinitialized, the endpoint + must restore the previous env vars and reinitialize from them, so this pod + does not keep serving credentials that were never committed to the DB.""" + mock_prisma, mock_db = _make_mock_db() + mock_cfg = _make_mock_proxy_config() + mock_cfg._last_cyberark_config = None + mock_db.upsert = AsyncMock(side_effect=Exception("db write failed")) + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + monkeypatch.setattr(ps, "proxy_config", mock_cfg) + old_client, old_kms = litellm.secret_manager_client, litellm._key_management_system + _set_admin() + + try: + monkeypatch.setenv("CYBERARK_API_BASE", "https://conjur.prev.com") + monkeypatch.setenv("CYBERARK_API_KEY", "prev-key") + + r = client.post( + CYBERARK_URL, + json={"cyberark_api_base": "https://conjur.new.com", "cyberark_api_key": "new-key"}, + ) + assert r.status_code == 500 + assert "persist" in r.json()["detail"].lower() + assert os.environ["CYBERARK_API_BASE"] == "https://conjur.prev.com" + assert os.environ["CYBERARK_API_KEY"] == "prev-key" + # last call must be the rollback reinit against the restored env + assert ( + mock_cfg.initialize_secret_manager.call_args_list[-1].kwargs["key_management_system"] == "cyberark" + ) + assert os.environ.get("CYBERARK_API_BASE") != "https://conjur.new.com" + finally: + litellm.secret_manager_client = old_client # test-quality-ok: endpoint hot-reloads litellm globals; test must set and restore them + litellm._key_management_system = old_kms # test-quality-ok: endpoint hot-reloads litellm globals; test must set and restore them + _cleanup() + + +@pytest.mark.asyncio +async def test_cyberark_persist_failure_restores_hashicorp_manager(client, monkeypatch): + """If CyberArk init displaced an env-configured Hashicorp manager and the DB + upsert then fails, rollback must bring the Hashicorp manager back.""" + mock_prisma, mock_db = _make_mock_db() + mock_cfg = _make_mock_proxy_config() + mock_cfg._last_cyberark_config = None + mock_db.upsert = AsyncMock(side_effect=Exception("db write failed")) + + def _fake_init(key_management_system): + litellm._key_management_system = ( # test-quality-ok: endpoint hot-reloads litellm globals; test must set and restore them + KeyManagementSystem.CYBERARK + if key_management_system == "cyberark" + else KeyManagementSystem.HASHICORP_VAULT + ) + + mock_cfg.initialize_secret_manager = MagicMock(side_effect=_fake_init) + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + monkeypatch.setattr(ps, "proxy_config", mock_cfg) + old_client, old_kms = litellm.secret_manager_client, litellm._key_management_system + _set_admin() + + try: + monkeypatch.setenv("HCP_VAULT_ADDR", "https://vault.example.com") + litellm._key_management_system = KeyManagementSystem.HASHICORP_VAULT # test-quality-ok: endpoint hot-reloads litellm globals; test must set and restore them + + r = client.post( + CYBERARK_URL, + json={"cyberark_api_base": "https://conjur.new.com", "cyberark_api_key": "new-key"}, + ) + assert r.status_code == 500 + assert litellm._key_management_system == KeyManagementSystem.HASHICORP_VAULT + assert ( + mock_cfg.initialize_secret_manager.call_args_list[-1].kwargs["key_management_system"] == "hashicorp_vault" + ) + finally: + litellm.secret_manager_client = old_client # test-quality-ok: endpoint hot-reloads litellm globals; test must set and restore them + litellm._key_management_system = old_kms # test-quality-ok: endpoint hot-reloads litellm globals; test must set and restore them + os.environ.pop("HCP_VAULT_ADDR", None) + _cleanup() + + +@pytest.mark.asyncio +async def test_cyberark_audit_log_redacts_values(client, monkeypatch): + monkeypatch.setattr(litellm, "store_audit_logs", True) + mock_prisma, mock_db = _make_mock_db() + mock_cfg = _make_mock_proxy_config() + mock_cfg._last_cyberark_config = None + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + monkeypatch.setattr(ps, "proxy_config", mock_cfg) + _set_admin() + + audit_calls = [] + + async def capture(request_data): + audit_calls.append(request_data) + + try: + with patch( # test-quality-ok: patching proxy-internal collaborator to isolate the endpoint + "litellm.proxy.management_helpers.audit_logs.create_audit_log_for_update", + new=capture, + ): + r = client.post( + CYBERARK_URL, + json={ + "cyberark_api_base": "https://conjur.example.com", + "cyberark_api_key": "my-very-secret-key", + }, + ) + assert r.status_code == 200 + for _ in range(3): + await asyncio.sleep(0) + + assert len(audit_calls) == 1 + log = audit_calls[0] + assert log.action == "created" + assert log.object_id == "cyberark" + assert "my-very-secret-key" not in log.updated_values + assert "conjur.example.com" not in log.updated_values + after = json.loads(log.updated_values) + assert "cyberark_api_key" in after["config"] + assert "cyberark_api_base" in after["config"] + finally: + _cleanup() + + +@pytest.mark.asyncio +async def test_cyberark_test_connection(client, monkeypatch): + """400 when not configured; success path authenticates and hits /whoami.""" + from litellm.secret_managers.cyberark_secret_manager import CyberArkSecretManager + + old_client, old_kms = litellm.secret_manager_client, litellm._key_management_system + _set_admin() + + try: + # Not configured → 400 + litellm.secret_manager_client = None # test-quality-ok: endpoint hot-reloads litellm globals; test must set and restore them + r = client.post(CYBERARK_URL + "/test_connection") + assert r.status_code == 400 + assert "not configured" in r.json()["detail"].lower() + + # Configured → authenticates and calls /whoami + mock_manager = MagicMock(spec=CyberArkSecretManager) + mock_manager.conjur_addr = "https://conjur.example.com" + mock_manager.ssl_verify = True + mock_manager._get_request_headers = MagicMock( + return_value={"Authorization": "Token abc"} + ) + litellm.secret_manager_client = mock_manager # test-quality-ok: endpoint hot-reloads litellm globals; test must set and restore them + + mock_response = MagicMock() + mock_response.raise_for_status = MagicMock() + mock_http = MagicMock() + mock_http.get = AsyncMock(return_value=mock_response) + with patch( # test-quality-ok: patching proxy-internal collaborator to isolate the endpoint + "litellm.proxy.management_endpoints.config_override_endpoints.get_async_httpx_client", + return_value=mock_http, + ): + r = client.post(CYBERARK_URL + "/test_connection") + assert r.status_code == 200 + assert "conjur.example.com" in r.json()["message"] + called_url = mock_http.get.call_args.args[0] + assert called_url == "https://conjur.example.com/whoami" + + # Auth failure → 502 + mock_manager._get_request_headers = MagicMock( + side_effect=Exception("bad credentials") + ) + r = client.post(CYBERARK_URL + "/test_connection") + assert r.status_code == 502 + assert "authentication failed" in r.json()["detail"].lower() + finally: + litellm.secret_manager_client = old_client # test-quality-ok: endpoint hot-reloads litellm globals; test must set and restore them + litellm._key_management_system = old_kms # test-quality-ok: endpoint hot-reloads litellm globals; test must set and restore them + _cleanup() + + # ── Audit-log emission for /config_overrides/hashicorp_vault ───────────────── diff --git a/ui/litellm-dashboard/src/app/(dashboard)/admin-panel/_components/AdminPanel.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/admin-panel/_components/AdminPanel.test.tsx index 8c3ca7bd9ff..ee55c568bac 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/admin-panel/_components/AdminPanel.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/admin-panel/_components/AdminPanel.test.tsx @@ -9,6 +9,8 @@ const mockAddAllowedIP = vi.fn(); const mockDeleteAllowedIP = vi.fn(); vi.mock("@/components/networking", () => ({ + getProxyBaseUrl: () => "http://localhost:4000", + getGlobalLitellmHeaderName: () => "Authorization", getSSOSettings: (...args: unknown[]) => mockGetSSOSettings(...args), getAllowedIPs: (...args: unknown[]) => mockGetAllowedIPs(...args), addAllowedIP: (...args: unknown[]) => mockAddAllowedIP(...args), diff --git a/ui/litellm-dashboard/src/app/(dashboard)/admin-panel/_components/AdminPanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/admin-panel/_components/AdminPanel.tsx index a98eb50ce77..1f35f46dcd4 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/admin-panel/_components/AdminPanel.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/admin-panel/_components/AdminPanel.tsx @@ -18,6 +18,7 @@ import LoggingSettings from "@/components/Settings/AdminSettings/LoggingSettings import SSOSettings from "@/components/Settings/AdminSettings/SSOSettings/SSOSettings"; import UISettings from "@/components/Settings/AdminSettings/UISettings/UISettings"; import UserBannerSettings from "@/components/Settings/AdminSettings/UserBannerSettings/UserBannerSettings"; +import CyberArk from "@/components/Settings/AdminSettings/CyberArk/CyberArk"; import HashicorpVault from "@/components/Settings/AdminSettings/HashicorpVault/HashicorpVault"; import PluginSettings from "@/components/Settings/AdminSettings/PluginSettings/PluginSettings"; import SSOModals from "@/components/SSOModals"; @@ -395,6 +396,11 @@ const AdminPanel: React.FC = ({ proxySettings }) => { label: "Hashicorp Vault", children: , }, + { + key: "cyberark", + label: "CyberArk Conjur", + children: , + }, { key: "plugins", label: "Plugins", diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/configOverrides/cyberArkApi.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/configOverrides/cyberArkApi.ts new file mode 100644 index 00000000000..910fe2b17e5 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/configOverrides/cyberArkApi.ts @@ -0,0 +1,38 @@ +import { getProxyBaseUrl, getGlobalLitellmHeaderName } from "@/components/networking"; +import { createApiClient } from "@/lib/http/client"; + +export interface CyberArkFieldSchema { + description?: string; + properties: Record; +} + +export interface CyberArkConfigResponse { + config_type: string; + values: Record; + field_schema: CyberArkFieldSchema; +} + +export interface CyberArkStatusResponse { + status: string; + message: string; +} + +const apiClient = createApiClient({ + getBaseUrl: getProxyBaseUrl, + getAuthHeaderName: getGlobalLitellmHeaderName, +}); + +export const getCyberArkConfig = async (accessToken: string): Promise => + apiClient.get("/config_overrides/cyberark", { accessToken }); + +export const updateCyberArkConfig = async ( + accessToken: string, + config: Record, +): Promise => + apiClient.post("/config_overrides/cyberark", { accessToken, body: config }); + +export const deleteCyberArkConfig = async (accessToken: string): Promise => + apiClient.delete("/config_overrides/cyberark", { accessToken }); + +export const testCyberArkConnection = async (accessToken: string): Promise => + apiClient.post("/config_overrides/cyberark/test_connection", { accessToken }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/configOverrides/useCyberArkConfig.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/configOverrides/useCyberArkConfig.ts new file mode 100644 index 00000000000..cfc3acdfe26 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/configOverrides/useCyberArkConfig.ts @@ -0,0 +1,24 @@ +import { getCyberArkConfig, type CyberArkConfigResponse } from "./cyberArkApi"; +import { useQuery } from "@tanstack/react-query"; +import useAuthorized from "../useAuthorized"; +import { createQueryKeys } from "../common/queryKeysFactory"; + +export const cyberArkKeys = createQueryKeys("cyberArkConfig"); + +export const useCyberArkConfig = () => { + const { accessToken } = useAuthorized(); + + const queryOptions = { + queryKey: cyberArkKeys.list({}), + queryFn: async () => { + if (!accessToken) { + throw new Error("Access token is required"); + } + return getCyberArkConfig(accessToken); + }, + enabled: !!accessToken, + staleTime: 60 * 60 * 1000, + gcTime: 60 * 60 * 1000, + }; + return useQuery(queryOptions); +}; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/configOverrides/useDeleteCyberArkConfig.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/configOverrides/useDeleteCyberArkConfig.ts new file mode 100644 index 00000000000..cebba3a202d --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/configOverrides/useDeleteCyberArkConfig.ts @@ -0,0 +1,19 @@ +import { deleteCyberArkConfig } from "./cyberArkApi"; +import { useMutation, useQueryClient } from "@tanstack/react-query"; +import { cyberArkKeys } from "./useCyberArkConfig"; + +export const useDeleteCyberArkConfig = (accessToken: string | null) => { + const queryClient = useQueryClient(); + + return useMutation({ + mutationFn: async () => { + if (!accessToken) { + throw new Error("Access token is required"); + } + return deleteCyberArkConfig(accessToken); + }, + onSuccess: () => { + queryClient.invalidateQueries({ queryKey: cyberArkKeys.all }); + }, + }); +}; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/configOverrides/useUpdateCyberArkConfig.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/configOverrides/useUpdateCyberArkConfig.ts new file mode 100644 index 00000000000..f5c88e833f7 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/configOverrides/useUpdateCyberArkConfig.ts @@ -0,0 +1,19 @@ +import { updateCyberArkConfig } from "./cyberArkApi"; +import { useMutation, useQueryClient } from "@tanstack/react-query"; +import { cyberArkKeys } from "./useCyberArkConfig"; + +export const useUpdateCyberArkConfig = (accessToken: string | null) => { + const queryClient = useQueryClient(); + + return useMutation({ + mutationFn: async (config: Record) => { + if (!accessToken) { + throw new Error("Access token is required"); + } + return updateCyberArkConfig(accessToken, config); + }, + onSuccess: () => { + queryClient.invalidateQueries({ queryKey: cyberArkKeys.all }); + }, + }); +}; diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/CyberArk/CyberArk.test.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/CyberArk/CyberArk.test.tsx new file mode 100644 index 00000000000..d9bd9b1ebc8 --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/CyberArk/CyberArk.test.tsx @@ -0,0 +1,78 @@ +import { screen } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { renderWithProviders } from "../../../../../tests/test-utils"; +import CyberArk from "./CyberArk"; + +const mockUseAuthorized = vi.hoisted(() => vi.fn()); +const mockUseCyberArkConfig = vi.hoisted(() => vi.fn()); + +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: mockUseAuthorized, +})); + +vi.mock("@/app/(dashboard)/hooks/configOverrides/useCyberArkConfig", () => ({ + useCyberArkConfig: mockUseCyberArkConfig, +})); + +vi.mock("@/app/(dashboard)/hooks/configOverrides/useDeleteCyberArkConfig", () => ({ + useDeleteCyberArkConfig: () => ({ mutate: vi.fn(), isPending: false }), +})); + +vi.mock("@/app/(dashboard)/hooks/configOverrides/useUpdateCyberArkConfig", () => ({ + useUpdateCyberArkConfig: () => ({ mutate: vi.fn(), isPending: false }), +})); + +vi.mock("./EditCyberArkModal", () => ({ + default: ({ isVisible }: { isVisible: boolean }) => (isVisible ?
Edit CyberArk Configuration
: null), +})); + +vi.mock("@/components/common_components/DeleteResourceModal", () => ({ + default: () => null, +})); + +describe("CyberArk", () => { + beforeEach(() => { + vi.clearAllMocks(); + mockUseAuthorized.mockReturnValue({ accessToken: "test-token" }); + const emptyConfigResult = { + data: { values: {} }, + isLoading: false, + isError: false, + error: null, + }; + mockUseCyberArkConfig.mockReturnValue(emptyConfigResult); + }); + + it("should render", () => { + renderWithProviders(); + + expect(screen.getByRole("heading", { name: "CyberArk Conjur" })).toBeInTheDocument(); + }); + + it("should open the configuration editor from the empty state", async () => { + const user = userEvent.setup(); + renderWithProviders(); + + await user.click(screen.getByRole("button", { name: /configure cyberark/i })); + + expect(screen.getByText("Edit CyberArk Configuration")).toBeInTheDocument(); + }); + + it("should display configured values and management actions", () => { + const configuredResult = { + data: { values: { cyberark_api_base: "https://conjur.example.com", cyberark_api_key: "secret" } }, + isLoading: false, + isError: false, + error: null, + }; + mockUseCyberArkConfig.mockReturnValue(configuredResult); + + renderWithProviders(); + + expect(screen.getByText("https://conjur.example.com")).toBeInTheDocument(); + expect(screen.getByText("Auth Method")).toBeInTheDocument(); + expect(screen.getAllByText("API Key")).toHaveLength(2); + expect(screen.getByRole("button", { name: /test connection/i })).toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/CyberArk/CyberArk.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/CyberArk/CyberArk.tsx new file mode 100644 index 00000000000..30026d0a834 --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/CyberArk/CyberArk.tsx @@ -0,0 +1,232 @@ +"use client"; + +import { Edit, ExternalLink, Info, KeyRound, PlugZap, Trash2 } from "lucide-react"; +import { useState } from "react"; + +import { testCyberArkConnection } from "@/app/(dashboard)/hooks/configOverrides/cyberArkApi"; +import { useCyberArkConfig } from "@/app/(dashboard)/hooks/configOverrides/useCyberArkConfig"; +import { useDeleteCyberArkConfig } from "@/app/(dashboard)/hooks/configOverrides/useDeleteCyberArkConfig"; +import { useUpdateCyberArkConfig } from "@/app/(dashboard)/hooks/configOverrides/useUpdateCyberArkConfig"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import DeleteResourceModal from "@/components/common_components/DeleteResourceModal"; +import { toast } from "@/lib/toast"; +import { Alert, AlertDescription, AlertTitle } from "@/components/shared/Alert"; +import { Button } from "@/components/ui/button"; +import { Card, CardAction, CardContent, CardDescription, CardHeader, CardTitle } from "@/components/ui/card"; +import { Skeleton } from "@/components/ui/skeleton"; + +import CyberArkEmptyPlaceholder from "./CyberArkEmptyPlaceholder"; +import EditCyberArkModal from "./EditCyberArkModal"; +import { FIELD_LABELS, SENSITIVE_FIELDS } from "./constants"; + +function detectAuthMethod(values: Record): string { + if (values.cyberark_api_key) return "API Key"; + if (values.client_cert && values.client_key) return "TLS Certificate"; + return "None"; +} + +function DetailRow({ children, label }: { children: React.ReactNode; label: string }) { + return ( +
+
{label}
+
{children}
+
+ ); +} + +export default function CyberArk() { + const { accessToken } = useAuthorized(); + const { data, isLoading, isError, error } = useCyberArkConfig(); + const { mutate: deleteConfig, isPending: isDeleting } = useDeleteCyberArkConfig(accessToken); + const { mutate: updateConfig, isPending: isClearingField } = useUpdateCyberArkConfig(accessToken); + const [isEditModalVisible, setIsEditModalVisible] = useState(false); + const [isDeleteModalOpen, setIsDeleteModalOpen] = useState(false); + const [clearingField, setClearingField] = useState(null); + const [isTesting, setIsTesting] = useState(false); + const rawValues = data?.values ?? {}; + const isConfigured = Boolean(rawValues.cyberark_api_base); + + const handleTestConnection = async () => { + if (!accessToken) return; + setIsTesting(true); + try { + const result = await testCyberArkConnection(accessToken); + toast.success(result.message || "Connection to CyberArk Conjur successful!"); + } catch (err) { + toast.fromError(err); + } finally { + setIsTesting(false); + } + }; + + const handleDelete = () => { + deleteConfig(undefined, { + onSuccess: () => { + toast.success("CyberArk configuration deleted"); + setIsDeleteModalOpen(false); + }, + onError: (err) => toast.fromError(err), + }); + }; + + const handleClearField = () => { + if (!clearingField) return; + updateConfig( + { [clearingField]: "" }, + { + onSuccess: () => { + toast.success(`${FIELD_LABELS[clearingField] ?? clearingField} cleared`); + setClearingField(null); + }, + onError: (err) => toast.fromError(err), + }, + ); + }; + + const renderValue = (key: string) => { + const value = rawValues[key]; + if (!value) return Not configured; + if (!SENSITIVE_FIELDS.has(key)) return {value}; + + return ( +
+ {value} + +
+ ); + }; + + const fieldsToShow = Object.entries(rawValues).filter(([, value]) => value != null && value !== ""); + + const renderCard = () => { + if (isLoading) { + return ( + + + + + + + ); + } + if (isError) { + return ( + + + + Could not load CyberArk configuration + {error instanceof Error && {error.message}} + + + + ); + } + return ( + + +
+ +
+ +

CyberArk Conjur

+ + Manage secret manager configuration +
+
+ {isConfigured && ( + + + + + + )} + + + {isConfigured && ( + + + Configuration changes are hot-reloaded across all proxy instances + + + View documentation + + + + + )} + + {isConfigured ? ( + fieldsToShow.length > 0 && ( +
+ {detectAuthMethod(rawValues)} + {fieldsToShow.map(([key]) => ( + + {renderValue(key)} + + ))} +
+ ) + ) : ( + setIsEditModalVisible(true)} /> + )} + + + ); + }; + + return ( + <> + {renderCard()} + + setIsEditModalVisible(false)} + onSuccess={() => setIsEditModalVisible(false)} + /> + setIsDeleteModalOpen(false)} + onOk={handleDelete} + confirmLoading={isDeleting} + /> + setClearingField(null)} + onOk={handleClearField} + confirmLoading={isClearingField} + /> + + ); +} diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/CyberArk/CyberArkEmptyPlaceholder.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/CyberArk/CyberArkEmptyPlaceholder.tsx new file mode 100644 index 00000000000..513d34d0f65 --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/CyberArk/CyberArkEmptyPlaceholder.tsx @@ -0,0 +1,24 @@ +import { KeyRound } from "lucide-react"; + +import { Button } from "@/components/ui/button"; + +interface CyberArkEmptyPlaceholderProps { + onAdd: () => void; +} + +export default function CyberArkEmptyPlaceholder({ onAdd }: CyberArkEmptyPlaceholderProps) { + return ( +
+
+ +
+

No CyberArk Configuration Found

+

+ Configure CyberArk Conjur to securely manage provider API keys and secrets for your LiteLLM deployment. +

+ +
+ ); +} diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/CyberArk/EditCyberArkModal.test.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/CyberArk/EditCyberArkModal.test.tsx new file mode 100644 index 00000000000..3f716451173 --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/CyberArk/EditCyberArkModal.test.tsx @@ -0,0 +1,176 @@ +import { fireEvent, screen, waitFor } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { beforeEach, describe, expect, it, vi } from "vitest"; + +import { renderWithProviders } from "../../../../../tests/test-utils"; +import EditCyberArkModal from "./EditCyberArkModal"; +import { useCyberArkConfig } from "@/app/(dashboard)/hooks/configOverrides/useCyberArkConfig"; +import { useUpdateCyberArkConfig } from "@/app/(dashboard)/hooks/configOverrides/useUpdateCyberArkConfig"; + +vi.mock("@/app/(dashboard)/hooks/configOverrides/useCyberArkConfig", () => ({ + useCyberArkConfig: vi.fn(), +})); + +vi.mock("@/app/(dashboard)/hooks/configOverrides/useUpdateCyberArkConfig", () => ({ + useUpdateCyberArkConfig: vi.fn(), +})); + +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: () => ({ accessToken: "sk-access-token" }), +})); + +vi.mock("@/lib/toast", () => ({ + toast: { success: vi.fn(), fromError: vi.fn() }, +})); + +const ALL_FIELDS = [ + "cyberark_api_base", + "cyberark_account", + "cyberark_username", + "cyberark_api_key", + "client_cert", + "client_key", + "ssl_verify", + "refresh_interval", +] as const; + +const propertiesFor = (fields: readonly string[]) => + Object.fromEntries(fields.map((name) => [name, { description: `${name} description` }])); + +const mutate = vi.fn(); + +const setup = (options?: { values?: Record; fields?: readonly string[] }) => { + vi.mocked(useCyberArkConfig).mockReturnValue({ + data: { + field_schema: { properties: propertiesFor(options?.fields ?? ALL_FIELDS) }, + values: options?.values ?? {}, + }, + } as unknown as ReturnType); + + vi.mocked(useUpdateCyberArkConfig).mockReturnValue({ + mutate, + isPending: false, + } as unknown as ReturnType); +}; + +const renderModal = (onSuccess = vi.fn(), onCancel = vi.fn()) => + renderWithProviders(); + +const save = async (user: ReturnType) => + user.click(screen.getByRole("button", { name: "Save" })); + +describe("EditCyberArkModal", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("clears untouched non-sensitive fields and omits untouched sensitive fields", async () => { + setup({ + values: { + cyberark_api_base: "https://conjur.example.com", + cyberark_account: "myorg", + cyberark_api_key: "super-secret-key", + client_key: "super-secret-pem", + }, + }); + const user = userEvent.setup(); + renderModal(); + + await save(user); + + await waitFor(() => { + expect(mutate).toHaveBeenCalledTimes(1); + }); + const expectedPayload = { + cyberark_api_base: "https://conjur.example.com", + cyberark_account: "myorg", + cyberark_username: "", + client_cert: "", + ssl_verify: "", + refresh_interval: "", + }; + expect(mutate.mock.calls[0][0]).toEqual(expectedPayload); + }); + + it("sends a sensitive field only once it is typed into", async () => { + setup({ values: { cyberark_api_base: "https://conjur.example.com", cyberark_api_key: "super-secret-key" } }); + const user = userEvent.setup(); + renderModal(); + + fireEvent.change(screen.getByLabelText("API Key"), { target: { value: "rotated-key" } }); + await save(user); + + await waitFor(() => { + expect(mutate).toHaveBeenCalledTimes(1); + }); + expect(mutate.mock.calls[0][0]).toMatchObject({ cyberark_api_key: "rotated-key" }); + }); + + it("never seeds a stored secret into its input", () => { + setup({ values: { cyberark_api_key: "super-secret-key", client_key: "super-secret-pem" } }); + renderModal(); + + expect(screen.getByLabelText("API Key")).toHaveValue(""); + expect(screen.getByLabelText("Client Key")).toHaveValue(""); + }); + + it("renders only the fields the schema declares, and sends only those", async () => { + setup({ + fields: ["cyberark_api_base", "cyberark_api_key"], + values: { cyberark_api_base: "https://conjur.example.com" }, + }); + const user = userEvent.setup(); + renderModal(); + + expect(screen.queryByLabelText("Account")).not.toBeInTheDocument(); + expect(screen.queryByLabelText("Client Key")).not.toBeInTheDocument(); + + await save(user); + + await waitFor(() => { + expect(mutate).toHaveBeenCalledTimes(1); + }); + expect(mutate.mock.calls[0][0]).toEqual({ cyberark_api_base: "https://conjur.example.com" }); + }); + + it("blocks the submit when the server url does not start with http", async () => { + setup({ values: {} }); + const user = userEvent.setup(); + renderModal(); + + fireEvent.change(screen.getByLabelText("Conjur Server URL"), { target: { value: "conjur.example.com" } }); + await save(user); + + expect(await screen.findByText("Must start with http:// or https://")).toBeInTheDocument(); + expect(mutate).not.toHaveBeenCalled(); + }); + + it("tells the admin a stored secret is kept when the field is left blank", () => { + setup({ values: { cyberark_api_key: "super-secret-key" } }); + renderModal(); + + expect(screen.getByLabelText("API Key")).toHaveAttribute( + "placeholder", + "Leave blank to keep existing (super-secret-key)", + ); + }); + + it("falls back to the schema description when no secret is stored yet", () => { + setup({ values: {} }); + renderModal(); + + expect(screen.getByLabelText("API Key")).toHaveAttribute("placeholder", "cyberark_api_key description"); + }); + + it("closes without saving when cancelled", async () => { + setup({ values: {} }); + const onCancel = vi.fn(); + const user = userEvent.setup(); + renderModal(vi.fn(), onCancel); + + await user.click(screen.getByRole("button", { name: "Cancel" })); + + expect(onCancel).toHaveBeenCalledTimes(1); + expect(mutate).not.toHaveBeenCalled(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/CyberArk/EditCyberArkModal.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/CyberArk/EditCyberArkModal.tsx new file mode 100644 index 00000000000..7a093e63f2b --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/CyberArk/EditCyberArkModal.tsx @@ -0,0 +1,176 @@ +"use client"; + +import { useCyberArkConfig } from "@/app/(dashboard)/hooks/configOverrides/useCyberArkConfig"; +import { useUpdateCyberArkConfig } from "@/app/(dashboard)/hooks/configOverrides/useUpdateCyberArkConfig"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import { toast } from "@/lib/toast"; +import React, { useMemo } from "react"; +import { z } from "zod/v4"; +import { FieldGroup } from "@/components/ui/field"; +import { FormField } from "@/components/shared/form/FormField"; +import { PasswordInput } from "@/components/shared/PasswordInput"; +import { Button } from "@/components/ui/button"; +import { Input } from "@/components/ui/input"; +import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner"; +import { Separator } from "@/components/ui/separator"; +import { useZodForm } from "@/lib/forms/useZodForm"; +import { SENSITIVE_FIELDS, FIELD_LABELS } from "./constants"; +import { Dialog, DialogContent, DialogFooter, DialogHeader, DialogTitle } from "@/components/ui/dialog"; + +interface CyberArkFieldGroup { + title: string; + subtitle?: string; + fields: string[]; +} + +const FIELD_GROUPS: CyberArkFieldGroup[] = [ + { + title: "Connection", + fields: ["cyberark_api_base", "cyberark_account", "cyberark_username"], + }, + { + title: "API Key Authentication", + subtitle: "Use a Conjur API key to authenticate. Only one auth method is required.", + fields: ["cyberark_api_key"], + }, + { + title: "Certificate Authentication", + subtitle: "Use a client TLS certificate and key to authenticate. Only one auth method is required.", + fields: ["client_cert", "client_key"], + }, + { + title: "Advanced", + subtitle: "Optional TLS and token caching settings.", + fields: ["ssl_verify", "refresh_interval"], + }, +]; + +type CyberArkFormValues = Record; + +const buildSchema = (fields: readonly string[]): z.ZodType => + z.object( + Object.fromEntries( + fields.map((name) => [ + name, + name === "cyberark_api_base" + ? z.string().refine((value) => value.length === 0 || /^https?:\/\/.+/.test(value), { + message: "Must start with http:// or https://", + }) + : z.string(), + ]), + ), + ) as unknown as z.ZodType; + +interface EditCyberArkModalProps { + isVisible: boolean; + onCancel: () => void; + onSuccess: () => void; +} + +const EditCyberArkModal: React.FC = ({ isVisible, onCancel, onSuccess }) => { + const { accessToken } = useAuthorized(); + const { data } = useCyberArkConfig(); + const { mutate, isPending } = useUpdateCyberArkConfig(accessToken); + + const properties: Record = useMemo( + () => data?.field_schema?.properties ?? {}, + [data], + ); + const rawValues: Record = useMemo(() => data?.values ?? {}, [data]); + + const visibleFields = useMemo( + () => FIELD_GROUPS.flatMap((group) => group.fields).filter((name) => properties[name] !== undefined), + [properties], + ); + + const seededValues = useMemo( + () => + Object.fromEntries( + visibleFields.map((name) => [name, SENSITIVE_FIELDS.has(name) ? "" : ((rawValues[name] ?? "") as string)]), + ), + [visibleFields, rawValues], + ); + + const schema = useMemo(() => buildSchema(visibleFields), [visibleFields]); + const form = useZodForm(schema, { values: seededValues }); + + const handleSubmit = (formValues: CyberArkFormValues) => { + const config: Record = Object.fromEntries( + Object.entries(formValues).flatMap(([key, value]) => { + if (value !== undefined && value !== null && value !== "") return [[key, value]]; + if (!SENSITIVE_FIELDS.has(key)) return [[key, ""]]; + return []; + }), + ); + + mutate(config, { + onSuccess: () => { + toast.success("CyberArk configuration updated successfully"); + onSuccess(); + }, + onError: (err) => { + toast.fromError(err); + }, + }); + }; + + const handleCancel = () => { + form.reset(seededValues); + onCancel(); + }; + + const renderField = (fieldName: string) => { + const fieldSchema = properties[fieldName]; + if (!fieldSchema) return null; + + const isSensitive = SENSITIVE_FIELDS.has(fieldName); + const existingValue = rawValues[fieldName]; + const hasExistingValue = isSensitive && existingValue != null && existingValue !== ""; + const placeholder = hasExistingValue ? `Leave blank to keep existing (${existingValue})` : fieldSchema?.description; + + return ( + + {({ ref, ...field }) => + isSensitive ? ( + + ) : ( + + ) + } + + ); + }; + + return ( + !open && handleCancel()}> + + + Edit CyberArk Configuration + +
+ {FIELD_GROUPS.map((group, index) => ( +
+ {index > 0 && } +
{group.title}
+ {group.subtitle &&

{group.subtitle}

} + {group.fields.map(renderField)} +
+ ))} +
+ +
+ + +
+ + + + ); +}; + +export default EditCyberArkModal; diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/CyberArk/constants.ts b/ui/litellm-dashboard/src/components/Settings/AdminSettings/CyberArk/constants.ts new file mode 100644 index 00000000000..5835f93d19a --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/CyberArk/constants.ts @@ -0,0 +1,12 @@ +export const SENSITIVE_FIELDS = new Set(["cyberark_api_key", "client_key"]); + +export const FIELD_LABELS: Record = { + cyberark_api_base: "Conjur Server URL", + cyberark_account: "Account", + cyberark_username: "Username", + cyberark_api_key: "API Key", + client_cert: "Client Certificate", + client_key: "Client Key", + ssl_verify: "SSL Verification", + refresh_interval: "Token Refresh Interval (seconds)", +}; diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 525642702ca..6c5319f0361 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -2858,6 +2858,59 @@ export interface paths { patch?: never; trace?: never; }; + "/config_overrides/cyberark": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Get Cyberark Config + * @description Get current CyberArk Conjur configuration. + * Returns decrypted values from DB, or falls back to current env vars. + * Sensitive fields are masked before leaving the server. + */ + get: operations["get_cyberark_config_config_overrides_cyberark_get"]; + put?: never; + /** + * Update Cyberark Config + * @description Update CyberArk Conjur secret manager configuration. + * Sets environment variables, encrypts sensitive fields, and stores in DB. + * Reinitializes the secret manager on this pod. + */ + post: operations["update_cyberark_config_config_overrides_cyberark_post"]; + /** + * Delete Cyberark Config + * @description Delete CyberArk Conjur configuration. Idempotent. + */ + delete: operations["delete_cyberark_config_config_overrides_cyberark_delete"]; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + "/config_overrides/cyberark/test_connection": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + get?: never; + put?: never; + /** + * Test Cyberark Connection + * @description Test the connection to the currently configured CyberArk Conjur server. + * Uses the already-initialized secret manager client. Does not modify any state. + */ + post: operations["test_cyberark_connection_config_overrides_cyberark_test_connection_post"]; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/config_overrides/hashicorp_vault": { parameters: { query?: never; @@ -25841,6 +25894,52 @@ export interface components { /** User Id */ user_id: string; }; + /** + * CyberArkConfig + * @description Configuration for CyberArk Conjur secret manager integration. + */ + CyberArkConfig: { + /** + * Client Cert + * @description Path to the client TLS certificate for certificate-based authentication + */ + client_cert?: string | null; + /** + * Client Key + * @description Path to the client TLS private key for certificate-based authentication + */ + client_key?: string | null; + /** + * Cyberark Account + * @description The Conjur organization account name + */ + cyberark_account?: string | null; + /** + * Cyberark Api Base + * @description The address of the CyberArk Conjur server (e.g., https://conjur.example.com) + */ + cyberark_api_base?: string | null; + /** + * Cyberark Api Key + * @description API key for Conjur API-key authentication + */ + cyberark_api_key?: string | null; + /** + * Cyberark Username + * @description The Conjur username (login) to authenticate as + */ + cyberark_username?: string | null; + /** + * Refresh Interval + * @description Auth token cache TTL in seconds (default: 300) + */ + refresh_interval?: string | null; + /** + * Ssl Verify + * @description Set to false to disable SSL verification (e.g., for self-signed certificates) + */ + ssl_verify?: string | null; + }; /** DailySpendData */ DailySpendData: { breakdown?: components["schemas"]["BreakdownMetrics"]; @@ -42786,6 +42885,120 @@ export interface operations { }; }; }; + get_cyberark_config_config_overrides_cyberark_get: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["ConfigOverrideSettingsResponse"]; + }; + }; + }; + }; + update_cyberark_config_config_overrides_cyberark_post: { + parameters: { + query?: never; + header?: { + /** @description The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability */ + "litellm-changed-by"?: string | null; + }; + path?: never; + cookie?: never; + }; + requestBody: { + content: { + "application/json": components["schemas"]["CyberArkConfig"]; + }; + }; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": { + [key: string]: string; + }; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + delete_cyberark_config_config_overrides_cyberark_delete: { + parameters: { + query?: never; + header?: { + /** @description The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability */ + "litellm-changed-by"?: string | null; + }; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": { + [key: string]: string; + }; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; + test_cyberark_connection_config_overrides_cyberark_test_connection_post: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": { + [key: string]: string; + }; + }; + }; + }; + }; get_hashicorp_vault_config_config_overrides_hashicorp_vault_get: { parameters: { query?: never; From 194a3cc202dd239365425cd45a4d8790ee6ee0f4 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 29 Aug 2026 13:58:04 -0700 Subject: [PATCH 60/68] ci: build the benchmark environment outside the CodSpeed runner (#38426) Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/workflows/codspeed.yml | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/.github/workflows/codspeed.yml b/.github/workflows/codspeed.yml index a69e50b5753..7e013b7bb0b 100644 --- a/.github/workflows/codspeed.yml +++ b/.github/workflows/codspeed.yml @@ -12,6 +12,7 @@ on: - "uv.lock" - ".github/workflows/codspeed.yml" - ".github/actions/setup-uv-with-retries/**" + - ".github/actions/cache-cargo-build/**" pull_request: branches: - main @@ -23,6 +24,7 @@ on: - "uv.lock" - ".github/workflows/codspeed.yml" - ".github/actions/setup-uv-with-retries/**" + - ".github/actions/cache-cargo-build/**" # Allow CodSpeed to trigger backtest performance analysis # in order to generate initial data workflow_dispatch: @@ -55,6 +57,26 @@ jobs: with: version: "0.10.9" + - name: Cache the Rust build + uses: ./.github/actions/cache-cargo-build + + # Build the wheel and resolve every dependency outside the CodSpeed + # runner: the same maturin build took 42 minutes inside `codspeed run` + # versus under 3 minutes as a plain step (LIT-6183) + - name: Build environment + run: > + env PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 + uv run --frozen --no-default-groups + --with pytest==8.3.5 + --with pytest-codspeed==4.3.0 + --with "mcp>=1.26.0,<2.0" + --with "a2a-sdk>=1.1.0,<2.0" + pytest + -p pytest_codspeed.plugin + tests/benchmarks/ + --codspeed + --collect-only -q + - name: Run benchmarks uses: CodSpeedHQ/action@1c8ae4843586d3ba879736b7f6b7b0c990757fab # v4.12.1 with: From 128d7e52784b4a8fee4b171e7528ff6841cd2a08 Mon Sep 17 00:00:00 2001 From: Mateo Wang Date: Sat, 29 Aug 2026 14:09:18 -0700 Subject: [PATCH 61/68] refactor(batches): make count_error_file_failed_requests public for the poller import --- .../proxy/common_utils/check_batch_cost.py | 4 ++-- litellm/batches/batch_utils.py | 6 +++--- 2 files changed, 5 insertions(+), 5 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 d34635fa253..3b09dc9272e 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -627,7 +627,7 @@ class CheckBatchCost: later poll. """ from litellm.batches.batch_utils import ( - _count_error_file_failed_requests, + count_error_file_failed_requests, _get_file_content_as_dictionary, calculate_batch_cost_and_usage, ) @@ -772,7 +772,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_failed_requests: Final = await count_error_file_failed_requests( response, custom_llm_provider=batch_file_provider, litellm_params={ diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index ea2a95c1717..3831f57a10d 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -89,13 +89,13 @@ async def _handle_completed_batch( 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( + failed_requests=await count_error_file_failed_requests( batch, custom_llm_provider=custom_llm_provider, litellm_params=litellm_params ), ) 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( + error_file_failed_requests: Final = await count_error_file_failed_requests( batch, custom_llm_provider=custom_llm_provider, litellm_params=litellm_params ) @@ -450,7 +450,7 @@ async def _fetch_batch_output_file_content( ) -async def _count_error_file_failed_requests( +async def count_error_file_failed_requests( batch: Batch, custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"], litellm_params: dict | None, From 2963b47cda666efc363cd497463e827cfe4d4706 Mon Sep 17 00:00:00 2001 From: Mateo Wang Date: Sat, 29 Aug 2026 14:09:33 -0700 Subject: [PATCH 62/68] test: patch the Logging handler instead of the class in the poller error-file test --- tests/proxy_unit_tests/test_check_batch_cost.py | 17 ++++++++--------- 1 file changed, 8 insertions(+), 9 deletions(-) diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index 757762eac87..ff5e8f89d64 100644 --- a/tests/proxy_unit_tests/test_check_batch_cost.py +++ b/tests/proxy_unit_tests/test_check_batch_cost.py @@ -1388,6 +1388,8 @@ class TestCheckBatchCost: import httpx import respx + from litellm.litellm_core_utils.litellm_logging import Logging + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=1) mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) @@ -1456,9 +1458,9 @@ class TestCheckBatchCost: with ( respx.mock(assert_all_called=True) as provider, - patch( # test-quality-ok: the poller builds Logging inline, the only seam to its handler kwargs - "litellm.litellm_core_utils.litellm_logging.Logging" - ) as mock_logging_cls, + patch.object( # test-quality-ok: the poller builds Logging inline, the only seam to its handler kwargs + Logging, "async_success_handler", new_callable=AsyncMock + ) as success_handler, ): 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()) @@ -1466,14 +1468,11 @@ class TestCheckBatchCost: 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()) ) - mock_logging_obj = MagicMock() - mock_logging_obj.async_success_handler = AsyncMock() - mock_logging_cls.return_value = mock_logging_obj - await check_batch_cost_instance.check_batch_cost() - mock_logging_obj.async_success_handler.assert_awaited_once() - handler_kwargs = mock_logging_obj.async_success_handler.await_args.kwargs + spend_log_calls = [call.kwargs for call in success_handler.await_args_list if "batch_cost" in call.kwargs] + assert len(spend_log_calls) == 1 + handler_kwargs = spend_log_calls[0] assert handler_kwargs["batch_successful_requests"] == 1 assert handler_kwargs["batch_failed_requests"] == 3, ( "2 error-file lines must add to the output file's 1 rejected request" From 5c034fda749391880e3b431be36a0d193faadc56 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Sat, 29 Aug 2026 14:43:20 -0700 Subject: [PATCH 63/68] fix(ui): allow in-place editing of classifier numeric inputs (#38803) Backspacing the last digit of Context Window Size instantly refilled the default (3), since onChange mapped empty input to null and the handler coalesced null back to the default. The same defect affected Timeout (ms) and Context Character Budget. Add per-field raw draft state so an empty or partial value stays visible while focused, commit only finite values (rounded, clamped to each field's minimum), and clear the draft on blur so an abandoned edit falls back to the committed value. 0 stays a valid committed value for both context controls. Add stable ids and label associations; update tests to query by label --- .../add_model/ClassificationMethodConfig.tsx | 104 ++++++++++++++---- .../add_model/ComplexityRouterConfig.test.tsx | 50 ++++++--- .../edit_auto_router_modal.test.tsx | 3 +- 3 files changed, 114 insertions(+), 43 deletions(-) diff --git a/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx index c48d15adecb..2947e29319b 100644 --- a/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx @@ -40,6 +40,10 @@ const DEFAULT_SCORING_EXPLANATION = "The router scores each request across 7 dimensions: token count, code presence, reasoning markers, technical " + "terms, simple indicators, multi-step patterns, and question complexity. The weighted score determines the tier:"; +const CLASSIFIER_TIMEOUT_ID = "classifier-timeout-ms"; +const CLASSIFIER_CONTEXT_WINDOW_SIZE_ID = "classifier-context-window-size"; +const CLASSIFIER_CONTEXT_BUDGET_CHARS_ID = "classifier-context-budget-chars"; + const CUSTOM_PROMPT_WITH_HEURISTIC_FALLBACK = "This router classifies with your own prompt, so the tier comes from whatever rubric it states. The four tier " + "names stay fixed. The scoring below is the heuristic, which now runs only when the classifier call fails:"; @@ -204,6 +208,7 @@ const ClassificationMethodConfig: React.FC = ({ showValidationErrors = false, defaultModel, }) => { + const [draft, setDraft] = React.useState<{ id: string; raw: string } | null>(null); const hasDefaultModel = Boolean(defaultModel); const classifierType = effectiveClassifierType(value); const classifierModelMissing = @@ -261,13 +266,13 @@ const ClassificationMethodConfig: React.FC = ({ }); }; - const handleClassifierTimeoutChange = (timeoutMs: number | null) => { + const handleClassifierTimeoutChange = (timeoutMs: number) => { onChange({ ...value, classifier_llm_config: { ...value.classifier_llm_config, model: value.classifier_llm_config?.model ?? "", - timeout_ms: timeoutMs ?? DEFAULT_CLASSIFIER_TIMEOUT_MS, + timeout_ms: timeoutMs, }, }); }; @@ -300,20 +305,32 @@ const ClassificationMethodConfig: React.FC = ({ onChange({ ...value, classifier_fallback: fallback }); }; - const handleClassifierContextWindowSizeChange = (windowSize: number | null) => { + const handleClassifierContextWindowSizeChange = (windowSize: number) => { onChange({ ...value, - classifier_context_window_size: windowSize ?? DEFAULT_CLASSIFIER_CONTEXT_WINDOW_SIZE, + classifier_context_window_size: windowSize, }); }; - const handleClassifierContextBudgetCharsChange = (budgetChars: number | null) => { + const handleClassifierContextBudgetCharsChange = (budgetChars: number) => { onChange({ ...value, - classifier_context_budget_chars: budgetChars ?? DEFAULT_CLASSIFIER_CONTEXT_BUDGET_CHARS, + classifier_context_budget_chars: budgetChars, }); }; + const handleClassifierIntegerChange = ( + id: string, + raw: string, + minimum: number, + onCommit: (value: number) => void, + ) => { + setDraft({ id, raw }); + const parsed = Number(raw); + if (raw.trim() === "" || !Number.isFinite(parsed)) return; + onCommit(Math.max(minimum, Math.round(parsed))); + }; + const handleClassifierContextIncludeAssistantTurnsChange = (includeAssistantTurns: boolean) => { onChange({ ...value, @@ -366,14 +383,27 @@ const ClassificationMethodConfig: React.FC = ({ {classifierModelMissing && A classifier model is required}
- Timeout (ms) + - handleClassifierTimeoutChange(event.target.value === "" ? null : event.target.valueAsNumber) + id={CLASSIFIER_TIMEOUT_ID} + type="text" + inputMode="numeric" + value={ + draft?.id === CLASSIFIER_TIMEOUT_ID + ? draft.raw + : String(value.classifier_llm_config?.timeout_ms ?? DEFAULT_CLASSIFIER_TIMEOUT_MS) } - min={1} + onChange={(event) => + handleClassifierIntegerChange( + CLASSIFIER_TIMEOUT_ID, + event.target.value, + 1, + handleClassifierTimeoutChange, + ) + } + onBlur={() => setDraft(null)} className="w-full" /> @@ -480,14 +510,27 @@ const ClassificationMethodConfig: React.FC = ({
- Context Window Size + - handleClassifierContextWindowSizeChange(event.target.value === "" ? null : event.target.valueAsNumber) + id={CLASSIFIER_CONTEXT_WINDOW_SIZE_ID} + type="text" + inputMode="numeric" + value={ + draft?.id === CLASSIFIER_CONTEXT_WINDOW_SIZE_ID + ? draft.raw + : String(value.classifier_context_window_size ?? DEFAULT_CLASSIFIER_CONTEXT_WINDOW_SIZE) } - min={0} + onChange={(event) => + handleClassifierIntegerChange( + CLASSIFIER_CONTEXT_WINDOW_SIZE_ID, + event.target.value, + 0, + handleClassifierContextWindowSizeChange, + ) + } + onBlur={() => setDraft(null)} className="w-full" /> @@ -497,14 +540,27 @@ const ClassificationMethodConfig: React.FC = ({
- Context Character Budget + - handleClassifierContextBudgetCharsChange(event.target.value === "" ? null : event.target.valueAsNumber) + id={CLASSIFIER_CONTEXT_BUDGET_CHARS_ID} + type="text" + inputMode="numeric" + value={ + draft?.id === CLASSIFIER_CONTEXT_BUDGET_CHARS_ID + ? draft.raw + : String(value.classifier_context_budget_chars ?? DEFAULT_CLASSIFIER_CONTEXT_BUDGET_CHARS) } - min={0} + onChange={(event) => + handleClassifierIntegerChange( + CLASSIFIER_CONTEXT_BUDGET_CHARS_ID, + event.target.value, + 0, + handleClassifierContextBudgetCharsChange, + ) + } + onBlur={() => setDraft(null)} className="w-full" /> diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx index 91201b51663..33ce1169c46 100644 --- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx @@ -131,10 +131,8 @@ describe("ComplexityRouterConfig", () => { fireEvent.click(screen.getByText("Advanced: Classification Method")); expect(screen.getByText("Classifier Model")).toBeInTheDocument(); - expect(screen.getByText("Timeout (ms)")).toBeInTheDocument(); - expect(screen.getByDisplayValue("750")).toBeInTheDocument(); - expect(screen.getByText("Context Window Size")).toBeInTheDocument(); - expect(screen.getByDisplayValue("5")).toBeInTheDocument(); + expect(screen.getByLabelText("Timeout (ms)")).toHaveValue("750"); + expect(screen.getByLabelText("Context Window Size")).toHaveValue("5"); expect(screen.queryByText("Context Per-Turn Character Limit")).not.toBeInTheDocument(); }); @@ -148,11 +146,8 @@ describe("ComplexityRouterConfig", () => { fireEvent.click(screen.getByText("Advanced: Classification Method")); - const windowSizeSection = screen.getByText("Context Window Size").closest("div") as HTMLElement; - expect(within(windowSizeSection).getByDisplayValue("3")).toBeInTheDocument(); - - const budgetSection = screen.getByText("Context Character Budget").closest("div") as HTMLElement; - expect(within(budgetSection).getByDisplayValue("8000")).toBeInTheDocument(); + expect(screen.getByLabelText("Context Window Size")).toHaveValue("3"); + expect(screen.getByLabelText("Context Character Budget")).toHaveValue("8000"); }); it("should warn when the budget is too small to quote any turn that does not already fit", () => { @@ -247,7 +242,11 @@ describe("ComplexityRouterConfig", () => { expect(screen.queryByText("Context Per-Turn Character Limit")).not.toBeInTheDocument(); }); - it("should call onChange with the updated classifier_context_window_size when edited", () => { + it.each([ + ["Timeout (ms)", "7", { classifier_llm_config: { model: "gpt-3.5-turbo", timeout_ms: 7 } }], + ["Context Window Size", "0", { classifier_context_window_size: 0 }], + ["Context Character Budget", "7", { classifier_context_budget_chars: 7 }], + ])("keeps %s empty while it is being edited, then commits %s", (label, replacement, expected) => { const onChange = vi.fn(); const llmValue: ComplexityRouterConfigValue = { ...defaultValue, @@ -257,14 +256,31 @@ describe("ComplexityRouterConfig", () => { renderWithProviders(); fireEvent.click(screen.getByText("Advanced: Classification Method")); - const windowSizeSection = screen.getByText("Context Window Size").closest("div") as HTMLElement; - const input = within(windowSizeSection).getByRole("spinbutton"); - fireEvent.change(input, { target: { value: "7" } }); + const input = screen.getByLabelText(label); + fireEvent.change(input, { target: { value: "" } }); - expect(onChange).toHaveBeenCalledWith({ - ...llmValue, - classifier_context_window_size: 7, - }); + expect(input).toHaveValue(""); + expect(onChange).not.toHaveBeenCalled(); + + fireEvent.change(input, { target: { value: replacement } }); + + expect(onChange).toHaveBeenLastCalledWith({ ...llmValue, ...expected }); + }); + + it("restores the committed context window size after an empty field loses focus", () => { + const llmValue: ComplexityRouterConfigValue = { + ...defaultValue, + classifier_type: "llm", + classifier_llm_config: { model: "gpt-3.5-turbo", timeout_ms: 3000 }, + }; + renderWithProviders(); + fireEvent.click(screen.getByText("Advanced: Classification Method")); + + const input = screen.getByLabelText("Context Window Size"); + fireEvent.change(input, { target: { value: "" } }); + fireEvent.blur(input); + + expect(input).toHaveValue("3"); }); it("should render the custom technical keywords field", () => { diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx index 022dd0b7cad..a4921fcfcb5 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx +++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx @@ -295,8 +295,7 @@ describe("EditAutoRouterModal classifier context window", () => { renderLlmModal(); await user.click(await screen.findByText("Advanced: Classification Method")); - const windowSizeSection = (await screen.findByText("Context Window Size")).closest("div") as HTMLElement; - const input = within(windowSizeSection).getByRole("spinbutton"); + const input = await screen.findByLabelText("Context Window Size"); fireEvent.change(input, { target: { value: "8" } }); await user.click(screen.getByRole("button", { name: /save changes/i })); From 36ea28b0922345b2a59f3f809ea75c70819181ad Mon Sep 17 00:00:00 2001 From: tin-berri Date: Sat, 29 Aug 2026 15:04:22 -0700 Subject: [PATCH 64/68] fix(anthropic): emit signature-only thinking blocks on the /v1/messages bridge (#38809) --- litellm/llms/anthropic/common_utils.py | 19 ++++++ .../adapters/streaming_iterator.py | 11 ++-- .../adapters/transformation.py | 4 +- ...al_pass_through_adapters_transformation.py | 23 ++++--- .../test_streaming_iterator_first_delta.py | 60 ++++++++++++++----- .../anthropic/test_anthropic_common_utils.py | 17 ++++++ 6 files changed, 104 insertions(+), 30 deletions(-) diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index 681a8397f66..d8d6a7fc9f8 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -1099,6 +1099,25 @@ def is_empty_thinking_block(block: object) -> bool: return not isinstance(thinking, str) or not thinking.strip() +def is_empty_unsigned_thinking_block(block: object) -> bool: + """ + True for an empty ``{"type": "thinking"}`` block carrying no signature. + + The emit-side predicate: response paths drop a thinking block only when it + holds nothing the client could need. A signature-only block is a real + provider response (Bedrock Converse under adaptive thinking emits a + reasoning block with empty text and only a signature) and the client needs + the signature to replay reasoning across tool-use turns, so it must be + emitted. Request paths keep using :func:`is_empty_thinking_block`: + Anthropic rejects empty thinking blocks in request history regardless of + signature, and the inbound strip self-heals a replayed signature-only + block. + """ + if not isinstance(block, dict) or not is_empty_thinking_block(block): + return False + return not block.get("signature") + + def normalize_anthropic_tool_use_id(raw_id: str) -> str: """ Normalize a tool_use / tool_result id for Anthropic's ``^[a-zA-Z0-9_-]+$`` diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py index cefd4aa2d77..cc5879df56d 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/streaming_iterator.py @@ -1029,7 +1029,7 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): @staticmethod def _is_blank_delta(chunk: "ModelResponseStream") -> bool: - from litellm.llms.anthropic.common_utils import is_empty_thinking_block + from litellm.llms.anthropic.common_utils import is_empty_unsigned_thinking_block choice: Final = chunk.choices[0] if choice.finish_reason is not None: @@ -1041,11 +1041,14 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): return False if getattr(delta, "reasoning_content", None): return False - # thinking_blocks whose entries are all empty (even if signed) must not + # thinking_blocks whose entries are all empty AND unsigned must not # open a block: the emitted {"type": "thinking", "thinking": ""} gets - # replayed as history and Anthropic rejects it (LIT-6357). + # replayed as history and Anthropic rejects it (LIT-6357). A signed + # entry opens the block so the client receives the replay signature. thinking_blocks: Final = getattr(delta, "thinking_blocks", None) - if thinking_blocks and any(isinstance(b, dict) and not is_empty_thinking_block(b) for b in thinking_blocks): + if thinking_blocks and any( + isinstance(b, dict) and not is_empty_unsigned_thinking_block(b) for b in thinking_blocks + ): return False return True diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py index a9fa00c827a..411df267442 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py @@ -90,7 +90,7 @@ from litellm.litellm_core_utils.reasoning_effort_utils import ( reasoning_effort_from_thinking_budget, ) from litellm.llms.anthropic.common_utils import ( - is_empty_thinking_block, + is_empty_unsigned_thinking_block, normalize_anthropic_tool_use_id, ) from litellm.llms.anthropic.experimental_pass_through.context_management import ( @@ -1267,7 +1267,7 @@ class LiteLLMAnthropicMessagesAdapter: if hasattr(choice.message, "thinking_blocks") and choice.message.thinking_blocks: for thinking_block in choice.message.thinking_blocks: if thinking_block.get("type") == "thinking": - if is_empty_thinking_block(thinking_block): + if is_empty_unsigned_thinking_block(thinking_block): continue thinking_value = thinking_block.get("thinking", "") signature_value = thinking_block.get("signature", "") diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py index ea1813acb82..2d74c00071b 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py @@ -1013,12 +1013,15 @@ def test_translate_openai_content_to_anthropic_thinking_and_redacted_thinking(): assert result[1]["data"] == "REDACTED" -def test_translate_openai_content_to_anthropic_drops_empty_thinking_blocks(): - """LIT-6357 non-streaming producer half: a bridged reasoning model whose - thinking_blocks entry has empty or whitespace-only text (signed or not) - must not surface as {"type": "thinking", "thinking": ""} — clients replay - it as history and Anthropic 400s with "each thinking block must contain - thinking". Non-empty thinking and redacted_thinking pass through.""" +def test_translate_openai_content_to_anthropic_drops_empty_unsigned_thinking_blocks(): + """LIT-6357 non-streaming producer half, narrowed to unsigned blocks: a + bridged reasoning model whose thinking_blocks entry has empty or + whitespace-only text and no signature must not surface as + {"type": "thinking", "thinking": ""}. A signature-only block (Bedrock + Converse adaptive thinking) must be emitted so the client keeps the + signature for tool-use replay; the inbound strip self-heals it if the + client loops it back. Non-empty thinking and redacted_thinking pass + through.""" openai_choices = [ Choices( message=Message( @@ -1037,9 +1040,11 @@ def test_translate_openai_content_to_anthropic_drops_empty_thinking_blocks(): adapter = LiteLLMAnthropicMessagesAdapter() result = adapter._translate_openai_content_to_anthropic(choices=openai_choices) - assert [b["type"] for b in result] == ["thinking", "redacted_thinking", "text"] - assert result[0]["thinking"] == "real plan" - assert result[1]["data"] == "REDACTED" + assert [b["type"] for b in result] == ["thinking", "thinking", "redacted_thinking", "text"] + assert result[0]["thinking"] == "" + assert result[0]["signature"] == "sig_abc" + assert result[1]["thinking"] == "real plan" + assert result[2]["data"] == "REDACTED" def test_translate_streaming_openai_chunk_to_anthropic_thinking_delta(): diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_first_delta.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_first_delta.py index 6268cd01efe..17d42f55ae0 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_first_delta.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_streaming_iterator_first_delta.py @@ -1048,19 +1048,20 @@ def _empty_thinking_then_tool_chunks(thinking: str = "", signature: str = "") -> @pytest.mark.parametrize("is_async", [False, True]) @pytest.mark.parametrize( "thinking,signature", - [("", ""), (" \n\t ", ""), ("", "sig_abc")], - ids=["empty", "whitespace-only", "empty-but-signed"], + [("", ""), (" \n\t ", "")], + ids=["empty", "whitespace-only"], ) @pytest.mark.asyncio async def test_contentless_thinking_chunk_opens_no_thinking_block(is_async: bool, thinking: str, signature: str): """LIT-6357 producer half: a reasoning model that goes straight to tool - calls streams a ``thinking_blocks`` entry with no real thinking text; the - wrapper used to open ``{"type": "thinking", "thinking": ""}`` for it and - close the block with no delta. Clients (Claude Code) replay that block as - history and Anthropic rejects the next tool-loop request with - "each thinking block must contain thinking" — empty-but-signed included. - The contentless chunk must open nothing; the tool_use block must be - unaffected.""" + calls streams a ``thinking_blocks`` entry with no real thinking text and + no signature; the wrapper used to open ``{"type": "thinking", + "thinking": ""}`` for it and close the block with no delta. Clients + (Claude Code) replay that block as history and Anthropic rejects the next + tool-loop request with "each thinking block must contain thinking". + The contentless unsigned chunk must open nothing; the tool_use block must + be unaffected. A SIGNED contentless chunk is different: see + test_signature_only_thinking_chunk_opens_signed_block.""" chunks = _empty_thinking_then_tool_chunks(thinking, signature) if is_async: wrapper = AnthropicStreamWrapper(completion_stream=_AsyncStream(chunks), model="claude-x") @@ -1138,11 +1139,38 @@ async def test_early_signature_on_blank_thinking_chunk_is_carried_to_the_opened_ @pytest.mark.parametrize("is_async", [False, True]) @pytest.mark.asyncio -async def test_early_signature_discarded_when_first_block_is_not_thinking(is_async: bool): - """An early signature from a skipped blank thinking chunk must not leak - into a text or tool_use first block, and must not resurrect an empty - thinking block on its own (an empty-but-signed block is exactly what - Anthropic rejects).""" +async def test_signature_only_thinking_chunk_opens_signed_block(is_async: bool): + """Bedrock Converse under adaptive thinking emits a reasoning delta with + empty text and only a signature. The signed chunk must open a thinking + block that carries the signature to the client (needed to replay reasoning + across tool-use turns); the tool_use block must be unaffected. Dropping it + like the unsigned case regressed the claude_code thinking e2e cells.""" + chunks = _empty_thinking_then_tool_chunks("", "sig_bedrock") + if is_async: + wrapper = AnthropicStreamWrapper(completion_stream=_AsyncStream(chunks), model="claude-x") + events = await _drain_async(wrapper) + else: + wrapper = AnthropicStreamWrapper(completion_stream=iter(chunks), model="claude-x") + events = _drain_sync(wrapper) + + starts = _thinking_block_starts(events) + assert len(starts) == 1 + assert starts[0].get("signature") == "sig_bedrock" or _signature_deltas(events) == ["sig_bedrock"] + assert _thinking_deltas(events) == [] + tool_starts = [ + e["content_block"] + for e in events + if e.get("type") == "content_block_start" and e["content_block"].get("type") == "tool_use" + ] + assert [b["name"] for b in tool_starts] == ["get_weather"] + _assert_deltas_match_their_block_type(events) + + +@pytest.mark.parametrize("is_async", [False, True]) +@pytest.mark.asyncio +async def test_signature_only_thinking_chunk_before_text_leaks_no_signature(is_async: bool): + """The signed thinking block a signature-only chunk opens must stay its + own block: the text block that follows carries no signature.""" chunks = [ _thinking_chunk("", signature="sig_early"), _make_chunk(Delta(content="Hello")), @@ -1155,7 +1183,9 @@ async def test_early_signature_discarded_when_first_block_is_not_thinking(is_asy wrapper = AnthropicStreamWrapper(completion_stream=iter(chunks), model="claude-x") events = _drain_sync(wrapper) - assert _thinking_block_starts(events) == [] + starts = _thinking_block_starts(events) + assert len(starts) == 1 + assert starts[0].get("signature") == "sig_early" or _signature_deltas(events) == ["sig_early"] text_starts = [ e["content_block"] for e in events diff --git a/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py b/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py index a2da2cccb7c..794613942a1 100644 --- a/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py +++ b/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py @@ -1364,6 +1364,23 @@ class TestAnthropicThinkingSignatureSelfHeal: assert is_empty_thinking_block({"type": "text", "text": ""}) is False assert is_empty_thinking_block("not a dict") is False + def test_is_empty_unsigned_thinking_block(self): + """Emit-side predicate: a signature-only block must be kept (Bedrock + Converse adaptive thinking emits empty text with only a signature, and + the client needs it to replay reasoning in tool-use turns); only an + empty block with nothing to preserve is droppable.""" + from litellm.llms.anthropic.common_utils import is_empty_unsigned_thinking_block + + assert is_empty_unsigned_thinking_block({"type": "thinking", "thinking": ""}) is True + assert is_empty_unsigned_thinking_block({"type": "thinking", "thinking": " \n\t "}) is True + assert is_empty_unsigned_thinking_block({"type": "thinking"}) is True + assert is_empty_unsigned_thinking_block({"type": "thinking", "thinking": "", "signature": ""}) is True + assert is_empty_unsigned_thinking_block({"type": "thinking", "thinking": "", "signature": "sig_abc"}) is False + assert is_empty_unsigned_thinking_block({"type": "thinking", "thinking": " ", "signature": "sig_abc"}) is False + assert is_empty_unsigned_thinking_block({"type": "thinking", "thinking": "plan"}) is False + assert is_empty_unsigned_thinking_block({"type": "redacted_thinking", "data": "opaque"}) is False + assert is_empty_unsigned_thinking_block("not a dict") is False + def test_strip_empty_content_blocks_drops_empty_thinking_blocks(self): """LIT-6357 ingestion half: an assistant tool-loop turn carrying an empty (even signed) thinking block keeps its tool_use blocks and loses From 02046fb7b73be9859fb4d44fe9e108aac5713990 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sat, 29 Aug 2026 15:52:48 -0700 Subject: [PATCH 65/68] fix(ui): keep a deleted-from search query instead of blanking the box The paginated search select diffs the input against the selected option's label to recover what the user just typed, so a query started from a picked value searches for the new text rather than the label plus the new text. That diff can only express an insertion: it walks a common prefix and a common suffix and returns what sits between them. A deletion leaves nothing between them, so every deletion-only edit returned the empty string. Backspacing once in a field showing a selected label therefore threw the edit away. The empty result was stored as the query, the controlled input re-rendered blank, and the server was asked for the unfiltered page instead of the text the user left in the box. An edit that yields no insertion but did change the value is a deletion, and there the remaining text is the query the user means. Insertions and whole-selection replacements are untouched. --- .../shared/PaginatedSearchSelect.test.tsx | 28 +++++++++++++++++++ .../shared/PaginatedSearchSelect.tsx | 7 ++++- 2 files changed, 34 insertions(+), 1 deletion(-) diff --git a/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.test.tsx b/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.test.tsx index c37589f63f5..d248cf311cc 100644 --- a/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.test.tsx @@ -361,4 +361,32 @@ describe("PaginatedSearchSelect", () => { await user.click(screen.getByRole("combobox")); expect(await screen.findByTestId("paginated-search-select-loading-more")).toBeInTheDocument(); }); + + it("queries the trimmed label when a character is deleted from the end of the selection", async () => { + const user = userEvent.setup(); + const onSearchChange = vi.fn(); + renderSelect({ onSearchChange, value: "alias-alpha" }); + + const input = screen.getByRole("combobox") as HTMLInputElement; + input.focus(); + input.setSelectionRange(input.value.length, input.value.length); + await user.keyboard("{Backspace}"); + + expect(input).toHaveValue("alias-alph"); + await waitFor(() => expect(onSearchChange).toHaveBeenLastCalledWith("alias-alph")); + }); + + it("queries what is left when a character is deleted from inside the selection", async () => { + const user = userEvent.setup(); + const onSearchChange = vi.fn(); + renderSelect({ onSearchChange, value: "alias-alpha" }); + + const input = screen.getByRole("combobox") as HTMLInputElement; + input.focus(); + input.setSelectionRange(6, 6); + await user.keyboard("{Backspace}"); + + expect(input).toHaveValue("aliasalpha"); + await waitFor(() => expect(onSearchChange).toHaveBeenLastCalledWith("aliasalpha")); + }); }); diff --git a/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.tsx b/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.tsx index 0f25260aad0..4b7ef7401c7 100644 --- a/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.tsx +++ b/ui/litellm-dashboard/src/components/shared/PaginatedSearchSelect.tsx @@ -50,6 +50,11 @@ const typedInsertion = (previous: string, next: string): string => { return next.slice(start, next.length - end); }; +const editedQuery = (label: string, next: string): string => { + const inserted = typedInsertion(label, next); + return inserted === "" && next !== label ? next : inserted; +}; + export function PaginatedSearchSelect({ options, value, @@ -101,7 +106,7 @@ export function PaginatedSearchSelect({ const replacedWholeInput = wholeSelectionRef.current; wholeSelectionRef.current = false; handleInputValueChange( - typedQuery === null && !replacedWholeInput ? typedInsertion(selected?.label ?? "", next) : next, + typedQuery === null && !replacedWholeInput ? editedQuery(selected?.label ?? "", next) : next, reason, ); }; From d28621685a6b0038b52d52dc43701fce0cc71dc1 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Tue, 4 Aug 2026 16:56:37 -0700 Subject: [PATCH 66/68] feat(proxy): add LiteLLM_BudgetWindowSpend table for per-window budget spend Multi-window budgets (budget_limits on keys/teams) currently keep window spend only in cache. Every cold or expired counter recomputes the window by aggregating LiteLLM_SpendLogs, which has no usable index for that query and saturates the DB on large tables (#35766). This adds a LiteLLM_BudgetWindowSpend table holding one row per configured window, keyed (entity_type, entity_id, window_duration), with window_start identifying the period the spend belongs to. Follow-up PRs maintain these rows from the spend update writer and move window budget enforcement reads onto them. --- .../migration.sql | 13 +++++++++++++ .../litellm_proxy_extras/schema.prisma | 12 ++++++++++++ litellm/proxy/schema.prisma | 12 ++++++++++++ schema.prisma | 12 ++++++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 +- 5 files changed, 50 insertions(+), 1 deletion(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260804162853_add_budget_window_spend_table/migration.sql diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260804162853_add_budget_window_spend_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260804162853_add_budget_window_spend_table/migration.sql new file mode 100644 index 00000000000..45cc927a328 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260804162853_add_budget_window_spend_table/migration.sql @@ -0,0 +1,13 @@ +-- CreateTable +CREATE TABLE IF NOT EXISTS "LiteLLM_BudgetWindowSpend" ( + "entity_type" TEXT NOT NULL, + "entity_id" TEXT NOT NULL, + "window_duration" TEXT NOT NULL, + "window_start" TIMESTAMP(3) NOT NULL, + "spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0, + "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "updated_at" TIMESTAMP(3) NOT NULL, + + CONSTRAINT "LiteLLM_BudgetWindowSpend_pkey" PRIMARY KEY ("entity_type","entity_id","window_duration") +); + diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 2bb850139a2..7818e3acd56 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -649,6 +649,18 @@ model LiteLLM_SpendLogs { @@index([session_id]) } +model LiteLLM_BudgetWindowSpend { + entity_type String + entity_id String + window_duration String + window_start DateTime + spend Float @default(0.0) + created_at DateTime @default(now()) + updated_at DateTime @updatedAt + + @@id([entity_type, entity_id, window_duration]) +} + // View spend, model, api_key per request model LiteLLM_ErrorLogs { request_id String @id @default(uuid()) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 2bb850139a2..7818e3acd56 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -649,6 +649,18 @@ model LiteLLM_SpendLogs { @@index([session_id]) } +model LiteLLM_BudgetWindowSpend { + entity_type String + entity_id String + window_duration String + window_start DateTime + spend Float @default(0.0) + created_at DateTime @default(now()) + updated_at DateTime @updatedAt + + @@id([entity_type, entity_id, window_duration]) +} + // View spend, model, api_key per request model LiteLLM_ErrorLogs { request_id String @id @default(uuid()) diff --git a/schema.prisma b/schema.prisma index 2bb850139a2..7818e3acd56 100644 --- a/schema.prisma +++ b/schema.prisma @@ -649,6 +649,18 @@ model LiteLLM_SpendLogs { @@index([session_id]) } +model LiteLLM_BudgetWindowSpend { + entity_type String + entity_id String + window_duration String + window_start DateTime + spend Float @default(0.0) + created_at DateTime @default(now()) + updated_at DateTime @updatedAt + + @@id([entity_type, entity_id, window_duration]) +} + // View spend, model, api_key per request model LiteLLM_ErrorLogs { request_id String @id @default(uuid()) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 6c5319f0361..c1a85a297c9 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -26073,7 +26073,7 @@ export interface components { * @description Default role assigned to new users created * @default internal_user_viewer */ - user_role: ("proxy_admin" | "proxy_admin_viewer" | "internal_user" | "internal_user_viewer") | null; + user_role: ("internal_user" | "internal_user_viewer" | "proxy_admin" | "proxy_admin_viewer") | null; }; /** * DefaultTeamSSOParams From a923502132a05e0ea52f3430d0abd97de66048b5 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Sat, 29 Aug 2026 11:36:59 -0700 Subject: [PATCH 67/68] chore(migrations): drop the generated comment from the budget window spend migration --- .../20260804162853_add_budget_window_spend_table/migration.sql | 1 - 1 file changed, 1 deletion(-) diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260804162853_add_budget_window_spend_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260804162853_add_budget_window_spend_table/migration.sql index 45cc927a328..c3018006adb 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260804162853_add_budget_window_spend_table/migration.sql +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260804162853_add_budget_window_spend_table/migration.sql @@ -1,4 +1,3 @@ --- CreateTable CREATE TABLE IF NOT EXISTS "LiteLLM_BudgetWindowSpend" ( "entity_type" TEXT NOT NULL, "entity_id" TEXT NOT NULL, From 774954d19a578061ec43abd74d4da6d1a76b58ea Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Sat, 29 Aug 2026 12:27:49 -0700 Subject: [PATCH 68/68] chore(ui): drop unrelated schema.d.ts enum reorder from the window spend schema branch --- ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index c1a85a297c9..6c5319f0361 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -26073,7 +26073,7 @@ export interface components { * @description Default role assigned to new users created * @default internal_user_viewer */ - user_role: ("internal_user" | "internal_user_viewer" | "proxy_admin" | "proxy_admin_viewer") | null; + user_role: ("proxy_admin" | "proxy_admin_viewer" | "internal_user" | "internal_user_viewer") | null; }; /** * DefaultTeamSSOParams