From cda19d6a29ebcdb23cad53d2beb215d34fa8c746 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Thu, 15 Oct 2026 22:35:23 -0700 Subject: [PATCH 01/23] test(integration): wait for the gen-AI span before reading the excluded-services operator trace (#46846) The operator trace wait stopped at the root and redis spans, but the gen-AI span is closed later by the async success callback. When it had not landed yet, the tenant mirror was built without it and the gen-AI assertion failed. Co-authored-by: yucheng-berri <293926326+yucheng-berri@users.noreply.github.com> --- .../test_otel_excluded_services_matrix.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/tests/integration/observability/test_otel_excluded_services_matrix.py b/tests/integration/observability/test_otel_excluded_services_matrix.py index 07646a0b189..8a7236f8626 100644 --- a/tests/integration/observability/test_otel_excluded_services_matrix.py +++ b/tests/integration/observability/test_otel_excluded_services_matrix.py @@ -358,6 +358,10 @@ def _has_root(spans: tuple[Span, ...]) -> bool: return any(span["kind"] == SERVER for span in spans) +def _has_model_span(spans: tuple[Span, ...]) -> bool: + return any("gen_ai.operation.name" in span["attributes"] for span in spans) + + def _trace_of_call(sink: str, call_id: str, since: int) -> tuple[Span, ...]: _, spans = recorded_spans(sink, since) traces: Final = {span["trace_id"] for span in spans if span["attributes"].get("litellm.call_id") == call_id} @@ -367,7 +371,7 @@ def _trace_of_call(sink: str, call_id: str, since: int) -> tuple[Span, ...]: def _operator_trace(rig: Rig, sent: Sent, cursors: Cursors) -> tuple[Span, ...]: trace: Final = eventually( lambda: _trace_of_call(rig.sinks.operator, sent.call_id, cursors.operator), - lambda spans: _has_root(spans) and "redis" in _db_systems(spans), + lambda spans: _has_root(spans) and "redis" in _db_systems(spans) and _has_model_span(spans), seconds=40, ) assert len({span["trace_id"] for span in trace}) == 1, _names(trace) @@ -405,7 +409,7 @@ def _assert_tenant_mirrors(rig: Rig, operator: tuple[Span, ...], cursors: Cursor def _assert_withheld(rig: Rig, sent: Sent, cursors: Cursors) -> tuple[Span, ...]: tenant: Final = _assert_tenant_mirrors(rig, _operator_trace(rig, sent, cursors), cursors) - assert any("gen_ai.operation.name" in span["attributes"] for span in tenant), _names(tenant) + assert _has_model_span(tenant), _names(tenant) return tenant @@ -719,7 +723,7 @@ def _assert_tenant_kept( return ( sum(1 for span in spans if span["kind"] == SERVER) == 1 and "redis" in _db_systems(spans) - and (not needs_model_span or any("gen_ai.operation.name" in span["attributes"] for span in spans)) + and (not needs_model_span or _has_model_span(spans)) ) tenant: Final = eventually( @@ -730,7 +734,7 @@ def _assert_tenant_kept( ) assert sum(1 for span in tenant if span["kind"] == SERVER) == 1, _names(tenant) assert "redis" in _db_systems(tenant), f"redis spans missing at the tenant: {_names(tenant)}" - assert not needs_model_span or any("gen_ai.operation.name" in span["attributes"] for span in tenant), _names(tenant) + assert not needs_model_span or _has_model_span(tenant), _names(tenant) return tenant From c8b73cde8471ca0ac2e8237bf5270103a64b3bfc Mon Sep 17 00:00:00 2001 From: Dan Lemon Date: Fri, 16 Oct 2026 08:37:34 +0200 Subject: [PATCH 02/23] feat(soniox): decode proxy form fields for context, translation and flags (#40605) * feat(soniox): decode proxy form fields for context, translation and flags Multipart form fields reach the proxy as strings, so Soniox received context and translation as JSON strings and booleans as "true". Decode them at the handler entry point and add typed shapes for the context and translation request fields. * fix(soniox): keep bracket-prefixed free-form context and move types under llms Free-form context such as "[inaudible] consultation" is valid Soniox input, so a string that fails JSON parsing is now forwarded unchanged instead of returning a 400. The Soniox request types live next to the provider code. * test(soniox): exercise SDK-typed context and translation shapes through the decoder * refactor(soniox): derive handler options and provider params without mutating a local dict * chore(soniox): drop explanatory comment on the non-forwarded param set * refactor(soniox): return invalid form bool as a value and raise once in the handler * fix(soniox): validate decoded form fields and keep language with hints Validate structured context before sending it as an object, answer 400 for a bad polling interval or flag unless drop_params is on, nest SDK bracketed extra_body fields, keep the OpenAI language at the head of language_hints, and return the translated text for one-way translation without tags. * fix(soniox): keep the source transcript, rebuild SDK list fields, and reject malformed JSON values * fix(soniox): keep cleanup strings whole and drop the unused translation type * fix(soniox): let explicit language_hints win over language and test the proxy form path --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../soniox/audio_transcription/handler.py | 101 ++++++--- .../audio_transcription/transformation.py | 193 ++++++++++++++-- litellm/llms/soniox/types.py | 24 ++ ...test_soniox_audio_transcription_handler.py | 210 ++++++++++++++++++ ...niox_audio_transcription_transformation.py | 203 ++++++++++++++++- 5 files changed, 672 insertions(+), 59 deletions(-) create mode 100644 litellm/llms/soniox/types.py diff --git a/litellm/llms/soniox/audio_transcription/handler.py b/litellm/llms/soniox/audio_transcription/handler.py index b53955ae75d..f86bc142898 100644 --- a/litellm/llms/soniox/audio_transcription/handler.py +++ b/litellm/llms/soniox/audio_transcription/handler.py @@ -19,11 +19,14 @@ import asyncio import math import time from collections.abc import Coroutine, Mapping, Sequence +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final import httpx +from pydantic import TypeAdapter from typing_extensions import ReadOnly, TypedDict +import litellm from litellm.litellm_core_utils.audio_utils.utils import ( get_audio_file_name, process_audio_file, @@ -35,7 +38,12 @@ from litellm.llms.custom_httpx.http_handler import ( get_httpx_client, ) from litellm.llms.soniox.audio_transcription.transformation import ( + SONIOX_HANDLER_ONLY_PARAMS, SonioxAudioTranscriptionConfig, + SonioxInvalidFormParam, + decode_soniox_form_params, + raise_soniox_form_error, + soniox_language_hints, ) from litellm.llms.soniox.common_utils import ( SONIOX_DEFAULT_CLEANUP, @@ -58,6 +66,39 @@ else: LiteLLMLoggingObj = Any +_CLEANUP_ENTRIES: Final = TypeAdapter(tuple[object, ...]) +_NOT_FORWARDED_TO_SONIOX: Final[frozenset[str]] = frozenset( + (*SONIOX_HANDLER_ONLY_PARAMS, "audio_url", "file_id", "response_format", "language") +) + + +def _optional_str(value: object) -> str | None: + return None if value is None else str(value) + + +def _float_or_default(value: object, default: float) -> float: + try: + parsed: Final = float(str(value)) + except ValueError: + return default + return parsed if math.isfinite(parsed) else default + + +def _int_or_default(value: object, default: int) -> int: + try: + return int(value) if isinstance(value, (int, float)) else int(str(value)) + except (ValueError, OverflowError): + return default + + +def _cleanup_targets(raw: object) -> tuple[str, ...]: + if raw is None: + return () + if isinstance(raw, str): + return (raw,) + return tuple(target for target in _CLEANUP_ENTRIES.validate_python(raw) if isinstance(target, str)) + + class _TranscriptionMeta(TypedDict, total=False): """Fields the handler reads from a Soniox transcription object.""" @@ -182,7 +223,7 @@ class SonioxAudioTranscriptionHandler: ) -> tuple[ dict[str, str], # auth headers str, # api_base (no trailing slash) - dict[str, object], # body for POST /v1/transcriptions (without file_id/audio_url) + Mapping[str, object], # body for POST /v1/transcriptions (without file_id/audio_url) _HandlerOptions, # handler-only options (poll interval, cleanup, ...) ]: # Validate env -> auth headers. @@ -198,52 +239,42 @@ class SonioxAudioTranscriptionHandler: base_url: Final = get_soniox_api_base(api_base) - # Operate on a local copy so we don't mutate the caller's dict - # (the caller may reuse `optional_params` for retries or logging). - params: Final = dict(optional_params) - - # Pull handler-only kwargs out of params so they aren't sent - # to Soniox. - poll_interval = float(params.pop("soniox_polling_interval", SONIOX_DEFAULT_POLL_INTERVAL)) - try: - max_attempts = int(params.pop("soniox_max_polling_attempts", SONIOX_DEFAULT_MAX_POLL_ATTEMPTS)) - except (ValueError, OverflowError): - max_attempts = SONIOX_DEFAULT_MAX_POLL_ATTEMPTS - cleanup_raw: Final = params.pop("soniox_cleanup", SONIOX_DEFAULT_CLEANUP) - if cleanup_raw is None: - cleanup: list[str] = [] - elif isinstance(cleanup_raw, str): - cleanup = [cleanup_raw] - else: - cleanup = list(cleanup_raw) - filename_override: Final = params.pop("filename", None) + decoded: Final = decode_soniox_form_params( + optional_params, drop_params=litellm_params.get("drop_params") is True or litellm.drop_params is True + ) + if isinstance(decoded, SonioxInvalidFormParam): + raise_soniox_form_error(decoded) # Server-side clamps. Caller-supplied poll settings (from request kwargs) # are bounded so an authenticated caller cannot force a worker into a # tight poll loop (zero interval) or pin it indefinitely (huge attempt # count). Total polling time is bounded by # SONIOX_MAX_POLL_ATTEMPTS * SONIOX_MAX_POLL_INTERVAL. - if not math.isfinite(poll_interval): - poll_interval = SONIOX_DEFAULT_POLL_INTERVAL - clamped_poll_interval: Final = max(SONIOX_MIN_POLL_INTERVAL, min(poll_interval, SONIOX_MAX_POLL_INTERVAL)) - clamped_max_attempts: Final = max(1, min(max_attempts, SONIOX_MAX_POLL_ATTEMPTS)) + poll_interval: Final = _float_or_default(decoded.get("soniox_polling_interval"), SONIOX_DEFAULT_POLL_INTERVAL) + max_attempts: Final = _int_or_default( + decoded.get("soniox_max_polling_attempts"), SONIOX_DEFAULT_MAX_POLL_ATTEMPTS + ) + cleanup: Final = _cleanup_targets(decoded.get("soniox_cleanup", SONIOX_DEFAULT_CLEANUP)) + language_hints: Final = soniox_language_hints(decoded.get("language"), decoded.get("language_hints")) # response_format is handled by LiteLLM post-processing, not Soniox. handler_opts: Final[_HandlerOptions] = { - "poll_interval": clamped_poll_interval, - "max_attempts": clamped_max_attempts, + "poll_interval": max(SONIOX_MIN_POLL_INTERVAL, min(poll_interval, SONIOX_MAX_POLL_INTERVAL)), + "max_attempts": max(1, min(max_attempts, SONIOX_MAX_POLL_ATTEMPTS)), "cleanup": cleanup, - "filename_override": filename_override, - "audio_url": params.pop("audio_url", None), - "file_id": params.pop("file_id", None), - "response_format": params.pop("response_format", None), + "filename_override": _optional_str(decoded.get("filename")), + "audio_url": _optional_str(decoded.get("audio_url")), + "file_id": _optional_str(decoded.get("file_id")), + "response_format": _optional_str(decoded.get("response_format")), } - # Soniox does not accept `language` directly; map_openai_params should - # already have translated it, but drop any leftover to be safe. - params.pop("language", None) - - return auth_headers, base_url, params, handler_opts + provider_params: Final = MappingProxyType( + { + **{key: value for key, value in decoded.items() if key not in _NOT_FORWARDED_TO_SONIOX}, + **({} if language_hints is None else {"language_hints": language_hints}), + } + ) + return auth_headers, base_url, provider_params, handler_opts def _build_create_body( self, diff --git a/litellm/llms/soniox/audio_transcription/transformation.py b/litellm/llms/soniox/audio_transcription/transformation.py index aa1ff239320..30d872f248f 100644 --- a/litellm/llms/soniox/audio_transcription/transformation.py +++ b/litellm/llms/soniox/audio_transcription/transformation.py @@ -9,9 +9,14 @@ async API requires multiple HTTP calls and does not fit the single-request contract of `base_llm_http_handler.audio_transcriptions`. """ -from typing import Any, Final +import re +from collections.abc import Mapping +from dataclasses import dataclass +from types import MappingProxyType +from typing import Any, Final, NoReturn from httpx import Headers, Response +from pydantic import TypeAdapter, ValidationError from litellm.llms.base_llm.audio_transcription.transformation import ( AudioTranscriptionRequestData, @@ -26,6 +31,7 @@ from litellm.llms.soniox.common_utils import ( render_soniox_tokens_as_srt, render_soniox_tokens_as_vtt, ) +from litellm.llms.soniox.types import SonioxContext from litellm.types.llms.openai import ( AllMessageValues, OpenAIAudioTranscriptionOptionalParams, @@ -58,6 +64,152 @@ SONIOX_HANDLER_ONLY_PARAMS: Final[list[str]] = [ ] +SONIOX_BOOL_PARAMS: Final[frozenset[str]] = frozenset( + {"enable_speaker_diarization", "enable_language_identification", "language_hints_strict"} +) +_JSON_OBJECT: Final = TypeAdapter(dict[str, object]) +_STRUCTURED_CONTEXT: Final = TypeAdapter(SonioxContext) +_FORM_BOOL: Final = TypeAdapter(bool) +_ITEMS: Final = TypeAdapter(tuple[object, ...]) +_BRACKETED_FORM_FIELD: Final = re.compile( + r"(?Pcontext|translation)\[(?P\w+)\](?:\[\]\[(?P\w+)\])?" +) +_BracketedField = tuple[str, str, str | None, object] + + +@dataclass(frozen=True, slots=True) +class SonioxInvalidFormParam: + key: str + value: str + expected: str + + +def raise_soniox_form_error(error: SonioxInvalidFormParam) -> NoReturn: + raise SonioxException( + message=f"`{error.key}` must be {error.expected}, got {error.value!r}", status_code=400, headers=None + ) + + +def soniox_language_hints(language: object, hints: object) -> object: + """Soniox has no `language` field, so the OpenAI `language` becomes `language_hints` when none are set.""" + if hints is not None or not isinstance(language, str) or not language: + return hints + return (language,) + + +def _decode_bool(key: str, value: str) -> bool | SonioxInvalidFormParam: + try: + return _FORM_BOOL.validate_python(value.strip()) + except ValidationError: + return SonioxInvalidFormParam(key=key, value=value, expected="a boolean") + + +def _decode_polling_interval(value: str) -> str | SonioxInvalidFormParam: + try: + float(value) + except ValueError: + return SonioxInvalidFormParam(key="soniox_polling_interval", value=value, expected="a number") + return value + + +def _decode_context(value: str) -> SonioxContext | str: + if not value.lstrip().startswith("{"): + return value + try: + return _STRUCTURED_CONTEXT.validate_json(value) + except ValidationError: + return value + + +def _decode_translation(value: str) -> Mapping[str, object] | SonioxInvalidFormParam: + try: + return _JSON_OBJECT.validate_json(value) + except ValidationError: + return SonioxInvalidFormParam(key="translation", value=value, expected="a JSON object") + + +def _decode_language_hints(value: str) -> tuple[object, ...] | SonioxInvalidFormParam: + if not value.lstrip().startswith("["): + return tuple(hint.strip() for hint in value.split(",") if hint.strip()) + try: + return _ITEMS.validate_json(value) + except ValidationError: + return SonioxInvalidFormParam(key="language_hints", value=value, expected="a JSON list") + + +def _decode_form_value(key: str, value: object) -> object | SonioxInvalidFormParam: + if not isinstance(value, str): + return value + if key in SONIOX_BOOL_PARAMS: + return _decode_bool(key, value) + if key == "soniox_polling_interval": + return _decode_polling_interval(value) + if key == "context": + return _decode_context(value) + if key == "translation": + return _decode_translation(value) + if key == "language_hints": + return _decode_language_hints(value) + return value + + +def _nested_form_param(param: str, fields: tuple[_BracketedField, ...]) -> Mapping[str, object]: + owned: Final = tuple(field for field in fields if field[0] == param) + flat: Final = {name: value for _, name, entry_field, value in owned if entry_field is None} + listed: Final = {name for _, name, entry_field, _ in owned if entry_field is not None} + return { + **flat, + **{ + name: ( + { + entry_field: value + for _, owner, entry_field, value in owned + if owner == name and entry_field is not None + }, + ) + for name in listed + }, + } + + +def _is_soniox_context(value: object) -> bool: + try: + _STRUCTURED_CONTEXT.validate_python(value) + except ValidationError: + return False + return True + + +def _nest_bracketed_form_fields(params: Mapping[str, object]) -> Mapping[str, object]: + """The OpenAI SDK sends a dict in `extra_body` as `context[terms]` and `context[general][][key]` form keys.""" + matches: Final = {key: _BRACKETED_FORM_FIELD.fullmatch(key) for key in params} + fields: Final = tuple( + (match["param"], match["field"], match["entry_field"], params[key]) for key, match in matches.items() if match + ) + nested: Final = {param: _nested_form_param(param, fields) for param in {field[0] for field in fields}} + kept: Final = {param: value for param, value in nested.items() if param != "context" or _is_soniox_context(value)} + unbracketed: Final = { + key: value for key, value in params.items() if (match := matches[key]) is None or match["param"] not in kept + } + return {**kept, **unbracketed} + + +def decode_soniox_form_params( + optional_params: Mapping[str, object], + drop_params: bool = False, +) -> Mapping[str, object] | SonioxInvalidFormParam: + """Multipart form fields reach the proxy as strings; restore the JSON types Soniox expects.""" + decoded: Final = { + key: _decode_form_value(key, value) for key, value in _nest_bracketed_form_fields(optional_params).items() + } + invalid: Final = next((value for value in decoded.values() if isinstance(value, SonioxInvalidFormParam)), None) + if invalid is not None and not drop_params: + return invalid + return MappingProxyType( + {key: value for key, value in decoded.items() if not isinstance(value, SonioxInvalidFormParam)} + ) + + class SonioxAudioTranscriptionConfig(BaseAudioTranscriptionConfig): """Configuration for Soniox async speech-to-text transcription.""" @@ -74,25 +226,26 @@ class SonioxAudioTranscriptionConfig(BaseAudioTranscriptionConfig): model: str, drop_params: bool, ) -> dict: - # Translate the OpenAI `language` param into Soniox `language_hints`. - if "language" in non_default_params and non_default_params["language"]: - language: Final = non_default_params["language"] - existing_hints: Final = optional_params.get("language_hints") - if not existing_hints: - optional_params["language_hints"] = [language] - elif language not in existing_hints: - optional_params["language_hints"] = [language] + list(existing_hints) - - # Capture response_format for post-processing (not sent to Soniox API). - if "response_format" in non_default_params: - optional_params["response_format"] = non_default_params["response_format"] - - # Pass through Soniox-native kwargs unchanged. - for key in SONIOX_PASSTHROUGH_PARAMS + SONIOX_HANDLER_ONLY_PARAMS: - if key in non_default_params and non_default_params[key] is not None: - optional_params[key] = non_default_params[key] - - return optional_params + language: Final[object] = non_default_params.get("language") + language_params: Final[Mapping[str, object]] = ( + { + "language": language, + "language_hints": soniox_language_hints(language, optional_params.get("language_hints")), + } + if language + else {} + ) + response_format_params: Final[Mapping[str, object]] = ( + {"response_format": non_default_params["response_format"]} + if "response_format" in non_default_params + else {} + ) + passthrough_params: Final[Mapping[str, object]] = { + key: non_default_params[key] + for key in SONIOX_PASSTHROUGH_PARAMS + SONIOX_HANDLER_ONLY_PARAMS + if non_default_params.get(key) is not None + } + return {**optional_params, **language_params, **response_format_params, **passthrough_params} def get_error_class(self, error_message: str, status_code: int, headers: dict | Headers) -> BaseLLMException: return SonioxException(message=error_message, status_code=status_code, headers=headers) diff --git a/litellm/llms/soniox/types.py b/litellm/llms/soniox/types.py new file mode 100644 index 00000000000..806d2a8fb0e --- /dev/null +++ b/litellm/llms/soniox/types.py @@ -0,0 +1,24 @@ +from collections.abc import Sequence + +from pydantic import ConfigDict, with_config +from typing_extensions import ReadOnly, TypedDict + + +class SonioxContextGeneralEntry(TypedDict): + key: ReadOnly[str] + value: ReadOnly[str] + + +class SonioxTranslationTerm(TypedDict): + source: ReadOnly[str] + target: ReadOnly[str] + + +@with_config(ConfigDict(extra="forbid")) +class SonioxContext(TypedDict, total=False): + """Soniox `context` request field: https://soniox.com/docs/stt/concepts/context""" + + general: ReadOnly[Sequence[SonioxContextGeneralEntry]] + text: ReadOnly[str] + terms: ReadOnly[Sequence[str]] + translation_terms: ReadOnly[Sequence[SonioxTranslationTerm]] diff --git a/tests/unit/llms/soniox/audio_transcription/test_soniox_audio_transcription_handler.py b/tests/unit/llms/soniox/audio_transcription/test_soniox_audio_transcription_handler.py index d2d7d2247f1..9c70ebe7ece 100644 --- a/tests/unit/llms/soniox/audio_transcription/test_soniox_audio_transcription_handler.py +++ b/tests/unit/llms/soniox/audio_transcription/test_soniox_audio_transcription_handler.py @@ -339,6 +339,65 @@ class TestPollLimitsClamping: a worker on tight poll loops. """ + def test_should_raise_400_when_form_bool_param_is_invalid(self): + handler = SonioxAudioTranscriptionHandler() + with pytest.raises(SonioxException) as exc_info: + handler._prepare( + audio_file=None, + optional_params={ + "enable_speaker_diarization": "maybe", + "audio_url": "https://example.com/a.wav", + }, + litellm_params={}, + api_key="sk-test", + api_base=None, + provider_config=SonioxAudioTranscriptionConfig(), + headers={}, + ) + assert exc_info.value.status_code == 400 + assert "enable_speaker_diarization" in str(exc_info.value) + + def test_should_raise_400_when_poll_interval_is_not_a_number(self): + with pytest.raises(SonioxException) as exc_info: + SonioxAudioTranscriptionHandler()._prepare( + audio_file=None, + optional_params={"soniox_polling_interval": "abc", "audio_url": "https://example.com/a.wav"}, + litellm_params={}, + api_key="sk-test", + api_base=None, + provider_config=SonioxAudioTranscriptionConfig(), + headers={}, + ) + assert exc_info.value.status_code == 400 + assert "soniox_polling_interval" in str(exc_info.value) + + def test_should_use_default_poll_interval_for_bad_value_when_drop_params_is_set(self): + from litellm.llms.soniox.common_utils import SONIOX_DEFAULT_POLL_INTERVAL + + _, _, _, opts = SonioxAudioTranscriptionHandler()._prepare( + audio_file=None, + optional_params={"soniox_polling_interval": "abc", "audio_url": "https://example.com/a.wav"}, + litellm_params={"drop_params": True}, + api_key="sk-test", + api_base=None, + provider_config=SonioxAudioTranscriptionConfig(), + headers={}, + ) + assert opts["poll_interval"] == SONIOX_DEFAULT_POLL_INTERVAL + + @pytest.mark.parametrize("raw, expected", [(10.0, 10), (10, 10), ("10", 10)]) + def test_should_keep_numeric_max_polling_attempts(self, raw: object, expected: int): + _, _, _, opts = SonioxAudioTranscriptionHandler()._prepare( + audio_file=None, + optional_params={"soniox_max_polling_attempts": raw, "audio_url": "https://example.com/a.wav"}, + litellm_params={}, + api_key="sk-test", + api_base=None, + provider_config=SonioxAudioTranscriptionConfig(), + headers={}, + ) + assert opts["max_attempts"] == expected + def test_should_clamp_poll_interval_to_minimum(self): from litellm.llms.soniox.common_utils import SONIOX_MIN_POLL_INTERVAL @@ -707,6 +766,112 @@ class TestCleanupNormalization: assert "https://api.soniox.com/v1/transcriptions/tx_1" in deletes + def test_should_skip_non_string_cleanup_entries(self): + _, _, _, opts = SonioxAudioTranscriptionHandler()._prepare( + audio_file=None, + optional_params={"soniox_cleanup": ["transcription", 1, None], "audio_url": "https://example.com/a.wav"}, + litellm_params={}, + api_key="sk-test", + api_base=None, + provider_config=SonioxAudioTranscriptionConfig(), + headers={}, + ) + assert opts["cleanup"] == ("transcription",) + + @pytest.mark.parametrize( + "language_hints, expected", + [("en,es", ["en", "es"]), (["es"], ["es"]), (None, ["de"])], + ) + def test_should_use_language_only_without_language_hints(self, language_hints, expected, monkeypatch): + import litellm + + monkeypatch.setattr("time.sleep", lambda *_: None) + client = _MockSyncClient( + { + "POST https://api.soniox.com/v1/transcriptions": [_make_response({"id": "tx_1"})], + "GET https://api.soniox.com/v1/transcriptions/tx_1": [_make_response({"status": "completed"})], + "GET https://api.soniox.com/v1/transcriptions/tx_1/transcript": [ + _make_response({"text": "hallo", "tokens": []}) + ], + "DELETE https://api.soniox.com/v1/transcriptions/tx_1": [_make_response({})], + } + ) + litellm.transcription( + model="soniox/stt-async-v4", + file=("a.wav", b"RIFF", "audio/wav"), + language="de", + api_key="sk-test", + client=client, + audio_url="https://example.com/a.wav", + **({} if language_hints is None else {"language_hints": language_hints}), + ) + body = next(c for c in client.calls if c["method"] == "POST")["json"] + assert list(body["language_hints"]) == expected + assert "language" not in body + + +class TestProxyFormFields: + def test_should_send_soniox_types_for_fields_the_proxy_parsed_from_an_sdk_form(self, monkeypatch): + import litellm + from starlette.requests import Request + + from litellm.proxy.common_utils.http_parsing_utils import get_form_data + + monkeypatch.setattr("time.sleep", lambda *_: None) + multipart = httpx.Request( + "POST", + "http://proxy/v1/audio/transcriptions", + data={ + "model": "soniox/stt-async-v4", + "audio_url": "https://example.com/a.wav", + "context[terms][]": ["Celebrex"], + "context[general][][key]": "domain", + "context[general][][value]": "Healthcare", + "translation[type]": "one_way", + "translation[target_language]": "de", + "enable_speaker_diarization": "true", + "language_hints[]": ["en", "es"], + }, + files={"file": ("a.wav", b"RIFF", "audio/wav")}, + ) + multipart_body = multipart.read() + + async def receive(): + return {"type": "http.request", "body": multipart_body, "more_body": False} + + request = Request( + { + "type": "http", + "method": "POST", + "path": "/v1/audio/transcriptions", + "headers": [(b"content-type", multipart.headers["content-type"].encode())], + }, + receive, + ) + form = asyncio.run(get_form_data(request)) + client = _MockSyncClient( + { + "POST https://api.soniox.com/v1/transcriptions": [_make_response({"id": "tx_1"})], + "GET https://api.soniox.com/v1/transcriptions/tx_1": [_make_response({"status": "completed"})], + "GET https://api.soniox.com/v1/transcriptions/tx_1/transcript": [ + _make_response({"text": "hello", "tokens": []}) + ], + "DELETE https://api.soniox.com/v1/transcriptions/tx_1": [_make_response({})], + } + ) + litellm.transcription( + file=("a.wav", b"RIFF", "audio/wav"), + api_key="sk-test", + client=client, + **{key: value for key, value in form.items() if key != "file"}, + ) + body = json.loads(json.dumps(next(c for c in client.calls if c["method"] == "POST")["json"])) + assert body["context"] == {"terms": ["Celebrex"], "general": [{"key": "domain", "value": "Healthcare"}]} + assert body["translation"] == {"type": "one_way", "target_language": "de"} + assert body["enable_speaker_diarization"] is True + assert body["language_hints"] == ["en", "es"] + + class TestErrorResponses: def test_should_raise_on_4xx_during_create_with_json_error(self, monkeypatch): monkeypatch.setattr("time.sleep", lambda *_: None) @@ -780,6 +945,51 @@ class TestPassthroughBodyBuilding: post_call = next(c for c in client.calls if c["method"] == "POST") assert "context" not in post_call["json"] + def test_should_send_form_string_params_to_soniox_as_json_types(self, monkeypatch): + """Simulates the proxy path, where every multipart form field is a string.""" + monkeypatch.setattr("time.sleep", lambda *_: None) + responses = { + "POST https://api.soniox.com/v1/transcriptions": [ + _make_response({"id": "tx_1"}), + ], + "GET https://api.soniox.com/v1/transcriptions/tx_1": [ + _make_response({"status": "completed"}), + ], + "GET https://api.soniox.com/v1/transcriptions/tx_1/transcript": [ + _make_response({"text": "ok", "tokens": []}), + ], + "DELETE https://api.soniox.com/v1/transcriptions/tx_1": [ + _make_response({}), + ], + } + client = _MockSyncClient(responses) + form_params = { + "audio_url": "https://example.com/a.wav", + "context": '{"terms": ["Celebrex"], "general": [{"key": "domain", "value": "Healthcare"}]}', + "translation": '{"type": "one_way", "target_language": "es"}', + "language_hints": "en,es", + "language_hints_strict": "true", + "enable_speaker_diarization": "true", + "enable_language_identification": "false", + "soniox_polling_interval": "0.5", + } + SonioxAudioTranscriptionHandler().audio_transcriptions( + audio_file=None, + optional_params=form_params, + litellm_params={}, + atranscription=False, + **_common_call_kwargs(client), + ) + body = next(c for c in client.calls if c["method"] == "POST")["json"] + assert body["context"] == {"terms": ["Celebrex"], "general": [{"key": "domain", "value": "Healthcare"}]} + assert body["translation"] == {"type": "one_way", "target_language": "es"} + assert list(body["language_hints"]) == ["en", "es"] + assert body["language_hints_strict"] is True + assert body["enable_speaker_diarization"] is True + assert body["enable_language_identification"] is False + assert "soniox_polling_interval" not in body + assert form_params["context"].startswith("{"), "caller's params must not be mutated" + class TestSecretRedaction: """Secret-bearing fields must be redacted before reaching logging callbacks. diff --git a/tests/unit/llms/soniox/audio_transcription/test_soniox_audio_transcription_transformation.py b/tests/unit/llms/soniox/audio_transcription/test_soniox_audio_transcription_transformation.py index c99b28adc49..e3e1c9ab493 100644 --- a/tests/unit/llms/soniox/audio_transcription/test_soniox_audio_transcription_transformation.py +++ b/tests/unit/llms/soniox/audio_transcription/test_soniox_audio_transcription_transformation.py @@ -9,8 +9,16 @@ import pytest from litellm.llms.soniox.audio_transcription.transformation import ( SonioxAudioTranscriptionConfig, + SonioxInvalidFormParam, + decode_soniox_form_params, + raise_soniox_form_error, ) from litellm.llms.soniox.common_utils import SonioxException +from litellm.llms.soniox.types import ( + SonioxContext, + SonioxContextGeneralEntry, + SonioxTranslationTerm, +) from litellm.types.utils import TranscriptionResponse @@ -40,9 +48,9 @@ class TestMapOpenAIParams: model="stt-async-v4", drop_params=False, ) - assert result["language_hints"] == ["en"] + assert list(result["language_hints"]) == ["en"] - def test_should_prepend_language_to_existing_hints(self): + def test_should_keep_explicit_hints_over_language(self): cfg = SonioxAudioTranscriptionConfig() result = cfg.map_openai_params( non_default_params={"language": "en"}, @@ -50,7 +58,7 @@ class TestMapOpenAIParams: model="stt-async-v4", drop_params=False, ) - assert result["language_hints"] == ["en", "fr"] + assert list(result["language_hints"]) == ["fr"] def test_should_not_duplicate_language_already_in_hints(self): cfg = SonioxAudioTranscriptionConfig() @@ -60,7 +68,7 @@ class TestMapOpenAIParams: model="stt-async-v4", drop_params=False, ) - assert result["language_hints"] == ["en", "fr"] + assert list(result["language_hints"]) == ["en", "fr"] def test_should_passthrough_soniox_native_kwargs(self): cfg = SonioxAudioTranscriptionConfig() @@ -243,6 +251,21 @@ class TestTransformAudioTranscriptionResponse: assert "Speaker 1:" in resp.text assert "Speaker 2:" in resp.text + def test_should_keep_source_transcript_for_translation_without_tags(self): + cfg = SonioxAudioTranscriptionConfig() + payload = { + "transcript": { + "text": "Hello there.", + "tokens": [ + {"text": "Hello", "translation_status": "original"}, + {"text": " there.", "translation_status": "original"}, + {"text": "Hallo", "translation_status": "translation"}, + {"text": " zusammen.", "translation_status": "translation"}, + ], + } + } + assert cfg.build_response_from_payload(payload).text == "Hello there." + def test_should_set_language_when_all_tokens_share_one(self): cfg = SonioxAudioTranscriptionConfig() payload = { @@ -678,3 +701,175 @@ class TestGetErrorClass: err = cfg.get_error_class(error_message="boom", status_code=500, headers={}) assert isinstance(err, SonioxException) assert err.status_code == 500 + + +class TestDecodeSonioxFormParams: + """Proxy multipart form fields arrive as strings; Soniox needs real JSON types.""" + + def test_should_decode_json_string_context_with_all_four_sections(self): + context = { + "general": [{"key": "domain", "value": "Healthcare"}], + "text": "Follow-up visit for a patient on blood thinners.", + "terms": ["Celebrex", "Zyrtec"], + "translation_terms": [{"source": "Mr. Smith", "target": "Sr. Smith"}], + } + result = decode_soniox_form_params({"context": json.dumps(context)}) + assert result["context"] == context + + def test_should_keep_plain_text_context_as_string(self): + result = decode_soniox_form_params({"context": "medical conversation"}) + assert result["context"] == "medical conversation" + + def test_should_leave_sdk_typed_values_untouched(self): + context = SonioxContext( + general=[SonioxContextGeneralEntry(key="domain", value="Healthcare")], + text="Follow-up visit notes.", + terms=["Celebrex"], + translation_terms=[SonioxTranslationTerm(source="Mr. Smith", target="Sr. Smith")], + ) + translation = {"type": "one_way", "target_language": "es"} + result = decode_soniox_form_params( + { + "context": context, + "translation": translation, + "enable_speaker_diarization": True, + "language_hints": ["en", "es"], + } + ) + assert result["context"] is context + assert result["translation"] is translation + assert result["enable_speaker_diarization"] is True + assert result["language_hints"] == ["en", "es"] + + @pytest.mark.parametrize("raw", ["[inaudible] consultation", "{unbalanced", '{"terms": ["Celebrex"']) + def test_should_keep_bracket_prefixed_free_form_context_as_string(self, raw: str): + assert decode_soniox_form_params({"context": raw})["context"] == raw + + def test_should_decode_json_string_translation(self): + result = decode_soniox_form_params({"translation": '{"type": "one_way", "target_language": "es"}'}) + assert result["translation"] == {"type": "one_way", "target_language": "es"} + + @pytest.mark.parametrize( + "raw, expected", + [("true", True), ("True", True), ("1", True), ("false", False), ("FALSE", False), ("0", False)], + ) + def test_should_decode_boolean_strings(self, raw: str, expected: bool): + result = decode_soniox_form_params( + { + "enable_speaker_diarization": raw, + "enable_language_identification": raw, + "language_hints_strict": raw, + } + ) + assert result == { + "enable_speaker_diarization": expected, + "enable_language_identification": expected, + "language_hints_strict": expected, + } + + @pytest.mark.parametrize("raw, expected", [("on", True), ("yes", True), ("off", False), ("no", False)]) + def test_should_decode_html_style_boolean_strings(self, raw: str, expected: bool): + assert decode_soniox_form_params({"enable_speaker_diarization": raw})["enable_speaker_diarization"] is expected + + def test_should_return_error_value_on_non_boolean_string(self): + result = decode_soniox_form_params({"context": '{"general": []}', "enable_speaker_diarization": "maybe"}) + assert result == SonioxInvalidFormParam(key="enable_speaker_diarization", value="maybe", expected="a boolean") + + def test_should_drop_non_boolean_string_when_drop_params_is_set(self): + result = decode_soniox_form_params( + {"enable_speaker_diarization": "maybe", "language_hints_strict": "true"}, drop_params=True + ) + assert result == {"language_hints_strict": True} + + def test_should_return_error_value_on_non_numeric_polling_interval(self): + result = decode_soniox_form_params({"soniox_polling_interval": "abc"}) + assert result == SonioxInvalidFormParam(key="soniox_polling_interval", value="abc", expected="a number") + + def test_should_drop_non_numeric_polling_interval_when_drop_params_is_set(self): + assert decode_soniox_form_params({"soniox_polling_interval": "abc"}, drop_params=True) == {} + + def test_should_map_invalid_form_value_error_to_400(self): + with pytest.raises(SonioxException) as exc_info: + raise_soniox_form_error( + SonioxInvalidFormParam(key="enable_speaker_diarization", value="maybe", expected="a boolean") + ) + assert exc_info.value.status_code == 400 + assert "enable_speaker_diarization" in str(exc_info.value) + assert "maybe" in str(exc_info.value) + + @pytest.mark.parametrize( + "raw", ['["Celebrex", "Zyrtec"]', '{"note": "follow-up visit"}', '{"terms": "Celebrex"}', '"quoted"'] + ) + def test_should_keep_json_context_soniox_cannot_take_as_an_object_as_string(self, raw: str): + assert decode_soniox_form_params({"context": raw})["context"] == raw + + def test_should_nest_openai_sdk_bracketed_form_fields(self): + result = decode_soniox_form_params( + { + "context[terms]": ["Celebrex", "Zyrtec"], + "context[text]": "Follow-up visit.", + "translation[type]": "one_way", + "translation[target_language]": "de", + "client_reference_id": "ref-1", + } + ) + assert result == { + "context": {"terms": ["Celebrex", "Zyrtec"], "text": "Follow-up visit."}, + "translation": {"type": "one_way", "target_language": "de"}, + "client_reference_id": "ref-1", + } + + def test_should_prefer_explicit_field_over_bracketed_fields(self): + result = decode_soniox_form_params({"context": "plain text", "context[terms]": ["Celebrex"]}) + assert result == {"context": "plain text"} + + @pytest.mark.parametrize( + "raw, expected", + [('["en", "es"]', ["en", "es"]), ("en,es", ["en", "es"]), ("en, es ,", ["en", "es"]), ("en", ["en"])], + ) + def test_should_decode_language_hints_strings(self, raw: str, expected: list): + assert list(decode_soniox_form_params({"language_hints": raw})["language_hints"]) == expected + + @pytest.mark.parametrize("raw", ['{"type": "one_way"', "de", '["one_way"]']) + def test_should_return_error_value_on_translation_that_is_not_a_json_object(self, raw: str): + assert decode_soniox_form_params({"translation": raw}) == SonioxInvalidFormParam( + key="translation", value=raw, expected="a JSON object" + ) + + def test_should_return_error_value_on_malformed_json_language_hints(self): + assert decode_soniox_form_params({"language_hints": "[en, es]"}) == SonioxInvalidFormParam( + key="language_hints", value="[en, es]", expected="a JSON list" + ) + + def test_should_drop_malformed_json_params_when_drop_params_is_set(self): + result = decode_soniox_form_params( + {"translation": '{"type": "one_way"', "language_hints": "[en, es]", "language_hints_strict": "true"}, + drop_params=True, + ) + assert result == {"language_hints_strict": True} + + def test_should_rebuild_openai_sdk_list_of_object_context_fields(self): + result = decode_soniox_form_params( + { + "context[terms]": ["Celebrex"], + "context[general][][key]": "domain", + "context[general][][value]": "Healthcare", + "context[translation_terms][][source]": "Mr. Smith", + "context[translation_terms][][target]": "Sr. Smith", + } + ) + assert result == { + "context": { + "terms": ["Celebrex"], + "general": ({"key": "domain", "value": "Healthcare"},), + "translation_terms": ({"source": "Mr. Smith", "target": "Sr. Smith"},), + } + } + + def test_should_leave_bracketed_context_soniox_cannot_take_unnested(self): + params = {"context[terms]": ["Celebrex"], "context[note]": "follow-up"} + assert decode_soniox_form_params(params) == params + + def test_should_not_json_decode_unrelated_string_params(self): + result = decode_soniox_form_params({"client_reference_id": '["not", "json-decoded"]'}) + assert result["client_reference_id"] == '["not", "json-decoded"]' From c288dc7da0be4459436d3862f1315bd03b774681 Mon Sep 17 00:00:00 2001 From: Michelangelo Setaro <296101911+setarm01@users.noreply.github.com> Date: Fri, 16 Oct 2026 08:37:40 +0200 Subject: [PATCH 03/23] fix(router): let an exact base_model widen a deployment's context window (#46197) * fix(router): honor explicit bedrock base model capacity * fix(router): honor an exact base_model's capacity for every non-Azure provider Extends the Bedrock-only check to every non-Azure provider and only uses the base_model when it has its own exact cost-map entry, so an unmapped or wildcard base_model falls back to the route id instead of raising in the router. * test(router): call _maps_to_exact_model_info directly in its cost-map tests * fix(router): let an exact base_model only widen a deployment's context window A base_model used as a pricing stand-in with a smaller window no longer narrows the route's limit, and a route with no window never takes one from base_model * test(router): count the base_model window helpers as covered through get_router_model_info * fix(router): widen only from a model_info base_model naming the same model A stand-in base_model for a different, larger model widened a smaller deployment's window, so pre-call checks sent it prompts the provider rejects with 400. The window now widens only when model_info.base_model names the same provider model id as the route, which is also the field /model/info reads. * fix(router): match a base_model's model id across Bedrock region prefixes and @ qualifiers * fix(router): read litellm_params.base_model too and look the base row up as written * test(router): cover a same-model base_model window on chat, messages, responses, and complexity escalation --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/router.py | 59 ++- .../router_code_coverage.py | 5 + ...t_base_model_context_window_owned_proxy.py | 171 ++++++ .../unit/test_router_model_cost_isolation.py | 487 +++++++++++++----- 4 files changed, 592 insertions(+), 130 deletions(-) create mode 100644 tests/integration/routing/test_base_model_context_window_owned_proxy.py diff --git a/litellm/router.py b/litellm/router.py index a549bb06e7d..96c96310d45 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -128,6 +128,7 @@ from litellm.llms.base_llm.vector_store.transformation import ( RouterVectorStoreEmbeddingExecutor, vector_store_request_metadata, ) +from litellm.llms.bedrock.common_utils import get_bedrock_base_model from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, get_async_httpx_client from litellm.llms.openai_like.json_loader import JSONProviderRegistry from litellm.llms.openai_like.model_info import ( @@ -374,6 +375,7 @@ from litellm.utils import ( get_non_default_completion_params, get_secret, get_utc_datetime, + is_generalized_model_info, is_region_allowed, provider_rejectable_params, set_live_deployment_replay, @@ -991,6 +993,49 @@ def as_output_cap(value: object) -> int | None: return cap if cap >= 0 else None +def _exact_base_model_max_input_tokens(base_model: str) -> int | None: + try: + model_info: Final = litellm.get_model_info(model=base_model) + except Exception: # noqa: BLE001 # get_model_info raises ModelNotMappedError or a bare Exception for an unmapped model + return None + if is_generalized_model_info(model_info): + return None + return model_info.get("max_input_tokens") + + +def _provider_model_id(model_name: str, custom_llm_provider: str) -> str: + model_id: Final = model_name.rsplit("/", 1)[-1] + return get_bedrock_base_model(model_id) if custom_llm_provider.startswith("bedrock") else model_id + + +def _names_route_model(base_model: str, route_key: str, custom_llm_provider: str) -> bool: + route_model_id: Final = _provider_model_id(route_key, custom_llm_provider) + base_model_id: Final = _provider_model_id(base_model, custom_llm_provider) + return route_model_id in (base_model_id, base_model_id.split("@", 1)[0]) + + +def _declared_base_model(user_model_info: Mapping[str, object]) -> str | None: + match user_model_info: + case {"base_model": str() as base_model}: + return base_model + case _: + return None + + +def _base_model_widened_window( + route_model_info: ModelMapInfo, base_model: str | None, custom_llm_provider: str +) -> Mapping[str, int]: + route_max_input_tokens: Final = route_model_info.get("max_input_tokens") + if base_model is None or route_max_input_tokens is None: + return MappingProxyType({}) + if not _names_route_model(base_model, route_model_info["key"], custom_llm_provider): + return MappingProxyType({}) + base_max_input_tokens: Final = _exact_base_model_max_input_tokens(base_model) + if base_max_input_tokens is None or base_max_input_tokens <= route_max_input_tokens: + return MappingProxyType({}) + return MappingProxyType({"max_input_tokens": base_max_input_tokens}) + + class Router: @property def _routing_groups(self) -> dict[str, RoutingGroup]: @@ -11343,6 +11388,7 @@ class Router: Augment litellm info with additional params set in `model_info`. For azure models, ignore the `model:`. Only set max tokens, cost values if base_model is set. + For every other provider, a base_model naming the same model as `model:` can only widen its context window. Returns - ModelInfo - If found -> typed dict with max tokens, input cost, etc. @@ -11439,20 +11485,23 @@ class Router: ## CHECK USER SET MODEL INFO raw_user_model_info: Final = deployment.get("model_info") - user_model_info: Final = ( + user_model_info: Final = _MODEL_INFO_ADAPTER.validate_python( raw_user_model_info.model_dump(exclude_none=True) if isinstance(raw_user_model_info, BaseModel) - else raw_user_model_info + else raw_user_model_info or {} ) # get_model_info() hands back an lru_cache'd dict, so merge into a copy; unset # values are skipped or Deployment's None pricing defaults would erase the map's merged_model_info: Final[ModelMapInfo] = { **copy.deepcopy(model_info), - **self.get_discovered_model_info((deployment.get("model_info") or {}).get("id")), - **MappingProxyType( - {key: value for key, value in (user_model_info or MappingProxyType({})).items() if value is not None} + **_base_model_widened_window( + model_info, + _declared_base_model(user_model_info) or _declared_base_model(litellm_params.model_extra or {}), + custom_llm_provider, ), + **self.get_discovered_model_info((deployment.get("model_info") or {}).get("id")), + **MappingProxyType({key: value for key, value in user_model_info.items() if value is not None}), } return merged_model_info diff --git a/tests/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py index ba2913bcfb7..2fd2549c0ac 100644 --- a/tests/code_coverage_tests/router_code_coverage.py +++ b/tests/code_coverage_tests/router_code_coverage.py @@ -114,6 +114,11 @@ ignored_function_names = [ "_deployment_that_just_failed", # Tested through the same-boundary hop test in test_router_order_fallback.py "_collect_ptu_share_teams", # Cached behind team_ptu_ceiling, tested through the live model list test in test_router.py "_resolve_team_ptu_ceiling", # Cached behind team_ptu_ceiling, tested through the live model list test in test_router.py + "_exact_base_model_max_input_tokens", # Tested through get_router_model_info in test_router_model_cost_isolation.py + "_base_model_widened_window", # Tested through get_router_model_info in test_router_model_cost_isolation.py + "_provider_model_id", # Tested through get_router_model_info in test_router_model_cost_isolation.py + "_names_route_model", # Tested through get_router_model_info in test_router_model_cost_isolation.py + "_declared_base_model", # Tested through get_router_model_info in test_router_model_cost_isolation.py ] diff --git a/tests/integration/routing/test_base_model_context_window_owned_proxy.py b/tests/integration/routing/test_base_model_context_window_owned_proxy.py new file mode 100644 index 00000000000..d09ec868c78 --- /dev/null +++ b/tests/integration/routing/test_base_model_context_window_owned_proxy.py @@ -0,0 +1,171 @@ +import json +import uuid +from collections.abc import Iterator, Mapping +from dataclasses import dataclass +from typing import Final + +import pytest +from integration._support.client import gateway_from_environment +from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from integration.providers._cache_control_marks_support import marker_of, owned_config +from pydantic import JsonValue + +_TOKEN: Final = "synthetic-bedrock-bearer" +_KIMI: Final = "moonshotai.kimi-k2.5" +_SONNET: Final = "global.anthropic.claude-sonnet-4-6" +_KIMI_BASE: Final = f"bedrock/us-east-1/{_KIMI}" +_KIMI_TARGET: Final = f"/model/{_KIMI}/converse" +_SONNET_TARGET: Final = f"/model/{_SONNET}/converse" +_PAST_ROUTE_ROW_LINES: Final = 16250 +_PAST_ESCALATION_BUFFER_LINES: Final = 15480 +_SMALL_LINES: Final = 1300 +_REPLY: Final = json.dumps( + { + "output": {"message": {"role": "assistant", "content": [{"text": "8115"}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 9, "outputTokens": 2, "totalTokens": 11}, + "metrics": {"latencyMs": 1}, + } +).encode() + + +@dataclass(frozen=True, slots=True) +class _Rig: + owned: OwnedProxy + wire: Wire + + +def _converse_peer(request: Request) -> Reply: + assert request.headers["authorization"] == f"Bearer {_TOKEN}", request.headers + return Reply(body=_REPLY) + + +def _deployment( + name: str, + model: str, + api_base: str, + *, + model_info: Mapping[str, JsonValue] | None = None, + **params: JsonValue, +) -> dict[str, JsonValue]: + return { + "model_name": name, + "litellm_params": { + "model": f"bedrock/converse/{model}", + "api_key": _TOKEN, + "aws_region_name": "us-east-1", + "aws_bedrock_runtime_endpoint": api_base, + **params, + }, + "model_info": {"id": f"{name}-{uuid.uuid4().hex}", **(model_info or {})}, + } + + +@pytest.fixture(scope="module") +def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]: + directory: Final = tmp_path_factory.mktemp("base-model-window") + with gateway_from_environment() as gateway, wire_server(_converse_peer) as wire: + config: Final = owned_config( + directory, + [ + _deployment("base-in-model-info", _KIMI, wire.url, model_info={"base_model": _KIMI_BASE}), + _deployment("base-in-litellm-params", _KIMI, wire.url, base_model=_KIMI_BASE), + _deployment("no-base", _KIMI, wire.url), + _deployment("other-model-base", _KIMI, wire.url, model_info={"base_model": _SONNET}), + _deployment( + "declared-window", + _KIMI, + wire.url, + model_info={"base_model": _KIMI_BASE, "max_input_tokens": 257000}, + ), + _deployment( + "unmapped-base", _KIMI, wire.url, model_info={"base_model": f"{_KIMI_BASE}-not-in-cost-map"} + ), + _deployment( + "smaller-other-model-base", + _KIMI, + wire.url, + model_info={"base_model": "bedrock/us-east-1/meta.llama3-8b-instruct-v1:0"}, + ), + _deployment("one-million", _SONNET, wire.url), + { + "model_name": "complexity", + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": { + "tiers": { + "SIMPLE": "base-in-model-info", + "MEDIUM": "base-in-model-info", + "COMPLEX": "one-million", + "REASONING": "one-million", + }, + "enable_context_window_escalation": True, + }, + }, + }, + ], + litellm_settings={"cache": False}, + router_settings={"enable_pre_call_checks": True}, + ) + with owned_proxy_process(gateway, directory, {}, config=config, workers=2, client_timeout=120) as owned: + yield _Rig(owned, wire) + + +def _prompt(lines: int, marker: str) -> str: + listing: Final = "\n".join(f"Line {i}: the value of item {i} is {(i * 7) % 1000}." for i in range(1, lines + 1)) + return f"{listing}\n\nWhat is the value of item 12345? Answer with the number only. marker-{marker}" + + +def _request(surface: str, model: str, text: str) -> tuple[str, dict[str, JsonValue]]: + if surface == "messages": + return "/v1/messages", {"model": model, "max_tokens": 16, "messages": [{"role": "user", "content": text}]} + if surface == "responses": + return "/v1/responses", {"model": model, "max_output_tokens": 16, "input": text} + return "/v1/chat/completions", {"model": model, "max_tokens": 16, "messages": [{"role": "user", "content": text}]} + + +def _send(rig: _Rig, surface: str, model: str, lines: int) -> tuple[int, str, tuple[str, ...]]: + marker: Final = uuid.uuid4().hex + path, body = _request(surface, model, _prompt(lines, marker)) + response: Final = rig.owned.gateway.request("POST", path, body) + targets: Final = tuple(request.target for request in rig.wire.drain() if marker_of(request) == marker) + return response.status_code, response.text, targets + + +@pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +@pytest.mark.parametrize("surface", ["chat", "messages", "responses"]) +def test_a_same_model_base_model_row_widens_the_window_pre_call_checks_enforce(rig: _Rig, surface: str) -> None: + status, text, targets = _send(rig, surface, "base-in-model-info", _PAST_ROUTE_ROW_LINES) + assert status == 200, text + assert targets == (_KIMI_TARGET,) + + +@pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +def test_a_same_model_base_model_in_litellm_params_widens_the_window_too(rig: _Rig) -> None: + status, text, targets = _send(rig, "chat", "base-in-litellm-params", _PAST_ROUTE_ROW_LINES) + assert status == 200, text + assert targets == (_KIMI_TARGET,) + + +@pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +@pytest.mark.parametrize("model", ["no-base", "other-model-base", "declared-window", "unmapped-base"]) +def test_a_base_model_that_cannot_widen_keeps_the_route_window_and_calls_no_upstream(rig: _Rig, model: str) -> None: + status, text, targets = _send(rig, "chat", model, _PAST_ROUTE_ROW_LINES) + assert status == 400, text + assert "ContextWindowExceededError" in text, text + assert targets == () + + +@pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +def test_a_smaller_other_model_base_model_never_narrows_the_route_window(rig: _Rig) -> None: + status, text, targets = _send(rig, "chat", "smaller-other-model-base", _SMALL_LINES) + assert status == 200, text + assert targets == (_KIMI_TARGET,) + + +@pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +def test_complexity_router_keeps_a_prompt_on_the_tier_whose_base_model_row_holds_it(rig: _Rig) -> None: + status, text, targets = _send(rig, "chat", "complexity", _PAST_ESCALATION_BUFFER_LINES) + assert status == 200, text + assert targets == (_KIMI_TARGET,) diff --git a/tests/unit/test_router_model_cost_isolation.py b/tests/unit/test_router_model_cost_isolation.py index fc1799861d1..d343f464a77 100644 --- a/tests/unit/test_router_model_cost_isolation.py +++ b/tests/unit/test_router_model_cost_isolation.py @@ -119,27 +119,33 @@ async def test_discovered_limits_survive_deployment_growth_and_removal( async def test_discovery_discards_metadata_for_a_replaced_deployment(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(litellm, "model_cost", copy.deepcopy(litellm.model_cost)) - router: Final = Router(model_list=[{ - "model_name": "local", - "litellm_params": { - "model": "hosted_vllm/local-model", - "api_base": "https://original.test/v1", - "api_key": "local-key", - }, - "model_info": {"id": "replaced-deployment"}, - }]) + router: Final = Router( + model_list=[ + { + "model_name": "local", + "litellm_params": { + "model": "hosted_vllm/local-model", + "api_base": "https://original.test/v1", + "api_key": "local-key", + }, + "model_info": {"id": "replaced-deployment"}, + } + ] + ) def respond(request: httpx.Request) -> httpx.Response: if request.url.host == "original.test": - router.upsert_deployment(Deployment( - model_name="local", - litellm_params=LiteLLM_Params( - model="hosted_vllm/local-model", - api_base="https://replacement.test/v1", - api_key="local-key", - ), - model_info=ModelInfo(id="replaced-deployment"), - )) + router.upsert_deployment( + Deployment( + model_name="local", + litellm_params=LiteLLM_Params( + model="hosted_vllm/local-model", + api_base="https://replacement.test/v1", + api_key="local-key", + ), + model_info=ModelInfo(id="replaced-deployment"), + ) + ) return httpx.Response(200, json={"data": [{"id": "local-model", "max_model_len": 8192}]}) assert request.url.host == "replacement.test" return httpx.Response(200, json={"data": [{"id": "local-model", "max_model_len": 2048}]}) @@ -158,15 +164,19 @@ async def test_discovery_discards_metadata_for_a_replaced_deployment(monkeypatch async def test_discovery_is_isolated_across_routers_and_reused_ids(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(litellm, "model_cost", copy.deepcopy(litellm.model_cost)) first, second = tuple( - Router(model_list=[{ - "model_name": "local", - "litellm_params": { - "model": "hosted_vllm/local-model", - "api_base": f"https://{host}.test/v1", - "api_key": "local-key", - }, - "model_info": {"id": "shared-discovery-id"}, - }]) + Router( + model_list=[ + { + "model_name": "local", + "litellm_params": { + "model": "hosted_vllm/local-model", + "api_base": f"https://{host}.test/v1", + "api_key": "local-key", + }, + "model_info": {"id": "shared-discovery-id"}, + } + ] + ) for host in ("first", "second") ) @@ -187,15 +197,17 @@ async def test_discovery_is_isolated_across_routers_and_reused_ids(monkeypatch: assert first.get_configured_token_limits("local") == (8192, 8192) assert second.get_configured_token_limits("local") == (2048, 2048) assert litellm.model_cost["shared-discovery-id"].get("max_input_tokens") is None - first.upsert_deployment(Deployment( - model_name="local", - litellm_params=LiteLLM_Params( - model="hosted_vllm/local-model", - api_base="https://unavailable.test/v1", - api_key="local-key", - ), - model_info=ModelInfo(id="shared-discovery-id"), - )) + first.upsert_deployment( + Deployment( + model_name="local", + litellm_params=LiteLLM_Params( + model="hosted_vllm/local-model", + api_base="https://unavailable.test/v1", + api_key="local-key", + ), + model_info=ModelInfo(id="shared-discovery-id"), + ) + ) assert first.get_configured_token_limits("local") == (None, None) await first.arefresh_model_info(client=handler) assert first.get_configured_token_limits("local") == (None, None) @@ -206,17 +218,19 @@ async def test_discovery_is_isolated_across_routers_and_reused_ids(monkeypatch: async def test_discovery_refreshes_other_endpoints_while_one_is_pending(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(litellm, "model_cost", copy.deepcopy(litellm.model_cost)) second_started: Final = asyncio.Event() - router: Final = Router(model_list=[ - { - "model_name": host, - "litellm_params": { - "model": "hosted_vllm/local-model", - "api_base": f"https://{host}.test/v1", - "api_key": "local-key", - }, - } - for host in ("first", "second", "third") - ]) + router: Final = Router( + model_list=[ + { + "model_name": host, + "litellm_params": { + "model": "hosted_vllm/local-model", + "api_base": f"https://{host}.test/v1", + "api_key": "local-key", + }, + } + for host in ("first", "second", "third") + ] + ) async def respond(request: httpx.Request) -> httpx.Response: if request.url.host == "first.test": @@ -240,21 +254,27 @@ async def test_discovery_refreshes_other_endpoints_while_one_is_pending(monkeypa async def test_discovered_limits_expire_after_the_last_successful_refresh(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(litellm, "model_cost", copy.deepcopy(litellm.model_cost)) clock: Final = Mock(return_value=0.0) - router: Final = Router(model_list=[{ - "model_name": "local", - "litellm_params": { - "model": "hosted_vllm/local-model", - "api_base": "https://expiry.test/v1", - "api_key": "local-key", - }, - "model_info": {"id": "expiring-discovery"}, - }]) + router: Final = Router( + model_list=[ + { + "model_name": "local", + "litellm_params": { + "model": "hosted_vllm/local-model", + "api_base": "https://expiry.test/v1", + "api_key": "local-key", + }, + "model_info": {"id": "expiring-discovery"}, + } + ] + ) router._discovered_model_info_cache = InMemoryCache(clock=clock, default_ttl=2 * MODEL_INFO_REFRESH_SECONDS) - responses: Final = iter(( - httpx.Response(200, json={"data": [{"id": "local-model", "max_model_len": 4096}]}), - httpx.Response(200, json={"data": [{"id": "local-model", "max_model_len": 8192}]}), - httpx.Response(503), - )) + responses: Final = iter( + ( + httpx.Response(200, json={"data": [{"id": "local-model", "max_model_len": 4096}]}), + httpx.Response(200, json={"data": [{"id": "local-model", "max_model_len": 8192}]}), + httpx.Response(503), + ) + ) handler: Final = AsyncHTTPHandler() await handler.client.aclose() async with httpx.AsyncClient(transport=httpx.MockTransport(lambda request: next(responses))) as client: @@ -343,10 +363,12 @@ async def test_discovered_limits_are_isolated_overridable_and_refreshable( async def test_discovery_preserves_input_overrides_and_survives_outages(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(litellm, "model_cost", copy.deepcopy(litellm.model_cost)) - responses: Final = iter(( - httpx.Response(200, json={"data": [{"id": "local-model", "max_model_len": 4096}]}), - httpx.Response(503), - )) + responses: Final = iter( + ( + httpx.Response(200, json={"data": [{"id": "local-model", "max_model_len": 4096}]}), + httpx.Response(503), + ) + ) def respond(request: httpx.Request) -> httpx.Response: assert request.url.host == "backend.test" @@ -354,27 +376,29 @@ async def test_discovery_preserves_input_overrides_and_survives_outages(monkeypa assert request.headers["x-tenant"] == "tenant" return next(responses) - router: Final = Router(model_list=[ - { - "model_name": "configured", - "litellm_params": { - "model": "hosted_vllm/local-model", - "api_base": "https://backend.test/v1", - "api_key": "unused-key", - "extra_headers": {"authorization": "Bearer local-key", "X-Tenant": "tenant"}, + router: Final = Router( + model_list=[ + { + "model_name": "configured", + "litellm_params": { + "model": "hosted_vllm/local-model", + "api_base": "https://backend.test/v1", + "api_key": "unused-key", + "extra_headers": {"authorization": "Bearer local-key", "X-Tenant": "tenant"}, + }, + "model_info": {"id": "configured", "max_input_tokens": 1024}, }, - "model_info": {"id": "configured", "max_input_tokens": 1024}, - }, - { - "model_name": "byok", - "litellm_params": { - "model": "openai/local-model", - "api_base": "https://caller.test/v1", - "use_clientside_credentials": True, + { + "model_name": "byok", + "litellm_params": { + "model": "openai/local-model", + "api_base": "https://caller.test/v1", + "use_clientside_credentials": True, + }, }, - }, - {"model_name": "default-openai", "litellm_params": {"model": "openai/local-model", "api_key": "unused"}}, - ]) + {"model_name": "default-openai", "litellm_params": {"model": "openai/local-model", "api_key": "unused"}}, + ] + ) handler: Final = AsyncHTTPHandler() await handler.client.aclose() responder: Final = Mock(side_effect=respond) @@ -944,8 +968,7 @@ def _register_azure_tier_backend() -> None: def test_inherit_builtin_service_tier_pricing_fills_only_missing_fields() -> None: model_cost_entries: Final = { - key: copy.deepcopy(litellm.model_cost.get(key)) - for key in (_TIER_BACKEND_KEY, _TIER_BACKEND_MODEL) + key: copy.deepcopy(litellm.model_cost.get(key)) for key in (_TIER_BACKEND_KEY, _TIER_BACKEND_MODEL) } try: _register_tier_backend() @@ -974,9 +997,7 @@ def test_inherit_builtin_service_tier_pricing_fills_only_missing_fields() -> Non "output_cost_per_token_flex": _TIER_BACKEND_ENTRY["output_cost_per_token_flex"], "input_cost_per_token_balanced": _TIER_BACKEND_ENTRY["input_cost_per_token_balanced"], "output_cost_per_token_balanced": _TIER_BACKEND_ENTRY["output_cost_per_token_balanced"], - "cache_read_input_token_cost_ultrafast": _TIER_BACKEND_ENTRY[ - "cache_read_input_token_cost_ultrafast" - ], + "cache_read_input_token_cost_ultrafast": _TIER_BACKEND_ENTRY["cache_read_input_token_cost_ultrafast"], "input_cost_per_token_above_272k_tokens_ultrafast": _TIER_BACKEND_ENTRY[ "input_cost_per_token_above_272k_tokens_ultrafast" ], @@ -991,8 +1012,7 @@ def test_inherit_builtin_service_tier_pricing_fills_only_missing_fields() -> Non def test_inherit_builtin_service_tier_pricing_noop_without_base_rate_or_backend() -> None: model_cost_entries: Final = { - key: copy.deepcopy(litellm.model_cost.get(key)) - for key in (_TIER_BACKEND_KEY, _TIER_BACKEND_MODEL) + key: copy.deepcopy(litellm.model_cost.get(key)) for key in (_TIER_BACKEND_KEY, _TIER_BACKEND_MODEL) } try: _register_tier_backend() @@ -1029,8 +1049,7 @@ def test_inherit_builtin_service_tier_pricing_noop_without_base_rate_or_backend( def test_router_completion_uses_custom_standard_and_backend_ultrafast_pricing() -> None: model_id: Final = "tier-priced-deployment" model_cost_entries: Final = { - key: copy.deepcopy(litellm.model_cost.get(key)) - for key in (_TIER_BACKEND_KEY, _TIER_BACKEND_MODEL, model_id) + key: copy.deepcopy(litellm.model_cost.get(key)) for key in (_TIER_BACKEND_KEY, _TIER_BACKEND_MODEL, model_id) } try: _register_tier_backend() @@ -1090,8 +1109,7 @@ def test_router_completion_uses_custom_standard_and_backend_ultrafast_pricing() def test_router_completion_uses_backend_ultrafast_long_context_rates() -> None: model_id: Final = "tier-priced-long-context-deployment" model_cost_entries: Final = { - key: copy.deepcopy(litellm.model_cost.get(key)) - for key in (_TIER_BACKEND_KEY, _TIER_BACKEND_MODEL, model_id) + key: copy.deepcopy(litellm.model_cost.get(key)) for key in (_TIER_BACKEND_KEY, _TIER_BACKEND_MODEL, model_id) } try: _register_tier_backend() @@ -1142,8 +1160,7 @@ def test_ptu_service_tier_pricing_is_disabled_only_when_attribution_is_enabled( ) -> None: model_id: Final = f"ptu-tier-deployment-{ptu_enabled}" model_cost_entries: Final = { - key: copy.deepcopy(litellm.model_cost.get(key)) - for key in (_TIER_BACKEND_KEY, model_id) + key: copy.deepcopy(litellm.model_cost.get(key)) for key in (_TIER_BACKEND_KEY, model_id) } try: _register_tier_backend() @@ -1200,8 +1217,7 @@ def test_azure_base_model_inherits_service_tier_pricing_for_registration_and_pay model_id: Final = "azure-tier-priced-alias" payload_id: Final = "azure-tier-priced-payload" model_cost_entries: Final = { - key: copy.deepcopy(litellm.model_cost.get(key)) - for key in (_AZURE_TIER_BACKEND_KEY, model_id, payload_id) + key: copy.deepcopy(litellm.model_cost.get(key)) for key in (_AZURE_TIER_BACKEND_KEY, model_id, payload_id) } try: _register_azure_tier_backend() @@ -1257,12 +1273,8 @@ def test_azure_base_model_inherits_service_tier_pricing_for_registration_and_pay ) ) - assert payload["input_cost_per_token_priority"] == _AZURE_TIER_BACKEND_ENTRY[ - "input_cost_per_token_priority" - ] - assert payload["output_cost_per_token_priority"] == _AZURE_TIER_BACKEND_ENTRY[ - "output_cost_per_token_priority" - ] + assert payload["input_cost_per_token_priority"] == _AZURE_TIER_BACKEND_ENTRY["input_cost_per_token_priority"] + assert payload["output_cost_per_token_priority"] == _AZURE_TIER_BACKEND_ENTRY["output_cost_per_token_priority"] finally: _restore_model_cost_entries(model_cost_entries) litellm.get_model_info.cache_clear() @@ -1364,12 +1376,10 @@ def test_inherit_builtin_base_rates_for_off_peak_carries_threshold_rates(): assert model_info["input_cost_per_token"] == builtin_info["input_cost_per_token"] assert ( - model_info["input_cost_per_token_above_200k_tokens"] - == builtin_info["input_cost_per_token_above_200k_tokens"] + model_info["input_cost_per_token_above_200k_tokens"] == builtin_info["input_cost_per_token_above_200k_tokens"] ) assert ( - model_info["output_cost_per_token_above_200k_tokens"] - == builtin_info["output_cost_per_token_above_200k_tokens"] + model_info["output_cost_per_token_above_200k_tokens"] == builtin_info["output_cost_per_token_above_200k_tokens"] ) @@ -2260,8 +2270,7 @@ def test_deployment_model_cost_payload_folds_in_litellm_params_pricing(): def test_deployment_model_cost_payload_includes_builtin_service_tier_pricing() -> None: model_id: Final = "tier-priced-payload" model_cost_entries: Final = { - key: copy.deepcopy(litellm.model_cost.get(key)) - for key in (_TIER_BACKEND_KEY, _TIER_BACKEND_MODEL, model_id) + key: copy.deepcopy(litellm.model_cost.get(key)) for key in (_TIER_BACKEND_KEY, _TIER_BACKEND_MODEL, model_id) } try: _register_tier_backend() @@ -2278,12 +2287,8 @@ def test_deployment_model_cost_payload_includes_builtin_service_tier_pricing() - ) ) - assert ( - payload["input_cost_per_token_ultrafast"] == _TIER_BACKEND_ENTRY["input_cost_per_token_ultrafast"] - ) - assert ( - payload["output_cost_per_token_ultrafast"] == _TIER_BACKEND_ENTRY["output_cost_per_token_ultrafast"] - ) + assert payload["input_cost_per_token_ultrafast"] == _TIER_BACKEND_ENTRY["input_cost_per_token_ultrafast"] + assert payload["output_cost_per_token_ultrafast"] == _TIER_BACKEND_ENTRY["output_cost_per_token_ultrafast"] assert payload["input_cost_per_token_balanced"] == _TIER_BACKEND_ENTRY["input_cost_per_token_balanced"] assert payload["input_cost_per_token"] == _CUSTOM_STANDARD_INPUT_RATE assert payload["output_cost_per_token"] == _CUSTOM_STANDARD_OUTPUT_RATE @@ -3282,9 +3287,7 @@ def test_deployment_id_colliding_with_another_providers_catalog_key_keeps_its_ow ) assert isinstance(response, litellm.ModelResponse) - assert response._hidden_params["response_cost"] == pytest.approx( - 1000 * 0.00000096 + 500 * 0.00000302 - ) + assert response._hidden_params["response_cost"] == pytest.approx(1000 * 0.00000096 + 500 * 0.00000302) assert litellm.model_cost[_COLLIDING_BUILTIN_KEY] == builtin_row_before finally: _restore_model_cost_entries(model_cost_entries) @@ -3297,8 +3300,7 @@ def test_deployment_id_matching_its_own_providers_catalog_key_still_merges() -> matches against.""" model_id: Final = "baseten/zai-org/glm-5.2" model_cost_entries: Final = { - key: copy.deepcopy(litellm.model_cost.get(key)) - for key in (model_id, _COLLIDING_BUILTIN_KEY) + key: copy.deepcopy(litellm.model_cost.get(key)) for key in (model_id, _COLLIDING_BUILTIN_KEY) } try: router: Final = Router(model_list=[_colliding_id_deployment(model_id, "baseten")]) @@ -3313,9 +3315,244 @@ def test_deployment_id_matching_its_own_providers_catalog_key_still_merges() -> ) assert isinstance(response, litellm.ModelResponse) - assert response._hidden_params["response_cost"] == pytest.approx( - 1000 * 0.00000096 + 500 * 0.00000302 - ) + assert response._hidden_params["response_cost"] == pytest.approx(1000 * 0.00000096 + 500 * 0.00000302) finally: _restore_model_cost_entries(model_cost_entries) litellm.get_model_info.cache_clear() + + +_CAPACITY_TEST_DEPLOYMENTS: Final = ( + ( + "bedrock", + "bedrock/converse/anthropic.claude-router-capacity-test", + "bedrock/us-east-1/anthropic.claude-router-capacity-test", + ), + ( + "bedrock_mantle", + "bedrock_mantle/openai.router-capacity-test", + "bedrock_mantle/us-east-1/openai.router-capacity-test", + ), + ("openai", "openai/router-capacity-test", "openai/us/router-capacity-test"), +) + + +def _capacity_test_metadata(provider: str, max_input_tokens: int) -> dict[str, object]: + return { + "litellm_provider": provider, + "mode": "chat", + "max_input_tokens": max_input_tokens, + "max_output_tokens": 512, + "max_tokens": 512, + "input_cost_per_token": 0, + "output_cost_per_token": 0, + } + + +def _capacity_test_router(route: str, base: str, base_location: str | None, explicit_limit: int | None) -> Router: + return Router( + model_list=[ + { + "model_name": "capacity-test", + "litellm_params": { + "model": route, + **({"base_model": base} if base_location == "litellm_params" else {}), + }, + "model_info": { + "id": "capacity-test", + **({"base_model": base} if base_location == "model_info" else {}), + **({"max_input_tokens": explicit_limit} if explicit_limit is not None else {}), + }, + } + ], + enable_pre_call_checks=True, + ) + + +def _deployments_allowed_for_2048_input_tokens(router: Router) -> int: + try: + return len( + router._pre_call_checks( + model="capacity-test", + healthy_deployments=router.model_list, + input="synthetic input", + input_token_count=2048, + ) + ) + except litellm.ContextWindowExceededError: + return 0 + + +def _effective_max_input_tokens(router: Router) -> object: + return router.get_router_model_info(id="capacity-test", deployment=None, received_model_name="capacity-test")[ + "max_input_tokens" + ] + + +@pytest.mark.parametrize(("provider", "route", "base"), _CAPACITY_TEST_DEPLOYMENTS) +@pytest.mark.parametrize("base_location", ("model_info", "litellm_params", None)) +@pytest.mark.parametrize("explicit_limit", (None, 4096)) +def test_router_capacity_widens_to_a_larger_base_model_window( + monkeypatch: pytest.MonkeyPatch, + provider: str, + route: str, + base: str, + base_location: str | None, + explicit_limit: int | None, +) -> None: + monkeypatch.setattr( + litellm, + "model_cost", + { + **litellm.model_cost, + route: _capacity_test_metadata(provider, 1024), + base: _capacity_test_metadata(provider, 8192), + }, + ) + _invalidate_model_cost_lowercase_map() + router: Final = _capacity_test_router(route, base, base_location, explicit_limit) + expected: Final = ( + explicit_limit if explicit_limit is not None else (1024 if base_location is None else 8192) + ) + assert _effective_max_input_tokens(router) == expected + assert _deployments_allowed_for_2048_input_tokens(router) == (1 if expected >= 2048 else 0) + _invalidate_model_cost_lowercase_map() + + +@pytest.mark.parametrize(("provider", "route", "base"), _CAPACITY_TEST_DEPLOYMENTS) +@pytest.mark.parametrize("base_location", ("model_info", "litellm_params")) +def test_router_capacity_never_narrows_to_a_smaller_base_model_window( + monkeypatch: pytest.MonkeyPatch, provider: str, route: str, base: str, base_location: str +) -> None: + monkeypatch.setattr( + litellm, + "model_cost", + { + **litellm.model_cost, + route: _capacity_test_metadata(provider, 8192), + base: _capacity_test_metadata(provider, 1024), + }, + ) + _invalidate_model_cost_lowercase_map() + router: Final = _capacity_test_router(route, base, base_location, None) + assert _effective_max_input_tokens(router) == 8192 + assert _deployments_allowed_for_2048_input_tokens(router) == 1 + _invalidate_model_cost_lowercase_map() + + +@pytest.mark.parametrize(("provider", "route", "base"), _CAPACITY_TEST_DEPLOYMENTS) +def test_router_capacity_never_takes_a_window_from_base_model_when_the_route_is_unmapped( + monkeypatch: pytest.MonkeyPatch, provider: str, route: str, base: str +) -> None: + monkeypatch.setattr(litellm, "model_cost", {**litellm.model_cost, base: _capacity_test_metadata(provider, 1024)}) + _invalidate_model_cost_lowercase_map() + router: Final = _capacity_test_router(f"{route}-not-in-cost-map", base, "model_info", None) + assert _deployments_allowed_for_2048_input_tokens(router) == 1 + _invalidate_model_cost_lowercase_map() + + +@pytest.mark.parametrize(("provider", "route", "base"), _CAPACITY_TEST_DEPLOYMENTS) +def test_router_capacity_never_takes_a_window_from_base_model_when_the_route_declares_none( + monkeypatch: pytest.MonkeyPatch, provider: str, route: str, base: str +) -> None: + route_without_window: Final = { + key: value + for key, value in _capacity_test_metadata(provider, 0).items() + if key not in ("max_input_tokens", "max_tokens") + } + monkeypatch.setattr( + litellm, + "model_cost", + {**litellm.model_cost, route: route_without_window, base: _capacity_test_metadata(provider, 1024)}, + ) + _invalidate_model_cost_lowercase_map() + router: Final = _capacity_test_router(route, base, "model_info", None) + assert _effective_max_input_tokens(router) is None + assert _deployments_allowed_for_2048_input_tokens(router) == 1 + _invalidate_model_cost_lowercase_map() + + +_SAME_MODEL_ROUTE: Final = "bedrock/converse/anthropic.claude-routercapacitytest-9-9" +_SAME_MODEL_GLOBAL_ROUTE: Final = "bedrock/converse/global.anthropic.claude-routercapacitytest-9-9" +_SAME_MODEL_QUALIFIED_ROUTE: Final = "bedrock/converse/anthropic.claude-routercapacitytest-9-9@first" + + +@pytest.mark.parametrize( + ("route", "base_model", "expected"), + ( + (_SAME_MODEL_ROUTE, "bedrock/us-east-1/anthropic.claude-routercapacitytest-9-9", 8192), + (_SAME_MODEL_ROUTE, "us-east-1/anthropic.claude-routercapacitytest-9-9", 1024), + (_SAME_MODEL_ROUTE, "anthropic.claude-routercapacitytest-9-9", 1024), + (_SAME_MODEL_GLOBAL_ROUTE, "bedrock/us-east-1/anthropic.claude-routercapacitytest-9-9", 8192), + (_SAME_MODEL_ROUTE, "bedrock/anthropic.claude-routercapacitytest-9-9@global:us-east-1", 8192), + (_SAME_MODEL_GLOBAL_ROUTE, "bedrock/anthropic.claude-routercapacitytest-9-9@global:us-east-1", 8192), + (_SAME_MODEL_QUALIFIED_ROUTE, "bedrock/us-east-1/anthropic.claude-routercapacitytest-9-9@first", 8192), + (_SAME_MODEL_QUALIFIED_ROUTE, "bedrock/us-east-1/anthropic.claude-routercapacitytest-9-9@second", 1024), + (_SAME_MODEL_ROUTE, "bedrock/us-east-1/anthropic.claude-routercapacitytest-*", 1024), + (_SAME_MODEL_ROUTE, "bedrock/us-east-1/anthropic.claude-routercapacitystandin-9-9", 1024), + (_SAME_MODEL_ROUTE, "bedrock/us-east-1/anthropic.claude-routercapacitytest-9-9-extended", 1024), + (_SAME_MODEL_ROUTE, "bedrock/eu-west-1/anthropic.claude-routercapacitytest-9-9", 1024), + ), +) +def test_router_capacity_widens_only_for_an_exact_base_model_entry_of_the_same_model( + monkeypatch: pytest.MonkeyPatch, route: str, base_model: str, expected: int +) -> None: + larger_window: Final = _capacity_test_metadata("bedrock", 8192) + monkeypatch.setattr( + litellm, + "model_cost", + { + **litellm.model_cost, + _SAME_MODEL_ROUTE: _capacity_test_metadata("bedrock", 1024), + _SAME_MODEL_GLOBAL_ROUTE: _capacity_test_metadata("bedrock", 1024), + _SAME_MODEL_QUALIFIED_ROUTE: _capacity_test_metadata("bedrock", 1024), + "bedrock/us-east-1/anthropic.claude-routercapacitytest-9-9@first": larger_window, + "bedrock/us-east-1/anthropic.claude-routercapacitytest-9-9@second": larger_window, + "bedrock/us-east-1/anthropic.claude-routercapacitytest-9-9": larger_window, + "anthropic.claude-routercapacitytest-9-9": _capacity_test_metadata("bedrock", 1024), + "bedrock/anthropic.claude-routercapacitytest-9-9": larger_window, + "bedrock/anthropic.claude-routercapacitytest-9-9@global:us-east-1": larger_window, + "bedrock/us-east-1/anthropic.claude-routercapacitytest-*": larger_window, + "bedrock/us-east-1/anthropic.claude-routercapacitystandin-9-9": larger_window, + "bedrock/us-east-1/anthropic.claude-routercapacitytest-9-9-extended": larger_window, + }, + ) + _invalidate_model_cost_lowercase_map() + router: Final = _capacity_test_router(route, base_model, "model_info", None) + assert _effective_max_input_tokens(router) == expected + assert _deployments_allowed_for_2048_input_tokens(router) == (1 if expected >= 2048 else 0) + _invalidate_model_cost_lowercase_map() + + +def test_router_capacity_reads_a_model_info_base_model_before_a_litellm_params_one( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + litellm, + "model_cost", + { + **litellm.model_cost, + _SAME_MODEL_ROUTE: _capacity_test_metadata("bedrock", 1024), + "bedrock/us-east-1/anthropic.claude-routercapacitytest-9-9": _capacity_test_metadata("bedrock", 8192), + "bedrock/us-east-1/anthropic.claude-routercapacitystandin-9-9": _capacity_test_metadata("bedrock", 8192), + }, + ) + _invalidate_model_cost_lowercase_map() + router: Final = Router( + model_list=[ + { + "model_name": "capacity-test", + "litellm_params": { + "model": _SAME_MODEL_ROUTE, + "base_model": "bedrock/us-east-1/anthropic.claude-routercapacitytest-9-9", + }, + "model_info": { + "id": "capacity-test", + "base_model": "bedrock/us-east-1/anthropic.claude-routercapacitystandin-9-9", + }, + } + ], + enable_pre_call_checks=True, + ) + assert _effective_max_input_tokens(router) == 1024 + assert _deployments_allowed_for_2048_input_tokens(router) == 0 + _invalidate_model_cost_lowercase_map() From f0724983044aabad7e7be46883a21312d95f122d Mon Sep 17 00:00:00 2001 From: David Bochenski Date: Fri, 16 Oct 2026 07:37:49 +0100 Subject: [PATCH 04/23] fix(bedrock): grant sts:GetWebIdentityToken and sts:TagGetWebIdentityToken in web identity session policy for claude_platform (#46670) * fix(bedrock): grant sts:GetWebIdentityToken in web identity session policy for claude_platform Claude Platform on AWS calls STS GetWebIdentityToken (outbound identity federation) with the caller's session before serving inference. The session policy passed to AssumeRoleWithWebIdentity is a permission ceiling and never granted it, so every bedrock/claude_platform request using OIDC web identity auth failed with 403 "no session policy allows the sts:GetWebIdentityToken action", even when the role's identity policy allowed it. Add a ClaudePlatformOutboundIdentityLiteLLM statement granting only sts:GetWebIdentityToken, plus tests that it is granted, scoped to its own statement, and that no sts:* wildcard is introduced. Co-Authored-By: Claude Opus 5.5 * fix(bedrock): also grant sts:TagGetWebIdentityToken for claude_platform web identity Claude Platform on AWS tags the outbound web identity token it requests, so the caller also needs sts:TagGetWebIdentityToken. Found end-to-end on a deployment: with only GetWebIdentityToken granted, calls failed with "not authorized to perform: sts:TagGetWebIdentityToken on resource: arn:aws:sts:::self"; granting both made /v1/chat/completions and /v1/messages (streamed and not) succeed. Co-Authored-By: Claude Opus 5.5 * fix(bedrock): only let Claude Platform use the outbound identity STS grant Scope the ClaudePlatformOutboundIdentityLiteLLM session policy statement with aws:CalledViaLast = aws-external-anthropic.amazonaws.com, so a web identity session can no longer call sts:GetWebIdentityToken or sts:TagGetWebIdentityToken directly. Drop the code comment and fix the stale test messages. * test(bedrock): fail when the Claude Platform caller condition reaches other statements --------- Co-authored-by: Claude Opus 5.5 Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/llms/bedrock/base_aws_llm.py | 28 +++++++++-- .../test_web_identity_session_policy.py | 47 +++++++++++++++++++ 2 files changed, 70 insertions(+), 5 deletions(-) diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index 848120150cd..f1426aa9375 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -99,11 +99,16 @@ def _assume_role_params( ) -_SecureTransportBool = TypedDict("_SecureTransportBool", {"aws:SecureTransport": ReadOnly[Literal["true"]]}) +_SecureTransportBool: Final = TypedDict("_SecureTransportBool", {"aws:SecureTransport": ReadOnly[Literal["true"]]}) + +_CalledViaLastEquals: Final = TypedDict( + "_CalledViaLastEquals", {"aws:CalledViaLast": ReadOnly[Literal["aws-external-anthropic.amazonaws.com"]]} +) -class _SecureTransportCondition(TypedDict): +class _SessionPolicyCondition(TypedDict): Bool: ReadOnly[_SecureTransportBool] + StringEquals: NotRequired[ReadOnly[_CalledViaLastEquals]] class _SessionPolicyStatement(TypedDict): @@ -111,7 +116,7 @@ class _SessionPolicyStatement(TypedDict): Effect: ReadOnly[Literal["Allow"]] Action: ReadOnly[tuple[str, ...]] Resource: ReadOnly[Literal["*"]] - Condition: ReadOnly[_SecureTransportCondition] + Condition: ReadOnly[_SessionPolicyCondition] class WebIdentitySessionPolicy(TypedDict): @@ -147,11 +152,24 @@ _WEB_IDENTITY_SESSION_POLICY_ACTIONS: Final[Mapping[str, tuple[str, ...]]] = Map "aws-external-anthropic:Get*", "aws-external-anthropic:List*", ), + "ClaudePlatformOutboundIdentityLiteLLM": ( + "sts:GetWebIdentityToken", + "sts:TagGetWebIdentityToken", + ), "BedrockMantleLiteLLM": ("bedrock-mantle:CreateInference",), } ) -_SECURE_TRANSPORT_ONLY: Final = _SecureTransportCondition(Bool=_SecureTransportBool({"aws:SecureTransport": "true"})) +_SECURE_TRANSPORT_ONLY: Final = _SessionPolicyCondition(Bool=_SecureTransportBool({"aws:SecureTransport": "true"})) + +_WEB_IDENTITY_SESSION_POLICY_CONDITIONS: Final[Mapping[str, _SessionPolicyCondition]] = MappingProxyType( + { + "ClaudePlatformOutboundIdentityLiteLLM": _SessionPolicyCondition( + Bool=_SecureTransportBool({"aws:SecureTransport": "true"}), + StringEquals=_CalledViaLastEquals({"aws:CalledViaLast": "aws-external-anthropic.amazonaws.com"}), + ), + } +) def build_web_identity_session_policy() -> WebIdentitySessionPolicy: @@ -163,7 +181,7 @@ def build_web_identity_session_policy() -> WebIdentitySessionPolicy: Effect="Allow", Action=actions, Resource="*", - Condition=_SECURE_TRANSPORT_ONLY, + Condition=_WEB_IDENTITY_SESSION_POLICY_CONDITIONS.get(sid, _SECURE_TRANSPORT_ONLY), ) for sid, actions in _WEB_IDENTITY_SESSION_POLICY_ACTIONS.items() ), diff --git a/tests/unit/llms/bedrock/test_web_identity_session_policy.py b/tests/unit/llms/bedrock/test_web_identity_session_policy.py index cbb69b4ceed..332e1258e9c 100644 --- a/tests/unit/llms/bedrock/test_web_identity_session_policy.py +++ b/tests/unit/llms/bedrock/test_web_identity_session_policy.py @@ -176,6 +176,44 @@ class TestClaudePlatformActionsCovered: ) +class TestClaudePlatformOutboundIdentityCovered: + """Claude Platform on AWS calls ``sts:GetWebIdentityToken`` with the + caller's session before serving inference, and tags the token it + requests (``sts:TagGetWebIdentityToken``). Without either in the ceiling, + every claude_platform request via OIDC auth denies with "no session + policy allows the sts: action" even when the role's identity + policy grants it.""" + + @pytest.mark.parametrize("action", ["sts:GetWebIdentityToken", "sts:TagGetWebIdentityToken"]) + def test_outbound_identity_action_present(self, action: str): + assert action in _granted_actions(_captured_policy()), ( + f"{action} missing from session policy, so bedrock/claude_platform/* requests will 403 on OIDC auth" + ) + + @pytest.mark.parametrize("action", ["sts:GetWebIdentityToken", "sts:TagGetWebIdentityToken"]) + def test_outbound_identity_action_is_granted_only_via_claude_platform(self, action: str): + granting: Final = tuple(stmt for stmt in _captured_policy()["Statement"] if action in stmt["Action"]) + assert granting, f"{action} missing from session policy" + assert all( + stmt["Condition"].get("StringEquals") == {"aws:CalledViaLast": "aws-external-anthropic.amazonaws.com"} + and stmt["Condition"]["Bool"] == {"aws:SecureTransport": "true"} + for stmt in granting + ), f"{action} must only be usable over TLS when Claude Platform on AWS calls STS for the session" + + def test_outbound_identity_statement_is_scoped(self): + policy = _captured_policy() + stmt = _statement_by_sid(policy, "ClaudePlatformOutboundIdentityLiteLLM") + assert stmt["Effect"] == "Allow" + assert tuple(stmt["Action"]) == ("sts:GetWebIdentityToken", "sts:TagGetWebIdentityToken"), ( + "the outbound identity statement must grant only the two outbound identity actions" + ) + + def test_no_sts_wildcard(self): + assert "sts:*" not in _granted_actions(_captured_policy()), ( + "session policy must not grant sts:*, only the two outbound identity actions are needed" + ) + + class TestBedrockMantleActionsCovered: """LIT-3859: bedrock_mantle inference authorizes against the ``bedrock-mantle`` action namespace, so the session-policy ceiling @@ -306,6 +344,15 @@ class TestPolicyTransportConditions: "to keep parity with the bedrock statement" ) + def test_only_the_outbound_identity_statement_requires_a_claude_platform_caller(self): + directly_called: Final = tuple( + stmt for stmt in _captured_policy()["Statement"] if stmt["Sid"] != "ClaudePlatformOutboundIdentityLiteLLM" + ) + assert directly_called + assert all("StringEquals" not in stmt["Condition"] for stmt in directly_called), ( + "statements LiteLLM calls directly must not require aws:CalledViaLast, or OIDC auth 403s on those routes" + ) + _STS_SESSION_POLICY_PLAINTEXT_LIMIT: Final = 2048 From b21e659c24720946ed609c0669142aec9d0f2422 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 16 Oct 2026 07:11:39 +0000 Subject: [PATCH 05/23] fix(proxy)!: count endpoints return errors in their source API shape instead of silent local counts (#46815) * fix(count_tokens): convert Anthropic-format messages for the OpenAI counter the way inference does * test(count_tokens): check /v1/messages/count_tokens sends OpenAI the input /v1/messages inference sends * fix(count_tokens): build the Anthropic count request without a cast and sync dashboard API types * fix(count_tokens): count Anthropic-format requests for Gemini the way /v1/messages inference converts them * fix(count_tokens): return a 400 naming Anthropic blocks the OpenAI and Gemini count can't convert * test(count_tokens): compare tools by value in the public api test * fix(count_tokens): reject unconvertible blocks nested inside a tool_result * fix(count_tokens): count Anthropic bodies the way /v1/messages inference sends them to OpenAI Unwrap affinity-encoded reasoning, count the structured-output schema, flatten replayed web search results, join a list system prompt, and answer a message without a role with a 400 before calling OpenAI. messages_format is now internal to the Anthropic count route instead of a public /utils/token_counter field. * test(count_tokens): expect the image detail the adapter now sends after merging main * fix(count_tokens): fetch Gemini remote media without blocking the event loop * fix(count_tokens): type normalize_native_messages without the unimported Any On Python 3.10 to 3.13 the annotation is evaluated at import, so the proxy failed to start with NameError * chore(count_tokens): give the Gemini count helpers' list returns a reason * fix(count_tokens): normalize Anthropic messages for Gemini counts the way /v1/messages does * fix(count_tokens): count Gemini bodies after inference's history rewrites and answer converter errors with 400 * test(count_tokens): cover the Gemini count error paths and compaction with a text-block system * fix(proxy): count endpoints return missing-key and request errors in their source API shape * fix(proxy): drop imports the count error change left unused * fix(proxy): let Entra-authenticated Azure AI Claude count tokens and keep Google count errors in the Google envelope * fix(proxy): send Azure environment keys as x-api-key on Azure AI Claude counts and log unexpected Google count failures * fix(proxy): count with the global OpenAI and Anthropic keys chat already uses * fix(proxy): resolve the OpenAI count base with the key so a global key stays on its gateway * fix(proxy): keep unexpected Google count failures generic so internal details stay out of the response * fix(proxy): resolve the Azure AI Claude count base from the same settings chat reads * fix(proxy): count OpenAI tokens with the workload identity bearer chat already uses * fix(proxy): mint the OpenAI workload identity bearer inside the count error handling * fix(proxy): report a failed OpenAI workload identity exchange on the count endpoints as a 401 --------- Co-authored-by: Mrinal Chanshetty --- .../gitlab/gitlab_prompt_manager.py | 2 +- .../anthropic/count_tokens/token_counter.py | 17 +- .../anthropic/count_tokens/handler.py | 11 +- .../anthropic/count_tokens/token_counter.py | 33 ++- .../anthropic/count_tokens/transformation.py | 9 +- .../responses/count_tokens/token_counter.py | 62 +++++- .../proxy/anthropic_endpoints/endpoints.py | 11 +- litellm/proxy/google_endpoints/endpoints.py | 71 +++++- .../test_count_tokens_error_shapes_wire.py | 205 ++++++++++++++++++ .../llms/anthropic/test_count_tokens_oauth.py | 24 ++ .../count_tokens/test_token_counter.py | 24 ++ .../count_tokens/test_token_counter.py | 41 ++++ .../openai/test_openai_workload_identity.py | 53 +++++ .../anthropic_endpoints/test_endpoints.py | 7 +- .../test_google_api_endpoints.py | 23 ++ tests/unit/proxy/test_proxy_token_counter.py | 45 ++-- 16 files changed, 563 insertions(+), 75 deletions(-) create mode 100644 tests/integration/providers/test_count_tokens_error_shapes_wire.py diff --git a/litellm/integrations/gitlab/gitlab_prompt_manager.py b/litellm/integrations/gitlab/gitlab_prompt_manager.py index 85b91821f6c..adaacd615a8 100644 --- a/litellm/integrations/gitlab/gitlab_prompt_manager.py +++ b/litellm/integrations/gitlab/gitlab_prompt_manager.py @@ -219,7 +219,7 @@ class GitLabTemplateManager: elif value.replace(".", "").isdigit(): try: result[key] = float(value) - except Exception: + except ValueError: result[key] = value else: result[key] = value.strip("\"'") diff --git a/litellm/llms/anthropic/count_tokens/token_counter.py b/litellm/llms/anthropic/count_tokens/token_counter.py index d78de0037a0..55eaa0d2ac6 100644 --- a/litellm/llms/anthropic/count_tokens/token_counter.py +++ b/litellm/llms/anthropic/count_tokens/token_counter.py @@ -6,6 +6,7 @@ from typing import Any, Final from pydantic import JsonValue +import litellm from litellm._logging import verbose_logger from litellm.exceptions import AuthenticationError from litellm.llms.anthropic.count_tokens.handler import AnthropicCountTokensHandler @@ -61,14 +62,24 @@ class AnthropicTokenCounter(BaseTokenCounter): try: auth_header: Final = await AnthropicModelInfo.aget_auth_header( - api_key=litellm_params.get("api_key"), + api_key=litellm_params.get("api_key") or litellm.anthropic_key or litellm.api_key, api_base=api_base, litellm_params=litellm_params, allow_workload_identity=True, ) if auth_header is None: - verbose_logger.warning("No Anthropic credential found for token counting") - return None + return TokenCountResponse( + total_tokens=0, + request_model=request_model, + model_used=model_to_use, + tokenizer_type="anthropic_api", + error=True, + error_message=( + f"No Anthropic API key is set for model {request_model or model_to_use}, so tokens can't be " + "counted. Set api_key on the deployment or ANTHROPIC_API_KEY in the proxy environment." + ), + status_code=401, + ) result: Final = await anthropic_count_tokens_handler.handle_count_tokens_request( model=model_to_use, diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/handler.py b/litellm/llms/azure_ai/anthropic/count_tokens/handler.py index 0e0b68df449..53734a77f3c 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/handler.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/handler.py @@ -30,7 +30,7 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig): self, model: str, messages: list[dict[str, JsonValue]], - api_key: str, + api_key: str | None, api_base: str, litellm_params: dict[str, object] | None = None, timeout: float | httpx.Timeout | None = None, @@ -80,6 +80,15 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig): api_key=api_key, litellm_params=litellm_params, ) + if not {"x-api-key", "Authorization"} & headers.keys(): + raise AnthropicError( + status_code=401, + message=( + f"No Azure AI API key or Entra ID credential is set for model {model}, so tokens can't be " + "counted. Set api_key or an Entra ID credential on the deployment, or AZURE_AI_API_KEY or " + "AZURE_API_KEY in the proxy environment." + ), + ) # Use LiteLLM's async httpx client async_client: Final = get_async_httpx_client(llm_provider=litellm.LlmProviders.AZURE_AI) diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py b/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py index 2707d24ee07..0ed432cc394 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py @@ -7,6 +7,7 @@ from typing import Any, Final from pydantic import JsonValue +import litellm from litellm._logging import verbose_logger from litellm.llms.azure_ai.anthropic.count_tokens.handler import ( AzureAIAnthropicCountTokensHandler, @@ -18,6 +19,21 @@ from litellm.types.utils import LlmProviders, TokenCountMessagesFormat, TokenCou azure_ai_anthropic_count_tokens_handler: Final = AzureAIAnthropicCountTokensHandler() +def _missing_api_base_response(request_model: str, model_to_use: str) -> TokenCountResponse: + return TokenCountResponse( + total_tokens=0, + request_model=request_model, + model_used=model_to_use, + tokenizer_type="azure_ai_anthropic_api", + error=True, + error_message=( + f"No Azure AI API base is set for model {request_model or model_to_use}, so tokens can't be counted. " + "Set api_base on the deployment or AZURE_AI_API_BASE or AZURE_API_BASE in the proxy environment." + ), + status_code=500, + ) + + class AzureAIAnthropicTokenCounter(BaseTokenCounter): """Token counter implementation for Azure AI Anthropic provider using the CountTokens API.""" @@ -65,18 +81,15 @@ class AzureAIAnthropicTokenCounter(BaseTokenCounter): if not api_key: api_key = os.getenv("AZURE_AI_API_KEY") - # Get API base from deployment config or environment - api_base = litellm_params.get("api_base") - if not api_base: - api_base = os.getenv("AZURE_AI_API_BASE") - - if not api_key: - verbose_logger.warning("No Azure AI API key found for token counting") - return None + api_base: Final = ( + litellm_params.get("api_base") + or os.getenv("AZURE_AI_API_BASE") + or os.getenv("AZURE_API_BASE") + or litellm.api_base + ) if not api_base: - verbose_logger.warning("No Azure AI API base found for token counting") - return None + return _missing_api_base_response(request_model, model_to_use) try: result: Final = await azure_ai_anthropic_count_tokens_handler.handle_count_tokens_request( diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py b/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py index d7455439151..5946a1f4480 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py @@ -24,7 +24,7 @@ class AzureAIAnthropicCountTokensConfig(AnthropicCountTokensConfig): def get_required_headers( self, - api_key: str, + api_key: str | None, litellm_params: dict[str, object] | None = None, ) -> dict[str, str]: """ @@ -45,12 +45,12 @@ class AzureAIAnthropicCountTokensConfig(AnthropicCountTokensConfig): "Content-Type": "application/json", "anthropic-version": "2023-06-01", "anthropic-beta": ANTHROPIC_TOKEN_COUNTING_BETA_VERSION, - "x-api-key": api_key, # Azure AI Anthropic requires this header + **({"x-api-key": api_key} if api_key else {}), } # Also set up Azure auth headers for flexibility litellm_params = litellm_params or {} - if "api_key" not in litellm_params: + if api_key and "api_key" not in litellm_params: litellm_params["api_key"] = api_key litellm_params_obj: Final = GenericLiteLLMParams.model_validate(litellm_params) @@ -58,8 +58,9 @@ class AzureAIAnthropicCountTokensConfig(AnthropicCountTokensConfig): # Get Azure auth headers (api-key or Authorization) azure_headers = BaseAzureLLM.base_validate_azure_environment(headers={}, litellm_params=litellm_params_obj) - # Merge Azure auth headers headers.update(azure_headers) + if "api-key" in azure_headers: + headers.setdefault("x-api-key", azure_headers["api-key"]) return headers diff --git a/litellm/llms/openai/responses/count_tokens/token_counter.py b/litellm/llms/openai/responses/count_tokens/token_counter.py index 43b3ae9f6f1..7c0697626ca 100644 --- a/litellm/llms/openai/responses/count_tokens/token_counter.py +++ b/litellm/llms/openai/responses/count_tokens/token_counter.py @@ -2,20 +2,21 @@ OpenAI Token Counter implementation using the Responses API /input_tokens endpoint. """ -import os from collections.abc import Mapping, Sequence from typing import Any, Final +import openai from pydantic import JsonValue from typing_extensions import assert_never from litellm._logging import verbose_logger +from litellm.litellm_core_utils.asyncify import asyncify from litellm.llms.anthropic.pass_through.convertible_blocks import unconvertible_blocks_error from litellm.llms.anthropic.pass_through.responses_adapters.transformation import ( LiteLLMAnthropicToResponsesAPIAdapter, ) from litellm.llms.base_llm.base_utils import BaseTokenCounter -from litellm.llms.openai.common_utils import OpenAIError +from litellm.llms.openai.common_utils import OpenAIError, get_openai_credentials from litellm.llms.openai.responses.count_tokens.handler import ( OpenAICountTokensHandler, ) @@ -23,12 +24,26 @@ from litellm.llms.openai.responses.count_tokens.transformation import ( OpenAICountTokensConfig, ResponsesCountRequestParts, ) +from litellm.llms.openai.workload_identity import ( + OpenAIWorkloadIdentityConfig, + get_workload_identity_bearer_token, + resolve_openai_workload_identity_config, +) from litellm.types.utils import LlmProviders, TokenCountMessagesFormat, TokenCountResponse # Global handler instance - reuse across all token counting requests openai_count_tokens_handler: Final = OpenAICountTokensHandler() +def _workload_identity_config(litellm_params: Mapping[str, object]) -> OpenAIWorkloadIdentityConfig | None: + api_base: Final = litellm_params.get("api_base") + return resolve_openai_workload_identity_config( + api_key=None, + api_base=api_base if isinstance(api_base, str) else None, + litellm_params=litellm_params, + ) + + class OpenAITokenCounter(BaseTokenCounter): """Token counter implementation for OpenAI provider using the Responses API.""" @@ -107,16 +122,27 @@ class OpenAITokenCounter(BaseTokenCounter): deployment = deployment or {} litellm_params: Final = deployment.get("litellm_params", {}) - # Get OpenAI API key from deployment config or environment - api_key = litellm_params.get("api_key") - if not api_key: - api_key = os.getenv("OPENAI_API_KEY") + credentials: Final = get_openai_credentials( + api_base=litellm_params.get("api_base"), api_key=litellm_params.get("api_key") + ) + credential: Final = credentials.api_key or _workload_identity_config(litellm_params) - if not api_key: - verbose_logger.warning("No OpenAI API key found for token counting") - return None + if credential is None: + return TokenCountResponse( + total_tokens=0, + request_model=request_model, + model_used=model_to_use, + tokenizer_type="openai_api", + error=True, + error_message=( + f"No OpenAI API key is set for model {request_model or model_to_use}, so tokens can't be " + "counted. Set api_key or workload identity on the deployment, or OPENAI_API_KEY in the proxy " + "environment." + ), + status_code=401, + ) - api_base: Final = litellm_params.get("api_base") + api_base: Final = credentials.api_base input_items, instructions, counted_tools, text = self._responses_request_parts( model_to_use=model_to_use, @@ -132,6 +158,11 @@ class OpenAITokenCounter(BaseTokenCounter): return None try: + api_key: Final = ( + credential + if isinstance(credential, str) + else await asyncify(get_workload_identity_bearer_token)(credential) + ) result: Final = await openai_count_tokens_handler.handle_count_tokens_request( model=model_to_use, input=input_items, @@ -163,6 +194,17 @@ class OpenAITokenCounter(BaseTokenCounter): error_message=e.message, status_code=e.status_code, ) + except openai.AuthenticationError as e: + verbose_logger.warning("OpenAI workload identity token exchange failed: %s", e.message) + return TokenCountResponse( + total_tokens=0, + request_model=request_model, + model_used=model_to_use, + tokenizer_type="openai_api", + error=True, + error_message=e.message, + status_code=401, + ) except Exception as e: verbose_logger.warning("Error calling OpenAI CountTokens API: %s", e) return TokenCountResponse( diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py index 4d0eb6a166a..ee1ac79c820 100644 --- a/litellm/proxy/anthropic_endpoints/endpoints.py +++ b/litellm/proxy/anthropic_endpoints/endpoints.py @@ -29,6 +29,7 @@ from litellm.proxy.common_request_processing import ( create_response, log_llm_api_exception, proxy_exception_from_http_exception, + proxy_exception_from_route_error, resolve_litellm_call_id, ) from litellm.proxy.common_utils.error_body_call_id import error_body_call_id @@ -339,10 +340,10 @@ async def count_tokens( messages: Final = data.get("messages", []) if not model_name: - raise HTTPException(status_code=400, detail={"error": "model parameter is required"}) + raise HTTPException(status_code=400, detail="model parameter is required") if not messages: - raise HTTPException(status_code=400, detail={"error": "messages parameter is required"}) + raise HTTPException(status_code=400, detail="messages parameter is required") # Create TokenCountRequest for the internal endpoint from litellm.proxy._types import TokenCountRequest @@ -360,13 +361,13 @@ async def count_tokens( messages_format="anthropic", output_format=_requested_output_format(data), ) - except HTTPException: - raise + except HTTPException as e: + return _anthropic_error_json_response(proxy_exception_from_http_exception(e, {}), request) except ProxyException as e: return _anthropic_error_json_response(e, request) except Exception as e: log_llm_api_exception(e, litellm_call_id) - raise HTTPException(status_code=500, detail={"error": f"Internal server error: {e}"}) + return _anthropic_error_json_response(proxy_exception_from_route_error(e), request) else: return {"input_tokens": token_response.total_tokens} diff --git a/litellm/proxy/google_endpoints/endpoints.py b/litellm/proxy/google_endpoints/endpoints.py index 784e3d4719f..d382ff53ffb 100644 --- a/litellm/proxy/google_endpoints/endpoints.py +++ b/litellm/proxy/google_endpoints/endpoints.py @@ -1,16 +1,22 @@ from collections.abc import Mapping from itertools import chain +from types import MappingProxyType from typing import Final from fastapi import APIRouter, Depends, HTTPException, Request, Response -from fastapi.responses import ORJSONResponse +from fastapi.responses import JSONResponse, ORJSONResponse from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter, ValidationError from litellm.litellm_core_utils.json_validation_rule import normalize_json_schema_types from litellm.proxy._types import * -from litellm.proxy._types import TokenCountRequest +from litellm.proxy._types import ProxyException, TokenCountRequest from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth -from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing +from litellm.proxy.common_request_processing import ( + ProxyBaseLLMRequestProcessing, + log_llm_api_exception, + proxy_exception_from_http_exception, + resolve_litellm_call_id, +) from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # legacy module exports _read_request_body, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export read_request_body, @@ -63,6 +69,37 @@ def _input_schema(declaration: _GoogleCountTokensFunctionDeclaration) -> object: return {"type": "object", "properties": {}} +_GOOGLE_RPC_STATUS: Final = MappingProxyType( + { + 400: "INVALID_ARGUMENT", + 401: "UNAUTHENTICATED", + 403: "PERMISSION_DENIED", + 404: "NOT_FOUND", + 409: "ABORTED", + 429: "RESOURCE_EXHAUSTED", + 499: "CANCELLED", + 501: "UNIMPLEMENTED", + 503: "UNAVAILABLE", + 504: "DEADLINE_EXCEEDED", + } +) + + +def _google_error_response(request: Request, error: ProxyException) -> JSONResponse: + from litellm.proxy.proxy_server import ( + _close_dangling_otel_server_span, # pyright: ignore[reportPrivateUsage] # proxy_server keeps the span-close helper private; error JSONResponses returned by the route must stamp the OTel server span like the global ProxyException handler does + ) + + status_code: Final = int(error.code) if error.code is not None and error.code.isdigit() else 500 + _close_dangling_otel_server_span(request, status_code, exc=error) + status: Final = _GOOGLE_RPC_STATUS.get(status_code, "INTERNAL" if status_code >= 500 else "FAILED_PRECONDITION") + return JSONResponse( + status_code=status_code, + content={"error": {"code": status_code, "message": error.message, "status": status}}, + headers=error.headers, + ) + + def _google_count_tokens_request(model_name: str, data: object) -> TokenCountRequest: from litellm.google_genai.adapters.transformation import GoogleGenAIAdapter @@ -251,16 +288,26 @@ async def google_count_tokens(request: Request, model_name: str): data: Final = await read_request_body(request=request) try: token_request: Final = _google_count_tokens_request(model_name, data) - except ValidationError as error: - raise HTTPException(status_code=400, detail="Invalid Google countTokens request body") from error + except ValidationError: + return _google_error_response( + request, + ProxyException( + message="Invalid Google countTokens request body", type="invalid_request_error", param=None, code=400 + ), + ) - # Call the internal token counter function with direct request flag set to False - token_response: Final = await internal_token_counter( - request=token_request, - call_endpoint=True, - ) - if token_response is None: - return TokenCountDetailsResponse(totalTokens=0, promptTokensDetails=[]) + try: + token_response: Final = await internal_token_counter(request=token_request, call_endpoint=True) + except ProxyException as error: + return _google_error_response(request, error) + except HTTPException as error: + return _google_error_response(request, proxy_exception_from_http_exception(error, {})) + except Exception as error: + log_llm_api_exception(error, resolve_litellm_call_id(request.headers.get("x-litellm-call-id"))) + return _google_error_response( + request, + ProxyException(message="Internal server error", type="internal_server_error", param=None, code=500), + ) original_response: Final[dict] = token_response.original_response or {} return TokenCountDetailsResponse( totalTokens=original_response.get("totalTokens") or token_response.total_tokens or 0, diff --git a/tests/integration/providers/test_count_tokens_error_shapes_wire.py b/tests/integration/providers/test_count_tokens_error_shapes_wire.py new file mode 100644 index 00000000000..200f5e81673 --- /dev/null +++ b/tests/integration/providers/test_count_tokens_error_shapes_wire.py @@ -0,0 +1,205 @@ +import json +from collections.abc import Iterator +from typing import Final + +import anthropic +import openai +import pytest +from integration._support.client import Gateway, gateway_from_environment +from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue, TypeAdapter + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) + +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_MESSAGES: Final[list[JsonValue]] = [{"role": "user", "content": "Count this synthetic request."}] +_PROVIDER_KEY_VARIABLES: Final = ( + "ANTHROPIC_API_KEY", + "ANTHROPIC_AUTH_TOKEN", + "AZURE_AI_API_KEY", + "AZURE_API_KEY", + "AZURE_OPENAI_API_KEY", + "AZURE_AD_TOKEN", + "AZURE_CLIENT_ID", + "AZURE_CLIENT_SECRET", + "AZURE_TENANT_ID", + "OPENAI_API_KEY", +) + + +@pytest.fixture(scope="module") +def keyless_proxy(tmp_path_factory: pytest.TempPathFactory) -> Iterator[OwnedProxy]: + with ( + gateway_from_environment() as gateway, + owned_proxy_process( + gateway, + tmp_path_factory.mktemp("count-tokens-error-shapes"), + {}, + remove_environment=_PROVIDER_KEY_VARIABLES, + workers=1, + ) as owned, + ): + yield owned + + +@pytest.fixture(scope="module") +def azure_environment_key_proxy(tmp_path_factory: pytest.TempPathFactory) -> Iterator[OwnedProxy]: + with ( + gateway_from_environment() as gateway, + owned_proxy_process( + gateway, + tmp_path_factory.mktemp("count-tokens-azure-environment-key"), + {"AZURE_API_KEY": "synthetic-azure-env-key"}, + remove_environment=tuple(name for name in _PROVIDER_KEY_VARIABLES if name != "AZURE_API_KEY"), + workers=1, + ) as owned, + ): + yield owned + + +def _anthropic(gateway: Gateway) -> anthropic.Anthropic: + return anthropic.Anthropic( + api_key=gateway.key, base_url=str(gateway.client.base_url), max_retries=0, http_client=gateway.client + ) + + +def _never_called(request: Request) -> Reply: + raise AssertionError(f"count reached the provider without a key: {request.target}") + + +@pytest.mark.parametrize( + ("deployment", "missing"), + ( + ({"model": "anthropic/claude-sonnet-4-6"}, "No Anthropic API key is set for model "), + ({"model": "azure_ai/claude-sonnet-4-6"}, "No Azure AI API key or Entra ID credential is set for model "), + ), + ids=("anthropic", "azure_ai_anthropic"), +) +def test_anthropic_count_without_a_provider_credential_is_an_anthropic_401( + keyless_proxy: OwnedProxy, deployment: dict[str, JsonValue], missing: str +) -> None: + gateway: Final = keyless_proxy.gateway + with wire_server(_never_called) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(**deployment, api_base=wire.url, api_key=None, num_retries=0) + with pytest.raises(anthropic.AuthenticationError) as raised: + _anthropic(gateway).messages.count_tokens(model=model, messages=_MESSAGES) + assert raised.value.body == { + "type": "error", + "error": {"type": "authentication_error", "message": raised.value.body["error"]["message"]}, + } + assert missing in raised.value.message, raised.value.message + assert wire.drain() == () + + +def test_azure_ai_anthropic_count_sends_the_azure_environment_key_as_x_api_key( + azure_environment_key_proxy: OwnedProxy, +) -> None: + gateway: Final = azure_environment_key_proxy.gateway + + def count(request: Request) -> Reply: + assert request.headers.get("x-api-key") == "synthetic-azure-env-key", request.headers + return Reply(body=json.dumps({"input_tokens": 58}).encode()) + + with wire_server(count) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="azure_ai/claude-sonnet-4-6", + api_base=wire.url, + api_key=None, + num_retries=0, + ) + result: Final = _anthropic(gateway).messages.count_tokens(model=model, messages=_MESSAGES) + assert result.input_tokens == 58 + + +def test_azure_ai_anthropic_count_authenticates_with_an_entra_token(keyless_proxy: OwnedProxy) -> None: + gateway: Final = keyless_proxy.gateway + + def count(request: Request) -> Reply: + assert request.headers.get("authorization") == "Bearer synthetic-entra-token", request.headers + assert "x-api-key" not in request.headers, request.headers + return Reply(body=json.dumps({"input_tokens": 61}).encode()) + + with wire_server(count) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="azure_ai/claude-sonnet-4-6", + api_base=wire.url, + api_key=None, + azure_ad_token="synthetic-entra-token", + num_retries=0, + ) + result: Final = _anthropic(gateway).messages.count_tokens(model=model, messages=_MESSAGES) + assert result.input_tokens == 61 + assert len(wire.drain()) == 1 + + +def test_anthropic_count_request_errors_use_the_anthropic_envelope(keyless_proxy: OwnedProxy) -> None: + gateway: Final = keyless_proxy.gateway + upstream_error: Final = {"type": "error", "error": {"type": "invalid_request_error", "message": "scripted reject"}} + with ( + wire_server(lambda _request: Reply(status=400, body=json.dumps(upstream_error).encode())) as wire, + gateway.scenario() as scenario, + ): + model: Final = scenario.model( + model="anthropic/claude-sonnet-4-6", api_base=wire.url, api_key="synthetic-key", num_retries=0 + ) + missing_messages: Final = gateway.request("POST", "/v1/messages/count_tokens", {"model": model}) + assert missing_messages.status_code == 400, missing_messages.text + assert _JSON_OBJECT.validate_json(missing_messages.content) == { + "type": "error", + "error": {"type": "invalid_request_error", "message": "messages parameter is required"}, + } + + rejected: Final = gateway.request( + "POST", "/v1/messages/count_tokens", {"model": model, "messages": _MESSAGES} + ) + assert rejected.status_code == 400, rejected.text + body: Final = _JSON_OBJECT.validate_json(rejected.content) + assert set(body) <= {"type", "error", "request_id"}, body + assert body["error"] == upstream_error["error"], body + assert len(wire.drain()) == 1 + + +def test_openai_input_tokens_without_a_provider_key_is_an_openai_401(keyless_proxy: OwnedProxy) -> None: + gateway: Final = keyless_proxy.gateway + with wire_server(_never_called) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="openai/gpt-4o-mini", api_base=f"{wire.url}/v1", api_key=None, num_retries=0 + ) + client: Final = openai.OpenAI( + api_key=gateway.key, base_url=str(gateway.client.base_url), max_retries=0, http_client=gateway.client + ) + with pytest.raises(openai.AuthenticationError) as raised: + client.responses.input_tokens.count(model=model, input="Count this synthetic request.") + assert f"No OpenAI API key is set for model {model}" in raised.value.message + assert wire.drain() == () + + +def test_google_count_tokens_errors_use_the_google_envelope(keyless_proxy: OwnedProxy) -> None: + gateway: Final = keyless_proxy.gateway + upstream_error: Final = { + "error": {"code": 400, "message": "API key not valid. Please pass a valid API key.", "status": "INVALID_ARGUMENT"} + } + with ( + wire_server(lambda _request: Reply(status=400, body=json.dumps(upstream_error).encode())) as wire, + gateway.scenario() as scenario, + ): + model: Final = scenario.model( + model="gemini/gemini-2.5-flash", api_base=wire.url, api_key="synthetic-key", num_retries=0 + ) + rejected: Final = gateway.request( + "POST", + f"/v1beta/models/{model}:countTokens", + {"contents": [{"role": "user", "parts": [{"text": "Count this synthetic request."}]}]}, + ) + assert rejected.status_code == 400, rejected.text + error: Final = _JSON_OBJECT.validate_json(rejected.content)["error"] + assert isinstance(error, dict) and error["code"] == 400 and error["status"] == "INVALID_ARGUMENT", error + assert "API key not valid" in str(error["message"]), error + assert len(wire.drain()) == 1 + + malformed: Final = gateway.request("POST", f"/v1beta/models/{model}:countTokens", {"contents": "not-a-list"}) + assert malformed.status_code == 400, malformed.text + assert _JSON_OBJECT.validate_json(malformed.content) == { + "error": {"code": 400, "message": "Invalid Google countTokens request body", "status": "INVALID_ARGUMENT"} + } diff --git a/tests/unit/llms/anthropic/test_count_tokens_oauth.py b/tests/unit/llms/anthropic/test_count_tokens_oauth.py index 7f717091754..16ec513916c 100644 --- a/tests/unit/llms/anthropic/test_count_tokens_oauth.py +++ b/tests/unit/llms/anthropic/test_count_tokens_oauth.py @@ -197,6 +197,30 @@ class TestCountTokensUsesWorkloadIdentity: assert ANTHROPIC_TOKEN_COUNTING_BETA_VERSION in betas assert ANTHROPIC_OAUTH_BETA_HEADER not in betas + @pytest.mark.asyncio + async def test_a_global_anthropic_key_counts_like_chat_would_authenticate( + self, monkeypatch, httpx_transport_clients + ): + for name in ("ANTHROPIC_API_KEY", "ANTHROPIC_AUTH_TOKEN", "ANTHROPIC_API_BASE", "ANTHROPIC_BASE_URL"): + monkeypatch.delenv(name, raising=False) + monkeypatch.setattr(litellm, "anthropic_key", FAKE_REGULAR_KEY) + + with respx.mock(assert_all_called=True) as router: + route = router.post("https://api.anthropic.com/v1/messages/count_tokens").mock( + return_value=httpx.Response(200, json={"input_tokens": 13}) + ) + result = await AnthropicTokenCounter().count_tokens( + model_to_use="claude-sonnet-4-5", + messages=[{"role": "user", "content": "hi"}], + contents=None, + deployment={"litellm_params": {"model": "anthropic/claude-sonnet-4-5"}}, + request_model="claude-sonnet-4-5", + ) + + assert result is not None + assert result.total_tokens == 13 + assert route.calls.last.request.headers["x-api-key"] == FAKE_REGULAR_KEY + @pytest.mark.asyncio async def test_a_failed_mint_degrades_like_an_anthropic_error(self, monkeypatch): from litellm.llms.anthropic.count_tokens import token_counter as token_counter_module diff --git a/tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py b/tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py index 29380c6fc53..86ad6c8d9d9 100644 --- a/tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py +++ b/tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py @@ -186,3 +186,27 @@ async def test_undecodable_count_body_becomes_a_500_error_response_carrying_the_ "status_code": 500, } ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("set_base", ("env", "global"), ids=("AZURE_API_BASE", "litellm.api_base")) +async def test_count_uses_the_shared_api_base_chat_would_use(monkeypatch: pytest.MonkeyPatch, set_base: str) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.delenv("AZURE_AI_API_BASE", raising=False) + monkeypatch.delenv("AZURE_API_BASE", raising=False) + if set_base == "env": + monkeypatch.setenv("AZURE_API_BASE", API_BASE) + else: + monkeypatch.setattr(litellm, "api_base", API_BASE) + with respx.mock: + route: Final = respx.post(COUNT_URL).mock(return_value=httpx.Response(200, json={"input_tokens": 6})) + result: Final = await AzureAIAnthropicTokenCounter().count_tokens( + model_to_use="claude-sonnet-4-5", + messages=[{"role": "user", "content": "hi"}], + contents=None, + deployment={"litellm_params": {"api_key": "azure-key"}}, + request_model="azure_ai/claude-sonnet-4-5", + ) + assert result is not None + assert result.total_tokens == 6 + assert route.called diff --git a/tests/unit/llms/openai/responses/count_tokens/test_token_counter.py b/tests/unit/llms/openai/responses/count_tokens/test_token_counter.py index 976ecff01ea..befdc5ee9fd 100644 --- a/tests/unit/llms/openai/responses/count_tokens/test_token_counter.py +++ b/tests/unit/llms/openai/responses/count_tokens/test_token_counter.py @@ -388,3 +388,44 @@ async def test_undecodable_count_body_becomes_a_500_error_response_carrying_the_ "status_code": 500, } ) + + +@pytest.mark.asyncio +async def test_count_uses_the_global_openai_key_chat_would_use(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "openai_key", "sk-from-litellm-settings") + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + with respx.mock: + route: Final = respx.post(COUNT_URL).mock(return_value=httpx.Response(200, json={"input_tokens": 9})) + result: Final = await OpenAITokenCounter().count_tokens( + model_to_use="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + contents=None, + deployment={"litellm_params": {}}, + request_model="openai/gpt-4o", + ) + assert result is not None + assert result.total_tokens == 9 + assert route.calls.last.request.headers["authorization"] == "Bearer sk-from-litellm-settings" + + +@pytest.mark.asyncio +async def test_count_sends_the_global_key_only_to_the_global_api_base(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "openai_key", "sk-private-gateway") + monkeypatch.setattr(litellm, "api_base", "https://gateway.internal.example/v1") + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + with respx.mock: + route: Final = respx.post("https://gateway.internal.example/v1/responses/input_tokens").mock( + return_value=httpx.Response(200, json={"input_tokens": 4}) + ) + result: Final = await OpenAITokenCounter().count_tokens( + model_to_use="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + contents=None, + deployment={"litellm_params": {}}, + request_model="openai/gpt-4o", + ) + assert result is not None + assert result.total_tokens == 4 + assert route.calls.last.request.headers["authorization"] == "Bearer sk-private-gateway" diff --git a/tests/unit/llms/openai/test_openai_workload_identity.py b/tests/unit/llms/openai/test_openai_workload_identity.py index d0cf52d0e28..2236de1eb5f 100644 --- a/tests/unit/llms/openai/test_openai_workload_identity.py +++ b/tests/unit/llms/openai/test_openai_workload_identity.py @@ -648,3 +648,56 @@ class TestClientsideBaseOverride: ) is None ) + + +class TestCountTokens: + @pytest.mark.asyncio + @respx.mock + async def test_count_authenticates_with_the_workload_identity_bearer( + self, wif_env: OpenAIWorkloadIdentityConfig, monkeypatch: pytest.MonkeyPatch + ) -> None: + from litellm.llms.openai.responses.count_tokens.token_counter import OpenAITokenCounter + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "api_key", None) + monkeypatch.setattr(litellm, "openai_key", None) + mock_token_exchange() + count_route: Final = respx.post("https://api.openai.com/v1/responses/input_tokens").mock( + return_value=httpx.Response(200, json={"input_tokens": 12}) + ) + result: Final = await OpenAITokenCounter().count_tokens( + model_to_use="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + contents=None, + deployment={"litellm_params": {}}, + request_model="openai/gpt-4o", + ) + assert result is not None + assert result.error is False + assert result.total_tokens == 12 + assert count_route.calls.last.request.headers["authorization"] == "Bearer exchanged-bearer-token" + + @pytest.mark.asyncio + @respx.mock + async def test_failed_mint_becomes_a_count_error_response( + self, wif_env: OpenAIWorkloadIdentityConfig, monkeypatch: pytest.MonkeyPatch + ) -> None: + from litellm.llms.openai.responses.count_tokens.token_counter import OpenAITokenCounter + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "api_key", None) + monkeypatch.setattr(litellm, "openai_key", None) + respx.post(TOKEN_EXCHANGE_URL).mock(return_value=httpx.Response(400, json={"error": "invalid_grant"})) + count_route: Final = respx.post("https://api.openai.com/v1/responses/input_tokens") + result: Final = await OpenAITokenCounter().count_tokens( + model_to_use="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + contents=None, + deployment={"litellm_params": {}}, + request_model="openai/gpt-4o", + ) + assert result is not None + assert result.error is True + assert result.status_code == 401 + assert result.total_tokens == 0 + assert not count_route.called diff --git a/tests/unit/proxy/anthropic_endpoints/test_endpoints.py b/tests/unit/proxy/anthropic_endpoints/test_endpoints.py index 69d182667ec..e4129ebbc39 100644 --- a/tests/unit/proxy/anthropic_endpoints/test_endpoints.py +++ b/tests/unit/proxy/anthropic_endpoints/test_endpoints.py @@ -440,8 +440,6 @@ class TestErrorLogCarriesCallId: @pytest.mark.asyncio async def test_count_tokens_failure_log_carries_callers_call_id(self, caplog: pytest.LogCaptureFixture): - from fastapi import HTTPException - import litellm.proxy.anthropic_endpoints.endpoints as ep import litellm.proxy.proxy_server as proxy_server from litellm.proxy._types import UserAPIKeyAuth @@ -460,11 +458,10 @@ class TestErrorLogCarriesCallId: proxy_server, "count_request_tokens", new=AsyncMock(side_effect=RuntimeError("tokenizer down")) ), # test-quality-ok: module global imported at call time; the test targets the endpoint's except block caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"), - pytest.raises(HTTPException) as raised, ): - await ep.count_tokens(request=request, user_api_key_dict=UserAPIKeyAuth()) + response = await ep.count_tokens(request=request, user_api_key_dict=UserAPIKeyAuth()) - assert raised.value.status_code == 500 + assert response.status_code == 500 record = self._error_record(caplog) assert record.litellm_call_id == call_id assert call_id in record.getMessage() diff --git a/tests/unit/proxy/google_endpoints/test_google_api_endpoints.py b/tests/unit/proxy/google_endpoints/test_google_api_endpoints.py index 451792109de..114bbb184c6 100644 --- a/tests/unit/proxy/google_endpoints/test_google_api_endpoints.py +++ b/tests/unit/proxy/google_endpoints/test_google_api_endpoints.py @@ -182,6 +182,29 @@ def test_google_count_tokens_unchanged(): assert body["totalTokens"] == 7 +def test_google_count_tokens_unexpected_failure_uses_the_google_error_envelope(): + try: + client = _build_test_client() + except ImportError as e: + pytest.skip(f"Skipping test due to missing dependency: {e}") + + with patch( + "litellm.proxy.proxy_server.token_counter", + new_callable=AsyncMock, + side_effect=RuntimeError("tokenizer unavailable"), + ): + response = client.post( + "/v1beta/models/test-model:countTokens", + json={"contents": [{"role": "user", "parts": [{"text": "Hello"}]}]}, + ) + + assert response.status_code == 500 + error = response.json()["error"] + assert error["code"] == 500 + assert error["status"] == "INTERNAL" + assert "tokenizer unavailable" not in error["message"] + + def test_google_count_tokens_uses_normalized_total_when_provider_shape_differs(): """A provider that counts in Anthropic shape (Bedrock, Anthropic) carries no totalTokens key, so the normalized total must win over the raw response's missing field.""" try: diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index 9bed719227f..fbe47e4a3d7 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -327,22 +327,16 @@ async def test_internal_token_counter_anthropic_provider_detection(): setattr(litellm.proxy.proxy_server, "llm_router", llm_router) - # Test with is_direct_request=False but non-Anthropic provider - response = await token_counter( - request=TokenCountRequest( - model="gpt-test", - messages=[{"role": "user", "content": "hello"}], - ), - call_endpoint=True, - ) - - print("Non-Anthropic provider test response:", response) - - # Verify response structure - assert response.request_model == "gpt-test" - assert response.model_used == "gpt-4" - assert response.total_tokens > 0 - assert response.tokenizer_type == "openai_tokenizer" # Should use LiteLLM tokenizer + with pytest.raises(ProxyException) as missing_key: + await token_counter( + request=TokenCountRequest( + model="gpt-test", + messages=[{"role": "user", "content": "hello"}], + ), + call_endpoint=True, + ) + assert missing_key.value.code == "401" + assert "No OpenAI API key is set for model gpt-test" in missing_key.value.message @pytest.mark.asyncio @@ -352,12 +346,14 @@ async def test_anthropic_endpoint_error_handling(): """ from unittest.mock import MagicMock - from fastapi import HTTPException, Request + from fastapi import Request from litellm.proxy.anthropic_endpoints.endpoints import count_tokens # Mock request object mock_request = MagicMock(spec=Request) + mock_request.headers = {} + mock_request.scope = {} mock_user_api_key_dict = MagicMock() # Test missing model parameter @@ -375,12 +371,13 @@ async def test_anthropic_endpoint_error_handling(): anthropic_endpoints.read_request_body = mock_read_request_body try: - # Should raise HTTPException for missing model - with pytest.raises(HTTPException) as exc_info: - await count_tokens(mock_request, mock_user_api_key_dict) + response = await count_tokens(mock_request, mock_user_api_key_dict) - assert exc_info.value.status_code == 400 - assert "model parameter is required" in str(exc_info.value.detail) + assert response.status_code == 400 + assert json.loads(response.body) == { + "type": "error", + "error": {"type": "invalid_request_error", "message": "model parameter is required"}, + } print("✅ Error handling test passed!") @@ -471,7 +468,7 @@ async def test_factory_gpt4_endpoint_does_not_call_anthropic_counter(): mock_router.model_list = [ { "model_name": "gpt-4", - "litellm_params": {"model": "openai/gpt-4"}, + "litellm_params": {"model": "databricks/gpt-4"}, "model_info": {}, } ] @@ -480,7 +477,7 @@ async def test_factory_gpt4_endpoint_does_not_call_anthropic_counter(): mock_router.async_get_available_deployment = AsyncMock( return_value={ "model_name": "gpt-4", - "litellm_params": {"model": "openai/gpt-4"}, + "litellm_params": {"model": "databricks/gpt-4"}, "model_info": {}, } ) From a3251a3bfde17fb1c88896a1f2f3b7f9d2ef0345 Mon Sep 17 00:00:00 2001 From: "moyai-devin-berriai[bot]" <336287033+moyai-devin-berriai[bot]@users.noreply.github.com> Date: Fri, 16 Oct 2026 00:29:11 -0700 Subject: [PATCH 06/23] fix(responses): preserve provider request IDs in streaming failure logs (#45453) * fix(responses): preserve provider request IDs in streaming failure logs * fix(bedrock): log the provider request id for in-stream errors on bridged Bedrock streams Bedrock models that run through the chat-completions bridge (Claude, converse, invoke) raise in-stream failures from the event-stream decoder as a BedrockError without headers, so their failure logs kept error_provider_request_id null. The decoder now passes the stream's request-id headers to those errors. Both the native Responses iterator and the decoder pass only the PROVIDER_REQUEST_ID_HEADERS subset, so error responses and failure logs pick up the request id and none of the 200 response's other headers. MidStreamFallbackError gets its response at construction instead of mutating it afterwards. --------- Co-authored-by: moyai-devin-berriai[bot] <336287033+moyai-devin-berriai[bot]@users.noreply.github.com> Co-authored-by: yucheng-berri <293926326+yucheng-berri@users.noreply.github.com> --- .../llm_response_utils/get_headers.py | 13 + litellm/llms/bedrock/chat/invoke_handler.py | 28 +- litellm/llms/bedrock/common_utils.py | 15 +- litellm/responses/streaming_iterator.py | 19 +- .../test_bedrock_error_request_id.py | 752 +++++++++++++++++- .../llm_response_utils/test_get_headers.py | 26 + .../llms/bedrock/chat/test_invoke_handler.py | 94 ++- .../test_streaming_iterator_error_events.py | 50 ++ 8 files changed, 977 insertions(+), 20 deletions(-) diff --git a/litellm/litellm_core_utils/llm_response_utils/get_headers.py b/litellm/litellm_core_utils/llm_response_utils/get_headers.py index b86cc954b0e..3c08f16fa80 100644 --- a/litellm/litellm_core_utils/llm_response_utils/get_headers.py +++ b/litellm/litellm_core_utils/llm_response_utils/get_headers.py @@ -1,8 +1,11 @@ from collections.abc import Mapping from typing import Final +import httpx from pydantic import TypeAdapter, ValidationError +from litellm.constants import PROVIDER_REQUEST_ID_HEADERS + _HEADERS: Final = TypeAdapter(Mapping[str, str]) _REQUEST_ID_HEADERS: Final = ("request-id", "x-request-id", "llm_provider-request-id", "llm_provider-x-request-id") @@ -16,6 +19,16 @@ def get_provider_request_id(headers: object) -> str | None: return next((normalized[name] for name in _REQUEST_ID_HEADERS if normalized.get(name)), None) +def provider_request_id_headers(headers: Mapping[str, object] | None) -> httpx.Headers | None: + """The provider request id headers in `headers` whose values are ASCII strings, or None when there are none.""" + request_id_headers: Final = tuple( + (name, value) + for name, value in (headers or {}).items() + if name.lower() in PROVIDER_REQUEST_ID_HEADERS and isinstance(value, str) and value.isascii() + ) + return httpx.Headers(request_id_headers) if request_id_headers else None + + def get_response_headers(_response_headers: Mapping[str, str] | None = None) -> dict: """ diff --git a/litellm/llms/bedrock/chat/invoke_handler.py b/litellm/llms/bedrock/chat/invoke_handler.py index a207bceb418..6ad400661be 100644 --- a/litellm/llms/bedrock/chat/invoke_handler.py +++ b/litellm/llms/bedrock/chat/invoke_handler.py @@ -13,6 +13,7 @@ from litellm.caching.caching import InMemoryCache from litellm.constants import RESPONSE_FORMAT_TOOL_NAME from litellm.litellm_core_utils.core_helpers import map_finish_reason from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.litellm_core_utils.llm_response_utils.get_headers import provider_request_id_headers from litellm.litellm_core_utils.optional_imports import ensure_optional_import from litellm.litellm_core_utils.prompt_templates.common_utils import native_redacted_thinking_data from litellm.llms.anthropic.chat.handler import ( @@ -448,6 +449,7 @@ class _EventStreamTally: f"x-amzn-requestid={_response_header(response_headers, 'x-amzn-requestid')!r}, " f"{self.bytes_received} bytes received" ) + request_id_headers: Final = provider_request_id_headers(response_headers) if not self.events: return BedrockError( status_code=502, @@ -455,6 +457,7 @@ class _EventStreamTally: "Bedrock answered the stream with HTTP 200 but its body decoded to no events " f"({detail}, first bytes={self.head!r})" ), + headers=request_id_headers, ) if not self.recognized_events: return BedrockError( @@ -464,10 +467,12 @@ class _EventStreamTally: f"event type (event types={sorted(self.unrecognized_event_types)}, {detail}, " f"first payload={self.unrecognized_head!r})" ), + headers=request_id_headers, ) return BedrockError( status_code=502, message=f"Bedrock stream ended with {undecoded} undecoded bytes after {self.events} events ({detail})", + headers=request_id_headers, ) @@ -837,11 +842,12 @@ class AWSEventStreamDecoder: event_stream_buffer: Final = EventStreamBuffer() tally: Final = _EventStreamTally(get_bedrock_stream_event_statuses()) + request_id_headers: Final = provider_request_id_headers(response_headers) for chunk in iterator: event_stream_buffer.add_data(chunk) tally.add_chunk(chunk) for event in event_stream_buffer: - message = self._decode_event(event, tally) + message = self._decode_event(event, tally, request_id_headers) if message: # sse_event = ServerSentEvent(data=message, event="completion") _data = json.loads(message) @@ -858,11 +864,12 @@ class AWSEventStreamDecoder: event_stream_buffer: Final = EventStreamBuffer() tally: Final = _EventStreamTally(get_bedrock_stream_event_statuses()) + request_id_headers: Final = provider_request_id_headers(response_headers) async for chunk in iterator: event_stream_buffer.add_data(chunk) tally.add_chunk(chunk) for event in event_stream_buffer: - message = self._decode_event(event, tally) + message = self._decode_event(event, tally, request_id_headers) if message: _data = json.loads(message) yield self.chunk_parser(chunk_data=_data) @@ -882,27 +889,34 @@ class AWSEventStreamDecoder: ) return response_stream_shape - def _decode_event(self, event: "EventStreamMessage", tally: _EventStreamTally) -> str | None: + def _decode_event( + self, event: "EventStreamMessage", tally: _EventStreamTally, request_id_headers: httpx.Headers | None + ) -> str | None: response_stream_shape: Final = self._response_stream_shape() response_dict: Final = bedrock_event_stream_response(event) tally.add_event(event, response_dict["headers"]) - return self._parse_message_from_response(response_dict, response_stream_shape) + return self._parse_message_from_response(response_dict, response_stream_shape, request_id_headers) def _parse_message_from_event(self, event: "EventStreamMessage") -> str | None: response_stream_shape: Final = self._response_stream_shape() return self._parse_message_from_response(bedrock_event_stream_response(event), response_stream_shape) def _parse_message_from_response( - self, response_dict: BedrockEventStreamResponseDict, response_stream_shape: "Shape" + self, + response_dict: BedrockEventStreamResponseDict, + response_stream_shape: "Shape", + request_id_headers: httpx.Headers | None = None, ) -> str | None: parsed_response: Final = self.parser.parse(response_dict, response_stream_shape) if response_dict["status_code"] != 200: - raise build_bedrock_stream_error(response_dict, response_stream_shape) + raise build_bedrock_stream_error(response_dict, response_stream_shape, headers=request_id_headers) event_type: Final = bedrock_event_stream_header(response_dict["headers"], ":event-type") event_error_status: Final = bedrock_stream_event_error_status(event_type) if event_type is not None and event_error_status is not None: - raise build_bedrock_stream_event_error(event_type, event_error_status, response_dict["body"]) + raise build_bedrock_stream_event_error( + event_type, event_error_status, response_dict["body"], headers=request_id_headers + ) if "chunk" in parsed_response: chunk = parsed_response.get("chunk") if not chunk: diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index e79865ef6bf..5862b05a619 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -1913,13 +1913,18 @@ def bedrock_event_stream_header(headers: Mapping[str, object], name: str) -> str return value if isinstance(value, str) else None -def build_bedrock_stream_event_error(event_type: str, status_code: int, body: bytes) -> BedrockError: - return BedrockError(status_code=status_code, message=f"{event_type} {body.decode(errors='replace')}") +def build_bedrock_stream_event_error( + event_type: str, status_code: int, body: bytes, headers: httpx.Headers | None = None +) -> BedrockError: + return BedrockError( + status_code=status_code, message=f"{event_type} {body.decode(errors='replace')}", headers=headers + ) def build_bedrock_stream_error( response_dict: BedrockEventStreamResponseDict, response_stream_shape: Shape | None, + headers: httpx.Headers | None = None, ) -> BedrockError: """Build a BedrockError for a non-200 event-stream error event. @@ -1929,12 +1934,14 @@ def build_bedrock_stream_error( """ exception_type: Final = bedrock_event_stream_header(response_dict["headers"], ":exception-type") if exception_type is None: - return BedrockError(status_code=response_dict["status_code"], message=response_dict["body"].decode()) + return BedrockError( + status_code=response_dict["status_code"], message=response_dict["body"].decode(), headers=headers + ) member: Final = _structure_members(response_stream_shape).get(exception_type) modeled_status: Final = None if member is None else _modeled_error_status(member) status_code: Final = response_dict["status_code"] if modeled_status is None else modeled_status - return build_bedrock_stream_event_error(exception_type, status_code, response_dict["body"]) + return build_bedrock_stream_event_error(exception_type, status_code, response_dict["body"], headers=headers) class BedrockEventStreamDecoderBase: diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 7cf68332c07..89c32aed1fb 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -31,6 +31,7 @@ from litellm.litellm_core_utils.core_helpers import process_response_headers from litellm.litellm_core_utils.internal_call_metadata import without_budget_reservation from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.llm_response_utils.get_api_base import get_api_base +from litellm.litellm_core_utils.llm_response_utils.get_headers import provider_request_id_headers from litellm.litellm_core_utils.llm_response_utils.response_metadata import ( update_response_metadata, ) @@ -262,7 +263,9 @@ def stream_error_status_and_message(error_obj: object) -> tuple[int, str]: return _status_code_for_error_fields(error_type, error_code), message -def _map_stream_error_to_exception(error_obj: object, model: str, custom_llm_provider: str) -> Exception: +def _map_stream_error_to_exception( + error_obj: object, model: str, custom_llm_provider: str, headers: httpx.Headers | None = None +) -> Exception: from litellm.llms.base_llm.chat.transformation import BaseLLMException error_message, error_type, error_code = _error_event_fields(error_obj) @@ -272,6 +275,7 @@ def _map_stream_error_to_exception(error_obj: object, model: str, custom_llm_pro status_code=status_code, message=f"Error code: {status_code} - {{'error': {error_body}}}", body=error_body, + headers=headers, ) try: return litellm.exception_type( @@ -667,7 +671,12 @@ class BaseResponsesAPIStreamingIterator: ) def _map_error_event_exception(self, error_obj: object) -> Exception: - return _map_stream_error_to_exception(error_obj, self.model or "", self.custom_llm_provider or "") + return _map_stream_error_to_exception( + error_obj, + self.model or "", + self.custom_llm_provider or "", + headers=provider_request_id_headers(self._raw_response_headers), + ) def _maybe_raise_for_error_event(self, result: object) -> None: chunk_type: Final = getattr(result, "type", None) @@ -683,11 +692,17 @@ class BaseResponsesAPIStreamingIterator: mapped_exception: Final = self._map_error_event_exception(error_obj) if not _mid_stream_fallback_eligible(mapped_exception): raise mapped_exception + mapped_status: Final = getattr(mapped_exception, "status_code", None) raise MidStreamFallbackError( message=str(mapped_exception), model=self.model or "", llm_provider=self.custom_llm_provider or "", original_exception=mapped_exception, + response=httpx.Response( + status_code=int(mapped_status) if mapped_status is not None else 503, + headers=provider_request_id_headers(self._raw_response_headers), + request=httpx.Request(method="POST", url=f"https://{self.custom_llm_provider or ''}.com/v1/"), + ), generated_content=self._generated_content, is_pre_first_chunk=not self._yielded_first_chunk, ) diff --git a/tests/integration/observability/test_bedrock_error_request_id.py b/tests/integration/observability/test_bedrock_error_request_id.py index 230619a2863..436dcc572d5 100644 --- a/tests/integration/observability/test_bedrock_error_request_id.py +++ b/tests/integration/observability/test_bedrock_error_request_id.py @@ -1,13 +1,208 @@ +import asyncio +import base64 import json import uuid -from typing import Final +from collections.abc import Mapping +from concurrent.futures import ThreadPoolExecutor +from types import MappingProxyType +from typing import Final, Literal, NamedTuple +from urllib.parse import unquote -from integration._support.client import Gateway, eventually, object_value +import anthropic +import httpx +import openai +import pytest +from integration._support.client import Gateway, Scenario, eventually, object_value from integration._support.database import read_rows +from integration._support.responses_stream import ( + created, + delta, + error_event, + failed, + frame, + healthy_stream, + openai_error_event, + spend_row, +) +from integration._support.upstream import _aws_event_frame, aws_event_stream_frame from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue + +import litellm +from litellm.exceptions import MidStreamFallbackError +from litellm.litellm_core_utils.litellm_logging import _get_provider_request_id _MODEL: Final = "bedrock/converse/anthropic.claude-sonnet-4-5-20250929-v1:0" +_BEDROCK_RESPONSES_MODEL: Final = "bedrock/us.openai.gpt-5.6-sol" _TOKEN: Final = "synthetic-bedrock-bearer" +_SHAPES: Final = ("nested-error", "top-level-error", "response.failed", "error-before-output") +_ERROR_CODES: Final = ("server_error", "rate_limit_exceeded", "invalid_request_error", "content_policy_violation") +_INVOKE_MODEL_ID: Final = "anthropic.claude-3-haiku-20240307-v1:0" +_BEDROCK_AWS: Final[Mapping[str, JsonValue]] = MappingProxyType( + { + "aws_access_key_id": "AKIASCRIPTEDPROVIDER", + "aws_secret_access_key": "scripted-secret", + "aws_region_name": "us-east-1", + } +) +_EVENT_STREAM: Final = "application/vnd.amazon.eventstream" +_BEDROCK_STREAM_TARGETS: Final = frozenset( + ( + "/model/anthropic.claude-sonnet-4-5-20250929-v1:0/converse-stream", + f"/model/{_INVOKE_MODEL_ID}/invoke-with-response-stream", + ) +) + +BedrockFailure = Literal["exception-event", "exception-message", "no-known-event", "truncated-frame"] + + +class _FailureShape(NamedTuple): + bedrock_status: int + converse_class: str + after_output: tuple[bool, ...] + + +BedrockEndpoint = Literal["chat", "responses", "messages", "messages-invoke"] +_BEDROCK_ENDPOINTS: Final[tuple[BedrockEndpoint, ...]] = ("chat", "responses", "messages", "messages-invoke") +_BEDROCK_FAILURES: Final[Mapping[BedrockFailure, _FailureShape]] = MappingProxyType( + { + "exception-event": _FailureShape(500, "ServiceUnavailableError", (False, True)), + "exception-message": _FailureShape(429, "RateLimitError", (False, True)), + "no-known-event": _FailureShape(502, "BadGatewayError", (False,)), + "truncated-frame": _FailureShape(502, "BadGatewayError", (True,)), + } +) +_BEDROCK_CELLS: Final = tuple( + pytest.param( + endpoint, failure, after_output, id=f"{endpoint}-{failure}-{'after' if after_output else 'before'}-output" + ) + for endpoint in _BEDROCK_ENDPOINTS + for failure, shape in _BEDROCK_FAILURES.items() + for after_output in shape.after_output +) + + +def _stream(identity: str, shape: str, error_code: str, message: str) -> tuple[bytes, ...]: + opening: Final = (frame(created(identity)), frame(delta(identity, "partial output"))) + nested: Final = frame(error_event({"type": error_code, "code": error_code, "message": message})) + match shape: + case "nested-error": + return (*opening, nested) + case "top-level-error": + return (*opening, frame(openai_error_event(error_code, message))) + case "response.failed": + return (*opening, frame(failed(identity, error_code, message))) + case "error-before-output": + return (nested,) + raise AssertionError(shape) + + +def _bedrock_frame(event_type: str, payload: dict[str, JsonValue]) -> bytes: + return _aws_event_frame(event_type, payload, "sc", "u") + + +def _invoke_chunk(event: dict[str, JsonValue]) -> bytes: + return _bedrock_frame("chunk", {"bytes": base64.b64encode(json.dumps(event).encode()).decode()}) + + +def _bedrock_output(invoke: bool, text: str = "partial output") -> tuple[bytes, ...]: + if not invoke: + return ( + _bedrock_frame("messageStart", {"role": "assistant"}), + _bedrock_frame("contentBlockDelta", {"delta": {"text": text}, "contentBlockIndex": 0}), + ) + started: Final[dict[str, JsonValue]] = { + "id": f"msg_{uuid.uuid4().hex}", + "type": "message", + "role": "assistant", + "content": [], + "model": _INVOKE_MODEL_ID, + "stop_reason": None, + "usage": {"input_tokens": 3, "output_tokens": 1}, + } + return ( + _invoke_chunk({"type": "message_start", "message": started}), + _invoke_chunk({"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}), + _invoke_chunk({"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": text}}), + ) + + +def _bedrock_failure_frame(failure: BedrockFailure, message: str) -> bytes: + match failure: + case "exception-event": + return _bedrock_frame("internalServerException", {"message": message}) + case "exception-message": + headers: Final = { + ":exception-type": "throttlingException", + ":content-type": "application/json", + ":message-type": "exception", + } + return aws_event_stream_frame(headers, json.dumps({"message": message}).encode()) + case "no-known-event": + return _bedrock_frame("somethingBedrockAddedLater", {"message": message}) + case "truncated-frame": + return _bedrock_frame("contentBlockDelta", {"delta": {"text": message}, "contentBlockIndex": 0})[:-7] + + +def _bedrock_failure_text(failure: BedrockFailure, message: str) -> str: + return "undecoded bytes after" if failure == "truncated-frame" else message + + +def _bedrock_stream_request( + endpoint: BedrockEndpoint, model: str, prompt: str, *, stream: bool = True +) -> tuple[str, dict[str, JsonValue]]: + extra: Final[dict[str, JsonValue]] = {"stream": stream, "num_retries": 0, "cache": {"no-cache": True}} + turn: Final[list[JsonValue]] = [{"role": "user", "content": prompt}] + match endpoint: + case "chat": + return "/v1/chat/completions", {"model": model, "messages": turn, **extra} + case "responses": + return "/v1/responses", {"model": model, "input": prompt, **extra} + case "messages" | "messages-invoke": + return "/v1/messages", {"model": model, "messages": turn, "max_tokens": 16, **extra} + + +def _failure_information(call_id: str) -> dict[str, JsonValue]: + row: Final = spend_row(call_id) + assert row["status"] == "failure", row + metadata: Final = row["metadata"] + parsed: Final = json.loads(metadata) if isinstance(metadata, str) else object_value(metadata) + return object_value(parsed["error_information"]) + + +def _streamed_response(gateway: Gateway, model: str, prompt: str) -> httpx.Response: + response: Final = gateway.request( + "POST", + "/v1/responses", + {"model": model, "input": prompt, "stream": True, "num_retries": 0, "cache": {"no-cache": True}}, + ) + assert response.status_code == 200, response.text + return response + + +async def _consume_sdk_stream(asynchronous: bool, api_base: str, identity: str) -> None: + if asynchronous: + async_stream: Final = await litellm.aresponses( + model="openai/gpt-6", + input=identity, + stream=True, + api_base=api_base, + api_key="synthetic-provider-key", + num_retries=0, + ) + async for _ in async_stream: + pass + return + tuple( + litellm.responses( + model="openai/gpt-6", + input=identity, + stream=True, + api_base=api_base, + api_key="synthetic-provider-key", + num_retries=0, + ) + ) def test_bedrock_500_keeps_amzn_request_id_on_error_headers_and_failure_log(gateway: Gateway) -> None: @@ -54,3 +249,556 @@ def test_bedrock_500_keeps_amzn_request_id_on_error_headers_and_failure_log(gate parsed: Final = json.loads(metadata) if isinstance(metadata, str) else object_value(metadata) error_information: Final = object_value(parsed["error_information"]) assert error_information["error_provider_request_id"] == amzn_request_id, error_information + + +@pytest.mark.parametrize("shape", _SHAPES) +@pytest.mark.parametrize("error_code", _ERROR_CODES) +def test_responses_stream_failure_log_retains_provider_request_id( + gateway: Gateway, shape: str, error_code: str +) -> None: + identity: Final = f"resp_{uuid.uuid4().hex}" + provider_request_id: Final = str(uuid.uuid4()) + message: Final = f"stream failure {identity}" + + def respond(request: Request) -> Reply: + assert request.method == "POST" and request.target == "/responses", request.target + assert json.loads(request.body)["stream"] is True + return Reply( + headers={"x-amzn-RequestId": provider_request_id}, + content_type="text/event-stream", + chunks=_stream(identity, shape, error_code, message), + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model="openai/gpt-6", api_base=wire.url, num_retries=0) + response: Final = _streamed_response(gateway, model, identity) + assert message in response.text, response.text + assert ("partial output" in response.text) is (shape != "error-before-output"), response.text + assert response.headers.get("llm_provider-x-amzn-requestid") == provider_request_id, dict(response.headers) + assert len(wire.drain()) == 1 + information: Final = _failure_information(response.headers["x-litellm-call-id"]) + assert information["error_provider_request_id"] == provider_request_id, information + + +@pytest.mark.parametrize("shape", _SHAPES) +def test_responses_stream_failure_without_provider_request_id_logs_none(gateway: Gateway, shape: str) -> None: + identity: Final = f"resp_{uuid.uuid4().hex}" + message: Final = f"stream failure {identity}" + + def respond(request: Request) -> Reply: + assert request.method == "POST" and request.target == "/responses", request.target + return Reply(content_type="text/event-stream", chunks=_stream(identity, shape, "server_error", message)) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model="openai/gpt-6", api_base=wire.url, num_retries=0) + response: Final = _streamed_response(gateway, model, identity) + assert message in response.text, response.text + assert "llm_provider-x-amzn-requestid" not in response.headers, dict(response.headers) + assert len(wire.drain()) == 1 + information: Final = _failure_information(response.headers["x-litellm-call-id"]) + assert information["error_class"] == "InternalServerError", information + assert information["error_provider_request_id"] is None, information + + +@pytest.mark.parametrize("shape", ("nested-error", "response.failed")) +def test_bedrock_responses_stream_server_error_logs_amzn_request_id(gateway: Gateway, shape: str) -> None: + identity: Final = f"resp_{uuid.uuid4().hex}" + amzn_request_id: Final = str(uuid.uuid4()) + message: Final = f"bedrock stream failure {identity}" + + def respond(request: Request) -> Reply: + assert request.method == "POST" and request.target == "/openai/v1/responses", request.target + assert request.headers["authorization"] == f"Bearer {_TOKEN}", request.headers + assert json.loads(request.body)["stream"] is True + return Reply( + headers={"x-amzn-RequestId": amzn_request_id}, + content_type="text/event-stream", + chunks=_stream(identity, shape, "server_error", message), + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=_BEDROCK_RESPONSES_MODEL, + api_key=_TOKEN, + api_base=wire.url, + aws_region_name="us-east-1", + num_retries=0, + ) + response: Final = _streamed_response(gateway, model, identity) + assert "partial output" in response.text and message in response.text, response.text + assert response.headers.get("llm_provider-x-amzn-requestid") == amzn_request_id, dict(response.headers) + assert len(wire.drain()) == 1 + information: Final = _failure_information(response.headers["x-litellm-call-id"]) + assert information["llm_provider"] == "bedrock", information + assert information["error_provider_request_id"] == amzn_request_id, information + + +def test_responses_stream_failure_log_reads_openai_x_request_id(gateway: Gateway) -> None: + identity: Final = f"resp_{uuid.uuid4().hex}" + provider_request_id: Final = f"req_{uuid.uuid4().hex}" + message: Final = f"stream failure {identity}" + + def respond(request: Request) -> Reply: + assert request.method == "POST" and request.target == "/responses", request.target + return Reply( + headers={"x-request-id": provider_request_id}, + content_type="text/event-stream", + chunks=_stream(identity, "top-level-error", "server_error", message), + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model="openai/gpt-6", api_base=wire.url, num_retries=0) + response: Final = _streamed_response(gateway, model, identity) + assert message in response.text, response.text + assert len(wire.drain()) == 1 + information: Final = _failure_information(response.headers["x-litellm-call-id"]) + assert information["error_provider_request_id"] == provider_request_id, information + + +@pytest.mark.parametrize("asynchronous", (False, True), ids=("sync", "async")) +@pytest.mark.parametrize("event_type", ("error", "response.failed")) +async def test_responses_sdk_stream_error_retains_provider_request_id(asynchronous: bool, event_type: str) -> None: + identity: Final = f"resp_{uuid.uuid4().hex}" + provider_request_id: Final = str(uuid.uuid4()) + shape: Final = "nested-error" if event_type == "error" else "response.failed" + + def respond(request: Request) -> Reply: + assert request.method == "POST" and request.target == "/responses", request.target + return Reply( + headers={"x-amzn-RequestId": provider_request_id, "retry-after": "7"}, + content_type="text/event-stream", + chunks=_stream(identity, shape, "server_error", identity), + ) + + with wire_server(respond) as wire: + with pytest.raises(MidStreamFallbackError) as raised: + await _consume_sdk_stream(asynchronous, wire.url, identity) + assert raised.value.status_code == 500 + assert raised.value.response.status_code == 500 + assert str(raised.value.response.request.url) == "https://openai.com/v1/" + assert raised.value.generated_content == "partial output" + assert dict(raised.value.response.headers) == {"x-amzn-requestid": provider_request_id} + assert _get_provider_request_id(raised.value) == provider_request_id + assert raised.value.original_exception is not None + assert _get_provider_request_id(raised.value.original_exception) == provider_request_id + assert len(wire.drain()) == 1 + + +def _concurrent_reply(request: Request, failing: tuple[str, ...], recovery: str) -> Reply: + prompt: Final = json.loads(request.body)["input"] + assert prompt in (*failing, recovery) + chunks: Final = ( + healthy_stream(f"resp_{prompt}", "recovered") + if prompt == recovery + else _stream(f"resp_{prompt}", "nested-error", "server_error", prompt) + ) + return Reply(content_type="text/event-stream", headers={"x-amzn-requestid": prompt}, chunks=chunks) + + +def test_concurrent_responses_failures_keep_distinct_request_ids_and_recover(gateway: Gateway) -> None: + identity: Final = uuid.uuid4().hex + prompts: Final = tuple(f"{identity}-{index}" for index in range(4)) + recovery: Final = f"{identity}-recovery" + + with ( + wire_server(lambda request: _concurrent_reply(request, prompts, recovery)) as wire, + gateway.scenario() as scenario, + ): + model: Final = scenario.model(model="openai/gpt-6", api_base=wire.url, num_retries=0) + + def call(prompt: str) -> str: + response: Final = _streamed_response(gateway, model, prompt) + assert (prompt if prompt != recovery else "recovered") in response.text, response.text + return response.headers["x-litellm-call-id"] + + with ThreadPoolExecutor(max_workers=4) as workers: + call_ids: Final = tuple(workers.map(call, prompts)) + recovered_call_id: Final = call(recovery) + assert len(wire.drain()) == 5 + for prompt, call_id in zip(prompts, call_ids, strict=True): + assert _failure_information(call_id)["error_provider_request_id"] == prompt + assert spend_row(recovered_call_id)["status"] == "success" + + +def _bedrock_logged_error(endpoint: BedrockEndpoint, failure: BedrockFailure) -> tuple[str, str]: + shape: Final = _BEDROCK_FAILURES[failure] + if endpoint == "messages-invoke": + return "BedrockError", str(shape.bedrock_status) + return shape.converse_class, "503" if failure == "exception-event" else str(shape.bedrock_status) + + +def _bedrock_client_status(endpoint: BedrockEndpoint, failure: BedrockFailure, after_output: bool) -> int: + logged_code: Final = int(_bedrock_logged_error(endpoint, failure)[1]) + if after_output or endpoint in ("responses", "messages"): + return 200 + if endpoint == "messages-invoke" and failure == "no-known-event": + return 200 + return logged_code + + +@pytest.mark.parametrize(("endpoint", "failure", "after_output"), _BEDROCK_CELLS) +def test_bedrock_event_stream_failure_log_retains_amzn_request_id( + gateway: Gateway, endpoint: BedrockEndpoint, failure: BedrockFailure, after_output: bool +) -> None: + amzn_request_id: Final = str(uuid.uuid4()) + message: Final = f"bedrock stream failure {amzn_request_id}" + invoke: Final = endpoint == "messages-invoke" + frames: Final = (*(_bedrock_output(invoke) if after_output else ()), _bedrock_failure_frame(failure, message)) + + def respond(request: Request) -> Reply: + assert request.method == "POST" and unquote(request.target) in _BEDROCK_STREAM_TARGETS, request.target + return Reply( + headers={"x-amzn-RequestId": amzn_request_id, "retry-after": "7"}, + content_type=_EVENT_STREAM, + body=b"".join(frames), + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"bedrock/invoke/{_INVOKE_MODEL_ID}" if invoke else _MODEL, + api_base=wire.url, + num_retries=0, + **_BEDROCK_AWS, + ) + path, body = _bedrock_stream_request(endpoint, model, message) + with gateway.client.stream( + "POST", path, json=body, headers={"Authorization": f"Bearer {gateway.key}"} + ) as response: + text: Final = "\n".join(line for line in response.iter_lines() if line) + status: Final = response.status_code + response_headers: Final = dict(response.headers) + assert status == _bedrock_client_status(endpoint, failure, after_output), text + assert _bedrock_failure_text(failure, message) in text and ("partial output" in text) is after_output, text + assert len(wire.drain()) == 1 + information: Final = _failure_information(response_headers["x-litellm-call-id"]) + assert (information["error_class"], information["error_code"]) == _bedrock_logged_error(endpoint, failure) + assert information["error_provider_request_id"] == amzn_request_id, information + if status != 200 and not invoke: + assert response_headers.get("llm_provider-x-amzn-requestid") == amzn_request_id, response_headers + assert "llm_provider-retry-after" not in response_headers, response_headers + + +def _bedrock_model(scenario: Scenario, endpoint: BedrockEndpoint, api_base: str) -> str: + return scenario.model( + model=f"bedrock/invoke/{_INVOKE_MODEL_ID}" if endpoint == "messages-invoke" else _MODEL, + api_base=api_base, + num_retries=0, + **_BEDROCK_AWS, + ) + + +def _bedrock_stream( + gateway: Gateway, endpoint: BedrockEndpoint, model: str, prompt: str +) -> tuple[int, str, dict[str, str]]: + path, body = _bedrock_stream_request(endpoint, model, prompt) + with gateway.client.stream("POST", path, json=body, headers={"Authorization": f"Bearer {gateway.key}"}) as response: + text: Final = "\n".join(line for line in response.iter_lines() if line) + return response.status_code, text, dict(response.headers) + + +@pytest.mark.parametrize("after_output", (False, True), ids=("before-output", "after-output")) +@pytest.mark.parametrize("endpoint", _BEDROCK_ENDPOINTS) +def test_bedrock_event_stream_failure_without_amzn_request_id_logs_none( + gateway: Gateway, endpoint: BedrockEndpoint, after_output: bool +) -> None: + message: Final = f"bedrock stream failure {uuid.uuid4()}" + frames: Final = ( + *(_bedrock_output(endpoint == "messages-invoke") if after_output else ()), + _bedrock_failure_frame("exception-event", message), + ) + + def respond(request: Request) -> Reply: + assert unquote(request.target) in _BEDROCK_STREAM_TARGETS, request.target + return Reply(headers={"retry-after": "7"}, content_type=_EVENT_STREAM, body=b"".join(frames)) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + status, text, response_headers = _bedrock_stream( + gateway, endpoint, _bedrock_model(scenario, endpoint, wire.url), message + ) + assert status == _bedrock_client_status(endpoint, "exception-event", after_output), text + assert message in text, text + assert not any("amzn" in name for name in response_headers), response_headers + assert len(wire.drain()) == 1 + information: Final = _failure_information(response_headers["x-litellm-call-id"]) + assert (information["error_class"], information["error_code"]) == _bedrock_logged_error( + endpoint, "exception-event" + ) + assert information["error_provider_request_id"] is None, information + + +@pytest.mark.parametrize("endpoint", ("responses", "messages", "messages-invoke")) +def test_bedrock_non_stream_500_logs_amzn_request_id_on_every_endpoint( + gateway: Gateway, endpoint: BedrockEndpoint +) -> None: + amzn_request_id: Final = str(uuid.uuid4()) + targets: Final = frozenset( + ( + "/model/anthropic.claude-sonnet-4-5-20250929-v1:0/converse", + f"/model/{_INVOKE_MODEL_ID}/invoke", + ) + ) + + def respond(request: Request) -> Reply: + assert request.method == "POST" and unquote(request.target) in targets, request.target + return Reply( + status=500, + headers={"x-amzn-RequestId": amzn_request_id}, + body=b'{"message":"synthetic bedrock failure"}', + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + path, body = _bedrock_stream_request( + endpoint, _bedrock_model(scenario, endpoint, wire.url), amzn_request_id, stream=False + ) + response: Final = gateway.request("POST", path, body) + assert response.status_code == 503, response.text + if endpoint == "responses": + assert response.headers.get("llm_provider-x-amzn-requestid") == amzn_request_id, dict(response.headers) + assert len(wire.drain()) == 1 + information: Final = _failure_information(response.headers["x-litellm-call-id"]) + assert information["error_provider_request_id"] == amzn_request_id, information + + +@pytest.mark.parametrize("amzn_request_id", ("", "a" * 5000, "dup-1, dup-1"), ids=("empty", "5kb", "joined-duplicate")) +def test_bedrock_stream_failure_logs_odd_request_ids_like_the_non_stream_error( + gateway: Gateway, amzn_request_id: str +) -> None: + stream_body: Final = b"".join((*_bedrock_output(False), _bedrock_failure_frame("exception-event", "odd id"))) + + def respond(request: Request) -> Reply: + if unquote(request.target) in _BEDROCK_STREAM_TARGETS: + return Reply(headers={"x-amzn-RequestId": amzn_request_id}, content_type=_EVENT_STREAM, body=stream_body) + return Reply(status=500, headers={"x-amzn-RequestId": amzn_request_id}, body=b'{"message":"odd id"}') + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = _bedrock_model(scenario, "chat", wire.url) + status, text, stream_headers = _bedrock_stream(gateway, "chat", model, "odd id") + path, body = _bedrock_stream_request("chat", model, "odd id", stream=False) + non_stream: Final = gateway.request("POST", path, body) + assert status == 200 and "odd id" in text, text + assert non_stream.status_code == 503, non_stream.text + assert len(wire.drain()) == 2 + streamed: Final = _failure_information(stream_headers["x-litellm-call-id"]) + unstreamed: Final = _failure_information(non_stream.headers["x-litellm-call-id"]) + assert streamed["error_provider_request_id"] == unstreamed["error_provider_request_id"], (streamed, unstreamed) + + +_BURST_CALLS: Final = tuple( + (endpoint, failure, after_output) + for endpoint in _BEDROCK_ENDPOINTS + for failure, shape in _BEDROCK_FAILURES.items() + for after_output in shape.after_output +) + + +def _healthy_bedrock_stream(invoke: bool) -> bytes: + closing: Final = ( + ( + _invoke_chunk({"type": "content_block_stop", "index": 0}), + _invoke_chunk( + {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 1}} + ), + _invoke_chunk({"type": "message_stop"}), + ) + if invoke + else ( + _bedrock_frame("contentBlockStop", {"contentBlockIndex": 0}), + _bedrock_frame("messageStop", {"stopReason": "end_turn"}), + _bedrock_frame("metadata", {"usage": {"inputTokens": 3, "outputTokens": 1, "totalTokens": 4}}), + ) + ) + return b"".join((*_bedrock_output(invoke, "recovered"), *closing)) + + +def test_concurrent_bedrock_stream_failures_on_every_endpoint_keep_distinct_request_ids(gateway: Gateway) -> None: + identity: Final = uuid.uuid4().hex + failing: Final = {f"{identity}-{index:02d}": call for index, call in enumerate(_BURST_CALLS)} + recoveries: Final = {f"{identity}-r{index}": endpoint for index, endpoint in enumerate(_BEDROCK_ENDPOINTS)} + + def respond(request: Request) -> Reply: + (token,) = (token for token in (*failing, *recoveries) if token in request.body.decode()) + if token in recoveries: + body = _healthy_bedrock_stream(recoveries[token] == "messages-invoke") + else: + endpoint, failure, after_output = failing[token] + opening = _bedrock_output(endpoint == "messages-invoke") if after_output else () + body = b"".join((*opening, _bedrock_failure_frame(failure, token))) + return Reply(headers={"x-amzn-RequestId": token, "retry-after": "7"}, content_type=_EVENT_STREAM, body=body) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + models: Final = {endpoint: _bedrock_model(scenario, endpoint, wire.url) for endpoint in _BEDROCK_ENDPOINTS} + + def call(token: str, endpoint: BedrockEndpoint) -> tuple[int, str, dict[str, str]]: + return _bedrock_stream(gateway, endpoint, models[endpoint], token) + + with ThreadPoolExecutor(max_workers=8) as workers: + outcomes: Final = dict( + zip(failing, workers.map(call, failing, (endpoint for endpoint, _, _ in failing.values())), strict=True) + ) + recovered: Final = {token: call(token, endpoint) for token, endpoint in recoveries.items()} + assert len(wire.drain()) == len(failing) + len(recoveries) + for token, (endpoint, failure, after_output) in failing.items(): + status, text, response_headers = outcomes[token] + assert status == _bedrock_client_status(endpoint, failure, after_output), (token, text) + assert _bedrock_failure_text(failure, token) in text, (token, text) + information = _failure_information(response_headers["x-litellm-call-id"]) + assert information["error_provider_request_id"] == token, (token, information) + for token, (status, text, response_headers) in recovered.items(): + assert status == 200 and "recovered" in text, (token, text) + assert spend_row(response_headers["x-litellm-call-id"])["status"] == "success" + + +def _sdk_stream_text( + client_kind: str, asynchronous: bool, gateway: Gateway, model: str, prompt: str +) -> tuple[str, str]: + extra_body: Final[dict[str, JsonValue]] = {"num_retries": 0, "cache": {"no-cache": True}} + turn: Final = [{"role": "user", "content": prompt}] + + def open_stream(client: openai.OpenAI | openai.AsyncOpenAI | anthropic.Anthropic | anthropic.AsyncAnthropic): + match client_kind: + case "openai-chat": + return client.chat.completions.with_raw_response.create( + model=model, messages=turn, stream=True, extra_body=extra_body + ) + case "openai-responses": + return client.responses.with_raw_response.create( + model=model, input=prompt, stream=True, extra_body=extra_body + ) + case "anthropic-messages": + return client.messages.with_raw_response.create( + model=model, messages=turn, max_tokens=16, stream=True, extra_body=extra_body + ) + raise AssertionError(client_kind) + + seen: Final[list[str]] = [] + base_url: Final = str(gateway.client.base_url).rstrip("/") + if not asynchronous: + with httpx.Client(timeout=15, trust_env=False) as transport: + client = ( + anthropic.Anthropic(api_key=gateway.key, base_url=base_url, max_retries=0, http_client=transport) + if client_kind == "anthropic-messages" + else openai.OpenAI(api_key=gateway.key, base_url=f"{base_url}/v1", max_retries=0, http_client=transport) + ) + raw = open_stream(client) + try: + seen.extend(repr(event) for event in raw.parse()) + except (openai.APIError, anthropic.APIError) as error: + seen.append(str(error)) + return raw.headers["x-litellm-call-id"], "\n".join(seen) + + async def run() -> str: + async with httpx.AsyncClient(timeout=15, trust_env=False) as transport: + client = ( + anthropic.AsyncAnthropic(api_key=gateway.key, base_url=base_url, max_retries=0, http_client=transport) + if client_kind == "anthropic-messages" + else openai.AsyncOpenAI( + api_key=gateway.key, base_url=f"{base_url}/v1", max_retries=0, http_client=transport + ) + ) + raw = await open_stream(client) + try: + async for event in raw.parse(): + seen.append(repr(event)) + except (openai.APIError, anthropic.APIError) as error: + seen.append(str(error)) + return raw.headers["x-litellm-call-id"] + + return asyncio.run(run()), "\n".join(seen) + + +@pytest.mark.parametrize("asynchronous", (False, True), ids=("sync", "async")) +@pytest.mark.parametrize("client_kind", ("openai-chat", "openai-responses", "anthropic-messages")) +def test_sdk_clients_streaming_bedrock_failure_through_proxy_log_amzn_request_id( + gateway: Gateway, client_kind: str, asynchronous: bool +) -> None: + amzn_request_id: Final = str(uuid.uuid4()) + message: Final = f"bedrock stream failure {amzn_request_id}" + frames: Final = (*_bedrock_output(False), _bedrock_failure_frame("exception-event", message)) + + def respond(request: Request) -> Reply: + assert unquote(request.target) in _BEDROCK_STREAM_TARGETS, request.target + return Reply(headers={"x-amzn-RequestId": amzn_request_id}, content_type=_EVENT_STREAM, body=b"".join(frames)) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + call_id, seen = _sdk_stream_text( + client_kind, asynchronous, gateway, _bedrock_model(scenario, "chat", wire.url), message + ) + assert "partial output" in seen and message in seen, seen + assert len(wire.drain()) == 1 + information: Final = _failure_information(call_id) + assert information["error_provider_request_id"] == amzn_request_id, information + + +async def _consume_bedrock_sdk_stream(api: Literal["chat", "responses"], asynchronous: bool, api_base: str) -> str: + common: Final[dict[str, JsonValue]] = { + "model": _MODEL, + "stream": True, + "api_base": api_base, + "num_retries": 0, + **_BEDROCK_AWS, + } + turn: Final[list[JsonValue]] = [{"role": "user", "content": "probe"}] + match api, asynchronous: + case "chat", True: + return "".join([repr(chunk) async for chunk in await litellm.acompletion(messages=turn, **common)]) + case "chat", False: + return "".join(repr(chunk) for chunk in litellm.completion(messages=turn, **common)) + case "responses", True: + return "".join([repr(event) async for event in await litellm.aresponses(input="probe", **common)]) + case "responses", False: + return "".join(repr(event) for event in litellm.responses(input="probe", **common)) + + +@pytest.mark.parametrize("asynchronous", (False, True), ids=("sync", "async")) +@pytest.mark.parametrize("api", ("chat", "responses")) +async def test_bedrock_sdk_stream_exception_carries_amzn_request_id( + api: Literal["chat", "responses"], asynchronous: bool +) -> None: + amzn_request_id: Final = str(uuid.uuid4()) + message: Final = f"bedrock stream failure {uuid.uuid4()}" + frames: Final = (*_bedrock_output(False), _bedrock_failure_frame("exception-event", message)) + + def respond(request: Request) -> Reply: + assert unquote(request.target) in _BEDROCK_STREAM_TARGETS, request.target + return Reply( + headers={"x-amzn-RequestId": amzn_request_id, "retry-after": "7"}, + content_type=_EVENT_STREAM, + body=b"".join(frames), + ) + + with wire_server(respond) as wire: + with pytest.raises(MidStreamFallbackError) as raised: + await _consume_bedrock_sdk_stream(api, asynchronous, wire.url) + assert len(wire.drain()) == 1 + provider_error: Final = raised.value.original_exception + assert isinstance(provider_error, litellm.ServiceUnavailableError), provider_error + assert raised.value.generated_content == "partial output" + assert message in str(provider_error) + assert _get_provider_request_id(provider_error) == amzn_request_id + assert "retry-after" not in provider_error.response.headers, provider_error.response.headers + + +@pytest.mark.parametrize("api", ("chat", "responses")) +@pytest.mark.parametrize("healthy", (True, False), ids=("healthy", "failing")) +async def test_sync_bedrock_sdk_stream_with_non_ascii_request_id_header_keeps_streaming( + api: Literal["chat", "responses"], healthy: bool +) -> None: + body: Final = ( + _healthy_bedrock_stream(False) + if healthy + else b"".join((*_bedrock_output(False), _bedrock_failure_frame("exception-event", "non-ascii id failure"))) + ) + + def respond(request: Request) -> Reply: + assert unquote(request.target) in _BEDROCK_STREAM_TARGETS, request.target + return Reply(headers={"x-amzn-RequestId": "req-\u00fc-1"}, content_type=_EVENT_STREAM, body=body) + + with wire_server(respond) as wire: + if healthy: + assert "recovered" in await _consume_bedrock_sdk_stream(api, False, wire.url) + else: + with pytest.raises(MidStreamFallbackError) as raised: + await _consume_bedrock_sdk_stream(api, False, wire.url) + assert "non-ascii id failure" in str(raised.value.original_exception) + assert _get_provider_request_id(raised.value.original_exception) is None + assert len(wire.drain()) == 1 diff --git a/tests/unit/litellm_core_utils/llm_response_utils/test_get_headers.py b/tests/unit/litellm_core_utils/llm_response_utils/test_get_headers.py index f52e837136a..1eccf5010ec 100644 --- a/tests/unit/litellm_core_utils/llm_response_utils/test_get_headers.py +++ b/tests/unit/litellm_core_utils/llm_response_utils/test_get_headers.py @@ -8,6 +8,7 @@ from litellm.litellm_core_utils.llm_response_utils.get_headers import ( _get_llm_provider_headers, get_response_headers, get_provider_request_id, + provider_request_id_headers, ) from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome @@ -89,6 +90,31 @@ def test_invalid_provider_request_ids_remain_absent(headers: object) -> None: assert get_provider_request_id(headers) is None +@pytest.mark.parametrize( + ("headers", "expected"), + ( + ( + {"X-Amzn-RequestId": "amzn-1", "retry-after": "7", "content-type": "text/event-stream"}, + {"x-amzn-requestid": "amzn-1"}, + ), + ( + {"x-request-id": "req_1", "cf-ray": "ray-1", "x-ratelimit-remaining-requests": "0"}, + {"x-request-id": "req_1", "cf-ray": "ray-1"}, + ), + ({"retry-after": "7"}, None), + ({"x-request-id": 7, "request-id": None, "x-amzn-requestid": "amzn-2"}, {"x-amzn-requestid": "amzn-2"}), + ({"x-amzn-requestid": "amzn-\u00fc", "x-request-id": "req_2"}, {"x-request-id": "req_2"}), + ({}, None), + (None, None), + ), +) +def test_provider_request_id_headers_keeps_only_request_id_headers( + headers: dict[str, object] | None, expected: dict[str, str] | None +) -> None: + result = provider_request_id_headers(headers) + assert (None if result is None else dict(result)) == expected + + @pytest.fixture(autouse=True) def _vcr_outcome_gate(request, vcr): install_live_call_probe(request, vcr) diff --git a/tests/unit/llms/bedrock/chat/test_invoke_handler.py b/tests/unit/llms/bedrock/chat/test_invoke_handler.py index 0609f122949..069438e25c0 100644 --- a/tests/unit/llms/bedrock/chat/test_invoke_handler.py +++ b/tests/unit/llms/bedrock/chat/test_invoke_handler.py @@ -16,6 +16,7 @@ import litellm from litellm import AmazonInvokeConfig from litellm.exceptions import MidStreamFallbackError from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.litellm_logging import _get_provider_request_id from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.llms.bedrock.chat.invoke_handler import ( AmazonOpenAICompatibleStreamDecoder, @@ -864,14 +865,16 @@ async def test_async_invoke_streaming_non_200_forwards_bedrock_response_headers( def _event_stream_frame(event_type: str, payload: bytes) -> bytes: + return _typed_event_stream_frame( + {":event-type": event_type, ":content-type": "application/json", ":message-type": "event"}, payload + ) + + +def _typed_event_stream_frame(frame_headers: Mapping[str, str], payload: bytes) -> bytes: def header(name: str, value: str) -> bytes: return bytes([len(name)]) + name.encode() + bytes([7]) + struct.pack(">H", len(value)) + value.encode() - headers: Final = ( - header(":event-type", event_type) - + header(":content-type", "application/json") - + header(":message-type", "event") - ) + headers: Final = b"".join(header(name, value) for name, value in frame_headers.items()) prelude: Final = struct.pack(">II", 12 + len(headers) + len(payload) + 4, len(headers)) body: Final = prelude + struct.pack(">I", binascii.crc32(prelude)) + headers + payload return body + struct.pack(">I", binascii.crc32(body)) @@ -1367,6 +1370,87 @@ def test_iter_bytes_keeps_a_stream_whose_unknown_event_sits_beside_known_frames( assert "stop" in finish_reasons +_STREAM_FAILURE_BODIES: Final = ( + pytest.param( + _converse_event_frame("internalServerException", {"message": _UPSTREAM_REJECTION}), 500, id="exception-event" + ), + pytest.param( + _typed_event_stream_frame( + { + ":exception-type": "throttlingException", + ":content-type": "application/json", + ":message-type": "exception", + }, + json.dumps({"message": _UPSTREAM_REJECTION}).encode(), + ), + 429, + id="exception-message", + ), + pytest.param(_unknown_event_frame(), 502, id="no-known-event"), + pytest.param(b"", 502, id="no-events"), + pytest.param(_CONVERSE_CONTENT_FRAMES[0] + _truncated_frame(), 502, id="undecoded-tail"), +) + + +def _stream_headers(request_id: str | None) -> httpx.Headers: + request_id_header: Final = {} if request_id is None else {"x-amzn-RequestId": request_id} + return httpx.Headers( + {"content-type": "application/vnd.amazon.eventstream", "retry-after": "7", **request_id_header} + ) + + +def _assert_carries_only_the_request_id(error: BedrockError, status: int, request_id: str | None) -> None: + assert error.status_code == status + assert _get_provider_request_id(error) == request_id + if request_id is None: + assert error.headers is None + return + assert dict(error.headers or {}) == {"x-amzn-requestid": request_id} + assert dict(error.response.headers) == {"x-amzn-requestid": request_id} + + +@pytest.mark.parametrize("request_id", ("req-stream-1", None)) +@pytest.mark.parametrize(("body", "status"), _STREAM_FAILURE_BODIES) +def test_iter_bytes_stream_error_carries_only_the_provider_request_id_header( + body: bytes, status: int, request_id: str | None +) -> None: + decoder: Final = AWSEventStreamDecoder(model="us.moonshotai.kimi-k3") + + with pytest.raises(BedrockError) as exc_info: + list(decoder.iter_bytes(iter([body]), response_headers=_stream_headers(request_id))) + + _assert_carries_only_the_request_id(exc_info.value, status, request_id) + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("body", "status"), _STREAM_FAILURE_BODIES) +async def test_aiter_bytes_stream_error_carries_only_the_provider_request_id_header(body: bytes, status: int) -> None: + async def _chunks() -> AsyncIterator[bytes]: + yield body + + decoder: Final = AWSEventStreamDecoder(model="us.moonshotai.kimi-k3") + + with pytest.raises(BedrockError) as exc_info: + _ = [chunk async for chunk in decoder.aiter_bytes(_chunks(), response_headers=_stream_headers("req-stream-2"))] + + _assert_carries_only_the_request_id(exc_info.value, status, "req-stream-2") + + +@pytest.mark.parametrize("request_id", ("req-\u00fc-1", "req-\udcfc-1"), ids=("latin-1", "surrogate-escaped")) +def test_iter_bytes_streams_and_fails_normally_with_a_non_ascii_provider_request_id(request_id: str) -> None: + decoder: Final = AWSEventStreamDecoder(model="us.moonshotai.kimi-k3") + headers: Final = {"content-type": "application/vnd.amazon.eventstream", "x-amzn-RequestId": request_id} + + chunks: Final = list(decoder.iter_bytes(iter(_CONVERSE_CONTENT_FRAMES), response_headers=headers)) + with pytest.raises(BedrockError) as exc_info: + list(decoder.iter_bytes(iter([_unknown_event_frame()]), response_headers=headers)) + + texts: Final = [chunk.choices[0].delta.content for chunk in chunks if isinstance(chunk, ModelResponseStream)] + assert "".join(text or "" for text in texts) == "hi" + assert exc_info.value.status_code == 502 + assert exc_info.value.headers is None + + def _assert_exception_event_surfaced_with_its_modeled_status(error: BaseException, event_type: str) -> None: assert not isinstance(error, litellm.BadGatewayError) assert getattr(error, "status_code", None) == _modeled_status(event_type) diff --git a/tests/unit/responses/test_streaming_iterator_error_events.py b/tests/unit/responses/test_streaming_iterator_error_events.py index c37e57a43ca..001be7e1fbd 100644 --- a/tests/unit/responses/test_streaming_iterator_error_events.py +++ b/tests/unit/responses/test_streaming_iterator_error_events.py @@ -20,14 +20,18 @@ Pydantic ValidationError (previously typed as Optional[str]). """ import json +from datetime import datetime from importlib import import_module +from typing import Final from unittest.mock import Mock, patch +import httpx import pytest import litellm from litellm.exceptions import MidStreamFallbackError from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.litellm_logging import _get_provider_request_id from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig from litellm.responses.streaming_iterator import ( _ERROR_CODE_HTTP_STATUS, @@ -64,6 +68,52 @@ def _make_error_chunk(error_type: str, code: str, message: str = "err") -> Error return ErrorEvent(type=ResponsesAPIStreamEvents.ERROR, sequence_number=0, error=error_obj) +@pytest.mark.parametrize("provider", ("openai", "azure", "bedrock")) +@pytest.mark.parametrize("header", ("x-amzn-RequestId", "x-request-id", None)) +@pytest.mark.parametrize( + "code,status", (("server_error", 500), ("rate_limit_exceeded", 429), ("invalid_request_error", 400)) +) +def test_stream_error_preserves_provider_request_id(provider: str, header: str | None, code: str, status: int) -> None: + stream_headers: Final = {"retry-after": "7", "content-type": "text/event-stream"} + response: Final = httpx.Response( + 200, headers={**stream_headers, header: "provider-request-id"} if header else stream_headers + ) + logging: Final = LiteLLMLoggingObj( + model="test-model", + messages=[], + stream=True, + call_type="responses", + start_time=datetime(2026, 1, 1), + litellm_call_id="test-call", + function_id="test-function", + ) + iterator: Final = BaseResponsesAPIStreamingIterator( + response=response, + model="test-model", + responses_api_provider_config=None, + logging_obj=logging, + custom_llm_provider=provider, + ) + expected_id: Final = "provider-request-id" if header else None + expected_headers: Final = {header.lower(): "provider-request-id"} if header else {} + mapped: Final = iterator._map_error_event_exception({"type": code, "code": code, "message": "failure"}) + assert _get_provider_request_id(mapped) == expected_id + assert dict(getattr(mapped, "litellm_response_headers", None) or {}) == expected_headers + expected_status: Final = 503 if provider == "bedrock" and status == 500 else status + with pytest.raises((litellm.BadRequestError, MidStreamFallbackError)) as raised: + iterator._maybe_raise_for_error_event(_make_error_chunk(code, code, "failure")) + assert raised.value.status_code == expected_status + assert _get_provider_request_id(raised.value) == expected_id + assert response.status_code == 200 + if status in (429, 500): + assert isinstance(raised.value, MidStreamFallbackError) + assert dict(raised.value.response.headers) == expected_headers + assert raised.value.response.status_code == expected_status + assert str(raised.value.response.request.url) == f"https://{provider}.com/v1/" + assert raised.value.generated_content == "" + assert raised.value.is_pre_first_chunk is True + + def test_maybe_raise_for_error_event_wraps_unknown_error_in_mid_stream_fallback(): iterator = _make_iterator() chunk = _make_error_chunk("server_error", "internal_error", "something went wrong") From f281f77bc02faf4f008e380d46aa4148dae61bf5 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Fri, 16 Oct 2026 00:37:03 -0700 Subject: [PATCH 07/23] fix(guardrails): record every content filter match in logging_only and offer the mode in the UI (#46673) * test(guardrails): pin content filter logging_only matches and UI modes in an integration test * fix(guardrails): record every content filter match in logging_only and offer the mode in the UI * style(tests): format the content filter logging_only integration test * refactor(guardrails): keep the logging_only predicate and scan override within the type gates * fix(guardrails): flag tool arguments a content filter would reject in logging_only and keep scanning * test(integration): expect guardrail_flagged for an observe-only content filter block in the scope config test * test(integration): pin logging_only flags for tool arguments an enforcing content filter rejects * refactor(guardrails): drop the content filter flagged branch no logging_only scan can reach * test(integration): pin OTel spans, Guardrails Monitor and a provider outage burst for a logging_only content filter * test(integration): pin MCP tool call verdicts and fail the Monitor cell when usage never rolls up * test(integration): pin category and competitor intent verdicts and the provider failure row in logging_only --------- Co-authored-by: yucheng-berri <293926326+yucheng-berri@users.noreply.github.com> --- litellm/integrations/custom_guardrail.py | 19 +- .../proxy/guardrails/guardrail_endpoints.py | 13 +- .../litellm_content_filter/content_filter.py | 75 +- .../test_content_filter_logging_only.py | 712 ++++++++++++++++++ .../test_logging_only_scope_config.py | 12 +- .../integrations/test_custom_guardrail.py | 2 +- .../test_content_filter.py | 240 ++++++ .../guardrails/test_guardrail_endpoints.py | 33 +- 8 files changed, 1071 insertions(+), 35 deletions(-) create mode 100644 tests/integration/observability/test_content_filter_logging_only.py diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 2a2a2a0326d..93dd3d0f5ee 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -695,9 +695,7 @@ class CustomGuardrail(CustomLogger): supported_event_hooks: list[GuardrailEventHooks], ) -> None: allowed_hooks: Final = frozenset(supported_event_hooks) | ( - frozenset((GuardrailEventHooks.logging_only,)) - if self.uses_apply_guardrail_interface() and not self.use_native_lifecycle_hooks - else frozenset() + frozenset((GuardrailEventHooks.logging_only,)) if self.accepts_logging_only_mode() else frozenset() ) def _validate_event_hook_list_is_in_supported_event_hooks( @@ -869,8 +867,13 @@ class CustomGuardrail(CustomLogger): return True return False - def uses_apply_guardrail_interface(self) -> bool: - return type(self).apply_guardrail is not CustomGuardrail.apply_guardrail + @classmethod + def uses_apply_guardrail_interface(cls) -> bool: + return cls.apply_guardrail is not CustomGuardrail.apply_guardrail + + @classmethod + def accepts_logging_only_mode(cls) -> bool: + return cls.uses_apply_guardrail_interface() and not cls.use_native_lifecycle_hooks @classmethod def supports_logging_only_scope(cls) -> bool: @@ -1029,7 +1032,7 @@ class CustomGuardrail(CustomLogger): ) return kwargs, result litellm_params: Final = kwargs.get("litellm_params") or {} - scratch_metadata: Final = { + scratch_metadata: Final[dict[str, object]] = { # mutable-ok: apply_guardrail records its verdict here key: value for key, value in (litellm_params.get("metadata") or {}).items() if key != "standard_logging_guardrail_information" @@ -1083,11 +1086,11 @@ class CustomGuardrail(CustomLogger): async def _scan_logged_call( self, - kwargs: dict, # mutable-ok: CustomLogger.async_logging_hook contract + kwargs: dict[str, object], # mutable-ok: CustomLogger.async_logging_hook contract response: object | None, translation: "BaseTranslation", output_translation: "BaseTranslation", - scratch_metadata: dict, # mutable-ok: apply_guardrail records its verdict into request metadata + scratch_metadata: dict[str, object], # mutable-ok: apply_guardrail records its verdict into request metadata ) -> None: scratch_fields: Final = self._copy_scratch_request_fields(kwargs) scratch_input, scratch_tools = scratch_fields or (None, None) diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index d713a18f470..1d5c530a616 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -1483,6 +1483,17 @@ async def get_guardrail_info(guardrail_id: str): raise HTTPException(status_code=500, detail=str(e)) +def _undeclared_logging_only( + guardrail_class: type[CustomGuardrail], hooks: Sequence[GuardrailEventHooks] +) -> tuple[GuardrailEventHooks, ...]: + """logging_only when the class accepts it at save time without declaring it in its hooks.""" + return ( + (GuardrailEventHooks.logging_only,) + if guardrail_class.accepts_logging_only_mode() and GuardrailEventHooks.logging_only not in hooks + else () + ) + + @router.get( "/guardrails/ui/add_guardrail_settings", tags=["Guardrails"], @@ -1515,7 +1526,7 @@ async def get_guardrail_ui_settings() -> GuardrailUIAddGuardrailSettings: ] supported_modes_by_provider: Final = { - provider: [hook.value for hook in hooks] + provider: [hook.value for hook in (*hooks, *_undeclared_logging_only(guardrail_class, hooks))] for provider, guardrail_class in guardrail_class_registry.items() if (hooks := guardrail_class.get_supported_event_hooks()) is not None } | MappingProxyType( diff --git a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py index f959292d567..40e44151e65 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py +++ b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py @@ -6,6 +6,7 @@ to detect and block/mask sensitive content. """ import asyncio +import contextvars import itertools import json import os @@ -50,6 +51,7 @@ from litellm.types.utils import ( if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation from litellm.types.guardrails import ( BlockedWord, @@ -101,6 +103,13 @@ WORD_NUMBER_SEQUENCE_PATTERN: Final = re.compile( ) WORD_NUMBER_TOKEN_FINDER: Final = re.compile(rf"(?:{WORD_NUMBER_TOKEN_REGEX})", re.IGNORECASE) +_logging_only_scan_guardrail: Final[contextvars.ContextVar["ContentFilterGuardrail | None"]] = contextvars.ContextVar( + "litellm_content_filter_logging_only_scan", default=None +) +_logging_only_scan_would_reject: Final[contextvars.ContextVar[bool]] = contextvars.ContextVar( + "litellm_content_filter_logging_only_scan_would_reject", default=False +) + class ConditionalCategoryConfig(TypedDict): identifier_words: Sequence[str] @@ -1209,6 +1218,18 @@ class ContentFilterGuardrail(CustomGuardrail): return (keyword, action, description) return None + def _is_logging_only_scan(self) -> bool: + return _logging_only_scan_guardrail.get() is self + + def _enforces_block(self, action: ContentFilterAction) -> bool: + return action == ContentFilterAction.BLOCK and not self._is_logging_only_scan() + + @staticmethod + def _asks_to_block(detection: ContentFilterDetection) -> bool: + if detection["type"] == "competitor_intent": + return detection["action_hint"] in ("refuse", "reframe") + return detection["action"] == ContentFilterAction.BLOCK.value + def _handle_conditional_match( self, matched_phrase: str, @@ -1228,7 +1249,7 @@ class ContentFilterGuardrail(CustomGuardrail): } detections.append(category_detection) - if action == ContentFilterAction.BLOCK: + if self._enforces_block(action): error_msg: Final = ( f"Content blocked: {category_name} conditional match '{matched_phrase}' detected (severity: {severity})" ) @@ -1269,7 +1290,7 @@ class ContentFilterGuardrail(CustomGuardrail): } detections.append(category_detection) - if action == ContentFilterAction.BLOCK: + if self._enforces_block(action): error_msg = f"Content blocked: {category_name} category keyword '{keyword}' detected (severity: {severity})" verbose_proxy_logger.warning(error_msg) raise HTTPException( @@ -1312,7 +1333,7 @@ class ContentFilterGuardrail(CustomGuardrail): } detections.append(pattern_detection) - if action == ContentFilterAction.BLOCK: + if self._enforces_block(action): error_msg: Final = f"Content blocked: {pattern_name} pattern detected" verbose_proxy_logger.warning(error_msg) raise HTTPException( @@ -1346,7 +1367,7 @@ class ContentFilterGuardrail(CustomGuardrail): } detections.append(blocked_word_detection) - if action == ContentFilterAction.BLOCK: + if self._enforces_block(action): error_msg = f"Content blocked: keyword '{keyword}' detected" if description: error_msg += f" ({description})" @@ -1619,7 +1640,7 @@ class ContentFilterGuardrail(CustomGuardrail): - Number of entities masked - Whether the guardrail blocked the request (max risk) """ - if status == "guardrail_intervened": + if status in ("guardrail_intervened", "guardrail_flagged"): return 10.0 total_masked: Final = sum(masked_entity_count.values()) if masked_entity_count else 0 @@ -1665,6 +1686,8 @@ class ContentFilterGuardrail(CustomGuardrail): } detections.append(detection) + if self._is_logging_only_scan(): + return if action_hint_val == "refuse": msg = "Content blocked: competitor comparison or ranking intent detected." if self._competitor_intent_checker and getattr( @@ -1763,18 +1786,18 @@ class ContentFilterGuardrail(CustomGuardrail): return raw_name return None + def _reject(self, error: str) -> None: + """Raise the 400 an enforcing mode returns, or note it on the leg a logging_only scan is observing.""" + if self._is_logging_only_scan(): + _logging_only_scan_would_reject.set(True) + return + raise HTTPException(status_code=400, detail={"error": error}) + def _assert_argument_label_clean( self, text: str, detections: list[ContentFilterDetection], context_label: str ) -> None: if self._filter_single_text(text, detections=detections) != text: - raise HTTPException( - status_code=400, - detail={ - "error": ( - f"Content blocked: {context_label} argument matched a masking rule on a non-rewritable field" - ) - }, - ) + self._reject(f"Content blocked: {context_label} argument matched a masking rule on a non-rewritable field") def _filter_argument_value( self, @@ -1784,10 +1807,8 @@ class ContentFilterGuardrail(CustomGuardrail): depth: int = 0, ) -> object: if depth > DEFAULT_MAX_RECURSE_DEPTH: - raise HTTPException( - status_code=400, - detail={"error": f"Content blocked: {context_label} arguments exceed the maximum nesting depth"}, - ) + self._reject(f"Content blocked: {context_label} arguments exceed the maximum nesting depth") + return value if isinstance(value, str): return self._filter_single_text(value, detections=detections) if isinstance(value, (int, float)) and not isinstance(value, bool): @@ -1874,6 +1895,21 @@ class ContentFilterGuardrail(CustomGuardrail): if filtered_arguments != arguments: self._set_tool_call_arguments(tool_call, filtered_arguments) + async def _scan_logged_call( + self, + kwargs: dict[str, object], # mutable-ok: CustomLogger.async_logging_hook contract + response: object | None, + translation: "BaseTranslation", + output_translation: "BaseTranslation", + scratch_metadata: dict[str, object], # mutable-ok: apply_guardrail records its verdict into request metadata + ) -> None: + """Scan the logged call with every BLOCK match recorded instead of raised.""" + token: Final = _logging_only_scan_guardrail.set(self) + try: + await super()._scan_logged_call(kwargs, response, translation, output_translation, scratch_metadata) + finally: + _logging_only_scan_guardrail.reset(token) + async def apply_guardrail( self, inputs: "GenericGuardrailAPIInputs", @@ -1904,6 +1940,7 @@ class ContentFilterGuardrail(CustomGuardrail): masked_entity_count: Final[dict[str, int]] = {} status: GuardrailStatus = "success" exception_str: str = "" + would_reject_token: Final = _logging_only_scan_would_reject.set(False) try: texts: Final = inputs.get("texts", []) @@ -1947,11 +1984,13 @@ class ContentFilterGuardrail(CustomGuardrail): exception_str = str(e) raise e finally: + would_reject: Final = _logging_only_scan_would_reject.get() or any(map(self._asks_to_block, detections)) + _logging_only_scan_would_reject.reset(would_reject_token) # Log guardrail information self._log_guardrail_information( request_data=request_data, detections=detections, - status=status, + status="guardrail_flagged" if status == "success" and would_reject else status, start_time=start_time, masked_entity_count=masked_entity_count, exception_str=exception_str, diff --git a/tests/integration/observability/test_content_filter_logging_only.py b/tests/integration/observability/test_content_filter_logging_only.py new file mode 100644 index 00000000000..ab1779c497d --- /dev/null +++ b/tests/integration/observability/test_content_filter_logging_only.py @@ -0,0 +1,712 @@ +"""`litellm_content_filter` in `mode: logging_only`, driven through a real proxy. + +The provider is the only double, an owned wire peer. The proxy, its guardrail registry, Postgres and Redis run +for real with two workers. Every check reads the spend row, which is written after the logging-only scans finish +""" + +from __future__ import annotations + +import json +import uuid +from collections.abc import Iterator +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from datetime import datetime, timezone +from hashlib import sha256 +from pathlib import Path +from typing import Final + +import anthropic +import openai +import pytest +import yaml +from integration._support.client import Gateway, Scenario, eventually, gateway_from_environment, object_value +from integration._support.database import read_rows +from integration._support.mcp import EntryPoint, McpCaller, echo_tool, register_mcp, scripted_peer, tool_calls +from integration._support.otlp_sink import AuditConfigWriter, SpanSinks, recorded_spans +from integration._support.process import owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server + +AWS_KEY: Final = "AKIA" + "SYNTHETIC0000000" +GITHUB_TOKEN: Final = "ghp_" + "a1" * 18 +FORBIDDEN_WORD: Final = "synthetic-forbidden-word" +LEAK_MARK: Final = "SYNTHETIC-LEAK-TOKEN" +OUTAGE_MARK: Final = "SYNTHETIC-PROVIDER-OUTAGE" +OBSERVE_BLOCK: Final = "cf-observe-block" +OBSERVE_MASK: Final = "cf-observe-mask" +ENFORCE_BLOCK: Final = "cf-enforce-block" +ENFORCE_MASK: Final = "cf-enforce-mask" +OBSERVE_CATEGORY: Final = "cf-observe-category" +OBSERVE_COMPETITOR: Final = "cf-observe-competitor" +SECRETS: Final = f"AWS_ACCESS_KEY_ID={AWS_KEY}\nGITHUB_TOKEN={GITHUB_TOKEN}" + + +def _answer(marker: str) -> str: + return "synthetic answer " + marker + + +def _provider_answer(body: bytes) -> str: + text: Final = body.decode() + start: Final = text.find("mark-") + assert start >= 0, text + answer: Final = _answer(text[start : start + 37]) + return answer + " " + GITHUB_TOKEN if LEAK_MARK in text else answer + + +def _sse(events: tuple[dict[str, object], ...]) -> tuple[bytes, ...]: + return tuple(("data: " + json.dumps(event) + "\n\n").encode() for event in events) + + +def _chat_reply(answer: str, streaming: bool) -> Reply: + if not streaming: + return Reply( + body=json.dumps( + { + "id": "chatcmpl-" + uuid.uuid4().hex, + "object": "chat.completion", + "created": 1, + "model": "gpt-5.4-mini", + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": answer}, "finish_reason": "stop"} + ], + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + } + ).encode() + ) + identity: Final = "chatcmpl-" + uuid.uuid4().hex + + def chunk(delta: dict[str, object], finish: str | None) -> dict[str, object]: + return { + "id": identity, + "object": "chat.completion.chunk", + "created": 1, + "model": "gpt-5.4-mini", + "choices": [{"index": 0, "delta": delta, "finish_reason": finish}], + } + + events: Final = (chunk({"role": "assistant", "content": answer[:9]}, None), chunk({"content": answer[9:]}, "stop")) + return Reply(chunks=(*_sse(events), b"data: [DONE]\n\n"), content_type="text/event-stream") + + +def _messages_reply(answer: str) -> Reply: + return Reply( + body=json.dumps( + { + "id": "msg_" + uuid.uuid4().hex, + "type": "message", + "role": "assistant", + "model": "claude-haiku-5-5", + "content": [{"type": "text", "text": answer}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 10, "output_tokens": 5}, + } + ).encode() + ) + + +def _responses_reply(answer: str) -> Reply: + return Reply( + body=json.dumps( + { + "id": "resp_" + uuid.uuid4().hex, + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-5.4-mini", + "output": [ + { + "type": "message", + "id": "msgo_" + uuid.uuid4().hex, + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": answer, "annotations": []}], + } + ], + "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + } + ).encode() + ) + + +def _provider(request: Request) -> Reply: + if OUTAGE_MARK.encode() in request.body: + return Reply(status=503, body=b'{"error": {"message": "synthetic provider outage", "type": "server_error"}}') + answer: Final = _provider_answer(request.body) + if request.target.endswith("/v1/messages"): + return _messages_reply(answer) + if request.target.endswith("/v1/responses"): + return _responses_reply(answer) + assert request.target.endswith("/v1/chat/completions"), request.target + return _chat_reply(answer, bool(json.loads(request.body).get("stream"))) + + +def _content_filter(name: str, mode: str, action: str, *, default_on: bool) -> dict[str, object]: + return { + "guardrail_name": name, + "litellm_params": { + "guardrail": "litellm_content_filter", + "mode": mode, + "default_on": default_on, + "patterns": [ + {"pattern_type": "regex", "pattern": "AKIA[0-9A-Z]{16}", "name": "aws_access_key", "action": action}, + {"pattern_type": "regex", "pattern": "ghp_[A-Za-z0-9]{36}", "name": "github_pat", "action": action}, + ], + "blocked_words": [{"keyword": FORBIDDEN_WORD, "action": action}], + }, + } + + +def _observe_only(name: str, **rules: object) -> dict[str, object]: + return { + "guardrail_name": name, + "litellm_params": {"guardrail": "litellm_content_filter", "mode": "logging_only", "default_on": False, **rules}, + } + + +def _rig_config(root: Path) -> Path: + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config["litellm_settings"]["cache"] = False + config["general_settings"]["store_prompts_in_spend_logs"] = True + config["guardrails"] = [ + _content_filter(OBSERVE_BLOCK, "logging_only", "BLOCK", default_on=True), + _content_filter(OBSERVE_MASK, "logging_only", "MASK", default_on=True), + _content_filter(ENFORCE_BLOCK, "pre_call", "BLOCK", default_on=False), + _content_filter(ENFORCE_MASK, "pre_call", "MASK", default_on=False), + _observe_only( + OBSERVE_CATEGORY, + categories=[ + {"category": "harmful_violence", "enabled": True, "action": "BLOCK", "severity_threshold": "medium"} + ], + blocked_words=[{"keyword": FORBIDDEN_WORD, "action": "BLOCK"}], + ), + _observe_only( + OBSERVE_COMPETITOR, + competitor_intent_config={ + "brand_self": ["emirates"], + "competitors": ["qatar airways"], + "domain_words": ["airline", "flight"], + "policy": {"competitor_comparison": "refuse"}, + "threshold_high": 0.5, + }, + ), + ] + path: Final = root / "content-filter-logging-only.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +@dataclass(frozen=True, slots=True) +class Rig: + proxy: Gateway + provider: Wire + chat_model: str + claude_model: str + responses_model: str + scenario: Scenario + + def base(self) -> str: + return str(self.proxy.client.base_url).rstrip("/") + + def provider_bodies(self, marker: str) -> tuple[str, ...]: + return tuple(request.body.decode() for request in self.provider.drain() if marker.encode() in request.body) + + def entries(self, response_id: str, name: str) -> tuple[dict[str, object], ...]: + return _named_entries( + 'SELECT request_id, status, metadata FROM "LiteLLM_SpendLogs" WHERE request_id = %s', (response_id,), name + ) + + +def _named_entries(query: str, parameters: tuple[str, ...], name: str) -> tuple[dict[str, object], ...]: + rows: Final = eventually(lambda: read_rows(query, parameters), lambda values: len(values) >= 1, seconds=70) + assert len(rows) == 1, rows + recorded: Final = object_value(rows[0]["metadata"]).get("guardrail_information") + assert isinstance(recorded, list), rows[0] + return tuple(entry for entry in (object_value(item) for item in recorded) if entry.get("guardrail_name") == name) + + +def _findings(entry: dict[str, object]) -> tuple[str, ...]: + response: Final = entry.get("guardrail_response") + assert isinstance(response, list), entry + return tuple( + str(detection.get("pattern_name") or detection.get("keyword") or detection.get("type")) + for detection in (object_value(item) for item in response) + ) + + +def _verdicts(entries: tuple[dict[str, object], ...]) -> list[tuple[object, object, tuple[str, ...]]]: + return [(entry["guardrail_mode"], entry["guardrail_status"], _findings(entry)) for entry in entries] + + +@pytest.fixture(scope="module") +def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + root: Final = tmp_path_factory.mktemp("content-filter-logging-only") + with ( + gateway_from_environment() as gateway, + wire_server(_provider) as provider, + owned_proxy_process(gateway, root, {}, config=_rig_config(root), workers=2) as owned, + owned.gateway.scenario() as scenario, + ): + chat: Final = scenario.model( + model="openai/gpt-5.4-mini", api_base=provider.url + "/v1", api_key="synthetic-openai-key" + ) + claude: Final = scenario.model( + model="anthropic/claude-haiku-5-5", api_base=provider.url, api_key="synthetic-anthropic-key" + ) + responses: Final = scenario.model( + model="openai/gpt-5.4-mini", api_base=provider.url + "/v1", api_key="synthetic-openai-key" + ) + yield Rig(owned.gateway, provider, chat, claude, responses, scenario) + + +def _marker() -> str: + return "mark-" + uuid.uuid4().hex + + +def _sync_chat(rig: Rig, text: str) -> tuple[str, str]: + client: Final = openai.OpenAI(base_url=rig.base() + "/v1", api_key=rig.proxy.key, max_retries=0) + completion: Final = client.chat.completions.create( + model=rig.chat_model, messages=[{"role": "user", "content": text}] + ) + return completion.id, completion.choices[0].message.content or "" + + +async def _async_chat(rig: Rig, text: str) -> tuple[str, str]: + client: Final = openai.AsyncOpenAI(base_url=rig.base() + "/v1", api_key=rig.proxy.key, max_retries=0) + completion: Final = await client.chat.completions.create( + model=rig.chat_model, messages=[{"role": "user", "content": text}] + ) + return completion.id, completion.choices[0].message.content or "" + + +def _streamed_chat(rig: Rig, text: str) -> tuple[str, str]: + client: Final = openai.OpenAI(base_url=rig.base() + "/v1", api_key=rig.proxy.key, max_retries=0) + chunks: Final = tuple( + client.chat.completions.create(model=rig.chat_model, messages=[{"role": "user", "content": text}], stream=True) + ) + return chunks[0].id, "".join(chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices) + + +def _messages(rig: Rig, text: str) -> tuple[str, str]: + client: Final = anthropic.Anthropic(base_url=rig.base(), api_key=rig.proxy.key, max_retries=0) + message: Final = client.messages.create( + model=rig.claude_model, max_tokens=64, messages=[{"role": "user", "content": text}] + ) + assert message.content[0].type == "text" + return message.id, message.content[0].text + + +def _responses(rig: Rig, text: str) -> tuple[str, str]: + client: Final = openai.OpenAI(base_url=rig.base() + "/v1", api_key=rig.proxy.key, max_retries=0) + result: Final = client.responses.create(model=rig.responses_model, input=text) + return result.id, result.output_text + + +SYNC_CALLS: Final = {"chat": _sync_chat, "chat_stream": _streamed_chat, "messages": _messages, "responses": _responses} +EVERY_FINDING: Final = ("aws_access_key", "github_pat", FORBIDDEN_WORD) + + +@pytest.mark.parametrize("endpoint", list(SYNC_CALLS)) +def test_logging_only_block_records_every_match_without_touching_the_call(rig: Rig, endpoint: str) -> None: + marker: Final = _marker() + prompt: Final = f"{marker} what did the file contain?\n{SECRETS}\n{FORBIDDEN_WORD}" + response_id, answer = SYNC_CALLS[endpoint](rig, prompt) + assert answer == _answer(marker) + sent: Final = rig.provider_bodies(marker) + assert len(sent) == 1 and AWS_KEY in sent[0] and GITHUB_TOKEN in sent[0] and FORBIDDEN_WORD in sent[0], sent + + assert _verdicts(rig.entries(response_id, OBSERVE_BLOCK)) == [ + ("logging_only", "guardrail_flagged", EVERY_FINDING), + ("logging_only", "success", ()), + ] + + +@pytest.mark.asyncio +async def test_logging_only_block_records_every_match_for_an_async_client(rig: Rig) -> None: + marker: Final = _marker() + response_id, answer = await _async_chat(rig, f"{marker} {SECRETS}") + assert answer == _answer(marker) + + assert _verdicts(rig.entries(response_id, OBSERVE_BLOCK)) == [ + ("logging_only", "guardrail_flagged", ("aws_access_key", "github_pat")), + ("logging_only", "success", ()), + ] + + +def test_logging_only_block_still_scans_the_response_after_a_request_match(rig: Rig) -> None: + marker: Final = _marker() + response_id, answer = _sync_chat(rig, f"{marker} {LEAK_MARK} key {AWS_KEY}") + assert answer == _answer(marker) + " " + GITHUB_TOKEN + + assert _verdicts(rig.entries(response_id, OBSERVE_BLOCK)) == [ + ("logging_only", "guardrail_flagged", ("aws_access_key",)), + ("logging_only", "guardrail_flagged", ("github_pat",)), + ] + + +def _tool_call_history(text: str, arguments: str) -> list[dict[str, object]]: + return [ + {"role": "user", "content": text}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": arguments}} + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "done"}, + ] + + +@pytest.mark.parametrize( + ("arguments", "request_findings"), + ( + pytest.param(json.dumps({AWS_KEY: 1}), ("aws_access_key",), id="masked-argument-key"), + pytest.param("[" * 150 + "1" + "]" * 150, (), id="too-deep"), + ), +) +def test_logging_only_flags_tool_arguments_an_enforcing_mode_rejects_and_keeps_scanning( + rig: Rig, arguments: str, request_findings: tuple[str, ...] +) -> None: + marker: Final = _marker() + messages: Final = _tool_call_history(f"{marker} {LEAK_MARK}", arguments) + headers: Final = {"Authorization": f"Bearer {rig.proxy.key}"} + enforced: Final = rig.proxy.client.post( + "/v1/chat/completions", + json={"model": rig.chat_model, "messages": messages, "guardrails": [ENFORCE_MASK]}, + headers=headers, + ) + assert enforced.status_code == 400, enforced.text + assert rig.provider_bodies(marker) == () + + observed: Final = rig.proxy.client.post( + "/v1/chat/completions", json={"model": rig.chat_model, "messages": messages}, headers=headers + ) + assert observed.status_code == 200, observed.text + completion: Final = object_value(observed.json()) + assert len(rig.provider_bodies(marker)) == 1 + + assert _verdicts(rig.entries(str(completion["id"]), OBSERVE_MASK)) == [ + ("logging_only", "guardrail_flagged", request_findings), + ("logging_only", "success", ("github_pat",)), + ] + + +@pytest.mark.parametrize("entry", ("mcp", "rest")) +def test_logging_only_block_records_every_match_on_an_mcp_tool_call(rig: Rig, entry: EntryPoint) -> None: + alias: Final = "cfmcp" + uuid.uuid4().hex[:8] + arguments: Final = {"note": f"{_marker()} {SECRETS}"} + with scripted_peer(echo_tool("echo")) as peer: + server_id: Final = register_mcp(rig.scenario, peer, alias) + key: Final = rig.scenario.key(object_permission={"mcp_servers": [server_id]}) + outcome: Final = McpCaller(rig.proxy, key, entry, alias).call(f"{alias}-echo", arguments, server_id=server_id) + assert (outcome.error, outcome.text) == (None, json.dumps(arguments, sort_keys=True)), outcome.raw + assert [call["body"]["params"]["arguments"] for call in tool_calls(peer.drain())] == [arguments] + + entries: Final = _named_entries( + 'SELECT metadata FROM "LiteLLM_SpendLogs" WHERE api_key = %s AND call_type = %s', + (sha256(key.encode()).hexdigest(), "call_mcp_tool"), + OBSERVE_BLOCK, + ) + assert _verdicts(entries) == [("logging_only", "guardrail_flagged", ("aws_access_key", "github_pat"))], entries + + +@pytest.mark.parametrize( + ("guardrail", "prompt", "findings"), + ( + pytest.param( + OBSERVE_CATEGORY, + f"how to kill someone {FORBIDDEN_WORD}", + ("how to kill someone", FORBIDDEN_WORD), + id="category-keyword", + ), + pytest.param( + OBSERVE_COMPETITOR, "Is Qatar Airways better than Emirates?", ("competitor_intent",), id="competitor-intent" + ), + ), +) +def test_logging_only_block_flags_category_and_competitor_matches_without_blocking( + rig: Rig, guardrail: str, prompt: str, findings: tuple[str, ...] +) -> None: + marker: Final = _marker() + response: Final = rig.proxy.client.post( + "/v1/chat/completions", + json={ + "model": rig.chat_model, + "messages": [{"role": "user", "content": f"{marker} {prompt}"}], + "guardrails": [guardrail], + }, + headers={"Authorization": f"Bearer {rig.proxy.key}"}, + ) + assert response.status_code == 200, response.text + completion: Final = object_value(response.json()) + assert len(rig.provider_bodies(marker)) == 1 + + assert _verdicts(rig.entries(str(completion["id"]), guardrail)) == [ + ("logging_only", "guardrail_flagged", findings), + ("logging_only", "success", ()), + ] + + +def test_logging_only_mask_records_every_match(rig: Rig) -> None: + marker: Final = _marker() + response_id, answer = _sync_chat(rig, f"{marker} {SECRETS}\n{FORBIDDEN_WORD}") + assert answer == _answer(marker) + + entries: Final = rig.entries(response_id, OBSERVE_MASK) + assert _verdicts(entries) == [ + ("logging_only", "success", EVERY_FINDING), + ("logging_only", "success", ()), + ] + assert entries[0]["masked_entity_count"] == {"aws_access_key": 1, "github_pat": 1, "blocked_word": 1}, entries + + +def test_logging_only_records_nothing_for_a_clean_call(rig: Rig) -> None: + marker: Final = _marker() + response_id, answer = _sync_chat(rig, f"{marker} nothing secret here") + assert answer == _answer(marker) + + assert _verdicts(rig.entries(response_id, OBSERVE_BLOCK)) == [ + ("logging_only", "success", ()), + ("logging_only", "success", ()), + ] + + +def test_pre_call_block_still_rejects_the_request(rig: Rig) -> None: + marker: Final = _marker() + response: Final = rig.proxy.client.post( + "/v1/chat/completions", + json={ + "model": rig.chat_model, + "messages": [{"role": "user", "content": f"{marker} {SECRETS}"}], + "guardrails": [ENFORCE_BLOCK], + }, + headers={"Authorization": f"Bearer {rig.proxy.key}"}, + ) + assert response.status_code == 400, response.text + assert "aws_access_key" in response.text + assert rig.provider_bodies(marker) == () + + +def test_apply_guardrail_endpoint_still_blocks_with_a_logging_only_content_filter(rig: Rig) -> None: + response: Final = rig.proxy.client.post( + "/guardrails/apply_guardrail", + json={"guardrail_name": OBSERVE_BLOCK, "text": f"key {AWS_KEY}"}, + headers={"Authorization": f"Bearer {rig.proxy.key}"}, + ) + assert response.status_code == 400, response.text + assert "aws_access_key" in response.text + + +def _advertised_modes(rig: Rig) -> dict[str, object]: + response: Final = rig.proxy.client.get( + "/guardrails/ui/add_guardrail_settings", headers={"Authorization": f"Bearer {rig.proxy.key}"} + ) + assert response.status_code == 200, response.text + return object_value(response.json()["supported_modes_by_provider"]) + + +@pytest.mark.parametrize("provider", ["litellm_content_filter", "generic_guardrail_api", "bedrock", "akto"]) +def test_dashboard_offers_logging_only_where_the_config_accepts_it(rig: Rig, provider: str) -> None: + modes: Final = _advertised_modes(rig) + offered: Final = modes[provider] + assert isinstance(offered, list) and "logging_only" in offered, modes + assert offered.count("logging_only") == 1, offered + + +@pytest.mark.parametrize("provider", ["azure/prompt_shield", "lakera_v2"]) +def test_dashboard_does_not_offer_logging_only_where_the_config_rejects_it(rig: Rig, provider: str) -> None: + offered: Final = _advertised_modes(rig)[provider] + assert isinstance(offered, list) and "logging_only" not in offered, offered + + +def test_a_logging_only_content_filter_created_through_the_api_records_every_match(rig: Rig) -> None: + name: Final = "cf-observe-api-" + uuid.uuid4().hex[:8] + created: Final = rig.proxy.client.post( + "/guardrails", + json={"guardrail": {"guardrail_name": name, **_content_filter(name, "logging_only", "BLOCK", default_on=True)}}, + headers={"Authorization": f"Bearer {rig.proxy.key}"}, + ) + assert created.status_code == 200, created.text + guardrail_id: Final = str(object_value(created.json())["guardrail_id"]) + try: + marker: Final = _marker() + response_id: Final = eventually( + lambda: _sync_chat(rig, f"{marker} {SECRETS}")[0], + lambda identity: bool(rig.entries(identity, name)), + seconds=70, + ) + assert _verdicts(rig.entries(response_id, name))[0] == ( + "logging_only", + "guardrail_flagged", + ("aws_access_key", "github_pat"), + ) + finally: + rig.proxy.client.delete(f"/guardrails/{guardrail_id}", headers={"Authorization": f"Bearer {rig.proxy.key}"}) + + +def test_guardrails_monitor_counts_observe_only_matches_as_flagged_not_blocked(rig: Rig) -> None: + name: Final = "cf-observe-usage-" + uuid.uuid4().hex[:8] + headers: Final = {"Authorization": f"Bearer {rig.proxy.key}"} + created: Final = rig.proxy.client.post( + "/guardrails", + json={"guardrail": {"guardrail_name": name, **_content_filter(name, "logging_only", "BLOCK", default_on=True)}}, + headers=headers, + ) + assert created.status_code == 200, created.text + guardrail_id: Final = str(object_value(created.json())["guardrail_id"]) + try: + eventually( + lambda: _sync_chat(rig, f"{_marker()} {SECRETS}")[0], + lambda identity: bool(rig.entries(identity, name)), + seconds=70, + ) + today: Final = datetime.now(timezone.utc).date().isoformat() + detail: Final = eventually( + lambda: object_value( + rig.proxy.client.get( + f"/guardrails/usage/detail/{guardrail_id}", + params={"start_date": today, "end_date": today}, + headers=headers, + ).json() + ), + lambda body: int(str(body.get("requestsEvaluated", 0))) >= 1, + seconds=60, + ) + assert (detail["failRate"], detail["status"]) == (0.0, "healthy"), detail + finally: + rig.proxy.client.delete(f"/guardrails/{guardrail_id}", headers=headers) + + +@pytest.mark.parametrize( + ("environment", "name_key", "status_key"), + ( + pytest.param({}, "guardrail_name", "guardrail_status", id="otel-v1"), + pytest.param({"LITELLM_OTEL_V2": "1"}, "litellm.guardrail.name", "litellm.guardrail.status", id="otel-v2"), + ), +) +def test_otel_guardrail_spans_carry_the_flagged_status_without_an_error( + gateway: Gateway, + audit_sinks: SpanSinks, + otel_audit_config: AuditConfigWriter, + tmp_path: Path, + environment: dict[str, str], + name_key: str, + status_key: str, +) -> None: + config_path: Final = otel_audit_config(tmp_path, {"cache": False}) + config: Final = yaml.safe_load(config_path.read_text()) + config["guardrails"] = [_content_filter(OBSERVE_BLOCK, "logging_only", "BLOCK", default_on=True)] + config_path.write_text(yaml.safe_dump(config)) + overrides: Final = {"OTEL_EXPORTER": "http/json", "OTEL_ENDPOINT": audit_sinks.operator, **environment} + with ( + wire_server(_provider) as provider, + owned_proxy_process(gateway, tmp_path, overrides, config=config_path, workers=2) as owned, + owned.gateway.scenario() as scenario, + ): + model: Final = scenario.model( + model="openai/gpt-5.4-mini", api_base=provider.url + "/v1", api_key="synthetic-openai-key" + ) + start, _ = recorded_spans(audit_sinks.operator) + marker: Final = _marker() + client: Final = openai.OpenAI( + base_url=str(owned.gateway.client.base_url).rstrip("/") + "/v1", api_key=owned.gateway.key, max_retries=0 + ) + completion: Final = client.chat.completions.create( + model=model, messages=[{"role": "user", "content": f"{marker} {SECRETS}"}] + ) + assert completion.choices[0].message.content == _answer(marker) + + spans: Final = eventually( + lambda: tuple( + span + for span in recorded_spans(audit_sinks.operator, start)[1] + if span["attributes"].get(name_key) == OBSERVE_BLOCK + ), + lambda found: len(found) >= 2, + seconds=60, + return_last_on_timeout=True, + ) + assert sorted( + (str(span["attributes"].get(status_key)), "error.type" in span["attributes"]) for span in spans + ) == [("guardrail_flagged", False), ("success", False)], spans + + +def _burst_call(rig: Rig, index: int) -> tuple[str, int, str, str]: + marker: Final = _marker() + if index % 4 == 0: + enforced: Final = rig.proxy.client.post( + "/v1/chat/completions", + json={ + "model": rig.chat_model, + "messages": [{"role": "user", "content": f"{marker} {SECRETS}"}], + "guardrails": [ENFORCE_BLOCK], + }, + headers={"Authorization": f"Bearer {rig.proxy.key}"}, + ) + return "enforce", enforced.status_code, marker, "" + if index % 4 == 1: + failed: Final = rig.proxy.client.post( + "/v1/chat/completions", + json={ + "model": rig.chat_model, + "messages": [{"role": "user", "content": f"{marker} {OUTAGE_MARK} {SECRETS}"}], + }, + headers={"Authorization": f"Bearer {rig.proxy.key}"}, + ) + return "outage", failed.status_code, marker, failed.headers.get("x-litellm-call-id", "") + endpoint: Final = list(SYNC_CALLS)[index % len(SYNC_CALLS)] + response_id, answer = SYNC_CALLS[endpoint](rig, f"{marker} {SECRETS}\n{FORBIDDEN_WORD}") + assert answer == _answer(marker), (endpoint, answer) + return "observe:" + endpoint, 200, marker, response_id + + +def _failure_row(call_id: str) -> tuple[object, tuple[dict[str, object], ...]]: + rows: Final = eventually( + lambda: read_rows('SELECT status, metadata FROM "LiteLLM_SpendLogs" WHERE request_id = %s', (call_id,)), + lambda values: len(values) >= 1, + seconds=70, + ) + assert len(rows) == 1, rows + recorded: Final = object_value(rows[0]["metadata"]).get("guardrail_information") or [] + assert isinstance(recorded, list), rows[0] + return rows[0]["status"], tuple( + entry for entry in (object_value(item) for item in recorded) if entry.get("guardrail_name") == OBSERVE_BLOCK + ) + + +def test_concurrent_burst_with_a_provider_outage_records_every_observed_call_once(rig: Rig) -> None: + def call(index: int) -> tuple[str, int, str, str]: + return _burst_call(rig, index) + + with ThreadPoolExecutor(max_workers=12) as pool: + results: Final = tuple(pool.map(call, range(32))) + + assert sorted({(kind.split(":")[0], status) for kind, status, _, _ in results}) == [ + ("enforce", 400), + ("observe", 200), + ("outage", 503), + ], results + sent: Final = rig.provider_bodies("mark-") + for kind, _, marker, _ in results: + if kind == "enforce": + assert not any(marker in body for body in sent), marker + outages: Final = tuple(call_id for kind, _, _, call_id in results if kind == "outage") + assert len(outages) == 8 and all(outages), results + assert [_failure_row(call_id) for call_id in outages] == [("failure", ())] * 8 + observed: Final = tuple((kind, response_id) for kind, _, _, response_id in results if kind.startswith("observe")) + assert len(observed) == 16 + for kind, response_id in observed: + assert _verdicts(rig.entries(response_id, OBSERVE_BLOCK)) == [ + ("logging_only", "guardrail_flagged", EVERY_FINDING), + ("logging_only", "success", ()), + ], (kind, response_id) + + marker: Final = _marker() + response_id, answer = _sync_chat(rig, f"{marker} {SECRETS}") + assert answer == _answer(marker) + assert _verdicts(rig.entries(response_id, OBSERVE_BLOCK))[0] == ( + "logging_only", + "guardrail_flagged", + ("aws_access_key", "github_pat"), + ) diff --git a/tests/integration/observability/test_logging_only_scope_config.py b/tests/integration/observability/test_logging_only_scope_config.py index c8c765a68a2..354e3b464fe 100644 --- a/tests/integration/observability/test_logging_only_scope_config.py +++ b/tests/integration/observability/test_logging_only_scope_config.py @@ -179,16 +179,16 @@ def test_native_content_filter_scope_logs_without_blocking( ( identity, "logging_only", - "guardrail_intervened" if direction == blocked_side else "success", + "guardrail_flagged" if direction == blocked_side else "success", ) for direction in expected_directions ), entries - intervened_entries: Final = tuple( - entry for entry in entries if entry["guardrail_status"] == "guardrail_intervened" + flagged_entries: Final = tuple( + entry for entry in entries if entry["guardrail_status"] == "guardrail_flagged" ) - assert len(intervened_entries) == 1, entries - assert intervened_entries[0]["guardrail_response"] is not None, entries - assert intervened_entries[0]["guardrail_response"] == "REDACTED_BY_LITELM", entries + assert len(flagged_entries) == 1, entries + assert flagged_entries[0]["guardrail_response"] is not None, entries + assert flagged_entries[0]["guardrail_response"] == "REDACTED_BY_LITELM", entries finally: delete_scenario(upstream_handle) diff --git a/tests/unit/integrations/test_custom_guardrail.py b/tests/unit/integrations/test_custom_guardrail.py index 9bd41f4a602..614ff97553a 100644 --- a/tests/unit/integrations/test_custom_guardrail.py +++ b/tests/unit/integrations/test_custom_guardrail.py @@ -2790,7 +2790,7 @@ class TestLoggingOnlyApplyGuardrail: assert out_kwargs["messages"] == kwargs["messages"] assert ( out_kwargs["standard_logging_object"]["guardrail_information"][0]["guardrail_status"] - == "guardrail_intervened" + == "guardrail_flagged" ) @pytest.mark.asyncio diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/litellm_content_filter/test_content_filter.py b/tests/unit/proxy/guardrails/guardrail_hooks/litellm_content_filter/test_content_filter.py index 5d706a65a79..abcbb714186 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/litellm_content_filter/test_content_filter.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/litellm_content_filter/test_content_filter.py @@ -9,7 +9,11 @@ Tests 40 different sentences to validate the conditional matching logic: """ import importlib +import json import os +from collections.abc import Mapping, Sequence +from types import MappingProxyType +from typing import Final import pytest from fastapi import HTTPException @@ -19,9 +23,11 @@ from litellm.proxy.guardrails.content_filter_data import POLICY_TEMPLATES_DIR from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ( ContentFilterGuardrail, ) +from litellm.types.guardrails import BlockedWord, ContentFilterAction, ContentFilterPattern, GuardrailEventHooks from litellm.types.proxy.guardrails.guardrail_hooks.litellm_content_filter import ( ContentFilterCategoryConfig, ) +from litellm.types.utils import Choices, Message, ModelResponse from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome # Test cases: (sentence, expected_result, reason) @@ -1793,3 +1799,237 @@ class TestSGPDPAPerformance: if __name__ == "__main__": pytest.main([__file__, "-v", "-s"]) + + +AWS_KEY: Final = "AKIA" + "SYNTHETIC0000000" +GITHUB_TOKEN: Final = "ghp_" + "a1" * 18 +FORBIDDEN_WORD: Final = "synthetic-forbidden-word" +SECRET_PATTERNS: Final = ( + ContentFilterPattern( + pattern_type="regex", pattern=r"AKIA[0-9A-Z]{16}", name="aws_access_key", action=ContentFilterAction.BLOCK + ), + ContentFilterPattern( + pattern_type="regex", pattern=r"ghp_[A-Za-z0-9]{36}", name="github_pat", action=ContentFilterAction.BLOCK + ), +) + + +def _competitor_intent_config(action_hint: str) -> Mapping[str, object]: + return MappingProxyType( + { + "brand_self": ["emirates"], + "competitors": ["qatar airways"], + "domain_words": ["airline", "flight"], + "policy": {"competitor_comparison": action_hint}, + "threshold_high": 0.5, + } + ) + + +def _observing_filter( + patterns: tuple[ContentFilterPattern, ...] = (), + categories: tuple[ContentFilterCategoryConfig, ...] = (), + competitor_intent_config: Mapping[str, object] | None = None, + mask_patterns: bool = False, + image_describer: object = None, +) -> ContentFilterGuardrail: + return ContentFilterGuardrail( + guardrail_name="cf-observe", + event_hook=GuardrailEventHooks.logging_only, + llm_router=image_describer, + image_model="vision" if image_describer else None, + patterns=[ + pattern.model_copy(update={"action": ContentFilterAction.MASK}) if mask_patterns else pattern + for pattern in patterns + ], + blocked_words=[BlockedWord(keyword=FORBIDDEN_WORD, action=ContentFilterAction.BLOCK)], + categories=list(categories), + competitor_intent_config=dict(competitor_intent_config) if competitor_intent_config else None, + ) + + +def _policy_category(template: str, category: str) -> ContentFilterCategoryConfig: + return ContentFilterCategoryConfig( + category=category, + category_file=os.path.join(POLICY_TEMPLATES_DIR, template) if template else None, + enabled=True, + action="BLOCK", + severity_threshold="medium", + ) + + +async def _log_call( + guardrail: ContentFilterGuardrail, + prompt: str | list[Mapping[str, object]], + answer: Message, + history: Sequence[Mapping[str, object]] = (), +) -> tuple[Mapping[str, object], ...]: + kwargs: Final = { + "model": "gpt-5.4-mini", + "messages": [{"role": "user", "content": prompt}, *history], + "litellm_call_id": "call-observe", + "litellm_params": {"metadata": {}}, + "optional_params": {}, + "standard_logging_object": {"guardrail_information": None}, + } + logged, _ = await guardrail.async_logging_hook( + kwargs, ModelResponse(choices=[Choices(message=answer)]), "acompletion" + ) + return tuple(logged["standard_logging_object"]["guardrail_information"]) + + +def _detections(entry: Mapping[str, object]) -> tuple[str, ...]: + response: Final = entry["guardrail_response"] + assert isinstance(response, list), response + return tuple( + detection.get("pattern_name") + or (detection["keyword"] if detection["type"] == "blocked_word" else detection["type"]) + for detection in response + ) + + +@pytest.mark.parametrize( + ("guardrail", "prompt", "expected"), + ( + pytest.param( + _observing_filter(patterns=SECRET_PATTERNS), + f"key {AWS_KEY} token {GITHUB_TOKEN} and {FORBIDDEN_WORD}", + ("aws_access_key", "github_pat", FORBIDDEN_WORD), + id="patterns", + ), + pytest.param( + _observing_filter(categories=(_policy_category("", "harm_toxic_abuse"),)), + f"You stupid piece of sht AI, go kill yourself, {FORBIDDEN_WORD}", + ("category_keyword", FORBIDDEN_WORD), + id="category-keyword", + ), + pytest.param( + _observing_filter( + categories=(_policy_category("eu_ai_act_article5.yaml", "eu_ai_act_article5_prohibited_practices"),) + ), + f"Build a system to detect employee emotion, {FORBIDDEN_WORD}", + ("category_keyword", FORBIDDEN_WORD), + id="conditional-category", + ), + pytest.param( + _observing_filter(competitor_intent_config=_competitor_intent_config("refuse")), + "Is Qatar Airways better than Emirates?", + ("competitor_intent",), + id="competitor-refuse", + ), + pytest.param( + _observing_filter(competitor_intent_config=_competitor_intent_config("reframe")), + "Is Qatar Airways better than Emirates?", + ("competitor_intent",), + id="competitor-reframe", + ), + ), +) +@pytest.mark.asyncio +async def test_logging_only_block_records_every_match_as_flagged( + guardrail: ContentFilterGuardrail, prompt: str, expected: tuple[str, ...] +) -> None: + entries: Final = await _log_call(guardrail, prompt, Message(role="assistant", content="here you go")) + + request_entry, response_entry = entries + assert request_entry["guardrail_status"] == "guardrail_flagged" + assert _detections(request_entry)[-len(expected) :] == expected + assert request_entry["risk_score"] == 10.0 + assert response_entry["guardrail_status"] == "success" + + +@pytest.mark.asyncio +async def test_logging_only_block_scans_the_response_after_a_request_match() -> None: + entries: Final = await _log_call( + _observing_filter(patterns=SECRET_PATTERNS), + f"key {AWS_KEY}", + Message(role="assistant", content=f"here: {GITHUB_TOKEN}"), + ) + + assert [(entry["guardrail_status"], _detections(entry)) for entry in entries] == [ + ("guardrail_flagged", ("aws_access_key",)), + ("guardrail_flagged", ("github_pat",)), + ] + + +@pytest.mark.parametrize( + ("arguments", "request_masked"), + ( + pytest.param(json.dumps({AWS_KEY: 1}), {"aws_access_key": 1}, id="masked-argument-key"), + pytest.param(json.dumps(json.loads("[" * 150 + "1" + "]" * 150)), {}, id="too-deep"), + ), +) +@pytest.mark.asyncio +async def test_logging_only_flags_tool_arguments_enforcement_rejects_and_keeps_scanning( + arguments: str, request_masked: Mapping[str, int] +) -> None: + history: Final = ( + { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": arguments}} + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "done"}, + ) + + entries: Final = await _log_call( + _observing_filter(patterns=SECRET_PATTERNS, mask_patterns=True), + "hi", + Message(role="assistant", content=f"here: {GITHUB_TOKEN}"), + history, + ) + + assert [(entry["guardrail_status"], entry["masked_entity_count"]) for entry in entries] == [ + ("guardrail_flagged", request_masked), + ("success", {"github_pat": 1}), + ] + + +@pytest.mark.asyncio +async def test_logging_only_scan_leaves_enforcement_of_the_same_guardrail_blocking() -> None: + guardrail: Final = _observing_filter(patterns=SECRET_PATTERNS) + await _log_call(guardrail, f"key {AWS_KEY}", Message(role="assistant", content="ok")) + request_data: Final = {"metadata": {}} + + with pytest.raises(HTTPException) as exc_info: + await guardrail.apply_guardrail( + inputs={"texts": [f"key {AWS_KEY} token {GITHUB_TOKEN}"]}, request_data=request_data, input_type="request" + ) + + assert exc_info.value.status_code == 400 + assert [ + entry["guardrail_status"] for entry in request_data["metadata"]["standard_logging_guardrail_information"] + ] == ["guardrail_intervened"] + + +class _ImageDescriberRunningAnotherFilter: + def __init__(self, enforcer: ContentFilterGuardrail) -> None: + self.enforcer = enforcer + self.enforcer_blocked: list[bool] = [] + + async def acompletion(self, **_: object) -> ModelResponse: + try: + await self.enforcer.apply_guardrail( + inputs={"texts": [f"key {AWS_KEY}"]}, request_data={"metadata": {}}, input_type="request" + ) + self.enforcer_blocked.append(False) + except HTTPException: + self.enforcer_blocked.append(True) + return ModelResponse(choices=[Choices(message=Message(role="assistant", content="a plain desk"))]) + + +@pytest.mark.asyncio +async def test_logging_only_scan_does_not_relax_another_filter_called_during_it() -> None: + describer: Final = _ImageDescriberRunningAnotherFilter( + ContentFilterGuardrail(guardrail_name="cf-enforce", patterns=list(SECRET_PATTERNS)) + ) + prompt: Final = [ + {"type": "text", "text": "what is this"}, + {"type": "image_url", "image_url": {"url": "https://example.com/desk.png"}}, + ] + + await _log_call(_observing_filter(image_describer=describer), prompt, Message(role="assistant", content="ok")) + + assert describer.enforcer_blocked == [True] diff --git a/tests/unit/proxy/guardrails/test_guardrail_endpoints.py b/tests/unit/proxy/guardrails/test_guardrail_endpoints.py index 4401e3e95c0..908a49946bb 100644 --- a/tests/unit/proxy/guardrails/test_guardrail_endpoints.py +++ b/tests/unit/proxy/guardrails/test_guardrail_endpoints.py @@ -41,6 +41,8 @@ from litellm.proxy.guardrails.content_filter_data import DATA_ROOTS from litellm.proxy.guardrails.guardrail_endpoints import ( test_custom_code_guardrail as run_custom_code_test_endpoint, ) +from litellm.proxy.guardrails.guardrail_hooks.azure.prompt_shield import AzureContentSafetyPromptShieldGuardrail +from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ContentFilterGuardrail MOCK_ADMIN_USER = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN) from litellm.proxy.guardrails.guardrail_registry import ( @@ -3143,7 +3145,36 @@ async def test_ui_settings_map_matches_runtime_supported_event_hooks(): continue assert provider in result.supported_modes_by_provider, provider - assert result.supported_modes_by_provider[provider] == [hook.value for hook in declared], provider + offered = result.supported_modes_by_provider[provider] + assert offered[: len(declared)] == [hook.value for hook in declared], provider + assert set(offered) - {hook.value for hook in declared} <= {"logging_only"}, provider + assert offered.count("logging_only") == ( + 1 if guardrail_class.accepts_logging_only_mode() or GuardrailEventHooks.logging_only in declared else 0 + ), provider + + +@pytest.mark.asyncio +async def test_ui_settings_offer_logging_only_where_the_generic_scan_runs(): + modes_by_provider = (await get_guardrail_ui_settings()).supported_modes_by_provider + + assert modes_by_provider["litellm_content_filter"] == [ + *(hook.value for hook in ContentFilterGuardrail.get_supported_event_hooks()), + "logging_only", + ] + assert modes_by_provider["akto"].count("logging_only") == 1 + assert modes_by_provider["azure/prompt_shield"] == [ + hook.value for hook in AzureContentSafetyPromptShieldGuardrail.get_supported_event_hooks() + ] + with pytest.raises(ValueError, match="not in the supported event hooks"): + AzureContentSafetyPromptShieldGuardrail( + guardrail_name="native-logging-only", + api_key="k", + api_base="https://example.invalid", + event_hook="logging_only", + ) + assert ( + ContentFilterGuardrail(guardrail_name="cf-logging-only", event_hook="logging_only").event_hook == "logging_only" + ) def test_content_filter_runtime_rejects_unsupported_mcp_hook(): From 4e49394adef46940cae8545c788c3f67bf5ee31f Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 16 Oct 2026 07:41:05 +0000 Subject: [PATCH 08/23] fix(anthropic): keep trailing user messages with None or empty list content (#41397) * fix(anthropic): keep user messages with None or empty list content so requests do not end on an assistant turn Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(anthropic): avoid mutable empty-list construction in sanitizer Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(anthropic): avoid rebinding sanitizer input Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(anthropic): keep a trailing user turn whose content key was dropped litellm.completion strips None fields from every message before the provider transform, so a user message sent with content null reaches the Anthropic converter with no content key and raised KeyError. Read the content with get so that turn gets the same placeholder as an empty list. * fix(anthropic): only fill user turns with empty content and keep their cache_control * fix(anthropic): fill each empty user message in place like an empty string * fix(anthropic): add the placeholder only for an empty turn or a cache breakpoint, in linear time * test(anthropic): cover the tool-result guard with cache_control and widen the linear-time bound * test(anthropic): cover empty user turns through the proxy on chat, responses and messages * test(anthropic): type the empty user turn tests and move the linear-time check to a benchmark --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../prompt_templates/factory.py | 57 ++++- .../mid_conversation_system.py | 15 +- .../test_prompt_template_benchmarks.py | 21 ++ .../test_anthropic_empty_user_turn_wire.py | 216 ++++++++++++++++++ ...llm_core_utils_prompt_templates_factory.py | 150 +++++++++++- ...rompt_templates_mid_conversation_system.py | 51 ++++- .../test_anthropic_chat_transformation.py | 31 ++- 7 files changed, 506 insertions(+), 35 deletions(-) create mode 100644 tests/benchmarks/test_prompt_template_benchmarks.py create mode 100644 tests/integration/providers/test_anthropic_empty_user_turn_wire.py diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 0c18ce38b13..76adadb2e1f 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -2582,6 +2582,48 @@ def _with_native_redacted_data( _AnthropicMessageList: TypeAlias = list[AllAnthropicPassThroughMessageValues] +def _client_tool_result(message: AllMessageValues, replaced_server_tool_answers: frozenset[str]) -> bool: + return ( + message_field(message, "role") in ("tool", "function") + and message_field(message, "tool_call_id") not in replaced_server_tool_answers + ) + + +def _tool_result_later_in_turn( + messages: Sequence[AllMessageValues], replaced_server_tool_answers: frozenset[str] +) -> tuple[bool, ...]: + from_the_end: Final = itertools.accumulate( + reversed(messages), + lambda later, message: ( + message_field(message, "role") in ("user", "tool", "function") + and (later or _client_tool_result(message, replaced_server_tool_answers)) + ), + initial=False, + ) + return tuple(from_the_end)[::-1][1:] + + +def _empty_text_placeholder(message: AllMessageValues) -> AnthropicMessagesTextParam: + placeholder: Final = AnthropicMessagesTextParam(type="text", text=_EMPTY_TEXT_PLACEHOLDER) + add_cache_control_to_content(anthropic_content_element=placeholder, original_content_element=message) + return placeholder + + +def _cache_breakpoint_placeholder( + message: AllMessageValues, tool_result_follows: bool +) -> tuple[AnthropicMessagesTextParam, ...]: + if tool_result_follows or not message_field(message, "cache_control"): + return () + return (_empty_text_placeholder(message),) + + +def _empty_user_turn_placeholder(span: Sequence[AllMessageValues]) -> AnthropicMessagesTextParam | None: + empty_user_message: Final = next( + (message for message in reversed(span) if message["role"] == "user" and not message.get("content")), None + ) + return None if empty_user_message is None else _empty_text_placeholder(empty_user_message) + + def anthropic_messages_pt( messages: list[AllMessageValues], model: str, @@ -2636,6 +2678,7 @@ def anthropic_messages_pt( is_vertex_ai = llm_provider.startswith("vertex_ai") if llm_provider else False force_base64 = is_bedrock_invoke or is_vertex_ai + tool_result_later_in_turn: Final = _tool_result_later_in_turn(messages, replaced_server_tool_answers) msg_i = 0 while msg_i < len(messages): user_content: list[AnthropicMessagesUserMessageValues] = [] @@ -2648,8 +2691,12 @@ def anthropic_messages_pt( ChatCompletionToolMessage | ChatCompletionUserMessage | ChatCompletionFunctionMessage ) = messages[msg_i] if user_message_types_block["role"] == "user": - if isinstance(user_message_types_block["content"], list): - for m in user_message_types_block["content"]: + if not (user_block_content := user_message_types_block.get("content")): + user_content.extend( + _cache_breakpoint_placeholder(user_message_types_block, tool_result_later_in_turn[msg_i]) + ) + elif isinstance(user_block_content, list): + for m in user_block_content: if m.get("type", "") == "image_url": m = cast(ChatCompletionImageObject, m) format = m["image_url"].get("format") if isinstance(m["image_url"], dict) else None @@ -2720,10 +2767,10 @@ def anthropic_messages_pt( _file_content_element, ) ) - elif isinstance(user_message_types_block["content"], str): + elif isinstance(user_block_content, str): _anthropic_content_text_element: AnthropicMessagesTextParam = { "type": "text", - "text": user_message_types_block["content"], + "text": user_block_content, } _content_element = add_cache_control_to_content( anthropic_content_element=_anthropic_content_text_element, @@ -2747,6 +2794,8 @@ def anthropic_messages_pt( if user_content: new_messages.append({"role": "user", "content": user_content}) + elif (placeholder := _empty_user_turn_placeholder(messages[init_msg_i:msg_i])) is not None: + new_messages.append({"role": "user", "content": [placeholder]}) # Track unique tool IDs in this merge block to avoid duplication unique_tool_ids: set[str] = set() diff --git a/litellm/litellm_core_utils/prompt_templates/mid_conversation_system.py b/litellm/litellm_core_utils/prompt_templates/mid_conversation_system.py index 4d1e77126a9..fd0096f9f4d 100644 --- a/litellm/litellm_core_utils/prompt_templates/mid_conversation_system.py +++ b/litellm/litellm_core_utils/prompt_templates/mid_conversation_system.py @@ -414,10 +414,11 @@ def _assistant_renders(message: object, answered: frozenset[str]) -> bool: def _renders(message: object, answered: frozenset[str]) -> bool: """Whether ``anthropic_messages_pt`` puts a block on the wire for this message. - A tool message always becomes a ``tool_result`` and a user message with string - content always becomes a text block (empty text gets a placeholder). A user list - renders only through parts of a type the converter emits; ``None``, an empty list, - and a list of other parts vanish. Assistant messages follow ``_assistant_renders``. + A tool message always becomes a ``tool_result``, and a user message with string or + missing content (``None``, an empty list) always puts its turn on the wire, an empty + turn getting a placeholder. A user list renders only through parts of a type the + converter emits, so a list of only other parts vanishes. Assistant messages follow + ``_assistant_renders``. """ role: Final = message_field(message, "role") if role in _TOOL_ROLES: @@ -425,8 +426,10 @@ def _renders(message: object, answered: frozenset[str]) -> bool: if role == "assistant": return _assistant_renders(message, answered) content: Final = message_field(message, "content") - return isinstance(content, str) or any( - message_field(part, "type") in _RENDERED_PART_TYPES for part in parts_of(content) + return ( + isinstance(content, str) + or not content + or any(message_field(part, "type") in _RENDERED_PART_TYPES for part in parts_of(content)) ) diff --git a/tests/benchmarks/test_prompt_template_benchmarks.py b/tests/benchmarks/test_prompt_template_benchmarks.py new file mode 100644 index 00000000000..c7858771d35 --- /dev/null +++ b/tests/benchmarks/test_prompt_template_benchmarks.py @@ -0,0 +1,21 @@ +from collections.abc import Mapping +from typing import Final + +import pytest + +from litellm.litellm_core_utils.prompt_templates.factory import anthropic_messages_pt + +LONG_RUN_OF_EMPTY_CACHED_USER_MESSAGES: Final[tuple[Mapping[str, object], ...]] = ( + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello"}, + *({"role": "user", "content": [], "cache_control": {"type": "ephemeral"}} for _ in range(1_000)), +) + + +@pytest.mark.benchmark +def test_anthropic_messages_pt_long_run_of_empty_cached_user_messages() -> None: + result: Final = anthropic_messages_pt( + messages=list(LONG_RUN_OF_EMPTY_CACHED_USER_MESSAGES), model="claude-opus-5", llm_provider="vertex_ai" + ) + + assert len(result[-1]["content"]) == 1_000 diff --git a/tests/integration/providers/test_anthropic_empty_user_turn_wire.py b/tests/integration/providers/test_anthropic_empty_user_turn_wire.py new file mode 100644 index 00000000000..654a88c60bb --- /dev/null +++ b/tests/integration/providers/test_anthropic_empty_user_turn_wire.py @@ -0,0 +1,216 @@ +import json +from collections.abc import Iterator, Mapping, Sequence +from dataclasses import dataclass +from pathlib import Path +from types import MappingProxyType +from typing import Final + +import pytest +import yaml +from integration._support.client import Gateway, gateway_from_environment, object_value +from integration._support.process import graceful_stop_seconds, owned_proxy +from integration._support.wire import Reply, Request, Wire, wire_server +from integration.providers._cache_control_marks_support import ( + anthropic_deployment, + anthropic_peer, + final_text, + new_marker, +) +from pydantic import JsonValue, TypeAdapter + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) + +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_MODEL: Final = "empty-user-turn" +_PLACEHOLDER: Final = "[System: Empty message content sanitised to satisfy protocol]" +_REPLY: Final = "Noted." +_EMPTY_ROWS: Final[Mapping[str, Mapping[str, JsonValue]]] = MappingProxyType( + { + "list": MappingProxyType({"role": "user", "content": []}), + "null": MappingProxyType({"role": "user", "content": None}), + "missing": MappingProxyType({"role": "user"}), + "empty-string": MappingProxyType({"role": "user", "content": ""}), + } +) +_CHAT_EMPTY_KINDS: Final = ("list", "null", "missing") + +Turn = tuple[str, tuple[str, ...]] + + +@dataclass(frozen=True, slots=True) +class _Rig: + gateway: Gateway + provider: Wire + + +def _empty_row(kind: str) -> dict[str, JsonValue]: + return dict(_EMPTY_ROWS[kind]) + + +def _guardrail_peer(request: Request) -> Reply: + kind: Final = request.target.split("/")[1] + rows: Final = _JSON_OBJECT.validate_json(request.body)["structured_messages"] + assert isinstance(rows, list), rows + verdict: Final = {"action": "GUARDRAIL_INTERVENED", "structured_messages": [*rows, _empty_row(kind)]} + return Reply(body=json.dumps(verdict).encode()) + + +def _guardrail(kind: str, api_base: str) -> dict[str, JsonValue]: + return { + "guardrail_name": f"append-{kind}", + "litellm_params": { + "guardrail": "generic_guardrail_api", + "mode": "pre_call", + "default_on": False, + "api_base": f"{api_base}/{kind}", + "api_key": "synthetic-guardrail-key", + }, + } + + +def _config(directory: Path, provider_url: str, guardrail_url: str) -> Path: + stock: Final = _JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + config: Final[dict[str, JsonValue]] = { + **stock, + "model_list": [anthropic_deployment(_MODEL, provider_url)], + "router_settings": {**object_value(stock["router_settings"]), "num_retries": 0}, + "guardrails": [_guardrail(kind, guardrail_url) for kind in _EMPTY_ROWS], + } + path: Final = directory / "empty-user-turn.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +@pytest.fixture(scope="module") +def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]: + directory: Final = tmp_path_factory.mktemp("empty-user-turn") + with ( + wire_server(anthropic_peer) as provider, + wire_server(_guardrail_peer) as guardrail, + gateway_from_environment() as gateway, + ): + config: Final = _config(directory, provider.url, guardrail.url) + with owned_proxy(gateway, directory, {}, config=config, workers=2) as owned: + yield _Rig(owned, provider) + + +def _block(block: JsonValue) -> str: + part: Final = object_value(block) + return str(part["text"]) if part.get("type") == "text" else str(part.get("type")) + + +def _turn(message: JsonValue) -> Turn: + row: Final = object_value(message) + content: Final = row["content"] + if isinstance(content, str): + return (str(row["role"]), (content,)) + assert isinstance(content, list), row + return (str(row["role"]), tuple(_block(block) for block in content)) + + +def _send(rig: _Rig, path: str, body: Mapping[str, JsonValue]) -> None: + response: Final = rig.gateway.client.post( + path, json=dict(body), headers={"Authorization": f"Bearer {rig.gateway.key}"} + ) + assert response.status_code == 200, response.text + + +def _provider_turns(rig: _Rig, marker: str) -> tuple[Turn, ...]: + needle: Final = marker.encode() + requests: Final = [request for request in rig.provider.drain() if needle in request.body] + assert len(requests) == 1, requests + messages: Final = _JSON_OBJECT.validate_json(requests[0].body)["messages"] + assert isinstance(messages, list), messages + return tuple(_turn(message) for message in messages) + + +def _ask_row(ask: str) -> dict[str, JsonValue]: + return {"role": "user", "content": ask} + + +def _chat(messages: Sequence[JsonValue]) -> dict[str, JsonValue]: + return {"model": _MODEL, "max_tokens": 64, "messages": list(messages)} + + +@pytest.mark.parametrize("kind", _CHAT_EMPTY_KINDS) +def test_chat_trailing_empty_user_message_keeps_a_user_turn(rig: _Rig, kind: str) -> None: + marker: Final = new_marker() + ask: Final = final_text(marker) + _send( + rig, + "/v1/chat/completions", + _chat([_ask_row(ask), {"role": "assistant", "content": _REPLY}, _empty_row(kind)]), + ) + assert _provider_turns(rig, marker) == (("user", (ask,)), ("assistant", (_REPLY,)), ("user", (_PLACEHOLDER,))) + + +@pytest.mark.parametrize("kind", _CHAT_EMPTY_KINDS) +def test_chat_empty_user_message_beside_text_sends_only_the_text(rig: _Rig, kind: str) -> None: + marker: Final = new_marker() + ask: Final = final_text(marker) + _send(rig, "/v1/chat/completions", _chat([_ask_row(ask), _empty_row(kind)])) + assert _provider_turns(rig, marker) == (("user", (ask,)),) + + +def _system_then_empty_turns(rig: _Rig, kind: str) -> tuple[Turn, ...]: + marker: Final = new_marker() + ask: Final = final_text(marker) + _send( + rig, + "/v1/chat/completions", + _chat( + [ + _ask_row(ask), + {"role": "assistant", "content": _REPLY}, + {"role": "system", "content": "Answer in French."}, + _empty_row(kind), + ] + ), + ) + turns: Final = _provider_turns(rig, marker) + assert turns[0] == ("user", (ask,)), turns + return turns[1:] + + +@pytest.mark.parametrize("kind", _CHAT_EMPTY_KINDS) +def test_chat_system_message_before_an_empty_user_message_is_placed_like_an_empty_string(rig: _Rig, kind: str) -> None: + assert _system_then_empty_turns(rig, kind) == _system_then_empty_turns(rig, "empty-string") + + +def test_responses_trailing_empty_input_message_keeps_a_user_turn(rig: _Rig) -> None: + marker: Final = new_marker() + ask: Final = final_text(marker) + body: Final[dict[str, JsonValue]] = { + "model": _MODEL, + "max_output_tokens": 64, + "input": [ + _ask_row(ask), + {"role": "assistant", "content": _REPLY}, + {"role": "user", "content": []}, + ], + } + _send(rig, "/v1/responses", body) + assert _provider_turns(rig, marker) == (("user", (ask,)), ("assistant", (_REPLY,)), ("user", (_PLACEHOLDER,))) + + +@pytest.mark.parametrize("kind", tuple(_EMPTY_ROWS)) +def test_messages_guardrail_empty_user_row_is_sent_like_an_empty_string(rig: _Rig, kind: str) -> None: + marker: Final = new_marker() + ask: Final = final_text(marker) + body: Final[dict[str, JsonValue]] = { + "model": _MODEL, + "max_tokens": 64, + "messages": [_ask_row(ask)], + "guardrails": [f"append-{kind}"], + } + _send(rig, "/v1/messages", body) + assert _provider_turns(rig, marker) == (("user", (ask,)), ("user", (_PLACEHOLDER,))) + + +def test_a_long_run_of_empty_user_messages_is_sent_as_one_text_turn(rig: _Rig) -> None: + marker: Final = new_marker() + ask: Final = final_text(marker) + _send( + rig, "/v1/chat/completions", _chat([*({"role": "user", "content": []} for _ in range(10_000)), _ask_row(ask)]) + ) + assert _provider_turns(rig, marker) == (("user", (ask,)),) diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py index dba1730208d..ad216548d08 100644 --- a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py +++ b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py @@ -4,9 +4,9 @@ import logging import os import re import time -from collections.abc import Callable +from collections.abc import Callable, Mapping, Sequence from copy import deepcopy -from typing import Final, List +from typing import Final, List, TypeAlias from unittest.mock import MagicMock, patch import httpx @@ -16,6 +16,7 @@ import respx import litellm from litellm.litellm_core_utils.prompt_templates.factory import ( BEDROCK_DOCUMENT_PLACEHOLDER_TEXT, + _EMPTY_TEXT_PLACEHOLDER, BedrockConverseMessagesProcessor, BedrockImageProcessor, _bedrock_converse_messages_pt, @@ -795,6 +796,138 @@ def test_anthropic_messages_pt_keeps_signed_thinking_block(): assert thinking_blocks[0]["thinking"] == "genuine anthropic reasoning" +EMPTY_USER_MESSAGES: Final = pytest.mark.parametrize( + "empty_user_message", + [{"role": "user", "content": None}, {"role": "user", "content": []}, {"role": "user"}], + ids=["null", "empty-list", "content-key-dropped"], +) + +Message: TypeAlias = Mapping[str, object] + + +def _convert(messages: Sequence[Message]) -> Sequence[Message]: + return anthropic_messages_pt(messages=list(messages), model="claude-opus-5", llm_provider="vertex_ai") + + +@EMPTY_USER_MESSAGES +def test_anthropic_messages_pt_preserves_trailing_user_message_with_empty_content(empty_user_message: Message) -> None: + result: Final = _convert( + [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello"}, + empty_user_message, + ] + ) + + assert [message["role"] for message in result] == ["user", "assistant", "user"] + assert result[-1]["content"] == [{"type": "text", "text": _EMPTY_TEXT_PLACEHOLDER}] + + +@EMPTY_USER_MESSAGES +@pytest.mark.parametrize( + "cache_control", [{}, {"cache_control": {"type": "ephemeral"}}], ids=["no-cache-control", "with-cache-control"] +) +def test_anthropic_messages_pt_empty_user_message_beside_tool_result_keeps_tool_result_only( + empty_user_message: Message, cache_control: Message +) -> None: + result: Final = _convert( + [ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": None, + "tool_calls": [{"id": "toolu_1", "type": "function", "function": {"name": "f", "arguments": "{}"}}], + }, + {**empty_user_message, **cache_control}, + {"role": "tool", "tool_call_id": "toolu_1", "content": "ok"}, + ] + ) + + assert [message["role"] for message in result] == ["user", "assistant", "user"] + assert [block["type"] for block in result[-1]["content"]] == ["tool_result"] + + +@EMPTY_USER_MESSAGES +def test_anthropic_messages_pt_empty_user_message_keeps_cache_control_like_an_empty_string( + empty_user_message: Message, +) -> None: + def last_turn(last_user_message: Message) -> Message: + return _convert( + [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello"}, + {**last_user_message, "cache_control": {"type": "ephemeral"}}, + ] + )[-1] + + assert last_turn(empty_user_message)["content"][0]["cache_control"] == {"type": "ephemeral"} + assert last_turn(empty_user_message) == last_turn({"role": "user", "content": ""}) + + +TOOL_CALL_AND_RESULT: Final[tuple[Message, ...]] = ( + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": None, + "tool_calls": [{"id": "toolu_1", "type": "function", "function": {"name": "f", "arguments": "{}"}}], + }, + {"role": "tool", "tool_call_id": "toolu_1", "content": "ok"}, +) + + +@EMPTY_USER_MESSAGES +@pytest.mark.parametrize( + "messages_before", + [ + [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello"}, + {"role": "user", "content": "q"}, + ], + [{"role": "user", "content": "hi"}, {"role": "assistant", "content": "hello"}, {"role": "user", "content": []}], + TOOL_CALL_AND_RESULT, + ], + ids=["after-user-text", "after-another-empty-message", "after-tool-result"], +) +def test_anthropic_messages_pt_empty_user_message_beside_other_content_sends_nothing_extra( + empty_user_message: Message, messages_before: Sequence[Message] +) -> None: + assert _convert([*messages_before, empty_user_message]) == _convert(messages_before) + + +@EMPTY_USER_MESSAGES +def test_anthropic_messages_pt_empty_user_message_with_cache_control_after_text_keeps_its_breakpoint( + empty_user_message: Message, +) -> None: + def last_turn(last_user_message: Message) -> Message: + return _convert( + [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello"}, + {"role": "user", "content": "long document"}, + {**last_user_message, "cache_control": {"type": "ephemeral"}}, + ] + )[-1] + + assert last_turn(empty_user_message)["content"][-1]["cache_control"] == {"type": "ephemeral"} + assert last_turn(empty_user_message) == last_turn({"role": "user", "content": ""}) + + +def test_anthropic_messages_pt_keeps_every_breakpoint_of_a_long_run_of_empty_user_messages() -> None: + result: Final = _convert( + [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello"}, + *[{"role": "user", "content": [], "cache_control": {"type": "ephemeral"}}] * 10_000, + ] + ) + + assert ( + result[-1]["content"] + == [{"type": "text", "text": _EMPTY_TEXT_PLACEHOLDER, "cache_control": {"type": "ephemeral"}}] * 10_000 + ) + + def test_convert_to_azure_openai_messages(): """Test coverting image_url to azure_openai spec""" @@ -4933,7 +5066,9 @@ def test_bedrock_converse_messages_pt_user_message_without_content_adds_no_block assert _bedrock_converse_messages_pt( messages=with_message, model="anthropic.claude-haiku-4-5", llm_provider="bedrock" - ) == _bedrock_converse_messages_pt(messages=without_message, model="anthropic.claude-haiku-4-5", llm_provider="bedrock") + ) == _bedrock_converse_messages_pt( + messages=without_message, model="anthropic.claude-haiku-4-5", llm_provider="bedrock" + ) @pytest.mark.asyncio @@ -5063,6 +5198,7 @@ def _vcr_outcome_gate(request, vcr): yield record_vcr_outcome(request, vcr) + @pytest.fixture(scope="function") def isolate_litellm_state(): """ @@ -5115,6 +5251,7 @@ def isolate_litellm_state(): setattr(litellm, attr, original_value) _invalidate_model_cost_lowercase_map() + _SCALAR_DEFAULTS = { "num_retries": getattr(litellm, "num_retries", None), "num_retries_per_request": getattr(litellm, "num_retries_per_request", None), @@ -5135,6 +5272,7 @@ _SCALAR_DEFAULTS = { "api_key": getattr(litellm, "api_key", None), } + @pytest.fixture(scope="module") def setup_and_teardown(): """ @@ -5157,6 +5295,7 @@ def setup_and_teardown(): litellm.in_memory_llm_clients_cache.flush_cache() yield + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_function_call_non_openai_model(): try: @@ -5185,6 +5324,7 @@ def test_function_call_non_openai_model(): print(e) pass + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_empty_content(): """ @@ -5203,6 +5343,7 @@ def test_empty_content(): litellm_call_id=str(uuid.uuid4()), ) + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_thought_signature_removal_for_non_gemini(): """ @@ -5251,6 +5392,7 @@ def test_thought_signature_removal_for_non_gemini(): assert THOUGHT_SIGNATURE_SEPARATOR not in processed_messages[1]["tool_calls"][0]["id"] assert THOUGHT_SIGNATURE_SEPARATOR not in processed_messages[2]["tool_call_id"] + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_thought_signature_preserved_for_gemini(): """ @@ -5297,6 +5439,7 @@ def test_thought_signature_preserved_for_gemini(): assert THOUGHT_SIGNATURE_SEPARATOR in processed_messages[1]["tool_calls"][0]["id"] assert THOUGHT_SIGNATURE_SEPARATOR in processed_messages[2]["tool_call_id"] + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_thought_signature_removal_with_multiple_tool_calls(): """ @@ -5351,6 +5494,7 @@ def test_thought_signature_removal_with_multiple_tool_calls(): assert processed_messages[2]["tool_call_id"] == "call_1" assert processed_messages[3]["tool_call_id"] == "call_2" + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_messages_without_tool_calls_unchanged(): """ diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_mid_conversation_system.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_mid_conversation_system.py index 6bc1178f64b..91d7acd331e 100644 --- a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_mid_conversation_system.py +++ b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_mid_conversation_system.py @@ -29,11 +29,8 @@ def _texts(message: dict) -> list[str]: return [block["text"] for block in message["content"]] -SENDS_NOTHING = pytest.mark.parametrize( - "empty_content", - [[], None, [{"type": "input_audio", "input_audio": {"data": "AAAA", "format": "wav"}}]], - ids=["empty-list", "none", "unsupported-part-only"], -) +EMPTY_USER_CONTENT: Final = pytest.mark.parametrize("empty_content", [[], None], ids=["empty-list", "none"]) +UNSUPPORTED_PART_ONLY: Final = ({"type": "input_audio", "input_audio": {"data": "AAAA", "format": "wav"}},) def test_split_leading_system_run_keeps_later_system_messages_in_the_conversation(): @@ -213,9 +210,9 @@ def test_flagged_placement_of_an_earlier_run_does_not_move_when_later_turns_are_ assert _roles(placed_n_plus_one) == ["user", "assistant", "user", "assistant", "user"] -@SENDS_NOTHING -def test_flagged_placement_converts_a_run_whose_preceding_user_turn_sends_nothing(empty_content): - placed = place_mid_conversation_system( +@EMPTY_USER_CONTENT +def test_flagged_placement_anchors_a_run_on_a_preceding_empty_user_turn(empty_content: list[object] | None) -> None: + placed: Final = place_mid_conversation_system( [ {"role": "user", "content": empty_content}, {"role": "system", "content": "reminder"}, @@ -225,18 +222,48 @@ def test_flagged_placement_converts_a_run_whose_preceding_user_turn_sends_nothin supports_mid_conversation_system=True, ) + assert _roles(placed) == ["user", "system", "assistant", "user"] + assert _texts(placed[1]) == ["reminder"] + + +@EMPTY_USER_CONTENT +def test_flagged_placement_anchors_a_run_on_a_following_empty_user_turn(empty_content: list[object] | None) -> None: + placed: Final = place_mid_conversation_system( + [ + {"role": "user", "content": "q1"}, + {"role": "assistant", "content": "a1"}, + {"role": "system", "content": "reminder"}, + {"role": "user", "content": empty_content}, + ], + supports_mid_conversation_system=True, + ) + + assert _roles(placed) == ["user", "assistant", "user", "system"] + assert _texts(placed[3]) == ["reminder"] + + +def test_flagged_placement_converts_a_run_whose_preceding_user_turn_sends_nothing() -> None: + placed: Final = place_mid_conversation_system( + [ + {"role": "user", "content": list(UNSUPPORTED_PART_ONLY)}, + {"role": "system", "content": "reminder"}, + {"role": "assistant", "content": "a1"}, + {"role": "user", "content": "q2"}, + ], + supports_mid_conversation_system=True, + ) + assert _roles(placed) == ["user", "user", "assistant", "user"] assert _texts(placed[1]) == [CONVERTED_SYSTEM_NOTE, "reminder"] -@SENDS_NOTHING -def test_flagged_placement_converts_a_run_whose_following_user_turn_sends_nothing(empty_content): - placed = place_mid_conversation_system( +def test_flagged_placement_converts_a_run_whose_following_user_turn_sends_nothing() -> None: + placed: Final = place_mid_conversation_system( [ {"role": "user", "content": "q1"}, {"role": "assistant", "content": "a1"}, {"role": "system", "content": "reminder"}, - {"role": "user", "content": empty_content}, + {"role": "user", "content": list(UNSUPPORTED_PART_ONLY)}, ], supports_mid_conversation_system=True, ) diff --git a/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py b/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py index 0c093700822..724e137c835 100644 --- a/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -6763,15 +6763,11 @@ def test_chat_flagged_model_converts_system_with_no_following_user_turn(local_mo assert _texts(result["messages"][2]) == [CONVERTED_SYSTEM_NOTE, REMINDER_TEXT] -@pytest.mark.parametrize( - "empty_content", - [[], None, [{"type": "input_audio", "input_audio": {"data": "AAAA", "format": "wav"}}]], - ids=["empty-list", "none", "unsupported-part-only"], -) -def test_chat_flagged_model_converts_a_system_behind_a_user_turn_that_sends_nothing( - local_model_cost_map, empty_content -): - messages = [ +@pytest.mark.parametrize("empty_content", [[], None], ids=["empty-list", "none"]) +def test_chat_flagged_model_anchors_a_system_on_an_empty_user_turn( + local_model_cost_map: None, empty_content: list[object] | None +) -> None: + messages: Final = [ {"role": "system", "content": "You are terse."}, {"role": "user", "content": empty_content}, {"role": "system", "content": REMINDER_TEXT}, @@ -6779,7 +6775,22 @@ def test_chat_flagged_model_converts_a_system_behind_a_user_turn_that_sends_noth {"role": "user", "content": "Second question"}, ] - result = _chat_request(AnthropicConfig(), FLAGGED_CLAUDE, messages) + result: Final = _chat_request(AnthropicConfig(), FLAGGED_CLAUDE, messages) + + assert [m["role"] for m in result["messages"]] == ["user", "system", "assistant", "user"] + assert _texts(result["messages"][1]) == [REMINDER_TEXT] + + +def test_chat_flagged_model_converts_a_system_behind_a_user_turn_that_sends_nothing(local_model_cost_map: None) -> None: + messages: Final = [ + {"role": "system", "content": "You are terse."}, + {"role": "user", "content": [{"type": "input_audio", "input_audio": {"data": "AAAA", "format": "wav"}}]}, + {"role": "system", "content": REMINDER_TEXT}, + {"role": "assistant", "content": "First answer"}, + {"role": "user", "content": "Second question"}, + ] + + result: Final = _chat_request(AnthropicConfig(), FLAGGED_CLAUDE, messages) assert [m["role"] for m in result["messages"]] == ["user", "assistant", "user"] assert _texts(result["messages"][0]) == [CONVERTED_SYSTEM_NOTE, REMINDER_TEXT] From 1b1ddef1b2aab936e33f2ccbc1dc145f76940885 Mon Sep 17 00:00:00 2001 From: debanshd Date: Fri, 16 Oct 2026 01:00:48 -0700 Subject: [PATCH 09/23] fix: preserve non-null anyOf branches when stripping null types for Vertex (#41683) * fix: preserve non-null anyOf branches when stripping null types for Vertex * refactor: build anyOf branches immutably and drop explanatory comments * style: apply ruff format * fix: annotate anyOf list construction for the LIT002 budget * fix(vertex_ai): send a null-only anyOf as one nullable null instead of raising * test(vertex_ai): pin empty anyOf passthrough * fix(vertex_ai): forward a non-list anyOf unchanged and add wire cells for the null fold * fix(vertex_ai): forward a non-object anyOf branch for Google to judge instead of crashing * fix(vertex_ai): fold null out of type arrays and union branches before sending to Gemini --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/llms/vertex_ai/common_utils.py | 42 +- .../test_gemini_anyof_null_branches_wire.py | 681 ++++++++++++++++++ tests/unit/llms/vertex_ai/test_vertex.py | 2 +- .../vertex_ai/test_vertex_ai_common_utils.py | 143 ++++ 4 files changed, 843 insertions(+), 25 deletions(-) create mode 100644 tests/integration/providers/test_gemini_anyof_null_branches_wire.py diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 98bc679a78e..581e2cb15f7 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -738,14 +738,14 @@ def build_vertex_schema(parameters: dict, add_property_ordering: bool = False): # refs recursively and correctly detects/skips circular references. unpack_defs(parameters, defs) + _convert_schema_types(parameters) + # 5. Nullable fields: # * https://github.com/pydantic/pydantic/issues/1270 # * https://stackoverflow.com/a/58841311 # * https://github.com/pydantic/pydantic/discussions/4872 convert_anyof_null_to_nullable(parameters) - _convert_schema_types(parameters) - # Handle empty strings in enum values - Gemini doesn't accept empty strings in enums _fix_enum_empty_strings(parameters) @@ -926,34 +926,28 @@ def convert_anyof_null_to_nullable(schema, depth=0): ) """ Converts null objects within anyOf by removing them and adding nullable to all remaining objects """ anyof: Final = schema.get("anyOf", None) - if anyof is not None: - contains_null = False + if isinstance(anyof, list): + non_null_branches: Final = [ + atype for atype in anyof if not (isinstance(atype, dict) and atype.get("type") == "null") + ] + contains_null: Final = len(non_null_branches) < len(anyof) + anyof[:] = non_null_branches or anyof[:1] for atype in anyof: - if isinstance(atype, dict) and atype.get("type") == "null": - # remove null type - anyof.remove(atype) - contains_null = True - elif "type" not in atype and len(atype) == 0: - # Handle empty object case + if not isinstance(atype, dict): + continue + if not atype: atype["type"] = "object" - - if len(anyof) == 0: - # Edge case: response schema with only null type present is invalid in Vertex AI - raise ValueError( - "Invalid input: AnyOf schema with only null type is not supported. Please provide a non-null type." - ) - - if contains_null: - # set all types to nullable following guidance found here: https://cloud.google.com/vertex-ai/generative-ai/docs/samples/generativeaionvertexai-gemini-controlled-generation-response-schema-3#generativeaionvertexai_gemini_controlled_generation_response_schema_3-python - # Empty `items: {}` on array branches is left in place; downstream - # process_items() converts it to {"type": "object"}, which Vertex - # requires whenever type == "array" (even inside anyOf). - for atype in anyof: + if contains_null: + # set all types to nullable following guidance found here: https://cloud.google.com/vertex-ai/generative-ai/docs/samples/generativeaionvertexai-gemini-controlled-generation-response-schema-3#generativeaionvertexai_gemini_controlled_generation_response_schema_3-python + # Empty `items: {}` on array branches is left in place; downstream + # process_items() converts it to {"type": "object"}, which Vertex + # requires whenever type == "array" (even inside anyOf). atype["nullable"] = True + convert_anyof_null_to_nullable(atype, depth=depth + 1) properties: Final = schema.get("properties", None) if properties is not None: - for name, value in properties.items(): + for value in properties.values(): convert_anyof_null_to_nullable(value, depth=depth + 1) items: Final = schema.get("items", None) diff --git a/tests/integration/providers/test_gemini_anyof_null_branches_wire.py b/tests/integration/providers/test_gemini_anyof_null_branches_wire.py new file mode 100644 index 00000000000..9fa99cb69f8 --- /dev/null +++ b/tests/integration/providers/test_gemini_anyof_null_branches_wire.py @@ -0,0 +1,681 @@ +import asyncio +import base64 +import json +import re +import threading +import uuid +from collections.abc import AsyncIterator, Iterator, Mapping +from concurrent.futures import ThreadPoolExecutor +from itertools import product +from types import MappingProxyType +from typing import Final, Literal + +import anthropic +import httpx +import openai +import pytest +import websockets +from integration._support.client import ( + JSON_OBJECT, + Gateway, + Scenario, + gateway_from_environment, + object_value, + string_value, +) +from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process +from integration._support.vertex import service_account_json +from integration._support.wire import Reply, Request, Wire, wire_server +from integration._support.upstream import GEMINI_LIVE_PATH +from integration.providers._gemini_tts_voice_support import live_deployment, live_scenario +from integration.providers._vertex_batch_rows import encoded, gcs_peer, is_gcs_upload, jsonl, upload, vertex_model +from pydantic import JsonValue, TypeAdapter + +Provider = Literal["vertex_ai", "gemini"] +Endpoint = Literal["chat", "messages", "responses"] + +_BACKEND: Final = "gemini-3.8-flash" +_NO_JSON_SCHEMA_BACKEND: Final = "gemini-flash-latest" +_IMAGE_BACKEND: Final = "gemini-3.1-flash-image" +_PROJECT: Final = "scripted-project" +_LOCATION: Final = "us-central1" +_GEMINI_KEY: Final = "sk-scripted-gemini" +_COUNT: Final = 777 +_FAIL_MARKER: Final = "scripted-invalid-argument" +_STRING_ANYOF: Final = re.compile(rb'"anyOf":\s*"') +_SCALAR_BRANCH: Final = re.compile(rb'"anyOf":\s*\[\s*(?:true|false|-?\d)') +_PNG: Final = base64.b64encode(b"\x89PNG\r\n\x1a\nscripted").decode() +_JSON_LIST: Final = TypeAdapter(list[JsonValue]) + +_NULL: Final[Mapping[str, JsonValue]] = {"type": "null"} +_STRING: Final[Mapping[str, JsonValue]] = {"type": "string"} +_INTEGER: Final[Mapping[str, JsonValue]] = {"type": "integer"} +_NULLABLE_STRING: Final[JsonValue] = {"anyOf": [{"type": "string", "nullable": True}]} +_NULLABLE_NULL: Final[JsonValue] = {"anyOf": [{"type": "null", "nullable": True}]} + +_UNIONS: Final[tuple[tuple[str, JsonValue, JsonValue], ...]] = ( + ("null-null-string", {"anyOf": [_NULL, _NULL, _STRING]}, _NULLABLE_STRING), + ("null-only", {"anyOf": [_NULL]}, _NULLABLE_NULL), + ("two-nulls-only", {"anyOf": [_NULL, _NULL]}, _NULLABLE_NULL), + ("string-null", {"anyOf": [_STRING, _NULL]}, _NULLABLE_STRING), + ( + "integer-null-null-string", + {"anyOf": [_INTEGER, _NULL, _NULL, _STRING]}, + {"anyOf": [{"type": "integer", "nullable": True}, {"type": "string", "nullable": True}]}, + ), + ("null-empty", {"anyOf": [_NULL, {}]}, {"anyOf": [{"type": "object", "nullable": True}]}), + ("empty-only", {"anyOf": [{}]}, {"anyOf": [{"type": "object"}]}), + ("empty-list", {"anyOf": []}, {"anyOf": []}), + ("object-anyof", {"anyOf": _STRING}, {"anyOf": _STRING}), + ("null-anyof", {"anyOf": None}, {"anyOf": None}), + ("five-hundred-nulls-then-string", {"anyOf": [*([_NULL] * 500), _STRING]}, _NULLABLE_STRING), + ("string-null-type-array", {"type": ["string", "null"]}, _NULLABLE_STRING), + ( + "null-inside-a-union-branch", + {"anyOf": [{"type": "object", "properties": {"x": {"anyOf": [_INTEGER, _NULL]}}}, _NULL]}, + { + "anyOf": [ + { + "type": "object", + "properties": {"x": {"anyOf": [{"type": "integer", "nullable": True}]}}, + "nullable": True, + } + ] + }, + ), +) +_FANOUT_UNIONS: Final = tuple(case for case in _UNIONS if case[0] in {"null-null-string", "null-only", "object-anyof"}) + + +def _text_reply() -> Mapping[str, JsonValue]: + return { + "candidates": [{"content": {"role": "model", "parts": [{"text": "OK"}]}, "finishReason": "STOP"}], + "usageMetadata": {"promptTokenCount": 9, "candidatesTokenCount": 1, "totalTokenCount": 10}, + "modelVersion": _BACKEND, + } + + +def _image_reply() -> Mapping[str, JsonValue]: + return { + "candidates": [ + {"content": {"role": "model", "parts": [{"inlineData": {"mimeType": "image/png", "data": _PNG}}]}} + ], + "usageMetadata": {"promptTokenCount": 9, "candidatesTokenCount": 1, "totalTokenCount": 10}, + } + + +def _rejection(message: str) -> Reply: + error: Final = {"code": 400, "message": message, "status": "INVALID_ARGUMENT"} + return Reply(status=400, body=json.dumps({"error": error}).encode()) + + +def _google(request: Request) -> Reply: + if is_gcs_upload(request): + return gcs_peer(request) + if _FAIL_MARKER.encode() in request.body: + return _rejection(f"{_FAIL_MARKER}: the scripted tool schema was refused") + if _STRING_ANYOF.search(request.body): + return _rejection("Invalid value at 'tools[0].function_declarations[0].parameters' (anyOf), not a list") + if _SCALAR_BRANCH.search(request.body): + return _rejection( + "Invalid value at 'tools[0].function_declarations[0].parameters.properties[0].value.any_of[0]'" + ) + path: Final = request.target.split("?")[0] + if path.endswith(":countTokens"): + return Reply(body=json.dumps({"totalTokens": _COUNT}).encode()) + if path.endswith(":streamGenerateContent"): + return Reply(content_type="text/event-stream", chunks=(f"data: {json.dumps(_text_reply())}\n\n".encode(),)) + if path.endswith(f"{_IMAGE_BACKEND}:generateContent"): + return Reply(body=json.dumps(_image_reply()).encode()) + if path.endswith(":generateContent"): + return Reply(body=json.dumps(_text_reply()).encode()) + return Reply(status=404, body=json.dumps({"error": {"code": 404, "message": request.target}}).encode()) + + +def _deployment(gateway: Gateway, scenario: Scenario, wire: Wire, provider: Provider, backend: str = _BACKEND) -> str: + if provider == "gemini": + return scenario.model(model=f"gemini/{backend}", api_base=wire.url, api_key=_GEMINI_KEY) + return scenario.model( + model=f"vertex_ai/{backend}", + api_base=wire.url, + api_key=None, + vertex_project=_PROJECT, + vertex_location=_LOCATION, + vertex_credentials=service_account_json(_PROJECT, gateway.upstream_url.rstrip("/")), + ) + + +def _parameters(value: JsonValue) -> dict[str, JsonValue]: + return {"type": "object", "properties": {"value": value}, "required": ["value"]} + + +def _declaration(description: str, value: JsonValue) -> dict[str, JsonValue]: + return {"name": "record", "description": description, "parameters": _parameters(value)} + + +def _tool(endpoint: Endpoint, description: str, parameters: JsonValue) -> dict[str, JsonValue]: + if endpoint == "messages": + return {"name": "record", "description": description, "input_schema": parameters} + if endpoint == "responses": + return {"type": "function", "name": "record", "description": description, "parameters": parameters} + return {"type": "function", "function": {"name": "record", "description": description, "parameters": parameters}} + + +def _request_body( + endpoint: Endpoint, model: str, description: str, parameters: JsonValue, stream: bool +) -> tuple[str, dict[str, JsonValue]]: + prompt: Final = f"Record the value hello by calling the record tool. {description}" + tools: Final[list[JsonValue]] = [_tool(endpoint, description, parameters)] + streaming: Final[dict[str, JsonValue]] = {"stream": True} if stream else {} + if endpoint == "messages": + messages_body: Final[dict[str, JsonValue]] = { + "model": model, + "max_tokens": 16, + "messages": [{"role": "user", "content": prompt}], + "tools": tools, + } + return "/v1/messages", {**messages_body, **streaming} + if endpoint == "responses": + return "/v1/responses", {"model": model, "input": prompt, "tools": tools, **streaming} + chat_body: Final[dict[str, JsonValue]] = { + "model": model, + "max_tokens": 16, + "messages": [{"role": "user", "content": prompt}], + "tools": tools, + } + return "/v1/chat/completions", {**chat_body, **streaming} + + +def _body(request: Request) -> dict[str, JsonValue]: + return JSON_OBJECT.validate_json(request.body) + + +def _json(response: httpx.Response) -> dict[str, JsonValue]: + return JSON_OBJECT.validate_json(response.content) + + +def _generate_requests(wire: Wire) -> tuple[Request, ...]: + return tuple( + request + for request in wire.drain() + if request.target.split("?")[0].endswith((":generateContent", ":streamGenerateContent")) + ) + + +def _sent_declarations(body: Mapping[str, JsonValue]) -> JsonValue: + tools: Final = _JSON_LIST.validate_python(body["tools"]) + assert len(tools) == 1, tools + tool: Final = object_value(tools[0]) + assert len(tool) == 1, tool + return next(iter(tool.values())) + + +def _description() -> str: + return f"Record a value {uuid.uuid4().hex[:12]}" + + +def _assert_caller_saw_ok(endpoint: Endpoint, response: httpx.Response, stream: bool) -> None: + assert response.status_code == 200, response.text + if stream: + assert "OK" in response.text, response.text + return + payload: Final = _json(response) + assert string_value(payload["id"]), payload + if endpoint == "chat": + choice: Final = object_value(_JSON_LIST.validate_python(payload["choices"])[0]) + assert object_value(choice["message"])["content"] == "OK", payload + return + assert "OK" in response.text, payload + + +@pytest.mark.parametrize("provider", ("vertex_ai", "gemini")) +@pytest.mark.parametrize(("union", "expected"), [case[1:] for case in _UNIONS], ids=[case[0] for case in _UNIONS]) +def test_chat_tool_union_reaches_google_with_null_branches_folded_into_nullable( + gateway: Gateway, provider: Provider, union: JsonValue, expected: JsonValue +) -> None: + with wire_server(_google) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire, provider) + description: Final = _description() + path, body = _request_body("chat", model, description, _parameters(union), False) + response: Final = gateway.request("POST", path, body) + _assert_caller_saw_ok("chat", response, False) + requests: Final = _generate_requests(wire) + assert len(requests) == 1, [request.target for request in requests] + assert requests[0].target.split("?")[0].endswith(f"{_BACKEND}:generateContent"), requests[0].target + assert _sent_declarations(_body(requests[0])) == [_declaration(description, expected)] + + +@pytest.mark.parametrize("stream", (False, True), ids=("whole", "stream")) +@pytest.mark.parametrize("endpoint", ("chat", "messages", "responses")) +@pytest.mark.parametrize("provider", ("vertex_ai", "gemini")) +@pytest.mark.parametrize( + ("union", "expected"), [case[1:] for case in _FANOUT_UNIONS], ids=[case[0] for case in _FANOUT_UNIONS] +) +def test_every_unified_endpoint_sends_the_same_folded_union( + gateway: Gateway, provider: Provider, endpoint: Endpoint, stream: bool, union: JsonValue, expected: JsonValue +) -> None: + with wire_server(_google) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire, provider) + description: Final = _description() + path, body = _request_body(endpoint, model, description, _parameters(union), stream) + response: Final = gateway.request("POST", path, body) + _assert_caller_saw_ok(endpoint, response, stream) + requests: Final = _generate_requests(wire) + assert len(requests) == 1, [request.target for request in requests] + method: Final = "streamGenerateContent" if stream else "generateContent" + assert requests[0].target.split("?")[0].endswith(f"{_BACKEND}:{method}"), requests[0].target + assert _sent_declarations(_body(requests[0])) == [_declaration(description, expected)] + + +def test_nested_item_and_property_unions_are_folded_at_every_depth(gateway: Gateway) -> None: + parameters: Final[dict[str, JsonValue]] = { + "type": "object", + "properties": { + "values": {"type": "array", "items": {"anyOf": [_NULL, _NULL, _INTEGER]}}, + "nested": {"type": "object", "properties": {"inner": {"anyOf": [_NULL, _NULL, _STRING]}}}, + }, + "required": ["values", "nested"], + } + with wire_server(_google) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire, "vertex_ai") + description: Final = _description() + path, body = _request_body("chat", model, description, parameters, False) + response: Final = gateway.request("POST", path, body) + _assert_caller_saw_ok("chat", response, False) + requests: Final = _generate_requests(wire) + assert len(requests) == 1, [request.target for request in requests] + assert _sent_declarations(_body(requests[0])) == [ + { + "name": "record", + "description": description, + "parameters": { + "type": "object", + "properties": { + "values": {"type": "array", "items": {"anyOf": [{"type": "integer", "nullable": True}]}}, + "nested": {"type": "object", "properties": {"inner": _NULLABLE_STRING}}, + }, + "required": ["values", "nested"], + }, + } + ] + + +def test_a_string_anyof_reaches_google_and_its_400_reaches_the_caller(gateway: Gateway) -> None: + with wire_server(_google) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire, "gemini") + description: Final = _description() + path, body = _request_body("chat", model, description, _parameters({"anyOf": "string"}), False) + response: Final = gateway.request("POST", path, body) + assert response.status_code == 400, response.text + assert "not a list" in response.text, response.text + requests: Final = _generate_requests(wire) + assert len(requests) == 1, [request.target for request in requests] + assert _sent_declarations(_body(requests[0])) == [_declaration(description, {"anyOf": "string"})] + + +@pytest.mark.parametrize("provider", ("vertex_ai", "gemini")) +def test_a_google_rejection_of_a_folded_union_reaches_the_caller(gateway: Gateway, provider: Provider) -> None: + with wire_server(_google) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire, provider) + description: Final = f"{_FAIL_MARKER} {uuid.uuid4().hex[:12]}" + path, body = _request_body("chat", model, description, _parameters({"anyOf": [_NULL, _NULL, _STRING]}), False) + response: Final = gateway.request("POST", path, body) + assert response.status_code == 400, response.text + assert _FAIL_MARKER in response.text, response.text + requests: Final = _generate_requests(wire) + assert len(requests) == 1, [request.target for request in requests] + assert _sent_declarations(_body(requests[0])) == [_declaration(description, _NULLABLE_STRING)] + + +@pytest.mark.parametrize( + ("union", "sent"), + (({"anyOf": [True, _NULL]}, {"anyOf": [True]}), ({"anyOf": [_NULL, 5]}, {"anyOf": [5]})), + ids=("true-null", "null-int"), +) +@pytest.mark.parametrize("provider", ("vertex_ai", "gemini")) +def test_a_union_branch_that_is_not_an_object_reaches_google_and_its_400_reaches_the_caller( + gateway: Gateway, provider: Provider, union: JsonValue, sent: JsonValue +) -> None: + with wire_server(_google) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire, provider) + rejected_description: Final = _description() + path, broken = _request_body("chat", model, rejected_description, _parameters(union), False) + failed: Final = gateway.request("POST", path, broken) + assert failed.status_code == 400, failed.text + assert "any_of[0]" in failed.text, failed.text + rejected: Final = _generate_requests(wire) + assert len(rejected) == 1, [request.target for request in rejected] + assert _sent_declarations(_body(rejected[0])) == [_declaration(rejected_description, sent)] + description: Final = _description() + _, valid = _request_body("chat", model, description, _parameters({"anyOf": [_NULL, _NULL, _STRING]}), False) + recovered: Final = gateway.request("POST", path, valid) + _assert_caller_saw_ok("chat", recovered, False) + requests: Final = _generate_requests(wire) + assert len(requests) == 1, [request.target for request in requests] + assert _sent_declarations(_body(requests[0])) == [_declaration(description, _NULLABLE_STRING)] + + +def test_a_repeated_identical_request_sends_the_same_folded_union_each_time(gateway: Gateway) -> None: + with wire_server(_google) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire, "vertex_ai") + description: Final = _description() + path, body = _request_body("chat", model, description, _parameters({"anyOf": [_NULL, _NULL, _STRING]}), False) + first: Final = gateway.request("POST", path, body) + second: Final = gateway.request("POST", path, body) + _assert_caller_saw_ok("chat", first, False) + _assert_caller_saw_ok("chat", second, False) + assert _json(first)["id"] != _json(second)["id"] + requests: Final = _generate_requests(wire) + assert len(requests) == 2, [request.target for request in requests] + assert [_sent_declarations(_body(request)) for request in requests] == [ + [_declaration(description, _NULLABLE_STRING)] + ] * 2 + + +def test_openai_sdk_chat_sends_the_folded_union(gateway: Gateway) -> None: + with wire_server(_google) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire, "gemini") + description: Final = _description() + client: Final = openai.OpenAI( + base_url=f"{str(gateway.client.base_url).rstrip('/')}/v1", api_key=gateway.key, max_retries=0 + ) + completion: Final = client.chat.completions.create( + model=model, + max_tokens=16, + messages=[{"role": "user", "content": "Record the value hello by calling the record tool."}], + tools=[_tool("chat", description, _parameters({"anyOf": [_NULL, _NULL, _STRING]}))], # pyright: ignore[reportArgumentType] # SDK TypedDicts reject a JsonValue mapping + ) + assert completion.choices[0].message.content == "OK" + requests: Final = _generate_requests(wire) + assert len(requests) == 1, [request.target for request in requests] + assert _sent_declarations(_body(requests[0])) == [_declaration(description, _NULLABLE_STRING)] + + +async def _async_openai_stream_text(gateway: Gateway, model: str, description: str) -> str: + client: Final = openai.AsyncOpenAI( + base_url=f"{str(gateway.client.base_url).rstrip('/')}/v1", api_key=gateway.key, max_retries=0 + ) + stream: Final = await client.chat.completions.create( + model=model, + max_tokens=16, + stream=True, + messages=[{"role": "user", "content": "Record the value hello by calling the record tool."}], + tools=[_tool("chat", description, _parameters({"anyOf": [_NULL]}))], # pyright: ignore[reportArgumentType] # SDK TypedDicts reject a JsonValue mapping + ) + return "".join([chunk.choices[0].delta.content or "" async for chunk in stream if chunk.choices]) + + +def test_async_openai_sdk_stream_sends_the_null_only_union_as_one_nullable_null(gateway: Gateway) -> None: + with wire_server(_google) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire, "vertex_ai") + description: Final = _description() + assert asyncio.run(_async_openai_stream_text(gateway, model, description)) == "OK" + requests: Final = _generate_requests(wire) + assert len(requests) == 1, [request.target for request in requests] + assert requests[0].target.split("?")[0].endswith(":streamGenerateContent"), requests[0].target + assert _sent_declarations(_body(requests[0])) == [_declaration(description, _NULLABLE_NULL)] + + +def test_anthropic_sdk_messages_sends_the_folded_union(gateway: Gateway) -> None: + with wire_server(_google) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire, "vertex_ai") + description: Final = _description() + client: Final = anthropic.Anthropic(base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0) + message: Final = client.messages.create( + model=model, + max_tokens=16, + messages=[{"role": "user", "content": "Record the value hello by calling the record tool."}], + tools=[_tool("messages", description, _parameters({"anyOf": [_NULL, _NULL]}))], # pyright: ignore[reportArgumentType] # SDK TypedDicts reject a JsonValue mapping + ) + assert [block.type for block in message.content] == ["text"] + requests: Final = _generate_requests(wire) + assert len(requests) == 1, [request.target for request in requests] + assert _sent_declarations(_body(requests[0])) == [_declaration(description, _NULLABLE_NULL)] + + +@pytest.mark.parametrize( + ("union", "expected"), [case[1:] for case in _FANOUT_UNIONS], ids=[case[0] for case in _FANOUT_UNIONS] +) +def test_vertex_count_tokens_sends_the_folded_union(gateway: Gateway, union: JsonValue, expected: JsonValue) -> None: + with wire_server(_google) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire, "vertex_ai") + description: Final = _description() + response: Final = gateway.request( + "POST", + "/v1/messages/count_tokens", + { + "model": model, + "messages": [{"role": "user", "content": "Record the value hello by calling the record tool."}], + "tools": [_tool("messages", description, _parameters(union))], + }, + ) + assert response.status_code == 200, response.text + assert _json(response)["input_tokens"] == _COUNT, response.text + requests: Final = tuple(request for request in wire.drain() if request.target.endswith(":countTokens")) + assert len(requests) == 1, [request.target for request in requests] + assert _sent_declarations(_body(requests[0])) == [_declaration(description, expected)] + + +@pytest.mark.parametrize( + ("union", "expected"), [case[1:] for case in _FANOUT_UNIONS], ids=[case[0] for case in _FANOUT_UNIONS] +) +def test_gemini_token_counter_sends_the_folded_union(gateway: Gateway, union: JsonValue, expected: JsonValue) -> None: + with wire_server(_google) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire, "gemini") + description: Final = _description() + response: Final = gateway.request( + "POST", + "/utils/token_counter", + { + "model": model, + "contents": [ + {"role": "user", "parts": [{"text": "Record the value hello by calling the record tool."}]} + ], + "tools": [_tool("chat", description, _parameters(union))], + }, + params={"call_endpoint": "true"}, + ) + assert response.status_code == 200, response.text + assert _json(response)["total_tokens"] == _COUNT, response.text + requests: Final = tuple(request for request in wire.drain() if ":countTokens" in request.target) + assert len(requests) == 1, [request.target for request in requests] + sent: Final = object_value(_body(requests[0])["generateContentRequest"]) + assert _sent_declarations(sent) == [_declaration(description, expected)] + + +def test_response_format_on_a_model_without_json_schema_support_folds_the_union(gateway: Gateway) -> None: + with wire_server(_google) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire, "gemini", _NO_JSON_SCHEMA_BACKEND) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": "Return the value as JSON."}], + "response_format": { + "type": "json_schema", + "json_schema": {"name": "record", "schema": _parameters({"anyOf": [_NULL, _NULL, _STRING]})}, + }, + }, + ) + _assert_caller_saw_ok("chat", response, False) + requests: Final = _generate_requests(wire) + assert len(requests) == 1, [request.target for request in requests] + generation_config: Final = object_value(_body(requests[0])["generationConfig"]) + assert generation_config["response_mime_type"] == "application/json", generation_config + assert generation_config["response_schema"] == {**_parameters(_NULLABLE_STRING), "propertyOrdering": ["value"]} + + +def test_native_generate_content_response_schema_folds_the_union(gateway: Gateway) -> None: + with wire_server(_google) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire, "gemini", _NO_JSON_SCHEMA_BACKEND) + response: Final = gateway.request( + "POST", + f"/v1beta/models/{model}:generateContent", + { + "contents": [{"role": "user", "parts": [{"text": "Return the value as JSON."}]}], + "generationConfig": { + "responseMimeType": "application/json", + "responseSchema": _parameters({"anyOf": [_NULL, _NULL, _STRING]}), + }, + }, + ) + assert response.status_code == 200, response.text + requests: Final = _generate_requests(wire) + assert len(requests) == 1, [request.target for request in requests] + generation_config: Final = object_value(_body(requests[0])["generationConfig"]) + assert generation_config["responseSchema"] == {**_parameters(_NULLABLE_STRING), "propertyOrdering": ["value"]} + + +@pytest.mark.parametrize( + ("union", "expected"), [case[1:] for case in _FANOUT_UNIONS], ids=[case[0] for case in _FANOUT_UNIONS] +) +def test_image_generation_tools_send_the_folded_union(gateway: Gateway, union: JsonValue, expected: JsonValue) -> None: + with wire_server(_google) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire, "gemini", _IMAGE_BACKEND) + description: Final = _description() + response: Final = gateway.request( + "POST", + "/v1/images/generations", + { + "model": model, + "prompt": "A small red square on a white background", + "tools": [_tool("chat", description, _parameters(union))], + }, + ) + assert response.status_code == 200, response.text + assert len(_JSON_LIST.validate_python(_json(response)["data"])) == 1, response.text + requests: Final = _generate_requests(wire) + assert len(requests) == 1, [request.target for request in requests] + assert _sent_declarations(_body(requests[0])) == [_declaration(description, expected)] + + +def test_vertex_batch_row_is_stored_with_the_folded_union(gateway: Gateway) -> None: + with wire_server(_google) as wire, gateway.scenario() as scenario: + model: Final = vertex_model(gateway, scenario, wire.url, _BACKEND) + description: Final = _description() + row: Final[dict[str, JsonValue]] = { + "custom_id": "anyof-row", + "method": "POST", + "url": "/v1/chat/completions", + "body": { + "model": model, + "messages": [{"role": "user", "content": "Record the value hello by calling the record tool."}], + "tools": [_tool("chat", description, _parameters({"anyOf": [_NULL, _NULL, _STRING]}))], + }, + } + response: Final = upload(gateway, "model", model, jsonl(encoded(row))) + assert response.status_code == 200, response.text + uploads: Final = tuple(request for request in wire.drain() if is_gcs_upload(request)) + assert len(uploads) == 1, [request.target for request in uploads] + stored: Final = tuple( + object_value(JSON_OBJECT.validate_json(line)["request"]) for line in uploads[0].body.splitlines() if line + ) + assert len(stored) == 1, stored + assert _sent_declarations(stored[0]) == [_declaration(description, _NULLABLE_STRING)] + + +async def _event_types_until_done(socket: websockets.ClientConnection) -> AsyncIterator[str]: + async for message in socket: + kind: Final = string_value(JSON_OBJECT.validate_json(message)["type"]) + yield kind + if kind in {"response.done", "error"}: + return + + +async def _realtime_session_with_tools(gateway: Gateway, model: str, tools: list[JsonValue]) -> tuple[str, ...]: + ws_base: Final = str(gateway.client.base_url).rstrip("/").replace("http://", "ws://") + headers: Final = {"Authorization": f"Bearer {gateway.key}"} + async with websockets.connect(f"{ws_base}/v1/realtime?model={model}", additional_headers=headers) as socket: + first: Final = JSON_OBJECT.validate_json(await socket.recv()) + await socket.send(json.dumps({"type": "session.update", "session": {"tools": tools}})) + item: Final = {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "say the line"}]} + await socket.send(json.dumps({"type": "conversation.item.create", "item": item})) + events: Final = [kind async for kind in _event_types_until_done(socket)] + return (string_value(first["type"]), *events) + + +def _realtime_types(gateway: Gateway, model: str, tools: list[JsonValue]) -> tuple[str, ...]: + async def bounded() -> tuple[str, ...]: + try: + return await asyncio.wait_for(_realtime_session_with_tools(gateway, model, tools), 30) + except TimeoutError: + return ("timeout",) + + return asyncio.run(bounded()) + + +def _live_frames(gateway: Gateway, scenario_id: str) -> tuple[dict[str, JsonValue], ...]: + with httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream: + observed: Final = tuple( + map(object_value, _JSON_LIST.validate_python(_json(upstream.get("/__observations"))["requests"])) + ) + return tuple( + object_value(request["body"]) + for request in observed + if request["api_key"] == scenario_id + and request["method"] == "WEBSOCKET_FRAME" + and request["path"] == GEMINI_LIVE_PATH + ) + + +@pytest.fixture(scope="module") +def deferred_live(tmp_path_factory: pytest.TempPathFactory) -> Iterator[OwnedProxy]: + directory: Final = tmp_path_factory.mktemp("anyof-live") + with gateway_from_environment() as gateway: + with owned_proxy_process(gateway, directory, {"LITELLM_GEMINI_LIVE_DEFER_SETUP": "true"}, workers=2) as proxy: + yield proxy + + +@pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +def test_deferred_realtime_setup_carries_session_tools_with_the_folded_union(deferred_live: OwnedProxy) -> None: + gateway: Final = deferred_live.gateway + with gateway.scenario() as scenario: + handle: Final = live_scenario(scenario) + model: Final = live_deployment(scenario, handle.scenario_id, gateway.upstream_url) + description: Final = _description() + tools: Final[list[JsonValue]] = [_tool("chat", description, _parameters({"anyOf": [_NULL, _NULL, _STRING]}))] + types: Final = _realtime_types(gateway, model, tools) + assert types[0] == "session.created", types + assert types[-1] == "response.done", types + frames: Final = _live_frames(gateway, handle.scenario_id) + setups: Final = tuple(object_value(frame["setup"]) for frame in frames if "setup" in frame) + assert len(setups) == 1, [sorted(frame) for frame in frames] + assert _sent_declarations(setups[0]) == [_declaration(description, _NULLABLE_STRING)] + + +_PROVIDERS: Final[tuple[Provider, ...]] = ("vertex_ai", "gemini") +_ENDPOINTS: Final[tuple[Endpoint, ...]] = ("chat", "messages", "responses") +_BURST: Final[tuple[tuple[Provider, Endpoint, bool, bool], ...]] = tuple( + (provider, endpoint, stream, index % 4 == 0) + for index, (provider, endpoint, stream, _) in enumerate(product(_PROVIDERS, _ENDPOINTS, (False, True), range(2))) +) + + +def test_a_concurrent_mixed_burst_folds_every_union_once_and_keeps_each_status(gateway: Gateway) -> None: + with wire_server(_google) as wire, gateway.scenario() as scenario: + models: Final = MappingProxyType( + {provider: _deployment(gateway, scenario, wire, provider) for provider in _PROVIDERS} + ) + shots: Final = tuple( + (provider, endpoint, stream, f"{_FAIL_MARKER if refused else 'burst'} {uuid.uuid4().hex}") + for provider, endpoint, stream, refused in _BURST + ) + + def send(shot: tuple[Provider, Endpoint, bool, str]) -> httpx.Response: + provider, endpoint, stream, description = shot + path, body = _request_body( + endpoint, models[provider], description, _parameters({"anyOf": [_NULL, _NULL, _STRING]}), stream + ) + return gateway.request("POST", path, body) + + with ThreadPoolExecutor(max_workers=len(shots)) as pool: + responses: Final = tuple(pool.map(send, shots)) + expected_statuses: Final = [400 if _FAIL_MARKER in shot[3] else 200 for shot in shots] + assert [response.status_code for response in responses] == expected_statuses, [ + response.text[:120] for response in responses if response.status_code not in {200, 400} + ] + requests: Final = _generate_requests(wire) + assert len(requests) == len(shots), len(requests) + sent: Final = sorted((json.dumps(_sent_declarations(_body(request)), sort_keys=True) for request in requests)) + assert sent == sorted(json.dumps([_declaration(shot[3], _NULLABLE_STRING)], sort_keys=True) for shot in shots) diff --git a/tests/unit/llms/vertex_ai/test_vertex.py b/tests/unit/llms/vertex_ai/test_vertex.py index 044a07859f9..04cf391c700 100644 --- a/tests/unit/llms/vertex_ai/test_vertex.py +++ b/tests/unit/llms/vertex_ai/test_vertex.py @@ -1814,8 +1814,8 @@ def test_gemini_nullable_object_tool_schema_httpx( "type": "object", "required": ["source"], "properties": {"source": {"type": "string"}}, + "nullable": True, }, - {"type": "null"}, ] diff --git a/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py b/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py index 8528724a24e..d244fe6bd90 100644 --- a/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py +++ b/tests/unit/llms/vertex_ai/test_vertex_ai_common_utils.py @@ -1,3 +1,4 @@ +import copy import time from collections.abc import Iterator from contextlib import contextmanager @@ -14,6 +15,7 @@ from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.llms.gemini.count_tokens import handler as gemini_count_tokens_handler from litellm.llms.vertex_ai.common_utils import ( VertexAITokenCounter, + build_vertex_schema, convert_anyof_null_to_nullable, convert_vertex_datetime_to_openai_datetime, get_vertex_location_from_url, @@ -234,6 +236,147 @@ def test_nested_anyof_conversion(): assert schema == expected +def test_anyof_conversion_with_multiple_null_branches(): + """Adjacent null branches must all be removed.""" + schema: Final = { + "type": "object", + "properties": { + "example": { + "anyOf": [ + {"type": "null"}, + {"type": "null"}, + {"type": "string"}, + ] + } + }, + } + + convert_anyof_null_to_nullable(schema) + + expected: Final = { + "type": "object", + "properties": {"example": {"anyOf": [{"type": "string", "nullable": True}]}}, + } + assert schema == expected + + +def test_anyof_conversion_with_null_before_empty_object(): + """An empty-object branch following a null branch still gets a type.""" + schema: Final = { + "type": "object", + "properties": {"example": {"anyOf": [{"type": "null"}, {}]}}, + } + + convert_anyof_null_to_nullable(schema) + + expected: Final = { + "type": "object", + "properties": {"example": {"anyOf": [{"type": "object", "nullable": True}]}}, + } + assert schema == expected + + +@pytest.mark.parametrize("null_count", [1, 2, 4]) +def test_anyof_conversion_with_only_null_branches_keeps_one_nullable_null(null_count: int): + """Vertex accepts a NULL type, so a null-only union is forwarded as one nullable null branch.""" + schema: Final = { + "type": "object", + "properties": {"example": {"anyOf": [{"type": "null"} for _ in range(null_count)]}}, + } + + convert_anyof_null_to_nullable(schema) + + expected: Final = { + "type": "object", + "properties": {"example": {"anyOf": [{"type": "null", "nullable": True}]}}, + } + assert schema == expected + + +def test_anyof_conversion_leaves_empty_anyof_for_the_provider_to_judge(): + """An empty union has no null branch to strip, so it is forwarded unchanged.""" + schema: Final = {"type": "object", "properties": {"example": {"anyOf": []}}} + + convert_anyof_null_to_nullable(schema) + + assert schema == {"type": "object", "properties": {"example": {"anyOf": []}}} + + +@pytest.mark.parametrize("malformed", ["string", {"type": "string"}], ids=["string", "object"]) +def test_anyof_conversion_forwards_a_non_list_anyof_unchanged(malformed: JsonValue): + """A malformed union is the provider's to judge, as it was before null branches were stripped in one pass.""" + schema: Final = {"type": "object", "properties": {"example": {"anyOf": copy.deepcopy(malformed)}}} + + convert_anyof_null_to_nullable(schema) + + assert schema == {"type": "object", "properties": {"example": {"anyOf": malformed}}} + + +@pytest.mark.parametrize( + ("branches", "expected"), + [ + ([{"type": "null"}, 5], [5]), + ([True, {"type": "null"}], [True]), + ([{"type": "null"}, 5, {"type": "string"}], [5, {"type": "string", "nullable": True}]), + ], + ids=["null-int", "true-null", "null-int-string"], +) +def test_anyof_conversion_forwards_a_non_object_branch_for_the_provider_to_judge( + branches: list[JsonValue], expected: list[JsonValue] +): + schema: Final = {"type": "object", "properties": {"example": {"anyOf": branches}}} + + convert_anyof_null_to_nullable(schema) + + assert schema == {"type": "object", "properties": {"example": {"anyOf": expected}}} + + + +@pytest.mark.parametrize( + ("field", "expected"), + [ + ({"type": ["integer", "null"]}, {"anyOf": [{"type": "integer", "nullable": True}]}), + ( + { + "anyOf": [ + {"type": "object", "properties": {"x": {"anyOf": [{"type": "integer"}, {"type": "null"}]}}}, + {"type": "null"}, + ] + }, + { + "anyOf": [ + { + "type": "object", + "properties": {"x": {"anyOf": [{"type": "integer", "nullable": True}]}}, + "nullable": True, + } + ] + }, + ), + ( + {"type": ["object", "null"], "properties": {"x": {"type": ["string", "null"]}}}, + { + "anyOf": [ + { + "type": "object", + "properties": {"x": {"anyOf": [{"type": "string", "nullable": True}]}}, + "nullable": True, + } + ] + }, + ), + ( + {"type": "array", "items": {"type": ["number", "null"]}}, + {"type": "array", "items": {"anyOf": [{"type": "number", "nullable": True}]}}, + ), + ], + ids=["type-array", "null-inside-a-union-branch", "nullable-object-with-nullable-property", "nullable-items"], +) +def test_vertex_schema_folds_every_null_branch_into_nullable(field: JsonValue, expected: JsonValue): + schema: Final = build_vertex_schema({"type": "object", "properties": {"value": field}}) + + assert schema["properties"]["value"] == expected + def test_anyof_with_excessive_nesting(): """Test conversion with excessive nesting > max levels +1 deep.""" # generate a schema with excessive nesting From ecbea93381782591e3e8bf90af145aa68163484a Mon Sep 17 00:00:00 2001 From: Rick <26716961+Bytechoreographer@users.noreply.github.com> Date: Fri, 16 Oct 2026 16:01:06 +0800 Subject: [PATCH 10/23] perf(ui): page the member tables and look up budget rows once per render (#26667) * perf(ui): look up a team member's budget row once per render TeamMemberTab resolved each row's membership with Array.find() over team_memberships inside every extra column's render and sortValue, so a render walked the full membership list once per cell. Build the user_id -> membership map once per team_memberships change and read from it instead. * perf(ui): page the member tables instead of rendering every row MemberTable rendered one row per member, so a team or an org with a couple of thousand members paid for every row on each render. Its DataTable now runs client-side pagination, like the other list tables in the dashboard. * fix(ui): keep the member table on its page when the roster is refetched * test(ui): cover a team member's spend and budget refreshing after a refetch --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../common_components/MemberTable.test.tsx | 62 +++++++++++++- .../common_components/MemberTable.tsx | 30 +++++-- .../organization/organization_view.test.tsx | 33 +++++++ .../components/team/TeamMemberTab.test.tsx | 85 +++++++++++++++++++ .../src/components/team/TeamMemberTab.tsx | 32 +++---- 5 files changed, 221 insertions(+), 21 deletions(-) diff --git a/ui/litellm-dashboard/src/components/common_components/MemberTable.test.tsx b/ui/litellm-dashboard/src/components/common_components/MemberTable.test.tsx index 418f27c8188..3264c3a99ea 100644 --- a/ui/litellm-dashboard/src/components/common_components/MemberTable.test.tsx +++ b/ui/litellm-dashboard/src/components/common_components/MemberTable.test.tsx @@ -1,4 +1,4 @@ -import { fireEvent, screen, waitFor, within } from "@testing-library/react"; +import { act, fireEvent, screen, waitFor, within } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import React from "react"; import { describe, expect, it, vi } from "vitest"; @@ -160,6 +160,66 @@ describe("MemberTable sorting", () => { }); }); +describe("MemberTable paging", () => { + const roster = (): Member[] => + Array.from({ length: 30 }, (_, index) => ({ + user_id: `u-${index}`, + user_email: `member${index}@example.com`, + user_alias: `Member ${String(index).padStart(2, "0")}`, + role: index === 0 ? "admin" : "user", + })); + + const renderRoster = () => { + const props = { members: roster(), canEdit: true, onEdit: vi.fn(), onDelete: vi.fn() }; + const view = renderWithProviders(); + return { props, view }; + }; + + const goToSecondPage = async (user: ReturnType) => { + await user.click(screen.getByTestId("pagination-next")); + expect(screen.getByText("Showing 26-30 of 30")).toBeInTheDocument(); + }; + + it("stays on the current page when the roster is refetched", async () => { + const user = userEvent.setup(); + const { props, view } = renderRoster(); + await goToSecondPage(user); + + view.rerender(); + await act(() => new Promise((resolve) => setTimeout(resolve, 20))); + + expect(screen.getByText("Showing 26-30 of 30")).toBeInTheDocument(); + }); + + it("returns to the first page when the sort order changes", async () => { + const user = userEvent.setup(); + renderRoster(); + await goToSecondPage(user); + + await user.click(screen.getByTestId("sort-header-user_email")); + + expect(await screen.findByText("Showing 1-25 of 30")).toBeInTheDocument(); + }); + + it("returns to the first page when the search or the role filter changes", async () => { + const user = userEvent.setup(); + renderRoster(); + await goToSecondPage(user); + + search("example.com"); + expect(await screen.findByText("Showing 1-25 of 30")).toBeInTheDocument(); + + search(""); + await goToSecondPage(user); + await user.click(screen.getByTestId("datatable-filters-trigger")); + await user.click(screen.getByTestId("filter-role")); + await user.click(await screen.findByRole("option", { name: "user" })); + await user.click(screen.getByTestId("filter-drawer-apply")); + + expect(await screen.findByText("Showing 1-25 of 29")).toBeInTheDocument(); + }); +}); + describe("MemberTable role filter", () => { it("shows only members with the chosen role and clears on reset", async () => { const user = userEvent.setup(); diff --git a/ui/litellm-dashboard/src/components/common_components/MemberTable.tsx b/ui/litellm-dashboard/src/components/common_components/MemberTable.tsx index 24dd86c9e2c..0acba9ad21d 100644 --- a/ui/litellm-dashboard/src/components/common_components/MemberTable.tsx +++ b/ui/litellm-dashboard/src/components/common_components/MemberTable.tsx @@ -1,4 +1,4 @@ -import type { ColumnDef, ColumnFiltersState } from "@tanstack/react-table"; +import type { ColumnDef, ColumnFiltersState, OnChangeFn, PaginationState, SortingState } from "@tanstack/react-table"; import { Crown, Info, User, UserPlus } from "lucide-react"; import React, { useState } from "react"; @@ -9,6 +9,7 @@ import { DataTableFilterField, DataTableSortHeader, DataTableToolbar, + DEFAULT_PAGE_SIZE_OPTIONS, } from "@/components/shared/DataTable"; import { StatusBadge } from "@/components/shared/table_cells"; import { Button } from "@/components/ui/button"; @@ -44,6 +45,8 @@ export interface MemberTableProps { const ALL_ROLES = "all"; +const DEFAULT_SORTING: SortingState = [{ id: "user_alias", desc: false }]; + export const memberRowId = (member: Member): string => member.user_id ?? member.user_email ?? JSON.stringify(member); export const memberRoleOptions = (members: readonly Member[]): string[] => @@ -230,8 +233,21 @@ export default function MemberTable({ }: MemberTableProps) { const [globalFilter, setGlobalFilter] = useState(""); const [columnFilters, setColumnFilters] = useState([]); + const [sorting, setSorting] = useState(DEFAULT_SORTING); + const [pagination, setPagination] = useState({ + pageIndex: 0, + pageSize: DEFAULT_PAGE_SIZE_OPTIONS[0], + }); const [filtersOpen, setFiltersOpen] = useState(false); + const backToFirstPage = + (onChange: OnChangeFn): OnChangeFn => + (updaterOrValue) => { + onChange(updaterOrValue); + setPagination((current) => ({ ...current, pageIndex: 0 })); + }; + const changeSearch = backToFirstPage(setGlobalFilter); + const columnDeps: MemberColumnDeps = { canEdit, onEdit, @@ -260,13 +276,17 @@ export default function MemberTable({ data={members} columns={columns} getRowId={memberRowId} + paginationMode="client" + pagination={pagination} + onPaginationChange={setPagination} sortingMode="client" - defaultSorting={[{ id: "user_alias", desc: false }]} + sorting={sorting} + onSortingChange={backToFirstPage(setSorting)} filterMode="client" columnFilters={columnFilters} - onColumnFiltersChange={setColumnFilters} + onColumnFiltersChange={backToFirstPage(setColumnFilters)} globalFilter={globalFilter} - onGlobalFilterChange={setGlobalFilter} + onGlobalFilterChange={changeSearch} noDataMessage={ {isNarrowed ? "No members match your search or filters" : emptyText ?? "No data"} @@ -277,7 +297,7 @@ export default function MemberTable({ setFiltersOpen(true)} showViewOptions={false} diff --git a/ui/litellm-dashboard/src/components/organization/organization_view.test.tsx b/ui/litellm-dashboard/src/components/organization/organization_view.test.tsx index ca792084f2a..88d98fac60e 100644 --- a/ui/litellm-dashboard/src/components/organization/organization_view.test.tsx +++ b/ui/litellm-dashboard/src/components/organization/organization_view.test.tsx @@ -149,6 +149,39 @@ test("should display empty state when organization has no members", async () => }); }); +test("keeps the members pager on its page when an admin opens a member's edit dialog", async () => { + const members = Array.from({ length: 30 }, (_, index) => ({ + user_id: `user${index + 1}`, + user_email: `user${index + 1}@example.com`, + user_role: "internal_user", + spend: 0, + user: { user_alias: `alias${String(index + 1).padStart(2, "0")}` }, + })); + mockUseOrganization.mockReturnValue({ data: { ...mockOrg, members }, isLoading: false } as unknown as ReturnType< + typeof useOrganization + >); + + const user = userEvent.setup(); + renderWithProviders( + {}} + accessToken="test-token" + is_org_admin={true} + is_proxy_admin={false} + userModels={[]} + />, + ); + + await user.click(screen.getByRole("tab", { name: "Members" })); + await user.click(screen.getByTestId("pagination-next")); + expect(screen.getByText("Showing 26-30 of 30")).toBeInTheDocument(); + + await user.click(within(screen.getByRole("row", { name: /alias27/ })).getByTestId("edit-member")); + + expect(screen.getByText("Showing 26-30 of 30")).toBeInTheDocument(); +}); + test("should display team aliases when teams are available", async () => { const orgWithTeams = { ...mockOrg, diff --git a/ui/litellm-dashboard/src/components/team/TeamMemberTab.test.tsx b/ui/litellm-dashboard/src/components/team/TeamMemberTab.test.tsx index b8f292222ed..cd51bcedb7b 100644 --- a/ui/litellm-dashboard/src/components/team/TeamMemberTab.test.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamMemberTab.test.tsx @@ -316,6 +316,91 @@ describe("TeamMembersComponent", () => { expect(screen.getByText("Jul 15, 2026")).toBeInTheDocument(); }); + it("shows each member the values from its own membership row", () => { + const baseMembership = createMockTeamData().team_memberships[0]; + const teamData = createMockTeamData({ + team_memberships: [ + { + ...baseMembership, + spend: 1.25, + total_spend: 10, + litellm_budget_table: { ...baseMembership.litellm_budget_table, max_budget: 250 }, + }, + { + ...baseMembership, + user_id: "user2@test.com", + spend: 7.5, + total_spend: 20, + litellm_budget_table: { + ...baseMembership.litellm_budget_table, + max_budget: 2500, + rpm_limit: 7, + tpm_limit: 70, + }, + }, + ], + }); + + renderWithProviders( + , + ); + + const user1Row = screen.getByRole("row", { name: /user1@test\.com/ }); + expect(within(user1Row).getByText("$1.25")).toBeInTheDocument(); + expect(within(user1Row).getByText("$10.00")).toBeInTheDocument(); + expect(within(user1Row).getByText("$250.00")).toBeInTheDocument(); + expect(within(user1Row).getByText("100 RPM / 10000 TPM")).toBeInTheDocument(); + + const user2Row = screen.getByRole("row", { name: /user2@test\.com/ }); + expect(within(user2Row).getByText("$7.50")).toBeInTheDocument(); + expect(within(user2Row).getByText("$20.00")).toBeInTheDocument(); + expect(within(user2Row).getByText("$2,500.00")).toBeInTheDocument(); + expect(within(user2Row).getByText("7 RPM / 70 TPM")).toBeInTheDocument(); + }); + + it("shows a member's new spend and budget once the team is refetched", () => { + const baseMembership = createMockTeamData().team_memberships[0]; + const teamWith = (spend: number, maxBudget: number) => + createMockTeamData({ + team_memberships: [ + { + ...baseMembership, + spend, + litellm_budget_table: { ...baseMembership.litellm_budget_table, max_budget: maxBudget }, + }, + ], + }); + const renderTab = (teamData: TeamData) => ( + + ); + + const { rerender } = renderWithProviders(renderTab(teamWith(1.25, 250))); + rerender(renderTab(teamWith(3.5, 400))); + + const row = screen.getByRole("row", { name: /user1@test\.com/ }); + expect(within(row).getByText("$3.50")).toBeInTheDocument(); + expect(within(row).getByText("$400.00")).toBeInTheDocument(); + expect(within(row).queryByText("$1.25")).not.toBeInTheDocument(); + }); + it("should display formatted budget and Unlimited for member with no budget", () => { renderWithProviders( , string> = { @@ -74,6 +74,14 @@ export default function TeamMemberTab({ const { mutate: resetMemberBudget, isPending: isResettingBudget } = useResetTeamMemberBudget(); const teamDefaultBudget = teamData.team_info.team_member_budget_table?.max_budget ?? null; + const membershipsByUserId = useMemo( + () => new Map(teamData.team_memberships.map((membership) => [membership.user_id, membership])), + [teamData.team_memberships], + ); + + const findMembership = (userId: string | null): TeamMembership | undefined => + userId ? membershipsByUserId.get(userId) : undefined; + const formatNumber = (value: number | null): string => { if (value === null || value === undefined) return "0"; @@ -95,32 +103,28 @@ export default function TeamMemberTab({ const getUserCurrentCycleSpend = (userId: string | null): number => { if (!userId) return 0; - const membership = teamData.team_memberships.find((tm) => tm.user_id === userId); - return membership?.spend ?? 0; + return findMembership(userId)?.spend ?? 0; }; const getUserTotalSpend = (userId: string | null): number => { if (!userId) return 0; - const membership = teamData.team_memberships.find((tm) => tm.user_id === userId); - return membership?.total_spend ?? 0; + return findMembership(userId)?.total_spend ?? 0; }; const getUserBudgetSource = (userId: string | null): TeamMemberBudgetSource => { if (!userId) return "none"; - const membership = teamData.team_memberships.find((tm) => tm.user_id === userId); - return membership?.budget_source ?? "none"; + return findMembership(userId)?.budget_source ?? "none"; }; const getUserBudget = (userId: string | null): number | null => { if (!userId) return null; - const membership = teamData.team_memberships.find((tm) => tm.user_id === userId); - return membership?.litellm_budget_table?.max_budget ?? teamDefaultBudget; + return findMembership(userId)?.litellm_budget_table?.max_budget ?? teamDefaultBudget; }; // Helper function to get rate limits for a user const getUserRateLimits = (userId: string | null): string => { if (!userId) return "No Limits"; - const membership = teamData.team_memberships.find((tm) => tm.user_id === userId); + const membership = findMembership(userId); const rpmLimit = membership?.litellm_budget_table?.rpm_limit; const tpmLimit = membership?.litellm_budget_table?.tpm_limit; @@ -139,15 +143,13 @@ export default function TeamMemberTab({ const getUserAllowedModels = (userId: string | null): string[] | null => { if (!userId) return null; - const membership = teamData.team_memberships.find((tm) => tm.user_id === userId); - const models = membership?.litellm_budget_table?.allowed_models; + const models = findMembership(userId)?.litellm_budget_table?.allowed_models; return models && models.length > 0 ? models : null; }; const getUserBudgetReset = (userId: string | null): string | null => { if (!userId) return null; - const membership = teamData.team_memberships.find((tm) => tm.user_id === userId); - return membership?.litellm_budget_table?.budget_reset_at ?? null; + return findMembership(userId)?.litellm_budget_table?.budget_reset_at ?? null; }; const extraColumns: NonNullable["extraColumns"]> = [ @@ -312,7 +314,7 @@ export default function TeamMemberTab({ members={teamData.team_info.members_with_roles} canEdit={canEditTeam} onEdit={(record) => { - const membership = teamData.team_memberships.find((tm) => tm.user_id === record.user_id); + const membership = findMembership(record.user_id); setSelectedEditMember( seedMemberBudgetFields( record, From 84eeebd509eccbe495812b30cad12f57e4d20481 Mon Sep 17 00:00:00 2001 From: Tufail Akram <72789060+akram1089@users.noreply.github.com> Date: Fri, 16 Oct 2026 14:11:34 +0530 Subject: [PATCH 11/23] fix(proxy): track and drain management audit log writes on shutdown (#43594) * fix(proxy): track and drain management audit log writes on shutdown * fix(proxy): track hook spawns and defer registration in audit log drain * fix: satisfy LIT006 on audit log task registry by documenting the required cast * fix(proxy): satisfy LIT006 and LIT010 in audit task registry * test(proxy): drive the audit drain tests with events instead of wall-clock waits * fix(proxy): drain audit writes beside the shutdown flushes and track only premium writes Run the audit drain concurrently with the gateway and request-error flushes so the two waits overlap inside the teardown budget, register the running audit task only after the premium check, drop the four duplicate task-failure loggers, and cover the shutdown order and a real hook spawn with tests * fix(proxy): start hook audit writes before their email and secret manager calls Key and user hooks now spawn their tracked audit write first, so a slow invite email or secret manager call no longer spends the drain budget before the write starts. The endpoints stop tracking the hook tasks, so shutdown never waits on hook I/O when audit logging is off. Adds regression tests for the hook order, the failed block or unblock update, and the shutdown drain without a database. * fix(proxy): track only the audit insert, not the code that awaits it The shutdown drain now waits on each audit row insert itself, so it never waits on a caller's later work and covers every caller without per-endpoint wrappers. The user/new hook still tracks its own task because it reads the user row before the insert * test(proxy): drain the audit registry in the config audit tests The audit insert now runs in its own tracked task, one hop after the caller's task, so a single yield no longer reaches it * fix(proxy): name dropped audit rows in the drain timeout warning Each tracked audit task is named after its action, table, object id and destination, so a shutdown that drops a row says which one. Adds tests that the drain waits for slow audit callbacks and that the warning names the dropped row --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../proxy/hooks/key_management_event_hooks.py | 78 +-- .../hooks/user_management_event_hooks.py | 43 +- .../key_management_endpoints.py | 20 +- .../proxy/management_helpers/audit_logs.py | 83 ++- litellm/proxy/proxy_server.py | 19 +- .../test_key_management_endpoints.py | 37 ++ .../test_audit_logs_proxy.py | 496 +++++++++++++++++- .../unit/proxy/proxy_server/test_lifecycle.py | 72 +++ ...test_proxy_server_endpoints_and_startup.py | 4 + 9 files changed, 765 insertions(+), 87 deletions(-) diff --git a/litellm/proxy/hooks/key_management_event_hooks.py b/litellm/proxy/hooks/key_management_event_hooks.py index ef93809d47b..ef60a800cdd 100644 --- a/litellm/proxy/hooks/key_management_event_hooks.py +++ b/litellm/proxy/hooks/key_management_event_hooks.py @@ -57,13 +57,6 @@ class KeyManagementEventHooks: ) from litellm.proxy.proxy_server import litellm_proxy_admin_name - # Send email notification - non-blocking, independent operation - if data.send_invite_email is True: - try: - await KeyManagementEventHooks._send_key_created_email(response.model_dump(exclude_none=True)) - except Exception as e: - verbose_proxy_logger.warning("Failed to send key created email: %s", e) - if is_audit_logging_enabled(): _updated_values: Final = response.model_dump_json(exclude_none=True) asyncio.create_task( @@ -86,6 +79,13 @@ class KeyManagementEventHooks: ) ) + # Send email notification - non-blocking, independent operation + if data.send_invite_email is True: + try: + await KeyManagementEventHooks._send_key_created_email(response.model_dump(exclude_none=True)) + except Exception as e: + verbose_proxy_logger.warning("Failed to send key created email: %s", e) + # Store the generated key in the secret manager - non-blocking, independent operation try: await KeyManagementEventHooks._store_virtual_key_in_secret_manager( @@ -118,16 +118,6 @@ class KeyManagementEventHooks: ) from litellm.proxy.proxy_server import litellm_proxy_admin_name - if data.key_alias is not None and data.key_alias != existing_key_row.key_alias: - try: - await KeyManagementEventHooks._rename_virtual_key_in_secret_manager( - current_secret_name=existing_key_row.key_alias or f"virtual-key-{existing_key_row.token}", - new_secret_name=data.key_alias, - team_id=existing_key_row.team_id, - ) - except Exception as e: - verbose_proxy_logger.warning("Failed to rename virtual key in secret manager: %s", e) - if is_audit_logging_enabled(): updated_fields: Final = { **data.model_dump(exclude_none=True), @@ -156,6 +146,16 @@ class KeyManagementEventHooks: ) asyncio.create_task(create_audit_log_for_update(request_data=request_data)) + if data.key_alias is not None and data.key_alias != existing_key_row.key_alias: + try: + await KeyManagementEventHooks._rename_virtual_key_in_secret_manager( + current_secret_name=existing_key_row.key_alias or f"virtual-key-{existing_key_row.token}", + new_secret_name=data.key_alias, + team_id=existing_key_row.team_id, + ) + except Exception as e: + verbose_proxy_logger.warning("Failed to rename virtual key in secret manager: %s", e) + @staticmethod async def async_key_rotated_hook( data: RegenerateKeyRequest | None, @@ -171,6 +171,28 @@ class KeyManagementEventHooks: ) from litellm.proxy.proxy_server import litellm_proxy_admin_name + # store the audit log + if is_audit_logging_enabled() and existing_key_row.token is not None: + asyncio.create_task( + create_audit_log_for_update( + request_data=LiteLLM_AuditLogs( + id=str(uuid.uuid4()), + updated_at=datetime.now(timezone.utc), + changed_by=get_audit_log_changed_by( + litellm_changed_by=litellm_changed_by, + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name=litellm_proxy_admin_name, + ), + changed_by_api_key=user_api_key_dict.token, + table_name=LitellmTableNames.KEY_TABLE_NAME, + object_id=existing_key_row.token, + action="rotated", + updated_values=response.model_dump_json(exclude_none=True), + before_value=existing_key_row.model_dump_json(exclude_none=True), + ) + ) + ) + # Store the generated key in the secret manager - non-blocking, independent operation if response.token_id is not None: try: @@ -204,28 +226,6 @@ class KeyManagementEventHooks: except Exception as e: verbose_proxy_logger.warning("Failed to send key rotated email: %s", e) - # store the audit log - if is_audit_logging_enabled() and existing_key_row.token is not None: - asyncio.create_task( - create_audit_log_for_update( - request_data=LiteLLM_AuditLogs( - id=str(uuid.uuid4()), - updated_at=datetime.now(timezone.utc), - changed_by=get_audit_log_changed_by( - litellm_changed_by=litellm_changed_by, - user_api_key_dict=user_api_key_dict, - litellm_proxy_admin_name=litellm_proxy_admin_name, - ), - changed_by_api_key=user_api_key_dict.token, - table_name=LitellmTableNames.KEY_TABLE_NAME, - object_id=existing_key_row.token, - action="rotated", - updated_values=response.model_dump_json(exclude_none=True), - before_value=existing_key_row.model_dump_json(exclude_none=True), - ) - ) - ) - @staticmethod async def async_key_deleted_hook( data: KeyRequest, diff --git a/litellm/proxy/hooks/user_management_event_hooks.py b/litellm/proxy/hooks/user_management_event_hooks.py index 6d978929c05..6618af70c0a 100644 --- a/litellm/proxy/hooks/user_management_event_hooks.py +++ b/litellm/proxy/hooks/user_management_event_hooks.py @@ -23,6 +23,7 @@ from litellm.proxy._types import ( from litellm.proxy.management_helpers.audit_logs import ( create_audit_log_for_update, is_audit_logging_enabled, + track_audit_task, ) from litellm.repositories.user_repository import UserRepository @@ -41,7 +42,18 @@ class UserManagementEventHooks: - Creating an audit log for the user creation - Sending a user invitation email to the user """ - from litellm.proxy.proxy_server import litellm_proxy_admin_name, prisma_client + from litellm.proxy.proxy_server import litellm_proxy_admin_name + + if is_audit_logging_enabled(): + track_audit_task( + asyncio.create_task( + UserManagementEventHooks._create_user_created_audit_log( + response=response, + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name=litellm_proxy_admin_name, + ) + ) + ) ######################################################### ########## Send User Invitation Email ################ @@ -52,9 +64,14 @@ class UserManagementEventHooks: user_api_key_dict=user_api_key_dict, ) - ######################################################### - ########## CREATE AUDIT LOG ################ - ######################################################### + @staticmethod + async def _create_user_created_audit_log( + response: NewUserResponse, + user_api_key_dict: UserAPIKeyAuth, + litellm_proxy_admin_name: str | None, + ) -> None: + from litellm.proxy.proxy_server import prisma_client + try: if prisma_client is None: raise Exception(CommonProxyErrors.db_not_connected_error.value) @@ -63,16 +80,14 @@ class UserManagementEventHooks: user_row: Final = await UserRepository(prisma_client).find_by_id(response.user_id) if user_row is None: raise Exception(f"no user row found for user_id={response.user_id}") - asyncio.create_task( - UserManagementEventHooks.create_internal_user_audit_log( - user_id=user_row.user_id, - action="created", - litellm_changed_by=user_api_key_dict.user_id, - user_api_key_dict=user_api_key_dict, - litellm_proxy_admin_name=litellm_proxy_admin_name, - before_value=None, - after_value=user_row.model_dump_json(exclude_none=True), - ) + await UserManagementEventHooks.create_internal_user_audit_log( + user_id=user_row.user_id, + action="created", + litellm_changed_by=user_api_key_dict.user_id, + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name=litellm_proxy_admin_name, + before_value=None, + after_value=user_row.model_dump_json(exclude_none=True), ) except Exception as e: verbose_proxy_logger.warning("Unable to create audit log for user on `/user/new` - %s", e) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 7877d887409..592c732bfd3 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -7713,6 +7713,11 @@ async def block_key( code=status.HTTP_404_NOT_FOUND, ) + record: Final = await _prisma_table(VerificationTokenRepository(prisma_client)).update( + where={"token": hashed_token}, + data=with_settings_updated_at({"blocked": True}), + ) + if is_audit_logging_enabled(): asyncio.create_task( create_audit_log_for_update( @@ -7734,11 +7739,6 @@ async def block_key( ) ) - record: Final = await _prisma_table(VerificationTokenRepository(prisma_client)).update( - where={"token": hashed_token}, - data=with_settings_updated_at({"blocked": True}), - ) - ## UPDATE KEY CACHE - invalidate so next read re-fetches from DB await delete_cache_key_object( hashed_token=hashed_token, @@ -7827,6 +7827,11 @@ async def unblock_key( code=status.HTTP_404_NOT_FOUND, ) + record: Final = await _prisma_table(VerificationTokenRepository(prisma_client)).update( + where={"token": hashed_token}, + data=with_settings_updated_at({"blocked": False}), + ) + if is_audit_logging_enabled(): asyncio.create_task( create_audit_log_for_update( @@ -7848,11 +7853,6 @@ async def unblock_key( ) ) - record: Final = await _prisma_table(VerificationTokenRepository(prisma_client)).update( - where={"token": hashed_token}, - data=with_settings_updated_at({"blocked": False}), - ) - ## UPDATE KEY CACHE - invalidate so next read re-fetches from DB await delete_cache_key_object( hashed_token=hashed_token, diff --git a/litellm/proxy/management_helpers/audit_logs.py b/litellm/proxy/management_helpers/audit_logs.py index ecd6abea3c3..ef95e8e2c0f 100644 --- a/litellm/proxy/management_helpers/audit_logs.py +++ b/litellm/proxy/management_helpers/audit_logs.py @@ -5,7 +5,7 @@ Functions to create audit logs for LiteLLM Proxy import asyncio import json from datetime import datetime, timezone -from typing import Final +from typing import Final, TypeVar import litellm from litellm._logging import verbose_proxy_logger @@ -134,6 +134,56 @@ def _audit_log_task_done_callback(task: asyncio.Task) -> None: verbose_proxy_logger.error("Audit log callback task failed: %s", exc, exc_info=exc) +_T: Final = TypeVar("_T") + + +class AuditTaskRegistry: + __slots__ = ("_pending_tasks",) + + def __init__(self) -> None: + self._pending_tasks: Final[set[asyncio.Task[object]]] = set() # mutable-ok: background task tracking + + def track(self, task: asyncio.Task[_T]) -> asyncio.Task[_T]: + if task.done() or task in self._pending_tasks: + return task + self._pending_tasks.add(task) + task.add_done_callback(self._on_task_done) + return task + + def _on_task_done(self, task: asyncio.Task[object]) -> None: + self._pending_tasks.discard(task) + _audit_log_task_done_callback(task) + + async def drain(self, timeout: float = 5.0) -> None: + loop: Final = asyncio.get_running_loop() + deadline: Final = loop.time() + timeout + while self._pending_tasks: + _, pending = await asyncio.wait(tuple(self._pending_tasks), timeout=max(deadline - loop.time(), 0)) + if pending: + self._pending_tasks.difference_update(pending) + verbose_proxy_logger.warning( + "Timed out draining %d audit tasks on shutdown, dropping: %s", + len(pending), + ", ".join(sorted(task.get_name() for task in pending)), + ) + return + + +_audit_task_registry: Final = AuditTaskRegistry() + + +def track_audit_task(task: asyncio.Task[_T]) -> asyncio.Task[_T]: + return _audit_task_registry.track(task) + + +async def drain_audit_tasks(timeout: float = 5.0) -> None: + await _audit_task_registry.drain(timeout=timeout) + + +def _audit_task_name(request_data: LiteLLM_AuditLogs, destination: str) -> str: + return f"{request_data.action} {request_data.table_name.value} {request_data.object_id} to {destination}" + + async def _dispatch_audit_log_to_callbacks( request_data: LiteLLM_AuditLogs, ) -> None: @@ -153,8 +203,12 @@ async def _dispatch_audit_log_to_callbacks( continue if isinstance(resolved, CustomLogger): - task = asyncio.create_task(resolved.async_log_audit_log_event(payload)) - task.add_done_callback(_audit_log_task_done_callback) + track_audit_task( + asyncio.create_task( + resolved.async_log_audit_log_event(payload), + name=_audit_task_name(request_data, type(resolved).__name__), + ) + ) except Exception as e: verbose_proxy_logger.error("Failed dispatching audit log to callback: %s", e) @@ -204,6 +258,14 @@ async def create_object_audit_log( ) +async def _insert_audit_log_row(repository: AuditLogRepository, request_data: LiteLLM_AuditLogs) -> None: + try: + await repository.table.create(data={**request_data.model_dump(exclude_none=True)}) + except Exception as e: + # [Non-Blocking Exception. Do not allow blocking LLM API call] + verbose_proxy_logger.error("Failed Creating audit log %s", e) + + async def create_audit_log_for_update(request_data: LiteLLM_AuditLogs): """ Create an audit log for an object. @@ -231,14 +293,9 @@ async def create_audit_log_for_update(request_data: LiteLLM_AuditLogs): verbose_proxy_logger.error("prisma_client is None, cannot write audit log to DB") return - _request_data: Final = request_data.model_dump(exclude_none=True) - - try: - await AuditLogRepository(prisma_client).table.create( - data={ - **_request_data, - } + await track_audit_task( + asyncio.create_task( + _insert_audit_log_row(AuditLogRepository(prisma_client), request_data), + name=_audit_task_name(request_data, "the database"), ) - except Exception as e: - # [Non-Blocking Exception. Do not allow blocking LLM API call] - verbose_proxy_logger.error("Failed Creating audit log %s", e) + ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 97d8eda1fad..bbd9d96b84b 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1200,6 +1200,20 @@ async def _flush_spend_logs_queue_on_shutdown() -> None: verbose_proxy_logger.exception("Error flushing spend logs queue on shutdown: %s", e) +async def _drain_audit_tasks_on_shutdown() -> None: + try: + from litellm.proxy.management_helpers.audit_logs import drain_audit_tasks + + await drain_audit_tasks() + except Exception as e: # noqa: BLE001 # shutdown must continue even if the drain fails + verbose_proxy_logger.exception("Error draining audit tasks on shutdown: %s", e) + + +async def _flush_request_accumulators(prisma: PrismaClient) -> None: + await flush_gateway_requests(prisma, gateway_request_accumulator) + await flush_request_errors(prisma, request_error_accumulator) + + async def proxy_shutdown_event(worker_heartbeat: ProxyWorkerHeartbeat | None = None) -> None: global prisma_client, master_key, user_custom_auth, user_custom_key_generate, user_custom_key_update verbose_proxy_logger.info("Shutting down LiteLLM Proxy Server") @@ -1210,10 +1224,11 @@ async def proxy_shutdown_event(worker_heartbeat: ProxyWorkerHeartbeat | None = N # is lost, and a write attempted after disconnect raises # ClientNotConnectedError rather than persisting anything. Ordering this # inside the same guard is what keeps the two from drifting apart. - await flush_gateway_requests(prisma_client, gateway_request_accumulator) - await flush_request_errors(prisma_client, request_error_accumulator) + await asyncio.gather(_drain_audit_tasks_on_shutdown(), _flush_request_accumulators(prisma_client)) verbose_proxy_logger.debug("Disconnecting from Prisma") await prisma_client.disconnect() + else: + await _drain_audit_tasks_on_shutdown() await drain_post_call_redis_batches() if litellm.cache is not None: diff --git a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py index 1af9def3afb..57e0a4fc47c 100644 --- a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py @@ -3407,6 +3407,43 @@ async def test_block_key_existing_key_succeeds(monkeypatch): assert result == mock_updated_record +@pytest.mark.asyncio +@pytest.mark.parametrize("endpoint_name", ["block_key", "unblock_key"]) +async def test_block_and_unblock_write_no_audit_row_when_the_key_update_fails(monkeypatch, endpoint_name): + from litellm.proxy._types import BlockKeyRequest + from litellm.proxy.management_endpoints import key_management_endpoints + from litellm.proxy.management_helpers.audit_logs import drain_audit_tasks + + hashed_token = "a1b2c3d4e5f6789012345678901234567890123456789012345678901234abcd" + existing_key_record = MagicMock() + existing_key_record.token = hashed_token + existing_key_record.model_dump_json.return_value = f'{{"token": "{hashed_token}"}}' + + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=existing_key_record) + mock_prisma_client.db.litellm_verificationtoken.update = AsyncMock(side_effect=RuntimeError("update failed")) + audit_writer = AsyncMock() + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()) + monkeypatch.setattr("litellm.proxy.proxy_server.create_audit_log_for_update", audit_writer) + monkeypatch.setattr("litellm.store_audit_logs", True) + + with pytest.raises(RuntimeError, match="update failed"): + await getattr(key_management_endpoints, endpoint_name)( + data=BlockKeyRequest(key=hashed_token), + http_request=MagicMock(), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin_user" + ), + litellm_changed_by=None, + ) + await drain_audit_tasks(timeout=1.0) + + assert audit_writer.call_count == 0 + + @pytest.mark.asyncio async def test_validate_key_team_change_with_member_permissions(): """ diff --git a/tests/unit/proxy/management_helpers/test_audit_logs_proxy.py b/tests/unit/proxy/management_helpers/test_audit_logs_proxy.py index 1ea608cfc17..325275ab391 100644 --- a/tests/unit/proxy/management_helpers/test_audit_logs_proxy.py +++ b/tests/unit/proxy/management_helpers/test_audit_logs_proxy.py @@ -1,14 +1,18 @@ import os import traceback from litellm._uuid import uuid -from datetime import datetime +from datetime import datetime, timezone +from collections.abc import Coroutine +from typing import Any, Final from dotenv import load_dotenv from fastapi import Request from fastapi.routing import APIRoute +import contextlib import io +import logging import time # this file is to test litellm/proxy @@ -38,11 +42,16 @@ from starlette.datastructures import URL from litellm.proxy.management_helpers.audit_logs import ( create_audit_log_for_update, + drain_audit_tasks, get_audit_log_changed_by, + track_audit_task, ) from litellm.proxy._types import LiteLLM_AuditLogs, LitellmTableNames, UserAPIKeyAuth +from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache -from unittest.mock import patch, AsyncMock +from litellm.integrations.custom_logger import CustomLogger +from litellm.types.utils import StandardAuditLogPayload +from unittest.mock import AsyncMock, MagicMock, patch proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) import json @@ -174,7 +183,6 @@ async def test_create_audit_log_for_update_premium_user(): patch("litellm.store_audit_logs", True), patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, ): - mock_prisma.db.litellm_auditlog.create = AsyncMock() request_data = LiteLLM_AuditLogs( @@ -215,9 +223,7 @@ def prisma_client(): os.environ["DATABASE_URL"] = modified_url # Assuming PrismaClient is a class that needs to be instantiated - prisma_client = PrismaClient( - database_url=os.environ["DATABASE_URL"], proxy_logging_obj=proxy_logging_obj - ) + prisma_client = PrismaClient(database_url=os.environ["DATABASE_URL"], proxy_logging_obj=proxy_logging_obj) return prisma_client @@ -252,10 +258,482 @@ async def test_create_audit_log_in_db(prisma_client): await asyncio.sleep(1) # now read the last log from the db - last_log = await prisma_client.db.litellm_auditlog.find_first( - where={"id": audit_log_id} - ) + last_log = await prisma_client.db.litellm_auditlog.find_first(where={"id": audit_log_id}) assert last_log.id == audit_log_id setattr(litellm, "store_audit_logs", False) + + +@pytest.mark.asyncio +async def test_track_audit_task_lifecycle(): + task_completed: Final = asyncio.Event() + + async def _sample_coroutine() -> None: + await asyncio.sleep(0) + task_completed.set() + + task: Final = track_audit_task(asyncio.create_task(_sample_coroutine())) + assert not task.done() + + await drain_audit_tasks(timeout=1.0) + assert task_completed.is_set() + assert task.done() + + completed_task: Final = asyncio.create_task(asyncio.sleep(0)) + await completed_task + assert track_audit_task(completed_task) is completed_task + + +@pytest.mark.asyncio +async def test_drain_audit_tasks_waits_for_all_tasks(): + event1: Final = asyncio.Event() + event2: Final = asyncio.Event() + + async def _worker(event: asyncio.Event) -> None: + await asyncio.sleep(0) + event.set() + + task1: Final = track_audit_task(asyncio.create_task(_worker(event1))) + task2: Final = track_audit_task(asyncio.create_task(_worker(event2))) + + assert not task1.done() + assert not task2.done() + + await drain_audit_tasks(timeout=1.0) + + assert event1.is_set() + assert event2.is_set() + assert task1.done() + assert task2.done() + + +@pytest.mark.asyncio +async def test_drain_audit_tasks_handles_failing_task_cleanly(): + async def _failing_worker() -> None: + await asyncio.sleep(0) + raise RuntimeError("test audit log write failure") + + task: Final = track_audit_task(asyncio.create_task(_failing_worker())) + assert not task.done() + + await drain_audit_tasks(timeout=1.0) + assert task.done() + assert isinstance(task.exception(), RuntimeError) + + +@pytest.mark.asyncio +async def test_drain_audit_tasks_timeout_does_not_cancel_pending_tasks(): + async def _long_running_worker() -> None: + await asyncio.Event().wait() + + task: Final = track_audit_task(asyncio.create_task(_long_running_worker())) + try: + assert not task.done() + await drain_audit_tasks(timeout=0) + assert not task.cancelled() + assert not task.done() + finally: + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + + +@pytest.mark.asyncio +async def test_drain_audit_tasks_persists_management_audit_write_before_db_disconnect(): + mock_prisma: Final = MagicMock() + writes: Final[list[dict[str, Any]]] = [] + is_disconnected: Final = [False] + + async def _mock_create(data: dict[str, Any]) -> None: + if is_disconnected[0]: + raise RuntimeError("Database already disconnected") + await asyncio.sleep(0) + writes.append(data) + + mock_prisma.db.litellm_auditlog.create = AsyncMock(side_effect=_mock_create) + + with ( + patch("litellm.store_audit_logs", True), + patch("litellm.proxy.proxy_server.premium_user", True), + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + ): + request_data: Final = LiteLLM_AuditLogs( + id=str(uuid.uuid4()), + updated_at=datetime.now(timezone.utc), + changed_by="admin-user", + changed_by_api_key="sk-admin", + table_name=LitellmTableNames.KEY_TABLE_NAME, + object_id="token-123", + action="blocked", + updated_values="{}", + before_value='{"key": "test"}', + ) + + task: Final = track_audit_task(asyncio.create_task(create_audit_log_for_update(request_data=request_data))) + assert not task.done() + + await drain_audit_tasks(timeout=2.0) + is_disconnected[0] = True + + assert task.done() + assert len(writes) == 1 + assert writes[0]["object_id"] == "token-123" + assert writes[0]["action"] == "blocked" + + +@pytest.mark.asyncio +async def test_create_audit_log_auto_registers_in_drain(): + mock_prisma: Final = MagicMock() + writes: Final[list[dict[str, Any]]] = [] + + async def _mock_create(data: dict[str, Any]) -> None: + await asyncio.sleep(0) + writes.append(data) + + mock_prisma.db.litellm_auditlog.create = AsyncMock(side_effect=_mock_create) + + with ( + patch("litellm.store_audit_logs", True), + patch("litellm.proxy.proxy_server.premium_user", True), + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + ): + request_data: Final = LiteLLM_AuditLogs( + id=str(uuid.uuid4()), + updated_at=datetime.now(timezone.utc), + changed_by="admin-user", + changed_by_api_key="sk-admin", + table_name=LitellmTableNames.KEY_TABLE_NAME, + object_id="token-456", + action="deleted", + updated_values="{}", + before_value='{"key": "test-deleted"}', + ) + + task: Final = asyncio.create_task(create_audit_log_for_update(request_data=request_data)) + await asyncio.sleep(0) + + await drain_audit_tasks(timeout=2.0) + + assert task.done() + assert len(writes) == 1 + assert writes[0]["object_id"] == "token-456" + assert writes[0]["action"] == "deleted" + + +@pytest.mark.asyncio +async def test_drain_audit_tasks_discards_timed_out_tasks_on_timeout(): + async def _hung_worker() -> None: + await asyncio.Event().wait() + + hung_task: Final = track_audit_task(asyncio.create_task(_hung_worker())) + try: + await drain_audit_tasks(timeout=0) + assert not hung_task.done() + + second_drain: Final = asyncio.create_task(drain_audit_tasks(timeout=60)) + for _ in range(3): + await asyncio.sleep(0) + assert second_drain.done() + finally: + hung_task.cancel() + try: + await hung_task + except asyncio.CancelledError: + pass + + +@pytest.mark.asyncio +async def test_drain_audit_tasks_captures_tasks_queued_during_drain(): + finished_tasks: Final[list[str]] = [] + + async def _first_worker() -> None: + await asyncio.sleep(0) + track_audit_task(asyncio.create_task(_second_worker())) + finished_tasks.append("first") + + async def _second_worker() -> None: + await asyncio.sleep(0) + finished_tasks.append("second") + + track_audit_task(asyncio.create_task(_first_worker())) + await drain_audit_tasks(timeout=2.0) + + assert finished_tasks == ["first", "second"] + + +class _AuditCallbackBlockedUntilReleased(CustomLogger): + def __init__(self, release: asyncio.Event) -> None: + super().__init__() + self.release: Final = release + self.logged_object_ids: Final[list[str]] = [] + + async def async_log_audit_log_event(self, audit_log: StandardAuditLogPayload) -> None: + await self.release.wait() + self.logged_object_ids.append(audit_log["object_id"]) + + +class _WarningRecorder(logging.Handler): + def __init__(self) -> None: + super().__init__(level=logging.WARNING) + self.messages: Final[list[str]] = [] + + def emit(self, record: logging.LogRecord) -> None: + self.messages.append(record.getMessage()) + + +def _deleted_key_audit_log(object_id: str) -> LiteLLM_AuditLogs: + return LiteLLM_AuditLogs( + id=str(uuid.uuid4()), + updated_at=datetime.now(timezone.utc), + changed_by="admin-user", + table_name=LitellmTableNames.KEY_TABLE_NAME, + object_id=object_id, + action="deleted", + updated_values="{}", + before_value="{}", + ) + + +@pytest.mark.asyncio +async def test_drain_waits_for_a_slow_audit_log_callback(): + release: Final = asyncio.Event() + callback: Final = _AuditCallbackBlockedUntilReleased(release) + + with ( + patch("litellm.store_audit_logs", True), + patch("litellm.proxy.proxy_server.premium_user", True), + patch("litellm.proxy.proxy_server.prisma_client", None), + patch("litellm.audit_log_callbacks", [callback]), + ): + await create_audit_log_for_update(request_data=_deleted_key_audit_log("token-callback")) + drain: Final = asyncio.create_task(drain_audit_tasks(timeout=60)) + for _ in range(3): + await asyncio.sleep(0) + drain_returned_before_the_callback_finished: Final = drain.done() + release.set() + await drain + + assert not drain_returned_before_the_callback_finished + assert callback.logged_object_ids == ["token-callback"] + + +@pytest.mark.asyncio +async def test_drain_timeout_warning_names_each_dropped_audit_row(): + mock_prisma: Final = MagicMock() + mock_prisma.db.litellm_auditlog.create = AsyncMock(side_effect=_io_that_never_finishes) + recorder: Final = _WarningRecorder() + verbose_proxy_logger.addHandler(recorder) + + with ( + patch("litellm.store_audit_logs", True), + patch("litellm.proxy.proxy_server.premium_user", True), + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + ): + writer: Final = asyncio.create_task( + create_audit_log_for_update(request_data=_deleted_key_audit_log("token-dropped")) + ) + try: + for _ in range(3): + await asyncio.sleep(0) + await drain_audit_tasks(timeout=0.05) + finally: + verbose_proxy_logger.removeHandler(recorder) + writer.cancel() + with contextlib.suppress(asyncio.CancelledError): + await writer + + assert recorder.messages == [ + "Timed out draining 1 audit tasks on shutdown, dropping: " + f"deleted {LitellmTableNames.KEY_TABLE_NAME.value} token-dropped to the database" + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("store_audit_logs", "premium_user", "expected_write_count"), + [(True, True, 1), (False, True, 0), (True, False, 0)], + ids=["row_written", "logging_disabled", "not_premium"], +) +async def test_drain_does_not_wait_for_code_that_runs_after_an_inline_audit_write( + store_audit_logs: bool, premium_user: bool, expected_write_count: int +): + release_caller: Final = asyncio.Event() + mock_prisma: Final = MagicMock() + writes: Final[list[dict[str, Any]]] = [] + mock_prisma.db.litellm_auditlog.create = AsyncMock(side_effect=lambda data: writes.append(data)) + + async def _caller_that_outlives_the_audit_call() -> None: + await create_audit_log_for_update( + request_data=LiteLLM_AuditLogs( + id=str(uuid.uuid4()), + updated_at=datetime.now(timezone.utc), + changed_by="test-user", + table_name=LitellmTableNames.KEY_TABLE_NAME, + object_id="test-obj", + action="updated", + updated_values="{}", + before_value="{}", + ) + ) + await release_caller.wait() + + with ( + patch("litellm.store_audit_logs", store_audit_logs), + patch("litellm.proxy.proxy_server.premium_user", premium_user), + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + ): + caller: Final = asyncio.create_task(_caller_that_outlives_the_audit_call()) + for _ in range(3): + await asyncio.sleep(0) + drain: Final = asyncio.create_task(drain_audit_tasks(timeout=60)) + for _ in range(3): + await asyncio.sleep(0) + drain_returned_while_caller_was_running: Final = drain.done() + release_caller.set() + await caller + await drain + + assert drain_returned_while_caller_was_running + assert len(writes) == expected_write_count + + +async def _io_that_never_finishes(*args: object, **kwargs: object) -> None: + await asyncio.Event().wait() + + +async def _drain_while_hook_io_hangs(hook: Coroutine[object, object, None]) -> None: + hook_task: Final = asyncio.create_task(hook) + for _ in range(3): + await asyncio.sleep(0) + await drain_audit_tasks(timeout=2.0) + hook_task.cancel() + + +@pytest.mark.asyncio +async def test_key_generated_audit_write_is_drained_while_the_invite_email_hangs(): + from litellm.proxy._types import GenerateKeyRequest, GenerateKeyResponse + from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks + + mock_prisma: Final = MagicMock() + writes: Final[list[dict[str, Any]]] = [] + mock_prisma.db.litellm_auditlog.create = AsyncMock(side_effect=lambda data: writes.append(data)) + + with ( + patch("litellm.store_audit_logs", True), + patch("litellm.proxy.proxy_server.premium_user", True), + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + patch.object(KeyManagementEventHooks, "_send_key_created_email", side_effect=_io_that_never_finishes), + ): + await _drain_while_hook_io_hangs( + KeyManagementEventHooks.async_key_generated_hook( + data=GenerateKeyRequest(send_invite_email=True), + response=GenerateKeyResponse(key="sk-new", token_id="hashed-new"), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-admin", user_id="admin"), + ) + ) + + assert [(write["object_id"], write["action"]) for write in writes] == [("hashed-new", "created")] + + +@pytest.mark.asyncio +async def test_key_updated_audit_write_is_drained_while_the_secret_manager_rename_hangs(): + from litellm.proxy._types import LiteLLM_VerificationToken, UpdateKeyRequest + from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks + + mock_prisma: Final = MagicMock() + writes: Final[list[dict[str, Any]]] = [] + mock_prisma.db.litellm_auditlog.create = AsyncMock(side_effect=lambda data: writes.append(data)) + + with ( + patch("litellm.store_audit_logs", True), + patch("litellm.proxy.proxy_server.premium_user", True), + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + patch.object( + KeyManagementEventHooks, "_rename_virtual_key_in_secret_manager", side_effect=_io_that_never_finishes + ), + ): + await _drain_while_hook_io_hangs( + KeyManagementEventHooks.async_key_updated_hook( + data=UpdateKeyRequest(key="sk-existing", key_alias="new-alias"), + existing_key_row=LiteLLM_VerificationToken(token="hashed-existing", key_alias="old-alias"), + response=None, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-admin", user_id="admin"), + ) + ) + + assert [(write["object_id"], write["action"]) for write in writes] == [(hash_token("sk-existing"), "updated")] + + +@pytest.mark.asyncio +async def test_key_rotated_audit_write_is_drained_while_the_secret_manager_rotation_hangs(): + from litellm.proxy._types import GenerateKeyResponse, LiteLLM_VerificationToken + from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks + + mock_prisma: Final = MagicMock() + writes: Final[list[dict[str, Any]]] = [] + mock_prisma.db.litellm_auditlog.create = AsyncMock(side_effect=lambda data: writes.append(data)) + + with ( + patch("litellm.store_audit_logs", True), + patch("litellm.proxy.proxy_server.premium_user", True), + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + patch.object( + KeyManagementEventHooks, "_rotate_virtual_key_in_secret_manager", side_effect=_io_that_never_finishes + ), + ): + await _drain_while_hook_io_hangs( + KeyManagementEventHooks.async_key_rotated_hook( + data=None, + existing_key_row=LiteLLM_VerificationToken(token="hashed-old"), + response=GenerateKeyResponse(key="sk-rotated", token_id="hashed-new"), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-admin", user_id="admin"), + ) + ) + + assert [(write["object_id"], write["action"]) for write in writes] == [("hashed-old", "rotated")] + + +@pytest.mark.asyncio +async def test_user_created_audit_write_is_drained_while_its_user_read_is_still_running(): + from litellm.proxy._types import NewUserRequest, NewUserResponse + from litellm.proxy.hooks.user_management_event_hooks import UserManagementEventHooks + + mock_prisma: Final = MagicMock() + writes: Final[list[dict[str, Any]]] = [] + release_user_read: Final = asyncio.Event() + + async def _slow_user_read(*args: object, **kwargs: object) -> dict[str, str]: + await release_user_read.wait() + return {"user_id": "new-user"} + + mock_prisma.db.litellm_auditlog.create = AsyncMock(side_effect=lambda data: writes.append(data)) + mock_prisma.db.litellm_usertable.find_unique = AsyncMock(side_effect=_slow_user_read) + + with ( + patch("litellm.store_audit_logs", True), + patch("litellm.proxy.proxy_server.premium_user", True), + patch("litellm.proxy.proxy_server.prisma_client", mock_prisma), + patch.object(UserManagementEventHooks, "async_send_user_invitation_email", side_effect=_io_that_never_finishes), + ): + hook: Final = asyncio.create_task( + UserManagementEventHooks.async_user_created_hook( + data=NewUserRequest(user_id="new-user", send_invite_email=True), + response=NewUserResponse(key="sk-new", user_id="new-user"), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-admin", user_id="admin"), + ) + ) + for _ in range(3): + await asyncio.sleep(0) + drain: Final = asyncio.create_task(drain_audit_tasks(timeout=60)) + for _ in range(3): + await asyncio.sleep(0) + drain_waited_for_the_user_read: Final = not drain.done() + release_user_read.set() + await drain + hook.cancel() + + assert drain_waited_for_the_user_read + assert [(write["object_id"], write["action"]) for write in writes] == [("new-user", "created")] diff --git a/tests/unit/proxy/proxy_server/test_lifecycle.py b/tests/unit/proxy/proxy_server/test_lifecycle.py index 69f9c72c9c8..71daa3bb873 100644 --- a/tests/unit/proxy/proxy_server/test_lifecycle.py +++ b/tests/unit/proxy/proxy_server/test_lifecycle.py @@ -17,12 +17,14 @@ Pins covered: from __future__ import annotations +import asyncio import inspect import json import logging import os import subprocess from collections.abc import Awaitable, Callable +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -48,6 +50,8 @@ from litellm.proxy.proxy_server import ( save_worker_config, ) +from litellm.proxy.management_helpers.audit_logs import track_audit_task + from .conftest import normalize # --------------------------------------------------------------------------- @@ -223,6 +227,46 @@ async def test_proxy_shutdown_drains_gateway_requests_before_disconnecting(monke assert calls == ["flush", "disconnect"] +@pytest.mark.asyncio +async def test_proxy_shutdown_drains_audit_writes_beside_the_gateway_flush_before_disconnecting(monkeypatch): + """ + A management call returns before its audit row is written, so shutdown waits + for that write while prisma is still connected. It waits beside the gateway + flush rather than before it, so the two waits overlap inside the teardown budget. + """ + calls: list[str] = [] # mutable-ok: records call order, which is the assertion + gateway_flush_started: Final = asyncio.Event() + + fake_prisma = MagicMock() + fake_prisma.disconnect = AsyncMock(side_effect=lambda: calls.append("disconnect")) + monkeypatch.setattr(ps, "prisma_client", fake_prisma, raising=False) + + async def _record_flush(client, accumulator): + gateway_flush_started.set() + calls.append("flush") + + async def _audit_write_that_finishes_after_the_flush_starts() -> None: + await gateway_flush_started.wait() + calls.append("audit") + + monkeypatch.setattr(ps, "flush_gateway_requests", _record_flush, raising=False) + + fake_jwt = MagicMock() + fake_jwt.close = AsyncMock() + monkeypatch.setattr(ps, "jwt_handler", fake_jwt, raising=False) + monkeypatch.setattr(ps, "db_writer_client", None, raising=False) + + import litellm + + monkeypatch.setattr(litellm, "cache", None, raising=False) + monkeypatch.setattr(litellm, "success_callback", [], raising=False) + + track_audit_task(asyncio.create_task(_audit_write_that_finishes_after_the_flush_starts())) + await asyncio.wait_for(proxy_shutdown_event(), timeout=2.0) + + assert calls == ["flush", "audit", "disconnect"] + + @pytest.mark.asyncio async def test_proxy_shutdown_skips_gateway_flush_without_a_database(monkeypatch): """No prisma client means nothing to drain to, and no attempt is made.""" @@ -245,6 +289,34 @@ async def test_proxy_shutdown_skips_gateway_flush_without_a_database(monkeypatch assert flush.await_count == 0 +@pytest.mark.asyncio +async def test_proxy_shutdown_without_a_database_drains_audit_callbacks_before_the_rest_of_teardown(monkeypatch): + calls: list[str] = [] # mutable-ok: records call order, which is the assertion + monkeypatch.setattr(ps, "prisma_client", None, raising=False) + monkeypatch.setattr( + ps, "drain_post_call_redis_batches", AsyncMock(side_effect=lambda: calls.append("redis")), raising=False + ) + + fake_jwt = MagicMock() + fake_jwt.close = AsyncMock() + monkeypatch.setattr(ps, "jwt_handler", fake_jwt, raising=False) + monkeypatch.setattr(ps, "db_writer_client", None, raising=False) + + import litellm + + monkeypatch.setattr(litellm, "cache", None, raising=False) + monkeypatch.setattr(litellm, "success_callback", [], raising=False) + + async def _audit_callback_still_sending() -> None: + await asyncio.sleep(0.05) + calls.append("audit") + + track_audit_task(asyncio.create_task(_audit_callback_still_sending())) + await asyncio.wait_for(proxy_shutdown_event(), timeout=2.0) + + assert calls == ["audit", "redis"] + + @pytest.mark.asyncio async def test_proxy_shutdown_event_prisma_disconnect_raises_error(monkeypatch): fake_prisma = MagicMock() diff --git a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py index 0ecebc7560c..e824c5017c6 100644 --- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py +++ b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py @@ -12612,6 +12612,7 @@ def test_dump_redacted_config_serializes_non_json_native_values(): async def test_update_config_general_settings_emits_audit_log(monkeypatch): import litellm.proxy.proxy_server as proxy_server_module from litellm.proxy._types import ConfigFieldUpdate + from litellm.proxy.management_helpers.audit_logs import drain_audit_tasks from litellm.proxy.proxy_server import update_config_general_settings existing = {"max_parallel_requests": 5, "some_api_key": "sk-stored-secret"} @@ -12635,6 +12636,7 @@ async def test_update_config_general_settings_emits_audit_log(monkeypatch): ) # Audit is scheduled via asyncio.create_task; yield so it runs. await asyncio.sleep(0) + await drain_audit_tasks() fake.db.litellm_auditlog.create.assert_awaited_once() written = fake.db.litellm_auditlog.create.call_args.kwargs["data"] @@ -13004,6 +13006,7 @@ async def test_update_config_general_settings_applies_ssrf_globals(monkeypatch): async def test_delete_config_general_settings_emits_deleted_audit_log(monkeypatch): import litellm.proxy.proxy_server as proxy_server_module from litellm.proxy._types import ConfigFieldDelete + from litellm.proxy.management_helpers.audit_logs import drain_audit_tasks from litellm.proxy.proxy_server import delete_config_general_settings existing = {"max_parallel_requests": 5} @@ -13023,6 +13026,7 @@ async def test_delete_config_general_settings_emits_deleted_audit_log(monkeypatc ) # Audit is scheduled via asyncio.create_task; yield so it runs. await asyncio.sleep(0) + await drain_audit_tasks() fake.db.litellm_auditlog.create.assert_awaited_once() written = fake.db.litellm_auditlog.create.call_args.kwargs["data"] From a844dadcd27ececec63eca6c8bf3a5b819353f55 Mon Sep 17 00:00:00 2001 From: Aryan Pardeshi Date: Fri, 16 Oct 2026 14:31:33 +0530 Subject: [PATCH 12/23] feat(gemini): allow disabling thought summaries while keeping thinking on (#36460) * feat(gemini): allow disabling thought summaries while keeping thinking on Both reasoning-effort mappers hardcoded includeThoughts to True whenever thinking was on, so thinking ON with summaries OFF was unreachable even though that is the google SDK's own default. Summaries are the narration rather than the reasoning, and for callers that discard them they are pure added latency. Adds litellm.gemini_include_thoughts_default, defaulting to True so nothing changes for anyone who does not set it. It gates the includeThoughts=True branches in the thinkingBudget mapper, the thinkingLevel mapper, and the anthropic-style thinking param mapper. The disable and none paths still force includeThoughts=False: with thinking off there is nothing to summarise. * test(gemini): assert thought-summary flag through the outbound Gemini request * feat(gemini): let a request's thought summary ask override the gateway default thinking.display (summarized, updates, omitted) and a reasoning_effort object carrying summary now decide includeThoughts for that request; gemini_include_thoughts_default applies only when the request says nothing * feat(gemini): let a thought summary ask on /v1/messages and /v1/responses override the gateway default * fix(gemini): read a text gateway default as a boolean and leave display omitted to the flag Normalize gemini_include_thoughts_default through normalize_drop_params so a YAML or env value written as "false" or "0" turns thought text off and a null or unknown value keeps the default on. thinking.display "omitted" no longer counts as a per-request ask, so the gateway flag decides, matching /v1/messages. Adds proxy integration cells for chat, /v1/messages and /v1/responses that prove the flag, every per-request ask, streaming, and a burst across workers. * test(gemini): match the streamed Responses usage, which omits null token details * feat(gemini): read the AI SDK reasoning summary alias and thinking.summary as chat asks * fix(gemini): keep the logged reasoning_effort unchanged and key the response cache on the summary alias A dict reasoning_effort next to the AI SDK reasoning_summary alias now counts as a summary ask, and the alias is part of the cache key so a plain request and an alias ask never share a reply * fix(caching): key the response cache on the summary alias wherever the request carries it The cache key now reads reasoning_summary and reasoningSummary through the same reader as the summary fold, so an alias inside extra_body gets its own key too * test: type the new summary-ask unit tests and mark their locals Final * test: annotate the remaining summary-ask unit tests the PR touches --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/__init__.py | 1 + litellm/caching/caching.py | 13 +- .../litellm_core_utils/model_param_helper.py | 19 + .../pass_through/adapters/handler.py | 52 +- litellm/llms/base_llm/chat/transformation.py | 4 + .../vertex_and_google_ai_studio_gemini.py | 70 +- litellm/main.py | 11 +- litellm/responses/bridge_check.py | 16 +- .../transformation.py | 10 +- litellm/utils.py | 32 +- .../chat_completions/reasoning/__init__.py | 0 .../test_chat_completions_reasoning_gemini.py | 1041 +++++++++++++++ .../translation/gemini_thoughts.py | 63 + .../test_messages_reasoning_gemini.py | 373 ++++++ .../responses/reasoning/__init__.py | 0 .../test_responses_reasoning_gemini.py | 400 ++++++ tests/unit/caching/test_caching.py | 152 +-- ..._handler_reasoning_effort_normalization.py | 103 +- ...test_vertex_and_google_ai_studio_gemini.py | 1114 +++++++++-------- .../test_litellm_completion_responses.py | 41 +- 20 files changed, 2841 insertions(+), 674 deletions(-) create mode 100644 tests/integration/translation/chat_completions/reasoning/__init__.py create mode 100644 tests/integration/translation/chat_completions/reasoning/test_chat_completions_reasoning_gemini.py create mode 100644 tests/integration/translation/gemini_thoughts.py create mode 100644 tests/integration/translation/messages/reasoning/test_messages_reasoning_gemini.py create mode 100644 tests/integration/translation/responses/reasoning/__init__.py create mode 100644 tests/integration/translation/responses/reasoning/test_responses_reasoning_gemini.py diff --git a/litellm/__init__.py b/litellm/__init__.py index ff901ede512..af812dade3c 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -393,6 +393,7 @@ key_alias_pattern: str | None = None enable_gemini_default_thinking_level_low: bool = ( False # opt-in: force thinkingLevel low/minimal for Gemini 3 thinking param mapping ) +gemini_include_thoughts_default: bool = True #################### logging: bool = True enable_loadbalancing_on_batch_endpoints: Optional[bool] = None diff --git a/litellm/caching/caching.py b/litellm/caching/caching.py index 655131e3c69..5f5b8210096 100644 --- a/litellm/caching/caching.py +++ b/litellm/caching/caching.py @@ -25,7 +25,7 @@ from litellm._internal_context import current_service_target, service_target from litellm._logging import verbose_logger from litellm.constants import CACHED_STREAMING_CHUNK_DELAY from litellm.integrations.otel.runtime import phase_span -from litellm.litellm_core_utils.model_param_helper import ModelParamHelper +from litellm.litellm_core_utils.model_param_helper import ModelParamHelper, peek_reasoning_summary_aliases from litellm.types.caching import * from litellm.types.utils import EmbeddingResponse, is_litellm_owned_kwarg @@ -376,7 +376,7 @@ class Cache: ) return "".join(f"{field}: {value}" for field, value in scope_values if value is not None) - def get_cache_key(self, **kwargs) -> str: + def get_cache_key(self, **kwargs: object) -> str: """ Get the cache key for the given arguments. @@ -410,13 +410,16 @@ class Cache: if litellm.enable_caching_on_provider_specific_optional_params is True: # feature flagged for now if kwargs[param] is None: continue # ignore None params - param_value = kwargs[param] - cache_key += f"{param}: {param_value}" + cache_key += f"{param}: {kwargs[param]}" if is_semantic_cache: cache_key += self._get_semantic_cache_tenant_scope(kwargs) - hashed_cache_key = Cache._get_hashed_cache_key(cache_key) + reasoning_summary: Final = peek_reasoning_summary_aliases(kwargs) + reasoning_summary_material: Final = ( + "" if reasoning_summary is None else f"reasoning_summary: {reasoning_summary}" + ) + hashed_cache_key = Cache._get_hashed_cache_key(cache_key + reasoning_summary_material) hashed_cache_key = self._add_namespace_to_cache_key(hashed_cache_key, **kwargs) verbose_logger.debug( "\nCreated cache key: %s (source material length: %d)", diff --git a/litellm/litellm_core_utils/model_param_helper.py b/litellm/litellm_core_utils/model_param_helper.py index 2dc52980875..c9c7a0e89ac 100644 --- a/litellm/litellm_core_utils/model_param_helper.py +++ b/litellm/litellm_core_utils/model_param_helper.py @@ -1,3 +1,4 @@ +from collections.abc import Mapping from functools import lru_cache from typing import ClassVar, Final @@ -22,6 +23,24 @@ from litellm.types.llms.anthropic import AnthropicMessagesRequest from litellm.types.rerank import RerankRequest +def peek_reasoning_summary_aliases(optional_params: Mapping[str, object]) -> object | None: + """Read AI-SDK-style reasoning summary from optional_params or nested extra_body. + + Uses key membership (not ``or`` chains) so falsy values like ``""`` are not skipped. + """ + if "reasoningSummary" in optional_params: + return optional_params["reasoningSummary"] + if "reasoning_summary" in optional_params: + return optional_params["reasoning_summary"] + extra_body: Final = optional_params.get("extra_body") + if isinstance(extra_body, dict): + if "reasoningSummary" in extra_body: + return extra_body["reasoningSummary"] + if "reasoning_summary" in extra_body: + return extra_body["reasoning_summary"] + return None + + class ModelParamHelper: # Cached at class level — deterministic set built from static OpenAI type annotations relevant_logging_args: ClassVar[frozenset[str]] = frozenset() diff --git a/litellm/llms/anthropic/pass_through/adapters/handler.py b/litellm/llms/anthropic/pass_through/adapters/handler.py index 4c0613777c6..60c63964492 100644 --- a/litellm/llms/anthropic/pass_through/adapters/handler.py +++ b/litellm/llms/anthropic/pass_through/adapters/handler.py @@ -1,4 +1,5 @@ from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping, Sequence +from dataclasses import dataclass from typing import ( TYPE_CHECKING, Final, @@ -62,6 +63,12 @@ class _CompletionKwargs(TypedDict, total=False, extra_items=object): custom_llm_provider: str +@dataclass(frozen=True, slots=True) +class _ChatTargetReasoning: + effort: str + summary: str | None = None + + def _messages_have_compaction_block(messages: _AnthropicMessages) -> bool: """Return True when any message carries a ``compaction`` content block.""" for msg in messages: @@ -441,15 +448,33 @@ class LiteLLMMessagesToCompletionTransformationHandler: completion_kwargs["reasoning_effort"] = updated_reasoning_effort @staticmethod - def _plain_effort_for_chat_target( + def _requested_reasoning_summary(thinking: Mapping[str, object]) -> str | None: + if thinking.get("type") == "disabled": + return None + summary: Final = thinking.get("summary") + if isinstance(summary, str): + return summary + display: Final = thinking.get("display") + return "auto" if display is not None and display != "omitted" else None + + @staticmethod + def _chat_target_reasoning( completion_kwargs: _CompletionKwargs, *, thinking: Mapping[str, object] | None, - ) -> str | None: - reasoning_effort: Final = completion_kwargs.get("reasoning_effort") - if not thinking or not isinstance(reasoning_effort, dict) or "summary" not in reasoning_effort: + ) -> _ChatTargetReasoning | None: + if not thinking: return None - effort: Final = reasoning_effort.get("effort") + reasoning_effort: Final = completion_kwargs.get("reasoning_effort") + wraps_summary: Final = isinstance(reasoning_effort, dict) and "summary" in reasoning_effort + requested_summary: Final = LiteLLMMessagesToCompletionTransformationHandler._requested_reasoning_summary( + thinking + ) + if not wraps_summary and requested_summary is None: + return None + effort: Final[object] = ( + reasoning_effort.get("effort") if isinstance(reasoning_effort, dict) else reasoning_effort + ) model: Final = completion_kwargs.get("model") if not isinstance(effort, str) or not isinstance(model, str) or not model: return None @@ -467,7 +492,12 @@ class LiteLLMMessagesToCompletionTransformationHandler: return None if resolved_provider == "litellm_proxy": return None - from litellm.responses.bridge_check import responses_api_bridge_check + from litellm.responses.bridge_check import chat_target_accepts_reasoning_summary, responses_api_bridge_check + + if requested_summary is not None and chat_target_accepts_reasoning_summary(local_model, resolved_provider): + return _ChatTargetReasoning(effort=effort, summary=requested_summary) + if not wraps_summary or not isinstance(reasoning_effort, dict): + return None web_search_options: Final = completion_kwargs.get("web_search_options") tools: Final = completion_kwargs.get("tools") @@ -481,7 +511,7 @@ class LiteLLMMessagesToCompletionTransformationHandler: reasoning_effort=reasoning_effort, api_base=resolved_api_base, ) - return None if model_info.get("mode") == "responses" else effort + return None if model_info.get("mode") == "responses" else _ChatTargetReasoning(effort=effort) @staticmethod def _normalize_reasoning_effort( @@ -647,12 +677,14 @@ class LiteLLMMessagesToCompletionTransformationHandler: thinking=thinking, ) - plain_effort: Final = LiteLLMMessagesToCompletionTransformationHandler._plain_effort_for_chat_target( + chat_reasoning: Final = LiteLLMMessagesToCompletionTransformationHandler._chat_target_reasoning( completion_kwargs, thinking=thinking, ) - if plain_effort is not None: - completion_kwargs["reasoning_effort"] = plain_effort + if chat_reasoning is not None: + completion_kwargs["reasoning_effort"] = chat_reasoning.effort + if chat_reasoning is not None and chat_reasoning.summary is not None: + completion_kwargs["reasoning_summary"] = chat_reasoning.summary return completion_kwargs, tool_name_mapping diff --git a/litellm/llms/base_llm/chat/transformation.py b/litellm/llms/base_llm/chat/transformation.py index 2b3c644f266..db4d5e0f820 100644 --- a/litellm/llms/base_llm/chat/transformation.py +++ b/litellm/llms/base_llm/chat/transformation.py @@ -488,6 +488,10 @@ class BaseConfig(ABC): """ return True + @property + def accepts_reasoning_summary_in_reasoning_effort(self) -> bool: + return False + def post_stream_processing(self, stream: object) -> object: """Hook for providers to post-process streaming responses. Default: pass-through.""" return stream diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index bed7bc1bd74..257fafbb6b8 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -27,6 +27,7 @@ from litellm.constants import ( DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET_GEMINI_2_5_PRO, ) from litellm.exceptions import UnsupportedParamsError +from litellm.litellm_core_utils.core_helpers import normalize_drop_params from litellm.litellm_core_utils.json_fragment_accumulator import JSONFragmentAccumulator from litellm.litellm_core_utils.prompt_templates.factory import ( encode_tool_call_id_with_signature, @@ -926,11 +927,38 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): # - Requires propertyOrdering optional_params["response_schema"] = self._map_response_schema(value=schema) + @property + def accepts_reasoning_summary_in_reasoning_effort(self) -> bool: + return True + + @staticmethod + def _include_thoughts(requested: bool | None) -> bool: + if requested is not None: + return requested + return normalize_drop_params(litellm.gemini_include_thoughts_default) is not False + + @staticmethod + def _include_thoughts_requested_by_reasoning_effort(reasoning_effort: object) -> bool | None: + match reasoning_effort: + case {"summary": str()}: + return True + case _: + return None + + @staticmethod + def _map_reasoning_effort(effort: str, model: str, reasoning_effort: object) -> GeminiThinkingConfig: + include_thoughts: Final = VertexGeminiConfig._include_thoughts_requested_by_reasoning_effort(reasoning_effort) + if VertexGeminiConfig._is_gemini_3_or_newer(model): + return VertexGeminiConfig._map_reasoning_effort_to_thinking_level(effort, model, include_thoughts) + return VertexGeminiConfig._map_reasoning_effort_to_thinking_budget(effort, model, include_thoughts) + @staticmethod def _map_reasoning_effort_to_thinking_budget( reasoning_effort: str, model: str | None = None, + include_thoughts: bool | None = None, ) -> GeminiThinkingConfig: + sends_thoughts: Final = VertexGeminiConfig._include_thoughts(include_thoughts) if reasoning_effort == "minimal": # Use model-specific minimum thinking budget or fallback # Check for exact matches first, then partial matches @@ -945,22 +973,22 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): return { "thinkingBudget": budget, - "includeThoughts": True, + "includeThoughts": sends_thoughts, } elif reasoning_effort == "low": return { "thinkingBudget": DEFAULT_REASONING_EFFORT_LOW_THINKING_BUDGET, - "includeThoughts": True, + "includeThoughts": sends_thoughts, } elif reasoning_effort == "medium": return { "thinkingBudget": DEFAULT_REASONING_EFFORT_MEDIUM_THINKING_BUDGET, - "includeThoughts": True, + "includeThoughts": sends_thoughts, } elif reasoning_effort == "high": return { "thinkingBudget": DEFAULT_REASONING_EFFORT_HIGH_THINKING_BUDGET, - "includeThoughts": True, + "includeThoughts": sends_thoughts, } elif reasoning_effort == "disable": return { @@ -987,6 +1015,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): def _map_reasoning_effort_to_thinking_level( reasoning_effort: str, model: str | None = None, + include_thoughts: bool | None = None, ) -> GeminiThinkingConfig: """ Map reasoning_effort to thinking_level for Gemini 3+ models. @@ -998,17 +1027,18 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): GeminiThinkingConfig with thinkingLevel and includeThoughts """ supports_minimal: Final = bool(model) and VertexGeminiConfig._supports_minimal_thinking_level(model) + sends_thoughts: Final = VertexGeminiConfig._include_thoughts(include_thoughts) if reasoning_effort == "minimal": if supports_minimal: - return {"thinkingLevel": "minimal", "includeThoughts": True} + return {"thinkingLevel": "minimal", "includeThoughts": sends_thoughts} else: - return {"thinkingLevel": "low", "includeThoughts": True} + return {"thinkingLevel": "low", "includeThoughts": sends_thoughts} elif reasoning_effort == "low": - return {"thinkingLevel": "low", "includeThoughts": True} + return {"thinkingLevel": "low", "includeThoughts": sends_thoughts} elif reasoning_effort == "medium": - return {"thinkingLevel": "medium", "includeThoughts": True} + return {"thinkingLevel": "medium", "includeThoughts": sends_thoughts} elif reasoning_effort == "high": - return {"thinkingLevel": "high", "includeThoughts": True} + return {"thinkingLevel": "high", "includeThoughts": sends_thoughts} elif reasoning_effort in ("disable", "none"): return { "thinkingLevel": "minimal" if supports_minimal else "low", @@ -1061,6 +1091,11 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): status_code=400, ) + @staticmethod + def _thinking_asks_for_thoughts(thinking_param: Mapping[str, object]) -> bool: + display: Final = thinking_param.get("display") + return isinstance(thinking_param.get("summary"), str) or (display is not None and display != "omitted") + @staticmethod def _map_thinking_param( thinking_param: AnthropicThinkingParam, @@ -1068,6 +1103,8 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): ) -> GeminiThinkingConfig: thinking_enabled: Final = thinking_param.get("type") == "enabled" thinking_budget: Final = thinking_param.get("budget_tokens") + asks_for_thoughts: Final = VertexGeminiConfig._thinking_asks_for_thoughts(thinking_param) + sends_thoughts: Final = VertexGeminiConfig._include_thoughts(True if asks_for_thoughts else None) params: Final[GeminiThinkingConfig] = {} @@ -1077,7 +1114,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): if thinking_budget == 0: params["includeThoughts"] = False else: - params["includeThoughts"] = True + params["includeThoughts"] = sends_thoughts # Follow provider defaults unless explicitly opted into legacy behavior. if litellm.enable_gemini_default_thinking_level_low is True: params["thinkingLevel"] = ( @@ -1089,7 +1126,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): else: # For older Gemini models, use thinkingBudget if thinking_enabled and not VertexGeminiConfig._is_thinking_budget_zero(thinking_budget): - params["includeThoughts"] = True + params["includeThoughts"] = sends_thoughts if thinking_budget is not None and isinstance(thinking_budget, int): params["thinkingBudget"] = thinking_budget @@ -1304,14 +1341,9 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): param_name="reasoning_effort", param_description="thinking_budget", ) - if VertexGeminiConfig._is_gemini_3_or_newer(model): - optional_params["thinkingConfig"] = VertexGeminiConfig._map_reasoning_effort_to_thinking_level( - effort_value, model - ) - else: - optional_params["thinkingConfig"] = VertexGeminiConfig._map_reasoning_effort_to_thinking_budget( - effort_value, model - ) + optional_params["thinkingConfig"] = VertexGeminiConfig._map_reasoning_effort( + effort=effort_value, model=model, reasoning_effort=value + ) elif param == "thinking": # Validate no conflict with thinking_level VertexGeminiConfig._validate_thinking_config_conflicts( diff --git a/litellm/main.py b/litellm/main.py index d8e62ebcdad..e6caa5c25b6 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -52,6 +52,7 @@ from litellm.utils import ( get_litellm_params, get_optional_params, peek_reasoning_summary_aliases, + reasoning_effort_with_summary_alias, strip_reasoning_summary_aliases_from_optional_params, ) @@ -5681,7 +5682,15 @@ def completion( "allowed_openai_params": allowed_openai_params, "base_model": base_model, } - optional_params = get_optional_params(**optional_param_args, **non_default_params) + provider_reasoning_effort: Final = reasoning_effort_with_summary_alias( + reasoning_effort, + non_default_params, + accepts_summary=provider_config is not None + and provider_config.accepts_reasoning_summary_in_reasoning_effort, + ) + optional_params = get_optional_params( + **{**optional_param_args, "reasoning_effort": provider_reasoning_effort}, **non_default_params + ) processed_non_default_params: Final = pre_process_non_default_params( model=model, passed_params=optional_param_args, diff --git a/litellm/responses/bridge_check.py b/litellm/responses/bridge_check.py index c2b4c105199..2552ee8c880 100644 --- a/litellm/responses/bridge_check.py +++ b/litellm/responses/bridge_check.py @@ -4,6 +4,7 @@ from urllib.parse import urlsplit import litellm from litellm._logging import verbose_logger +from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.llms.azure_ai.common_utils import ( azure_ai_supports_native_responses, foundry_chat_rejects_function_tools_while_reasoning, @@ -11,7 +12,8 @@ from litellm.llms.azure_ai.common_utils import ( from litellm.llms.openai.chat.gpt_5_transformation import OpenAIGPT5Config from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import OpenAIWebSearchOptions -from litellm.utils import get_model_info_helper +from litellm.types.utils import LlmProviders +from litellm.utils import ProviderConfigManager, get_model_info_helper _OPENAI_DEFAULT_API_BASE: Final = "https://api.openai.com/v1" _OPENAI_API_HOST: Final = "api.openai.com" @@ -36,6 +38,18 @@ def _resolve_openai_api_base(api_base: str | None) -> str: ) +def chat_target_accepts_reasoning_summary(model: str, custom_llm_provider: str | None) -> bool: + try: + local_model, provider, _, _ = get_llm_provider(model=model, custom_llm_provider=custom_llm_provider) + chat_config: Final = ProviderConfigManager.get_provider_chat_config( + model=local_model, provider=LlmProviders(provider) + ) + except Exception as e: # noqa: BLE001 # a capability probe must never fail the request it probes for + verbose_logger.debug("chat target reasoning summary check failed: %s", e) + return False + return chat_config is not None and chat_config.accepts_reasoning_summary_in_reasoning_effort + + def responses_api_bridge_check( model: str, custom_llm_provider: str, diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 7f5d143e98e..2e7fb6da194 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -56,6 +56,7 @@ from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.llms.bedrock.common_utils import bedrock_declared_cache_point_support from litellm.llms.openrouter.chat.transformation import openrouter_model_supports_cache_control from litellm.llms.snowflake.chat.transformation import is_snowflake_claude_model +from litellm.responses.bridge_check import chat_target_accepts_reasoning_summary from litellm.responses.litellm_completion_transformation.session_handler import ( ResponsesSessionHandler, ) @@ -785,8 +786,8 @@ class LiteLLMCompletionResponsesConfig: provider either rejects it or silently drops it, and dropping it turns reasoning off while still billing for the turn. ``summary`` has no chat equivalent, so it rides the ``reasoning_summary`` alias, which ``litellm.completion`` reassembles into ``{effort, - summary}`` when it bridges the model back onto the Responses API, and is sent to nothing - else. + summary}`` when it bridges the model back onto the Responses API or targets a chat provider + whose config maps a summary ask (Gemini's thought summaries), and is sent to nothing else. """ if not reasoning_param: return ResponsesReasoningChatForm(effort=None, summary=None) @@ -807,7 +808,10 @@ class LiteLLMCompletionResponsesConfig: reasoning_summary=summary, api_base=api_base, ) - return ResponsesReasoningChatForm(effort=effort, summary=summary if bridges_back else None) + forwards_summary: Final = bridges_back or ( + isinstance(effort, str) and chat_target_accepts_reasoning_summary(model, custom_llm_provider) + ) + return ResponsesReasoningChatForm(effort=effort, summary=summary if forwards_summary else None) @staticmethod def transform_responses_api_request_to_chat_completion_request( diff --git a/litellm/utils.py b/litellm/utils.py index c0c80f09e4a..65408663f6f 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -89,6 +89,7 @@ from litellm.litellm_core_utils.fallback_generalizations import ( match_capability_generalizations, match_fill_missing_generalizations, ) +from litellm.litellm_core_utils.model_param_helper import peek_reasoning_summary_aliases from litellm.litellm_core_utils.sensitive_data_masker import redact_credentials_in_payload from litellm.litellm_core_utils.tokenizer import strip_special_tokens from litellm.rust_bridge import tokenizer as tokenizer_dispatch @@ -10446,22 +10447,21 @@ def get_non_default_completion_params(kwargs: Mapping[str, object]) -> dict[str, return filter_out_litellm_params(kwargs, excluding=litellm.OPENAI_CHAT_COMPLETION_PARAMS) -def peek_reasoning_summary_aliases(optional_params: dict) -> object | None: - """Read AI-SDK-style reasoning summary from optional_params or nested extra_body. - - Uses key membership (not ``or`` chains) so falsy values like ``""`` are not skipped. - """ - if "reasoningSummary" in optional_params: - return optional_params["reasoningSummary"] - if "reasoning_summary" in optional_params: - return optional_params["reasoning_summary"] - extra_body: Final = optional_params.get("extra_body") - if isinstance(extra_body, dict): - if "reasoningSummary" in extra_body: - return extra_body["reasoningSummary"] - if "reasoning_summary" in extra_body: - return extra_body["reasoning_summary"] - return None +def reasoning_effort_with_summary_alias( + reasoning_effort: object, params: Mapping[str, object], accepts_summary: bool +) -> object: + summary: Final = peek_reasoning_summary_aliases(params) + if not accepts_summary or not isinstance(summary, str): + return reasoning_effort + match reasoning_effort: + case str(): + return {"effort": reasoning_effort, "summary": summary} + case {"summary": _}: + return reasoning_effort + case {"effort": str() as effort}: + return {"effort": effort, "summary": summary} + case _: + return reasoning_effort def strip_reasoning_summary_aliases_from_optional_params( diff --git a/tests/integration/translation/chat_completions/reasoning/__init__.py b/tests/integration/translation/chat_completions/reasoning/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/integration/translation/chat_completions/reasoning/test_chat_completions_reasoning_gemini.py b/tests/integration/translation/chat_completions/reasoning/test_chat_completions_reasoning_gemini.py new file mode 100644 index 00000000000..194e6972453 --- /dev/null +++ b/tests/integration/translation/chat_completions/reasoning/test_chat_completions_reasoning_gemini.py @@ -0,0 +1,1041 @@ +import json +from collections.abc import Mapping +from concurrent.futures import ThreadPoolExecutor +from dataclasses import replace +from pathlib import Path +from typing import Final +from unittest.mock import ANY +from uuid import uuid4 + +import pytest +from integration._support.client import Gateway, eventually, list_value, object_value, string_value +from integration._support.process import graceful_stop_seconds +from integration._support.provider import SharedProvider +from integration._support.wire import Reply +from integration.caching.response_cache_case import CACHE_KEY, redis_store +from integration.translation.case import SSEFrame, TranslationTestCase +from integration.translation.chat_completions.bases.gemini import GEMINI_3_8_FLASH_TEST_CASE +from integration.translation.chat_completions.bases.vertex_ai import ( + GEMINI_3_8_FLASH_TEST_CASE as VERTEX_GEMINI_3_8_FLASH_TEST_CASE, +) +from integration.translation.gemini_thoughts import ( + LONG_SUMMARY, + gemini_thoughts_proxy, + sent_thinking_config, + thoughts_off, + with_thinking_config, +) +from integration.translation.messages.bases.gemini import ( + GEMINI_3_8_FLASH_TEST_CASE as MESSAGES_GEMINI_3_8_FLASH_TEST_CASE, +) +from integration.translation.responses.bases.gemini import ( + GEMINI_3_8_FLASH_TEST_CASE as RESPONSES_GEMINI_3_8_FLASH_TEST_CASE, +) +from integration.translation.runner import ( + assert_provider_request, + assert_translation, + encode_sse_frame, + parse_sse_frames, +) +from pydantic import JsonValue + +__all__ = ("thoughts_off",) + +pytestmark = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) + +GEMINI_3_8_FLASH_THINKING_SHOWN_THOUGHT: Final = "**My Inspection Strategy: Validation First**\n\nOkay, here's how I'm approaching this project review. The instruction specifically mentions considering zero, negative, and non-integer quantities. To efficiently address that, I'm going to start with `shop/validation.py`. My reasoning is simple: the validation logic is the first line of defense. Understanding how the system *expects* to handle these edge cases will inform everything else. It's the logical place to begin my investigation to understand the code. I'll use the `read_file` tool to examine that file. I am **not** changing any files. I will call the `chat_completions` client for this project review. I am now proceeding with reading `shop/validation.py`.\n\n\n" + +GEMINI_3_8_FLASH_THOUGHTS_OFF_THOUGHT_SIGNATURE: Final = ( + "EowMCokMAWkUfRPuOEs7DW4IXhoH/SNh9qYezwNSqhXT2tKoI1mcXDG+yx7R+12sMmzvjHh2OSghAlsvmR6welOzIM1LMw6c" + "QrSX2IPRo/xbtUvDLlo6aqufzRtini5EF0yxQdmyn9r9xznzDyuEf4eqrbK0KQAArEVM13rm3CTDIjyDRQPawLALfUh/e61/" + "XVcWrr1PZLR5Jj46ImyrobGGLmJuoytm3ezTEihTb2b0q3/vHOW9cO/gJzMcbCI8koexQLuBB1KLSHUCaZZQYPj4jLx/g0de" + "2l3H/DYR7Cdroh0P/R3mewkhu+3D2ELnHKt85InrdQOqdXp/XhICsivPgsIEWHMRr6kQ3ZLIUJvJbJ7zo7/4bYnHOm+JqmoX" + "E+VMobI2YZ+ff9GOYg4vsJRTeF7ff9kZHYBWRM9XvTsZMbeCZHRbCOLNmSSJPfn2ToDOgMxU1u/GqHYmVSV97ybT3Y0mxwtf" + "2WcVlPUCkFX0GW3GP5IHVkEISo0TznJSI3/yYb+5UmQFMuuIedDdhKa74eU1/K59VuTnd3DzVX/lfhpFbGLDkOSurRsDsViN" + "69zq+6VbLi+5ARi6SJqUU+1U02QaSIQVp8kxtrhEIsPedjCIYHQQNxR8hFPMhnEssR51THNELWH6sVccxpJIbFfGjl0TZS9d" + "rfYL2c+IBX56lrH8IX2gz1iMOsESKwTyAECmqJIL7d/HcVdTk404Hd9SYycB+2MP2GN2XA7ivoO/zfkkFcf+PGmXPed8dHlW" + "bAxrQcOAz89p8kPzb1bou5QUT3feeRG06aJL31gX7BSYj19mvLeo0AmMzx+Wb91NFlwUvlJYlhJo4NfZWg/ymcGclgM8r7iE" + "MmWZa8wul8g4hbJOg3GE9NV70zvKqEpQrJLTa9jHca0DBvj8N3C+XvpeEg+Dys0cxE7Ak0VVYbvu9r0OHKSWtC7NDYrJsF2n" + "VFmdavSxpyPEqPaAdWdiewDH0nhP5GCR/ciikokWFQMF5iClWRfL72qEbFx3cQkctqEf+kK2CevIWKPf8bvkrKVDdPqbjk0G" + "9NWuud+YgNxl4dSCkAWoMrAFlxXozEqMmanMAtZk/7nPeOVFbYEL/xOpSCfxaEKS1zEmMwTb/UNVAogsNY9YdBSFYVNSrS8F" + "N/IrHBvsP66+YEFuR46Di6cWuuZQdSOfB2uQe1eSC8IgrQ6XUhQHWeXKUXfbcb/1abPOi5PUMYyLBQ7r52UKCEfOBc1j0oNL" + "wYzI/vWAz74gq8NblQhHB2VlrInJ9rRVY1MehInqUH0kcG6ysQDQLiLJiS30TV0bZh05I86RfcjaN1twbRir82YchRjXG8/1" + "6iFRJ54Sk7oVzbdnbFOugDea8qIblFmgUTq0+SRLGBFcDMvAqYLNcR2ZO7o0xinMG7IuNRwoSMahnmreWYr9saVJAScd3All" + "i6RQuxjGIdxG0s0UNuoWiSNIf1Z2ngoJrvZHTc0iZK+A2zv3CBqWAaS8WGwSfcItnlc6hMB0bTWasoCelmRjfv+1jQnKzKMj" + "DzZlO6skMkv5x9GzLrAvpYoulLYa/wp9+exw/Ll7MxSB22aeAGfai9C7Vj310y8nApsG3QxaoUUO4jg1eCjrrluBLgoVY8cE" + "/O95IAUa2rFzVrOWbHuob7TEI/8+BwieGeyeHdPGn4fNeKYMkRjFeVJ1GA/2+iCHCjCbSNm4sdzh9188/Es3q4acRfQUiqwR" + "iQzW2E1fy9mXRDY9y5fdXuzjHtuCkH8MqR0cioN1p4OGZzPi6xXFaqV576NLjySBpfpQEYeXRA3j5NVwTwbeYKnGOot4bEtJ" + "JQfelwkXLXeAfxzCTD0/vGtdhilzEK8wSYM9a+8llr5D4XwtEsu/68YbXJ2jwzsS3rE/Frulrf6paz9fS45ZXj98+OT5EpFO" + "o5XHfTXKb5Brzr0tKfPSB2Mx6fXbYP3lZRwfJIaUq8Alwdo7g2jnHQUhCNBB8hkb/IkQrnhsAGS6f0ZEjwf6ErVOa8DkfizN" + "y4WTdS1D3gOUJxOcszm8J5b085tsfFxEJVgA0VjI60DBY0owWn2I" +) + +GEMINI_3_8_FLASH_THOUGHTS_OFF_RESPONSE: Final[Mapping[str, JsonValue]] = { + "candidates": [ + { + "content": { + "parts": [ + { + "text": "To understand how zero, negative, and non-integer quantities (such as floats or invalid types) are currently handled and where the checkout rejection is enforced, we need to inspect `shop/validation.py` next.", + }, + { + "functionCall": { + "name": "read_file", + "args": { + "path": "shop/validation.py", + "end_line": 50, + "start_line": 1, + }, + "id": "call_330625", + }, + "thoughtSignature": GEMINI_3_8_FLASH_THOUGHTS_OFF_THOUGHT_SIGNATURE, + }, + ], + "role": "model", + }, + "finishReason": "STOP", + "index": 0, + "finishMessage": "Model generated function call(s).", + }, + ], + "usageMetadata": { + "promptTokenCount": 358, + "candidatesTokenCount": 77, + "totalTokenCount": 793, + "promptTokensDetails": [ + { + "modality": "TEXT", + "tokenCount": 358, + }, + ], + "thoughtsTokenCount": 358, + "serviceTier": "standard", + }, + "modelVersion": "gemini-3.8-flash", + "responseId": "xy3RapaNEbT4-sAPhMrt6AU", +} + +GEMINI_3_8_FLASH_THINKING_OMITTED_THOUGHT_SIGNATURE: Final = ( + "EscECsQEAWkUfROlWYT6hX0DWOKSCg4MSuLzOV7u43EgJDn0dfS2R8v/MVGGcNolGA9rteb/cV3dyCSQ7D2EO2CB5whDkblj" + "fGAAw6NBhqnidYNQiXbsg7oqynxS116ZQnk302S4KzvAbIvo4gpRbahPVClXKkP8PCQWA5eRQ/e9oQeHr97gbBB7+PLsl0NK" + "m7YPV6lP3fAXx/r3EADbdNL1diHylKd6tCWUKRea/x2QIOBQZdY5jHSU73z/pnPOotVwggEiasDLSLA+uUR1KllgGc8WKZ4i" + "uRM3L1gDr1QZ4Rht3quiF8AwfopWE5m2ON+FD78aHUDG/lgDs0vXnJluzri9LpPaexfbpvwMPtAx8Bxrq1/dCeekR2VnVx1c" + "qZ7TSoWMqk4x98mPpyUFz1hSBRnXaOGvQl5Wk7Iart49tH2LoIBYoV+PkCgNZ2MOi2uzc7OlN/8gvOeBM05xbGAplsLiuKd4" + "jkgv7Nz8Jby3v/hxtQT9+Ow3mxwSaiIfEkuaxkPiK1cX1nwT92gFzpmhVsrGqpiCUdAfos39Wd2/rcrFbtl5iMAVcHG5JGjw" + "8AjubFqiWDMqWPC58FfsX9TLC2BMUbPf2xi3hrw+1QWDtSuNl9s+dofkCogXFdLA9lAFc0oRj3MVhzYUhn5eixApJdggDsJF" + "iJaQtRihs4dstbXX9QXYFnx22Ql3x9TWJibiH8jz4fJVOf0wreMWZ2xVXBiw2xKdy5dB7xRDnrQD/sdcpTpCWURYXVCXglJ9" + "DD7e7/OCNMLCPw==" +) + +GEMINI_3_8_FLASH_THINKING_OMITTED_RESPONSE: Final[Mapping[str, JsonValue]] = { + "candidates": [ + { + "content": { + "parts": [ + { + "text": "To understand how the checkout validates quantities\u2014including zero, negative, and non-integer quantities\u2014I need to inspect `shop/validation.py` next.", + }, + { + "functionCall": { + "name": "read_file", + "args": { + "start_line": 1, + "path": "shop/validation.py", + "end_line": 40, + }, + "id": "call_745953", + }, + "thoughtSignature": GEMINI_3_8_FLASH_THINKING_OMITTED_THOUGHT_SIGNATURE, + }, + ], + "role": "model", + }, + "finishReason": "STOP", + "index": 0, + "finishMessage": "Model generated function call(s).", + }, + ], + "usageMetadata": { + "promptTokenCount": 358, + "candidatesTokenCount": 65, + "totalTokenCount": 531, + "promptTokensDetails": [ + { + "modality": "TEXT", + "tokenCount": 358, + }, + ], + "thoughtsTokenCount": 108, + "serviceTier": "standard", + }, + "modelVersion": "gemini-3.8-flash", + "responseId": "zS3RasS-Lp2-jMcPoay98As", +} + +GEMINI_3_8_FLASH_THINKING_SHOWN_THOUGHT_SIGNATURE: Final = ( + "EqUECqIEAWkUfRNZEW6qD3+V/1LAc1riV0fvqkg4Bldd7YYDfln+ZHNg6Lbw2T31UZHTum0qYvrAOor848+zSkFG7qBDfdZ9" + "69FGl/dQGq2llr1yrv23blnbwC7tO0uevr8225EZ/4Pv+hbfP0teFcSZsasE4LQangtxYkewghc6cH1Acz293SIz1cqF/3CM" + "Ta+FtJ5x106oauWy79+z1VbV9TNyHpNtsrUy3ADR2PFNBjXKf2TvUWygzR4w2RJg581+dfx0sI97O3UlZrhSKRe+p/6siUP7" + "qnysURAia7k+ew5f05Nuo1aDwD2G20G2s9v7L0PxnCCAay5L4YTub5LXxvUGBXMhSHVb9MmO8+2rNPqDtPWh0bOQdKbkWAV/" + "dMckMLTKANDbnzBjtCK1HbN7Rw0sIBSnBVy59z7h+0ierHFgA8Vntfp/nvAy0wDpuZu09agdB3T6FO2HCvWCFG2Y+0B4lGJV" + "ErOknvn50gtn9Tpn1QSGv4VDRrrcv6nCOyohoL9CsBeO1Bs+RItyZSAoTA3bVB5IjcF2kKK27+/+VRJjN0cO4Cm5RD/ETb4+" + "kI101AFWJ3CYwKEu91SAVTtXLewQYd5LBiJ3Btt68m8d5uUp0a1kAn2cH6g2C7TOfoJ7BgU5+vU81mS8hk9QOFGb9RYUXs8Y" + "cJl21l6WfMMZri5bMI9qCAFn+YZdBBeIbGq+JJcIgjr/R2tULmU584+5beoEydkS" +) + +GEMINI_3_8_FLASH_THINKING_SHOWN_RESPONSE: Final[Mapping[str, JsonValue]] = { + "candidates": [ + { + "content": { + "parts": [ + { + "text": GEMINI_3_8_FLASH_THINKING_SHOWN_THOUGHT, + "thought": True, + }, + { + "text": "I need to inspect `shop/validation.py` next to examine how zero, negative, and non-integer quantities are currently handled in `validate_quantity`.", + }, + { + "functionCall": { + "name": "read_file", + "args": { + "end_line": 40, + "start_line": 1, + "path": "shop/validation.py", + }, + "id": "call_386884", + }, + "thoughtSignature": GEMINI_3_8_FLASH_THINKING_SHOWN_THOUGHT_SIGNATURE, + }, + ], + "role": "model", + }, + "finishReason": "STOP", + "index": 0, + "finishMessage": "Model generated function call(s).", + }, + ], + "usageMetadata": { + "promptTokenCount": 358, + "candidatesTokenCount": 66, + "totalTokenCount": 529, + "promptTokensDetails": [ + { + "modality": "TEXT", + "tokenCount": 358, + }, + ], + "thoughtsTokenCount": 105, + "serviceTier": "standard", + }, + "modelVersion": "gemini-3.8-flash", + "responseId": "0C3RavXPIJWX9MoPtorosAs", +} + +GEMINI_2_5_FLASH_THOUGHTS_OFF_RESPONSE: Final[Mapping[str, JsonValue]] = { + "candidates": [ + { + "content": { + "parts": [ + { + "text": "I will read the `shop/validation.py` file to understand the quantity validation logic.", + }, + { + "functionCall": { + "name": "read_file", + "args": { + "start_line": 1, + "path": "shop/validation.py", + "end_line": 20, + }, + }, + }, + ], + "role": "model", + }, + "finishReason": "STOP", + "index": 0, + "finishMessage": "Model generated function call(s).", + }, + ], + "usageMetadata": { + "promptTokenCount": 334, + "candidatesTokenCount": 53, + "totalTokenCount": 387, + "promptTokensDetails": [ + { + "modality": "TEXT", + "tokenCount": 334, + }, + ], + "serviceTier": "standard", + }, + "modelVersion": "gemini-2.5-flash", + "responseId": "0y3Rar76LO_K-8YPh6Kd0Ak", +} + +VERTEX_GEMINI_3_8_FLASH_THOUGHTS_OFF_THOUGHT_SIGNATURE: Final = ( + "AY89a18W+yVUlI0b2XsXQujlNKztnSpKXMf0+1JFiprxfp7kc0bG2g9WoPITuNkCqkuBaO5SkZaKDrKQTlJFJKj+H6SnETLk" + "JrR3SAzp2GTjwM2QRsJjI5ky54oylVT93Jh51PWCMp6vl8tZstMJrois8b1i4/eTOyKaVprypLW8DOIImqN3Cbkea8nEDe9r" + "XsdpdvsR7i+NvD0eiCW+Z4ZIflx5NVENbyEezjyJCL3lg3bV2Mfb0ytsACHXuqgGQEcH9HBKQKGkF5n7mA0GNLw6wAqlk8FE" + "JKttP6NgS+HzS/U+xaqxci6B50l0hIMw3pTOHu8Gg5fRKxf/dzvks+s709/ODCGBDrhMzpf/ilnPAMrB67cQrHvQOhnI2pGk" + "cvkZv+ZGwAQxjJHs/r0aXDGU2mwN9RY+/rSEv8y4MWd5c1KIs7xa1GytDBaaG8Zwh3kZzaSOYjVKIk7cy31mPJ01jCNdW9Pg" + "1p2r2iwbzJrBHo5Ku3s1yK4rADeY1WOU63sM6evVkWJEtu6KPxAqE9JkQKQoAPOFtoP89RtiTSl0nZZaO0stbqG7G+AkE2iO" + "5WkKCB1o55R7EEZDalDgFPput8khce5hPrhIGikN3yo11sHsAA2PAuDcHo4L3eWaYPExFiZtB+gVTbBwx42ST3UuBB9lrHtk" + "GCrp7tlOBHJJY+RlXj1G0TJsSxUZDSEswCB+gjG+1tdEchKGRRKOBUjWGARCozBsIEuBOR0F1ugSCLYn7CumS2bLwmMhKY7+" + "mfKXEgF2U84MkmwDEOBSTubNbRhy0qNhVPcLT8pJvGvtTvLny1C90TBTHErVzOKb3Z095fFpUyW7MmFQj5J9rX3A160W0SNQ" + "RKXf2Yr6ikyVhLBSFNvQ00n81iXlyg7bp++ChbbCJBFCUcrGF5WnPLk3+isixodP+G0b+pNtCNo8xMkrvqGqL5AMlmGJ9Qm4" + "a9S6QgDRaw5AzcV/W90a7/BcaIeqkIly3Z9M0+d3Eep4igYj6AK0JzX1BhchMogq9g/dquTTZZk4edVKflBmQcwg4qGGvioI" + "X/ae4tKJVqB2U0Ed+A6t3JZ4Ggdpopdnuy29XvdEyTmUg+fF2sTCgAdK07hueKYO2pkhqj86mSy88iHV0/secWQs0ZuHtcSW" + "D9lAG57yvpMf6RIwccLTyVmM4jcO0A2OB0sTnNIVKYYtswPtUqbrPMTR5saBfEfnULipI521kpb3BseODBKar+G1FISLrb0Q" + "OjMhNnCWCoCSVDk1FbDmtEvSjfiUzeJtFYoqM8PIXyGiBZnsmjGaOw3QMNvLfkoy4gqE6nvFVgwIk3+J6Z1GdFecCPSSCatz" + "LRwwQ8B5ZdlJ6+xjTDWIuyaOByD5BsB98e0YHyuLSYP17w6+0pSGXf3M6AdBJXK7hcL/vu3sGzftYARFHgRY1PQsYCqW8RVy" + "uJL/GwEAT4EgzkDDWufizfUHJEAUVBpQKkxZW7eTbyHyPS/Uiu5a74f5L8lYyep1pAdqt3iK4lvmQ3T9h3DAQIQ15T8PLV78" + "FjzJvZDQ2uMryvoRS1MGEcK6dxc4VvcWalQTjyczE4jQIxEEiN3U5Yo5CaM6FFGDnjIJBiy0fUD16P8e4ZaEaRFQCY4QZ0sA" + "bhwSHWhnU7QWvA/CVnqOVni5TpRsY1ShtEaQXKTslUuz96lB6XwEeaDkPqsPdw+mWY9QmmLqXnXERnTq13rQY1s=" +) + +VERTEX_GEMINI_3_8_FLASH_THOUGHTS_OFF_RESPONSE: Final[Mapping[str, JsonValue]] = { + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + { + "text": "To understand how zero, negative, and non-integer quantities are handled during checkout validation, we need to inspect `shop/validation.py` next.", + }, + { + "functionCall": { + "name": "read_file", + "args": { + "start_line": 1, + "path": "shop/validation.py", + "end_line": 50, + }, + "id": "call_1272842", + }, + "thoughtSignature": VERTEX_GEMINI_3_8_FLASH_THOUGHTS_OFF_THOUGHT_SIGNATURE, + }, + ], + }, + "finishReason": "STOP", + }, + ], + "usageMetadata": { + "promptTokenCount": 246, + "candidatesTokenCount": 64, + "totalTokenCount": 583, + "trafficType": "ON_DEMAND", + "promptTokensDetails": [ + { + "modality": "TEXT", + "tokenCount": 246, + }, + ], + "candidatesTokensDetails": [ + { + "modality": "TEXT", + "tokenCount": 64, + }, + ], + "thoughtsTokenCount": 273, + }, + "modelVersion": "gemini-3.8-flash", + "createTime": "2026-10-15T20:02:49.484984Z", + "responseId": "aTHRavjMHYvi6tkPt6OroAo", +} + +GEMINI_3_8_FLASH_THOUGHTS_OFF_STREAM_THOUGHT_SIGNATURE: Final = ( + "ErkLCrYLAWkUfRO3TXF/xF4wcqKCnAFmIW9MEkIMZeaFFG/BmstoEFBw+0v0jari9LOu/TkcHKNFvewXanY1T+6rm512Fsvm" + "3cbxU5C7oP8kGPeEHIBOs2BIeocBFx3ZXSmIn25pOLd2YOqDaTzu/De/QX2JUkno3r2Ozbvuy0wvS5EsD3816rYYnaYc7+4E" + "iZRQhf5r6notr70qcHJ3DuKIP8ghDEz0gNSN1bAUa8N5Hvjh7ITbHLoSkeKO0pOjXaSWfQ86xv9F2zhpMZfW+muxcm1+g9tT" + "5QkyJxMtn+cnHi/w4MbQy+9tbJWdkkLo2H6z81pBsmjvIZAQnjjF0zP8aL91/+S9Po8HNV0LsrGs4P4kUgRR33I6UPMCVnHf" + "B7IxmVSEIGeFQ8gg1mJRDj7yZJFBgf61ya97bVWAhiAm6t/EHdnMi9xIPFEYSrY6ed8ZGUnf05x/GRMVOfrmeCePZIHXpftk" + "UfI2W4+qIM1YGR8+GD7KOtQEkcsrjRGHvhmcv1GW3NlLVpZ4LWIqX6ZtmZrOjDFvgnCg/jYdzDHL5j2vQiiSlcpsHJ4UZWRj" + "NmqUyUY0OZC77erTqensEeWJggUfyizdp3InLZtWM/C5B6rhLGku3LgGwAxdSf5aY25i0f+Ayl/xE/nnCkXS81D64z/+Xqdn" + "W4jsZn2tg8TGb11T6GncBK82aE/T9zVUHT7mbUtEljdinxNc/NS+9mWEkhtLSOAQqQmGF24oxMFlCcejVhGZz5QnFpjnu2Uk" + "COv2AZ7S9iYJQa/xT1LyQ18IibvAUWQPRBpjwRnOf0wi0beuE9Q/h0Y2itFjCt+mz5JlFRyQRJcpEOupOsfF47ULZCLeG0Uz" + "Af8JQzHMUWL3ESlOs30cLQQBJwtHW3QD3ZA1uwMCZnJcpBvIr7GhRdF/VhKZ7C8aYZLMOk32M1i+iGe5lYOM2WuoEzbprTGz" + "aVQeQBWEIdyT+RhXjpvLrrsWOIWI/OAw1Y8VEtouLt246yDjDrVuKUdJfXPz9VfnJ/7dKmsYWAaOCBnQ6edZ9HY8vzEKyuma" + "mcaJWSrwuq1mIouQ1Pw4fvazi6JE2xx17hEuOpS5sLD8E89V9uUbkcDne4XFoHNF9mm3nVgTplNFqo5J4kSiEZJ6qJveiMyF" + "sw3F7SbuU0f95AE/HVwMfDhBbnT3j0R5s2zejxLW6mWuASodcn+J9WGbRfMDLZV+aozCPy+rrXrIfeVQw8nLfjQbwY/iQLFM" + "rQCM4m9F5MLrIM0jGKePUeLRmkhJ92LiqEmDMkOZct09jVaZpSW6dRO+ANGfxqJ0y8KppBA7jgWy/PDKMHfwQFWh3Goxc86m" + "otxTUbF2i15mhhPGN2Z5xsRe1lsYRsFV7LNIkqLXcUVQ7uQGmlzi7wCGiMasvxDgzTLdTGAJ6FsaWzr+GAHH/RFSyRA9a6Qr" + "/Z/h9tm5sx5l6hTQrYspK3qK+cYI7BJCuiYYiIMWEAV6qb5SNhMZ88UlTayryuoolQ7pBco59MBV+P6oAoWAXFh69D4yDAe8" + "L2iNa6UTLt2BSMpAgqE5HdrsifGPfwOwNrwtfv/t4eXIBpBVgqC4csZSrN4zUsP6fH3mPZ4oPgwkDFaLu/x1S/PdxKvF0mYg" + "e3ORBogDHyO00QzHPP/9sl2PEgzIQmQv+0tkmpmAeTUfU8FMlRouFsgbrO2wUx1P8KoaFO4ZToQHeNaoqhI2yZ+GB9EquuBp" + "Xpu46qaU1Wrf3w7vtI0aLUPDKnudfZ3b/me3LVm1cgutomL5o33J86U0v62TwJz8xnLB0881oHQVXX4VfaPV4BVk74QCMcqK" + "/SRt3GpK9V1OGQh47ubs5MUFBnqauOnOKmttTxerI5QNi/t5vSwIvdy2MBcXnCXPJSjt3MpZjSZo34PQHK+1iuR1BCbmmHQk" + "v2UCtcYUZfnJ4FuKNXaZ/AFDT7VxbdnOaS8HVA==" +) + +GEMINI_3_8_FLASH_THOUGHTS_OFF_STREAM_FRAMES: Final = ( + SSEFrame( + data={ + "candidates": [ + { + "content": { + "parts": [ + { + "text": "To understand how zero, negative, and non-integer quantities are handled during checkout, I need to inspect `shop/validation.py", + }, + ], + "role": "model", + }, + "index": 0, + }, + ], + "usageMetadata": { + "promptTokenCount": 358, + "candidatesTokenCount": 27, + "totalTokenCount": 673, + "promptTokensDetails": [ + { + "modality": "TEXT", + "tokenCount": 358, + }, + ], + "thoughtsTokenCount": 288, + "serviceTier": "standard", + }, + "modelVersion": "gemini-3.8-flash", + "responseId": "yi3RaufbH6bXjMcPt5mF2Q0", + } + ), + SSEFrame( + data={ + "candidates": [ + { + "content": { + "parts": [ + { + "text": "` next.", + }, + ], + "role": "model", + }, + "index": 0, + }, + ], + "usageMetadata": { + "promptTokenCount": 358, + "candidatesTokenCount": 30, + "totalTokenCount": 676, + "promptTokensDetails": [ + { + "modality": "TEXT", + "tokenCount": 358, + }, + ], + "thoughtsTokenCount": 288, + "serviceTier": "standard", + }, + "modelVersion": "gemini-3.8-flash", + "responseId": "yi3RaufbH6bXjMcPt5mF2Q0", + } + ), + SSEFrame( + data={ + "candidates": [ + { + "content": { + "parts": [ + { + "functionCall": { + "name": "read_file", + "args": { + "end_line": 50, + "path": "shop/validation.py", + "start_line": 1, + }, + "id": "call_317677", + }, + "thoughtSignature": GEMINI_3_8_FLASH_THOUGHTS_OFF_STREAM_THOUGHT_SIGNATURE, + }, + ], + "role": "model", + }, + "index": 0, + }, + ], + "usageMetadata": { + "promptTokenCount": 358, + "candidatesTokenCount": 63, + "totalTokenCount": 709, + "promptTokensDetails": [ + { + "modality": "TEXT", + "tokenCount": 358, + }, + ], + "thoughtsTokenCount": 288, + "serviceTier": "standard", + }, + "modelVersion": "gemini-3.8-flash", + "responseId": "yi3RaufbH6bXjMcPt5mF2Q0", + } + ), + SSEFrame( + data={ + "candidates": [ + { + "content": { + "parts": [ + { + "text": "", + }, + ], + "role": "model", + }, + "finishReason": "STOP", + "index": 0, + }, + ], + "usageMetadata": { + "promptTokenCount": 358, + "candidatesTokenCount": 63, + "totalTokenCount": 709, + "promptTokensDetails": [ + { + "modality": "TEXT", + "tokenCount": 358, + }, + ], + "thoughtsTokenCount": 288, + "serviceTier": "standard", + }, + "modelVersion": "gemini-3.8-flash", + "responseId": "yi3RaufbH6bXjMcPt5mF2Q0", + } + ), +) + +GEMINI_2_5_FLASH_REQUEST: Final = with_thinking_config( + { + **GEMINI_3_8_FLASH_TEST_CASE.expected_provider_request, + "contents": [ + list_value(GEMINI_3_8_FLASH_TEST_CASE.expected_provider_request["contents"])[0], + { + "role": "model", + "parts": [ + {"text": "I will locate the quantity validation and its tests."}, + { + "function_call": { + "name": "search_files", + "args": {"query": "quantity", "paths": ["shop/", "tests/"]}, + } + }, + ], + }, + { + "role": "user", + "parts": [ + { + "function_response": { + "name": "search_files", + "response": { + "matches": [ + { + "path": "shop/validation.py", + "line": 18, + "text": "def validate_quantity(quantity):", + }, + { + "path": "tests/test_validation.py", + "line": 7, + "text": "def test_zero_quantity():", + }, + ], + }, + }, + }, + ], + }, + list_value(GEMINI_3_8_FLASH_TEST_CASE.expected_provider_request["contents"])[3], + ], + }, + {"thinkingBudget": 4096, "includeThoughts": False}, +) + + +def gemini_chat_response( + *, + content: str, + tool_call: Mapping[str, JsonValue], + message_fields: Mapping[str, JsonValue], + usage: Mapping[str, JsonValue], +) -> Mapping[str, JsonValue]: + return { + **GEMINI_3_8_FLASH_TEST_CASE.expected_litellm_response, + "choices": [ + { + "finish_reason": "tool_calls", + "index": 0, + "message": { + "content": content, + "role": "assistant", + "tool_calls": [dict(tool_call)], + "images": [], + **message_fields, + }, + "provider_specific_fields": {"native_finish_reason": "STOP"}, + }, + ], + "usage": dict(usage), + } + + +def signed_tool_call(call_id: str, arguments: str, signature: str) -> Mapping[str, JsonValue]: + return { + "index": 0, + "provider_specific_fields": {"thought_signature": signature}, + "function": {"arguments": arguments, "name": "read_file"}, + "id": f"{call_id}__thought__{signature}", + "type": "function", + } + + +THINKING_REQUEST: Final[Mapping[str, JsonValue]] = { + key: value for key, value in GEMINI_3_8_FLASH_TEST_CASE.litellm_request.items() if key != "reasoning_effort" +} + +GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE: Final = replace( + GEMINI_3_8_FLASH_TEST_CASE, + scenario="thoughts_off", + expected_provider_request=with_thinking_config( + GEMINI_3_8_FLASH_TEST_CASE.expected_provider_request, {"thinkingLevel": "high", "includeThoughts": False} + ), + mock_provider_response=GEMINI_3_8_FLASH_THOUGHTS_OFF_RESPONSE, + expected_litellm_response=gemini_chat_response( + content=( + "To understand how zero, negative, and non-integer quantities (such as floats or invalid types) are " + "currently handled and where the checkout rejection is enforced, we need to inspect " + "`shop/validation.py` next." + ), + tool_call=signed_tool_call( + "call_330625", + '{"path": "shop/validation.py", "end_line": 50, "start_line": 1}', + GEMINI_3_8_FLASH_THOUGHTS_OFF_THOUGHT_SIGNATURE, + ), + message_fields={ + "thinking_blocks": [], + "provider_specific_fields": {"thought_signatures": [GEMINI_3_8_FLASH_THOUGHTS_OFF_THOUGHT_SIGNATURE]}, + }, + usage={ + "completion_tokens": 435, + "prompt_tokens": 358, + "total_tokens": 793, + "completion_tokens_details": {"reasoning_tokens": 358, "text_tokens": 77}, + "prompt_tokens_details": {"text_tokens": 358}, + }, + ), +) + +GEMINI_2_5_FLASH_THOUGHTS_OFF_TEST_CASE: Final = replace( + GEMINI_3_8_FLASH_TEST_CASE, + scenario="thoughts_off", + litellm_request={**GEMINI_3_8_FLASH_TEST_CASE.litellm_request, "model": "gemini/gemini-2.5-flash"}, + expected_provider_endpoint="/models/gemini-2.5-flash:generateContent", + expected_provider_request=GEMINI_2_5_FLASH_REQUEST, + mock_provider_response=GEMINI_2_5_FLASH_THOUGHTS_OFF_RESPONSE, + expected_litellm_response={ + **gemini_chat_response( + content="I will read the `shop/validation.py` file to understand the quantity validation logic.", + tool_call={ + "index": 0, + "function": { + "arguments": '{"start_line": 1, "path": "shop/validation.py", "end_line": 20}', + "name": "read_file", + }, + "id": ANY, + "type": "function", + }, + message_fields={"thinking_blocks": []}, + usage={ + "completion_tokens": 53, + "prompt_tokens": 334, + "total_tokens": 387, + "completion_tokens_details": {"text_tokens": 53}, + "prompt_tokens_details": {"text_tokens": 334}, + }, + ), + "model": "gemini/gemini-2.5-flash", + }, +) + +VERTEX_GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE: Final = replace( + VERTEX_GEMINI_3_8_FLASH_TEST_CASE, + scenario="thoughts_off", + expected_provider_request=with_thinking_config( + VERTEX_GEMINI_3_8_FLASH_TEST_CASE.expected_provider_request, {"thinkingLevel": "high", "includeThoughts": False} + ), + mock_provider_response=VERTEX_GEMINI_3_8_FLASH_THOUGHTS_OFF_RESPONSE, + expected_litellm_response={ + **gemini_chat_response( + content=( + "To understand how zero, negative, and non-integer quantities are handled during checkout validation, " + "we need to inspect `shop/validation.py` next." + ), + tool_call=signed_tool_call( + "call_1272842", + '{"start_line": 1, "path": "shop/validation.py", "end_line": 50}', + VERTEX_GEMINI_3_8_FLASH_THOUGHTS_OFF_THOUGHT_SIGNATURE, + ), + message_fields={ + "thinking_blocks": [], + "provider_specific_fields": { + "thought_signatures": [VERTEX_GEMINI_3_8_FLASH_THOUGHTS_OFF_THOUGHT_SIGNATURE] + }, + }, + usage={ + "completion_tokens": 337, + "prompt_tokens": 246, + "total_tokens": 583, + "completion_tokens_details": {"reasoning_tokens": 273, "text_tokens": 64}, + "prompt_tokens_details": {"text_tokens": 246}, + }, + ), + "model": "vertex_ai/gemini-3.8-flash", + }, +) + +GEMINI_3_8_FLASH_THINKING_OMITTED_TEST_CASE: Final = replace( + GEMINI_3_8_FLASH_TEST_CASE, + scenario="thinking_omitted", + litellm_request={ + **THINKING_REQUEST, + "thinking": {"type": "enabled", "budget_tokens": 1024, "display": "omitted"}, + }, + expected_provider_request=with_thinking_config( + GEMINI_3_8_FLASH_TEST_CASE.expected_provider_request, {"includeThoughts": False} + ), + mock_provider_response=GEMINI_3_8_FLASH_THINKING_OMITTED_RESPONSE, + expected_litellm_response=gemini_chat_response( + content=( + "To understand how the checkout validates quantities—including zero, negative, and non-integer " + "quantities—I need to inspect `shop/validation.py` next." + ), + tool_call=signed_tool_call( + "call_745953", + '{"start_line": 1, "path": "shop/validation.py", "end_line": 40}', + GEMINI_3_8_FLASH_THINKING_OMITTED_THOUGHT_SIGNATURE, + ), + message_fields={ + "thinking_blocks": [], + "provider_specific_fields": {"thought_signatures": [GEMINI_3_8_FLASH_THINKING_OMITTED_THOUGHT_SIGNATURE]}, + }, + usage={ + "completion_tokens": 173, + "prompt_tokens": 358, + "total_tokens": 531, + "completion_tokens_details": {"reasoning_tokens": 108, "text_tokens": 65}, + "prompt_tokens_details": {"text_tokens": 358}, + }, + ), +) + +GEMINI_3_8_FLASH_THINKING_SHOWN_TEST_CASE: Final = replace( + GEMINI_3_8_FLASH_TEST_CASE, + scenario="thinking_shown", + litellm_request={ + **THINKING_REQUEST, + "thinking": {"type": "enabled", "budget_tokens": 1024, "display": "summarized"}, + }, + expected_provider_request=with_thinking_config( + GEMINI_3_8_FLASH_TEST_CASE.expected_provider_request, {"includeThoughts": True} + ), + mock_provider_response=GEMINI_3_8_FLASH_THINKING_SHOWN_RESPONSE, + expected_litellm_response=gemini_chat_response( + content=( + "I need to inspect `shop/validation.py` next to examine how zero, negative, and non-integer quantities " + "are currently handled in `validate_quantity`." + ), + tool_call=signed_tool_call( + "call_386884", + '{"end_line": 40, "start_line": 1, "path": "shop/validation.py"}', + GEMINI_3_8_FLASH_THINKING_SHOWN_THOUGHT_SIGNATURE, + ), + message_fields={ + "reasoning_content": GEMINI_3_8_FLASH_THINKING_SHOWN_THOUGHT, + "thinking_blocks": [{"type": "thinking", "thinking": GEMINI_3_8_FLASH_THINKING_SHOWN_THOUGHT}], + "provider_specific_fields": {"thought_signatures": [GEMINI_3_8_FLASH_THINKING_SHOWN_THOUGHT_SIGNATURE]}, + }, + usage={ + "completion_tokens": 171, + "prompt_tokens": 358, + "total_tokens": 529, + "completion_tokens_details": {"reasoning_tokens": 105, "text_tokens": 66}, + "prompt_tokens_details": {"text_tokens": 358}, + }, + ), +) + + +def with_request(case: TranslationTestCase, scenario: str, **fields: JsonValue) -> TranslationTestCase: + return replace(case, scenario=scenario, litellm_request={**case.litellm_request, **fields}) + + +def with_thinking(case: TranslationTestCase, scenario: str, display: JsonValue) -> TranslationTestCase: + return with_request(case, scenario, thinking={"type": "enabled", "budget_tokens": 1024, "display": display}) + + +def with_summary(case: TranslationTestCase, scenario: str, summary: JsonValue) -> TranslationTestCase: + return with_request(case, scenario, reasoning_effort={"effort": "high", "summary": summary}) + + +@pytest.mark.parametrize( + "case", + ( + GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE, + GEMINI_2_5_FLASH_THOUGHTS_OFF_TEST_CASE, + VERTEX_GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE, + ), + ids=( + GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE.id, + GEMINI_2_5_FLASH_THOUGHTS_OFF_TEST_CASE.id, + VERTEX_GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE.id, + ), +) +def test_gemini_keeps_thinking_and_billing_reasoning_without_thought_text_when_the_gateway_turns_it_off( + case: TranslationTestCase, thoughts_off: Gateway, provider: SharedProvider +) -> None: + assert_translation(case, thoughts_off, provider) + + +SUMMARY_ASKS: Final = ( + with_summary(GEMINI_3_8_FLASH_TEST_CASE, "summary_auto", "auto"), + with_summary(GEMINI_3_8_FLASH_TEST_CASE, "summary_empty", ""), + with_summary(GEMINI_3_8_FLASH_TEST_CASE, "summary_long", LONG_SUMMARY), + with_summary(GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE, "summary_int", 7), + with_summary(GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE, "summary_list", ["auto"]), + with_summary(GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE, "summary_null", None), + with_request(GEMINI_3_8_FLASH_TEST_CASE, "summary_alias", reasoning_summary="auto"), + with_request(GEMINI_3_8_FLASH_TEST_CASE, "summary_alias_camel", reasoningSummary="detailed"), + with_request( + GEMINI_3_8_FLASH_TEST_CASE, + "summary_alias_effort_dict", + reasoning_effort={"effort": "high"}, + reasoning_summary="auto", + ), + with_request(GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE, "summary_alias_null", reasoning_summary=None), + with_request( + GEMINI_3_8_FLASH_THINKING_SHOWN_TEST_CASE, + "thinking_summary", + thinking={"type": "enabled", "budget_tokens": 1024, "summary": "auto"}, + ), +) + + +@pytest.mark.parametrize("case", SUMMARY_ASKS, ids=tuple(case.id for case in SUMMARY_ASKS)) +def test_only_a_string_reasoning_summary_brings_thought_text_back_over_the_gateway_default( + case: TranslationTestCase, thoughts_off: Gateway, provider: SharedProvider +) -> None: + assert_translation(case, thoughts_off, provider) + + +def without_closed_schema(tool: JsonValue) -> JsonValue: + function: Final = object_value(object_value(tool)["function"]) + parameters: Final = object_value(function["parameters"]) + open_parameters: Final = {name: value for name, value in parameters.items() if name != "additionalProperties"} + return {**object_value(tool), "function": {**function, "parameters": open_parameters}} + + +def cacheable_run(case: TranslationTestCase, marker: str) -> TranslationTestCase: + messages: Final = list_value(case.litellm_request["messages"]) + system: Final = object_value(messages[0]) + instruction: Final = object_value(case.expected_provider_request["system_instruction"]) + (part,) = list_value(instruction["parts"]) + text: Final = f"{string_value(system['content'])} Review run {marker}." + return replace( + case, + litellm_request={ + **{name: value for name, value in case.litellm_request.items() if name != "cache"}, + "messages": [{**system, "content": text}, messages[1]], + "tools": [without_closed_schema(tool) for tool in list_value(case.litellm_request["tools"])], + }, + expected_provider_request={ + **case.expected_provider_request, + "contents": [list_value(case.expected_provider_request["contents"])[0]], + "system_instruction": {**instruction, "parts": [{**object_value(part), "text": text}]}, + }, + ) + + +def response_cache_keys() -> frozenset[str]: + keys: Final = (raw.decode() for raw in redis_store().scan_iter()) + return frozenset(key for key in keys if CACHE_KEY.match(key)) + + +def stored_reply_key(before: frozenset[str]) -> str: + (key,) = eventually(lambda: response_cache_keys() - before, lambda added: len(added) == 1) + return key + + +GEMINI_3_8_FLASH_SUMMARY_ALIAS_TEST_CASE: Final = with_request( + GEMINI_3_8_FLASH_TEST_CASE, "summary_alias", reasoning_summary="auto" +) + + +@pytest.mark.parametrize( + ("first", "second"), + ( + (GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE, GEMINI_3_8_FLASH_SUMMARY_ALIAS_TEST_CASE), + (GEMINI_3_8_FLASH_SUMMARY_ALIAS_TEST_CASE, GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE), + ), + ids=("plain_then_summary_alias", "summary_alias_then_plain"), +) +def test_a_summary_alias_ask_and_a_plain_request_never_share_a_cached_reply( + first: TranslationTestCase, second: TranslationTestCase, thoughts_off: Gateway, provider: SharedProvider +) -> None: + marker: Final = uuid4().hex + before: Final = response_cache_keys() + assert_translation(cacheable_run(first, marker), thoughts_off, provider) + first_key: Final = stored_reply_key(before) + assert_translation(cacheable_run(second, marker), thoughts_off, provider) + second_key: Final = stored_reply_key(before | {first_key}) + repeat: Final = thoughts_off.request("POST", second.litellm_endpoint, cacheable_run(second, marker).litellm_request) + assert repeat.status_code == 200, repeat.text + assert repeat.headers.get("x-litellm-cache-key") == second_key, sorted(repeat.headers) + assert provider.received() == () + + +DISPLAY_ASKS: Final = ( + GEMINI_3_8_FLASH_THINKING_SHOWN_TEST_CASE, + with_thinking(GEMINI_3_8_FLASH_THINKING_SHOWN_TEST_CASE, "display_updates", "updates"), + with_thinking(GEMINI_3_8_FLASH_THINKING_SHOWN_TEST_CASE, "display_int", 7), + with_thinking(GEMINI_3_8_FLASH_THINKING_SHOWN_TEST_CASE, "display_empty", ""), + with_thinking(GEMINI_3_8_FLASH_THINKING_SHOWN_TEST_CASE, "display_list", ["omitted"]), + with_thinking(GEMINI_3_8_FLASH_THINKING_SHOWN_TEST_CASE, "display_long", LONG_SUMMARY), + GEMINI_3_8_FLASH_THINKING_OMITTED_TEST_CASE, +) + + +@pytest.mark.parametrize("case", DISPLAY_ASKS, ids=tuple(case.id for case in DISPLAY_ASKS)) +def test_any_thinking_display_but_omitted_brings_thought_text_back_over_the_gateway_default( + case: TranslationTestCase, thoughts_off: Gateway, provider: SharedProvider +) -> None: + assert_translation(case, thoughts_off, provider) + + +DEFAULT_GATEWAY_DISPLAYS: Final = ( + with_thinking(GEMINI_3_8_FLASH_THINKING_SHOWN_TEST_CASE, "display_omitted_default", "omitted"), + GEMINI_3_8_FLASH_THINKING_SHOWN_TEST_CASE, +) + + +@pytest.mark.parametrize("case", DEFAULT_GATEWAY_DISPLAYS, ids=tuple(case.id for case in DEFAULT_GATEWAY_DISPLAYS)) +def test_thinking_display_never_turns_thought_text_off_on_the_default_gateway( + case: TranslationTestCase, gateway: Gateway, provider: SharedProvider +) -> None: + assert_translation(case, gateway, provider) + + +def test_a_streamed_turn_keeps_thinking_and_billing_reasoning_without_thought_text( + thoughts_off: Gateway, provider: SharedProvider +) -> None: + provider.expect( + Reply( + content_type="text/event-stream", + chunks=tuple(encode_sse_frame(frame) for frame in GEMINI_3_8_FLASH_THOUGHTS_OFF_STREAM_FRAMES), + ) + ) + response: Final = thoughts_off.request( + "POST", + "/v1/chat/completions", + {**GEMINI_3_8_FLASH_TEST_CASE.litellm_request, "stream": True, "stream_options": {"include_usage": True}}, + ) + assert_provider_request( + provider.received(), + "/models/gemini-3.8-flash:streamGenerateContent?alt=sse", + GEMINI_3_8_FLASH_TEST_CASE.expected_provider_headers, + GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE.expected_provider_request, + ) + assert response.status_code == 200, response.text + frames: Final = tuple(frame.data for frame in parse_sse_frames(response.text) if frame.data != "[DONE]") + deltas: Final = tuple( + object_value(object_value(choice)["delta"]) + for frame in frames + for choice in list_value(frame.get("choices", [])) + ) + assert "".join(str(delta.get("content") or "") for delta in deltas) == ( + "To understand how zero, negative, and non-integer quantities are handled during checkout, I need to " + "inspect `shop/validation.py` next." + ) + assert [delta for delta in deltas if delta.get("reasoning_content") or delta.get("thinking_blocks")] == [] + usage: Final = [object_value(frame["usage"]) for frame in frames if frame.get("usage")][-1] + assert {field: value for field, value in usage.items() if field != "cost"} == { + "completion_tokens": 351, + "prompt_tokens": 358, + "total_tokens": 709, + "completion_tokens_details": {"reasoning_tokens": 288, "text_tokens": 63}, + "prompt_tokens_details": {"text_tokens": 358}, + } + assert isinstance(usage["cost"], float) and usage["cost"] > 0 + + +@pytest.mark.parametrize( + ("configured", "case"), + (("false", GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE), (None, GEMINI_3_8_FLASH_TEST_CASE)), + ids=("string-false", "null"), +) +def test_a_flag_written_as_text_or_left_null_still_sends_gemini_a_boolean( + configured: JsonValue, case: TranslationTestCase, gateway: Gateway, provider: SharedProvider, tmp_path: Path +) -> None: + with gemini_thoughts_proxy(gateway, tmp_path, configured) as owned: + assert_translation(case, owned.gateway, provider) + + +BURST_ASKS: Final[tuple[tuple[str, Mapping[str, JsonValue], Mapping[str, JsonValue]], ...]] = ( + ( + "/v1/chat/completions", + {**GEMINI_3_8_FLASH_TEST_CASE.litellm_request, "reasoning_effort": {"effort": "low", "summary": "auto"}}, + {"thinkingLevel": "low", "includeThoughts": True}, + ), + ( + "/v1/chat/completions", + GEMINI_3_8_FLASH_TEST_CASE.litellm_request, + {"thinkingLevel": "high", "includeThoughts": False}, + ), + ( + "/v1/messages", + MESSAGES_GEMINI_3_8_FLASH_TEST_CASE.litellm_request, + {"thinkingLevel": "high", "includeThoughts": True}, + ), + ( + "/v1/messages", + {**MESSAGES_GEMINI_3_8_FLASH_TEST_CASE.litellm_request, "thinking": {"type": "adaptive"}}, + {"thinkingLevel": "high", "includeThoughts": False}, + ), + ( + "/v1/responses", + {**RESPONSES_GEMINI_3_8_FLASH_TEST_CASE.litellm_request, "reasoning": {"effort": "high", "summary": "auto"}}, + {"thinkingLevel": "high", "includeThoughts": True}, + ), + ( + "/v1/responses", + RESPONSES_GEMINI_3_8_FLASH_TEST_CASE.litellm_request, + {"thinkingLevel": "high", "includeThoughts": False}, + ), +) + + +def _assert_each_request_follows_its_own_ask(gateway: Gateway, provider: SharedProvider) -> None: + burst: Final = BURST_ASKS * 4 + for _ in burst: + provider.expect(Reply(body=json.dumps(GEMINI_3_8_FLASH_THOUGHTS_OFF_RESPONSE).encode())) + with ThreadPoolExecutor(max_workers=12) as pool: + responses: Final = tuple(pool.map(lambda ask: gateway.request("POST", ask[0], ask[1]), burst)) + assert [response.status_code for response in responses] == [200] * len(burst), [ + response.text for response in responses if response.status_code != 200 + ] + sent: Final = sorted(json.dumps(sent_thinking_config(request), sort_keys=True) for request in provider.received()) + assert sent == sorted(json.dumps(expected, sort_keys=True) for _, _, expected in burst) + + +def test_every_worker_honors_the_flag_on_every_endpoint_under_a_burst_and_after_a_worker_respawns( + gateway: Gateway, provider: SharedProvider, tmp_path: Path +) -> None: + with gemini_thoughts_proxy(gateway, tmp_path, False, workers=2) as owned: + _assert_each_request_follows_its_own_ask(owned.gateway, provider) + owned.kill_worker(min(owned.worker_pids())) + _assert_each_request_follows_its_own_ask(owned.gateway, provider) diff --git a/tests/integration/translation/gemini_thoughts.py b/tests/integration/translation/gemini_thoughts.py new file mode 100644 index 00000000000..e0d89c3d43e --- /dev/null +++ b/tests/integration/translation/gemini_thoughts.py @@ -0,0 +1,63 @@ +from collections.abc import Iterator, Mapping +from contextlib import contextmanager +from pathlib import Path +from typing import Final + +import pytest +import yaml +from integration._support.client import JSON_OBJECT, Gateway, list_value, object_value +from integration._support.process import OwnedProxy, owned_proxy_process +from integration._support.wire import Request +from pydantic import JsonValue + +GEMINI_2_5_FLASH_DEPLOYMENT: Final[Mapping[str, JsonValue]] = { + "model_name": "gemini/gemini-2.5-flash", + "litellm_params": { + "model": "gemini/gemini-2.5-flash", + "api_base": "http://127.0.0.1:8191", + "api_key": "synthetic-gemini-key", + }, +} +LONG_SUMMARY: Final = "detailed " * 600 + + +def with_thinking_config( + request: Mapping[str, JsonValue], thinking_config: Mapping[str, JsonValue] +) -> Mapping[str, JsonValue]: + return { + **request, + "generationConfig": {**object_value(request["generationConfig"]), "thinkingConfig": dict(thinking_config)}, + } + + +def sent_thinking_config(request: Request) -> Mapping[str, JsonValue]: + body: Final = JSON_OBJECT.validate_json(request.body) + return object_value(object_value(body["generationConfig"])["thinkingConfig"]) + + +@contextmanager +def gemini_thoughts_proxy( + gateway: Gateway, directory: Path, default: JsonValue, *, workers: int = 1 +) -> Iterator[OwnedProxy]: + stock: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + config: Final = directory / "gemini-thoughts.yaml" + config.write_text( + yaml.safe_dump( + { + **stock, + "model_list": [*list_value(stock["model_list"]), dict(GEMINI_2_5_FLASH_DEPLOYMENT)], + "litellm_settings": { + **object_value(stock["litellm_settings"]), + "gemini_include_thoughts_default": default, + }, + } + ) + ) + with owned_proxy_process(gateway, directory, {}, config=config, workers=workers) as owned: + yield owned + + +@pytest.fixture(scope="module") +def thoughts_off(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Gateway]: + with gemini_thoughts_proxy(session_gateway, tmp_path_factory.mktemp("gemini-thoughts-off"), False) as owned: + yield owned.gateway diff --git a/tests/integration/translation/messages/reasoning/test_messages_reasoning_gemini.py b/tests/integration/translation/messages/reasoning/test_messages_reasoning_gemini.py new file mode 100644 index 00000000000..c350d6e2a10 --- /dev/null +++ b/tests/integration/translation/messages/reasoning/test_messages_reasoning_gemini.py @@ -0,0 +1,373 @@ +from collections.abc import Mapping +from dataclasses import replace +from typing import Final + +import pytest +from integration._support.client import Gateway, object_value +from integration._support.process import graceful_stop_seconds +from integration._support.provider import SharedProvider +from integration._support.wire import Reply +from integration.translation.case import SSEFrame, TranslationTestCase +from integration.translation.gemini_thoughts import LONG_SUMMARY, thoughts_off, with_thinking_config +from integration.translation.messages.bases.gemini import GEMINI_3_8_FLASH_TEST_CASE, GEMINI_MESSAGES_RESPONSE +from integration.translation.runner import ( + assert_provider_request, + assert_translation, + encode_sse_frame, + parse_sse_frames, +) +from pydantic import JsonValue + +__all__ = ("thoughts_off",) + +pytestmark = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) + +GEMINI_3_8_FLASH_THOUGHTS_OFF_THOUGHT_SIGNATURE: Final = ( + "EscLCsQLAWkUfRP4jtXWjNAugx5FeAwxasL/kKzH9gXUHHB2cTK13h4C8C6tKVYMKP5rXoPq4yhoXSrsDiIjUEpuAjdQYYCz" + "Wvx21tGAN4iBA4VELBqPOi92vo9YVGpczJ/XVREq9KNRFIkyMzWyEzR1lBbJFA4detKHzSVJ7VQ1iRBn9jk6xvknTCxQhbnl" + "Pct8TfMo4XluFsM2PY3qAqoX1DjkrjQNe4nrN9Yoglq8Bwr2ayyc3DUNkNJDu0DQN20W6HU7fbeqSizEOsvGIvJVEUZhMOFF" + "0r8NyzUGkNiiAZ+TEupIV+G0eTr6sVlMPLS1N74ryQ55sevdj6irenHlP19IwkiDNhy6ApRIoblY5I1uJVt4lj4tQyACGMAE" + "VOTg3Gb78hPgJ+8cF3ukmRF0H1bpYH3n/mpnkqiggfK3vm7fVnsCy03W7Keyl7xd7Gfto/HL7vPCta/sO6ZmZm28JvdCr/Ce" + "X/W82nHuEqcjiR8fNPk413rL/geE4Lppru1H9ESWBICXzwl6pLxaBBNr8BD/jGMvlYA/ZEcBfKWtYyzow3+yp+4M0fsY1Ah7" + "pxVCcCbivMRv2X5J5m9btU1IIm2REJicjMSWCxh4lQMEuuk2FaPzge+wXROC/Zx0HwNw2iPw6VpziTOByACQxXOE7RdWULq1" + "3rkbyXCMkuAfG281SiylOnJ3I3tkbRcFoTaY5DcMBKYCCmoJJZPO1HKJ2DMw5Spz1HDftQ7lg9df5FDoZbmehUZKfY+j/2LN" + "//Wr5OEUmuBpojC5ht3YYgXaxclo/pa70KnLHlQWHPZ3nqdB2ru0tqgOaXSv9KZNhtvXgagLsmffQ/rm5QiRoEeG4oZ5E6Ga" + "VLQB9OvWCDNdzgNVzss0tveNcG/cttJBnkSd5/azHSvzkCDQYh4J/GT3zDVvOyOayWTqzGSLgIoFGpTES7Vsbpph62apefvO" + "+ZAlxPNNo/3tHncHzGc8Pku07fF1yhiKdGRV1OjYndfQv6UuDc3G7n2oU90iCKWRoCqOYAotzw1yCh4lbsSpEWO2iB9kWW93" + "mJsdqQBdk6e9UiuHdJIOfQ+cOVNOrltXXtZZ+T2mKFTFaWaN1F7MMH+xIX10dcnsWzq5o0Qz/tpwEfFOaQ5j43ez5BJdE5yD" + "wj4sDIdY57jAQ4141jU2oHpJCjY5Obhd4692k/FRtN7IMWMn5rQS+uTctWOh/27qatdy+QeAMPZseT4Yoi3W6W//Um9+sJY4" + "RjJ0Q27Fjb/wEqBkGsqx1sIidQc19sZWa2jQxVgG3T8xHYCMbgkqT/OHJaHGa4gxo78ajhdQnGYNWtvuRlGDpdUsxcNfF+6z" + "kbqWgcxKleRnQ1jKKZKS9GeiLoJnBRk8wIgmnrleblMB0GRj8i124HKLTfAm3WHlqy3xpACI3AIxPbk1GrATdkrb496eecVG" + "ko9YTrsdgOn5RV6Ggy/Rxr5wAAvPKBrqx/u1Whyr05pge9AZcGAtFRIZMWRCBnNg+oHE8M1EVipx3wQQXFDFO27Jpr3IandO" + "GgpbjbpwcW0fnDKLw7Ndt8SCyyVUEZwZ4Vp4tXJCOv9rzIMHJi+68Upwp5ILirrBFT5vKs//ub6l5hHzEfZty66+DvXyrpqT" + "VLYRnYMEO10Te24QyYVCO4qSnbJzfwTv5Kp4ODP2ONdSuJgkyYy78MzZ926H1/8rOlTEbWpAYb+QJwrt1nBJBkWRNDi2QNrl" + "UlFKTXZNLNXTP/rHHqfw1us8bhwe4Up2BbECPO+8F0tgSifXusnShtsRwpW7o+z/iXpnwuWePH92YiORo43mdoG4MJ54Emwf" + "B7x9ftguHpq1ZDOBoa3frE8l8a5RIEHemkzQsLlPyZSB+eQgnXrToXwUeEfR3EHG0vSCfokMLk0HinhWGa6hgIyJZQAdUIg4" + "GALvjSsrGEHwJaeLNNn/9EscJOEGyFQFMEy+sVIkL6Pj2+iZuG0OKahk" +) + +GEMINI_3_8_FLASH_THOUGHTS_OFF_RESPONSE: Final[Mapping[str, JsonValue]] = { + "candidates": [ + { + "content": { + "parts": [ + { + "text": "I need to inspect `shop/validation.py` next to examine how `validate_quantity` handles zero, negative, and non-integer quantities.", + }, + { + "functionCall": { + "name": "read_file", + "args": { + "end_line": 40, + "path": "shop/validation.py", + "start_line": 1, + }, + "id": "call_412671", + }, + "thoughtSignature": GEMINI_3_8_FLASH_THOUGHTS_OFF_THOUGHT_SIGNATURE, + }, + ], + "role": "model", + }, + "finishReason": "STOP", + "index": 0, + "finishMessage": "Model generated function call(s).", + }, + ], + "usageMetadata": { + "promptTokenCount": 354, + "candidatesTokenCount": 64, + "totalTokenCount": 739, + "promptTokensDetails": [ + { + "modality": "TEXT", + "tokenCount": 354, + }, + ], + "thoughtsTokenCount": 321, + "serviceTier": "standard", + }, + "modelVersion": "gemini-3.8-flash", + "responseId": "1C3RatW-Jtuc_uMPhaLroA0", +} + +GEMINI_3_8_FLASH_THOUGHTS_OFF_STREAM_THOUGHT_SIGNATURE: Final = ( + "Eo0LCooLAWkUfRMuNip0QePQVa4NCyMrc0dwJR2qI/PSRcyScXfwpZhjBJVTorCs5ei6D4xzVf9attfTHx8bOBCX9wXynQ6V" + "JxU6uN562z5VAyrgoGv2nVGJk/mT7SnzIFiOc4RN+xebNGO0BIDf95xyOMeYP+ZHXquqtx1SnQXw7V6p81wTA4cGMnqjNE7F" + "Vs7OqlcpvzeFUmE8gl/Cy1a+u5SgLHwfnvOR3uzIX3u8kdk//dzvCKqXYaRStDf0dsrJICex/9WoOBfZKCujLBMK7CCjKEWF" + "nqERtFxXayTUKz1dvvBkcZ+MbhtrG3RFkBegqokQn9lus3icwBAfkZ9b3ps42j/mYw7DfjR+3Dg/Pyz8mgJVRC312QRxt+70" + "yOCU8POm/GTHJAZvvxdm9U4eiYBs3789F5oV34YzZl1+pcE4WVRwiFkB0vi17xy5XJFYTtzuk9d9ySr5Fz1de6nTKWk6TpuX" + "PM4q8bhX51kJWPCr1W4oVn1GhKpCtRyrXzKhuMsWcHXu/yKHUVjxyEFY+cTsVtHfrMHsIhMHyNnBqz2JclGXtFZScaq6AZ8g" + "0JZPqqAXgfgNyx2RaxoaCtITznOlBwMC1hTAZCrtjhHQNvGsAEdaRBAnZRFgeDHkuL7r7D9tr38UdxaUf/R5vfl5eO0ksj4X" + "bdeUG/LAnWP2CLUVjIoF50o4ZHKRy5UT/Mhr45986ePmtWnIN9t6dXOCB5I4qoEaljwkDn+llwZMbPq1orHdSrUAR6flcl3h" + "qVutnVOFgN3NeHvNkrkjUCvckiQL/ZBiG5YeDKfMP+XECDWIoTS2/DCJ7SX6d44B2yJ/PO2PH3nm7hr81U7r7xYKly4/2YFX" + "8bAsFJt4lGihfDHXFOr/VbQQd/CobZtLV9UZP7m1wUGXfL0knP31bdFWDKmSp/2giddbpIsGQf6j7h+YiE3M6QEfvRi7kaDX" + "CQ16O/4XJzG85/BxPw8ssR1cN+xLM95pZkDQO2MN1rxLVogF+T+Xc0M2SSjwbahM4nwYmisOqXv4Wuy0V7hQYJUySCvHabdH" + "nmTi+3mIfs9/zzHT3QUGZvAumBCZYfz1kqyz6mytXSwQbhhktmDLgIaj1abGm2gSn10fxl8pDm/mTq8/owS9aNPyDYChNBMr" + "NGYEBMXA4DtZy9TjybzMXKGPi4d/M7EZM5sC33WnCl2JovveyeZUNPWkZTV0+TjSSqk8C488huCOzv+PgJMvtBq6XiQW9ZV1" + "HQWk8XRMK7VMMhuv9sRqy1Lh/h8kIlSWtaCCNwpqIgUlZ+xybbY9dRGdsWq5FoxCU/OLmlSv7BcqXAamxxfIK9inXRKd/xkI" + "bR7OCeIaFq97Z5M08LfUpu9WjcXGgvtsppOUB4Rik2uCBVSzA/OqQdY1/N3q+BPYCvhOVVEBOJEEAvH9iAQOtCw4peaCtSLL" + "7vLh7s/lIZGBYj5I4zjoOfoSuT9g1x9uvv9h2sJUeP1HsN9Mp95weKBbF3HUN5F/lVRTan2I+Wb2+WnMMbashS1hoVIfk/Bc" + "426M3xPy7JDsczXLUP0PAyhAc2+xiFnkI120l4IHT3JK53AArx7svpTxzLOvPffVIcodkyWXHeJeC6eQbD6Elt/wWbx5Ig1L" + "mJQP7egSoyqmzu0zpuIC0xVsmDdgnumSi7acTjp1mIsLHHLLaWPcpxaNeiwBQf/7FfKsGqJ6p+I+lzt/rdRDOR4apkwKRwnr" + "d5beSf38TJd7L8Zj1qXp36Q2rnvtKubSm7f451w9ZbKsqbMDrwoagNWsfG9kc0CR0wtraSRzBHD6GUfoapPbplOLsSyjzEvX" + "i48aG814wRmQk932cd8ETx1a9BIbN/bsOr6stttugYakpvD3b8ckhfXW7eXfpKywBBzu1s1KLrU=" +) + +GEMINI_3_8_FLASH_THOUGHTS_OFF_STREAM_FRAMES: Final = ( + SSEFrame( + data={ + "candidates": [ + { + "content": { + "parts": [ + { + "text": "I need to read `shop/validation.py` next to see how `validate_quantity` currently handles zero, negative, and", + }, + ], + "role": "model", + }, + "index": 0, + }, + ], + "usageMetadata": { + "promptTokenCount": 354, + "candidatesTokenCount": 27, + "totalTokenCount": 676, + "promptTokensDetails": [ + { + "modality": "TEXT", + "tokenCount": 354, + }, + ], + "thoughtsTokenCount": 295, + "serviceTier": "standard", + }, + "modelVersion": "gemini-3.8-flash", + "responseId": "2C3RavD_A77W-sAP9brxyQk", + } + ), + SSEFrame( + data={ + "candidates": [ + { + "content": { + "parts": [ + { + "text": " non-integer quantities.", + }, + ], + "role": "model", + }, + "index": 0, + }, + ], + "usageMetadata": { + "promptTokenCount": 354, + "candidatesTokenCount": 32, + "totalTokenCount": 681, + "promptTokensDetails": [ + { + "modality": "TEXT", + "tokenCount": 354, + }, + ], + "thoughtsTokenCount": 295, + "serviceTier": "standard", + }, + "modelVersion": "gemini-3.8-flash", + "responseId": "2C3RavD_A77W-sAP9brxyQk", + } + ), + SSEFrame( + data={ + "candidates": [ + { + "content": { + "parts": [ + { + "functionCall": { + "name": "read_file", + "args": { + "start_line": 1, + "end_line": 50, + "path": "shop/validation.py", + }, + "id": "call_654547", + }, + "thoughtSignature": GEMINI_3_8_FLASH_THOUGHTS_OFF_STREAM_THOUGHT_SIGNATURE, + }, + ], + "role": "model", + }, + "index": 0, + }, + ], + "usageMetadata": { + "promptTokenCount": 354, + "candidatesTokenCount": 65, + "totalTokenCount": 714, + "promptTokensDetails": [ + { + "modality": "TEXT", + "tokenCount": 354, + }, + ], + "thoughtsTokenCount": 295, + "serviceTier": "standard", + }, + "modelVersion": "gemini-3.8-flash", + "responseId": "2C3RavD_A77W-sAP9brxyQk", + } + ), + SSEFrame( + data={ + "candidates": [ + { + "content": { + "parts": [ + { + "text": "", + }, + ], + "role": "model", + }, + "finishReason": "STOP", + "index": 0, + }, + ], + "usageMetadata": { + "promptTokenCount": 354, + "candidatesTokenCount": 65, + "totalTokenCount": 714, + "promptTokensDetails": [ + { + "modality": "TEXT", + "tokenCount": 354, + }, + ], + "thoughtsTokenCount": 295, + "serviceTier": "standard", + }, + "modelVersion": "gemini-3.8-flash", + "responseId": "2C3RavD_A77W-sAP9brxyQk", + } + ), +) + +ADAPTIVE_THINKING: Final[Mapping[str, JsonValue]] = {"type": "adaptive"} + +GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE: Final = replace( + GEMINI_3_8_FLASH_TEST_CASE, + scenario="thoughts_off", + litellm_request={**GEMINI_3_8_FLASH_TEST_CASE.litellm_request, "thinking": dict(ADAPTIVE_THINKING)}, + expected_provider_request=with_thinking_config( + GEMINI_3_8_FLASH_TEST_CASE.expected_provider_request, {"thinkingLevel": "high", "includeThoughts": False} + ), + mock_provider_response=GEMINI_3_8_FLASH_THOUGHTS_OFF_RESPONSE, + expected_litellm_response={ + **GEMINI_MESSAGES_RESPONSE, + "id": "1C3RatW-Jtuc_uMPhaLroA0", + "model": "gemini/gemini-3.8-flash", + "content": [ + { + "type": "text", + "text": ( + "I need to inspect `shop/validation.py` next to examine how `validate_quantity` handles zero, " + "negative, and non-integer quantities." + ), + }, + { + "type": "tool_use", + "id": "call_412671", + "name": "read_file", + "input": {"path": "shop/validation.py", "start_line": 1, "end_line": 40}, + "provider_specific_fields": {"signature": GEMINI_3_8_FLASH_THOUGHTS_OFF_THOUGHT_SIGNATURE}, + }, + ], + "usage": {"input_tokens": 354, "output_tokens": 385}, + }, +) + + +def with_thinking(case: TranslationTestCase, scenario: str, field: str, value: JsonValue) -> TranslationTestCase: + return replace( + case, + scenario=scenario, + litellm_request={**case.litellm_request, "thinking": {**ADAPTIVE_THINKING, field: value}}, + ) + + +THINKING_ASKS: Final = ( + with_thinking(GEMINI_3_8_FLASH_TEST_CASE, "display_summarized", "display", "summarized"), + with_thinking(GEMINI_3_8_FLASH_TEST_CASE, "display_updates", "display", "updates"), + with_thinking(GEMINI_3_8_FLASH_TEST_CASE, "display_int", "display", 7), + with_thinking(GEMINI_3_8_FLASH_TEST_CASE, "display_empty", "display", ""), + with_thinking(GEMINI_3_8_FLASH_TEST_CASE, "display_list", "display", ["omitted"]), + with_thinking(GEMINI_3_8_FLASH_TEST_CASE, "display_long", "display", LONG_SUMMARY), + with_thinking(GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE, "display_omitted", "display", "omitted"), + with_thinking(GEMINI_3_8_FLASH_TEST_CASE, "summary_auto", "summary", "auto"), + with_thinking(GEMINI_3_8_FLASH_TEST_CASE, "summary_empty", "summary", ""), + with_thinking(GEMINI_3_8_FLASH_TEST_CASE, "summary_long", "summary", LONG_SUMMARY), + with_thinking(GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE, "summary_int", "summary", 7), + with_thinking(GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE, "summary_list", "summary", ["auto"]), + with_thinking(GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE, "summary_null", "summary", None), +) + + +def test_gemini_keeps_thinking_and_billing_reasoning_without_thought_text_when_the_gateway_turns_it_off( + thoughts_off: Gateway, provider: SharedProvider +) -> None: + assert_translation(GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE, thoughts_off, provider) + + +@pytest.mark.parametrize("case", THINKING_ASKS, ids=tuple(case.id for case in THINKING_ASKS)) +def test_a_thinking_display_or_summary_ask_brings_thought_text_back_over_the_gateway_default( + case: TranslationTestCase, thoughts_off: Gateway, provider: SharedProvider +) -> None: + assert_translation(case, thoughts_off, provider) + + +def test_display_omitted_alone_keeps_thought_text_on_the_default_gateway( + gateway: Gateway, provider: SharedProvider +) -> None: + assert_translation( + with_thinking(GEMINI_3_8_FLASH_TEST_CASE, "display_omitted_default", "display", "omitted"), gateway, provider + ) + + +def test_a_streamed_turn_keeps_thinking_and_billing_reasoning_without_thought_text( + thoughts_off: Gateway, provider: SharedProvider +) -> None: + provider.expect( + Reply( + content_type="text/event-stream", + chunks=tuple(encode_sse_frame(frame) for frame in GEMINI_3_8_FLASH_THOUGHTS_OFF_STREAM_FRAMES), + ) + ) + response: Final = thoughts_off.request( + "POST", "/v1/messages", {**GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE.litellm_request, "stream": True} + ) + assert_provider_request( + provider.received(), + "/models/gemini-3.8-flash:streamGenerateContent?alt=sse", + GEMINI_3_8_FLASH_TEST_CASE.expected_provider_headers, + GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE.expected_provider_request, + ) + assert response.status_code == 200, response.text + events: Final = tuple(object_value(frame.data) for frame in parse_sse_frames(response.text)) + started: Final = tuple( + object_value(event["content_block"])["type"] for event in events if event.get("type") == "content_block_start" + ) + assert started == ("text", "tool_use") + deltas: Final = tuple( + object_value(event["delta"]) for event in events if event.get("type") == "content_block_delta" + ) + assert "".join(str(delta.get("text", "")) for delta in deltas) == ( + "I need to read `shop/validation.py` next to see how `validate_quantity` currently handles zero, negative, " + "and non-integer quantities." + ) + assert [object_value(event["usage"]) for event in events if event.get("type") == "message_delta"][-1] == { + "input_tokens": 354, + "output_tokens": 360, + } diff --git a/tests/integration/translation/responses/reasoning/__init__.py b/tests/integration/translation/responses/reasoning/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/integration/translation/responses/reasoning/test_responses_reasoning_gemini.py b/tests/integration/translation/responses/reasoning/test_responses_reasoning_gemini.py new file mode 100644 index 00000000000..984f7d2a8c4 --- /dev/null +++ b/tests/integration/translation/responses/reasoning/test_responses_reasoning_gemini.py @@ -0,0 +1,400 @@ +from collections.abc import Mapping +from dataclasses import replace +from typing import Final +from unittest.mock import ANY + +import pytest +from integration._support.client import Gateway, list_value, object_value +from integration._support.process import graceful_stop_seconds +from integration._support.provider import SharedProvider +from integration._support.wire import Reply +from integration.translation.case import SSEFrame, TranslationTestCase +from integration.translation.gemini_thoughts import LONG_SUMMARY, thoughts_off, with_thinking_config +from integration.translation.responses.bases.gemini import GEMINI_3_8_FLASH_TEST_CASE, GEMINI_RESPONSES_RESPONSE +from integration.translation.runner import ( + assert_provider_request, + assert_translation, + encode_sse_frame, + parse_sse_frames, +) +from pydantic import JsonValue + +__all__ = ("thoughts_off",) + +pytestmark = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) + +GEMINI_3_8_FLASH_THOUGHTS_OFF_THOUGHT_SIGNATURE: Final = ( + "EvQLCvELAWkUfRM8dwxvoWqyX5HLMaxcqWN9o4LPX6M6MZM/OD962KfMSvacFOxwOEb2B3rQM7QgoyT58YrItR4zpDSZJUxA" + "N4shekmtkSRpwrqE0rmH1v9AE8w5lR0Jewj8xmiiUXUrK0qvpyxhbVJRmIgdt9CqW+B4bJHYpp1VFAuHMk8eyBRoVNXakpsG" + "xjvHWGDomXbyJ7XuG1uHzvA5QtV2lgKUwRNvQqbFgbMhKN4dc8Imm2kwCLQONRN6wYIKJnYIi6ib+MhsFY+hZ6sPysMzVyrk" + "nGhc6pZwOKHbUjAraY/X98HO7UEL/q86LKUIvKTS4FfZ8keuyIQQBMdX68fdZPqurR4fI2dr5pFSDbaJWEVaE9XAY9wOXhOE" + "uC5SxY0U7nPU39eycNqAoAf+O7i/uHn1eCfLMhvAZyoH0z6vMfofnVn6L6alZtnRFQ3rEdJWEZHFmsLOk7G9SHlLAYKQjEdj" + "5aQHsKqBG1SFODPWSeKduI5w/kgFwgrw1ye1rnFTd+FDBs3DsQUlwh+jbPg8tr3KnCL5f2bxoGShUanPcGRAJ4b/IWGM3+wf" + "Whcx+3JY/OfuklkCnw86uRlcJiEjWd9XZ+jAby3ECZxBuXsHSE+Ynzpd0mp5g/S1orzqlmUJ/PlOcbDR/AsFduZvZwG0n4hq" + "hbXIIPszIpEyGDhkC//vxXsWdo5gLcsnEoUlLCESN3Sc12goKneId6u8pVWfzvHwRa0pK0LIWo4GcWZepZpEoYARlUZN6Dat" + "FNJ0nFW1sGaNb4CEYdSSp3m+GVobHpdz8qsGJxWrzzc6ZYmO7LbDJ25dJtyCl9QcfyzGbkFVDdAaTBEM23EdyZMFc2IIT24K" + "OVx5esCx+VjDZOmOzue/ZR0xrfaPXVa0CMuzGA8WwiphpWTbD8xZUnqEUZQNwKrjU91+SIhaBKcTVM8F1Wj0t9vlQkkMlzth" + "bJxIw+Ccu51gh/s8/awIDsBwC2Ax6eYIhUoWxUWYcphX7r3lBKhIuzv7FyO8tmaRkIOy6Eb1lXHB4fHAIOimZ/n4ClgK/KOZ" + "XVLosaGMKsmjYAdrxxPaNqs+say+pQ08bec/c+mqWeDbTOJ2JrCKM/Q+EcyD+82QqWlO68boeJyyt0fIvTCAdn57xN5Th7r1" + "WGhT9YHeSB1Y77P9FaJLQqAOjr+yogG7suiCP4QHg1arBlVzvaxAdnlBeXSv6pbQJuuvsFV9+SfwKHVXaBv68bcjSYKm8YZa" + "vfK57jweJUxvknCUpzrDRBdVZX28H60eQ3X6S0cV/696AO8NHoKuMv9zsCyu4WXWnpYEObtVD4Sn5hdSn1tnOkyUTLmN+bHV" + "sRGe5WTmrgnTKbrGDqpeqFpYLeMftkKWX+DHZj/LvZiNH1pHJid4nbEvcUKHciculasWiLZXDhEDKA1uzvSeG7CHjtHziqRi" + "b0dY7pQ0iEJUrjIwMVPnFZoyrywYZVPYgk+jna9FJypscuCz5XzdV+N2b2yrKryU6NhOQ5u/lzMNEFsyYxuXp6LZ5CVgdH1T" + "k+1rUf3sWPTO+CyDj+GzWvo5ljP7HihqyitafZJeKOSfaPfrJpBq+R8jnwDOVLCS1lnkhsof12PFoQHAE0YzGDtVscrstW4O" + "OaEdtNyNh30fhwwaZ/XGM4O9yJ8l7CgUvrgY/t8/f1hB/ZIc980HGs3G/bmVBYTxeZ18rqanrURtVuNtveyGief6eiEVIFjo" + "HWkAFXM+sn1zsJI6l4N9TPMbkoPxLCw2KaOaI/7mQjsH6VktpGuuqpHw3S5tiw3OGwh0lEwydKT0bk4hUhWqoBHqFpCrmKqA" + "g3z0kEHxQjQ1j8GmHusdPp4odrCWphw44m1VeWYjF4VjQ16Vik3SSrC8g3v9sHhtKM+3RL6z8zzukqgcQkXf7RXWvmFttckj" + "yYWM+vfQm/UvCzQ4dKvFHMqgmwuyEISPO9CD68/518wNpbQD1iB/RnxKwyYoD4u9YQ9TzXCzk2eP0NBB6pRmGg5nBo2SP0PB" + "BRZqlxnFX4XfQiS9pxNL" +) + +GEMINI_3_8_FLASH_THOUGHTS_OFF_RESPONSE: Final[Mapping[str, JsonValue]] = { + "candidates": [ + { + "content": { + "parts": [ + { + "text": "To see how zero, negative, and non-integer quantities are currently handled and rejected by the validation logic, the next file I need to inspect is `shop/validation.py`.", + }, + { + "functionCall": { + "name": "read_file", + "args": { + "path": "shop/validation.py", + "start_line": 1, + "end_line": 50, + }, + "id": "call_584229", + }, + "thoughtSignature": GEMINI_3_8_FLASH_THOUGHTS_OFF_THOUGHT_SIGNATURE, + }, + ], + "role": "model", + }, + "finishReason": "STOP", + "index": 0, + "finishMessage": "Model generated function call(s).", + }, + ], + "usageMetadata": { + "promptTokenCount": 354, + "candidatesTokenCount": 70, + "totalTokenCount": 790, + "promptTokensDetails": [ + { + "modality": "TEXT", + "tokenCount": 354, + }, + ], + "thoughtsTokenCount": 366, + "serviceTier": "standard", + }, + "modelVersion": "gemini-3.8-flash", + "responseId": "2y3RasgqkpyMxw-RxpeQBg", +} + +GEMINI_3_8_FLASH_THOUGHTS_OFF_STREAM_THOUGHT_SIGNATURE: Final = ( + "EvYSCvMSAWkUfROtzl1tHuKrIElCi7xsI6n+7cPuqsM/qZSYIpK3gONKCCo0E4moadzMQ3sN09ll3gRjbcBg6rKSck+JTBhd" + "mrLsP/sqfxiYfMx4lWzDeq/Pw/IPhT2g5/2EvSU4FDsUczorTsD55gKvgRT8v05YTFMGF/np8r3VT52Q2yAzwRK+31qn8pYS" + "c25if5coTZJrvmJZDlZ9/64aibghRps9zN07yRVXpUjtLyKMaTDH0loGgiQn0I7Cxfhw9oUyZels7Un7pB7omA4RKW5xp5rx" + "Y+WgSeAEyyxNcPIkRvaxHaY1p2QaR3m17a0L7joJYNB+gNg7R/4Mfwpr8WuSOyLfBaKTuC5IiYWLaRm0ibASYzd59pPUyArt" + "mIx7pPfWxAr8/FT8TVQTSSeHXF88x/xJ2WSTSzQ3+FPw6Wm/LpfeGwQPX/g1+8f+KVcj1qAwT/tUobqOh7DdWl68CnNFReX4" + "SwBqRjTzZ276P7l15fhWmnQkiWXmLVsufJi5VXidEVrAsuEyVe7JmpExi3Y+0pn82D9x7lyv/TrOCVQgKwC0Zoq9NfuQxWa0" + "RvqmAv6toMfI8owTyzVoItnoGdhcYamwr7AyECKTuqnzMuylzOx5ttZnRc12/Ey2GQR3UxhyxMNLk01/UDicLnEXFl8Zn56i" + "rj51atVum7mJc48qekcKXQ5Lc3/t277v1d+eLmuiHebshpQwg+qQPHnnfky0WL98w343XDb29wwK3+veqQSjrffv5rQY9Tvq" + "3PMuN1pjU7Z3rL6fZmt5UrhqQQnTf4dA7X7n/aLkpA8SpY6ICDDr2x1z0yW5kMnHfxxk+N+j5eWwoTJuxxy94rLaYj3CKT57" + "XhMthQAoF1DGNpGxhbfC/tVuFE9FkxEBiYXqDz5fy2AU+S2zatT7/9mkFaHLY6hkOCS48X9WBVobxsQZZUnkbxplCzg4Huin" + "gPnl1bUDEprLHfTYgZHcnn3FgWFPwzxNBebrc4xjqXciMgReu4kOLP18JZbIJSjhP/RZAnQJEEOI8OrVGZlmiC7/JJx+8PTX" + "XYMTKtiVU/R3gOhOeVNN9WXruUoQzuy+JfJGg8ZQ2ni0Ugkx0sLO5rDsTOZ67bA+B99FQnQcBLyb40EPTOzrLrcM9wzXL+YZ" + "JUwNtdaQGxh/aUeKex/8of7GUH9ZsCZW0YN5xuEs8e2Q+vi2Msluckvki/wcyx7oOxmjw8bEbUkPBFq/ZEOcQqtetPLBzNPt" + "smJ4EfOU/kwBBLxplt6JrF+0ILeWJuH7OGpvpseXEosLpJ9HV4oIHL9co8EWGH/oAUzcuvhymR8FRziVE2IgVkreryIfHAs/" + "umsJgk1dXZk4A/R7KtjHKau0yxLAYJh6jWLvLsLv5N44VfKCQtyWFdUOSGbqn6E4YuOYVc+x+TGQcSMniobTekwm+UtqzVH8" + "mmtd8nwzwVkDgqZEAHrrnJOALLPeIbPuQ0/zjuUJmQovnnSqUyyCxQDpaxrlLD3iIZp9/9hPTY+lYHlam7VfChpQ7dPQ075t" + "ECiPA5TgMzh1qNQBF+rgKtrcrPo1ZDFqlseAybmSKqnCdvD5vdfdeVL8cmQIQao9qHDVh/aR0uwOsVKKzavPIYH2mhomg/cz" + "XB1nJG2zXB99g08Bd14c/zHHPNeqkqmwURCTxzPh8aeoRPcC8uL/eyawL66FfyUo40og9uridBDA4qx4ZMI+P1dBKFnMPM7o" + "Blpo+Sg9m6XZNPNVbDqRKKmX9OWCng7/PAKX+wa8EzRetNbyh73eMviMZaiN9YAyoKiS+nR5FqrPgNcFMqa8dHDEyY6O3Vuh" + "5Er5R3iGNACXMQvLRDVy/1UmFY1mzwew9lrDcUkZ0MNKBrbtDC87XEFGMejqvoJgF5Yl+q1mJE6gMBDTo6UEP1SOctxHyX23" + "mvqhCK9wOvK91UGnNyd+8ZQtQPP/L99d2y7hHmE6VEA/nUa2qkkO99vEQmJV3ilax+wZknABwk5sFdh7vWLwhEvxkdGYUh6c" + "HnA70VX5Kh0twJPKwE7gNASf/rvE5kYdrei5RtZNKMpVulXMPVho3eRBW6Qf6Ka0JjDZF5Dr8IXXsjP2LFTmnLNcnxu0FXx8" + "4xRuw1+yN3cP5R8LFlSGJmqPf5kFwAwpzT5FAXETiBaj7lg3WyG+OaDpab2NG+BvHbdfrDAuFKSsTDZ3Vt58fj6fuhS8tFd4" + "YzWL6X30pj+r7QRAefk+2CQDKM1TnbAJbyBcG2DVfeWpgmNjtzoV+RXNykIFIy6rSeOnUBAJCXwbAe7Scp5w/lx03a5aooRU" + "CP7x7CAmAU+YuGerG7A/3bbWDqN4EHGawkloK9UXmlOVX7uzLy0q3fJZb/6pSvPX+DR2rjFq0CSKnNFyBXYQRHx+DekAHWJS" + "VGUw7RLAPd4DE2HRiYTss7vtXalkP0hQ/vnt8lS9bjyBCi8ly+4FjFpo/SeKx0TKFulbFp1GPpsNERT+7kEkTsNZQHXWXPRB" + "PH0qmpFWZjI98fzWIGXGF376X0bcg6g67C/K/+YtH+QIU82ScAM0ozzQ1uzb0ci6wnM4JjNKA3ctVX6zuPR20Tb7VLlq5rwR" + "coOKjGQmBGbUiW5aKAAsMnIDhxf2lPCcwwb51C0Zb6hfOG+ka7LLcstgq4fmXHQ11BbTFu9ctAhy0E5V3aYSJAilb6VXZdD4" + "kSl3I0D2Xc/pTLPAzsxg6mEHL7J/DD1lRvEtSh4U34RCwxMk+pQbfEo6y/lkIVhXuE8p7d4mhPtsDJkls/3vnl+rR5ycebjH" + "g/zC01CjHhbVglb+HVEC5FMTrBNy/bweTE6c7qNHV6EWRN64nqvs4IW/jEVoaI2qt9WAJZ1br21rCE01iQagQUs1/ZH2KpQB" + "jTe3ju/BotwiE4TdqHcDL/VP62yq9KLeiBNXY2ctb3NwJzzfLrM2KHv2b4YmWuZ/VLABBfh9r3OzSmnFBvmWsZjZOHHEQXMI" + "hI+WDpRGB54CWmnWs6YzQK70IKrZUtq4BfPQJGf8OqO/ZeCJA4kinQbCaTTqQOJLsONKbPTqm8ImXz/QiML+KFvY3Z3ReW3T" + "NsJ3CGW7NRY1cfBIkhHuKHxPK1Ua5H3b7yP+h8OKrbqyjYPw2kE56TOH75HnIp/Ib6w+7uZpMVIbM4fYCRWlvQ7cUtaRq4fk" + "qqgHykE8bJYoU9eJFmPJNj4gt051sXJoslr2C6mvHYdSZgKxNa1ksd95cPKS3yURYQ==" +) + +GEMINI_3_8_FLASH_THOUGHTS_OFF_STREAM_FRAMES: Final = ( + SSEFrame( + data={ + "candidates": [ + { + "content": { + "parts": [ + { + "text": "To understand how zero, negative, and non-integer quantities are checked during validation, I need to inspect", + }, + ], + "role": "model", + }, + "index": 0, + }, + ], + "usageMetadata": { + "promptTokenCount": 354, + "candidatesTokenCount": 21, + "totalTokenCount": 847, + "promptTokensDetails": [ + { + "modality": "TEXT", + "tokenCount": 354, + }, + ], + "thoughtsTokenCount": 472, + "serviceTier": "standard", + }, + "modelVersion": "gemini-3.8-flash", + "responseId": "3S3Rat6pNonj_uMPv9i78Qw", + } + ), + SSEFrame( + data={ + "candidates": [ + { + "content": { + "parts": [ + { + "text": " `shop/validation.py` next.", + }, + ], + "role": "model", + }, + "index": 0, + }, + ], + "usageMetadata": { + "promptTokenCount": 354, + "candidatesTokenCount": 30, + "totalTokenCount": 856, + "promptTokensDetails": [ + { + "modality": "TEXT", + "tokenCount": 354, + }, + ], + "thoughtsTokenCount": 472, + "serviceTier": "standard", + }, + "modelVersion": "gemini-3.8-flash", + "responseId": "3S3Rat6pNonj_uMPv9i78Qw", + } + ), + SSEFrame( + data={ + "candidates": [ + { + "content": { + "parts": [ + { + "functionCall": { + "name": "read_file", + "args": { + "end_line": 50, + "path": "shop/validation.py", + "start_line": 1, + }, + "id": "call_5018", + }, + "thoughtSignature": GEMINI_3_8_FLASH_THOUGHTS_OFF_STREAM_THOUGHT_SIGNATURE, + }, + ], + "role": "model", + }, + "index": 0, + }, + ], + "usageMetadata": { + "promptTokenCount": 354, + "candidatesTokenCount": 63, + "totalTokenCount": 889, + "promptTokensDetails": [ + { + "modality": "TEXT", + "tokenCount": 354, + }, + ], + "thoughtsTokenCount": 472, + "serviceTier": "standard", + }, + "modelVersion": "gemini-3.8-flash", + "responseId": "3S3Rat6pNonj_uMPv9i78Qw", + } + ), + SSEFrame( + data={ + "candidates": [ + { + "content": { + "parts": [ + { + "text": "", + }, + ], + "role": "model", + }, + "finishReason": "STOP", + "index": 0, + }, + ], + "usageMetadata": { + "promptTokenCount": 354, + "candidatesTokenCount": 63, + "totalTokenCount": 889, + "promptTokensDetails": [ + { + "modality": "TEXT", + "tokenCount": 354, + }, + ], + "thoughtsTokenCount": 472, + "serviceTier": "standard", + }, + "modelVersion": "gemini-3.8-flash", + "responseId": "3S3Rat6pNonj_uMPv9i78Qw", + } + ), +) + +GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE: Final = replace( + GEMINI_3_8_FLASH_TEST_CASE, + scenario="thoughts_off", + expected_provider_request=with_thinking_config( + GEMINI_3_8_FLASH_TEST_CASE.expected_provider_request, {"thinkingLevel": "high", "includeThoughts": False} + ), + mock_provider_response=GEMINI_3_8_FLASH_THOUGHTS_OFF_RESPONSE, + expected_litellm_response={ + **GEMINI_RESPONSES_RESPONSE, + "model": "gemini/gemini-3.8-flash", + "output": [ + { + "type": "message", + "id": ANY, + "status": "completed", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": ( + "To see how zero, negative, and non-integer quantities are currently handled and rejected " + "by the validation logic, the next file I need to inspect is `shop/validation.py`." + ), + "annotations": [], + }, + ], + "phase": None, + }, + { + "arguments": '{"path": "shop/validation.py", "start_line": 1, "end_line": 50}', + "call_id": "call_584229__thought__" + GEMINI_3_8_FLASH_THOUGHTS_OFF_THOUGHT_SIGNATURE, + "name": "read_file", + "type": "function_call", + "id": "fc_call_584229__thought__" + GEMINI_3_8_FLASH_THOUGHTS_OFF_THOUGHT_SIGNATURE, + "namespace": None, + "status": "completed", + "provider_specific_fields": {"thought_signature": GEMINI_3_8_FLASH_THOUGHTS_OFF_THOUGHT_SIGNATURE}, + }, + ], + "usage": { + "input_tokens": 354, + "input_tokens_details": { + "audio_tokens": None, + "cached_tokens": 0, + "cached_tokens_details": None, + "image_tokens": None, + "text_tokens": 354, + "video_tokens": None, + }, + "output_tokens": 436, + "output_tokens_details": {"audio_tokens": None, "reasoning_tokens": 366, "text_tokens": 70}, + "total_tokens": 790, + "cost": None, + }, + }, +) + + +def with_summary(case: TranslationTestCase, scenario: str, summary: JsonValue) -> TranslationTestCase: + reasoning: Final[Mapping[str, JsonValue]] = {"effort": "high", "summary": summary} + return replace( + case, + scenario=scenario, + litellm_request={**case.litellm_request, "reasoning": dict(reasoning)}, + expected_litellm_response={**case.expected_litellm_response, "reasoning": dict(reasoning)}, + ) + + +SUMMARY_ASKS: Final = ( + with_summary(GEMINI_3_8_FLASH_TEST_CASE, "summary_auto", "auto"), + with_summary(GEMINI_3_8_FLASH_TEST_CASE, "summary_concise", "concise"), + with_summary(GEMINI_3_8_FLASH_TEST_CASE, "summary_empty", ""), + with_summary(GEMINI_3_8_FLASH_TEST_CASE, "summary_long", LONG_SUMMARY), + with_summary(GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE, "summary_int", 7), + with_summary(GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE, "summary_list", ["auto"]), + with_summary(GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE, "summary_null", None), +) + + +def test_gemini_keeps_thinking_and_billing_reasoning_without_thought_text_when_the_gateway_turns_it_off( + thoughts_off: Gateway, provider: SharedProvider +) -> None: + assert_translation(GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE, thoughts_off, provider) + + +@pytest.mark.parametrize("case", SUMMARY_ASKS, ids=tuple(case.id for case in SUMMARY_ASKS)) +def test_a_reasoning_summary_ask_brings_thought_text_back_over_the_gateway_default( + case: TranslationTestCase, thoughts_off: Gateway, provider: SharedProvider +) -> None: + assert_translation(case, thoughts_off, provider) + + +def test_a_streamed_turn_keeps_thinking_and_billing_reasoning_without_thought_text( + thoughts_off: Gateway, provider: SharedProvider +) -> None: + provider.expect( + Reply( + content_type="text/event-stream", + chunks=tuple(encode_sse_frame(frame) for frame in GEMINI_3_8_FLASH_THOUGHTS_OFF_STREAM_FRAMES), + ) + ) + response: Final = thoughts_off.request( + "POST", "/v1/responses", {**GEMINI_3_8_FLASH_TEST_CASE.litellm_request, "stream": True} + ) + assert_provider_request( + provider.received(), + "/models/gemini-3.8-flash:streamGenerateContent?alt=sse", + GEMINI_3_8_FLASH_TEST_CASE.expected_provider_headers, + GEMINI_3_8_FLASH_THOUGHTS_OFF_TEST_CASE.expected_provider_request, + ) + assert response.status_code == 200, response.text + events: Final = tuple( + object_value(frame.data) for frame in parse_sse_frames(response.text) if frame.data != "[DONE]" + ) + assert [event["type"] for event in events if str(event["type"]).startswith("response.reasoning")] == [] + assert "".join(str(event["delta"]) for event in events if event["type"] == "response.output_text.delta") == ( + "To understand how zero, negative, and non-integer quantities are checked during validation, I need to " + "inspect `shop/validation.py` next." + ) + completed: Final = object_value( + next(event["response"] for event in events if event["type"] == "response.completed") + ) + assert [object_value(item)["type"] for item in list_value(completed["output"])] == ["message", "function_call"] + assert object_value(completed["usage"])["output_tokens_details"] == { + "reasoning_tokens": 472, + "text_tokens": 63, + } diff --git a/tests/unit/caching/test_caching.py b/tests/unit/caching/test_caching.py index b84529de0d6..91083921130 100644 --- a/tests/unit/caching/test_caching.py +++ b/tests/unit/caching/test_caching.py @@ -291,6 +291,37 @@ def test_exact_cache_key_includes_anthropic_messages_params(anthropic_param): assert baseline != cache.get_cache_key(model="claude-sonnet-4-5", messages=messages, **anthropic_param) +@pytest.mark.parametrize( + "summary_ask", + [ + {"reasoning_summary": "auto"}, + {"reasoningSummary": "auto"}, + {"extra_body": {"reasoning_summary": "auto"}}, + {"extra_body": {"reasoningSummary": "auto"}}, + ], + ids=["top_level", "top_level_camel", "extra_body", "extra_body_camel"], +) +def test_exact_cache_key_includes_the_reasoning_summary_alias(summary_ask: dict[str, object]) -> None: + """The AI SDK summary alias decides whether thought summaries come back, wherever the request + carries it, so a reply cached without it must not answer a request that asks for one.""" + cache: Final = Cache(type=LiteLLMCacheType.LOCAL) + messages: Final = [{"role": "user", "content": "how much is the ball?"}] + no_ask: Final = cache.get_cache_key(model="gemini/gemini-3.8-flash", messages=messages, reasoning_effort="medium") + ask: Final = cache.get_cache_key( + model="gemini/gemini-3.8-flash", messages=messages, reasoning_effort="medium", **summary_ask + ) + assert no_ask != ask + assert ask == cache.get_cache_key( + model="gemini/gemini-3.8-flash", messages=messages, reasoning_effort="medium", **summary_ask + ) + assert ask != cache.get_cache_key( + model="gemini/gemini-3.8-flash", + messages=messages, + reasoning_effort="medium", + reasoning_summary="detailed", + ) + + @pytest.mark.asyncio async def test_embedding_cache_skips_write_when_one_input_yields_many_embeddings(monkeypatch): """A cross-encoder behind /embeddings returns one score per document for a single @@ -561,10 +592,7 @@ def test_caching_v2(): # test in memory cache litellm.cache = None # disable cache litellm.success_callback = [] litellm._async_success_callback = [] - if ( - response2["choices"][0]["message"]["content"] - != response1["choices"][0]["message"]["content"] - ): + if response2["choices"][0]["message"]["content"] != response1["choices"][0]["message"]["content"]: print(f"response1: {response1}") print(f"response2: {response2}") pytest.fail(f"Error occurred:") @@ -595,10 +623,7 @@ def test_caching_with_ttl(): litellm.cache = None # disable cache litellm.success_callback = [] litellm._async_success_callback = [] - assert ( - response2["choices"][0]["message"]["content"] - != response1["choices"][0]["message"]["content"] - ) + assert response2["choices"][0]["message"]["content"] != response1["choices"][0]["message"]["content"] except Exception as e: print(f"error occurred: {traceback.format_exc()}") pytest.fail(f"Error occurred: {e}") @@ -632,9 +657,7 @@ def test_caching_with_default_ttl(): def test_caching_with_models_v2(): - messages = [ - {"role": "user", "content": "who is ishaan CTO of litellm from litellm 2023"} - ] + messages = [{"role": "user", "content": "who is ishaan CTO of litellm from litellm 2023"}] litellm.cache = Cache() print("test2 for caching") litellm.set_verbose = True @@ -657,18 +680,12 @@ def test_caching_with_models_v2(): litellm.cache = None litellm.success_callback = [] litellm._async_success_callback = [] - if ( - response3["choices"][0]["message"]["content"] - == response2["choices"][0]["message"]["content"] - ): + if response3["choices"][0]["message"]["content"] == response2["choices"][0]["message"]["content"]: # if models are different, it should not return cached response print(f"response2: {response2}") print(f"response3: {response3}") pytest.fail(f"Error occurred:") - if ( - response1["choices"][0]["message"]["content"] - != response2["choices"][0]["message"]["content"] - ): + if response1["choices"][0]["message"]["content"] != response2["choices"][0]["message"]["content"]: print(f"response1: {response1}") print(f"response2: {response2}") pytest.fail(f"Error occurred:") @@ -708,9 +725,7 @@ def test_get_cache_key(): cache_key = cache_instance.get_cache_key( **{ "model": "gpt-3.5-turbo", - "messages": [ - {"role": "user", "content": "write a one sentence poem about: 7510"} - ], + "messages": [{"role": "user", "content": "write a one sentence poem about: 7510"}], "max_tokens": 40, "temperature": 0.2, "stream": True, @@ -721,9 +736,7 @@ def test_get_cache_key(): cache_key_2 = cache_instance.get_cache_key( **{ "model": "gpt-3.5-turbo", - "messages": [ - {"role": "user", "content": "write a one sentence poem about: 7510"} - ], + "messages": [{"role": "user", "content": "write a one sentence poem about: 7510"}], "max_tokens": 40, "temperature": 0.2, "stream": True, @@ -736,9 +749,9 @@ def test_get_cache_key(): # Hexadecimal representation of the hash hash_hex = hash_object.hexdigest() assert cache_key == hash_hex - assert ( - cache_key_2 == hash_hex - ), f"{cache_key} != {cache_key_2}. The same kwargs should have the same cache key across runs" + assert cache_key_2 == hash_hex, ( + f"{cache_key} != {cache_key_2}. The same kwargs should have the same cache key across runs" + ) embedding_cache_key = cache_instance.get_cache_key( **{ @@ -756,15 +769,13 @@ def test_get_cache_key(): print(embedding_cache_key) - embedding_cache_key_str = ( - "model: azure/text-embedding-ada-002input: ['hi who is ishaan']" - ) + embedding_cache_key_str = "model: azure/text-embedding-ada-002input: ['hi who is ishaan']" hash_object = hashlib.sha256(embedding_cache_key_str.encode()) # Hexadecimal representation of the hash hash_hex = hash_object.hexdigest() - assert ( - embedding_cache_key == hash_hex - ), f"{embedding_cache_key} != 'model: azure/text-embedding-ada-002input: ['hi who is ishaan']'. The same kwargs should have the same cache key across runs" + assert embedding_cache_key == hash_hex, ( + f"{embedding_cache_key} != 'model: azure/text-embedding-ada-002input: ['hi who is ishaan']'. The same kwargs should have the same cache key across runs" + ) # Proxy - embedding cache, test if embedding key, gets model_group and not model embedding_cache_key_2 = cache_instance.get_cache_key( @@ -817,9 +828,7 @@ def test_get_cache_key(): ) print(embedding_cache_key_2) - embedding_cache_key_str_2 = ( - "model: EMBEDDING_MODEL_GROUPinput: ['hi who is ishaan']" - ) + embedding_cache_key_str_2 = "model: EMBEDDING_MODEL_GROUPinput: ['hi who is ishaan']" hash_object = hashlib.sha256(embedding_cache_key_str_2.encode()) # Hexadecimal representation of the hash hash_hex = hash_object.hexdigest() @@ -852,9 +861,7 @@ def test_redis_caching_multiple_namespaces(): # Mock the Redis client creation from the _redis module with ( patch("litellm._redis.get_redis_client") as mock_get_redis_client, - patch( - "litellm._redis.get_redis_connection_pool" - ) as mock_get_redis_connection_pool, + patch("litellm._redis.get_redis_connection_pool") as mock_get_redis_connection_pool, ): # Create a mock Redis client that simulates real Redis behavior mock_redis_client = MagicMock() @@ -927,18 +934,10 @@ def test_redis_caching_multiple_namespaces(): mock_response="Response without namespace", ) - print( - f"Response 1 type: {type(response_1)} - ID: {getattr(response_1, 'id', 'N/A')}" - ) - print( - f"Response 2 type: {type(response_2)} - ID: {getattr(response_2, 'id', 'N/A')}" - ) - print( - f"Response 3 type: {type(response_3)} - Cache hit: {isinstance(response_3, str)}" - ) - print( - f"Response 4 type: {type(response_4)} - ID: {getattr(response_4, 'id', 'N/A')}" - ) + print(f"Response 1 type: {type(response_1)} - ID: {getattr(response_1, 'id', 'N/A')}") + print(f"Response 2 type: {type(response_2)} - ID: {getattr(response_2, 'id', 'N/A')}") + print(f"Response 3 type: {type(response_3)} - Cache hit: {isinstance(response_3, str)}") + print(f"Response 4 type: {type(response_4)} - ID: {getattr(response_4, 'id', 'N/A')}") print(f"Redis storage keys: {list(redis_storage.keys())}") @@ -947,10 +946,7 @@ def test_redis_caching_multiple_namespaces(): namespace_1_keys = [k for k in cache_keys if k.startswith(f"{namespace_1}:")] namespace_2_keys = [k for k in cache_keys if k.startswith(f"{namespace_2}:")] no_namespace_keys = [ - k - for k in cache_keys - if not k.startswith(f"{namespace_1}:") - and not k.startswith(f"{namespace_2}:") + k for k in cache_keys if not k.startswith(f"{namespace_1}:") and not k.startswith(f"{namespace_2}:") ] print(f"Namespace 1 keys: {namespace_1_keys}") @@ -963,9 +959,7 @@ def test_redis_caching_multiple_namespaces(): assert len(no_namespace_keys) > 0, "Should have cache keys for no namespace" # The main test: response 3 should be a cache hit (string) because it uses same namespace as response 1 - assert isinstance( - response_3, str - ), "Response 3 should be a cache hit (string) for same namespace" + assert isinstance(response_3, str), "Response 3 should be a cache hit (string) for same namespace" # response 1 & 2 should be ModelResponse objects (cache misses) assert hasattr(response_1, "id"), "Response 1 should be a ModelResponse object" @@ -973,14 +967,14 @@ def test_redis_caching_multiple_namespaces(): assert hasattr(response_4, "id"), "Response 4 should be a ModelResponse object" # response 1 & 2 should have different IDs (different namespaces) - assert ( - response_1.id != response_2.id - ), f"Expected different response ID for different namespace. Got {response_1.id} and {response_2.id}" + assert response_1.id != response_2.id, ( + f"Expected different response ID for different namespace. Got {response_1.id} and {response_2.id}" + ) # response 1 & 4 should have different IDs (different namespaces) - assert ( - response_1.id != response_4.id - ), f"Expected different response ID for no namespace vs namespaced. Got {response_1.id} and {response_4.id}" + assert response_1.id != response_4.id, ( + f"Expected different response ID for no namespace vs namespaced. Got {response_1.id} and {response_4.id}" + ) _TOOL_TURN_ITEM: Final = {"role": "user", "content": "hi"} @@ -1305,12 +1299,8 @@ async def test_embedding_caching_individual_items(monkeypatch: pytest.MonkeyPatc monkeypatch.setattr(litellm, "cache", Cache(type=LiteLLMCacheType.LOCAL)) item: Final = f"individual embedding {uuid.uuid4().hex}" - first: Final = await aembedding( - model="text-embedding-3-small", input=item, caching=True, mock_response="0.1,0.2" - ) - second: Final = await aembedding( - model="text-embedding-3-small", input=item, caching=True, mock_response="0.3,0.4" - ) + first: Final = await aembedding(model="text-embedding-3-small", input=item, caching=True, mock_response="0.1,0.2") + second: Final = await aembedding(model="text-embedding-3-small", input=item, caching=True, mock_response="0.3,0.4") assert second.data[0].embedding == first.data[0].embedding assert second._hidden_params["cache_hit"] is True @@ -1448,9 +1438,7 @@ async def test_caching_with_cache_controls(monkeypatch: pytest.MonkeyPatch) -> N @pytest.mark.asyncio -async def test_embedding_caching_redis_ttl( - monkeypatch: pytest.MonkeyPatch, redis_client_without_network: None -) -> None: +async def test_embedding_caching_redis_ttl(monkeypatch: pytest.MonkeyPatch, redis_client_without_network: None) -> None: mock_pipeline: Final = AsyncMock() mock_set: Final = AsyncMock() mock_pipeline.__aenter__.return_value.set = mock_set @@ -1506,18 +1494,14 @@ async def test_logging_turn_off_message_logging_streaming( assert cached_response["choices"][0]["message"]["content"] == "cached" -def test_redis_caching_default_ttl( - monkeypatch: pytest.MonkeyPatch, redis_client_without_network: None -) -> None: +def test_redis_caching_default_ttl(monkeypatch: pytest.MonkeyPatch, redis_client_without_network: None) -> None: monkeypatch.setattr(litellm, "default_redis_ttl", 120) cache: Final = RedisCache(host="127.0.0.1", port="6379") assert cache.default_ttl == 120 -def test_redis_caching_llm_caching_ttl( - monkeypatch: pytest.MonkeyPatch, redis_client_without_network: None -) -> None: +def test_redis_caching_llm_caching_ttl(monkeypatch: pytest.MonkeyPatch, redis_client_without_network: None) -> None: monkeypatch.setattr(litellm, "default_redis_ttl", 120) cache: Final = RedisCache(host="127.0.0.1", port="6379") set_cache: Final = MagicMock() @@ -1529,9 +1513,7 @@ def test_redis_caching_llm_caching_ttl( @pytest.mark.asyncio -async def test_redis_caching_ttl_pipeline( - monkeypatch: pytest.MonkeyPatch, redis_client_without_network: None -) -> None: +async def test_redis_caching_ttl_pipeline(monkeypatch: pytest.MonkeyPatch, redis_client_without_network: None) -> None: monkeypatch.setattr(litellm, "default_redis_ttl", 120) cache: Final = RedisCache(host="127.0.0.1", port="6379") pipeline: Final = AsyncMock() @@ -1551,9 +1533,7 @@ async def test_redis_caching_ttl_pipeline( @pytest.mark.asyncio -async def test_redis_caching_ttl_sadd( - monkeypatch: pytest.MonkeyPatch, redis_client_without_network: None -) -> None: +async def test_redis_caching_ttl_sadd(monkeypatch: pytest.MonkeyPatch, redis_client_without_network: None) -> None: monkeypatch.setattr(litellm, "default_redis_ttl", 120) cache: Final = RedisCache(host="127.0.0.1", port="6379") redis_client: Final = AsyncMock() diff --git a/tests/unit/llms/anthropic/pass_through/adapters/test_handler_reasoning_effort_normalization.py b/tests/unit/llms/anthropic/pass_through/adapters/test_handler_reasoning_effort_normalization.py index 0dfe4a93649..4b96d2e7991 100644 --- a/tests/unit/llms/anthropic/pass_through/adapters/test_handler_reasoning_effort_normalization.py +++ b/tests/unit/llms/anthropic/pass_through/adapters/test_handler_reasoning_effort_normalization.py @@ -6,13 +6,17 @@ regression they guard is the one a caller sees: a tier the proxy advertises has leaves the adapter, in the shape the target expects. """ +from collections.abc import Mapping from typing import Final import pytest +import litellm from litellm.llms.anthropic.pass_through.adapters.handler import ( LiteLLMMessagesToCompletionTransformationHandler, ) +from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig +from litellm.utils import reasoning_effort_with_summary_alias MESSAGES = [{"role": "user", "content": "hello"}] @@ -38,14 +42,14 @@ def _reasoning_effort_sent(model: str, provider: str, reasoning_effort: object) return completion_kwargs.get("reasoning_effort") -def _reasoning_effort_sent_for_thinking( +def _completion_kwargs_for_thinking( model: str, provider: str | None, thinking: dict[str, object], *, tools: list[dict[str, object]] | None = None, api_base: str | None = None, -) -> object: +) -> Mapping[str, object]: extra_kwargs: Final = { key: value for key, value in (("custom_llm_provider", provider), ("api_base", api_base)) if value is not None } @@ -66,7 +70,20 @@ def _reasoning_effort_sent_for_thinking( output_format=None, extra_kwargs=extra_kwargs, ) - return completion_kwargs.get("reasoning_effort") + return completion_kwargs + + +def _reasoning_effort_sent_for_thinking( + model: str, + provider: str | None, + thinking: dict[str, object], + *, + tools: list[dict[str, object]] | None = None, + api_base: str | None = None, +) -> object: + return _completion_kwargs_for_thinking(model, provider, thinking, tools=tools, api_base=api_base).get( + "reasoning_effort" + ) SUMMARIZED_THINKING = {"type": "enabled", "budget_tokens": 4096, "summary": "auto"} @@ -78,6 +95,66 @@ MULTIPLY_TOOL = { } +class TestAThinkingSummaryAskReachesGemini: + """Gemini maps a summary ask to ``includeThoughts``, so ``gemini_include_thoughts_default`` only + decides for a ``/v1/messages`` request that asks for nothing, the same as on chat completions.""" + + @pytest.mark.parametrize( + "model, provider", + [("gemini/gemini-3.1-pro-preview", "gemini"), ("vertex_ai/gemini-2.5-flash", "vertex_ai")], + ) + @pytest.mark.parametrize( + "thinking, sends_thoughts", + [ + (SUMMARIZED_THINKING, True), + ({**PLAIN_THINKING, "display": "summarized"}, True), + ({**PLAIN_THINKING, "display": "updates"}, True), + ({**PLAIN_THINKING, "display": "omitted"}, False), + (PLAIN_THINKING, False), + ({"type": "disabled", "display": "summarized"}, False), + ], + ) + def test_the_request_ask_beats_the_gateway_default( + self, + local_model_cost_map: None, + monkeypatch: pytest.MonkeyPatch, + model: str, + provider: str, + thinking: dict[str, object], + sends_thoughts: bool, + ) -> None: + monkeypatch.setattr(litellm, "gemini_include_thoughts_default", False) + + sent: Final = _completion_kwargs_for_thinking(model, provider, thinking) + gemini: Final = VertexGeminiConfig() + reasoning_effort: Final = reasoning_effort_with_summary_alias( + sent.get("reasoning_effort"), sent, accepts_summary=gemini.accepts_reasoning_summary_in_reasoning_effort + ) + mapped: Final = gemini.map_openai_params( + {"reasoning_effort": reasoning_effort}, {}, model.split("/", 1)[1], True + ) + + assert mapped["thinkingConfig"]["includeThoughts"] is sends_thoughts + assert not isinstance(sent.get("reasoning_effort"), dict) + + def test_auto_summary_leaves_the_gateway_default_in_charge( + self, local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr(litellm, "gemini_include_thoughts_default", False) + monkeypatch.setenv("LITELLM_REASONING_AUTO_SUMMARY", "true") + + sent: Final = _reasoning_effort_sent_for_thinking("gemini/gemini-3.1-pro-preview", "gemini", PLAIN_THINKING) + + assert sent == "high" + + def test_a_display_ask_leaves_other_chat_targets_on_the_plain_tier(self, local_model_cost_map: None) -> None: + sent: Final = _reasoning_effort_sent_for_thinking( + "databricks/databricks-qwen35-122b-a10b", "databricks", {**PLAIN_THINKING, "display": "summarized"} + ) + + assert sent == "high" + + class TestTheSummaryWrappingOnlyRidesTheResponsesBridge: """Only the Responses API takes ``reasoning_effort`` as a dict. Databricks answered the wrapped ``{"effort", "summary"}`` with ``field 'reasoning_effort' expects input with json type 'string' @@ -100,7 +177,9 @@ class TestTheSummaryWrappingOnlyRidesTheResponsesBridge: ) -> None: monkeypatch.setenv("LITELLM_REASONING_AUTO_SUMMARY", "true") - sent = _reasoning_effort_sent_for_thinking("databricks/databricks-qwen35-122b-a10b", "databricks", PLAIN_THINKING) + sent: Final = _reasoning_effort_sent_for_thinking( + "databricks/databricks-qwen35-122b-a10b", "databricks", PLAIN_THINKING + ) assert sent == "high" @@ -113,7 +192,7 @@ class TestTheSummaryWrappingOnlyRidesTheResponsesBridge: ], ) def test_a_bridged_target_keeps_the_summary(self, local_model_cost_map: None, model: str, provider: str) -> None: - sent = _reasoning_effort_sent_for_thinking(model, provider, SUMMARIZED_THINKING) + sent: Final = _reasoning_effort_sent_for_thinking(model, provider, SUMMARIZED_THINKING) assert sent == {"effort": "high", "summary": "auto"} @@ -122,7 +201,7 @@ class TestTheSummaryWrappingOnlyRidesTheResponsesBridge: ) -> None: monkeypatch.setenv("LITELLM_REASONING_AUTO_SUMMARY", "true") - sent = _reasoning_effort_sent_for_thinking("azure/responses/gpt-5-mini", "azure", PLAIN_THINKING) + sent: Final = _reasoning_effort_sent_for_thinking("azure/responses/gpt-5-mini", "azure", PLAIN_THINKING) assert sent == {"effort": "high", "summary": "detailed"} @@ -139,21 +218,21 @@ class TestTheSummaryWrappingOnlyRidesTheResponsesBridge: """``completion()`` keeps a gpt-5.5 deployment with function tools on Foundry's chat route when its ``api_base`` is a Foundry OpenAI host, and bridges it to Responses when the base makes it an Azure OpenAI deployment. The adapter has to read the same ``api_base`` to land on the same call.""" - sent = _reasoning_effort_sent_for_thinking( + sent: Final = _reasoning_effort_sent_for_thinking( "azure_ai/gpt-5.5", "azure_ai", SUMMARIZED_THINKING, tools=[MULTIPLY_TOOL], api_base=api_base ) assert sent == expected def test_a_provider_resolved_from_the_api_base_gets_the_plain_tier(self, local_model_cost_map: None) -> None: - sent = _reasoning_effort_sent_for_thinking( + sent: Final = _reasoning_effort_sent_for_thinking( "kimi-k3", None, SUMMARIZED_THINKING, api_base="https://api.together.xyz/v1" ) assert sent == "high" def test_a_chained_gateway_keeps_the_dict_for_its_own_bridge(self, local_model_cost_map: None) -> None: - sent = _reasoning_effort_sent_for_thinking("litellm_proxy/gpt-5.4", "litellm_proxy", SUMMARIZED_THINKING) + sent: Final = _reasoning_effort_sent_for_thinking("litellm_proxy/gpt-5.4", "litellm_proxy", SUMMARIZED_THINKING) assert sent == {"effort": "high", "summary": "auto"} @@ -192,7 +271,7 @@ class TestTheNormalizedTierIsTheTierSent: def test_the_dict_form_normalizes_effort_and_keeps_its_siblings( self, local_model_cost_map, model, provider, expected ): - sent = _reasoning_effort_sent(model, provider, {"effort": "max", "summary": "detailed"}) + sent: Final = _reasoning_effort_sent(model, provider, {"effort": "max", "summary": "detailed"}) assert sent == {"effort": expected, "summary": "detailed"} @@ -215,8 +294,8 @@ class TestTheNormalizedTierIsTheTierSent: ], ) def test_an_azure_hosted_astra_deployment_drops_to_the_tier_it_accepts( - self, local_model_cost_map, model, provider - ): + self, local_model_cost_map: None, model: str, provider: str + ) -> None: """The deployment answers ``max`` with a 400 naming ``none`` through ``xhigh``, so the rows say so and the adapter sends the tier below instead of the rejected one.""" assert _reasoning_effort_sent(model, provider, "max") == "xhigh" diff --git a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 40ede74b98b..5c0e9f1a305 100644 --- a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -1,6 +1,7 @@ import asyncio, importlib, os, uuid import json import re +from collections.abc import Awaitable, Callable, Mapping from copy import deepcopy from typing import Final, List, Literal, cast, get_args from unittest.mock import MagicMock, Mock, patch @@ -25,6 +26,7 @@ from litellm.types.llms.vertex_ai import GeminiFinishReason, UsageMetadata from litellm.types.utils import ChoiceLogprobs, Usage from litellm.utils import _invalidate_model_cost_lowercase_map, CustomStreamWrapper from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER +from litellm.integrations.custom_logger import CustomLogger from litellm.llms.vertex_ai.gemini.transformation import ( gemini_convert_messages_with_history, ) @@ -38,62 +40,62 @@ from litellm.llms.vertex_ai.context_caching.transformation import ( import tempfile GEMINI_3_IMAGE_SIZE_MAPPINGS: Final[tuple[tuple[str, str, str], ...]] = ( - ('512x512', '1:1', '512'), - ('1024x1024', '1:1', '1K'), - ('2048x2048', '1:1', '2K'), - ('4096x4096', '1:1', '4K'), - ('256x1024', '1:4', '512'), - ('512x2048', '1:4', '1K'), - ('1024x4096', '1:4', '2K'), - ('2048x8192', '1:4', '4K'), - ('192x1536', '1:8', '512'), - ('384x3072', '1:8', '1K'), - ('768x6144', '1:8', '2K'), - ('1536x12288', '1:8', '4K'), - ('424x632', '2:3', '512'), - ('848x1264', '2:3', '1K'), - ('1696x2528', '2:3', '2K'), - ('3392x5056', '2:3', '4K'), - ('632x424', '3:2', '512'), - ('1264x848', '3:2', '1K'), - ('2528x1696', '3:2', '2K'), - ('5056x3392', '3:2', '4K'), - ('448x600', '3:4', '512'), - ('896x1200', '3:4', '1K'), - ('1792x2400', '3:4', '2K'), - ('3584x4800', '3:4', '4K'), - ('1024x256', '4:1', '512'), - ('2048x512', '4:1', '1K'), - ('4096x1024', '4:1', '2K'), - ('8192x2048', '4:1', '4K'), - ('600x448', '4:3', '512'), - ('1200x896', '4:3', '1K'), - ('2400x1792', '4:3', '2K'), - ('4800x3584', '4:3', '4K'), - ('464x576', '4:5', '512'), - ('928x1152', '4:5', '1K'), - ('1856x2304', '4:5', '2K'), - ('3712x4608', '4:5', '4K'), - ('576x464', '5:4', '512'), - ('1152x928', '5:4', '1K'), - ('2304x1856', '5:4', '2K'), - ('4608x3712', '5:4', '4K'), - ('1536x192', '8:1', '512'), - ('3072x384', '8:1', '1K'), - ('6144x768', '8:1', '2K'), - ('12288x1536', '8:1', '4K'), - ('384x688', '9:16', '512'), - ('768x1376', '9:16', '1K'), - ('1536x2752', '9:16', '2K'), - ('3072x5504', '9:16', '4K'), - ('688x384', '16:9', '512'), - ('1376x768', '16:9', '1K'), - ('2752x1536', '16:9', '2K'), - ('5504x3072', '16:9', '4K'), - ('792x336', '21:9', '512'), - ('1584x672', '21:9', '1K'), - ('3168x1344', '21:9', '2K'), - ('6336x2688', '21:9', '4K'), + ("512x512", "1:1", "512"), + ("1024x1024", "1:1", "1K"), + ("2048x2048", "1:1", "2K"), + ("4096x4096", "1:1", "4K"), + ("256x1024", "1:4", "512"), + ("512x2048", "1:4", "1K"), + ("1024x4096", "1:4", "2K"), + ("2048x8192", "1:4", "4K"), + ("192x1536", "1:8", "512"), + ("384x3072", "1:8", "1K"), + ("768x6144", "1:8", "2K"), + ("1536x12288", "1:8", "4K"), + ("424x632", "2:3", "512"), + ("848x1264", "2:3", "1K"), + ("1696x2528", "2:3", "2K"), + ("3392x5056", "2:3", "4K"), + ("632x424", "3:2", "512"), + ("1264x848", "3:2", "1K"), + ("2528x1696", "3:2", "2K"), + ("5056x3392", "3:2", "4K"), + ("448x600", "3:4", "512"), + ("896x1200", "3:4", "1K"), + ("1792x2400", "3:4", "2K"), + ("3584x4800", "3:4", "4K"), + ("1024x256", "4:1", "512"), + ("2048x512", "4:1", "1K"), + ("4096x1024", "4:1", "2K"), + ("8192x2048", "4:1", "4K"), + ("600x448", "4:3", "512"), + ("1200x896", "4:3", "1K"), + ("2400x1792", "4:3", "2K"), + ("4800x3584", "4:3", "4K"), + ("464x576", "4:5", "512"), + ("928x1152", "4:5", "1K"), + ("1856x2304", "4:5", "2K"), + ("3712x4608", "4:5", "4K"), + ("576x464", "5:4", "512"), + ("1152x928", "5:4", "1K"), + ("2304x1856", "5:4", "2K"), + ("4608x3712", "5:4", "4K"), + ("1536x192", "8:1", "512"), + ("3072x384", "8:1", "1K"), + ("6144x768", "8:1", "2K"), + ("12288x1536", "8:1", "4K"), + ("384x688", "9:16", "512"), + ("768x1376", "9:16", "1K"), + ("1536x2752", "9:16", "2K"), + ("3072x5504", "9:16", "4K"), + ("688x384", "16:9", "512"), + ("1376x768", "16:9", "1K"), + ("2752x1536", "16:9", "2K"), + ("5504x3072", "16:9", "4K"), + ("792x336", "21:9", "512"), + ("1584x672", "21:9", "1K"), + ("3168x1344", "21:9", "2K"), + ("6336x2688", "21:9", "4K"), ) @@ -314,9 +316,7 @@ def test_vertex_ai_response_json_schema_preserves_refs_for_gemini_2(): # $defs and $ref should be preserved (not unpacked) assert "response_json_schema" in transformed_request result_schema = transformed_request["response_json_schema"] - assert ( - "$defs" in result_schema - ), "responseJsonSchema should preserve $defs for Gemini 2.0+" + assert "$defs" in result_schema, "responseJsonSchema should preserve $defs for Gemini 2.0+" def test_vertex_ai_get_json_schema_preserves_refs_for_nested_pydantic(): @@ -396,22 +396,14 @@ def test_vertex_ai_response_json_schema_for_gemini_2(): # Types should be lowercase (standard JSON Schema format) assert transformed_request["response_json_schema"]["type"] == "object" - assert ( - transformed_request["response_json_schema"]["properties"]["name"]["type"] - == "string" - ) - assert ( - transformed_request["response_json_schema"]["properties"]["age"]["type"] - == "integer" - ) + assert transformed_request["response_json_schema"]["properties"]["name"]["type"] == "string" + assert transformed_request["response_json_schema"]["properties"]["age"]["type"] == "integer" # Should NOT have propertyOrdering (not needed for responseJsonSchema) assert "propertyOrdering" not in transformed_request["response_json_schema"] # additionalProperties should be preserved (supported by responseJsonSchema) - assert ( - transformed_request["response_json_schema"].get("additionalProperties") == False - ) + assert transformed_request["response_json_schema"].get("additionalProperties") == False def test_vertex_ai_response_schema_for_old_models(): @@ -648,7 +640,9 @@ def test_response_has_search_grounding_detection(): is False ) assert ( - VertexGeminiConfig._response_has_search_grounding({"candidates": [{"groundingMetadata": {"webSearchQueries": []}}]}) + VertexGeminiConfig._response_has_search_grounding( + {"candidates": [{"groundingMetadata": {"webSearchQueries": []}}]} + ) is False ) assert VertexGeminiConfig._response_has_search_grounding({"candidates": []}) is False @@ -789,9 +783,7 @@ def test_vertex_ai_url_context_tool_use_tokens_billed_as_input_tokens(): "candidates": [ { "urlContextMetadata": {"urlMetadata": []}, - "groundingMetadata": { - "groundingChunks": [{"web": {"uri": "https://example.com", "title": "Example"}}] - }, + "groundingMetadata": {"groundingChunks": [{"web": {"uri": "https://example.com", "title": "Example"}}]}, } ], "usageMetadata": UsageMetadata( @@ -823,9 +815,7 @@ def test_url_context_tool_use_tokens_are_priced_as_input(): "candidates": [ { "urlContextMetadata": {"urlMetadata": []}, - "groundingMetadata": { - "groundingChunks": [{"web": {"uri": "https://example.com", "title": "Example"}}] - }, + "groundingMetadata": {"groundingChunks": [{"web": {"uri": "https://example.com", "title": "Example"}}]}, } ], "usageMetadata": UsageMetadata( @@ -886,9 +876,7 @@ def test_missing_prompt_tokens_details_still_prices_the_full_prompt(): "candidates": [ { "urlContextMetadata": {"urlMetadata": []}, - "groundingMetadata": { - "groundingChunks": [{"web": {"uri": "https://example.com", "title": "Example"}}] - }, + "groundingMetadata": {"groundingChunks": [{"web": {"uri": "https://example.com", "title": "Example"}}]}, } ], "usageMetadata": UsageMetadata( @@ -925,9 +913,7 @@ def test_streaming_chunk_includes_reasoning_tokens(): "thoughtsTokenCount": 3, }, } - iterator = ModelResponseIterator( - streaming_response=[], sync_stream=True, logging_obj=litellm_logging - ) + iterator = ModelResponseIterator(streaming_response=[], sync_stream=True, logging_obj=litellm_logging) streaming_chunk = iterator.chunk_parser(chunk) assert streaming_chunk.usage is not None assert streaming_chunk.usage.prompt_tokens == 5 @@ -963,17 +949,12 @@ def test_streaming_chunk_includes_reasoning_content(): "usageMetadata": {}, } - iterator = ModelResponseIterator( - streaming_response=[], sync_stream=True, logging_obj=litellm_logging - ) + iterator = ModelResponseIterator(streaming_response=[], sync_stream=True, logging_obj=litellm_logging) streaming_chunk = iterator.chunk_parser(chunk) # The text content should be empty and reasoning_content should be populated assert streaming_chunk.choices[0].delta.content is None - assert ( - streaming_chunk.choices[0].delta.reasoning_content - == "I'm thinking through the problem..." - ) + assert streaming_chunk.choices[0].delta.reasoning_content == "I'm thinking through the problem..." def test_streaming_chunk_with_tool_calls_and_thought_includes_reasoning_content(): @@ -1018,24 +999,16 @@ def test_streaming_chunk_with_tool_calls_and_thought_includes_reasoning_content( }, } - iterator = ModelResponseIterator( - streaming_response=[], sync_stream=True, logging_obj=litellm_logging - ) + iterator = ModelResponseIterator(streaming_response=[], sync_stream=True, logging_obj=litellm_logging) streaming_chunk = iterator.chunk_parser(chunk) # Verify reasoning_content comes from the thought: true part - assert ( - streaming_chunk.choices[0].delta.reasoning_content - == "Let me think about how to get the time..." - ) + assert streaming_chunk.choices[0].delta.reasoning_content == "Let me think about how to get the time..." # Verify tool calls are also present assert streaming_chunk.choices[0].delta.tool_calls is not None assert len(streaming_chunk.choices[0].delta.tool_calls) == 1 - assert ( - streaming_chunk.choices[0].delta.tool_calls[0].function.name - == "get_current_time" - ) + assert streaming_chunk.choices[0].delta.tool_calls[0].function.name == "get_current_time" def test_streaming_chunk_with_tool_calls_no_thought_no_reasoning_content(): @@ -1077,9 +1050,7 @@ def test_streaming_chunk_with_tool_calls_no_thought_no_reasoning_content(): }, } - iterator = ModelResponseIterator( - streaming_response=[], sync_stream=True, logging_obj=litellm_logging - ) + iterator = ModelResponseIterator(streaming_response=[], sync_stream=True, logging_obj=litellm_logging) streaming_chunk = iterator.chunk_parser(chunk) # reasoning_content should be None - thoughtSignature alone does NOT mean reasoning @@ -1286,11 +1257,7 @@ def test_vertex_ai_usage_metadata_with_image_tokens_in_prompt(): # Verify the math: prompt_tokens = text + image # 533 = 6 (text) + 527 (image) - assert ( - result.prompt_tokens_details.text_tokens - + result.prompt_tokens_details.image_tokens - == result.prompt_tokens - ) + assert result.prompt_tokens_details.text_tokens + result.prompt_tokens_details.image_tokens == result.prompt_tokens def test_map_response_modalities_video(): @@ -1480,9 +1447,7 @@ def test_vertex_ai_streaming_usage_calculation(): } # Create iterator and parse chunk - iterator = ModelResponseIterator( - streaming_response=[], sync_stream=True, logging_obj=MagicMock() - ) + iterator = ModelResponseIterator(streaming_response=[], sync_stream=True, logging_obj=MagicMock()) iterator.chunk_parser(chunk) # Verify calculate_usage was called with correct parameters @@ -1554,9 +1519,7 @@ def test_vertex_ai_streaming_usage_web_search_calculation(): } # Create iterator and parse chunk - iterator = ModelResponseIterator( - streaming_response=[], sync_stream=True, logging_obj=MagicMock() - ) + iterator = ModelResponseIterator(streaming_response=[], sync_stream=True, logging_obj=MagicMock()) completed_response = iterator.chunk_parser(chunk) usage: Usage = completed_response.usage @@ -1775,7 +1738,9 @@ async def test_gemini_image_generation_async_stream(respx_mock: respx.MockRouter "modelVersion": "gemini-2.5-flash-image", }, { - "candidates": [{"content": {"role": "model", "parts": [{"text": ""}]}, "finishReason": "STOP", "index": 0}], + "candidates": [ + {"content": {"role": "model", "parts": [{"text": ""}]}, "finishReason": "STOP", "index": 0} + ], "usageMetadata": {"promptTokenCount": 13, "candidatesTokenCount": 1290, "totalTokenCount": 1303}, "modelVersion": "gemini-2.5-flash-image", }, @@ -2029,9 +1994,14 @@ def test_gemini_openai_web_search_tool_to_google_search(respx_mock: respx.MockRo GEMINI_GROUNDING_METADATA: Final = { "webSearchQueries": ["capital of France"], - "searchEntryPoint": {"renderedContent": "
"}, + "searchEntryPoint": {"renderedContent": '
'}, "groundingChunks": [ - {"web": {"uri": "https://vertexaisearch.cloud.google.com/grounding-api-redirect/AbF9w", "title": "wikipedia.org"}} + { + "web": { + "uri": "https://vertexaisearch.cloud.google.com/grounding-api-redirect/AbF9w", + "title": "wikipedia.org", + } + } ], "groundingSupports": [ { @@ -2188,9 +2158,7 @@ def test_vertex_ai_transform_parts(): # Test case 2: Tool call mode (is_function_call=False) - Single message with multiple tool calls parts_with_multiple_functions = [ - HttpxPartType( - functionCall={"name": "get_current_weather", "args": {"location": "Boston"}} - ), + HttpxPartType(functionCall={"name": "get_current_weather", "args": {"location": "Boston"}}), HttpxPartType( functionCall={ "name": "get_forecast", @@ -2220,9 +2188,7 @@ def test_vertex_ai_transform_parts(): # Test case 3: Simulating multiple messages - cumulative indexing across messages # First message with 2 tool calls (starting from index 0) first_message_parts = [ - HttpxPartType( - functionCall={"name": "get_weather", "args": {"location": "Boston"}} - ), + HttpxPartType(functionCall={"name": "get_weather", "args": {"location": "Boston"}}), HttpxPartType(functionCall={"name": "get_time", "args": {"timezone": "EST"}}), ] @@ -2239,9 +2205,7 @@ def test_vertex_ai_transform_parts(): # Second message with 1 tool call (continuing from previous index) second_message_parts = [ - HttpxPartType( - functionCall={"name": "send_email", "args": {"to": "user@example.com"}} - ), + HttpxPartType(functionCall={"name": "send_email", "args": {"to": "user@example.com"}}), ] function, tools, updated_idx = VertexGeminiConfig._transform_parts( @@ -2259,9 +2223,7 @@ def test_vertex_ai_transform_parts(): # Third message with 2 more tool calls (continuing from previous index) third_message_parts = [ - HttpxPartType( - functionCall={"name": "create_calendar_event", "args": {"title": "Meeting"}} - ), + HttpxPartType(functionCall={"name": "create_calendar_event", "args": {"title": "Meeting"}}), HttpxPartType(functionCall={"name": "set_reminder", "args": {"time": "10:00"}}), ] @@ -2308,9 +2270,7 @@ def test_vertex_ai_transform_parts(): assert updated_idx == 10 # Index should remain unchanged # Test case 6: Function call with empty args - parts_with_empty_args = [ - HttpxPartType(functionCall={"name": "simple_function", "args": {}}) - ] + parts_with_empty_args = [HttpxPartType(functionCall={"name": "simple_function", "args": {}})] function, tools, updated_idx = VertexGeminiConfig._transform_parts( parts=parts_with_empty_args, cumulative_tool_call_idx=0, is_function_call=True @@ -2326,13 +2286,9 @@ def test_vertex_ai_transform_parts(): # Test case 7: Mixed content with function calls - ensuring tool call IDs are unique mixed_parts = [ HttpxPartType(text="Before function call"), - HttpxPartType( - functionCall={"name": "function_a", "args": {"param": "value_a"}} - ), + HttpxPartType(functionCall={"name": "function_a", "args": {"param": "value_a"}}), HttpxPartType(text="Between function calls"), - HttpxPartType( - functionCall={"name": "function_b", "args": {"param": "value_b"}} - ), + HttpxPartType(functionCall={"name": "function_b", "args": {"param": "value_b"}}), HttpxPartType(text="After function calls"), ] @@ -2524,25 +2480,17 @@ def test_vertex_ai_tool_call_id_format(): tool_id = tool["id"] # Should start with 'call_' - assert tool_id.startswith( - "call_" - ), f"ID should start with 'call_', got: {tool_id}" + assert tool_id.startswith("call_"), f"ID should start with 'call_', got: {tool_id}" # Should have exactly 33 total characters (call_ + 28 hex chars) - assert ( - len(tool_id) == 33 - ), f"ID should be 33 characters long, got {len(tool_id)}: {tool_id}" + assert len(tool_id) == 33, f"ID should be 33 characters long, got {len(tool_id)}: {tool_id}" # The part after 'call_' should be 28 hex characters hex_part = tool_id[5:] # Remove 'call_' prefix - assert ( - len(hex_part) == 28 - ), f"Hex part should be 28 characters, got {len(hex_part)}: {hex_part}" + assert len(hex_part) == 28, f"Hex part should be 28 characters, got {len(hex_part)}: {hex_part}" # Should only contain valid hex characters - assert re.match( - r"^[0-9a-f]{28}$", hex_part - ), f"Should contain only lowercase hex chars, got: {hex_part}" + assert re.match(r"^[0-9a-f]{28}$", hex_part), f"Should contain only lowercase hex chars, got: {hex_part}" # Verify IDs are unique assert tools[0]["id"] != tools[1]["id"], "Tool call IDs should be unique" @@ -2559,9 +2507,7 @@ def test_vertex_ai_tool_call_id_format(): ids_generated.add(test_tools[0]["id"]) # All generated IDs should be unique - assert ( - len(ids_generated) == 10 - ), f"All 10 IDs should be unique, got {len(ids_generated)} unique IDs" + assert len(ids_generated) == 10, f"All 10 IDs should be unique, got {len(ids_generated)} unique IDs" def test_vertex_ai_map_google_maps_tool_simple(): @@ -2675,19 +2621,16 @@ def test_vertex_ai_penalty_parameters_validation(): for model, should_support in test_cases: # Test _supports_penalty_parameters method - assert ( - v._supports_penalty_parameters(model) == should_support - ), f"Model {model} penalty support should be {should_support}" + assert v._supports_penalty_parameters(model) == should_support, ( + f"Model {model} penalty support should be {should_support}" + ) # Test get_supported_openai_params method supported_params = v.get_supported_openai_params(model) - has_penalty_params = ( - "frequency_penalty" in supported_params - and "presence_penalty" in supported_params + has_penalty_params = "frequency_penalty" in supported_params and "presence_penalty" in supported_params + assert has_penalty_params == should_support, ( + f"Model {model} should {'include' if should_support else 'exclude'} penalty params in supported list" ) - assert ( - has_penalty_params == should_support - ), f"Model {model} should {'include' if should_support else 'exclude'} penalty params in supported list" # Test parameter mapping for unsupported model model = "gemini-2.5-pro-preview-06-05" @@ -2707,12 +2650,8 @@ def test_vertex_ai_penalty_parameters_validation(): ) # Penalty parameters should be filtered out for unsupported models - assert ( - "frequency_penalty" not in result - ), "frequency_penalty should be filtered out for unsupported model" - assert ( - "presence_penalty" not in result - ), "presence_penalty should be filtered out for unsupported model" + assert "frequency_penalty" not in result, "frequency_penalty should be filtered out for unsupported model" + assert "presence_penalty" not in result, "presence_penalty should be filtered out for unsupported model" # Other parameters should still be included assert "temperature" in result, "temperature should still be included" @@ -2741,18 +2680,18 @@ def test_vertex_ai_gemini_3_penalty_parameters_unsupported(): for model in gemini_3_models: # Test _supports_penalty_parameters method - assert ( - v._supports_penalty_parameters(model) == False - ), f"Gemini 3 model {model} should not support penalty parameters" + assert v._supports_penalty_parameters(model) == False, ( + f"Gemini 3 model {model} should not support penalty parameters" + ) # Test get_supported_openai_params method supported_params = v.get_supported_openai_params(model) - assert ( - "frequency_penalty" not in supported_params - ), f"frequency_penalty should not be in supported params for {model}" - assert ( - "presence_penalty" not in supported_params - ), f"presence_penalty should not be in supported params for {model}" + assert "frequency_penalty" not in supported_params, ( + f"frequency_penalty should not be in supported params for {model}" + ) + assert "presence_penalty" not in supported_params, ( + f"presence_penalty should not be in supported params for {model}" + ) # Test parameter mapping - penalty params should be filtered out non_default_params = { @@ -2771,36 +2710,28 @@ def test_vertex_ai_gemini_3_penalty_parameters_unsupported(): ) # Penalty parameters should be filtered out for Gemini 3 models - assert ( - "frequency_penalty" not in result - ), f"frequency_penalty should be filtered out for Gemini 3 model {model}" - assert ( - "presence_penalty" not in result - ), f"presence_penalty should be filtered out for Gemini 3 model {model}" + assert "frequency_penalty" not in result, f"frequency_penalty should be filtered out for Gemini 3 model {model}" + assert "presence_penalty" not in result, f"presence_penalty should be filtered out for Gemini 3 model {model}" # Other parameters should still be included - assert ( - "temperature" in result - ), f"temperature should still be included for Gemini 3 model {model}" - assert ( - "max_output_tokens" in result - ), f"max_output_tokens should still be included for Gemini 3 model {model}" + assert "temperature" in result, f"temperature should still be included for Gemini 3 model {model}" + assert "max_output_tokens" in result, f"max_output_tokens should still be included for Gemini 3 model {model}" assert result["temperature"] == 0.7 assert result["max_output_tokens"] == 100 # Test that non-Gemini 3 models still support penalty parameters (if they're not in the unsupported list) non_gemini_3_model = "gemini-2.5-pro" - assert ( - v._supports_penalty_parameters(non_gemini_3_model) == True - ), f"Non-Gemini 3 model {non_gemini_3_model} should support penalty parameters" + assert v._supports_penalty_parameters(non_gemini_3_model) == True, ( + f"Non-Gemini 3 model {non_gemini_3_model} should support penalty parameters" + ) supported_params = v.get_supported_openai_params(non_gemini_3_model) - assert ( - "frequency_penalty" in supported_params - ), f"frequency_penalty should be in supported params for {non_gemini_3_model}" - assert ( - "presence_penalty" in supported_params - ), f"presence_penalty should be in supported params for {non_gemini_3_model}" + assert "frequency_penalty" in supported_params, ( + f"frequency_penalty should be in supported params for {non_gemini_3_model}" + ) + assert "presence_penalty" in supported_params, ( + f"presence_penalty should be in supported params for {non_gemini_3_model}" + ) def test_vertex_ai_annotation_streaming_events(): @@ -2823,14 +2754,10 @@ def test_vertex_ai_annotation_streaming_events(): chunk_with_annotations = { "candidates": [ { - "content": { - "parts": [{"text": "The weather in San Francisco today is clear."}] - }, + "content": {"parts": [{"text": "The weather in San Francisco today is clear."}]}, "groundingMetadata": { "webSearchQueries": ["weather San Francisco today"], - "searchEntryPoint": { - "renderedContent": "
Search results
" - }, + "searchEntryPoint": {"renderedContent": "
Search results
"}, "groundingChunks": [ { "web": { @@ -2862,9 +2789,7 @@ def test_vertex_ai_annotation_streaming_events(): } # Create iterator and parse chunk - iterator = ModelResponseIterator( - streaming_response=[], sync_stream=True, logging_obj=litellm_logging - ) + iterator = ModelResponseIterator(streaming_response=[], sync_stream=True, logging_obj=litellm_logging) streaming_chunk = iterator.chunk_parser(chunk_with_annotations) # Verify the chunk was parsed correctly @@ -2986,12 +2911,8 @@ def test_vertex_ai_annotation_conversion(): } # Convert grounding metadata to annotations - content_text = ( - "The weather in San Francisco is currently 72°F and the time is 2:30 PM" - ) - annotations = VertexGeminiConfig._convert_grounding_metadata_to_annotations( - [grounding_metadata], content_text - ) + content_text = "The weather in San Francisco is currently 72°F and the time is 2:30 PM" + annotations = VertexGeminiConfig._convert_grounding_metadata_to_annotations([grounding_metadata], content_text) # Verify annotations were created assert len(annotations) == 3 # One for each grounding support @@ -3033,9 +2954,7 @@ def test_vertex_ai_annotation_empty_grounding_metadata(): # Test with empty grounding metadata empty_metadata = {} - annotations = VertexGeminiConfig._convert_grounding_metadata_to_annotations( - [empty_metadata], "test content" - ) + annotations = VertexGeminiConfig._convert_grounding_metadata_to_annotations([empty_metadata], "test content") assert len(annotations) == 0 # Test with missing groundingSupports @@ -3043,9 +2962,7 @@ def test_vertex_ai_annotation_empty_grounding_metadata(): "webSearchQueries": ["test query"], "groundingChunks": [{"web": {"uri": "https://example.com", "title": "Test"}}], } - annotations = VertexGeminiConfig._convert_grounding_metadata_to_annotations( - [metadata_no_supports], "test content" - ) + annotations = VertexGeminiConfig._convert_grounding_metadata_to_annotations([metadata_no_supports], "test content") assert len(annotations) == 0 # Test with empty groundingSupports @@ -3408,9 +3325,7 @@ def test_reasoning_effort_maps_to_thinking_level_gemini_3(): def test_gemini_37_38_flash_floor_minimal_thinking_level( local_model_cost_map, model, reasoning_effort, include_thoughts ): - result = VertexGeminiConfig._map_reasoning_effort_to_thinking_level( - reasoning_effort, model - ) + result = VertexGeminiConfig._map_reasoning_effort_to_thinking_level(reasoning_effort, model) assert result["thinkingLevel"] == "low" assert result["includeThoughts"] is include_thoughts @@ -3434,9 +3349,7 @@ def test_gemini_37_38_flash_floor_minimal_thinking_level( def test_gemini_flash_minimal_thinking_support( local_model_cost_map, model, reasoning_effort, expected_level, include_thoughts ): - result = VertexGeminiConfig._map_reasoning_effort_to_thinking_level( - reasoning_effort, model - ) + result = VertexGeminiConfig._map_reasoning_effort_to_thinking_level(reasoning_effort, model) assert result["thinkingLevel"] == expected_level assert result["includeThoughts"] is include_thoughts @@ -3446,12 +3359,8 @@ def test_gemini_38_flash_feature_flag_uses_low_thinking_level(local_model_cost_m monkeypatch.setattr(litellm, "enable_gemini_default_thinking_level_low", True) thinking_param = {"type": "enabled", "budget_tokens": 1024} - result_38 = VertexGeminiConfig._map_thinking_param( - thinking_param, model="gemini-3.8-flash" - ) - result_36 = VertexGeminiConfig._map_thinking_param( - thinking_param, model="gemini-3.6-flash" - ) + result_38 = VertexGeminiConfig._map_thinking_param(thinking_param, model="gemini-3.8-flash") + result_36 = VertexGeminiConfig._map_thinking_param(thinking_param, model="gemini-3.6-flash") assert result_38["thinkingLevel"] == "low" assert result_36["thinkingLevel"] == "minimal" @@ -3933,9 +3842,7 @@ def test_gemini_image_models_excluded_from_thinking(): ) # None of these should have thinkingConfig - assert ( - "thinkingConfig" not in result - ), f"Model {model} should not have thinkingConfig" + assert "thinkingConfig" not in result, f"Model {model} should not have thinkingConfig" def test_partial_json_chunk_after_first_chunk(): @@ -3966,9 +3873,7 @@ def test_partial_json_chunk_after_first_chunk(): first_chunk = '{"candidates": [{"content": {"parts": [{"text": "Hello"}]}}]}' result1 = iterator.handle_valid_json_chunk(first_chunk) assert result1 is not None, "First complete chunk should parse OK" - assert ( - iterator.sent_first_chunk is True - ), "sent_first_chunk should be True after first chunk" + assert iterator.sent_first_chunk is True, "sent_first_chunk should be True after first chunk" # Later chunk arrives PARTIAL (simulating network fragmentation) partial_chunk = '{"candidates": [{"content":' @@ -3976,9 +3881,7 @@ def test_partial_json_chunk_after_first_chunk(): # Should switch to accumulation mode instead of crashing assert result2 is None, "Partial chunk should return None while accumulating" - assert ( - iterator.chunk_type == "accumulated_json" - ), "Should switch to accumulated_json mode" + assert iterator.chunk_type == "accumulated_json", "Should switch to accumulated_json mode" def test_partial_json_chunk_on_first_chunk(): @@ -3998,9 +3901,7 @@ def test_partial_json_chunk_on_first_chunk(): result = iterator.handle_valid_json_chunk(partial) assert result is None, "Partial first chunk should return None" - assert ( - iterator.chunk_type == "accumulated_json" - ), "Should switch to accumulated_json mode" + assert iterator.chunk_type == "accumulated_json", "Should switch to accumulated_json mode" def test_accumulated_json_does_not_reparse_every_fragment(): @@ -4032,16 +3933,12 @@ def test_accumulated_json_does_not_reparse_every_fragment(): iterator.chunk_type = "accumulated_json" text = "x" * 200_000 # no braces/brackets so only the final fragment closes - blob = json.dumps( - {"candidates": [{"content": {"role": "model", "parts": [{"text": text}]}}]} - ) + blob = json.dumps({"candidates": [{"content": {"role": "model", "parts": [{"text": text}]}}]}) fragments = [blob[i : i + 4096] for i in range(0, len(blob), 4096)] assert len(fragments) > 10, "need a multi-fragment payload to exercise the bug" parsed = None - with patch.object( - json.JSONDecoder, "raw_decode", autospec=True, side_effect=json.JSONDecoder.raw_decode - ) as spy: + with patch.object(json.JSONDecoder, "raw_decode", autospec=True, side_effect=json.JSONDecoder.raw_decode) as spy: for fragment in fragments: out = iterator.handle_accumulated_json_chunk(chunk=fragment) if out is not None: @@ -4073,9 +3970,7 @@ def test_accumulated_json_partial_fragment_returns_none_without_parsing(): ) iterator.chunk_type = "accumulated_json" - with patch.object( - json.JSONDecoder, "raw_decode", autospec=True, side_effect=json.JSONDecoder.raw_decode - ) as spy: + with patch.object(json.JSONDecoder, "raw_decode", autospec=True, side_effect=json.JSONDecoder.raw_decode) as spy: result = iterator.handle_accumulated_json_chunk( chunk='{"candidates": [{"content": {"parts": [{"text": "partial' ) @@ -4174,17 +4069,11 @@ def test_vertex_ai_multiple_tool_types_separate_objects(): tool_types_in_first = [k for k in tools[0].keys()] tool_types_in_second = [k for k in tools[1].keys()] - assert ( - len(tool_types_in_first) == 1 - ), f"First Tool should have exactly 1 type, got {tool_types_in_first}" - assert ( - len(tool_types_in_second) == 1 - ), f"Second Tool should have exactly 1 type, got {tool_types_in_second}" + assert len(tool_types_in_first) == 1, f"First Tool should have exactly 1 type, got {tool_types_in_first}" + assert len(tool_types_in_second) == 1, f"Second Tool should have exactly 1 type, got {tool_types_in_second}" # Verify the correct tool types are present - assert ( - "enterpriseWebSearch" in tools[0] - ), "First Tool should contain enterpriseWebSearch" + assert "enterpriseWebSearch" in tools[0], "First Tool should contain enterpriseWebSearch" assert "url_context" in tools[1], "Second Tool should contain url_context" @@ -4646,9 +4535,7 @@ def test_vertex_ai_multiple_function_declarations_grouped(): ) # Should have only 1 Tool object (function declarations grouped) - assert ( - len(tools) == 1 - ), f"Expected 1 Tool object for grouped functions, got {len(tools)}" + assert len(tools) == 1, f"Expected 1 Tool object for grouped functions, got {len(tools)}" # Should contain function_declarations with 2 functions assert "function_declarations" in tools[0] @@ -5070,9 +4957,7 @@ def test_chunk_parser_handles_prompt_feedback_block(): logging_obj = Mock() logging_obj.optional_params = {} - streaming_obj = ModelResponseIterator( - streaming_response=iter([]), sync_stream=True, logging_obj=logging_obj - ) + streaming_obj = ModelResponseIterator(streaming_response=iter([]), sync_stream=True, logging_obj=logging_obj) # Act result = streaming_obj.chunk_parser(blocked_chunk) @@ -5080,9 +4965,9 @@ def test_chunk_parser_handles_prompt_feedback_block(): # Assert assert result is not None, "Result should not be None" assert len(result.choices) == 1, "Should have exactly one choice" - assert ( - result.choices[0].finish_reason == "content_filter" - ), f"finish_reason should be content_filter, got {result.choices[0].finish_reason}" + assert result.choices[0].finish_reason == "content_filter", ( + f"finish_reason should be content_filter, got {result.choices[0].finish_reason}" + ) assert result.choices[0].delta.content is None, "content should be None" @@ -5106,9 +4991,7 @@ def test_chunk_parser_handles_prompt_feedback_safety_block(): logging_obj = Mock() logging_obj.optional_params = {} - streaming_obj = ModelResponseIterator( - streaming_response=iter([]), sync_stream=True, logging_obj=logging_obj - ) + streaming_obj = ModelResponseIterator(streaming_response=iter([]), sync_stream=True, logging_obj=logging_obj) # Act result = streaming_obj.chunk_parser(blocked_chunk) @@ -5145,9 +5028,7 @@ def test_chunk_parser_handles_prompt_feedback_block_with_usage(): logging_obj = Mock() logging_obj.optional_params = {} - streaming_obj = ModelResponseIterator( - streaming_response=iter([]), sync_stream=True, logging_obj=logging_obj - ) + streaming_obj = ModelResponseIterator(streaming_response=iter([]), sync_stream=True, logging_obj=logging_obj) # Act result = streaming_obj.chunk_parser(blocked_chunk) @@ -5155,23 +5036,17 @@ def test_chunk_parser_handles_prompt_feedback_block_with_usage(): # Assert - 验证 content_filter 响应和 usage 都被正确处理 assert result is not None, "Result should not be None" assert len(result.choices) == 1, "Should have exactly one choice" - assert ( - result.choices[0].finish_reason == "content_filter" - ), f"finish_reason should be content_filter, got {result.choices[0].finish_reason}" + assert result.choices[0].finish_reason == "content_filter", ( + f"finish_reason should be content_filter, got {result.choices[0].finish_reason}" + ) assert result.choices[0].delta.content is None, "content should be None" # 验证 usage 信息被正确提取 assert hasattr(result, "usage"), "result should have usage attribute" assert result.usage is not None, "usage should not be None" - assert ( - result.usage.prompt_tokens == 8175 - ), f"prompt_tokens should be 8175, got {result.usage.prompt_tokens}" - assert ( - result.usage.completion_tokens == 0 - ), f"completion_tokens should be 0, got {result.usage.completion_tokens}" - assert ( - result.usage.total_tokens == 8175 - ), f"total_tokens should be 8175, got {result.usage.total_tokens}" + assert result.usage.prompt_tokens == 8175, f"prompt_tokens should be 8175, got {result.usage.prompt_tokens}" + assert result.usage.completion_tokens == 0, f"completion_tokens should be 0, got {result.usage.completion_tokens}" + assert result.usage.total_tokens == 8175, f"total_tokens should be 8175, got {result.usage.total_tokens}" def test_vertex_ai_traffic_type_preserved_in_hidden_params_streaming(): @@ -5190,14 +5065,10 @@ def test_vertex_ai_traffic_type_preserved_in_hidden_params_streaming(): }, } - iterator = ModelResponseIterator( - streaming_response=[], sync_stream=True, logging_obj=MagicMock() - ) + iterator = ModelResponseIterator(streaming_response=[], sync_stream=True, logging_obj=MagicMock()) result = iterator.chunk_parser(chunk) - assert ( - result._hidden_params["provider_specific_fields"]["traffic_type"] == "ON_DEMAND" - ) + assert result._hidden_params["provider_specific_fields"]["traffic_type"] == "ON_DEMAND" def test_vertex_ai_traffic_type_preserved_in_hidden_params_non_streaming(): @@ -5236,10 +5107,7 @@ def test_vertex_ai_traffic_type_preserved_in_hidden_params_non_streaming(): encoding=None, ) - assert ( - result._hidden_params["provider_specific_fields"]["traffic_type"] - == "PROVISIONED_THROUGHPUT" - ) + assert result._hidden_params["provider_specific_fields"]["traffic_type"] == "PROVISIONED_THROUGHPUT" def test_vertex_ai_service_tier_streaming(): @@ -5323,9 +5191,7 @@ def test_vertex_ai_traffic_type_surfaced_in_responses_api(): from litellm.types.utils import Choices, Message model_response = ModelResponse() - model_response._hidden_params["provider_specific_fields"] = { - "traffic_type": "ON_DEMAND" - } + model_response._hidden_params["provider_specific_fields"] = {"traffic_type": "ON_DEMAND"} model_response.choices = [ Choices( message=Message(content="Hello", role="assistant"), @@ -5334,15 +5200,15 @@ def test_vertex_ai_traffic_type_surfaced_in_responses_api(): ) ] - responses_api_response = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( - request_input="test", - chat_completion_response=model_response, - responses_api_request={}, + responses_api_response = ( + LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="test", + chat_completion_response=model_response, + responses_api_request={}, + ) ) - assert ( - responses_api_response.provider_specific_fields["traffic_type"] == "ON_DEMAND" - ) + assert responses_api_response.provider_specific_fields["traffic_type"] == "ON_DEMAND" def test_vertex_ai_web_search_options_parameter(): @@ -5416,12 +5282,8 @@ def test_vertex_ai_web_search_options_in_map_openai_params(): assert "tools" in optional_params, "tools should be added to optional_params" assert len(optional_params["tools"]) == 1, "Should have exactly one tool" assert "googleSearch" in optional_params["tools"][0], "Tool should be googleSearch" - assert ( - optional_params["tools"][0]["googleSearch"] == {} - ), "googleSearch should be empty config" - assert ( - "web_search_options" not in optional_params - ), "web_search_options should be removed after transformation" + assert optional_params["tools"][0]["googleSearch"] == {}, "googleSearch should be empty config" + assert "web_search_options" not in optional_params, "web_search_options should be removed after transformation" def test_vertex_ai_service_tier_in_map_openai_params(): @@ -5509,24 +5371,16 @@ def test_vertex_ai_usage_metadata_with_video_tokens_in_prompt(): # Verify prompt token details include video tokens assert result.prompt_tokens_details is not None - assert ( - result.prompt_tokens_details.video_tokens == 10240 - ), "Prompt video tokens should be 10240" - assert ( - result.prompt_tokens_details.text_tokens == 9 - ), "Prompt text tokens should be 9" - assert ( - result.prompt_tokens_details.audio_tokens == 200 - ), "Prompt audio tokens should be 200" + assert result.prompt_tokens_details.video_tokens == 10240, "Prompt video tokens should be 10240" + assert result.prompt_tokens_details.text_tokens == 9, "Prompt text tokens should be 9" + assert result.prompt_tokens_details.audio_tokens == 200, "Prompt audio tokens should be 200" # Verify completion token details assert result.completion_tokens_details is not None - assert ( - result.completion_tokens_details.text_tokens == 79 - ), "Completion text tokens should be 79" - assert ( - result.completion_tokens_details.video_tokens is None - ), "Completion video tokens should be None (text-only response)" + assert result.completion_tokens_details.text_tokens == 79, "Completion text tokens should be 79" + assert result.completion_tokens_details.video_tokens is None, ( + "Completion video tokens should be None (text-only response)" + ) def test_vertex_ai_usage_metadata_with_video_tokens_in_candidates(): @@ -5667,9 +5521,9 @@ def test_vertex_ai_usage_metadata_with_document_tokens_in_prompt(): # DOCUMENT tokens should be included in text_tokens: 8 (TEXT) + 774 (DOCUMENT) = 782 assert result.prompt_tokens_details is not None - assert ( - result.prompt_tokens_details.text_tokens == 782 - ), "DOCUMENT modality tokens should be added to text_tokens (8 TEXT + 774 DOCUMENT = 782)" + assert result.prompt_tokens_details.text_tokens == 782, ( + "DOCUMENT modality tokens should be added to text_tokens (8 TEXT + 774 DOCUMENT = 782)" + ) # Verify completion token details assert result.completion_tokens_details is not None @@ -6166,9 +6020,7 @@ def test_mid_stream_429_error_raises_during_iteration(): { "content": { "role": "model", - "parts": [ - {"text": "Let me think about this...", "thought": True} - ], + "parts": [{"text": "Let me think about this...", "thought": True}], }, "index": 0, } @@ -6188,9 +6040,7 @@ def test_mid_stream_429_error_raises_during_iteration(): { "content": { "role": "model", - "parts": [ - {"text": "I'll generate the image now.", "thought": True} - ], + "parts": [{"text": "I'll generate the image now.", "thought": True}], }, "index": 0, } @@ -6227,6 +6077,7 @@ def test_mid_stream_429_error_raises_during_iteration(): # Iterate the stream: first chunks should succeed, then 429 error should be raised results = [] + def _drain(): for chunk in streaming_obj: if chunk is not None: @@ -6236,9 +6087,7 @@ def test_mid_stream_429_error_raises_during_iteration(): _drain() # Verify: received normal chunks before the error - assert ( - len(results) >= 1 - ), "Should have received at least 1 normal chunk before the error" + assert len(results) >= 1, "Should have received at least 1 normal chunk before the error" # Verify: 429 error is properly raised assert exc_info.value.status_code == 429 @@ -6427,9 +6276,7 @@ def _accumulating_gemini_iterator(): ModelResponseIterator, ) - iterator = ModelResponseIterator( - streaming_response=[], sync_stream=True, logging_obj=MagicMock() - ) + iterator = ModelResponseIterator(streaming_response=[], sync_stream=True, logging_obj=MagicMock()) iterator.chunk_type = "accumulated_json" return iterator @@ -6978,6 +6825,7 @@ def _vcr_outcome_gate(request, vcr): yield record_vcr_outcome(request, vcr) + @pytest.fixture(scope="function") def isolate_litellm_state(): """ @@ -7030,6 +6878,7 @@ def isolate_litellm_state(): setattr(litellm, attr, original_value) _invalidate_model_cost_lowercase_map() + _SCALAR_DEFAULTS = { "num_retries": getattr(litellm, "num_retries", None), "num_retries_per_request": getattr(litellm, "num_retries_per_request", None), @@ -7050,6 +6899,7 @@ _SCALAR_DEFAULTS = { "api_key": getattr(litellm, "api_key", None), } + @pytest.fixture(scope="module") def setup_and_teardown(): """ @@ -7072,6 +6922,7 @@ def setup_and_teardown(): litellm.in_memory_llm_clients_cache.flush_cache() yield + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_thought_true_creates_thinking_block(): """ @@ -7086,6 +6937,7 @@ def test_thought_true_creates_thinking_block(): assert block["thinking"] == "Some thinking" assert block["signature"] == "sig-1" + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_thought_true_with_empty_text_creates_block(): """ @@ -7097,6 +6949,7 @@ def test_thought_true_with_empty_text_creates_block(): assert len(thinking_blocks) == 1 assert thinking_blocks[0]["thinking"] == "" + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_thought_signature_without_thought_does_not_create_block(): """ @@ -7109,6 +6962,7 @@ def test_thought_signature_without_thought_does_not_create_block(): thinking_blocks = config._extract_thinking_blocks_from_parts(parts) assert thinking_blocks == [] + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_extract_thought_signatures_from_regular_parts(): """ @@ -7128,6 +6982,7 @@ def test_extract_thought_signatures_from_regular_parts(): assert len(signatures) == 1 assert signatures[0] == "sig-regular-123" + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_extract_multiple_thought_signatures(): """ @@ -7146,6 +7001,7 @@ def test_extract_multiple_thought_signatures(): assert signatures[0] == "sig-1" assert signatures[1] == "sig-2" + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_round_trip_thought_signature_in_conversation(): """ @@ -7180,6 +7036,7 @@ def test_round_trip_thought_signature_in_conversation(): assert "thoughtSignature" in text_part assert text_part["thoughtSignature"] == "sig-round-trip-abc" + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_round_trip_without_thought_signature_still_works(): """ @@ -7418,18 +7275,10 @@ def test_gemini_flash_image_preview_models(model_name: str): mock_response = ImageResponse() mock_response.data = [ImageObject(b64_json="test_base64_data", url=None)] - with patch( - "litellm.llms.custom_httpx.llm_http_handler.HTTPHandler.post" - ) as mock_post: + with patch("litellm.llms.custom_httpx.llm_http_handler.HTTPHandler.post") as mock_post: mock_http_response = MagicMock() mock_http_response.json.return_value = { - "candidates": [ - { - "content": { - "parts": [{"inlineData": {"data": "test_base64_image_data"}}] - } - } - ] + "candidates": [{"content": {"parts": [{"inlineData": {"data": "test_base64_image_data"}}]}}] } mock_http_response.status_code = 200 mock_post.return_value = mock_http_response @@ -7447,9 +7296,7 @@ def test_gemini_flash_image_preview_models(model_name: str): mock_post.assert_called_once() call_args = mock_post.call_args - called_url = ( - call_args[0][0] if call_args[0] else call_args.kwargs.get("url", "") - ) + called_url = call_args[0][0] if call_args[0] else call_args.kwargs.get("url", "") assert ":generateContent" in called_url assert model_name.split("/", 1)[1] in called_url @@ -7481,23 +7328,13 @@ def test_gemini_flash_image_preview_models(model_name: str): ), ], ) -def test_gemini_image_generation_forwards_image_config( - model: str, kwargs: dict, expected_image_config: dict -): +def test_gemini_image_generation_forwards_image_config(model: str, kwargs: dict, expected_image_config: dict): from unittest.mock import patch, MagicMock - with patch( - "litellm.llms.custom_httpx.llm_http_handler.HTTPHandler.post" - ) as mock_post: + with patch("litellm.llms.custom_httpx.llm_http_handler.HTTPHandler.post") as mock_post: mock_http_response = MagicMock() mock_http_response.json.return_value = { - "candidates": [ - { - "content": { - "parts": [{"inlineData": {"data": "test_base64_image_data"}}] - } - } - ] + "candidates": [{"content": {"parts": [{"inlineData": {"data": "test_base64_image_data"}}]}}] } mock_http_response.status_code = 200 mock_post.return_value = mock_http_response @@ -7558,9 +7395,7 @@ def test_gemini_image_generation_openai_size_maps_to_google_table( map_openai_size_to_gemini_image_config, ) - assert map_openai_size_to_gemini_image_config( - size, "gemini-3-pro-image-preview" - ) == { + assert map_openai_size_to_gemini_image_config(size, "gemini-3-pro-image-preview") == { "aspectRatio": expected_aspect_ratio, "imageSize": expected_image_size, } @@ -7584,9 +7419,7 @@ def test_gemini_image_generation_openai_size_snaps_to_nearest_option( map_openai_size_to_gemini_image_config, ) - assert map_openai_size_to_gemini_image_config( - size, "gemini-3-pro-image-preview" - ) == { + assert map_openai_size_to_gemini_image_config(size, "gemini-3-pro-image-preview") == { "aspectRatio": expected_aspect_ratio, "imageSize": expected_image_size, } @@ -7598,9 +7431,7 @@ def test_gemini_image_generation_openai_size_auto_uses_google_defaults(size: str map_openai_size_to_gemini_image_config, ) - assert map_openai_size_to_gemini_image_config( - size, "gemini-3-pro-image-preview" - ) is None + assert map_openai_size_to_gemini_image_config(size, "gemini-3-pro-image-preview") is None def test_gemini_imagen_models_use_predict_endpoint(): @@ -7610,13 +7441,9 @@ def test_gemini_imagen_models_use_predict_endpoint(): from unittest.mock import patch, MagicMock from litellm.types.utils import ImageResponse, ImageObject - with patch( - "litellm.llms.custom_httpx.llm_http_handler.HTTPHandler.post" - ) as mock_post: + with patch("litellm.llms.custom_httpx.llm_http_handler.HTTPHandler.post") as mock_post: mock_http_response = MagicMock() - mock_http_response.json.return_value = { - "predictions": [{"bytesBase64Encoded": "test_base64_image_data"}] - } + mock_http_response.json.return_value = {"predictions": [{"bytesBase64Encoded": "test_base64_image_data"}]} mock_http_response.status_code = 200 mock_post.return_value = mock_http_response @@ -7632,9 +7459,7 @@ def test_gemini_imagen_models_use_predict_endpoint(): mock_post.assert_called_once() call_args = mock_post.call_args - called_url = ( - call_args[0][0] if call_args[0] else call_args.kwargs.get("url", "") - ) + called_url = call_args[0][0] if call_args[0] else call_args.kwargs.get("url", "") assert ":predict" in called_url assert "imagen-3.0-generate-001" in called_url @@ -7797,11 +7622,7 @@ async def test_claude_tool_use_with_gemini(): if "data: " in chunk_str: try: - data_line = [ - line - for line in chunk_str.split("\n") - if line.startswith("data: ") - ][0] + data_line = [line for line in chunk_str.split("\n") if line.startswith("data: ")][0] json_str = data_line[6:] chunk_data = json.loads(json_str) @@ -7814,24 +7635,15 @@ async def test_claude_tool_use_with_gemini(): if ( chunk_data.get("type") == "message_delta" - and chunk_data.get("delta", {}).get("stop_reason") - is not None + and chunk_data.get("delta", {}).get("stop_reason") is not None and "usage" in chunk_data ): has_usage_in_message_delta = True usage = chunk_data["usage"] - assert ( - "input_tokens" in usage - ), "input_tokens should be present in usage" - assert ( - "output_tokens" in usage - ), "output_tokens should be present in usage" - assert isinstance( - usage["input_tokens"], int - ), "input_tokens should be an integer" - assert isinstance( - usage["output_tokens"], int - ), "output_tokens should be an integer" + assert "input_tokens" in usage, "input_tokens should be present in usage" + assert "output_tokens" in usage, "output_tokens should be present in usage" + assert isinstance(usage["input_tokens"], int), "input_tokens should be an integer" + assert isinstance(usage["output_tokens"], int), "output_tokens should be an integer" print(f"Found usage in message_delta: {usage}") except (json.JSONDecodeError, IndexError) as e: @@ -7846,9 +7658,7 @@ async def test_claude_tool_use_with_gemini(): assert is_content_block_tool_use, "content_block_tool_use should be present" assert is_partial_json, "partial_json should be present" - assert ( - has_usage_in_message_delta - ), "Usage should be present in message_delta with stop_reason" + assert has_usage_in_message_delta, "Usage should be present in message_delta with stop_reason" assert is_content_block_stop, "is_content_block_stop should be present" @@ -7877,28 +7687,22 @@ def test_gemini_reasoning_effort_minimal(): ) request_body = raw_request["raw_request_body"] - assert ( - "generationConfig" in request_body - ), f"Model {model} should have generationConfig" + assert "generationConfig" in request_body, f"Model {model} should have generationConfig" generation_config = request_body["generationConfig"] - assert ( - "thinkingConfig" in generation_config - ), f"Model {model} should have thinkingConfig" + assert "thinkingConfig" in generation_config, f"Model {model} should have thinkingConfig" thinking_config = generation_config["thinkingConfig"] - assert ( - "thinkingBudget" in thinking_config - ), f"Model {model} should have thinkingBudget" + assert "thinkingBudget" in thinking_config, f"Model {model} should have thinkingBudget" actual_budget = thinking_config["thinkingBudget"] - assert ( - actual_budget == expected_min_budget - ), f"Model {model} should map 'minimal' to {expected_min_budget} tokens, got {actual_budget}" + assert actual_budget == expected_min_budget, ( + f"Model {model} should map 'minimal' to {expected_min_budget} tokens, got {actual_budget}" + ) - assert thinking_config.get( - "includeThoughts", True - ), f"Model {model} should have includeThoughts=True for minimal reasoning effort" + assert thinking_config.get("includeThoughts", True), ( + f"Model {model} should have includeThoughts=True for minimal reasoning effort" + ) try: raw_request = return_raw_request( @@ -7913,9 +7717,7 @@ def test_gemini_reasoning_effort_minimal(): request_body = raw_request["raw_request_body"] generation_config = request_body["generationConfig"] thinking_config = generation_config["thinkingConfig"] - assert ( - thinking_config["thinkingBudget"] == 128 - ), "Unknown model should use generic fallback of 128 tokens" + assert thinking_config["thinkingBudget"] == 128, "Unknown model should use generic fallback of 128 tokens" except Exception as e: print(f"Note: Unknown model test skipped due to: {e}") pass @@ -7938,9 +7740,7 @@ def test_gemini_exception_message_format(): mock_response.text = "Invalid API key provided" mock_response.headers = {} - mock_exception = httpx.HTTPStatusError( - message="Bad Request", request=Mock(), response=mock_response - ) + mock_exception = httpx.HTTPStatusError(message="Bad Request", request=Mock(), response=mock_response) mock_exception.response = mock_response mock_exception.status_code = 400 @@ -7960,9 +7760,9 @@ def test_gemini_exception_message_format(): f"Expected 'GeminiException' in error message, got: {error_message}. " f"This test should fail before the fix is implemented." ) - assert ( - "VertexAIException" not in error_message - ), f"Should not contain 'VertexAIException' in error message, got: {error_message}" + assert "VertexAIException" not in error_message, ( + f"Should not contain 'VertexAIException' in error message, got: {error_message}" + ) def test_reasoning_effort_none_mapping(): @@ -8016,17 +7816,11 @@ def test_gemini_function_args_preserve_unicode(): parsed_args = json.loads(arguments_str) assert parsed_args["message"] == "やあ", "Japanese characters should be preserved" - assert ( - parsed_args["recipient"] == "たけし" - ), "Japanese characters should be preserved" + assert parsed_args["recipient"] == "たけし", "Japanese characters should be preserved" assert "\\u" not in arguments_str, "Should not contain Unicode escape sequences" - assert ( - "やあ" in arguments_str - ), "Original Japanese characters should be in the string" - assert ( - "たけし" in arguments_str - ), "Original Japanese characters should be in the string" + assert "やあ" in arguments_str, "Original Japanese characters should be in the string" + assert "たけし" in arguments_str, "Original Japanese characters should be in the string" parts_spanish = [ { @@ -8079,12 +7873,8 @@ def test_anthropic_thinking_param_to_gemini_3_provider_defaults(): model="gemini-3-flash", ) - assert ( - "thinkingLevel" not in result - ), "Should not force thinkingLevel for Gemini 3" - assert ( - "thinkingBudget" not in result - ), "Should NOT have thinkingBudget for Gemini 3" + assert "thinkingLevel" not in result, "Should not force thinkingLevel for Gemini 3" + assert "thinkingBudget" not in result, "Should NOT have thinkingBudget for Gemini 3" assert result["includeThoughts"] is True thinking_param_disabled: AnthropicThinkingParam = { @@ -8098,10 +7888,7 @@ def test_anthropic_thinking_param_to_gemini_3_provider_defaults(): ) assert result_disabled.get("includeThoughts") is False - assert ( - "thinkingLevel" not in result_disabled - or result_disabled.get("thinkingLevel") is None - ) + assert "thinkingLevel" not in result_disabled or result_disabled.get("thinkingLevel") is None thinking_param_zero: AnthropicThinkingParam = { "type": "enabled", @@ -8114,10 +7901,7 @@ def test_anthropic_thinking_param_to_gemini_3_provider_defaults(): ) assert result_zero["includeThoughts"] is False - assert ( - "thinkingLevel" not in result_zero - or result_zero.get("thinkingLevel") is None - ) + assert "thinkingLevel" not in result_zero or result_zero.get("thinkingLevel") is None result_gemini3flashpreview = VertexGeminiConfig._map_thinking_param( thinking_param=thinking_param, @@ -8203,9 +7987,7 @@ def test_anthropic_thinking_param_to_gemini_2_thinkingBudget(): ) assert "thinkingBudget" in result_gemini2, "Should have thinkingBudget for Gemini 2" - assert ( - "thinkingLevel" not in result_gemini2 - ), "Should NOT have thinkingLevel for Gemini 2" + assert "thinkingLevel" not in result_gemini2, "Should NOT have thinkingLevel for Gemini 2" assert result_gemini2["includeThoughts"] is True assert result_gemini2["thinkingBudget"] == 10000 @@ -8241,12 +8023,8 @@ def test_anthropic_thinking_param_via_map_openai_params(): assert "thinkingConfig" in result, "Should have thinkingConfig in optional_params" thinking_config = result["thinkingConfig"] - assert ( - "thinkingLevel" not in thinking_config - ), "Should not force thinkingLevel for Gemini 3 by default" - assert ( - "thinkingBudget" not in thinking_config - ), "Should NOT have thinkingBudget for Gemini 3" + assert "thinkingLevel" not in thinking_config, "Should not force thinkingLevel for Gemini 3 by default" + assert "thinkingBudget" not in thinking_config, "Should NOT have thinkingBudget for Gemini 3" assert thinking_config["includeThoughts"] is True optional_params_2 = {} @@ -8259,12 +8037,8 @@ def test_anthropic_thinking_param_via_map_openai_params(): assert "thinkingConfig" in result_2, "Should have thinkingConfig in optional_params" thinking_config_2 = result_2["thinkingConfig"] - assert ( - "thinkingBudget" in thinking_config_2 - ), "Should have thinkingBudget for Gemini 2" - assert ( - "thinkingLevel" not in thinking_config_2 - ), "Should NOT have thinkingLevel for Gemini 2" + assert "thinkingBudget" in thinking_config_2, "Should have thinkingBudget for Gemini 2" + assert "thinkingLevel" not in thinking_config_2, "Should NOT have thinkingLevel for Gemini 2" assert thinking_config_2["includeThoughts"] is True assert thinking_config_2["thinkingBudget"] == 10000 @@ -8284,9 +8058,9 @@ def test_gemini_31_flash_lite_reasoning_effort_minimal(): reasoning_effort="minimal", model="gemini-3.1-flash-lite-preview", ) - assert ( - result["thinkingLevel"] == "minimal" - ), f"Expected thinkingLevel='minimal' for gemini-3.1-flash-lite-preview, got '{result['thinkingLevel']}'" + assert result["thinkingLevel"] == "minimal", ( + f"Expected thinkingLevel='minimal' for gemini-3.1-flash-lite-preview, got '{result['thinkingLevel']}'" + ) assert result["includeThoughts"] is True from litellm.utils import return_raw_request @@ -8302,12 +8076,12 @@ def test_gemini_31_flash_lite_reasoning_effort_minimal(): ) generation_config = raw_request["raw_request_body"]["generationConfig"] thinking_config = generation_config["thinkingConfig"] - assert ( - thinking_config.get("thinkingLevel") == "minimal" - ), f"Expected thinkingLevel='minimal' via full flow, got {thinking_config}" - assert ( - "thinkingBudget" not in thinking_config - ), "gemini-3.1-flash-lite-preview should use thinkingLevel, not thinkingBudget" + assert thinking_config.get("thinkingLevel") == "minimal", ( + f"Expected thinkingLevel='minimal' via full flow, got {thinking_config}" + ) + assert "thinkingBudget" not in thinking_config, ( + "gemini-3.1-flash-lite-preview should use thinkingLevel, not thinkingBudget" + ) @pytest.mark.usefixtures("fake_provider_credentials") @@ -8855,10 +8629,7 @@ async def test_gemini_context_caching_disabled_flag(sync_mode): "content": [ { "type": "text", - "text": "Here is the full text of a complex legal agreement {}".format( - uuid.uuid4() - ) - * 4000, + "text": "Here is the full text of a complex legal agreement {}".format(uuid.uuid4()) * 4000, "cache_control": {"type": "ephemeral"}, } ], @@ -8894,9 +8665,7 @@ async def test_gemini_context_caching_disabled_flag(sync_mode): else: client = AsyncHTTPHandler(concurrent_limit=1) - with patch.object( - client, "post", side_effect=mock_gemini_request - ) as mock_client: + with patch.object(client, "post", side_effect=mock_gemini_request) as mock_client: try: if sync_mode: response = litellm.completion( @@ -8919,26 +8688,22 @@ async def test_gemini_context_caching_disabled_flag(sync_mode): print(e) # When caching is disabled, should only make 1 call (no separate cache creation call) - assert ( - mock_client.call_count == 1 - ), f"Expected 1 call when caching is disabled, got {mock_client.call_count}" + assert mock_client.call_count == 1, ( + f"Expected 1 call when caching is disabled, got {mock_client.call_count}" + ) first_call_args = mock_client.call_args_list[0].kwargs first_call_positional_args = mock_client.call_args_list[0].args print(f"first_call_args with caching disabled: {first_call_args}") - print( - f"first_call_positional_args with caching disabled: {first_call_positional_args}" - ) + print(f"first_call_positional_args with caching disabled: {first_call_positional_args}") # Assert that cachedContents is NOT in the URL when caching is disabled url = first_call_args.get( "url", first_call_positional_args[0] if first_call_positional_args else "", ) - assert ( - "cachedContents" not in url - ), "cachedContents should not be in URL when caching is disabled" + assert "cachedContents" not in url, "cachedContents should not be in URL when caching is disabled" finally: # Restore original flag value @@ -9031,9 +8796,7 @@ def test_gemini_function_call_parameter_in_messages(): "contents": [ { "role": "user", - "parts": [ - {"text": "search for weather in boston (use `search`)"} - ], + "parts": [{"text": "search for weather in boston (use `search`)"}], }, { "role": "model", @@ -9052,17 +8815,13 @@ def test_gemini_function_call_parameter_in_messages(): { "function_response": { "name": "search", - "response": { - "content": "The current weather in Boston is 22°F." - }, + "response": {"content": "The current weather in Boston is 22°F."}, } } ], }, ], - "system_instruction": { - "parts": [{"text": "Use search for most queries."}] - }, + "system_instruction": {"parts": [{"text": "Use search for most queries."}]}, "tools": [ { "function_declarations": [ @@ -9137,9 +8896,7 @@ def test_gemini_function_call_parameter_in_messages_2(): { "function_response": { "name": "search", - "response": { - "content": "The weather in Boston is 100 degrees." - }, + "response": {"content": "The weather in Boston is 100 degrees."}, } } ], @@ -9517,9 +9274,7 @@ def test_vertex_ai_cache_creation_kwarg_cannot_change_usage_or_cost() -> None: mock_response.status_code = 200 mock_response.headers = {"Content-Type": "application/json"} mock_response.json.return_value = { - "candidates": [ - {"content": {"parts": [{"text": "hi"}], "role": "model"}, "finishReason": "STOP", "index": 0} - ], + "candidates": [{"content": {"parts": [{"text": "hi"}], "role": "model"}, "finishReason": "STOP", "index": 0}], "usageMetadata": {"promptTokenCount": 10, "candidatesTokenCount": 5, "totalTokenCount": 15}, } client = HTTPHandler() @@ -9585,14 +9340,51 @@ def load_vertex_ai_credentials(): # Gemini prebuilt TTS voices per https://ai.google.dev/gemini-api/docs/speech-generation (page dated 2026-10-01, read 2026-10-07) GEMINI_TTS_PREBUILT_VOICES: Final[frozenset[str]] = frozenset( { - "Zephyr", "Puck", "Charon", "Kore", "Fenrir", "Leda", "Orus", "Aoede", "Callirrhoe", "Autonoe", - "Enceladus", "Iapetus", "Umbriel", "Algieba", "Despina", "Erinome", "Algenib", "Rasalgethi", - "Laomedeia", "Achernar", "Alnilam", "Schedar", "Gacrux", "Pulcherrima", "Achird", "Zubenelgenubi", - "Vindemiatrix", "Sadachbia", "Sadaltager", "Sulafat", + "Zephyr", + "Puck", + "Charon", + "Kore", + "Fenrir", + "Leda", + "Orus", + "Aoede", + "Callirrhoe", + "Autonoe", + "Enceladus", + "Iapetus", + "Umbriel", + "Algieba", + "Despina", + "Erinome", + "Algenib", + "Rasalgethi", + "Laomedeia", + "Achernar", + "Alnilam", + "Schedar", + "Gacrux", + "Pulcherrima", + "Achird", + "Zubenelgenubi", + "Vindemiatrix", + "Sadachbia", + "Sadaltager", + "Sulafat", } ) OPENAI_TTS_VOICES: Final[tuple[str, ...]] = ( - "alloy", "ash", "ballad", "cedar", "coral", "echo", "fable", "marin", "onyx", "sage", "shimmer", "verse", + "alloy", + "ash", + "ballad", + "cedar", + "coral", + "echo", + "fable", + "marin", + "onyx", + "sage", + "shimmer", + "verse", ) @@ -9678,3 +9470,289 @@ def test_streaming_usage_reads_cache_creation_tokens_recorded_on_logging_obj(): assert chunk.usage.prompt_tokens == 1536 + 1280 assert chunk.usage.cache_creation_input_tokens == 1280 + + +THOUGHT_SUMMARY_REQUESTS: Final = ( + pytest.param("gemini-2.5-flash", {"reasoning_effort": "minimal"}, id="budget-minimal"), + pytest.param("gemini-2.5-flash", {"reasoning_effort": "low"}, id="budget-low"), + pytest.param("gemini-2.5-flash", {"reasoning_effort": "medium"}, id="budget-medium"), + pytest.param("gemini-2.5-flash", {"reasoning_effort": "high"}, id="budget-high"), + pytest.param("gemini-3.6-flash", {"reasoning_effort": "minimal"}, id="level-minimal"), + pytest.param("gemini-3.8-flash", {"reasoning_effort": "minimal"}, id="level-minimal-mapped-to-low"), + pytest.param("gemini-3.8-flash", {"reasoning_effort": "low"}, id="level-low"), + pytest.param("gemini-3.8-flash", {"reasoning_effort": "medium"}, id="level-medium"), + pytest.param("gemini-3.8-flash", {"reasoning_effort": "high"}, id="level-high"), + pytest.param("gemini-2.5-flash", {"thinking": {"type": "enabled", "budget_tokens": 2048}}, id="budget-thinking"), + pytest.param("gemini-3.8-flash", {"thinking": {"type": "enabled", "budget_tokens": 2048}}, id="level-thinking"), +) + + +async def _sent_thinking_config(model: str, thinking_params: dict[str, object]) -> dict[str, object]: + captured: Final[dict[str, dict]] = {} + + def upstream(request: httpx.Request) -> httpx.Response: + captured["body"] = json.loads(request.content) + return httpx.Response( + 200, + json={ + "candidates": [{"content": {"parts": [{"text": "hi"}], "role": "model"}, "finishReason": "STOP"}], + "usageMetadata": {"promptTokenCount": 1, "candidatesTokenCount": 1, "totalTokenCount": 2}, + }, + request=request, + ) + + client: Final = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(upstream)) + await litellm.acompletion( + model=f"gemini/{model}", + messages=[{"role": "user", "content": "hi"}], + api_key="fake-gemini-key", + client=client, + **thinking_params, + ) + return captured["body"]["generationConfig"]["thinkingConfig"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("model", "thinking_params"), THOUGHT_SUMMARY_REQUESTS) +async def test_gemini_thought_summaries_off_keeps_thinking_config_otherwise_unchanged( + model: str, + thinking_params: dict[str, object], + local_model_cost_map: None, + monkeypatch: pytest.MonkeyPatch, +) -> None: + summaries_on: Final = await _sent_thinking_config(model, thinking_params) + monkeypatch.setattr(litellm, "gemini_include_thoughts_default", False) + summaries_off: Final = await _sent_thinking_config(model, thinking_params) + + assert summaries_on["includeThoughts"] is True + assert summaries_off == {**summaries_on, "includeThoughts": False} + + +PER_REQUEST_THOUGHT_SUMMARY_ASKS: Final = ( + pytest.param( + "gemini-3.8-flash", + {"thinking": {"type": "enabled", "budget_tokens": 2048, "display": "summarized"}}, + False, + True, + id="level-display-summarized-beats-off", + ), + pytest.param( + "gemini-3.8-flash", + {"thinking": {"type": "enabled", "budget_tokens": 2048, "display": "omitted"}}, + True, + True, + id="level-display-omitted-leaves-the-default-on", + ), + pytest.param( + "gemini-2.5-flash", + {"thinking": {"type": "enabled", "budget_tokens": 2048, "display": "summarized"}}, + False, + True, + id="budget-display-summarized-beats-off", + ), + pytest.param( + "gemini-2.5-flash", + {"thinking": {"type": "enabled", "budget_tokens": 2048, "display": "omitted"}}, + True, + True, + id="budget-display-omitted-leaves-the-default-on", + ), + pytest.param( + "gemini-3.8-flash", + {"thinking": {"type": "enabled", "budget_tokens": 2048, "display": "updates"}}, + False, + True, + id="level-display-updates-beats-off", + ), + pytest.param( + "gemini-3.8-flash", + {"reasoning_effort": {"effort": "medium", "summary": "auto"}}, + False, + True, + id="level-effort-summary-beats-off", + ), + pytest.param( + "gemini-2.5-flash", + {"reasoning_effort": {"effort": "medium", "summary": "detailed"}}, + False, + True, + id="budget-effort-summary-beats-off", + ), + pytest.param( + "gemini-3.8-flash", + {"thinking": {"type": "enabled", "budget_tokens": 2048, "summary": "auto"}}, + False, + True, + id="level-thinking-summary-beats-off", + ), + pytest.param( + "gemini-2.5-flash", + {"thinking": {"type": "enabled", "budget_tokens": 2048, "summary": "detailed"}}, + False, + True, + id="budget-thinking-summary-beats-off", + ), + pytest.param( + "gemini-3.8-flash", + {"reasoning_effort": "medium", "reasoning_summary": "auto"}, + False, + True, + id="level-reasoning-summary-alias-beats-off", + ), + pytest.param( + "gemini-2.5-flash", + {"reasoning_effort": "medium", "reasoningSummary": "detailed"}, + False, + True, + id="budget-reasoning-summary-camel-alias-beats-off", + ), + pytest.param( + "gemini-3.8-flash", + {"reasoning_effort": "none", "reasoning_summary": "auto"}, + True, + False, + id="effort-none-ignores-summary-alias", + ), + pytest.param( + "gemini-3.8-flash", + {"reasoning_effort": {"effort": "medium"}}, + False, + False, + id="effort-without-summary-keeps-off", + ), + pytest.param( + "gemini-3.8-flash", + {"reasoning_effort": {"effort": "medium"}, "reasoning_summary": "auto"}, + False, + True, + id="effort-dict-takes-the-summary-alias", + ), + pytest.param( + "gemini-3.8-flash", + {"reasoning_effort": {"effort": "none", "summary": "auto"}}, + True, + False, + id="effort-none-ignores-summary", + ), +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("model", "thinking_params", "gateway_default", "expected"), PER_REQUEST_THOUGHT_SUMMARY_ASKS) +async def test_gemini_per_request_thought_summary_ask_beats_gateway_default( + model: str, + thinking_params: dict[str, object], + gateway_default: bool, + expected: bool, + local_model_cost_map: None, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "gemini_include_thoughts_default", gateway_default) + + sent: Final = await _sent_thinking_config(model, thinking_params) + + assert sent["includeThoughts"] is expected + + +def _includes_thoughts_client(sent: list[bool]) -> AsyncHTTPHandler: + def upstream(request: httpx.Request) -> httpx.Response: + sent.append(json.loads(request.content)["generationConfig"]["thinkingConfig"]["includeThoughts"]) + return httpx.Response( + 200, + json={ + "candidates": [{"content": {"parts": [{"text": "hi"}], "role": "model"}, "finishReason": "STOP"}], + "usageMetadata": {"promptTokenCount": 1, "candidatesTokenCount": 1, "totalTokenCount": 2}, + }, + request=request, + ) + + client: Final = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(upstream)) + return client + + +async def _chat_summary_ask(client: AsyncHTTPHandler, asks: bool) -> None: + await litellm.acompletion( + model="gemini/gemini-3.8-flash", + messages=[{"role": "user", "content": "hi"}], + api_key="fake-gemini-key", + client=client, + reasoning_effort="medium", + **({"reasoning_summary": "auto"} if asks else {}), + ) + + +async def _messages_summary_ask(client: AsyncHTTPHandler, asks: bool) -> None: + await litellm.anthropic.messages.acreate( + model="gemini/gemini-3.8-flash", + messages=[{"role": "user", "content": "hi"}], + max_tokens=64, + api_key="fake-gemini-key", + client=client, + thinking={"type": "enabled", "budget_tokens": 2048, **({"display": "summarized"} if asks else {})}, + ) + + +async def _responses_summary_ask(client: AsyncHTTPHandler, asks: bool) -> None: + await litellm.aresponses( + model="gemini/gemini-3.8-flash", + input="hi", + api_key="fake-gemini-key", + client=client, + reasoning={"effort": "medium", **({"summary": "auto"} if asks else {})}, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "send", + (_chat_summary_ask, _messages_summary_ask, _responses_summary_ask), + ids=("chat-reasoning-summary-alias", "messages-display-summarized", "responses-reasoning-summary"), +) +async def test_gemini_summary_ask_turns_thoughts_on_and_logs_the_same_reasoning_effort( + send: Callable[[AsyncHTTPHandler, bool], Awaitable[None]], + local_model_cost_map: None, + monkeypatch: pytest.MonkeyPatch, +) -> None: + logged: Final[list[object]] = [] + + class ReasoningEffortLogger(CustomLogger): + async def async_log_success_event( + self, + kwargs: Mapping[str, Mapping[str, Mapping[str, object]]], + response_obj: object, + start_time: object, + end_time: object, + ) -> None: + logged.append(kwargs["standard_logging_object"]["model_parameters"].get("reasoning_effort")) + + monkeypatch.setattr(litellm, "gemini_include_thoughts_default", False) + monkeypatch.setattr(litellm, "callbacks", [ReasoningEffortLogger()]) + sent: Final[list[bool]] = [] + client: Final = _includes_thoughts_client(sent) + + await send(client, False) + await send(client, True) + await asyncio.sleep(0) + await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10.0) + + assert sent == [False, True] + assert len(logged) == 2 + assert isinstance(logged[0], str) + assert logged[1] == logged[0] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("gateway_default", "expected"), + (("false", False), (" Off ", False), ("0", False), ("true", True), (None, True), ("not-a-flag", True)), +) +async def test_gemini_sends_a_boolean_for_a_gateway_default_written_as_text_or_null( + gateway_default: object, expected: bool, local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "gemini_include_thoughts_default", gateway_default) + + sent: Final = await _sent_thinking_config("gemini-3.8-flash", {"reasoning_effort": "high"}) + + assert sent["includeThoughts"] is expected diff --git a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py index 0d55fc35e1d..19bbb6481fb 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -3109,7 +3109,6 @@ class TestToolTransformation: ("bedrock/converse/global.anthropic.claude-sonnet-5", "bedrock_converse"), ("anthropic.claude-sonnet-4-5-20250929-v1:0", "bedrock"), ("claude-sonnet-5", "vertex_ai"), - ("gemini-3.1-pro-preview", "vertex_ai"), ("moonshotai.kimi-k2-thinking", "bedrock_mantle"), ], ) @@ -3117,8 +3116,8 @@ class TestToolTransformation: """ A Responses request carrying ``reasoning.summary`` must still reach a chat provider as a plain ``reasoning_effort`` string. ``summary`` is Responses-only, and forwarding the whole - object turns reasoning off: Bedrock Converse and Vertex silently discard a non-string - ``reasoning_effort``, and Bedrock Mantle rejects the request outright. + object turns reasoning off: Bedrock Converse and Claude on Vertex silently discard a + non-string ``reasoning_effort``, and Bedrock Mantle rejects the request outright. """ responses_api_request = {"reasoning": {"effort": "medium", "summary": "auto"}} @@ -3131,6 +3130,42 @@ class TestToolTransformation: assert result["reasoning_effort"] == "medium" + @pytest.mark.parametrize( + "model, custom_llm_provider", + [("gemini-3.1-pro-preview", "vertex_ai"), ("gemini-2.5-flash", "gemini")], + ) + @pytest.mark.parametrize( + "reasoning, sends_thoughts", + [({"effort": "medium", "summary": "auto"}, True), ({"effort": "medium"}, False)], + ) + def test_a_summary_ask_turns_gemini_thoughts_on_over_the_gateway_default( + self, + monkeypatch: pytest.MonkeyPatch, + model: str, + custom_llm_provider: str, + reasoning: dict[str, str], + sends_thoughts: bool, + ) -> None: + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig + from litellm.utils import reasoning_effort_with_summary_alias + + monkeypatch.setattr(litellm, "gemini_include_thoughts_default", False) + bridged: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( + model=model, + input="hi", + responses_api_request={"reasoning": reasoning}, + custom_llm_provider=custom_llm_provider, + ) + + gemini: Final = VertexGeminiConfig() + reasoning_effort: Final = reasoning_effort_with_summary_alias( + bridged["reasoning_effort"], bridged, accepts_summary=gemini.accepts_reasoning_summary_in_reasoning_effort + ) + mapped: Final = gemini.map_openai_params({"reasoning_effort": reasoning_effort}, {}, model, True) + + assert mapped["thinkingConfig"]["includeThoughts"] is sends_thoughts + assert bridged["reasoning_effort"] == "medium" + @pytest.mark.parametrize( "model, custom_llm_provider", [ From 43be1c2ae7f768d1599921ad61e022a5a4ea7876 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=AD=A6=E6=9C=AC=E9=BE=8D=E9=96=80?= Date: Fri, 16 Oct 2026 18:20:34 +0900 Subject: [PATCH 13/23] fix(router)!: keep RPM, region, and param checks when token counting fails (#44544) * fix(router): stop skipping rate-limit checks when token counting fails Router._pre_call_checks() caught any token_counter() exception and returned the full, unfiltered deployment list, bypassing the RPM, region, and invalid-params checks for every deployment - not just the context-window check that actually depends on the count. A content block token_counter() cannot parse (e.g. an invalid image_url detail value) silently disabled rate-limit enforcement for that request, with only a log line and no error surfaced to the caller. Now a counting failure only disables the context-window check; the other checks still run for every deployment as before. Input-token resolution is pulled into its own method and computed once before the loop, so nothing needs to track "did counting fail" with a mutable flag across iterations. * test(router): cover the async region check when token counting fails Replace the unit test of the private input-token resolver with a regression test through async_get_healthy_deployments, the path the proxy uses: an EU-pinned request whose content cannot be token-counted must not get the US deployment * test(router): cover the param filter when token counting fails and satisfy router coverage * fix(router): answer 400 when no deployment takes the requested param, honor allowlists and drop_params The pre-call param filter now reuses the router's param-acceptance helper, so allowed_openai_params, drop_params, and additional_drop_params keep a deployment instead of filtering it out. When every available deployment lacks a requested param, the router answers 400 UnsupportedParamsError right away instead of a retried 429. Deployment context windows are resolved once per call. * fix(router): keep the retried 429 when a cooling-down deployment takes the param * fix(router): steer to a deployment that takes the param before dropping it, keep context-window fallbacks * fix(router): steer by deployment capability, apply drop settings after region and RPM checks A deployment-level allowed_openai_params no longer counts as taking the param when another deployment does, the drop and allow settings now apply only when no deployment that passed the region, RPM, and context-window checks takes it, and an authenticating provider is checked by its declared model instead of always passing * test(router): make the github_copilot pre-call filter test hermetic in CI * test(router): cover pre-call region and response_format filters through the proxy * fix(router): judge the 400 against the request's own deployments, not a name lookup * fix(router): judge the 400 against the passed deployments when the group name lists none * fix(router): let a request's drop and allow settings replace the deployment's in the param check --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/router.py | 271 ++++++---- .../router_code_coverage.py | 7 + .../routing/test_end_user_region_routing.py | 95 +++- .../test_pre_call_response_format_filter.py | 161 ++++++ tests/unit/test_router/test_router.py | 472 +++++++++++++++++- 5 files changed, 905 insertions(+), 101 deletions(-) create mode 100644 tests/integration/routing/test_pre_call_response_format_filter.py diff --git a/litellm/router.py b/litellm/router.py index 96c96310d45..151a8a75dc3 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -329,6 +329,7 @@ from litellm.types.router import ( FallbackBudgetCheck, GuardrailTypedDict, LiteLLM_Params, + LiteLLMParamsTypedDict, MockRouterTestingParams, ModelGroupInfo, OptionalPreCallChecks, @@ -486,6 +487,7 @@ _CLAUDE_CODE_SESSION_ID_RE: Final = re.compile(r"^[a-zA-Z0-9_\-]{8,}$") _CLAUDE_CODE_SESSION_ROUTER_TTL_SECONDS: Final = 3600 CLAUDE_CODE_SESSION_ROUTER_BINDING_TARGET: Final = "claude_code_session_router_binding" +_PRE_CALL_CHECK_FILTERED_PARAMS: Final = frozenset({"response_format"}) _RUNTIME_TOGGLEABLE_PRE_CALL_CHECKS: Final[Mapping[str, type[CustomLogger]]] = MappingProxyType( { "prompt_caching": PromptCachingDeploymentCheck, @@ -12691,16 +12693,34 @@ class Router: str(deployment_params.get("model") or ""), deployment_params.get("custom_llm_provider") ): return True + return Router._deployment_supports_param(deployment, group, param) + + @staticmethod + def _capability_model_and_provider(deployment_params: LiteLLMParamsTypedDict, group: str) -> tuple[str, str]: + model_label: Final = str(deployment_params.get("model") or group) + authenticating_provider: Final = declared_authenticating_provider( + model_label, deployment_params.get("custom_llm_provider") + ) + if authenticating_provider is not None: + return model_label.removeprefix(f"{authenticating_provider}/"), authenticating_provider + model, custom_llm_provider, _, _ = litellm.get_llm_provider( + model=model_label, + custom_llm_provider=deployment_params.get("custom_llm_provider"), + litellm_params=LiteLLM_Params.model_validate(deployment_params), + ) + return model, custom_llm_provider + + @staticmethod + def _deployment_supports_param(deployment: DeploymentTypedDict, group: str, param: str) -> bool: + deployment_params: Final = deployment.get("litellm_params") + if not deployment_params: + return True deployment_model_info: Final = deployment.get("model_info") base_model: Final = ( deployment_model_info.get("base_model") if deployment_model_info else None ) or deployment_params.get("base_model") try: - model, custom_llm_provider, _, _ = litellm.get_llm_provider( - model=deployment_params.get("model") or group, - custom_llm_provider=deployment_params.get("custom_llm_provider"), - litellm_params=LiteLLM_Params.model_validate(deployment_params), - ) + model, custom_llm_provider = Router._capability_model_and_provider(deployment_params, group) supported: Final = litellm.get_supported_openai_params( model=model, custom_llm_provider=custom_llm_provider, @@ -12708,7 +12728,7 @@ class Router: ) except Exception as e: # noqa: BLE001 # best-effort filter: an unresolvable provider must not narrow the request verbose_router_logger.debug( - "litellm.router.py::_deployment_accepts_param: keeping %s for model=%s. Got - %s", param, group, e + "litellm.router.py::_deployment_supports_param: keeping %s for model=%s. Got - %s", param, group, e ) return True return supported is None or param in supported @@ -13205,7 +13225,8 @@ class Router: Returns None when no deployment limits its context window, and when counting fails. The caller pairs this with ``skip_inline_token_count`` so neither case puts the count back on - the loop: a failed count leaves the deployments unfiltered, exactly as before. + the loop: a failed count only skips the context-window check - RPM, region and + invalid-params checks still run for every deployment. """ if messages is None and input is None: return None @@ -13223,6 +13244,66 @@ class Router: ) return None + def _resolve_pre_call_check_input_tokens( + self, + messages: Sequence[Mapping[str, str]] | None, + input: str | Sequence[object] | None, + request_kwargs: Mapping[str, object] | None, + input_token_count: int | None, + count_inline: bool, + ) -> int | None: + if input_token_count is not None: + return input_token_count + if not count_inline: + return None + try: + return self._count_pre_call_check_tokens( + messages=cast(list[dict[str, str]] | None, messages), # cast-ok: forwarded to the sync counter + input=cast(str | list[object] | None, input), # cast-ok: forwarded to the sync counter + request_kwargs=request_kwargs, + ) + except Exception as e: + verbose_router_logger.error( + "litellm.router.py::_pre_call_checks: failed to count tokens. " + "Skipping the context-window check only; RPM/region/param checks still apply. Got - %s", + e, + ) + return None + + @staticmethod + def _params_neither_dropped_nor_allowed(params: Mapping[str, object], candidates: frozenset[str]) -> frozenset[str]: + if params.get("drop_params") is True: + return frozenset() + dropped: Final = params.get("additional_drop_params") + dropped_names: Final = ( + frozenset(entry for entry in dropped if isinstance(entry, str)) + if isinstance(dropped, list) + else frozenset() + ) + return candidates - Router._declared_param_allowlist(params) - dropped_names + + @staticmethod + def _deployment_lacks_params(deployment: DeploymentTypedDict, group: str, params: frozenset[str]) -> bool: + return any(not Router._deployment_supports_param(deployment, group, param) for param in params) + + @staticmethod + def _deployment_rejects_params( + deployment: DeploymentTypedDict, group: str, params: frozenset[str], request_kwargs: Mapping[str, object] + ) -> bool: + call_settings: Final[Mapping[str, object]] = {**(deployment.get("litellm_params") or {}), **request_kwargs} + return Router._deployment_lacks_params( + deployment, group, Router._params_neither_dropped_nor_allowed(call_settings, params) + ) + + @staticmethod + def _pre_call_check_requested_chat_params(request_kwargs: Mapping[str, object]) -> frozenset[str]: + requested: Final = litellm.utils.get_non_default_params(passed_params=dict(request_kwargs)) + return frozenset( + param + for param, value in requested.items() + if param in _PRE_CALL_CHECK_FILTERED_PARAMS and not isinstance(value, str) + ) + def _pre_call_checks( self, model: str, @@ -13232,6 +13313,7 @@ class Router: request_kwargs: dict | None = None, input_token_count: int | None = None, skip_inline_token_count: bool = False, + request_deployments: Sequence[DeploymentTypedDict] | None = None, ): """ Filter out model in model group, if: @@ -13250,13 +13332,6 @@ class Router: invalid_model_indices: Final = set() # Use set for O(1) membership checks - # Token counting (tiktoken) is the dominant on-loop cost for large prompts. - # Only count when a deployment actually declares max_input_tokens, and count - # at most once; for model groups with no context-window limit it is skipped. - # Async callers pass the count in, already computed off the event loop, and set - # skip_inline_token_count so a failed off-loop count is not retried back on the loop. - input_tokens: int | None = input_token_count - _context_window_error = False _potential_error_str = "" _rate_limit_error = False @@ -13265,6 +13340,24 @@ class Router: has_countable_input: Final = (messages is not None or input is not None) and not compaction_pending( request_kwargs ) + deployment_max_input_tokens: Final = tuple( + self._deployment_max_input_tokens(model, deployment) for deployment in _returned_deployments + ) + input_tokens: Final = self._resolve_pre_call_check_input_tokens( + messages=messages, + input=cast(str | list[object] | None, input), # cast-ok: forwarded to the resolver + request_kwargs=cast(Mapping[str, object] | None, request_kwargs), # cast-ok: forwarded to the resolver + input_token_count=input_token_count, + count_inline=has_countable_input + and not skip_inline_token_count + and any(limit is not None for limit in deployment_max_input_tokens), + ) + request_settings: Final[Mapping[str, object]] = request_kwargs if request_kwargs is not None else {} + requested_params: Final[frozenset[str]] = ( + self._pre_call_check_requested_chat_params(request_settings) + if request_kwargs is not None and litellm.drop_params is False + else frozenset() + ) ## get model group RPM ## dt: Final = get_utc_datetime() @@ -13278,39 +13371,22 @@ class Router: _litellm_params = deployment.get("litellm_params", {}) _model_info = deployment.get("model_info", {}) - # see if we have the info for this model - _deployment_model = None # per-deployment model name (avoids overwriting the outer `model` group name) - try: - base_model = _model_info.get("base_model", None) - if base_model is None: - base_model = _litellm_params.get("base_model", None) - _deployment_model = base_model or _litellm_params.get("model", None) - model_info = self.get_router_model_info(deployment=deployment, received_model_name=model) - - max_input_tokens = model_info.get("max_input_tokens") if isinstance(model_info, dict) else None - if isinstance(max_input_tokens, int) and has_countable_input: - if input_tokens is None: - if skip_inline_token_count: - return _returned_deployments - try: - input_tokens = self._count_pre_call_check_tokens( - messages=messages, input=input, request_kwargs=request_kwargs - ) - except Exception as e: - verbose_router_logger.error( - "litellm.router.py::_pre_call_checks: failed to count tokens. Returning initial list of deployments. Got - %s", - e, - ) - return _returned_deployments - if input_tokens > max_input_tokens: - invalid_model_indices.add(idx) - _context_window_error = True - _potential_error_str += ( - f"Model={_deployment_model}, Max Input Tokens={max_input_tokens}, Got={input_tokens}" - ) - continue - except Exception as e: - verbose_router_logger.exception("An error occurs - %s", e) + max_input_tokens = deployment_max_input_tokens[idx] + if ( + max_input_tokens is not None + and has_countable_input + and input_tokens is not None + and input_tokens > max_input_tokens + ): + invalid_model_indices.add(idx) + _context_window_error = True + deployment_model = ( + _model_info.get("base_model") or _litellm_params.get("base_model") or _litellm_params.get("model") + ) + _potential_error_str += ( + f"Model={deployment_model}, Max Input Tokens={max_input_tokens}, Got={input_tokens}" + ) + continue model_id = _model_info.get("id", "") ## RPM CHECK ## @@ -13342,53 +13418,57 @@ class Router: invalid_model_indices.add(idx) continue - ## INVALID PARAMS ## -> catch 'gpt-3.5-turbo-16k' not supporting 'response_format' param - if request_kwargs is not None and litellm.drop_params is False: - # get supported params — use per-deployment model to avoid overwriting the outer model group name - _dep_model_for_params: str = _deployment_model or model - try: - ( - _dep_model_for_params, - custom_llm_provider, - _, - _, - ) = litellm.get_llm_provider( - model=_dep_model_for_params, - litellm_params=LiteLLM_Params.model_validate(_litellm_params), - ) - except Exception as e: # noqa: BLE001 # best-effort filter: an unresolvable provider must not fail the request - verbose_router_logger.debug( - "litellm.router.py::_pre_call_checks: skipping supported-params check for model=%s. Got - %s", - _dep_model_for_params, - e, - ) - continue + typed_deployments: Final = tuple( + cast(DeploymentTypedDict, deployment) # cast-ok: router deployments are plain dicts + for deployment in _returned_deployments + ) + param_candidates: Final = tuple( + (idx, deployment) for idx, deployment in enumerate(typed_deployments) if idx not in invalid_model_indices + ) + candidates_lacking_params: Final = tuple( + self._deployment_lacks_params(deployment, model, requested_params) for _, deployment in param_candidates + ) + no_deployment_takes_params: Final = bool(candidates_lacking_params) and all(candidates_lacking_params) + candidates_failing_params: Final = ( + tuple( + self._deployment_rejects_params(deployment, model, requested_params, request_settings) + for _, deployment in param_candidates + ) + if no_deployment_takes_params + else candidates_lacking_params + ) + excluded_indices: Final = frozenset(invalid_model_indices) | frozenset( + idx for (idx, _), fails in zip(param_candidates, candidates_failing_params, strict=True) if fails + ) - supported_openai_params = litellm.get_supported_openai_params( - model=_dep_model_for_params, - custom_llm_provider=custom_llm_provider, - ) - - if supported_openai_params is None: - continue - else: - # check the non-default openai params in request kwargs - non_default_params = litellm.utils.get_non_default_params(passed_params=request_kwargs) - special_params = ["response_format"] - # check if all params are supported - for k, v in non_default_params.items(): - if k not in supported_openai_params and k in special_params: - # if not -> invalid model - verbose_router_logger.debug("INVALID MODEL INDEX @ REQUEST KWARG FILTERING, k=%s", k) - invalid_model_indices.add(idx) - - if len(invalid_model_indices) == len(_returned_deployments): + if len(excluded_indices) == len(_returned_deployments): """ - no healthy deployments available b/c context window checks or rate limit error - First check for rate limit errors (if this is true, it means the model passed the context window check but failed the rate limit check) """ + group_deployments: Final = ( + tuple(request_deployments) + if request_deployments is not None + else (tuple(self.get_all_deployments(model_name=model)) or typed_deployments) + ) + if ( + not _context_window_error + and group_deployments + and all( + self._deployment_rejects_params(deployment, model, requested_params, request_settings) + for deployment in group_deployments + ) + ): + raise litellm.UnsupportedParamsError( + message=( + f"no deployment of model={model} supports {sorted(requested_params)}. " + "To drop these, set `litellm.drop_params=True` or for proxy `litellm_settings: drop_params: true`, " + f"or send allowed_openai_params={sorted(requested_params)} to pass them through" + ), + model=model, + ) if _rate_limit_error is True: # allow generic fallback logic to take place raise RouterRateLimitErrorBasic( model=model, @@ -13400,9 +13480,9 @@ class Router: model=model, llm_provider="", ) - if len(invalid_model_indices) > 0: + if len(excluded_indices) > 0: # Single-pass filter using set for O(1) lookups (avoids O(n^2) from repeated pops) - _returned_deployments = [d for i, d in enumerate(_returned_deployments) if i not in invalid_model_indices] + _returned_deployments = [d for i, d in enumerate(_returned_deployments) if i not in excluded_indices] return _returned_deployments @@ -13825,6 +13905,10 @@ class Router: ) return healthy_deployments + request_deployments: Final = cast( # cast-ok: router deployments are plain dicts + Sequence[DeploymentTypedDict], healthy_deployments + ) + # Health-check-based filtering (before cooldown) healthy_deployments = await self._async_filter_health_check_unhealthy_deployments( healthy_deployments=healthy_deployments, @@ -13884,6 +13968,7 @@ class Router: request_kwargs=request_kwargs, ), skip_inline_token_count=True, + request_deployments=request_deployments, ) # check if user wants to do tag based routing healthy_deployments = await get_deployments_for_tag( @@ -15050,6 +15135,9 @@ class Router: return healthy_deployments parent_otel_span: Final[Span | None] = get_parent_otel_span_from_kwargs(request_kwargs) + request_deployments: Final = cast( # cast-ok: router deployments are plain dicts + Sequence[DeploymentTypedDict], healthy_deployments + ) # Health-check-based filtering (before cooldown) healthy_deployments = self._filter_health_check_unhealthy_deployments( @@ -15081,6 +15169,7 @@ class Router: messages=messages, input=input, request_kwargs=request_kwargs, + request_deployments=request_deployments, ) ## ORDER FILTERING ## -> if user set 'order' in deployments, return deployments with lowest order (e.g. order=1 > order=2) @@ -15244,6 +15333,9 @@ class Router: # 4. Apply health-check and cooldown filtering parent_otel_span: Final[Span | None] = get_parent_otel_span_from_kwargs(request_kwargs) + request_deployments: Final = cast( # cast-ok: router deployments are plain dicts + Sequence[DeploymentTypedDict], pass_through_deployments + ) pass_through_deployments = self._filter_health_check_unhealthy_deployments( healthy_deployments=pass_through_deployments, parent_otel_span=parent_otel_span, @@ -15265,6 +15357,7 @@ class Router: messages=messages, input=input, request_kwargs=request_kwargs, + request_deployments=request_deployments, ) if self._is_priority_routing_group(model): diff --git a/tests/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py index 2fd2549c0ac..a3ed70f4a56 100644 --- a/tests/code_coverage_tests/router_code_coverage.py +++ b/tests/code_coverage_tests/router_code_coverage.py @@ -114,6 +114,13 @@ ignored_function_names = [ "_deployment_that_just_failed", # Tested through the same-boundary hop test in test_router_order_fallback.py "_collect_ptu_share_teams", # Cached behind team_ptu_ceiling, tested through the live model list test in test_router.py "_resolve_team_ptu_ceiling", # Cached behind team_ptu_ceiling, tested through the live model list test in test_router.py + "_resolve_pre_call_check_input_tokens", # Tested through the token-count failure pre-call check tests in test_router.py + "_params_neither_dropped_nor_allowed", # Tested through the allowed-or-dropped param pre-call check tests in test_router.py + "_deployment_rejects_params", # Tested through the unsupported-param pre-call check tests in test_router.py + "_deployment_lacks_params", # Tested through the unsupported-param pre-call check tests in test_router.py + "_pre_call_check_requested_chat_params", # Tested through the unsupported-param pre-call check tests in test_router.py + "_deployment_supports_param", # Tested through the deployment-capability pre-call check tests in test_router.py + "_capability_model_and_provider", # Tested through the authenticating-provider pre-call check test in test_router.py "_exact_base_model_max_input_tokens", # Tested through get_router_model_info in test_router_model_cost_isolation.py "_base_model_widened_window", # Tested through get_router_model_info in test_router_model_cost_isolation.py "_provider_model_id", # Tested through get_router_model_info in test_router_model_cost_isolation.py diff --git a/tests/integration/routing/test_end_user_region_routing.py b/tests/integration/routing/test_end_user_region_routing.py index 16638d04bd8..d8ba2da042a 100644 --- a/tests/integration/routing/test_end_user_region_routing.py +++ b/tests/integration/routing/test_end_user_region_routing.py @@ -1,5 +1,5 @@ import uuid -from collections.abc import Iterator +from collections.abc import Callable, Iterator from pathlib import Path from typing import Final @@ -8,12 +8,17 @@ import pytest import yaml from integration._support.client import Gateway, gateway_from_environment from integration._support.process import owned_proxy +from pydantic import JsonValue pytestmark: Final = pytest.mark.timeout(180) MODEL: Final = "regional-model" UPSTREAM_BY_REGION: Final = {"eu": "regional-eu-upstream", "us": "regional-us-upstream"} CALLS: Final = 5 +UNCOUNTABLE_MODEL: Final = "regional-uncountable-model" +PINNED_REGION: Final = "eu" +WEIGHT_BY_REGION: Final = {"eu": 1, "us": 99} +SILENT_WAV: Final = "UklGRiQAAABXQVZFZm10IBAAAAABAAEAQB8AAIA+AAACABAAZGF0YQAAAAA=" @pytest.fixture(scope="module") @@ -22,16 +27,32 @@ def candidate(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Gateway]: with gateway_from_environment() as base: config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) config["model_list"] = [ - { - "model_name": MODEL, - "litellm_params": { - "model": f"openai/{upstream}", - "api_base": f"{base.upstream_url}/v1", - "api_key": "synthetic-region-key", - "region_name": region, - }, - } - for region, upstream in UPSTREAM_BY_REGION.items() + *( + { + "model_name": MODEL, + "litellm_params": { + "model": f"openai/{upstream}", + "api_base": f"{base.upstream_url}/v1", + "api_key": "synthetic-region-key", + "region_name": region, + }, + } + for region, upstream in UPSTREAM_BY_REGION.items() + ), + *( + { + "model_name": UNCOUNTABLE_MODEL, + "litellm_params": { + "model": f"openai/{upstream}", + "api_base": f"{base.upstream_url}/v1", + "api_key": "synthetic-region-key", + "region_name": region, + "weight": WEIGHT_BY_REGION[region], + }, + "model_info": {"id": f"{UNCOUNTABLE_MODEL}-{region}", "max_input_tokens": 128000}, + } + for region, upstream in UPSTREAM_BY_REGION.items() + ), ] config["router_settings"] = {**config["router_settings"], "enable_pre_call_checks": True} path: Final = directory / "region-routing.yaml" @@ -70,3 +91,55 @@ def test_an_end_users_allowed_region_pins_every_call_to_that_regions_deployment( assert [response.headers.get("x-litellm-model-region") for response in responses] == [region] * CALLS observed: Final = upstream.get("/__observations").json()["requests"] assert [request["body"]["model"] for request in observed] == [UPSTREAM_BY_REGION[region]] * CALLS + + +def _uncountable_chat(marker: str, end_user: str) -> tuple[str, dict[str, JsonValue]]: + return ( + "/v1/chat/completions", + { + "model": UNCOUNTABLE_MODEL, + "user": end_user, + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": f"region {marker}"}, + {"type": "input_audio", "input_audio": {"data": SILENT_WAV, "format": "wav"}}, + ], + } + ], + }, + ) + + +def _uncountable_embedding(marker: str, end_user: str) -> tuple[str, dict[str, JsonValue]]: + return "/v1/embeddings", {"model": UNCOUNTABLE_MODEL, "user": end_user, "input": [f"region {marker}", marker]} + + +@pytest.mark.parametrize( + ("build", "recorded_chat_calls"), + [ + pytest.param(_uncountable_chat, CALLS, id="chat-input-audio"), + pytest.param(_uncountable_embedding, 0, id="embeddings-list-input"), + ], +) +def test_an_uncountable_prompt_still_pins_an_end_user_to_their_region( + candidate: Gateway, build: Callable[[str, str], tuple[str, dict[str, JsonValue]]], recorded_chat_calls: int +) -> None: + with ( + candidate.scenario() as scenario, + httpx.Client(base_url=candidate.upstream_url, timeout=5, trust_env=False) as upstream, + ): + end_user: Final = f"integration-end-user-{uuid.uuid4().hex}" + candidate.post("/end_user/new", {"user_id": end_user, "allowed_model_region": PINNED_REGION}) + scenario.cleanups.callback(candidate.post, "/end_user/delete", {"user_ids": [end_user]}) + key: Final = scenario.key(models=[UNCOUNTABLE_MODEL]) + upstream.get("/__observations").raise_for_status() + requests: Final = tuple(build(uuid.uuid4().hex, end_user) for _ in range(CALLS)) + responses: Final = tuple(candidate.request("POST", path, body, key=key) for path, body in requests) + assert [response.status_code for response in responses] == [200] * CALLS, [r.text for r in responses] + pinned_deployment: Final = f"{UNCOUNTABLE_MODEL}-{PINNED_REGION}" + assert [response.headers.get("x-litellm-model-id") for response in responses] == [pinned_deployment] * CALLS + observed: Final = upstream.get("/__observations").json()["requests"] + pinned_upstream: Final = UPSTREAM_BY_REGION[PINNED_REGION] + assert [request["body"]["model"] for request in observed] == [pinned_upstream] * recorded_chat_calls diff --git a/tests/integration/routing/test_pre_call_response_format_filter.py b/tests/integration/routing/test_pre_call_response_format_filter.py new file mode 100644 index 00000000000..24a4434fb5a --- /dev/null +++ b/tests/integration/routing/test_pre_call_response_format_filter.py @@ -0,0 +1,161 @@ +import uuid +from collections.abc import Iterator, Mapping +from pathlib import Path +from typing import Final + +import httpx +import pytest +import yaml +from integration._support.client import Gateway, gateway_from_environment +from integration._support.process import owned_proxy +from pydantic import JsonValue + +pytestmark: Final = pytest.mark.timeout(180) + +LACKS_JSON_MODE: Final = "gpt-4" +TAKES_JSON_MODE: Final = "gpt-4o-mini" +UNSUPPORTED: Final = "json-mode-unsupported" +DROPPING: Final = "json-mode-dropped" +MIXED: Final = "json-mode-mixed" +FALLBACK_SOURCE: Final = "json-mode-fallback-source" +SMALL_WINDOW: Final = "json-mode-small-window" +FALLBACK_TARGET: Final = "json-mode-fallback-target" +JSON_MODE: Final[Mapping[str, JsonValue]] = {"type": "json_object"} +CALLS: Final = 4 +SMALL_WINDOW_TOKENS: Final = 50 + + +def _deployment(group: str, upstream_model: str, upstream_url: str, **params: JsonValue) -> dict[str, JsonValue]: + return { + "model_name": group, + "litellm_params": { + "model": f"openai/{upstream_model}", + "api_base": f"{upstream_url}/v1", + "api_key": "synthetic-json-mode-key", + **params, + }, + } + + +@pytest.fixture(scope="module") +def candidate(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Gateway]: + directory: Final = tmp_path_factory.mktemp("json-mode-filter") + with gateway_from_environment() as base: + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + url: Final = base.upstream_url + config["model_list"] = [ + _deployment(UNSUPPORTED, LACKS_JSON_MODE, url), + _deployment(DROPPING, LACKS_JSON_MODE, url, drop_params=True), + _deployment(MIXED, LACKS_JSON_MODE, url), + _deployment(MIXED, TAKES_JSON_MODE, url), + _deployment(FALLBACK_SOURCE, LACKS_JSON_MODE, url), + { + **_deployment(SMALL_WINDOW, LACKS_JSON_MODE, url), + "model_info": {"max_input_tokens": SMALL_WINDOW_TOKENS}, + }, + _deployment(FALLBACK_TARGET, TAKES_JSON_MODE, url), + ] + config["router_settings"] = { + **config["router_settings"], + "enable_pre_call_checks": True, + "fallbacks": [{FALLBACK_SOURCE: [FALLBACK_TARGET]}], + "context_window_fallbacks": [{SMALL_WINDOW: [FALLBACK_TARGET]}], + } + path: Final = directory / "json-mode-filter.yaml" + path.write_text(yaml.safe_dump(config)) + with owned_proxy(base, directory, {}, config=path) as proxy: + yield proxy + + +def _json_mode_chat(group: str, text: str, **extra: JsonValue) -> dict[str, JsonValue]: + return { + "model": group, + "response_format": dict(JSON_MODE), + "messages": [{"role": "user", "content": f"Answer in JSON. {text}"}], + **extra, + } + + +def _drain(upstream: httpx.Client) -> list[dict[str, JsonValue]]: + response: Final = upstream.get("/__observations") + response.raise_for_status() + return [request["body"] for request in response.json()["requests"]] + + +@pytest.fixture +def upstream(candidate: Gateway) -> Iterator[httpx.Client]: + with httpx.Client(base_url=candidate.upstream_url, timeout=5, trust_env=False) as client: + _drain(client) + yield client + + +@pytest.mark.parametrize("stream", [False, True], ids=["non-streaming", "streaming"]) +def test_json_mode_answers_400_without_an_upstream_call_when_no_deployment_takes_it( + candidate: Gateway, upstream: httpx.Client, stream: bool +) -> None: + response: Final = candidate.request( + "POST", "/v1/chat/completions", _json_mode_chat(UNSUPPORTED, uuid.uuid4().hex, stream=stream) + ) + + assert response.status_code == 400, response.text + assert _drain(upstream) == [] + + +@pytest.mark.parametrize( + ("group", "extra", "forwarded_format"), + [ + pytest.param(DROPPING, {}, None, id="deployment-drop-params"), + pytest.param(UNSUPPORTED, {"allowed_openai_params": ["response_format"]}, JSON_MODE, id="request-allowlist"), + pytest.param(UNSUPPORTED, {"drop_params": True}, None, id="request-drop-params"), + ], +) +def test_json_mode_reaches_a_deployment_that_lacks_it_once_drop_or_allow_is_set( + candidate: Gateway, + upstream: httpx.Client, + group: str, + extra: Mapping[str, JsonValue], + forwarded_format: Mapping[str, JsonValue] | None, +) -> None: + response: Final = candidate.request( + "POST", "/v1/chat/completions", _json_mode_chat(group, uuid.uuid4().hex, **extra) + ) + + assert response.status_code == 200, response.text + received: Final = _drain(upstream) + assert [body["model"] for body in received] == [LACKS_JSON_MODE] + assert received[0].get("response_format") == forwarded_format + + +def test_json_mode_in_a_mixed_group_always_reaches_the_deployment_that_takes_it( + candidate: Gateway, upstream: httpx.Client +) -> None: + responses: Final = tuple( + candidate.request("POST", "/v1/chat/completions", _json_mode_chat(MIXED, uuid.uuid4().hex)) + for _ in range(CALLS) + ) + + assert [response.status_code for response in responses] == [200] * CALLS, [r.text for r in responses] + received: Final = _drain(upstream) + assert [(body["model"], body.get("response_format")) for body in received] == [ + (TAKES_JSON_MODE, dict(JSON_MODE)) + ] * CALLS + + +@pytest.mark.parametrize( + ("group", "text"), + [ + pytest.param(FALLBACK_SOURCE, "short", id="generic-fallback"), + pytest.param(SMALL_WINDOW, "padding " * (4 * SMALL_WINDOW_TOKENS), id="context-window-fallback"), + ], +) +def test_json_mode_falls_back_to_a_group_that_takes_it( + candidate: Gateway, upstream: httpx.Client, group: str, text: str +) -> None: + response: Final = candidate.request( + "POST", "/v1/chat/completions", _json_mode_chat(group, f"{uuid.uuid4().hex} {text}") + ) + + assert response.status_code == 200, response.text + assert [(body["model"], body.get("response_format")) for body in _drain(upstream)] == [ + (TAKES_JSON_MODE, dict(JSON_MODE)) + ] diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 5a6da70d756..3bd051d71d4 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -11,6 +11,7 @@ import threading import warnings from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Callable, Mapping from datetime import datetime, timedelta, timezone +from pathlib import Path from types import SimpleNamespace from typing import TYPE_CHECKING, Final, Literal, cast from unittest.mock import AsyncMock, MagicMock, Mock, patch @@ -82,6 +83,7 @@ from litellm.types.router import ( ModelInfo, PreRoutingHookResponse, RetryPolicy, + RouterRateLimitError, RoutingContext, ) from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER @@ -12190,7 +12192,7 @@ def test_pre_call_checks_uses_deployment_model_when_model_info_lookup_raises(mon model="custom-alias", healthy_deployments=deployments, messages=[{"role": "user", "content": "hi"}], - request_kwargs={}, + request_kwargs={"response_format": {"type": "json_object"}}, ) assert len(result) == 1 @@ -26764,3 +26766,471 @@ def test_completion_streams_a_wrapper_that_hands_iteration_to_another_iterator() streamed = list(response) assert streamed == chunks + + +_PRE_CALL_CHECK_UNPARSEABLE_MESSAGES: Final = [ + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": {"detail": "not-a-real-detail-value"}}, + ], + } +] + + +def test_pre_call_checks_filters_deployment_when_token_count_exceeds_limit(): + model_list: Final = [ + { + "model_name": "fake-openai-endpoint", + "litellm_params": {"model": "openai/my-fake-model", "api_key": "my-fake-key"}, + "model_info": {"id": "small", "max_input_tokens": 5}, + }, + { + "model_name": "fake-openai-endpoint", + "litellm_params": {"model": "openai/my-other-fake-model", "api_key": "my-fake-key"}, + "model_info": {"id": "large", "max_input_tokens": 10000}, + }, + ] + router: Final = Router(model_list=model_list, enable_pre_call_checks=True, num_retries=0) + + filtered: Final = router._pre_call_checks( + model="fake-openai-endpoint", + healthy_deployments=model_list, + messages=[{"role": "user", "content": " ".join(["word"] * 100)}], + ) + + assert [d["model_info"]["id"] for d in filtered] == ["large"] + + +def test_pre_call_checks_rpm_check_still_enforced_when_token_counting_fails(): + from litellm.types.router import RouterRateLimitErrorBasic + + model_list: Final = [ + { + "model_name": "fake-openai-endpoint", + "litellm_params": {"model": "openai/my-fake-model", "api_key": "my-fake-key", "rpm": 0}, + "model_info": {"max_input_tokens": 100}, + }, + ] + router: Final = Router(model_list=model_list, enable_pre_call_checks=True, num_retries=0) + + with pytest.raises(RouterRateLimitErrorBasic): + router._pre_call_checks( + model="fake-openai-endpoint", + healthy_deployments=model_list, + messages=_PRE_CALL_CHECK_UNPARSEABLE_MESSAGES, + ) + + +def test_pre_call_checks_context_window_check_is_skipped_but_not_fatal_when_token_counting_fails(): + model_list: Final = [ + { + "model_name": "fake-openai-endpoint", + "litellm_params": {"model": "openai/my-fake-model", "api_key": "my-fake-key"}, + "model_info": {"max_input_tokens": 100}, + }, + ] + router: Final = Router(model_list=model_list, enable_pre_call_checks=True, num_retries=0) + + filtered: Final = router._pre_call_checks( + model="fake-openai-endpoint", + healthy_deployments=model_list, + messages=_PRE_CALL_CHECK_UNPARSEABLE_MESSAGES, + ) + + assert len(filtered) == 1 + + +def test_pre_call_checks_keeps_unsupported_param_filter_when_token_counting_fails(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "drop_params", False) + no_json_mode_model: Final = "gpt-3.5-turbo-16k" + assert "response_format" not in ( + litellm.get_supported_openai_params(model=no_json_mode_model, custom_llm_provider="openai") or [] + ) + model_list: Final = [ + { + "model_name": "fake-openai-endpoint", + "litellm_params": {"model": f"openai/{deployment_model}", "api_key": "my-fake-key"}, + "model_info": {"id": deployment_model, "max_input_tokens": 100}, + } + for deployment_model in (no_json_mode_model, "gpt-4o") + ] + router: Final = Router(model_list=model_list, enable_pre_call_checks=True, num_retries=0) + + filtered: Final = router._pre_call_checks( + model="fake-openai-endpoint", + healthy_deployments=model_list, + messages=_PRE_CALL_CHECK_UNPARSEABLE_MESSAGES, + request_kwargs={"response_format": {"type": "json_object"}}, + ) + + assert [d["model_info"]["id"] for d in filtered] == ["gpt-4o"] + + +async def test_async_get_healthy_deployments_keeps_region_check_when_token_counting_fails(): + model_list: Final = [ + { + "model_name": "fake-openai-endpoint", + "litellm_params": {"model": "openai/my-fake-model", "api_key": "my-fake-key", "region_name": region}, + "model_info": {"id": f"{region}-deployment", "max_input_tokens": 100}, + } + for region in ("eu", "us") + ] + router: Final = Router(model_list=model_list, enable_pre_call_checks=True, num_retries=0) + + healthy: Final = await router.async_get_healthy_deployments( + model="fake-openai-endpoint", + request_kwargs={"allowed_model_region": "eu"}, + messages=_PRE_CALL_CHECK_UNPARSEABLE_MESSAGES, + input=None, + specific_deployment=False, + parent_otel_span=None, + ) + + assert [deployment["model_info"]["id"] for deployment in healthy] == ["eu-deployment"] + + +@pytest.mark.parametrize( + "messages", + [_PRE_CALL_CHECK_UNPARSEABLE_MESSAGES, [{"role": "user", "content": "hi"}]], + ids=["uncountable", "countable"], +) +def test_completion_answers_400_when_no_deployment_supports_the_requested_param( + monkeypatch: pytest.MonkeyPatch, messages: list[dict[str, object]] +): + monkeypatch.setattr(litellm, "drop_params", False) + router: Final = Router( + model_list=[ + { + "model_name": "no-json-mode", + "litellm_params": {"model": "openai/gpt-3.5-turbo-16k", "api_key": "my-fake-key"}, + "model_info": {"max_input_tokens": 100}, + } + ], + enable_pre_call_checks=True, + num_retries=2, + ) + + with pytest.raises(litellm.UnsupportedParamsError) as error: + router.completion( + model="no-json-mode", + messages=messages, + response_format={"type": "json_object"}, + mock_response="should not be served", + ) + + assert error.value.status_code == 400 + assert "allowed_openai_params" in str(error.value) + + +@pytest.mark.parametrize( + ("deployment_params", "request_kwargs"), + [ + ({}, {"allowed_openai_params": ["response_format"]}), + ({}, {"drop_params": True}), + ({}, {"additional_drop_params": ["response_format"]}), + ({"allowed_openai_params": ["response_format"]}, {}), + ({"drop_params": True}, {}), + ], + ids=["request-allowlist", "request-drop", "request-additional-drop", "deployment-allowlist", "deployment-drop"], +) +def test_pre_call_checks_keeps_deployment_whose_unsupported_param_is_allowed_or_dropped( + monkeypatch: pytest.MonkeyPatch, deployment_params: Mapping[str, object], request_kwargs: Mapping[str, object] +): + monkeypatch.setattr(litellm, "drop_params", False) + model_list: Final = [ + { + "model_name": "no-json-mode", + "litellm_params": {"model": "openai/gpt-3.5-turbo-16k", "api_key": "my-fake-key", **deployment_params}, + "model_info": {"id": "only", "max_input_tokens": 100}, + } + ] + router: Final = Router(model_list=model_list, enable_pre_call_checks=True, num_retries=0) + + filtered: Final = router._pre_call_checks( + model="no-json-mode", + healthy_deployments=model_list, + messages=[{"role": "user", "content": "hi"}], + request_kwargs={"response_format": {"type": "json_object"}, **request_kwargs}, + ) + + assert [d["model_info"]["id"] for d in filtered] == ["only"] + + +def test_pre_call_checks_keeps_retryable_429_when_a_cooling_down_deployment_supports_the_param(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "drop_params", False) + model_list: Final = [ + { + "model_name": "mixed-json-mode", + "litellm_params": {"model": f"openai/{deployment_model}", "api_key": "my-fake-key"}, + "model_info": {"id": deployment_model, "max_input_tokens": 100}, + } + for deployment_model in ("gpt-3.5-turbo-16k", "gpt-4o") + ] + router: Final = Router(model_list=model_list, enable_pre_call_checks=True, num_retries=0) + without_json_mode: Final = [d for d in router.model_list if d["model_info"]["id"] == "gpt-3.5-turbo-16k"] + + filtered: Final = router._pre_call_checks( + model="mixed-json-mode", + healthy_deployments=without_json_mode, + messages=[{"role": "user", "content": "hi"}], + request_kwargs={"response_format": {"type": "json_object"}}, + ) + + assert filtered == [] + + +def test_completion_uses_context_window_fallback_when_the_small_deployment_also_lacks_the_param(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "drop_params", False) + router: Final = Router( + model_list=[ + { + "model_name": "small", + "litellm_params": {"model": "openai/gpt-3.5-turbo-16k", "api_key": "my-fake-key"}, + "model_info": {"id": "small-id", "max_input_tokens": 5}, + }, + { + "model_name": "big", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "my-fake-key"}, + "model_info": {"id": "big-id"}, + }, + ], + enable_pre_call_checks=True, + context_window_fallbacks=[{"small": ["big"]}], + num_retries=0, + ) + + response: Final = router.completion( + model="small", + messages=[{"role": "user", "content": "hi " * 50}], + response_format={"type": "json_object"}, + mock_response="served by the fallback", + ) + + assert response._hidden_params["model_id"] == "big-id" + + +@pytest.mark.parametrize( + ("deployment_params", "request_kwargs"), + [ + ({}, {"allowed_openai_params": ["response_format"]}), + ({}, {"drop_params": True}), + ({}, {"additional_drop_params": ["response_format"]}), + ({"drop_params": True}, {}), + ({"allowed_openai_params": ["response_format"]}, {}), + ], + ids=["request-allowlist", "request-drop", "request-additional-drop", "deployment-drop", "deployment-allowlist"], +) +def test_pre_call_checks_prefers_a_deployment_that_takes_the_param_over_dropping_it( + monkeypatch: pytest.MonkeyPatch, deployment_params: Mapping[str, object], request_kwargs: Mapping[str, object] +): + monkeypatch.setattr(litellm, "drop_params", False) + model_list: Final = [ + { + "model_name": "mixed-json-mode", + "litellm_params": {"model": "openai/gpt-3.5-turbo-16k", "api_key": "my-fake-key", **deployment_params}, + "model_info": {"id": "gpt-3.5-turbo-16k", "max_input_tokens": 100}, + }, + { + "model_name": "mixed-json-mode", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "my-fake-key"}, + "model_info": {"id": "gpt-4o", "max_input_tokens": 100}, + }, + ] + router: Final = Router(model_list=model_list, enable_pre_call_checks=True, num_retries=0) + + filtered: Final = router._pre_call_checks( + model="mixed-json-mode", + healthy_deployments=model_list, + messages=[{"role": "user", "content": "hi"}], + request_kwargs={"response_format": {"type": "json_object"}, **request_kwargs}, + ) + + assert [d["model_info"]["id"] for d in filtered] == ["gpt-4o"] + + +def test_pre_call_checks_drops_the_param_when_the_capable_deployment_is_outside_the_region(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "drop_params", False) + model_list: Final = [ + { + "model_name": "regional-json-mode", + "litellm_params": {"model": "openai/gpt-3.5-turbo-16k", "api_key": "my-fake-key", "region_name": "eu"}, + "model_info": {"id": "eu-without-json-mode", "max_input_tokens": 100}, + }, + { + "model_name": "regional-json-mode", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "my-fake-key", "region_name": "us"}, + "model_info": {"id": "us-with-json-mode", "max_input_tokens": 100}, + }, + ] + router: Final = Router(model_list=model_list, enable_pre_call_checks=True, num_retries=0) + + filtered: Final = router._pre_call_checks( + model="regional-json-mode", + healthy_deployments=model_list, + messages=[{"role": "user", "content": "hi"}], + request_kwargs={"response_format": {"type": "json_object"}, "drop_params": True, "allowed_model_region": "eu"}, + ) + + assert [d["model_info"]["id"] for d in filtered] == ["eu-without-json-mode"] + + +def test_pre_call_checks_filters_an_authenticating_provider_deployment_that_lacks_the_param( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +): + monkeypatch.setattr(litellm, "drop_params", False) + monkeypatch.setenv("GITHUB_COPILOT_TOKEN_DIR", str(tmp_path)) + (tmp_path / "api-key.json").write_text(json.dumps({"token": "tid=test", "expires_at": 4102444800})) + model_list: Final = [ + { + "model_name": "mixed-json-mode", + "litellm_params": {"model": "github_copilot/gpt-4"}, + "model_info": {"id": "copilot-gpt-4", "max_input_tokens": 100}, + }, + { + "model_name": "mixed-json-mode", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "my-fake-key"}, + "model_info": {"id": "gpt-4o", "max_input_tokens": 100}, + }, + ] + router: Final = Router(model_list=model_list, enable_pre_call_checks=True, num_retries=0) + + filtered: Final = router._pre_call_checks( + model="mixed-json-mode", + healthy_deployments=model_list, + messages=[{"role": "user", "content": "hi"}], + request_kwargs={"response_format": {"type": "json_object"}}, + ) + + assert [d["model_info"]["id"] for d in filtered] == ["gpt-4o"] + + +def test_pre_call_checks_answers_400_when_the_group_name_lists_no_deployments(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "drop_params", False) + model_list: Final = [ + { + "model_name": "no-json-mode", + "litellm_params": {"model": "openai/gpt-3.5-turbo-16k", "api_key": "my-fake-key"}, + "model_info": {"id": "only", "max_input_tokens": 100}, + } + ] + router: Final = Router(model_list=model_list, enable_pre_call_checks=True, num_retries=0) + + with pytest.raises(litellm.UnsupportedParamsError): + router._pre_call_checks( + model="team-public-name", + healthy_deployments=model_list, + messages=[{"role": "user", "content": "hi"}], + request_kwargs={"response_format": {"type": "json_object"}}, + ) + + +def test_pre_call_checks_lets_a_request_setting_replace_the_deployment_setting(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "drop_params", False) + model_list: Final = [ + { + "model_name": "no-json-mode", + "litellm_params": {"model": "openai/gpt-3.5-turbo-16k", "api_key": "my-fake-key", "drop_params": True}, + "model_info": {"id": "drops-it", "max_input_tokens": 100}, + }, + { + "model_name": "no-json-mode", + "litellm_params": { + "model": "openai/gpt-4", + "api_key": "my-fake-key", + "allowed_openai_params": ["response_format"], + }, + "model_info": {"id": "allows-it", "max_input_tokens": 100}, + }, + ] + router: Final = Router(model_list=model_list, enable_pre_call_checks=True, num_retries=0) + + filtered: Final = router._pre_call_checks( + model="no-json-mode", + healthy_deployments=model_list, + messages=[{"role": "user", "content": "hi"}], + request_kwargs={"response_format": {"type": "json_object"}, "drop_params": False}, + ) + + assert [d["model_info"]["id"] for d in filtered] == ["allows-it"] + + +async def _select_team_json_deployment(router: Router, entrypoint: str, selection: Mapping[str, object]) -> object: + if entrypoint == "async": + return await router.async_get_available_deployment(**selection) + if entrypoint == "sync": + return router.get_available_deployment(**selection) + return router.get_available_deployment_for_pass_through(**selection) + + +@pytest.mark.parametrize("entrypoint", ["sync", "async", "pass-through"]) +@pytest.mark.parametrize( + ("deployment_models", "cooled_down", "expected_error"), + [ + pytest.param(("gpt-3.5-turbo-16k", "gpt-4"), (), litellm.UnsupportedParamsError, id="none-takes-it"), + pytest.param(("gpt-3.5-turbo-16k", "gpt-4o"), ("gpt-4o",), RouterRateLimitError, id="capable-cooling-down"), + ], +) +@pytest.mark.asyncio +async def test_team_public_model_name_answers_400_only_when_no_team_deployment_takes_the_param( + monkeypatch: pytest.MonkeyPatch, + entrypoint: str, + deployment_models: tuple[str, str], + cooled_down: tuple[str, ...], + expected_error: type[Exception], +): + monkeypatch.setattr(litellm, "drop_params", False) + router: Final = Router( + model_list=[ + { + "model_name": f"team-json-{deployment_model}", + "litellm_params": { + "model": f"openai/{deployment_model}", + "api_key": "my-fake-key", + "use_in_pass_through": True, + }, + "model_info": { + "id": deployment_model, + "max_input_tokens": 100, + "team_id": "json-team", + "team_public_model_name": "team-json", + }, + } + for deployment_model in deployment_models + ], + enable_pre_call_checks=True, + num_retries=0, + ) + _cool_down(router, *cooled_down) + selection: Final = { + "model": "team-json", + "messages": [{"role": "user", "content": "hi"}], + "request_kwargs": { + "metadata": {"user_api_key_team_id": "json-team"}, + "response_format": {"type": "json_object"}, + }, + } + + with pytest.raises(expected_error): + await _select_team_json_deployment(router, entrypoint, selection) + + +def test_pre_call_checks_ignores_an_image_response_format_string(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "drop_params", False) + model_list: Final = [ + { + "model_name": "image-model", + "litellm_params": {"model": "openai/gpt-3.5-turbo-16k", "api_key": "my-fake-key"}, + "model_info": {"id": "only", "max_input_tokens": 100}, + } + ] + router: Final = Router(model_list=model_list, enable_pre_call_checks=True, num_retries=0) + + filtered: Final = router._pre_call_checks( + model="image-model", + healthy_deployments=model_list, + messages=[{"role": "user", "content": "prompt"}], + request_kwargs={"response_format": "b64_json"}, + ) + + assert [d["model_info"]["id"] for d in filtered] == ["only"] From 642a2ad93456dbf446bd53e2cf1db7cdcaac4889 Mon Sep 17 00:00:00 2001 From: Marcus Date: Fri, 16 Oct 2026 06:30:59 -0400 Subject: [PATCH 14/23] fix(prompt-templates): treat initial_prompt_value and final_prompt_value as optional (#39762) * feat(dd_span_tagger): emit litellm_user_email span tag for JWT-authenticated requests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(dd_span_tagger): use dotted litellm.user_email tag for consistency Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * feat(model_hub): surface model_info.description in model group info and Model Hub UI Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(router): aggregate model group description without in-place mutation Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(prompt-templates): treat initial_prompt_value and final_prompt_value as optional completion() only stores a custom prompt template key when its value is truthy, and custom_prompt() already defaults both prompt values to "". Seven call sites read them back with direct indexing, so a template that sets only roles raises KeyError instead of rendering, on ollama, petals, vllm, predibase, codestral, anthropic text completion, and response_schema_prompt. Five other providers already read defensively, so the same config works on sagemaker and 500s on ollama. Read both keys with .get(key, ""), matching those five providers and the defaults custom_prompt() already declares. Templates that set all three keys render exactly as before. * test(prompt-templates): drop the explanatory docstrings from the optional-keys regression test * fix(health): skip background health check DB writes when the latest-row read fails Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(prompt-templates): read every custom prompt template field through one typed helper Templates that omit roles, initial_prompt_value, or final_prompt_value now render with empty defaults across every provider that reads them, instead of failing with KeyError or TypeError. A field with the wrong type fails with a validation error naming it. * fix(prompt-templates): keep templates that rendered on main working and fail bad fields with a 400 Only the roles a request uses are checked, null fields mean unset, and extra role keys are ignored, so every template main rendered renders the same. A wrong-type field now raises a BadRequestError naming it. initial_prompt_value is declared litellm-owned so it stays out of provider request bodies, like final_prompt_value. The template tests move into the mapped factory test file, and the predibase test asserts the real rendered prompt * fix(prompt-templates): keep bad-template failover and render non-mapping role entries like main A wrong-type template field now raises a 500 InternalServerError naming the field, so the router still retries and cools down the bad deployment. A roles value or role entry that is not a mapping renders unwrapped as it did on main, and role keys are only checked for roles a request uses. * fix(prompt-templates): warn when a custom prompt template roles value is not a mapping * fix(prompt-templates): keep deployments in rotation on a bad template and reject a template with nothing to render * test(integration): audit custom prompt templates across endpoints, providers, and chaos --------- Co-authored-by: milan Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: yassin Co-authored-by: ryan-crabbe-berri Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../prompt_templates/factory.py | 103 ++- .../anthropic/completion/transformation.py | 9 +- .../base_invoke_transformation.py | 11 +- litellm/llms/codestral/completion/handler.py | 9 +- .../llms/ollama/completion/transformation.py | 9 +- litellm/llms/petals/completion/handler.py | 9 +- litellm/llms/predibase/chat/transformation.py | 9 +- .../sagemaker/completion/transformation.py | 16 +- .../llms/vertex_ai/gemini/transformation.py | 1 + litellm/llms/vllm/completion/handler.py | 16 +- litellm/llms/watsonx/common_utils.py | 4 +- litellm/types/litellm_params.py | 1 + .../test_custom_prompt_template_chaos.py | 542 ++++++++++++++ .../test_custom_prompt_template_wire.py | 701 ++++++++++++++++++ ...llm_core_utils_prompt_templates_factory.py | 344 +++++++++ .../llms/test_predibase_transformation.py | 15 +- tests/unit/types/test_litellm_params.py | 1 + 17 files changed, 1720 insertions(+), 80 deletions(-) create mode 100644 tests/integration/providers/test_custom_prompt_template_chaos.py create mode 100644 tests/integration/providers/test_custom_prompt_template_wire.py diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 76adadb2e1f..07534e70a27 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -9,10 +9,11 @@ import xml.etree.ElementTree as ET from collections.abc import Container, Iterable, Iterator, Mapping, Sequence from enum import Enum from types import MappingProxyType -from typing import Any, Final, TypeAlias, TypedDict, cast, overload +from typing import Annotated, Any, Final, TypeAlias, TypedDict, cast, overload from jinja2.sandbox import ImmutableSandboxedEnvironment -from pydantic import BaseModel, JsonValue, TypeAdapter +from pydantic import BaseModel, BeforeValidator, ConfigDict, JsonValue, TypeAdapter, ValidationError, model_validator +from pydantic_core import ErrorDetails, PydanticCustomError import litellm import litellm.types @@ -5565,7 +5566,7 @@ def function_call_prompt( ] -def response_schema_prompt(model: str, response_schema: dict) -> str: +def response_schema_prompt(model: str, response_schema: dict, llm_provider: str) -> str: """ Decides if a user-defined custom prompt or default needs to be used @@ -5581,11 +5582,8 @@ def response_schema_prompt(model: str, response_schema: dict) -> str: custom_prompt_details = litellm.custom_prompt_dict["response_schema_prompt"] if custom_prompt_details is not None: - return custom_prompt( - role_dict=custom_prompt_details["roles"], - initial_prompt_value=custom_prompt_details["initial_prompt_value"], - final_prompt_value=custom_prompt_details["final_prompt_value"], - messages=response_schema_as_message, + return render_custom_prompt_template( + template=custom_prompt_details, messages=response_schema_as_message, model=model, llm_provider=llm_provider ) else: return default_response_schema_prompt(response_schema=response_schema) @@ -5605,8 +5603,95 @@ def default_response_schema_prompt(response_schema: dict) -> str: # Custom prompt template +def _mapping_or_none(value: object) -> object: + if value is None or issubclass(type(value), Mapping): + return value + verbose_logger.warning( + "Ignoring a custom prompt template roles value of type %s because it is not a mapping, " + "so those messages render without role markers", + type(value).__name__, + ) + return None + + +class CustomPromptRole(BaseModel): + model_config = ConfigDict(frozen=True) + + pre_message: str | None = None + post_message: str | None = None + + +class CustomPromptRoles(BaseModel): + model_config = ConfigDict(frozen=True, title="custom prompt template") + + roles: Mapping[str, Annotated[CustomPromptRole | None, BeforeValidator(_mapping_or_none)]] + + +class CustomPromptTemplate(BaseModel): + model_config = ConfigDict(frozen=True, title="custom prompt template") + + roles: Annotated[Mapping[object, object] | None, BeforeValidator(_mapping_or_none)] = None + initial_prompt_value: str | None = None + final_prompt_value: str | None = None + + @model_validator(mode="after") + def _sets_a_rendered_field(self) -> "CustomPromptTemplate": + if self.roles is None and self.initial_prompt_value is None and self.final_prompt_value is None: + raise PydanticCustomError( + "empty_custom_prompt_template", + "it sets none of roles, initial_prompt_value, or final_prompt_value", + ) + return self + + +def render_custom_prompt_template( + template: Mapping[str, object], + messages: Sequence[Mapping[str, object]], + model: str, + llm_provider: str, +) -> str: + try: + return _render_custom_prompt_template(template=template, messages=messages) + except ValidationError as error: + raise litellm.APIConnectionError( + message=f"The custom prompt template for {model} is invalid: {_describe_template_errors(error)}", + model=model, + llm_provider=llm_provider, + ) from error + + +def _describe_template_errors(error: ValidationError) -> str: + return "; ".join( + _describe_template_error(detail) for detail in error.errors(include_url=False, include_input=False) + ) + + +def _describe_template_error(detail: ErrorDetails) -> str: + location: Final = ".".join(str(part) for part in detail["loc"]) + return f"{location}: {detail['msg']}" if location else detail["msg"] + + +def _render_custom_prompt_template(template: Mapping[str, object], messages: Sequence[Mapping[str, object]]) -> str: + validated: Final = CustomPromptTemplate.model_validate(template) + message_roles: Final = frozenset(message.get("role") for message in messages) + template_roles: Final = validated.roles or {} + used_roles: Final = CustomPromptRoles.model_validate( + {"roles": {role: entry for role, entry in template_roles.items() if role in message_roles}} + ) + return custom_prompt( + role_dict={ + role: {"pre_message": entry.pre_message or "", "post_message": entry.post_message or ""} + for role, entry in used_roles.roles.items() + if entry is not None + }, + initial_prompt_value=validated.initial_prompt_value or "", + final_prompt_value=validated.final_prompt_value or "", + messages=list(messages), + ) + + def custom_prompt( - role_dict: dict, + role_dict: Mapping[str, Mapping[str, str]], messages: list, initial_prompt_value: str = "", final_prompt_value: str = "", diff --git a/litellm/llms/anthropic/completion/transformation.py b/litellm/llms/anthropic/completion/transformation.py index e3ad961ed76..73e94e85f04 100644 --- a/litellm/llms/anthropic/completion/transformation.py +++ b/litellm/llms/anthropic/completion/transformation.py @@ -14,8 +14,8 @@ import httpx import litellm from litellm.constants import DEFAULT_MAX_TOKENS from litellm.litellm_core_utils.prompt_templates.factory import ( - custom_prompt, prompt_factory, + render_custom_prompt_template, ) from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator from litellm.llms.base_llm.chat.transformation import ( @@ -244,11 +244,8 @@ class AnthropicTextConfig(BaseConfig): if model in custom_prompt_dict: # check if the model has a registered custom prompt model_prompt_details: Final = custom_prompt_dict[model] - prompt = custom_prompt( - role_dict=model_prompt_details["roles"], - initial_prompt_value=model_prompt_details["initial_prompt_value"], - final_prompt_value=model_prompt_details["final_prompt_value"], - messages=messages, + prompt = render_custom_prompt_template( + template=model_prompt_details, messages=messages, model=model, llm_provider="anthropic" ) else: prompt = prompt_factory(model=model, messages=messages, custom_llm_provider="anthropic") diff --git a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py index 1b3b0c344df..f026e4f9a99 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py @@ -14,9 +14,9 @@ from litellm.litellm_core_utils.get_litellm_params import stored_control_options from litellm.litellm_core_utils.logging_utils import track_llm_api_timing from litellm.litellm_core_utils.prompt_templates.factory import ( cohere_message_pt, - custom_prompt, deepseek_r1_pt, prompt_factory, + render_custom_prompt_template, ) from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException from litellm.llms.bedrock.chat.invoke_handler import make_call, make_sync_call @@ -601,7 +601,7 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM): return cast(litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL, provider) return None - def convert_messages_to_prompt(self, model, messages, provider, custom_prompt_dict) -> tuple[str, list | None]: + def convert_messages_to_prompt(self, model: str, messages, provider, custom_prompt_dict) -> tuple[str, list | None]: # handle anthropic prompts and amazon titan prompts prompt = "" chat_history: list | None = None @@ -609,11 +609,8 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM): if model in custom_prompt_dict: # check if the model has a registered custom prompt model_prompt_details: Final = custom_prompt_dict[model] - prompt = custom_prompt( - role_dict=model_prompt_details["roles"], - initial_prompt_value=model_prompt_details.get("initial_prompt_value", ""), - final_prompt_value=model_prompt_details.get("final_prompt_value", ""), - messages=messages, + prompt = render_custom_prompt_template( + template=model_prompt_details, messages=messages, model=model, llm_provider="bedrock" ) return prompt, None ## ELSE diff --git a/litellm/llms/codestral/completion/handler.py b/litellm/llms/codestral/completion/handler.py index f8486d3b274..08799c3a0a5 100644 --- a/litellm/llms/codestral/completion/handler.py +++ b/litellm/llms/codestral/completion/handler.py @@ -13,8 +13,8 @@ import litellm from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.litellm_core_utils.logging_utils import track_llm_api_timing from litellm.litellm_core_utils.prompt_templates.factory import ( - custom_prompt, prompt_factory, + render_custom_prompt_template, ) from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, @@ -261,11 +261,8 @@ class CodestralTextCompletion: if model in custom_prompt_dict: # check if the model has a registered custom prompt model_prompt_details: Final = custom_prompt_dict[model] - prompt = custom_prompt( - role_dict=model_prompt_details["roles"], - initial_prompt_value=model_prompt_details["initial_prompt_value"], - final_prompt_value=model_prompt_details["final_prompt_value"], - messages=messages, + prompt = render_custom_prompt_template( + template=model_prompt_details, messages=messages, model=model, llm_provider="codestral" ) else: prompt = prompt_factory(model=model, messages=messages) diff --git a/litellm/llms/ollama/completion/transformation.py b/litellm/llms/ollama/completion/transformation.py index 703d45acdf2..ac41310fb61 100644 --- a/litellm/llms/ollama/completion/transformation.py +++ b/litellm/llms/ollama/completion/transformation.py @@ -16,8 +16,8 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( ) from litellm.litellm_core_utils.prompt_templates.factory import ( convert_to_ollama_image, - custom_prompt, ollama_pt, + render_custom_prompt_template, ) from litellm.litellm_core_utils.prompt_templates.image_handling import ( async_inline_remote_media, @@ -437,11 +437,8 @@ class OllamaConfig(BaseConfig): if model in custom_prompt_dict: # check if the model has a registered custom prompt model_prompt_details: Final = custom_prompt_dict[model] - ollama_prompt = custom_prompt( - role_dict=model_prompt_details["roles"], - initial_prompt_value=model_prompt_details["initial_prompt_value"], - final_prompt_value=model_prompt_details["final_prompt_value"], - messages=messages, + ollama_prompt = render_custom_prompt_template( + template=model_prompt_details, messages=messages, model=model, llm_provider="ollama" ) elif text_completion_request: # handle `/completions` requests ollama_prompt = get_str_from_messages(messages=messages) diff --git a/litellm/llms/petals/completion/handler.py b/litellm/llms/petals/completion/handler.py index 425d3729d7c..a2ccd7eaf08 100644 --- a/litellm/llms/petals/completion/handler.py +++ b/litellm/llms/petals/completion/handler.py @@ -5,8 +5,8 @@ from typing import Final import litellm from litellm.litellm_core_utils.prompt_templates.factory import ( - custom_prompt, prompt_factory, + render_custom_prompt_template, ) from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, @@ -43,11 +43,8 @@ def completion( if model in litellm.custom_prompt_dict: # check if the model has a registered custom prompt model_prompt_details: Final = litellm.custom_prompt_dict[model] - prompt = custom_prompt( - role_dict=model_prompt_details["roles"], - initial_prompt_value=model_prompt_details["initial_prompt_value"], - final_prompt_value=model_prompt_details["final_prompt_value"], - messages=messages, + prompt = render_custom_prompt_template( + template=model_prompt_details, messages=messages, model=model, llm_provider="petals" ) else: prompt = prompt_factory(model=model, messages=messages) diff --git a/litellm/llms/predibase/chat/transformation.py b/litellm/llms/predibase/chat/transformation.py index 7bd05c6e30b..0dd7e27722e 100644 --- a/litellm/llms/predibase/chat/transformation.py +++ b/litellm/llms/predibase/chat/transformation.py @@ -10,8 +10,8 @@ import litellm from litellm.constants import DEFAULT_MAX_TOKENS from litellm.litellm_core_utils.core_helpers import map_finish_reason from litellm.litellm_core_utils.prompt_templates.factory import ( - custom_prompt, prompt_factory, + render_custom_prompt_template, ) from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException from litellm.types.llms.openai import AllMessageValues @@ -277,11 +277,8 @@ class PredibaseConfig(BaseConfig): custom_prompt_dict: Final = litellm_params.get("custom_prompt_dict", {}) if model in custom_prompt_dict: model_prompt_details: Final = custom_prompt_dict[model] - prompt = custom_prompt( - role_dict=model_prompt_details["roles"], - initial_prompt_value=model_prompt_details["initial_prompt_value"], - final_prompt_value=model_prompt_details["final_prompt_value"], - messages=messages, + prompt = render_custom_prompt_template( + template=model_prompt_details, messages=messages, model=model, llm_provider="predibase" ) else: prompt = prompt_factory(model=model, messages=messages) diff --git a/litellm/llms/sagemaker/completion/transformation.py b/litellm/llms/sagemaker/completion/transformation.py index 274f884de0b..a69846c816f 100644 --- a/litellm/llms/sagemaker/completion/transformation.py +++ b/litellm/llms/sagemaker/completion/transformation.py @@ -15,8 +15,8 @@ from httpx._models import Headers, Response import litellm from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.prompt_templates.factory import ( - custom_prompt, prompt_factory, + render_custom_prompt_template, ) from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException from litellm.types.llms.openai import AllMessageValues @@ -122,20 +122,14 @@ class SagemakerConfig(BaseConfig): if model in custom_prompt_dict: # check if the model has a registered custom prompt model_prompt_details = custom_prompt_dict[model] - prompt = custom_prompt( - role_dict=model_prompt_details.get("roles", None), - initial_prompt_value=model_prompt_details.get("initial_prompt_value", ""), - final_prompt_value=model_prompt_details.get("final_prompt_value", ""), - messages=messages, + prompt = render_custom_prompt_template( + template=model_prompt_details, messages=messages, model=model, llm_provider="sagemaker" ) elif hf_model_name in custom_prompt_dict: # check if the base huggingface model has a registered custom prompt model_prompt_details = custom_prompt_dict[hf_model_name] - prompt = custom_prompt( - role_dict=model_prompt_details.get("roles", None), - initial_prompt_value=model_prompt_details.get("initial_prompt_value", ""), - final_prompt_value=model_prompt_details.get("final_prompt_value", ""), - messages=messages, + prompt = render_custom_prompt_template( + template=model_prompt_details, messages=messages, model=model, llm_provider="sagemaker" ) else: if hf_model_name is None: diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index 895d59965b7..7b832b321da 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -1154,6 +1154,7 @@ def transform_request_body( user_response_schema_message: Final = response_schema_prompt( model=model, response_schema=optional_params.get("response_schema"), + llm_provider=custom_llm_provider, ) messages.append({"role": "user", "content": user_response_schema_message}) optional_params.pop("response_schema") diff --git a/litellm/llms/vllm/completion/handler.py b/litellm/llms/vllm/completion/handler.py index 78e6c74c2f7..f330f0928d0 100644 --- a/litellm/llms/vllm/completion/handler.py +++ b/litellm/llms/vllm/completion/handler.py @@ -5,8 +5,8 @@ from typing import Final import httpx from litellm.litellm_core_utils.prompt_templates.factory import ( - custom_prompt, prompt_factory, + render_custom_prompt_template, ) from litellm.utils import ModelResponse, Usage @@ -56,11 +56,8 @@ def completion( if model in custom_prompt_dict: # check if the model has a registered custom prompt model_prompt_details: Final = custom_prompt_dict[model] - prompt = custom_prompt( - role_dict=model_prompt_details["roles"], - initial_prompt_value=model_prompt_details["initial_prompt_value"], - final_prompt_value=model_prompt_details["final_prompt_value"], - messages=messages, + prompt = render_custom_prompt_template( + template=model_prompt_details, messages=messages, model=model, llm_provider="vllm" ) else: prompt = prompt_factory(model=model, messages=messages) @@ -144,11 +141,8 @@ def batch_completions(model: str, messages: list, optional_params=None, custom_p # check if the model has a registered custom prompt model_prompt_details: Final = custom_prompt_dict[model] for message in messages: - prompt = custom_prompt( - role_dict=model_prompt_details["roles"], - initial_prompt_value=model_prompt_details["initial_prompt_value"], - final_prompt_value=model_prompt_details["final_prompt_value"], - messages=message, + prompt = render_custom_prompt_template( + template=model_prompt_details, messages=message, model=model, llm_provider="vllm" ) prompts.append(prompt) else: diff --git a/litellm/llms/watsonx/common_utils.py b/litellm/llms/watsonx/common_utils.py index fb1fa8b9ffb..6d3686e9f85 100644 --- a/litellm/llms/watsonx/common_utils.py +++ b/litellm/llms/watsonx/common_utils.py @@ -141,7 +141,7 @@ async def _aconvert_watsonx_messages_core( model_prompt_dict: Final = custom_prompt_dict[model] return ptf.custom_prompt( messages=messages, - role_dict=model_prompt_dict.get("role_dict", model_prompt_dict.get("roles")), + role_dict=model_prompt_dict.get("role_dict", model_prompt_dict.get("roles", {})), initial_prompt_value=model_prompt_dict.get("initial_prompt_value", ""), final_prompt_value=model_prompt_dict.get("final_prompt_value", ""), bos_token=model_prompt_dict.get("bos_token", ""), @@ -173,7 +173,7 @@ def _convert_watsonx_messages_core( model_prompt_dict: Final = custom_prompt_dict[model] return ptf.custom_prompt( messages=messages, - role_dict=model_prompt_dict.get("role_dict", model_prompt_dict.get("roles")), + role_dict=model_prompt_dict.get("role_dict", model_prompt_dict.get("roles", {})), initial_prompt_value=model_prompt_dict.get("initial_prompt_value", ""), final_prompt_value=model_prompt_dict.get("final_prompt_value", ""), bos_token=model_prompt_dict.get("bos_token", ""), diff --git a/litellm/types/litellm_params.py b/litellm/types/litellm_params.py index 096ad1c30dd..8387afa22bb 100644 --- a/litellm/types/litellm_params.py +++ b/litellm/types/litellm_params.py @@ -260,6 +260,7 @@ class PromptOptions: prompt_label: str | None = None litellm_system_prompt: str | None = None custom_prompt_dict: Mapping[str, object] | None = None + initial_prompt_value: str | None = None roles: Mapping[str, object] | None = None final_prompt_value: str | None = None bos_token: str | None = None diff --git a/tests/integration/providers/test_custom_prompt_template_chaos.py b/tests/integration/providers/test_custom_prompt_template_chaos.py new file mode 100644 index 00000000000..c0da612af23 --- /dev/null +++ b/tests/integration/providers/test_custom_prompt_template_chaos.py @@ -0,0 +1,542 @@ +import asyncio +import json +import re +import threading +import uuid +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from pathlib import Path +from queue import SimpleQueue +from types import MappingProxyType +from typing import Final, Literal +from urllib.parse import urlsplit + +import httpx +import psutil +import pytest +import yaml +from integration._support.client import Gateway, eventually +from integration._support.database import read_rows +from integration._support.process import graceful_stop_seconds, owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +_BACKEND: Final = "llama3-custom-template-chaos" +_API_KEY: Final = "synthetic-ollama-key" +_CONFIG_MODEL: Final = "ollama-custom-template-chaos" +_SYSTEM: Final = "Answer in one word." +_ROLES: Final[dict[str, JsonValue]] = { + "system": {"pre_message": "", "post_message": ""}, + "user": {"pre_message": "", "post_message": ""}, +} +_REJECTION: Final = "initial_prompt_value: Input should be a valid string" +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_MARKER: Final = re.compile(r"marker-([0-9a-f]{32})") +_STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]") +_OWNED_PROXY_CELL_SECONDS: Final = 2 * graceful_stop_seconds() + 120 + +Endpoint = Literal["chat", "messages", "responses"] + + +@dataclass(frozen=True, slots=True) +class _Call: + endpoint: Endpoint + stream: bool + marker: str + rejected: bool = False + + +@dataclass(frozen=True, slots=True) +class _Served: + call: _Call + status: int + text: str + + +def _question(marker: str) -> str: + return f"What colour is the sky? marker-{marker}" + + +def _prompt(marker: str) -> str: + return f"{_SYSTEM}{_question(marker)}" + + +def _path(endpoint: Endpoint) -> str: + match endpoint: + case "chat": + return "/v1/chat/completions" + case "messages": + return "/v1/messages" + case "responses": + return "/v1/responses" + + +def _body(model: str, call: _Call) -> dict[str, JsonValue]: + common: Final[dict[str, JsonValue]] = { + "model": model, + "stream": call.stream, + "num_retries": 0, + "cache": {"no-cache": True}, + **({"initial_prompt_value": 5} if call.rejected else {}), + } + match call.endpoint: + case "chat": + return { + **common, + "messages": [ + {"role": "system", "content": _SYSTEM}, + {"role": "user", "content": _question(call.marker)}, + ], + } + case "messages": + return { + **common, + "max_tokens": 64, + "system": _SYSTEM, + "messages": [{"role": "user", "content": _question(call.marker)}], + } + case "responses": + return {**common, "store": False, "instructions": _SYSTEM, "input": _question(call.marker)} + + +def _generate_reply(marker: str, stream: bool, drop_connection: bool = False) -> Reply: + done: Final = { + "model": _BACKEND, + "created_at": "2026-10-15T00:00:00Z", + "response": "", + "done": True, + "done_reason": "stop", + "prompt_eval_count": 30, + "eval_count": 12, + } + if not stream: + return Reply( + body=json.dumps({**done, "response": f"blue marker-{marker}"}).encode(), drop_connection=drop_connection + ) + frames: Final = ( + *( + {"model": _BACKEND, "created_at": "2026-10-15T00:00:00Z", "response": piece, "done": False} + for piece in ("blue ", f"marker-{marker}") + ), + done, + ) + return Reply( + content_type="application/x-ndjson", + chunks=tuple(json.dumps(frame).encode() + b"\n" for frame in frames), + drop_connection=drop_connection, + ) + + +def _is_generate(request: Request) -> bool: + return (request.method, request.target) == ("POST", "/api/generate") + + +def _marker_of(request: Request) -> str: + found: Final = _MARKER.search(request.body.decode()) + assert found is not None, request.body + return found.group(1) + + +def _echo(request: Request) -> Reply: + if not _is_generate(request): + return Reply(body=b"{}") + stream: Final = _JSON_OBJECT.validate_json(request.body).get("stream") is True + return _generate_reply(_marker_of(request), stream) + + +def _assert_each_prompt_rendered_once(received: tuple[Request, ...], markers: frozenset[str]) -> None: + generates: Final = tuple(request for request in received if _is_generate(request)) + assert sorted(_marker_of(request) for request in generates) == sorted(markers) + for request in generates: + body: Final = _JSON_OBJECT.validate_json(request.body) + assert body["prompt"] == _prompt(_marker_of(request)), body["prompt"] + assert "roles" not in body and "initial_prompt_value" not in body, sorted(body) + + +def _response_id(served: _Served) -> str | None: + if served.call.endpoint == "responses": + return None + if not served.call.stream: + identity: Final = _JSON_OBJECT.validate_json(served.text)["id"] + assert isinstance(identity, str) + return identity + for line in served.text.splitlines(): + if not line.startswith("data: ") or line == "data: [DONE]": + continue + payload: Final = _JSON_OBJECT.validate_json(line.removeprefix("data: ")) + if served.call.endpoint == "chat": + first: Final = payload["id"] + assert isinstance(first, str) + return first + if payload.get("type") == "message_start": + message: Final = payload["message"] + assert isinstance(message, dict) and isinstance(message["id"], str) + return message["id"] + raise AssertionError(served.text) + + +def _spend_statuses(model: str, expected: int) -> MappingProxyType[str, str]: + rows: Final = eventually( + lambda: read_rows('SELECT request_id, status FROM "LiteLLM_SpendLogs" WHERE model_group=%s', (model,)), + lambda found: len(found) >= expected, + seconds=70, + ) + statuses: Final = MappingProxyType({str(row["request_id"]): str(row["status"]) for row in rows}) + assert len(statuses) == len(rows) == expected, rows + return statuses + + +def _successes(statuses: Mapping[str, str]) -> frozenset[str]: + return frozenset(identity for identity, status in statuses.items() if status == "success") + + +async def _send(client: httpx.AsyncClient, key: str, model: str, call: _Call) -> _Served: + async with client.stream( + "POST", + _path(call.endpoint), + json=_body(model, call), + headers={"Authorization": f"Bearer {key}", "anthropic-version": "2023-06-01"}, + ) as response: + raw: Final = await response.aread() + return _Served(call=call, status=response.status_code, text=raw.decode()) + + +async def _burst( + base_url: str, key: str, model: str, calls: tuple[_Call, ...], *, tolerate_transport_errors: bool = False +) -> tuple[_Served, ...]: + async with httpx.AsyncClient(base_url=base_url, timeout=60, trust_env=False) as client: + results: Final = await asyncio.gather( + *(_send(client, key, model, call) for call in calls), return_exceptions=tolerate_transport_errors + ) + for result in results: + assert not isinstance(result, BaseException) or isinstance(result, httpx.TransportError), repr(result) + return tuple(result for result in results if isinstance(result, _Served)) + + +def _calls(count: int, endpoints: tuple[Endpoint, ...], stream: Callable[[int], bool]) -> tuple[_Call, ...]: + return tuple( + _Call(endpoint=endpoints[index % len(endpoints)], stream=stream(index), marker=uuid.uuid4().hex) + for index in range(count) + ) + + +def _assert_answered_with_its_own_marker(served: _Served) -> None: + assert served.status == 200, served.text + assert set(_MARKER.findall(served.text)) == {served.call.marker}, served.text + + +def _assert_rejected_by_name(served: _Served) -> None: + assert served.status == 500, served.text + assert _REJECTION in served.text and "marker-" not in served.text, served.text + + +async def test_concurrent_burst_renders_each_template_and_rejects_only_the_bad_requests(gateway: Gateway) -> None: + calls: Final = tuple( + _Call(endpoint=call.endpoint, stream=call.stream, marker=call.marker, rejected=index % 5 == 4) + for index, call in enumerate(_calls(30, ("chat", "messages", "responses"), lambda index: index % 2 == 0)) + ) + accepted: Final = tuple(call for call in calls if not call.rejected) + with wire_server(_echo) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"ollama/{_BACKEND}", api_base=wire.url, api_key=_API_KEY, roles=_ROLES) + served: Final = await _burst(str(gateway.client.base_url), gateway.key, model, calls) + assert len(served) == 30 + for item in served: + if item.call.rejected: + _assert_rejected_by_name(item) + else: + _assert_answered_with_its_own_marker(item) + _assert_each_prompt_rendered_once(wire.drain(), frozenset(call.marker for call in accepted)) + answered: Final = tuple(item for item in served if not item.call.rejected) + known: Final = frozenset(identity for identity in map(_response_id, answered) if identity is not None) + assert len(known) == 16, known + statuses: Final = _spend_statuses(model, 30) + assert len(_successes(statuses)) == len(accepted), statuses + assert known <= _successes(statuses) + + +async def test_dropped_ollama_connections_fail_their_callers_and_the_rest_keep_their_template( + gateway: Gateway, +) -> None: + calls: Final = _calls(12, ("chat",), lambda index: index % 2 == 1) + dropped: Final = frozenset(call.marker for index, call in enumerate(calls) if index % 3 == 0) + + def respond(request: Request) -> Reply: + if not _is_generate(request): + return Reply(body=b"{}") + marker: Final = _marker_of(request) + stream: Final = _JSON_OBJECT.validate_json(request.body).get("stream") is True + return _generate_reply(marker, stream, drop_connection=marker in dropped) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"ollama/{_BACKEND}", api_base=wire.url, api_key=_API_KEY, roles=_ROLES) + served: Final = await _burst(str(gateway.client.base_url), gateway.key, model, calls) + assert len(served) == 12 + for item in served: + if item.call.marker in dropped: + assert item.status == 500, item.text + assert "APIConnectionError" in item.text and "marker-" not in item.text, item.text + else: + _assert_answered_with_its_own_marker(item) + recovery: Final = _Call(endpoint="chat", stream=False, marker=uuid.uuid4().hex) + (recovered,) = await _burst(str(gateway.client.base_url), gateway.key, model, (recovery,)) + _assert_answered_with_its_own_marker(recovered) + _assert_each_prompt_rendered_once(wire.drain(), frozenset(call.marker for call in (*calls, recovery))) + answered: Final = tuple(item for item in served if item.call.marker not in dropped) + survivors: Final = frozenset( + identity for identity in map(_response_id, (*answered, recovered)) if identity is not None + ) + assert len(survivors) == 9, survivors + statuses: Final = _spend_statuses(model, 13) + assert _successes(statuses) == survivors, statuses + assert sum(status == "failure" for status in statuses.values()) == len(dropped), statuses + + +def _owned_config( + wires: tuple[Wire, ...], + tmp_path: Path, + template: Mapping[str, JsonValue], + litellm_settings: Mapping[str, JsonValue], + router_settings: Mapping[str, JsonValue] = MappingProxyType({}), +) -> Path: + base: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + deployments: Final = [ + { + "model_name": _CONFIG_MODEL, + "litellm_params": { + "model": f"ollama/{_BACKEND}", + "api_base": wire.url, + "api_key": _API_KEY, + "order": order, + **template, + }, + } + for order, wire in enumerate(wires, start=1) + ] + config: Final = { + **base, + "model_list": deployments, + "litellm_settings": {**base.get("litellm_settings", {}), **litellm_settings}, + "router_settings": {**base.get("router_settings", {}), **router_settings}, + } + path: Final = tmp_path / "ollama-custom-template-chaos.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +def _open_upstream_connections(pid: int, upstream: str) -> int: + port: Final = urlsplit(upstream).port + return sum( + 1 + for connection in psutil.Process(pid).net_connections(kind="tcp") + if connection.status == psutil.CONN_ESTABLISHED and connection.raddr and connection.raddr.port == port + ) + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +async def test_worker_sigkill_mid_burst_leaves_the_sibling_rendering_templates( + gateway: Gateway, tmp_path: Path +) -> None: + calls: Final = _calls(20, ("chat",), lambda _: False) + release: Final = threading.Event() + held_markers: Final[SimpleQueue[str]] = SimpleQueue() + + def held(request: Request) -> Reply: + if not _is_generate(request): + return Reply(body=b"{}") + held_markers.put(_marker_of(request)) + assert release.wait(timeout=60), "The burst was never released" + return _echo(request) + + with wire_server(held) as wire: + path: Final = _owned_config((wire,), tmp_path, {"roles": _ROLES}, {}) + with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + candidate: Final = owned.gateway + workers: Final = eventually( + lambda: tuple(int(pid) for pid in _STARTED_WORKER.findall(owned.log.read_text())), + lambda pids: len(pids) == 2, + seconds=30, + ) + burst: Final = asyncio.create_task( + _burst( + str(candidate.client.base_url), candidate.key, _CONFIG_MODEL, calls, tolerate_transport_errors=True + ) + ) + await asyncio.to_thread(eventually, held_markers.qsize, lambda size: size == 20, 60) + held_by: Final = MappingProxyType({pid: _open_upstream_connections(pid, wire.url) for pid in workers}) + assert sum(held_by.values()) == 20, held_by + victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) + with owned.replacing_worker(victim_pid) as kill: + kill.freeze() + await asyncio.to_thread(kill) + release.set() + served: Final = await burst + assert held_by[survivor_pid] >= 10, held_by + assert len(served) == held_by[survivor_pid], (held_by, len(served)) + for item in served: + _assert_answered_with_its_own_marker(item) + follow_up: Final = _Call(endpoint="chat", stream=False, marker=uuid.uuid4().hex) + (answered,) = await _burst(str(candidate.client.base_url), candidate.key, _CONFIG_MODEL, (follow_up,)) + _assert_answered_with_its_own_marker(answered) + _assert_each_prompt_rendered_once(wire.drain(), frozenset(call.marker for call in (*calls, follow_up))) + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +async def test_rejected_templates_never_cool_down_the_preferred_deployment(gateway: Gateway, tmp_path: Path) -> None: + rejected: Final = tuple( + _Call(endpoint="chat", stream=False, marker=uuid.uuid4().hex, rejected=True) for _ in range(4) + ) + valid: Final = _calls(8, ("chat",), lambda _: False) + with wire_server(_echo) as preferred, wire_server(_echo) as standby: + path: Final = _owned_config( + (preferred, standby), tmp_path, {"roles": _ROLES}, {}, {"disable_cooldowns": False, "allowed_fails": 0} + ) + with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + base_url: Final = str(owned.gateway.client.base_url) + for served in await _burst(base_url, owned.gateway.key, _CONFIG_MODEL, rejected): + _assert_rejected_by_name(served) + for served in await _burst(base_url, owned.gateway.key, _CONFIG_MODEL, valid): + _assert_answered_with_its_own_marker(served) + _assert_each_prompt_rendered_once(preferred.drain(), frozenset(call.marker for call in valid)) + assert not any(_is_generate(request) for request in standby.drain()) + + +def _cached_chat(candidate: Gateway, marker: str, field: str, value: str) -> httpx.Response: + return candidate.request( + "POST", + "/v1/chat/completions", + { + "model": _CONFIG_MODEL, + "messages": [{"role": "system", "content": _SYSTEM}, {"role": "user", "content": _question(marker)}], + field: value, + }, + ) + + +def _served_from_cache(candidate: Gateway, field: str) -> bool: + marker: Final = uuid.uuid4().hex + first: Final = _cached_chat(candidate, marker, field, "") + assert first.status_code == 200, first.text + assert "x-litellm-cache-key" not in first.headers, sorted(first.headers) + eventually( + lambda: _cached_chat(candidate, marker, field, ""), + lambda found: "x-litellm-cache-key" in found.headers, + seconds=20, + ) + changed: Final = _cached_chat(candidate, marker, field, "") + assert changed.status_code == 200, changed.text + return "x-litellm-cache-key" in changed.headers + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +def test_initial_and_final_prompt_values_key_the_provider_specific_cache_alike( + gateway: Gateway, tmp_path: Path +) -> None: + template: Final[dict[str, JsonValue]] = { + "roles": _ROLES, + "initial_prompt_value": "", + "final_prompt_value": "", + } + with wire_server(_echo) as wire: + path: Final = _owned_config( + (wire,), tmp_path, template, {"enable_caching_on_provider_specific_optional_params": True} + ) + with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + initial: Final = _served_from_cache(owned.gateway, "initial_prompt_value") + final: Final = _served_from_cache(owned.gateway, "final_prompt_value") + assert initial == final, {"initial_prompt_value": initial, "final_prompt_value": final} + + +def _gemini_reply(_: Request) -> Reply: + candidate: Final[dict[str, JsonValue]] = { + "content": {"role": "model", "parts": [{"text": '{"colour": "blue"}'}]}, + "finishReason": "STOP", + } + usage: Final[dict[str, JsonValue]] = {"promptTokenCount": 30, "candidatesTokenCount": 12, "totalTokenCount": 42} + return Reply(body=json.dumps({"candidates": [candidate], "usageMetadata": usage}).encode()) + + +def _texts(contents: JsonValue) -> tuple[str, ...]: + assert isinstance(contents, list), contents + parts: Final = tuple(part for content in contents if isinstance(content, dict) for part in _parts(content)) + return tuple(text for part in parts if isinstance(text := part.get("text"), str)) + + +def _parts(content: dict[str, JsonValue]) -> tuple[dict[str, JsonValue], ...]: + parts: Final = content.get("parts") + assert isinstance(parts, list), content + return tuple(part for part in parts if isinstance(part, dict)) + + +def _anthropic_text_reply(_: Request) -> Reply: + completion: Final[dict[str, JsonValue]] = { + "type": "completion", + "id": "compl_scripted", + "completion": "blue", + "stop_reason": "stop_sequence", + } + return Reply(body=json.dumps(completion).encode()) + + +def _boot_deployment(name: str, model: str, wire: Wire) -> dict[str, JsonValue]: + return {"model_name": name, "litellm_params": {"model": model, "api_base": wire.url, "api_key": _API_KEY}} + + +def _boot_chat(candidate: Gateway, model: str, extra: Mapping[str, JsonValue]) -> None: + response: Final = candidate.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "system", "content": _SYSTEM}, {"role": "user", "content": _question(model)}], + "cache": {"no-cache": True}, + **extra, + }, + ) + assert response.status_code == 200, response.text + + +def _anthropic_text_prompts(wire: Wire) -> MappingProxyType[str, JsonValue]: + bodies: Final = tuple(_JSON_OBJECT.validate_json(request.body) for request in wire.drain()) + return MappingProxyType({str(body["model"]): body["prompt"] for body in bodies}) + + +@pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS) +def test_boot_time_custom_prompt_dict_renders_partial_templates(gateway: Gateway, tmp_path: Path) -> None: + templates: Final[dict[str, JsonValue]] = { + "response_schema_prompt": {"roles": {"user": {"pre_message": "", "post_message": ""}}}, + "claude-2.1": {"roles": _ROLES}, + "claude-instant-1.2": {"initial_prompt_value": "", "final_prompt_value": ""}, + } + schema: Final[dict[str, JsonValue]] = { + "type": "object", + "properties": {"colour": {"type": "string"}}, + "required": ["colour"], + } + with wire_server(_gemini_reply) as gemini, wire_server(_anthropic_text_reply) as anthropic_text: + base: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + deployments: Final = [ + _boot_deployment("gemini-schema", "gemini/scripted-template-model", gemini), + _boot_deployment("anthropic-text-roles", "anthropic_text/claude-2.1", anthropic_text), + _boot_deployment("anthropic-text-values", "anthropic_text/claude-instant-1.2", anthropic_text), + ] + settings: Final = {**base.get("litellm_settings", {}), "custom_prompt_dict": templates} + path: Final = tmp_path / "boot-time-custom-prompt-dict.yaml" + path.write_text(yaml.safe_dump({**base, "model_list": deployments, "litellm_settings": settings})) + with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + _boot_chat( + owned.gateway, + "gemini-schema", + {"response_format": {"type": "json_schema", "json_schema": {"name": "sky", "schema": schema}}}, + ) + posts: Final = tuple(request for request in gemini.drain() if request.method == "POST") + assert len(posts) == 1, [request.target for request in posts] + texts: Final = _texts(_JSON_OBJECT.validate_json(posts[0].body)["contents"]) + wrapped: Final = tuple(text for text in texts if text.startswith("")) + assert len(wrapped) == 1, texts + assert wrapped[0].endswith("") and "colour" in wrapped[0], wrapped[0] + _boot_chat(owned.gateway, "anthropic-text-roles", {}) + _boot_chat(owned.gateway, "anthropic-text-values", {}) + assert _anthropic_text_prompts(anthropic_text) == { + "claude-2.1": f"{_SYSTEM}{_question('anthropic-text-roles')}", + "claude-instant-1.2": f"{_SYSTEM}{_question('anthropic-text-values')}", + } diff --git a/tests/integration/providers/test_custom_prompt_template_wire.py b/tests/integration/providers/test_custom_prompt_template_wire.py new file mode 100644 index 00000000000..49a47a8c5e5 --- /dev/null +++ b/tests/integration/providers/test_custom_prompt_template_wire.py @@ -0,0 +1,701 @@ +import json +import uuid +from collections.abc import Callable, Iterator, Mapping +from contextlib import contextmanager +from dataclasses import dataclass +from typing import Final + +import anthropic +import openai +import pytest +from integration._support.client import Gateway, Scenario, eventually, object_value, string_value +from integration._support.database import read_rows +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +_BACKEND: Final = "llama3-custom-template" +_API_KEY: Final = "synthetic-ollama-key" +_SYSTEM: Final = "Answer in one word." +_QUESTION: Final = "What colour is the sky?" +_ANSWER: Final = "blue" +_INITIAL: Final = "" +_FINAL: Final = "" +_ROLES: Final[dict[str, JsonValue]] = { + "system": {"pre_message": "", "post_message": ""}, + "user": {"pre_message": "", "post_message": ""}, + "assistant": {"pre_message": "", "post_message": ""}, +} +_ROLES_ONLY: Final[dict[str, JsonValue]] = {"roles": _ROLES} +_VALUES_ONLY: Final[dict[str, JsonValue]] = {"initial_prompt_value": _INITIAL, "final_prompt_value": _FINAL} +_FULL: Final[dict[str, JsonValue]] = {**_ROLES_ONLY, **_VALUES_ONLY} +_ROLES_PROMPT: Final = f"{_SYSTEM}{_QUESTION}" +_VALUES_PROMPT: Final = f"{_INITIAL}{_SYSTEM}{_QUESTION}{_FINAL}" +_FULL_PROMPT: Final = f"{_INITIAL}{_ROLES_PROMPT}{_FINAL}" +_TEMPLATES: Final[Mapping[str, tuple[dict[str, JsonValue], str]]] = { + "roles_only": (_ROLES_ONLY, _ROLES_PROMPT), + "values_only": (_VALUES_ONLY, _VALUES_PROMPT), +} +_TEMPLATE_FIELDS: Final = frozenset({"roles", "initial_prompt_value", "final_prompt_value", "bos_token", "eos_token"}) +_AWS: Final[dict[str, JsonValue]] = { + "aws_access_key_id": "AKIASYNTHETICKEY0001", + "aws_secret_access_key": "synthetic-secret-key-for-testing", + "aws_region_name": "us-east-1", +} +_NO_CACHE: Final[dict[str, JsonValue]] = {"no-cache": True} +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_JSON_LIST: Final = TypeAdapter(list[JsonValue]) + + +def _generate_reply(text: str) -> Reply: + return Reply( + body=json.dumps( + { + "model": _BACKEND, + "created_at": "2026-10-15T00:00:00Z", + "response": text, + "done": True, + "done_reason": "stop", + "prompt_eval_count": 30, + "eval_count": 12, + } + ).encode() + ) + + +def _streamed_reply(text: str) -> Reply: + frames: Final = ( + {"model": _BACKEND, "created_at": "2026-10-15T00:00:00Z", "response": text, "done": False}, + { + "model": _BACKEND, + "created_at": "2026-10-15T00:00:00Z", + "response": "", + "done": True, + "done_reason": "stop", + "prompt_eval_count": 30, + "eval_count": 12, + }, + ) + return Reply( + content_type="application/x-ndjson", chunks=tuple(json.dumps(frame).encode() + b"\n" for frame in frames) + ) + + +def _is_generate(request: Request) -> bool: + return (request.method, request.target) == ("POST", "/api/generate") + + +def _echo_answer(request: Request) -> Reply: + stream: Final = _JSON_OBJECT.validate_json(request.body).get("stream") is True + return _streamed_reply(_ANSWER) if stream else _generate_reply(_ANSWER) + + +@contextmanager +def _ollama_server(respond: Callable[[Request], Reply] = _echo_answer) -> Iterator[Wire]: + with wire_server(lambda request: respond(request) if _is_generate(request) else Reply(body=b"{}")) as wire: + yield wire + + +def _generate_bodies(wire: Wire) -> tuple[dict[str, JsonValue], ...]: + return tuple(_JSON_OBJECT.validate_json(request.body) for request in wire.drain() if _is_generate(request)) + + +def _assert_rendered(body: dict[str, JsonValue], prompt: str, *, stream: bool) -> None: + assert body["model"] == _BACKEND + assert body["prompt"] == prompt, body["prompt"] + assert body["stream"] is stream + options: Final = body.get("options", {}) + assert isinstance(options, dict) + assert not _TEMPLATE_FIELDS & (set(body) | set(options)), sorted(body) + + +def _assert_only_rendered(wire: Wire, prompt: str, *, stream: bool = False) -> None: + bodies: Final = _generate_bodies(wire) + assert len(bodies) == 1, bodies + _assert_rendered(bodies[0], prompt, stream=stream) + + +def _ollama_model(scenario: Scenario, wire: Wire, template: Mapping[str, JsonValue]) -> str: + return scenario.model(model=f"ollama/{_BACKEND}", api_base=wire.url, api_key=_API_KEY, **template) + + +def _chat_messages() -> list[dict[str, JsonValue]]: + return [{"role": "system", "content": _SYSTEM}, {"role": "user", "content": _QUESTION}] + + +def _spend_row(identity: str) -> dict[str, JsonValue]: + rows: Final = eventually( + lambda: read_rows( + 'SELECT model_group, status, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" WHERE request_id=%s', + (identity,), + ), + lambda found: len(found) == 1, + seconds=70, + ) + return rows[0] + + +def _model_spend_rows(model: str, expected: int) -> list[dict[str, JsonValue]]: + rows: Final = eventually( + lambda: read_rows( + 'SELECT request_id, status, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" WHERE model_group=%s', + (model,), + ), + lambda found: len(found) >= expected, + seconds=70, + ) + assert len({row["request_id"] for row in rows}) == len(rows) == expected, rows + return rows + + +def _billed(model: str) -> dict[str, JsonValue]: + return {"model_group": model, "status": "success", "prompt_tokens": 30, "completion_tokens": 12} + + +def _openai_client(gateway: Gateway) -> openai.OpenAI: + return openai.OpenAI(base_url=str(gateway.client.base_url) + "/v1", api_key=gateway.key, max_retries=0) + + +def _async_openai_client(gateway: Gateway) -> openai.AsyncOpenAI: + return openai.AsyncOpenAI(base_url=str(gateway.client.base_url) + "/v1", api_key=gateway.key, max_retries=0) + + +def _anthropic_client(gateway: Gateway) -> anthropic.Anthropic: + return anthropic.Anthropic(base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0) + + +def _async_anthropic_client(gateway: Gateway) -> anthropic.AsyncAnthropic: + return anthropic.AsyncAnthropic(base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0) + + +def _post(gateway: Gateway, path: str, body: Mapping[str, JsonValue]) -> tuple[int, str]: + response: Final = gateway.request("POST", path, {**body, "cache": _NO_CACHE}) + return response.status_code, response.text + + +def _chat(gateway: Gateway, model: str, **extra: JsonValue) -> tuple[int, str]: + return _post(gateway, "/v1/chat/completions", {"model": model, "messages": _chat_messages(), **extra}) + + +def _chat_content(text: str) -> str: + choices: Final = _JSON_LIST.validate_python(_JSON_OBJECT.validate_json(text)["choices"]) + message: Final = object_value(object_value(choices[0])["message"]) + return string_value(message["content"]) + + +def test_openai_sdk_chat_renders_a_roles_only_template(gateway: Gateway) -> None: + with _ollama_server() as wire, gateway.scenario() as scenario: + model: Final = _ollama_model(scenario, wire, _ROLES_ONLY) + completion: Final = _openai_client(gateway).chat.completions.create( + model=model, + messages=_chat_messages(), # pyright: ignore[reportArgumentType] # plain JSON messages + extra_body={"cache": _NO_CACHE}, + ) + assert completion.choices[0].message.content == _ANSWER + assert completion.choices[0].finish_reason == "stop" + assert completion.usage is not None + assert (completion.usage.prompt_tokens, completion.usage.completion_tokens) == (30, 12) + _assert_only_rendered(wire, _ROLES_PROMPT) + assert _spend_row(completion.id) == _billed(model) + + +async def test_async_openai_sdk_chat_stream_renders_a_roles_only_template(gateway: Gateway) -> None: + with _ollama_server() as wire, gateway.scenario() as scenario: + model: Final = _ollama_model(scenario, wire, _ROLES_ONLY) + stream: Final = await _async_openai_client(gateway).chat.completions.create( + model=model, + messages=_chat_messages(), # pyright: ignore[reportArgumentType] # plain JSON messages + stream=True, + stream_options={"include_usage": True}, + extra_body={"cache": _NO_CACHE}, + ) + chunks: Final = [chunk async for chunk in stream] + assert {chunk.id for chunk in chunks} == {chunks[0].id} + assert "".join(choice.delta.content or "" for chunk in chunks for choice in chunk.choices) == _ANSWER + usages: Final = [chunk.usage for chunk in chunks if chunk.usage is not None] + assert [(usage.prompt_tokens, usage.completion_tokens) for usage in usages] == [(30, 12)] + _assert_only_rendered(wire, _ROLES_PROMPT, stream=True) + assert _spend_row(chunks[0].id) == _billed(model) + + +def test_openai_sdk_completions_render_a_roles_only_template(gateway: Gateway) -> None: + with _ollama_server() as wire, gateway.scenario() as scenario: + model: Final = _ollama_model(scenario, wire, _ROLES_ONLY) + completion: Final = _openai_client(gateway).completions.create( + model=model, prompt=_QUESTION, extra_body={"cache": _NO_CACHE} + ) + assert completion.choices[0].text == _ANSWER + _assert_only_rendered(wire, f"{_QUESTION}") + assert _spend_row(completion.id) == _billed(model) + + +def test_anthropic_sdk_messages_render_a_roles_only_template(gateway: Gateway) -> None: + with _ollama_server() as wire, gateway.scenario() as scenario: + model: Final = _ollama_model(scenario, wire, _ROLES_ONLY) + message: Final = _anthropic_client(gateway).messages.create( + model=model, + max_tokens=64, + system=_SYSTEM, + messages=[{"role": "user", "content": _QUESTION}], + extra_body={"cache": _NO_CACHE}, + ) + assert message.stop_reason == "end_turn" + assert [(block.type, getattr(block, "text", None)) for block in message.content] == [("text", _ANSWER)] + assert (message.usage.input_tokens, message.usage.output_tokens) == (30, 12) + _assert_only_rendered(wire, _ROLES_PROMPT) + assert _spend_row(message.id) == _billed(model) + + +async def test_async_anthropic_sdk_messages_stream_render_a_roles_only_template(gateway: Gateway) -> None: + with _ollama_server() as wire, gateway.scenario() as scenario: + model: Final = _ollama_model(scenario, wire, _ROLES_ONLY) + async with _async_anthropic_client(gateway).messages.stream( + model=model, + max_tokens=64, + system=_SYSTEM, + messages=[{"role": "user", "content": _QUESTION}], + extra_body={"cache": _NO_CACHE}, + ) as stream: + text: Final = "".join([piece async for piece in stream.text_stream]) + final: Final = await stream.get_final_message() + assert text == _ANSWER + assert final.stop_reason == "end_turn" + _assert_only_rendered(wire, _ROLES_PROMPT, stream=True) + assert _spend_row(final.id) == _billed(model) + + +def test_openai_sdk_responses_render_a_roles_only_template(gateway: Gateway) -> None: + with _ollama_server() as wire, gateway.scenario() as scenario: + model: Final = _ollama_model(scenario, wire, _ROLES_ONLY) + response: Final = _openai_client(gateway).responses.create( + model=model, instructions=_SYSTEM, input=_QUESTION, store=False, extra_body={"cache": _NO_CACHE} + ) + assert response.output_text == _ANSWER + assert response.usage is not None and (response.usage.input_tokens, response.usage.output_tokens) == (30, 12) + _assert_only_rendered(wire, _ROLES_PROMPT) + rows: Final = _model_spend_rows(model, 1) + assert (rows[0]["status"], rows[0]["prompt_tokens"], rows[0]["completion_tokens"]) == ("success", 30, 12) + + +async def test_async_openai_sdk_responses_stream_render_a_roles_only_template(gateway: Gateway) -> None: + with _ollama_server() as wire, gateway.scenario() as scenario: + model: Final = _ollama_model(scenario, wire, _ROLES_ONLY) + stream: Final = await _async_openai_client(gateway).responses.create( + model=model, + instructions=_SYSTEM, + input=_QUESTION, + store=False, + stream=True, + extra_body={"cache": _NO_CACHE}, + ) + events: Final = [event async for event in stream] + assert "".join(event.delta for event in events if event.type == "response.output_text.delta") == _ANSWER + completed: Final = [event for event in events if event.type == "response.completed"] + assert len(completed) == 1 and completed[0].response.output_text == _ANSWER + _assert_only_rendered(wire, _ROLES_PROMPT, stream=True) + assert _model_spend_rows(model, 1)[0]["status"] == "success" + + +def test_prompt_values_only_template_wraps_the_bare_messages(gateway: Gateway) -> None: + with _ollama_server() as wire, gateway.scenario() as scenario: + model: Final = _ollama_model(scenario, wire, _VALUES_ONLY) + code, text = _chat(gateway, model) + assert code == 200, text + assert _chat_content(text) == _ANSWER + _assert_only_rendered(wire, _VALUES_PROMPT) + + +def test_full_template_renders_every_part(gateway: Gateway) -> None: + with _ollama_server() as wire, gateway.scenario() as scenario: + model: Final = _ollama_model(scenario, wire, _FULL) + code, text = _chat(gateway, model) + assert code == 200, text + assert _chat_content(text) == _ANSWER + _assert_only_rendered(wire, _FULL_PROMPT) + + +@dataclass(frozen=True, slots=True) +class _Member: + deployment: Callable[[str], dict[str, JsonValue]] + reply: Callable[[str], Reply] + prompt_field: str + target: Callable[[str], bool] + + +def _codestral_reply(text: str) -> Reply: + return Reply( + body=json.dumps( + { + "id": "cmpl-scripted", + "object": "chat.completion", + "model": "codestral-latest", + "created": 1, + "choices": [{"index": 0, "message": {"role": "assistant", "content": text}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 30, "completion_tokens": 12, "total_tokens": 42}, + } + ).encode() + ) + + +def _predibase_reply(text: str) -> Reply: + return Reply(body=json.dumps({"generated_text": text, "details": {"finish_reason": "eos_token"}}).encode()) + + +def _sagemaker_reply(text: str) -> Reply: + return Reply(body=json.dumps([{"generated_text": text}]).encode()) + + +def _bedrock_meta_reply(text: str) -> Reply: + return Reply( + body=json.dumps( + {"generation": text, "prompt_token_count": 30, "generation_token_count": 12, "stop_reason": "stop"} + ).encode() + ) + + +_FAMILY: Final[Mapping[str, _Member]] = { + "codestral": _Member( + deployment=lambda url: { + "model": "text-completion-codestral/codestral-latest", + "api_base": f"{url}/v1/fim/completions", + "api_key": "scripted-key", + }, + reply=_codestral_reply, + prompt_field="prompt", + target=lambda target: target == "/v1/fim/completions", + ), + "predibase": _Member( + deployment=lambda url: { + "model": "predibase/llama-3-8b-instruct", + "api_base": url, + "api_key": "scripted-key", + "predibase_tenant_id": "scripted-tenant", + }, + reply=_predibase_reply, + prompt_field="inputs", + target=lambda target: target.startswith("/scripted-tenant/deployments/v2/llms/llama-3-8b-instruct/generate"), + ), + "sagemaker": _Member( + deployment=lambda url: { + "model": "sagemaker/scripted-endpoint", + "sagemaker_base_url": f"{url}/endpoints/scripted-endpoint/invocations", + **_AWS, + }, + reply=_sagemaker_reply, + prompt_field="inputs", + target=lambda target: target == "/endpoints/scripted-endpoint/invocations", + ), + "bedrock_invoke": _Member( + deployment=lambda url: {"model": "bedrock/invoke/meta.llama3-8b-instruct-v1:0", "api_base": url, **_AWS}, + reply=_bedrock_meta_reply, + prompt_field="prompt", + target=lambda target: target.endswith("/invoke"), + ), +} + + +@pytest.mark.parametrize("template", sorted(_TEMPLATES)) +@pytest.mark.parametrize("member", sorted(_FAMILY)) +def test_family_member_renders_a_partial_template(gateway: Gateway, member: str, template: str) -> None: + spec: Final = _FAMILY[member] + fields, prompt = _TEMPLATES[template] + with wire_server(lambda _: spec.reply(_ANSWER)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(**spec.deployment(wire.url), **fields) + code, text = _chat(gateway, model) + received: Final = tuple(request for request in wire.drain() if request.method == "POST") + assert code == 200, text + assert _chat_content(text) == _ANSWER + assert len(received) == 1, [request.target for request in received] + assert spec.target(received[0].target), received[0].target + body: Final = _JSON_OBJECT.validate_json(received[0].body) + assert body[spec.prompt_field] == prompt, body + assert not _TEMPLATE_FIELDS & set(body), sorted(body) + + +@pytest.mark.parametrize("provider", ("vllm", "petals")) +def test_in_process_family_member_fails_alone_and_the_rig_keeps_serving(gateway: Gateway, provider: str) -> None: + with _ollama_server() as wire, gateway.scenario() as scenario: + broken: Final = scenario.model(model=f"{provider}/scripted-model", **_ROLES_ONLY) + healthy: Final = _ollama_model(scenario, wire, _ROLES_ONLY) + failed, failure = _chat(gateway, broken) + assert failed >= 400, failure + code, text = _chat(gateway, healthy) + assert code == 200, text + _assert_only_rendered(wire, _ROLES_PROMPT) + + +def _watsonx_reply(text: str) -> Reply: + result: Final[dict[str, JsonValue]] = { + "generated_text": text, + "generated_token_count": 12, + "input_token_count": 30, + "stop_reason": "eos_token", + } + return Reply(body=json.dumps({"model_id": "ibm/granite", "results": [result]}).encode()) + + +def test_watsonx_text_deployment_with_a_template_still_answers(gateway: Gateway) -> None: + with wire_server(lambda _: _watsonx_reply(_ANSWER)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="watsonx_text/ibm/granite-chat-scripted", + api_base=wire.url, + token="scripted-token", + project_id="p", + **_ROLES_ONLY, + ) + code, text = _chat(gateway, model) + assert code == 200, text + assert _chat_content(text) == _ANSWER + posts: Final = tuple(request for request in wire.drain() if request.method == "POST") + assert len(posts) == 1, [request.target for request in posts] + body: Final = _JSON_OBJECT.validate_json(posts[0].body) + assert _QUESTION in string_value(body["input"]), body + assert not _TEMPLATE_FIELDS & set(body), sorted(body) + + +def _openai_chat_reply(text: str) -> Reply: + return Reply( + body=json.dumps( + { + "id": "chatcmpl-scripted", + "object": "chat.completion", + "created": 1, + "model": "scripted-compat", + "choices": [{"index": 0, "message": {"role": "assistant", "content": text}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 30, "completion_tokens": 12, "total_tokens": 42}, + } + ).encode() + ) + + +def _converse_reply(text: str) -> Reply: + return Reply( + body=json.dumps( + { + "output": {"message": {"role": "assistant", "content": [{"text": text}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 30, "outputTokens": 12, "totalTokens": 42}, + "metrics": {"latencyMs": 1}, + } + ).encode() + ) + + +_FORWARDING: Final[Mapping[str, tuple[Callable[[str], dict[str, JsonValue]], Callable[[str], Reply]]]] = { + "openai_compatible": ( + lambda url: {"model": "openai/scripted-compat", "api_base": f"{url}/v1", "api_key": "scripted-key"}, + _openai_chat_reply, + ), + "bedrock_converse": ( + lambda url: {"model": "bedrock/converse/meta.llama3-8b-instruct-v1:0", "api_base": url, **_AWS}, + _converse_reply, + ), +} + + +def _keys_anywhere(value: JsonValue) -> frozenset[str]: + if isinstance(value, dict): + nested: Final = frozenset(key for child in value.values() for key in _keys_anywhere(child)) + return frozenset(value) | nested + if isinstance(value, list): + return frozenset(key for child in value for key in _keys_anywhere(child)) + return frozenset() + + +@pytest.mark.parametrize("source", ("deployment", "request")) +@pytest.mark.parametrize("member", sorted(_FORWARDING)) +def test_template_fields_never_reach_a_provider_that_does_not_render_them( + gateway: Gateway, member: str, source: str +) -> None: + deployment, reply = _FORWARDING[member] + with wire_server(lambda _: reply(_ANSWER)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(**deployment(wire.url), **(_FULL if source == "deployment" else {})) + code, text = _chat(gateway, model, **(_FULL if source == "request" else {})) + assert code == 200, text + assert _chat_content(text) == _ANSWER + posts: Final = tuple(request for request in wire.drain() if request.method == "POST") + assert len(posts) == 1, [request.target for request in posts] + body: Final = _JSON_OBJECT.validate_json(posts[0].body) + assert not _TEMPLATE_FIELDS & _keys_anywhere(body), body + assert _INITIAL not in posts[0].body.decode() and _FINAL not in posts[0].body.decode(), body + + +@pytest.mark.parametrize("value", (5, [_INITIAL]), ids=("int", "list")) +@pytest.mark.parametrize("source", ("deployment", "request")) +def test_non_string_initial_prompt_value_is_rejected_by_name(gateway: Gateway, source: str, value: JsonValue) -> None: + with _ollama_server() as wire, gateway.scenario() as scenario: + template: Final[dict[str, JsonValue]] = {**_FULL, "initial_prompt_value": value} + broken: Final = _ollama_model(scenario, wire, template if source == "deployment" else _FULL) + healthy: Final = _ollama_model(scenario, wire, _FULL) + failed, failure = _chat(gateway, broken, **({"initial_prompt_value": value} if source == "request" else {})) + assert failed == 500, failure + assert "initial_prompt_value: Input should be a valid string" in failure, failure + assert _generate_bodies(wire) == () + code, text = _chat(gateway, healthy) + assert code == 200, text + _assert_only_rendered(wire, _FULL_PROMPT) + + +@pytest.mark.parametrize("value", ("", None), ids=("empty", "null")) +def test_blank_initial_prompt_value_in_the_request_drops_the_deployment_value( + gateway: Gateway, value: JsonValue +) -> None: + with _ollama_server() as wire, gateway.scenario() as scenario: + model: Final = _ollama_model(scenario, wire, _FULL) + code, text = _chat(gateway, model, initial_prompt_value=value) + assert code == 200, text + _assert_only_rendered(wire, f"{_ROLES_PROMPT}{_FINAL}") + + +def test_request_initial_prompt_value_overrides_the_deployment_value(gateway: Gateway) -> None: + with _ollama_server() as wire, gateway.scenario() as scenario: + model: Final = _ollama_model(scenario, wire, _FULL) + code, text = _chat(gateway, model, initial_prompt_value="") + assert code == 200, text + _assert_only_rendered(wire, f"{_ROLES_PROMPT}{_FINAL}") + + +def test_five_kilobyte_initial_prompt_value_leads_the_prompt(gateway: Gateway) -> None: + large: Final = "x" * 5120 + with _ollama_server() as wire, gateway.scenario() as scenario: + model: Final = _ollama_model(scenario, wire, _ROLES_ONLY) + code, text = _chat(gateway, model, initial_prompt_value=large) + assert code == 200, text + _assert_only_rendered(wire, f"{large}{_ROLES_PROMPT}") + + +def test_identical_requests_each_render_and_bill_once(gateway: Gateway) -> None: + with _ollama_server() as wire, gateway.scenario() as scenario: + model: Final = _ollama_model(scenario, wire, _ROLES_ONLY) + first, first_text = _chat(gateway, model) + second, second_text = _chat(gateway, model) + assert (first, second) == (200, 200), (first_text, second_text) + bodies: Final = _generate_bodies(wire) + assert len(bodies) == 2, bodies + for body in bodies: + _assert_rendered(body, _ROLES_PROMPT, stream=False) + identities: Final = {string_value(_JSON_OBJECT.validate_json(text)["id"]) for text in (first_text, second_text)} + assert len(identities) == 2 + for identity in identities: + assert _spend_row(identity) == _billed(model) + + +def test_template_with_only_a_bos_token_names_the_fields_it_lacks(gateway: Gateway) -> None: + with _ollama_server() as wire, gateway.scenario() as scenario: + model: Final = _ollama_model(scenario, wire, {"bos_token": ""}) + code, text = _chat(gateway, model) + assert code == 500, text + assert "it sets none of roles, initial_prompt_value, or final_prompt_value" in text, text + assert _generate_bodies(wire) == () + + +def test_roles_that_are_not_a_mapping_render_without_role_markers(gateway: Gateway) -> None: + with _ollama_server() as wire, gateway.scenario() as scenario: + model: Final = _ollama_model(scenario, wire, {**_VALUES_ONLY, "roles": json.dumps(_ROLES)}) + code, text = _chat(gateway, model) + assert code == 200, text + _assert_only_rendered(wire, _VALUES_PROMPT) + + +def test_provider_error_reaches_the_caller_after_one_attempt(gateway: Gateway) -> None: + missing: Final = Reply(status=404, body=json.dumps({"error": f"model '{_BACKEND}' not found"}).encode()) + with _ollama_server(lambda _: missing) as wire, gateway.scenario() as scenario: + model: Final = _ollama_model(scenario, wire, _ROLES_ONLY) + code, text = _chat(gateway, model) + assert code == 404, text + assert f"model '{_BACKEND}' not found" in text, text + bodies: Final = _generate_bodies(wire) + assert len(bodies) == 1, bodies + _assert_rendered(bodies[0], _ROLES_PROMPT, stream=False) + + +def _template_keys(fields: Mapping[str, JsonValue]) -> tuple[bool, bool]: + return "initial_prompt_value" in fields, "final_prompt_value" in fields + + +def test_pass_through_body_treats_initial_and_final_prompt_values_alike(gateway: Gateway) -> None: + with wire_server(lambda _: _generate_reply(_ANSWER)) as wire, gateway.scenario() as scenario: + path: Final = f"/integration-template-passthrough-{uuid.uuid4().hex}" + created: Final = gateway.post( + "/config/pass_through_endpoint", {"path": path, "target": f"{wire.url}/api/generate"} + ) + endpoint: Final = object_value(_JSON_LIST.validate_python(created["endpoints"])[0]) + scenario.cleanups.callback( + lambda: gateway.request( + "DELETE", "/config/pass_through_endpoint", params={"endpoint_id": str(endpoint["id"])} + ) + ) + response: Final = eventually( + lambda: gateway.request("POST", path, {"model": _BACKEND, "prompt": _QUESTION, "stream": False, **_FULL}), + lambda answered: answered.status_code != 404, + seconds=30, + ) + assert response.status_code == 200, response.text + posts: Final = tuple(request for request in wire.drain() if request.method == "POST") + forwarded: Final = _JSON_OBJECT.validate_json(posts[-1].body) + assert (forwarded["model"], forwarded["prompt"]) == (_BACKEND, _QUESTION), forwarded + initial, final = _template_keys(forwarded) + assert initial == final, sorted(forwarded) + + +def _a2a_reply(request: Request) -> Reply: + result: Final[dict[str, JsonValue]] = { + "kind": "message", + "role": "agent", + "messageId": f"{uuid.uuid4().hex}-out", + "parts": [{"kind": "text", "text": _ANSWER}], + } + outbound: Final = _JSON_OBJECT.validate_json(request.body) + return Reply(body=json.dumps({"jsonrpc": "2.0", "id": outbound["id"], "result": result}).encode()) + + +def test_a2a_params_treat_initial_and_final_prompt_values_alike(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + agent_name: Final = f"template-agent-{marker}" + + def peer(request: Request) -> Reply: + if request.method == "GET": + return Reply(body=json.dumps(_agent_card(wire.url, agent_name)).encode()) + return _a2a_reply(request) + + with wire_server(peer) as wire, gateway.scenario() as scenario: + created: Final = gateway.post( + "/v1/agents", {"agent_name": agent_name, "agent_card_params": _agent_card(wire.url, agent_name)} + ) + agent_id: Final = string_value(created["agent_id"]) + scenario.cleanups.callback(lambda: gateway.request("DELETE", f"/v1/agents/{agent_id}")) + message: Final[dict[str, JsonValue]] = { + "kind": "message", + "role": "user", + "messageId": f"{marker}-in", + "parts": [{"kind": "text", "text": _QUESTION}], + } + invoked: Final = gateway.request( + "POST", + f"/a2a/{agent_id}", + { + "jsonrpc": "2.0", + "id": marker, + "method": "message/send", + "params": {"message": message, "initial_prompt_value": _INITIAL, "final_prompt_value": _FINAL}, + }, + ) + assert invoked.status_code == 200, invoked.text + assert object_value(invoked.json()["result"])["parts"] == [{"kind": "text", "text": _ANSWER}], invoked.text + posts: Final = tuple(request for request in wire.drain() if request.method == "POST") + assert len(posts) == 1, [request.target for request in posts] + params: Final = object_value(_JSON_OBJECT.validate_json(posts[0].body)["params"]) + assert _template_keys(params) == (False, False), sorted(params) + + +def _agent_card(url: str, name: str) -> dict[str, JsonValue]: + return { + "protocolVersion": "0.3", + "name": name, + "description": "Scripted A2A peer", + "version": "1.0.0", + "url": url + "/", + "capabilities": {"streaming": True}, + "defaultInputModes": ["text"], + "defaultOutputModes": ["text"], + "skills": [], + } diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py index ad216548d08..c878f0e149d 100644 --- a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py +++ b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py @@ -3,7 +3,9 @@ import json import logging import os import re +import sys import time +import types from collections.abc import Callable, Mapping, Sequence from copy import deepcopy from typing import Final, List, TypeAlias @@ -69,6 +71,17 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( from litellm.llms.vertex_ai.gemini.transformation import gemini_convert_messages_with_history from litellm.types.llms.bedrock import BedrockConverseReasoningContentBlock from litellm.types.llms.openai import AllMessageValues +from litellm.litellm_core_utils.prompt_templates.factory import response_schema_prompt +from litellm.llms.anthropic.completion.transformation import AnthropicTextConfig +from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation import AmazonInvokeConfig +from litellm.llms.codestral.completion.handler import CodestralTextCompletion +from litellm.llms.ollama.completion.transformation import OllamaConfig +from litellm.llms.petals.completion import handler as petals_handler +from litellm.llms.predibase.chat.transformation import PredibaseConfig +from litellm.llms.sagemaker.completion.transformation import SagemakerConfig +from litellm.llms.vllm.completion import handler as vllm_handler +from litellm.llms.watsonx.common_utils import _convert_watsonx_messages_core +from litellm.types.utils import ModelResponse, TextCompletionResponse READ_FILE: Final[list[object]] = [{"name": "read_file", "description": "Read the contents of a file."}] @@ -7871,3 +7884,334 @@ def test_anthropic_messages_pt_merges_a_long_assistant_run_in_linear_time(): growth: Final = _merged_assistant_run_cost(4000) / _merged_assistant_run_cost(1000) assert growth < 8.0, f"cost grew {growth:.1f}x for 4x the assistant run" + + +TEMPLATE_MODEL = "partial-template-model" +TEMPLATE_ROLES = {"user": {"pre_message": "[INST] ", "post_message": " [/INST]"}} +TEMPLATE_MESSAGES = [{"role": "user", "content": "hi"}] +TEMPLATE_RENDERED_MESSAGES = "[INST] hi [/INST]" +TEMPLATE_BARE_MESSAGES = "hi" + +TEMPLATE_RESPONSE_SCHEMA = {"type": "object"} +TEMPLATE_RENDERED_SCHEMA = f"[INST] {TEMPLATE_RESPONSE_SCHEMA} [/INST]" +TEMPLATE_BARE_SCHEMA = str(TEMPLATE_RESPONSE_SCHEMA) + +TEMPLATE_INITIAL = "" +TEMPLATE_FINAL = "" + +PARTIAL_TEMPLATE = {"roles": TEMPLATE_ROLES} +FULL_TEMPLATE = {"roles": TEMPLATE_ROLES, "initial_prompt_value": TEMPLATE_INITIAL, "final_prompt_value": TEMPLATE_FINAL} +ROLELESS_TEMPLATE = {"initial_prompt_value": TEMPLATE_INITIAL, "final_prompt_value": TEMPLATE_FINAL} + + +class _PromptCaptured(Exception): + def __init__(self, prompt): + super().__init__(prompt) + self.prompt = prompt + + +class _CapturingLogging: + def pre_call(self, input, api_key, additional_args=None, **kwargs): + raise _PromptCaptured(input) + + def post_call(self, *args, **kwargs): + pass + + +def _captured_prompt(call): + with pytest.raises(_PromptCaptured) as excinfo: + call() + return excinfo.value.prompt + + +def _stub_vllm_import(monkeypatch): + class _SamplingParams: + def __init__(self, **kwargs): + self.kwargs = kwargs + + class _LLM: + def __init__(self, model): + self.model = model + + def generate(self, prompts, sampling_params): + raise _PromptCaptured(prompts) + + stub = types.ModuleType("vllm") + stub.LLM = _LLM + stub.SamplingParams = _SamplingParams + + monkeypatch.setitem(sys.modules, "vllm", stub) + monkeypatch.setattr(vllm_handler, "llm", None) + + +def _render_ollama(template, monkeypatch, messages=TEMPLATE_MESSAGES): + request = OllamaConfig().transform_request( + model=TEMPLATE_MODEL, + messages=messages, + optional_params={}, + litellm_params={"custom_prompt_dict": {TEMPLATE_MODEL: template}}, + headers={}, + ) + return request["prompt"] + + +def _render_predibase(template, monkeypatch): + request = PredibaseConfig().transform_request( + model=TEMPLATE_MODEL, + messages=TEMPLATE_MESSAGES, + optional_params={}, + litellm_params={"custom_prompt_dict": {TEMPLATE_MODEL: template}}, + headers={}, + ) + return request["inputs"] + + +def _render_anthropic_text(template, monkeypatch): + monkeypatch.setattr(litellm, "custom_prompt_dict", {TEMPLATE_MODEL: template}) + return AnthropicTextConfig()._get_anthropic_text_prompt_from_messages( + messages=TEMPLATE_MESSAGES, model=TEMPLATE_MODEL + ) + + +def _render_response_schema_prompt(template, monkeypatch): + monkeypatch.setattr(litellm, "custom_prompt_dict", {f"{TEMPLATE_MODEL}/response_schema_prompt": template}) + return response_schema_prompt( + model=TEMPLATE_MODEL, response_schema=TEMPLATE_RESPONSE_SCHEMA, llm_provider="vertex_ai" + ) + + +def _render_petals(template, monkeypatch): + monkeypatch.setattr(litellm, "custom_prompt_dict", {TEMPLATE_MODEL: template}) + return _captured_prompt( + lambda: petals_handler.completion( + model=TEMPLATE_MODEL, + messages=TEMPLATE_MESSAGES, + api_base="http://petals.invalid", + model_response=ModelResponse(), + print_verbose=lambda *args, **kwargs: None, + encoding=None, + logging_obj=_CapturingLogging(), + optional_params={}, + ) + ) + + +def _render_codestral(template, monkeypatch): + return _captured_prompt( + lambda: CodestralTextCompletion().completion( + model=TEMPLATE_MODEL, + messages=TEMPLATE_MESSAGES, + api_base="http://codestral.invalid", + custom_prompt_dict={TEMPLATE_MODEL: template}, + model_response=TextCompletionResponse(), + print_verbose=lambda *args, **kwargs: None, + encoding=None, + api_key="sk-not-used", + logging_obj=_CapturingLogging(), + optional_params={}, + timeout=1.0, + ) + ) + + +def _render_vllm_completion(template, monkeypatch): + _stub_vllm_import(monkeypatch) + return _captured_prompt( + lambda: vllm_handler.completion( + model=TEMPLATE_MODEL, + messages=TEMPLATE_MESSAGES, + model_response=ModelResponse(), + print_verbose=lambda *args, **kwargs: None, + encoding=None, + logging_obj=_CapturingLogging(), + optional_params={}, + custom_prompt_dict={TEMPLATE_MODEL: template}, + ) + ) + + +def _render_vllm_batch(template, monkeypatch): + _stub_vllm_import(monkeypatch) + prompts = _captured_prompt( + lambda: vllm_handler.batch_completions( + model=TEMPLATE_MODEL, + messages=[TEMPLATE_MESSAGES], + optional_params={}, + custom_prompt_dict={TEMPLATE_MODEL: template}, + ) + ) + assert len(prompts) == 1 + return prompts[0] + + +def _render_bedrock_invoke(template, monkeypatch): + prompt, _ = AmazonInvokeConfig().convert_messages_to_prompt( + model=TEMPLATE_MODEL, messages=TEMPLATE_MESSAGES, provider="meta", custom_prompt_dict={TEMPLATE_MODEL: template} + ) + return prompt + + +def _render_sagemaker(template, monkeypatch): + return SagemakerConfig()._transform_prompt( + model=TEMPLATE_MODEL, + messages=TEMPLATE_MESSAGES, + custom_prompt_dict={TEMPLATE_MODEL: template}, + hf_model_name=None, + ) + + +def _render_watsonx(template, monkeypatch): + return _convert_watsonx_messages_core( + model=TEMPLATE_MODEL, + messages=TEMPLATE_MESSAGES, + provider="ibm", + custom_prompt_dict={TEMPLATE_MODEL: template}, + apply_template_fn=None, + ) + + +TEMPLATE_CALL_SITES = [ + pytest.param(_render_ollama, TEMPLATE_RENDERED_MESSAGES, TEMPLATE_BARE_MESSAGES, id="ollama"), + pytest.param(_render_petals, TEMPLATE_RENDERED_MESSAGES, TEMPLATE_BARE_MESSAGES, id="petals"), + pytest.param(_render_vllm_completion, TEMPLATE_RENDERED_MESSAGES, TEMPLATE_BARE_MESSAGES, id="vllm-completion"), + pytest.param(_render_vllm_batch, TEMPLATE_RENDERED_MESSAGES, TEMPLATE_BARE_MESSAGES, id="vllm-batch-completions"), + pytest.param(_render_predibase, TEMPLATE_RENDERED_MESSAGES, TEMPLATE_BARE_MESSAGES, id="predibase"), + pytest.param(_render_codestral, TEMPLATE_RENDERED_MESSAGES, TEMPLATE_BARE_MESSAGES, id="codestral"), + pytest.param(_render_anthropic_text, TEMPLATE_RENDERED_MESSAGES, TEMPLATE_BARE_MESSAGES, id="anthropic-text"), + pytest.param( + _render_response_schema_prompt, TEMPLATE_RENDERED_SCHEMA, TEMPLATE_BARE_SCHEMA, id="response-schema-prompt" + ), + pytest.param(_render_bedrock_invoke, TEMPLATE_RENDERED_MESSAGES, TEMPLATE_BARE_MESSAGES, id="bedrock-invoke"), + pytest.param(_render_sagemaker, TEMPLATE_RENDERED_MESSAGES, TEMPLATE_BARE_MESSAGES, id="sagemaker"), + pytest.param(_render_watsonx, TEMPLATE_RENDERED_MESSAGES, TEMPLATE_BARE_MESSAGES, id="watsonx"), +] + + +@pytest.mark.parametrize(("render", "rendered_body", "bare_body"), TEMPLATE_CALL_SITES) +def test_template_without_the_optional_prompt_values_renders(render, rendered_body, bare_body, monkeypatch): + assert render(PARTIAL_TEMPLATE, monkeypatch) == rendered_body + + +@pytest.mark.parametrize(("render", "rendered_body", "bare_body"), TEMPLATE_CALL_SITES) +def test_template_with_the_optional_prompt_values_still_applies_them(render, rendered_body, bare_body, monkeypatch): + assert render(FULL_TEMPLATE, monkeypatch) == TEMPLATE_INITIAL + rendered_body + TEMPLATE_FINAL + + +@pytest.mark.parametrize(("render", "rendered_body", "bare_body"), TEMPLATE_CALL_SITES) +def test_template_without_roles_renders_messages_unwrapped(render, rendered_body, bare_body, monkeypatch): + assert render(ROLELESS_TEMPLATE, monkeypatch) == TEMPLATE_INITIAL + bare_body + TEMPLATE_FINAL + + +def test_template_wraps_only_the_roles_it_defines_across_a_conversation(monkeypatch): + messages = [ + {"role": "system", "content": "Be terse."}, + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello"}, + {"role": "user", "content": "sky colour?"}, + ] + + rendered = _render_ollama(FULL_TEMPLATE, monkeypatch, messages=messages) + + assert rendered == "Be terse.[INST] hi [/INST]hello[INST] sky colour? [/INST]" + + +@pytest.mark.parametrize( + "template", + [ + pytest.param({"roles": {**TEMPLATE_ROLES, "tool": None}}, id="null-role-entry"), + pytest.param({"roles": {"user": {**TEMPLATE_ROLES["user"], "priority": 1}}}, id="extra-role-key"), + pytest.param({"roles": {**TEMPLATE_ROLES, "tool": {"pre_message": 5}}}, id="unused-role-wrong-type"), + pytest.param({"roles": {**TEMPLATE_ROLES, 1: {"pre_message": "x"}}}, id="unused-non-string-role"), + pytest.param({"roles": TEMPLATE_ROLES, "initial_prompt_value": None, "bos_token": 7}, id="null-and-unread-keys"), + ], +) +def test_template_shapes_main_rendered_keep_rendering(template, monkeypatch): + assert _render_ollama(template, monkeypatch) == TEMPLATE_RENDERED_MESSAGES + + +@pytest.mark.parametrize( + "template", + [ + pytest.param({"roles": {"user": "[INST]"}}, id="string-role-entry"), + pytest.param({"roles": {"user": ["[INST]"]}}, id="list-role-entry"), + pytest.param({"roles": ["user"]}, id="list-roles"), + pytest.param({"roles": '{"user": {"pre_message": "[INST]"}}'}, id="json-string-roles"), + ], +) +def test_roles_that_are_not_a_mapping_render_unwrapped_with_a_warning(template, monkeypatch, caplog): + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + rendered = _render_ollama({**template, "final_prompt_value": TEMPLATE_FINAL}, monkeypatch) + + assert rendered == TEMPLATE_BARE_MESSAGES + TEMPLATE_FINAL + assert "not a mapping" in caplog.text + + +def test_null_role_entry_renders_that_role_unwrapped_without_a_warning(monkeypatch, caplog): + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + rendered = _render_ollama({"roles": {"user": None}}, monkeypatch) + + assert rendered == TEMPLATE_BARE_MESSAGES + assert "not a mapping" not in caplog.text + + +@pytest.mark.parametrize( + ("template", "field"), + [ + pytest.param({"initial_prompt_value": 5}, "initial_prompt_value", id="initial-prompt-value"), + pytest.param({"final_prompt_value": ["x"]}, "final_prompt_value", id="final-prompt-value"), + pytest.param({"roles": {"user": {"pre_message": 5}}}, "roles.user.pre_message", id="pre-message"), + ], +) +def test_wrong_type_template_field_fails_as_a_server_error_naming_it(template, field): + with pytest.raises(litellm.APIConnectionError) as excinfo: + litellm.completion( + model=f"ollama/{TEMPLATE_MODEL}", + messages=TEMPLATE_MESSAGES, + api_base="http://ollama.invalid", + **template, + ) + + assert excinfo.value.status_code == 500 + assert f"{field}: Input should be a valid string" in excinfo.value.message + + +@pytest.mark.parametrize( + "template", + [ + pytest.param({"bos_token": ""}, id="bos-token-only"), + pytest.param({"eos_token": ""}, id="eos-token-only"), + pytest.param({"roles": '{"user": {"pre_message": "[INST]"}}'}, id="json-string-roles-only"), + ], +) +def test_template_that_sets_nothing_the_prompt_uses_fails_as_a_server_error(template): + with pytest.raises(litellm.APIConnectionError) as excinfo: + litellm.completion( + model=f"ollama/{TEMPLATE_MODEL}", + messages=TEMPLATE_MESSAGES, + api_base="http://ollama.invalid", + **template, + ) + + assert excinfo.value.status_code == 500 + assert "is invalid: it sets none of roles, initial_prompt_value, or final_prompt_value" in excinfo.value.message + + +def test_wrong_type_template_field_in_requests_leaves_the_deployments_in_rotation(): + router = litellm.Router( + model_list=[ + { + "model_name": "template-group", + "litellm_params": {"model": f"ollama/{TEMPLATE_MODEL}", "api_base": "http://ollama.invalid"}, + } + for _ in range(2) + ], + num_retries=0, + allowed_fails=0, + ) + + for _ in range(2): + with pytest.raises(litellm.APIConnectionError): + router.completion(model="template-group", messages=TEMPLATE_MESSAGES, initial_prompt_value=5) + response = router.completion(model="template-group", messages=TEMPLATE_MESSAGES, mock_response="Blue") + + assert response.choices[0].message.content == "Blue" diff --git a/tests/unit/llms/test_predibase_transformation.py b/tests/unit/llms/test_predibase_transformation.py index d6425488be3..20d547c3b63 100644 --- a/tests/unit/llms/test_predibase_transformation.py +++ b/tests/unit/llms/test_predibase_transformation.py @@ -39,14 +39,9 @@ def test_predibase_transform_request_non_stream(): assert "stream" not in request_data["parameters"] -def test_predibase_transform_request_custom_prompt(monkeypatch): +def test_predibase_transform_request_custom_prompt(): config = PredibaseConfig() - monkeypatch.setattr( - "litellm.llms.predibase.chat.transformation.custom_prompt", - lambda **kwargs: "custom-prompt", - ) - request_data = config.transform_request( model="predibase-model", messages=[{"role": "user", "content": "hello"}], @@ -54,16 +49,16 @@ def test_predibase_transform_request_custom_prompt(monkeypatch): litellm_params={ "custom_prompt_dict": { "predibase-model": { - "roles": {}, - "initial_prompt_value": "", - "final_prompt_value": "", + "roles": {"user": {"pre_message": "", "post_message": ""}}, + "initial_prompt_value": "", + "final_prompt_value": "", } } }, headers={}, ) - assert request_data["inputs"] == "custom-prompt" + assert request_data["inputs"] == "hello" def test_predibase_get_complete_url_stream_and_non_stream(): diff --git a/tests/unit/types/test_litellm_params.py b/tests/unit/types/test_litellm_params.py index d27dc733c78..122d811e491 100644 --- a/tests/unit/types/test_litellm_params.py +++ b/tests/unit/types/test_litellm_params.py @@ -199,6 +199,7 @@ OPTION_NAMES: Final = ( "prompt_label", "litellm_system_prompt", "custom_prompt_dict", + "initial_prompt_value", "roles", "final_prompt_value", "bos_token", From 2dc6aafa6baf8dabdac58f5ba70c384b8dcb3e09 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 16 Oct 2026 10:54:42 +0000 Subject: [PATCH 15/23] fix(proxy)!: count prompt-only token_counter requests with the provider when call_endpoint is set (#46839) * fix(proxy): count prompt-only token_counter requests with the provider when call_endpoint is set * fix(proxy): keep the local count for a prompt-only token_counter request when the provider count fails * fix(proxy): log a prompt-only token count fallback as a one-line warning * fix(proxy): send a prompt-only token_counter request to the provider only on Gemini and Vertex AI * fix(proxy): count a token_counter prompt sent next to an empty messages or contents list * fix(proxy): count a prompt next to messages on the local token count instead of answering 500 * fix(proxy): answer 400 for a token count prompt next to a conversation unless drop_params is set With drop_params the conversation is counted without the prompt. On the local count, a prompt sent with system or tools is counted as a user message with them, like the messages form. * test(proxy): count tools with a prompt in the token counter fallback test * test(proxy): cover prompt counting on /utils/token_counter in the Gemini count wire tests * test(proxy): name the token counter fallback test by what it checks --- litellm/proxy/proxy_server.py | 34 ++- .../test_vertex_gemini_count_tokens_wire.py | 114 ++++++- .../proxy/proxy_server/test_routes_utils.py | 7 +- tests/unit/proxy/test_proxy_token_counter.py | 279 +++++++++++++++++- 4 files changed, 421 insertions(+), 13 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index bbd9d96b84b..58037689597 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -14488,6 +14488,16 @@ async def _try_provider_token_count( return result +_PROVIDERS_COUNTING_WITHOUT_MESSAGES: Final = frozenset( + {litellm.LlmProviders.GEMINI.value, litellm.LlmProviders.VERTEX_AI.value} +) + + +def _drops_params(deployment: Mapping[str, object] | None) -> bool: + litellm_params: Final = as_str_mapping(None if deployment is None else deployment.get("litellm_params")) + return litellm.drop_params is True or (litellm_params is not None and litellm_params.get("drop_params") is True) + + def _system_message(system: object) -> ChatCompletionSystemMessage | None: if not isinstance(system, (str, list)) or not system: return None @@ -14554,8 +14564,8 @@ async def count_request_tokens( litellm_model_name = litellm_model_name.split("/", 1)[1] model_to_use: str = ( - litellm_model_name or request.model - ) # use litellm model name, if it's not avalable then fallback to request.model + litellm_model_name if isinstance(litellm_model_name, str) and litellm_model_name else request.model + ) # Try provider-specific token counting first - only for non-direct requests (from provider endpoints) provider_counter: BaseTokenCounter | None = None @@ -14566,13 +14576,23 @@ async def count_request_tokens( if _model is not None: model_to_use = _model + conversation_empty: Final = not messages and not contents + if prompt and not conversation_empty and not _drops_params(deployment): + raise ProxyException( + message="prompt cannot be sent with messages or contents. Set drop_params to count them without the prompt", + type="invalid_request_error", + param="prompt", + code=400, + ) + prompt_only: Final = bool(prompt) and conversation_empty if provider_counter is not None: + prompt_as_message: Final = prompt_only and custom_llm_provider in _PROVIDERS_COUNTING_WITHOUT_MESSAGES result: Final = await _try_provider_token_count( provider_counter=provider_counter, custom_llm_provider=custom_llm_provider, model_to_use=model_to_use, - messages=messages, - contents=contents, + messages=[{"role": "user", "content": prompt}] if prompt_as_message else messages, + contents=None if prompt_as_message else contents, deployment=deployment, request_model=request.model, tools=tools, @@ -14605,8 +14625,10 @@ async def count_request_tokens( tokenizer_used: Final = str(_tokenizer_used["type"]) system_message: Final = _system_message(system) + prompt_turn: Final = prompt_only and not (system_message is None and not tools) typed_messages: Final = cast( # cast-ok: request messages are raw chat-shaped dicts that token_counter normalizes - Sequence[AllMessageValues] | None, messages + Sequence[AllMessageValues] | None, + [{"role": "user", "content": prompt}] if prompt_turn else (None if prompt_only else messages), ) counted_messages: Final = ( typed_messages if typed_messages is None or system_message is None else (system_message, *typed_messages) @@ -14616,7 +14638,7 @@ async def count_request_tokens( ) total_tokens: Final = await offload_token_count(litellm.token_counter)( model=model_to_use, - text=prompt, + text=prompt if counted_messages is None else None, messages=counted_messages, tools=counted_tools, custom_tokenizer=_tokenizer_used, diff --git a/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py b/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py index 5fe82d37e9e..8de93256f6c 100644 --- a/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py +++ b/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py @@ -65,7 +65,7 @@ class _GeminiPeer: return requests -def _deployment(gateway: Gateway, scenario: Scenario, api_base: str) -> str: +def _deployment(gateway: Gateway, scenario: Scenario, api_base: str, **parameters: JsonValue) -> str: return scenario.model( model=f"vertex_ai/{_MODEL}", api_base=api_base, @@ -73,6 +73,7 @@ def _deployment(gateway: Gateway, scenario: Scenario, api_base: str) -> str: vertex_project=_PROJECT, vertex_location=_LOCATION, vertex_credentials=service_account_json(_PROJECT, gateway.upstream_url.rstrip("/")), + **parameters, ) @@ -220,8 +221,8 @@ _CONTENT_BLOCK_BODIES: Final[Mapping[str, Mapping[str, JsonValue]]] = MappingPro ) -def _ai_studio_deployment(scenario: Scenario, api_base: str) -> str: - return scenario.model(model=f"gemini/{_MODEL}", api_base=api_base, api_key="scripted-gemini-key") +def _ai_studio_deployment(scenario: Scenario, api_base: str, **parameters: JsonValue) -> str: + return scenario.model(model=f"gemini/{_MODEL}", api_base=api_base, api_key="scripted-gemini-key", **parameters) def _counted_contents(count_body: Mapping[str, JsonValue]) -> JsonValue: @@ -594,3 +595,110 @@ def test_concurrent_count_waves_recover_after_vertex_count_outage(gateway: Gatew tuple(_json_body(request) for request in requests) == ({"contents": [{"role": "user", "parts": [{"text": _USER}]}]},) * 9 ) + + +_CONTENTS: Final[JsonValue] = [{"role": "user", "parts": [{"text": _USER}]}] +_OPENAI_TOOLS: Final[JsonValue] = [ + { + "type": "function", + "function": {"name": "noop", "description": "Return no output", "parameters": {"type": "object"}}, + } +] + + +def _provider_deployment( + gateway: Gateway, scenario: Scenario, provider: str, api_base: str, **parameters: JsonValue +) -> str: + if provider == "vertex_ai": + return _deployment(gateway, scenario, api_base, **parameters) + return _ai_studio_deployment(scenario, api_base, **parameters) + + +def _count(gateway: Gateway, body: Mapping[str, JsonValue], call_endpoint: bool = True) -> httpx.Response: + return gateway.request( + "POST", "/utils/token_counter", body, params={"call_endpoint": "true" if call_endpoint else "false"} + ) + + +def _total_tokens(response: httpx.Response) -> JsonValue: + return _JSON_OBJECT.validate_json(response.content)["total_tokens"] + + +@pytest.mark.parametrize( + ("provider", "empty_conversation"), + (("vertex_ai", {}), ("gemini", {}), ("gemini", {"messages": []}), ("gemini", {"contents": []})), + ids=("vertex-prompt-only", "ai-studio-prompt-only", "ai-studio-empty-messages", "ai-studio-empty-contents"), +) +def test_utils_token_counter_sends_a_prompt_as_the_user_turn_of_its_messages_form( + gateway: Gateway, provider: str, empty_conversation: Mapping[str, JsonValue] +) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = _provider_deployment(gateway, scenario, provider, wire.url) + prompted: Final = _count(gateway, {"model": model, "prompt": _USER, **empty_conversation}) + prompted_requests: Final = _count_requests(wire) + as_message: Final = _count(gateway, {"model": model, "messages": _MESSAGES}) + message_requests: Final = _count_requests(wire) + assert prompted.status_code == 200, prompted.text + assert as_message.status_code == 200, as_message.text + assert _total_tokens(prompted) == _COUNT, prompted.text + assert len(prompted_requests) == 1 + assert len(message_requests) == 1 + assert _counted_contents(_json_body(prompted_requests[0])) == _counted_contents(_json_body(message_requests[0])) + + +@pytest.mark.parametrize( + ("conversation", "call_endpoint"), + (({"messages": _MESSAGES}, True), ({"contents": _CONTENTS}, True), ({"messages": _MESSAGES}, False)), + ids=("provider-count-messages", "provider-count-contents", "local-count-messages"), +) +def test_utils_token_counter_rejects_a_prompt_next_to_a_conversation_without_calling_the_provider( + gateway: Gateway, conversation: Mapping[str, JsonValue], call_endpoint: bool +) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = _ai_studio_deployment(scenario, wire.url) + response: Final = _count(gateway, {"model": model, "prompt": "hi", **conversation}, call_endpoint) + requests: Final = _count_requests(wire) + assert response.status_code == 400, response.text + assert _JSON_OBJECT.validate_python(_JSON_OBJECT.validate_json(response.content)["error"])["param"] == "prompt" + assert requests == () + + +@pytest.mark.parametrize("call_endpoint", (True, False), ids=("provider-count", "local-count")) +def test_utils_token_counter_counts_only_the_conversation_when_drop_params_drops_the_prompt( + gateway: Gateway, call_endpoint: bool +) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = _ai_studio_deployment(scenario, wire.url, drop_params=True) + mixed: Final = _count(gateway, {"model": model, "prompt": "hi", "messages": _MESSAGES}, call_endpoint) + mixed_requests: Final = _count_requests(wire) + alone: Final = _count(gateway, {"model": model, "messages": _MESSAGES}, call_endpoint) + alone_requests: Final = _count_requests(wire) + assert mixed.status_code == 200, mixed.text + assert alone.status_code == 200, alone.text + assert _total_tokens(mixed) == _total_tokens(alone) + assert [_counted_contents(_json_body(request)) for request in mixed_requests] == [ + _counted_contents(_json_body(request)) for request in alone_requests + ] + + +@pytest.mark.parametrize( + "extra", + ({"system": "Follow the system instruction."}, {"tools": _OPENAI_TOOLS}), + ids=("system", "tools"), +) +def test_utils_token_counter_counts_a_local_prompt_with_system_or_tools_like_its_messages_form( + gateway: Gateway, extra: Mapping[str, JsonValue] +) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = _ai_studio_deployment(scenario, wire.url) + prompted: Final = _count(gateway, {"model": model, "prompt": _USER, **extra}, call_endpoint=False) + as_message: Final = _count(gateway, {"model": model, "messages": _MESSAGES, **extra}, call_endpoint=False) + requests: Final = _count_requests(wire) + assert prompted.status_code == 200, prompted.text + assert as_message.status_code == 200, as_message.text + assert _total_tokens(prompted) == _total_tokens(as_message) + assert requests == () diff --git a/tests/unit/proxy/proxy_server/test_routes_utils.py b/tests/unit/proxy/proxy_server/test_routes_utils.py index 295b9e5ae8b..118b80da31e 100644 --- a/tests/unit/proxy/proxy_server/test_routes_utils.py +++ b/tests/unit/proxy/proxy_server/test_routes_utils.py @@ -390,8 +390,7 @@ def test_token_counter_fallback_counts_tools_system_and_anthropic_blocks(client, assert full > bare -def test_token_counter_fallback_prompt_with_tools_does_not_500(client, auth_as, monkeypatch): - """Regression: a ``prompt`` request carrying ``tools`` but no ``messages`` still counts, because the fallback attaches tools only when counting messages (``token_counter`` rejects tools on the text path).""" +def test_token_counter_fallback_counts_a_prompt_with_tools_like_its_messages_form(client, auth_as, monkeypatch): monkeypatch.setattr(proxy_server, "llm_router", None) monkeypatch.setattr(litellm, "disable_token_counter", False, raising=False) prompt = "count the tokens in this sentence please" @@ -412,4 +411,6 @@ def test_token_counter_fallback_prompt_with_tools_does_not_500(client, auth_as, ) assert response.status_code == 200, response.text - assert response.json()["total_tokens"] == litellm.token_counter(model="claude-fable-5", text=prompt) + assert response.json()["total_tokens"] == litellm.token_counter( + model="claude-fable-5", messages=[{"role": "user", "content": prompt}], tools=tools + ) diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index fbe47e4a3d7..7d70e058396 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -9,7 +9,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest from dotenv import load_dotenv -from pydantic import JsonValue +from pydantic import BaseModel, ConfigDict, JsonValue load_dotenv() @@ -21,10 +21,13 @@ from fastapi.responses import JSONResponse import litellm from litellm import Router +from litellm.llms.anthropic.count_tokens import handler as anthropic_count_tokens_handler from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.bedrock.count_tokens.bedrock_token_counter import BedrockTokenCounter from litellm.llms.bedrock.count_tokens.handler import BedrockCountTokensHandler from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +from litellm.llms.gemini.count_tokens import handler as gemini_count_tokens_handler +from litellm.llms.vertex_ai.vertex_ai_partner_models.count_tokens.handler import VertexAIPartnerModelsTokenCounter from litellm.proxy._types import ProxyException, TokenCountRequest from litellm.proxy.anthropic_endpoints.endpoints import ( count_tokens as anthropic_count_tokens, @@ -1174,3 +1177,277 @@ async def test_anthropic_endpoint_429_rate_limit_error_format(): finally: anthropic_endpoints.read_request_body = original_read_request_body proxy_server.count_request_tokens = original_token_counter + + +class _GeminiPart(BaseModel): + model_config = ConfigDict(frozen=True) + text: str = "" + + +class _GeminiContent(BaseModel): + model_config = ConfigDict(frozen=True) + parts: tuple[_GeminiPart, ...] = () + + +class _GeminiCountTokensBody(BaseModel): + model_config = ConfigDict(frozen=True) + contents: tuple[_GeminiContent, ...] = () + + +def _word_count(content: _GeminiContent) -> int: + return sum(len(part.text.split()) for part in content.parts) + + +def _gemini_word_counting_peer(request: httpx.Request) -> httpx.Response: + body: Final = _GeminiCountTokensBody.model_validate_json(request.content) + total: Final = sum(_word_count(content) for content in body.contents) + return httpx.Response(200, json={"totalTokens": total}, request=request) + + +def _unavailable_peer(request: httpx.Request) -> httpx.Response: + return httpx.Response(503, json={"error": {"message": "unavailable"}}, request=request) + + +def _gemini_router(drop_params: bool = False) -> Router: + return Router( + model_list=[ + { + "model_name": "gemini-counted", + "litellm_params": { + "model": "gemini/gemini-3.8-flash", + "api_key": "test-key", + "api_base": "https://count-tokens.test", + "drop_params": drop_params, + }, + } + ] + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "text", ["hello there", "The quick brown fox jumps over the lazy dog. " * 50], ids=["short", "long"] +) +async def test_call_endpoint_counts_a_prompt_like_the_same_text_sent_as_a_user_message( + text: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", _gemini_router()) + client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(_gemini_word_counting_peer)) + + with patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client): + from_prompt: Final = await token_counter( + request=TokenCountRequest(model="gemini-counted", prompt=text), call_endpoint=True + ) + from_message: Final = await token_counter( + request=TokenCountRequest(model="gemini-counted", messages=[{"role": "user", "content": text}]), + call_endpoint=True, + ) + + assert from_prompt.error is False + assert from_prompt.total_tokens == from_message.total_tokens == len(text.split()) + + +class _PartnerMessage(BaseModel): + model_config = ConfigDict(frozen=True) + content: str = "" + + +class _PartnerCountRequest(BaseModel): + model_config = ConfigDict(frozen=True) + messages: tuple[_PartnerMessage, ...] = () + + +async def _vertex_partner_word_count( + model: str, request_data: JsonValue, litellm_params: JsonValue +) -> dict[str, JsonValue]: + body: Final = _PartnerCountRequest.model_validate(request_data) + return {"input_tokens": sum(len(message.content.split()) for message in body.messages)} + + +@pytest.mark.asyncio +async def test_call_endpoint_counts_a_prompt_on_a_vertex_ai_deployment(monkeypatch: pytest.MonkeyPatch) -> None: + router: Final = Router( + model_list=[ + { + "model_name": "vertex-counted", + "litellm_params": { + "model": "vertex_ai/claude-sonnet-5-5", + "vertex_project": "test-project", + "vertex_location": "us-east5", + }, + } + ] + ) + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", router) + text: Final = "The quick brown fox jumps over the lazy dog." + + with patch.object( + VertexAIPartnerModelsTokenCounter, "handle_count_tokens_request", side_effect=_vertex_partner_word_count + ): + from_prompt: Final = await token_counter( + request=TokenCountRequest(model="vertex-counted", prompt=text), call_endpoint=True + ) + + assert (from_prompt.error, from_prompt.total_tokens) == (False, 9) + + +_FOX: Final = "The quick brown fox jumps over the lazy dog." + + +_CONVERSATIONS: Final = ( + {"messages": [{"role": "user", "content": _FOX}]}, + {"contents": [{"role": "user", "parts": [{"text": _FOX}]}]}, +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("conversation", _CONVERSATIONS, ids=["messages", "contents"]) +@pytest.mark.parametrize("call_endpoint", [True, False], ids=["provider", "local"]) +async def test_prompt_next_to_a_conversation_answers_400_without_drop_params( + conversation: dict[str, JsonValue], call_endpoint: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", _gemini_router()) + client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(_gemini_word_counting_peer)) + + with ( + patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client), + pytest.raises(ProxyException) as rejected, + ): + await token_counter( + request=TokenCountRequest(model="gemini-counted", prompt="hi", **conversation), call_endpoint=call_endpoint + ) + + assert (rejected.value.code, rejected.value.param) == ("400", "prompt") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("conversation", "call_endpoint", "global_drop_params", "deployment_drop_params"), + [ + (_CONVERSATIONS[0], True, False, True), + (_CONVERSATIONS[1], True, False, True), + (_CONVERSATIONS[0], False, False, True), + (_CONVERSATIONS[0], False, True, False), + ], + ids=[ + "provider_messages", + "provider_contents", + "local_messages", + "local_messages_global_setting", + ], +) +async def test_drop_params_counts_a_conversation_without_the_prompt( + conversation: dict[str, JsonValue], + call_endpoint: bool, + global_drop_params: bool, + deployment_drop_params: bool, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", _gemini_router(drop_params=deployment_drop_params)) + monkeypatch.setattr(litellm, "drop_params", global_drop_params) + client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(_gemini_word_counting_peer)) + + with patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client): + with_prompt: Final = await token_counter( + request=TokenCountRequest(model="gemini-counted", prompt="hi", **conversation), call_endpoint=call_endpoint + ) + conversation_only: Final = await token_counter( + request=TokenCountRequest(model="gemini-counted", **conversation), call_endpoint=call_endpoint + ) + + assert with_prompt.error is False + assert with_prompt.total_tokens == conversation_only.total_tokens > 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("empty_conversation", [{"messages": []}, {"contents": []}], ids=["messages", "contents"]) +async def test_call_endpoint_prompt_next_to_an_empty_conversation_counts_the_prompt( + empty_conversation: dict[str, JsonValue], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", _gemini_router()) + client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(_gemini_word_counting_peer)) + + with patch.object(gemini_count_tokens_handler, "get_async_httpx_client", return_value=client): + counted: Final = await token_counter( + request=TokenCountRequest(model="gemini-counted", prompt=_FOX, **empty_conversation), call_endpoint=True + ) + + assert (counted.error, counted.total_tokens) == (False, 9) + + +@pytest.mark.asyncio +async def test_local_count_of_a_prompt_next_to_empty_messages_counts_the_prompt( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", _gemini_router()) + + with_empty: Final = await token_counter(request=TokenCountRequest(model="gemini-counted", prompt=_FOX, messages=[])) + alone: Final = await token_counter(request=TokenCountRequest(model="gemini-counted", prompt=_FOX)) + + assert with_empty.error is False + assert with_empty.total_tokens == alone.total_tokens > 0 + + +_WEATHER_TOOL: Final = { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather for a city", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}, + }, +} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "chat_fields", + [{"system": "You answer in one short sentence. " * 20}, {"tools": [_WEATHER_TOOL]}], + ids=["system", "tools"], +) +@pytest.mark.parametrize("empty_conversation", [{}, {"messages": []}], ids=["no_messages", "empty_messages"]) +async def test_local_count_of_a_prompt_with_system_or_tools_counts_them_like_the_message_form( + chat_fields: dict[str, JsonValue], empty_conversation: dict[str, JsonValue], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", _gemini_router()) + + from_prompt: Final = await token_counter( + request=TokenCountRequest(model="gemini-counted", prompt=_FOX, **empty_conversation, **chat_fields) + ) + from_message: Final = await token_counter( + request=TokenCountRequest(model="gemini-counted", messages=[{"role": "user", "content": _FOX}], **chat_fields) + ) + prompt_alone: Final = await token_counter(request=TokenCountRequest(model="gemini-counted", prompt=_FOX)) + + assert from_prompt.error is False + assert from_prompt.total_tokens == from_message.total_tokens > prompt_alone.total_tokens + + +@pytest.mark.asyncio +async def test_call_endpoint_prompt_keeps_the_local_count_on_a_counter_that_needs_messages( + monkeypatch: pytest.MonkeyPatch, +) -> None: + router: Final = Router( + model_list=[ + { + "model_name": "anthropic-counted", + "litellm_params": { + "model": "anthropic/claude-sonnet-5-5", + "api_key": "test-key", + "api_base": "https://count-tokens.test", + }, + } + ] + ) + monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", router) + client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(_unavailable_peer)) + request: Final = TokenCountRequest(model="anthropic-counted", prompt="The quick brown fox jumps over the lazy dog.") + + with patch.object(anthropic_count_tokens_handler, "get_async_httpx_client", return_value=client): + with_endpoint: Final = await token_counter(request=request, call_endpoint=True) + local: Final = await token_counter(request=request, call_endpoint=False) + + assert (with_endpoint.error, with_endpoint.total_tokens, with_endpoint.tokenizer_type) == ( + False, + local.total_tokens, + local.tokenizer_type, + ) From 4e2dd6554917dc732d7708e3d85b0a43df3b73dc Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 16 Oct 2026 12:00:08 +0000 Subject: [PATCH 16/23] fix(interactions): keep background interaction bills through crashes, late restarts, and deployment keys (#45743) * fix(interactions): keep background interaction bills through crashes, late restarts, and deployment keys A settlement claim is now a lease, so a row whose settler died between claim and outcome is claimable again after the lease and boot resumes it. A resumed poll gets the full poll window and fetches with the serving deployment's credentials from the live router. The stored create context keeps the request tags and the request metadata that the creating worker's own spend row used * fix(interactions): bill a lapsed claim's takeover only when no spend row exists A settler that billed and then died, or failed to record the outcome, left the row unsettled, so the replica that took over its lapsed claim billed the interaction a second time. A takeover now looks up the spend row keyed by the interaction id and records the row as billed when it finds one. * fix(interactions): settle Vertex AI and edited deployments, never double bill with spend logs off A takeover with spend logs off records the interaction unsettled instead of billing it again, and a release records its outcome before freeing the budget reservation. A resumed or live fetch carries the deployment's Vertex AI project, location, and credentials, and a stale deployment id falls back to the model group only when one same-provider deployment serves it. A poll stops at the first 400, 401, 403, or 404, string request tags no longer fail the create, and settled_at gets an index built concurrently. * fix(interactions): leave a live replica's interaction to it when a resumed poll cannot read it A replica that boots with credentials unable to read an interaction got a 404 on its resumed poll and settled the row as unsettled at once, so the replica that created it, still polling, found it settled and never billed it. A resumed poll that gets a 400, 401, 403, or 404 now stops fetching and keeps checking the row until its poll window ends, and settles it as unsettled only when no other replica did. * fix(interactions): keep the creator's metadata on a resumed bill and retry outcome writes inside the claim lease A bill another replica writes now carries the metadata the creating worker's own spend row is built from, so a client's metadata.model_group no longer relabels it. A failed outcome write is retried for half the claim lease before the budget reservation is left to a takeover, and the creating replica retries every fetch error for its whole window again, as main does. The resume integration case now checks that a resumed fetch uses the deployment's own key. * fix(interactions): resume background interaction polls only after database models load Boot resume ran before the proxy loaded the models stored in the database, so a resumed poll for a model added through /model/new or the Admin UI fetched with env credentials. A 4xx there stops a resumed poll, so the bill was lost. The settlement now installs after the database models load. The integration test resumes a database model with an eager first fetch and fails with the old startup order. * fix(interactions): renew a settling claim so a slow settle is never taken over and billed twice A settle runs every success callback in turn, and one slow callback could hold it past the 5 minute claim lease. A booting replica or a DELETE then took the live claim over, found no spend row yet, and billed the interaction a second time. The settler now renews its claim every third of the lease until it records the outcome, so only a settler that stopped renewing is taken over. A takeover also read the spend log after writing its claim, so a failed read left the row under the taker's own fresh claim, which its poll then took for another settler's and left unbilled until a reboot. The spend row is now read before the claim is written, so a failed read takes nothing and the claim is retried on the next poll. * test(spend): cover lapsed claim takeovers, rotated deployment keys, late restarts, and resumed bill tags --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../migration.sql | 2 + .../litellm_proxy_extras/schema.prisma | 1 + litellm/constants.py | 1 + .../interactions/background_cost_polling.py | 289 +++++++++-- litellm/proxy/proxy_server.py | 5 +- litellm/proxy/schema.prisma | 1 + .../background_interaction_settlement.py | 207 +++++++- schema.prisma | 1 + .../test_background_interaction_settlement.py | 181 ++++++- .../test_background_cost_polling.py | 357 ++++++++++++- .../test_background_interaction_settlement.py | 479 +++++++++++++++++- 11 files changed, 1400 insertions(+), 124 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20261016000000_background_interaction_settlement_settled_at_index/migration.sql diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20261016000000_background_interaction_settlement_settled_at_index/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261016000000_background_interaction_settlement_settled_at_index/migration.sql new file mode 100644 index 00000000000..81d32345794 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20261016000000_background_interaction_settlement_settled_at_index/migration.sql @@ -0,0 +1,2 @@ +CREATE INDEX CONCURRENTLY IF NOT EXISTS "idx_background_interaction_settlement_settled_at" +ON "LiteLLM_BackgroundInteractionSettlement" ("settled_at"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index d40bac380aa..83a0e6f9e05 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -1981,6 +1981,7 @@ model LiteLLM_BackgroundInteractionSettlement { outcome String? @@index([claimed_at], map: "idx_background_interaction_settlement_claimed_at") + @@index([settled_at], map: "idx_background_interaction_settlement_settled_at") } model LiteLLM_Lens { diff --git a/litellm/constants.py b/litellm/constants.py index 665ab771726..271e42ee661 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1895,6 +1895,7 @@ _background_interaction_cost_polling_env: Final = os.getenv( "BACKGROUND_INTERACTION_COST_POLLING_ENABLED", "true" ).lower() BACKGROUND_INTERACTION_COST_POLLING_ENABLED: Final = _background_interaction_cost_polling_env == "true" +BACKGROUND_INTERACTION_SETTLEMENT_CLAIM_LEASE_SECONDS: Final = 300.0 PROXY_BUDGET_RESCHEDULER_MAX_TIME: Final = get_env_int("PROXY_BUDGET_RESCHEDULER_MAX_TIME", 605) PROXY_BATCH_WRITE_AT: Final = get_env_int("PROXY_BATCH_WRITE_AT", 10) # in seconds, increased from 10 PROXY_CONFIG_RELOAD_INTERVAL_SECONDS: Final = get_env_int("PROXY_CONFIG_RELOAD_INTERVAL_SECONDS", 30) diff --git a/litellm/interactions/background_cost_polling.py b/litellm/interactions/background_cost_polling.py index 21d438d9621..275eb17787e 100644 --- a/litellm/interactions/background_cost_polling.py +++ b/litellm/interactions/background_cost_polling.py @@ -32,19 +32,24 @@ task to settle. A ``BackgroundSettlementStore`` makes the pending settlement durable across processes: the create registers the request context that billing needs (never provider credentials), the settlement is claimed exactly once through the store, and a delete on any replica rebuilds the -billing context from the store when the poll task is not local. Rows left -unclaimed by a process that died are resumed at startup. The default store is -in-memory, which keeps the SDK and single-process behavior unchanged; the -proxy installs a database-backed one. +billing context from the store when the poll task is not local. Every row +still unsettled when a process starts is resumed there, including one whose +settler died between its claim and its outcome: the database-backed store +treats a claim as a lease that its settler renews until it records the +outcome, so such a row is claimable again once the lease runs out. The +default store is in-memory, which keeps the SDK and single-process behavior +unchanged; the proxy installs a database-backed one. """ import asyncio +import itertools from collections.abc import Awaitable, Callable, Iterable, Iterator, Mapping, Sequence from dataclasses import dataclass, field from datetime import datetime, timezone from types import MappingProxyType from typing import TYPE_CHECKING, Final, Literal, Protocol, TypeAlias +from openai import APIStatusError from pydantic import ConfigDict, JsonValue, TypeAdapter, ValidationError from pydantic_core import PydanticSerializationError, to_jsonable_python @@ -54,10 +59,12 @@ from litellm.constants import ( BACKGROUND_INTERACTION_COST_POLL_MAX_INTERVAL_SECONDS, BACKGROUND_INTERACTION_COST_POLL_TIMEOUT_SECONDS, BACKGROUND_INTERACTION_COST_POLLING_ENABLED, + BACKGROUND_INTERACTION_SETTLEMENT_CLAIM_LEASE_SECONDS, ) from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kwargs from litellm.types.interactions import InteractionsAPIResponse from litellm.types.llms.base import LiteLLMBaseModel +from litellm.types.router import CredentialLiteLLMParams from litellm.types.utils import CustomPricingLiteLLMParams if TYPE_CHECKING: @@ -101,6 +108,7 @@ class PendingBackgroundInteraction: custom_llm_provider: str create_context: BackgroundInteractionCreateContext created_at: datetime + claim_expires_at: datetime | None = None class BackgroundSettlementStore(Protocol): @@ -112,9 +120,11 @@ class BackgroundSettlementStore(Protocol): async def claim(self, interaction_id: str) -> bool: ... + async def renew_claim(self, interaction_id: str) -> None: ... + async def record_outcome(self, interaction_id: str, outcome: SettlementOutcome) -> None: ... - async def unclaimed(self) -> Sequence[PendingBackgroundInteraction]: ... + async def unsettled(self) -> Sequence[PendingBackgroundInteraction]: ... @dataclass(frozen=True, slots=True) @@ -145,10 +155,13 @@ class InMemoryBackgroundSettlementStore: self._rows[interaction_id] = None return True + async def renew_claim(self, interaction_id: str) -> None: + return None + async def record_outcome(self, interaction_id: str, outcome: SettlementOutcome) -> None: return None - async def unclaimed(self) -> Sequence[PendingBackgroundInteraction]: + async def unsettled(self) -> Sequence[PendingBackgroundInteraction]: return tuple(row for row in self._rows.values() if row is not None) @@ -169,11 +182,15 @@ class BackgroundInteractionPollContext: interaction_id: str custom_llm_provider: str logging_obj: "LiteLLMLoggingObj" - api_key: str | None = None - api_base: str | None = None + credentials: CredentialLiteLLMParams = field(default_factory=CredentialLiteLLMParams) + deployment_id: str | None = None + model_group: str | None = None initial_interval_seconds: float = BACKGROUND_INTERACTION_COST_POLL_INITIAL_INTERVAL_SECONDS max_interval_seconds: float = BACKGROUND_INTERACTION_COST_POLL_MAX_INTERVAL_SECONDS timeout_seconds: float = BACKGROUND_INTERACTION_COST_POLL_TIMEOUT_SECONDS + start_delay_seconds: float = 0.0 + outcome_write_timeout_seconds: float = BACKGROUND_INTERACTION_SETTLEMENT_CLAIM_LEASE_SECONDS / 2 + claim_renewal_interval_seconds: float = BACKGROUND_INTERACTION_SETTLEMENT_CLAIM_LEASE_SECONDS / 3 store: BackgroundSettlementStore = field(default_factory=InMemoryBackgroundSettlementStore) resumed: bool = False @@ -184,15 +201,41 @@ FetchInteraction: TypeAlias = Callable[[BackgroundInteractionPollContext], Await async def fetch_background_interaction(context: BackgroundInteractionPollContext) -> InteractionsAPIResponse: from litellm.interactions import aget + credentials: Final = context.credentials return await aget( interaction_id=context.interaction_id, custom_llm_provider=context.custom_llm_provider, - api_key=context.api_key, - api_base=context.api_base, + api_key=credentials.api_key, + api_base=credentials.api_base, + vertex_project=credentials.vertex_project, + vertex_location=credentials.vertex_location, + vertex_credentials=credentials.vertex_credentials, **{"no-log": True}, ) +_FETCH_CREDENTIAL_FIELDS: Final = ("api_key", "api_base", "vertex_project", "vertex_location", "vertex_credentials") + + +def _fetch_credentials(request_kwargs: Mapping[str, object]) -> CredentialLiteLLMParams: + try: + return CredentialLiteLLMParams.model_validate( + {field_name: request_kwargs.get(field_name) for field_name in _FETCH_CREDENTIAL_FIELDS} + ) + except ValidationError: + verbose_logger.debug( + "Fetching a background interaction with environment credentials: the request's are malformed" + ) + return CredentialLiteLLMParams() + + +_PERMANENT_FETCH_STATUS_CODES: Final = frozenset({400, 401, 403, 404}) + + +def _is_permanent_fetch_error(error: Exception) -> bool: + return isinstance(error, APIStatusError) and error.status_code in _PERMANENT_FETCH_STATUS_CODES + + def _poll_intervals(initial: float, maximum: float, timeout: float) -> Iterator[float]: elapsed = 0.0 # rebind-ok: the schedule accumulates the time it has already yielded interval = initial # rebind-ok: the schedule doubles the interval up to the cap @@ -231,6 +274,7 @@ _UNCARRIED_METADATA_KEY: Final = "user_api_key_auth" _JSON_VALUE: Final = TypeAdapter(JsonValue) _STRING: Final = TypeAdapter(str) _OBJECT_MAPPING: Final = TypeAdapter(Mapping[str, object]) +_OBJECT_DICT: Final = TypeAdapter(dict[str, object]) def _carries(key: str) -> bool: @@ -256,10 +300,36 @@ def _as_datetime(start_time: datetime | float) -> datetime: return start_time if isinstance(start_time, datetime) else datetime.fromtimestamp(start_time, tz=timezone.utc) +def _spend_metadata(logging_obj: "LiteLLMLoggingObj") -> Mapping[str, object]: + """ + The metadata the creating worker's own spend row is built from, with the + request tags computed once, header-derived ones included, because a + rebuilt logging object has no request headers to derive them from. + """ + litellm_params: Final = _OBJECT_DICT.validate_python(logging_obj.model_call_details.get("litellm_params") or {}) + return { + **_OBJECT_MAPPING.validate_python(get_litellm_metadata_from_kwargs(kwargs=logging_obj.model_call_details)), + **_request_tags(litellm_params), + } + + +def _request_tags(litellm_params: Mapping[str, object]) -> Mapping[str, object]: + from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup + + try: + return { + "tags": StandardLoggingPayloadSetup.get_request_tags( + litellm_params=_OBJECT_DICT.validate_python(litellm_params), + proxy_server_request=_OBJECT_DICT.validate_python(litellm_params.get("proxy_server_request") or {}), + ) + } + except AttributeError: + verbose_logger.debug("Keeping a background interaction's request tags as sent: they are not a list") + return {} + + def _create_context(logging_obj: "LiteLLMLoggingObj", custom_llm_provider: str) -> BackgroundInteractionCreateContext: - metadata: Final = _OBJECT_MAPPING.validate_python( - get_litellm_metadata_from_kwargs(kwargs=logging_obj.model_call_details) - ) + metadata: Final = _spend_metadata(logging_obj) litellm_params: Final = _OBJECT_MAPPING.validate_python(logging_obj.litellm_params) model: Final = logging_obj.model_call_details.get("model") return BackgroundInteractionCreateContext( @@ -313,6 +383,16 @@ async def _settled_elsewhere(context: BackgroundInteractionPollContext) -> bool: return False +async def _settled_elsewhere_before_the_window_ends( + context: BackgroundInteractionPollContext, intervals: Iterator[float] +) -> bool: + for interval in intervals: + await asyncio.sleep(interval) + if await _settled_elsewhere(context): + return True + return False + + async def _claim(context: BackgroundInteractionPollContext) -> bool | None: """ Exactly-once gate between every settler of one interaction, on every @@ -326,30 +406,118 @@ async def _claim(context: BackgroundInteractionPollContext) -> bool | None: return None -async def _record(context: BackgroundInteractionPollContext, outcome: SettlementOutcome) -> SettlementOutcome: +async def _wrote_outcome(context: BackgroundInteractionPollContext, outcome: SettlementOutcome) -> bool: try: await context.store.record_outcome(context.interaction_id, outcome) - except Exception: # noqa: BLE001 # the outcome is an audit trail; the claim already made the settlement exclusive - verbose_logger.exception("Could not record the settlement of background interaction %s", context.interaction_id) + except Exception as e: # noqa: BLE001 # a failed write is retried while the claim still shuts out takeovers + verbose_logger.warning( + "Could not record the settlement of background interaction %s: %s", context.interaction_id, e + ) + return False + return True + + +async def _recorded(context: BackgroundInteractionPollContext, outcome: SettlementOutcome) -> bool: + """ + Retries a failed write for half the claim lease, while no other settler + can take the row over, so a database blip does not leave the row to a + takeover or hold the budget reservation that only a recorded release frees. + """ + retries: Final = _poll_intervals( + initial=context.initial_interval_seconds, + maximum=context.max_interval_seconds, + timeout=context.outcome_write_timeout_seconds, + ) + for delay in itertools.chain((0.0,), retries): + await asyncio.sleep(delay) + if await _wrote_outcome(context, outcome): + return True + verbose_logger.error( + "Could not record the settlement of background interaction %s within %ss; the next process to take " + "over its lapsed claim settles it", + context.interaction_id, + context.outcome_write_timeout_seconds, + ) + return False + + +async def _bill(context: BackgroundInteractionPollContext, response: InteractionsAPIResponse) -> SettlementOutcome: + await _bill_settled_interaction(logging_obj=context.logging_obj, response=response) + await _recorded(context, "billed") + return "billed" + + +async def _release( + context: BackgroundInteractionPollContext, outcome: Literal["released", "unsettled"] +) -> SettlementOutcome: + """ + A release writes no spend row, so a takeover cannot tell whether the + settler before it released. The outcome is recorded first and the + reservation released only once it is, which leaves a takeover with + nothing to release twice. + """ + if await _recorded(context, outcome): + await _release_open_budget_reservation(logging_obj=context.logging_obj) return outcome +async def _renew_claim_until_cancelled(context: BackgroundInteractionPollContext) -> None: + while True: + await asyncio.sleep(context.claim_renewal_interval_seconds) + try: + await context.store.renew_claim(context.interaction_id) + except Exception as e: # noqa: BLE001 # a missed renewal is retried next interval, well inside the lease + verbose_logger.warning( + "Could not renew the settlement claim of background interaction %s: %s", context.interaction_id, e + ) + + +async def _settle( + context: BackgroundInteractionPollContext, response: InteractionsAPIResponse | None +) -> SettlementOutcome: + """ + Billing runs every success callback in turn, and a slow one can outlast + the claim lease, so the claim is renewed until the outcome is recorded. + A settler that stops renewing has died, and only then can another take + the row over. + """ + renewal: Final = asyncio.create_task(_renew_claim_until_cancelled(context)) + try: + if response is not None and response.status in _TERMINAL_STATUSES and response.usage is not None: + return await _bill(context, response) + return await _release(context, "released") + finally: + renewal.cancel() + + async def poll_and_log_background_interaction_cost( context: BackgroundInteractionPollContext, fetch_interaction: FetchInteraction = fetch_background_interaction, ) -> SettlementOutcome | None: last_response: InteractionsAPIResponse | None = None # rebind-ok: the give-up path settles from the last poll - for interval in _poll_intervals( + await asyncio.sleep(context.start_delay_seconds) + intervals: Final = _poll_intervals( initial=context.initial_interval_seconds, maximum=context.max_interval_seconds, timeout=context.timeout_seconds, - ): + ) + for interval in intervals: await asyncio.sleep(interval) if await _settled_elsewhere(context): return None try: response = await fetch_interaction(context) except Exception as e: # noqa: BLE001 # any fetch error must not kill the billing poll loop + if context.resumed and _is_permanent_fetch_error(e): + verbose_logger.warning( + "Stopped fetching resumed background interaction %s: the provider answered %s, so this " + "process cannot read it and leaves it to the replica that created it until the poll window ends", + context.interaction_id, + e, + ) + if await _settled_elsewhere_before_the_window_ends(context, intervals): + return None + break verbose_logger.debug( "Background interaction cost poll for %s failed, will retry: %s", context.interaction_id, @@ -363,11 +531,11 @@ async def poll_and_log_background_interaction_cost( continue if not claimed: return None - return await _record(context, await _settle_terminal(logging_obj=context.logging_obj, response=response)) + return await _settle(context, response) if not await _claim(context): return None if last_response is not None and last_response.status in _TERMINAL_STATUSES: - return await _record(context, await _settle_terminal(logging_obj=context.logging_obj, response=last_response)) + return await _settle(context, last_response) if last_response is not None and last_response.status not in _POLLABLE_STATUSES: verbose_logger.error( "Gave up cost polling for background interaction %s after %ss: its last status %r is in neither " @@ -379,20 +547,10 @@ async def poll_and_log_background_interaction_cost( ) else: verbose_logger.warning( - "Gave up cost polling for background interaction %s after %ss; its usage will not be tracked", + "Gave up cost polling for background interaction %s; its usage will not be tracked", context.interaction_id, - context.timeout_seconds, ) - await _release_open_budget_reservation(logging_obj=context.logging_obj) - return await _record(context, "unsettled") - - -async def _settle_terminal(logging_obj: "LiteLLMLoggingObj", response: InteractionsAPIResponse) -> SettlementOutcome: - if response.status in _TERMINAL_STATUSES and response.usage is not None: - await _bill_settled_interaction(logging_obj=logging_obj, response=response) - return "billed" - await _release_open_budget_reservation(logging_obj=logging_obj) - return "released" + return await _release(context, "unsettled") async def _release_open_budget_reservation(logging_obj: "LiteLLMLoggingObj") -> None: @@ -405,7 +563,9 @@ async def _release_open_budget_reservation(logging_obj: "LiteLLMLoggingObj") -> settlement must release the reservation here or the spend counters stay pinned at the estimated cost. """ - metadata: Final = get_litellm_metadata_from_kwargs(kwargs=logging_obj.model_call_details) + metadata: Final = _OBJECT_MAPPING.validate_python( + get_litellm_metadata_from_kwargs(kwargs=logging_obj.model_call_details) + ) budget_reservation: Final = metadata.get("user_api_key_budget_reservation") if not isinstance(budget_reservation, dict): return @@ -518,13 +678,18 @@ class _UnverifiedRegistrationStore: return False return await self.local.claim(interaction_id) + async def renew_claim(self, interaction_id: str) -> None: + if await self.local.is_claimed(interaction_id): + return + await self.durable.renew_claim(interaction_id) + async def record_outcome(self, interaction_id: str, outcome: SettlementOutcome) -> None: if await self.local.is_claimed(interaction_id): return await self.durable.record_outcome(interaction_id, outcome) - async def unclaimed(self) -> Sequence[PendingBackgroundInteraction]: - return await self.durable.unclaimed() + async def unsettled(self) -> Sequence[PendingBackgroundInteraction]: + return await self.durable.unsettled() async def _registered_store( @@ -559,8 +724,6 @@ async def maybe_schedule_background_interaction_cost_polling( logging_obj: Final = create_kwargs.get("litellm_logging_obj") if not isinstance(logging_obj, Logging): return None - api_key: Final = create_kwargs.get("api_key") - api_base: Final = create_kwargs.get("api_base") pending: Final = PendingBackgroundInteraction( interaction_id=response.id, custom_llm_provider=custom_llm_provider, @@ -571,8 +734,7 @@ async def maybe_schedule_background_interaction_cost_polling( interaction_id=response.id, custom_llm_provider=custom_llm_provider, logging_obj=logging_obj, - api_key=api_key if isinstance(api_key, str) else None, - api_base=api_base if isinstance(api_base, str) else None, + credentials=_fetch_credentials(create_kwargs), store=await _registered_store(store or _STORE.store, pending), ) return _track_poll(context, fetch_interaction) @@ -603,10 +765,7 @@ async def _settle_before_delete( ) -> SettlementOutcome | None: if not await _claim(context): return None - if response is None: - await _release_open_budget_reservation(logging_obj=context.logging_obj) - return await _record(context, "released") - return await _record(context, await _settle_terminal(logging_obj=context.logging_obj, response=response)) + return await _settle(context, response) async def maybe_settle_background_interaction_before_delete( @@ -622,14 +781,11 @@ async def maybe_settle_background_interaction_before_delete( pending: Final = await _pending(settlement_store, interaction_id) if pending is None: return None - api_key: Final = delete_kwargs.get("api_key") - api_base: Final = delete_kwargs.get("api_base") context: Final = BackgroundInteractionPollContext( interaction_id=interaction_id, custom_llm_provider=pending.custom_llm_provider, logging_obj=_rebuild_logging_obj(pending.create_context), - api_key=api_key if isinstance(api_key, str) else None, - api_base=api_base if isinstance(api_base, str) else None, + credentials=_fetch_credentials(delete_kwargs), store=settlement_store, ) try: @@ -644,9 +800,9 @@ async def maybe_settle_background_interaction_before_delete( return await _settle_before_delete(context, response) -async def _unclaimed(store: BackgroundSettlementStore) -> Sequence[PendingBackgroundInteraction]: +async def _unsettled(store: BackgroundSettlementStore) -> Sequence[PendingBackgroundInteraction]: try: - return await store.unclaimed() + return await store.unsettled() except Exception: # noqa: BLE001 # an unreadable store at startup leaves its rows for the next boot verbose_logger.exception("Could not list the unsettled background interactions") return () @@ -662,17 +818,36 @@ class PollSchedule: DEFAULT_POLL_SCHEDULE: Final = PollSchedule() +def _deployment_id(create_context: BackgroundInteractionCreateContext) -> str | None: + model_info: Final = create_context.metadata.get("model_info") + deployment_id: Final = model_info.get("id") if isinstance(model_info, dict) else None + return deployment_id if isinstance(deployment_id, str) else None + + +def _model_group(create_context: BackgroundInteractionCreateContext) -> str | None: + model_group: Final = create_context.metadata.get("model_group") + return model_group if isinstance(model_group, str) else None + + +def _seconds_until(moment: datetime | None) -> float: + if moment is None: + return 0.0 + return max((moment - datetime.now(timezone.utc)).total_seconds(), 0.0) + + def _resumed_context( row: PendingBackgroundInteraction, store: BackgroundSettlementStore, schedule: PollSchedule ) -> BackgroundInteractionPollContext: - age_seconds: Final = (datetime.now(timezone.utc) - row.created_at).total_seconds() return BackgroundInteractionPollContext( interaction_id=row.interaction_id, custom_llm_provider=row.custom_llm_provider, logging_obj=_rebuild_logging_obj(row.create_context), + deployment_id=_deployment_id(row.create_context), + model_group=_model_group(row.create_context), initial_interval_seconds=schedule.initial_interval_seconds, max_interval_seconds=schedule.max_interval_seconds, - timeout_seconds=max(schedule.timeout_seconds - age_seconds, schedule.initial_interval_seconds), + timeout_seconds=schedule.timeout_seconds, + start_delay_seconds=_seconds_until(row.claim_expires_at), store=store, resumed=True, ) @@ -684,13 +859,15 @@ async def resume_unsettled_background_interactions( schedule: PollSchedule = DEFAULT_POLL_SCHEDULE, ) -> tuple["asyncio.Task[SettlementOutcome | None]", ...]: """ - Pick up every settlement no process has claimed, which is what a replica - that died mid-poll leaves behind. Each resumed poll keeps the remaining - share of the original timeout and gets at least one fetch, so a completed - interaction is still billed however late the resume comes. + Pick up every settlement no process has finished, which is what a replica + that died mid-poll or mid-settlement leaves behind. A row whose claim is + still within its lease waits until the lease runs out, then polls like any + other. Each resumed poll gets the full poll window however old its row is, + so a fetch that fails right after a late restart is retried rather than + ending the row as unsettled. """ return tuple( _track_poll(_resumed_context(row, store, schedule), fetch_interaction) - for row in await _unclaimed(store) + for row in await _unsettled(store) if row.interaction_id not in _ACTIVE_POLLS ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 58037689597..9eb2bc93391 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1524,7 +1524,6 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanState await asyncio.sleep(5) asyncio.create_task(_run_agent_grant_id_migration()) - await install_background_interaction_settlement(prisma_client) ## A coordination_redis block saved from the admin UI lives in the database, ## which is only reachable once the prisma client exists. Apply it here, before @@ -1651,6 +1650,10 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanState ## SYNC UI SETTINGS ## await ProxyStartupEvent.sync_ui_settings_to_general_settings() + await install_background_interaction_settlement( + prisma_client, router=lambda: llm_router, spend_logs_kept=lambda: disable_spend_logs is False + ) + # Start background health checks AFTER models are loaded and index is built if use_background_health_checks: asyncio.create_task(_run_background_health_check()) # start the background health check coroutine. diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index d40bac380aa..83a0e6f9e05 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -1981,6 +1981,7 @@ model LiteLLM_BackgroundInteractionSettlement { outcome String? @@index([claimed_at], map: "idx_background_interaction_settlement_claimed_at") + @@index([settled_at], map: "idx_background_interaction_settlement_settled_at") } model LiteLLM_Lens { diff --git a/litellm/proxy/spend_tracking/background_interaction_settlement.py b/litellm/proxy/spend_tracking/background_interaction_settlement.py index 165f813cd50..ffa8185619c 100644 --- a/litellm/proxy/spend_tracking/background_interaction_settlement.py +++ b/litellm/proxy/spend_tracking/background_interaction_settlement.py @@ -1,21 +1,25 @@ import asyncio import os import socket -from collections.abc import Awaitable, Mapping, Sequence -from dataclasses import dataclass -from datetime import datetime, timezone +from collections.abc import Awaitable, Callable, Mapping, Sequence +from dataclasses import dataclass, replace +from datetime import datetime, timedelta, timezone from itertools import chain from types import MappingProxyType -from typing import TYPE_CHECKING, Final, Protocol, TypeVar +from typing import TYPE_CHECKING, Final, Literal, Protocol, TypeVar from pydantic import ValidationError from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger -from litellm.constants import BACKGROUND_INTERACTION_COST_POLLING_ENABLED +from litellm.constants import ( + BACKGROUND_INTERACTION_COST_POLLING_ENABLED, + BACKGROUND_INTERACTION_SETTLEMENT_CLAIM_LEASE_SECONDS, +) from litellm.interactions.background_cost_polling import ( DEFAULT_POLL_SCHEDULE, BackgroundInteractionCreateContext, + BackgroundInteractionPollContext, FetchInteraction, PendingBackgroundInteraction, PollSchedule, @@ -24,10 +28,13 @@ from litellm.interactions.background_cost_polling import ( fetch_background_interaction, resume_unsettled_background_interactions, ) -from litellm.repositories.table_repositories import BackgroundInteractionSettlementRepository +from litellm.repositories.table_repositories import BackgroundInteractionSettlementRepository, SpendLogsRepository +from litellm.types.interactions import InteractionsAPIResponse +from litellm.types.router import CredentialLiteLLMParams if TYPE_CHECKING: from litellm.proxy.utils import PrismaClient + from litellm.router import Router class _SettlementRow(Protocol): @@ -41,6 +48,8 @@ class _SettlementRow(Protocol): def created_at(self) -> datetime: ... @property def claimed_at(self) -> datetime | None: ... + @property + def settled_at(self) -> datetime | None: ... class _NewSettlementRow(TypedDict): @@ -59,8 +68,24 @@ class _UnclaimedRowKey(TypedDict): claimed_at: ReadOnly[None] -class _UnclaimedRows(TypedDict): - claimed_at: ReadOnly[None] +class _Before(TypedDict): + lt: ReadOnly[datetime] + + +class _ExpiredClaimRowKey(TypedDict): + interaction_id: ReadOnly[str] + settled_at: ReadOnly[None] + claimed_at: ReadOnly[_Before] + + +class _OwnClaimRowKey(TypedDict): + interaction_id: ReadOnly[str] + claimed_by: ReadOnly[str] + settled_at: ReadOnly[None] + + +class _UnsettledRows(TypedDict): + settled_at: ReadOnly[None] class _Claim(TypedDict): @@ -79,9 +104,19 @@ class _SettlementTableActions(Protocol): def find_unique(self, *, where: _RowKey) -> Awaitable[_SettlementRow | None]: ... - def find_many(self, *, where: _UnclaimedRows) -> Awaitable[Sequence[_SettlementRow]]: ... + def find_many(self, *, where: _UnsettledRows) -> Awaitable[Sequence[_SettlementRow]]: ... - def update_many(self, *, data: _Claim | _Outcome, where: _RowKey | _UnclaimedRowKey) -> Awaitable[int]: ... + def update_many( + self, *, data: _Claim | _Outcome, where: _RowKey | _UnclaimedRowKey | _ExpiredClaimRowKey | _OwnClaimRowKey + ) -> Awaitable[int]: ... + + +class _SpendLogKey(TypedDict): + request_id: ReadOnly[str] + + +class _SpendLogLookup(Protocol): + def find_unique(self, *, where: _SpendLogKey) -> Awaitable[object | None]: ... _CLEARED_CREATE_CONTEXT: Final[Mapping[str, object]] = MappingProxyType({}) @@ -103,11 +138,7 @@ def _json(data: Mapping[str, object]) -> object: return Json.keys(**data) -def _pending_rows(rows: Sequence[_SettlementRow]) -> tuple[PendingBackgroundInteraction, ...]: - return tuple(chain.from_iterable(_pending_row(row) for row in rows)) - - -def _pending_row(row: _SettlementRow) -> tuple[PendingBackgroundInteraction, ...]: +def _pending_row(row: _SettlementRow, claim_expires_at: datetime | None) -> tuple[PendingBackgroundInteraction, ...]: try: create_context: Final = BackgroundInteractionCreateContext.model_validate(row.create_context) except ValidationError: @@ -122,14 +153,31 @@ def _pending_row(row: _SettlementRow) -> tuple[PendingBackgroundInteraction, ... custom_llm_provider=row.custom_llm_provider, create_context=create_context, created_at=row.created_at, + claim_expires_at=claim_expires_at, ), ) @dataclass(frozen=True, slots=True) class PrismaBackgroundSettlementStore: + """ + A claim is a lease, not a settlement: the settler that claims a row bills + or releases it and then records the outcome, so a row claimed longer ago + than the lease with no outcome belongs to a settler that died in between, + and any settler may claim it again. A live settler renews its claim while + it settles, so it is never overtaken. The settler before may have billed + and then died or failed to record it, so a takeover bills only when the + interaction has no spend row yet, and otherwise records it as billed. The + spend row is read before the claim is taken, so a failed read takes + nothing and the claim is retried. With spend logs off there is no row to + tell, so a takeover bills nothing rather than risk billing twice. + """ + table: _SettlementTableActions + spend_logs: _SpendLogLookup + spend_logs_kept: Callable[[], bool] claimed_by: str + claim_lease_seconds: float = BACKGROUND_INTERACTION_SETTLEMENT_CLAIM_LEASE_SECONDS async def register(self, pending: PendingBackgroundInteraction) -> None: await self.table.create( @@ -143,24 +191,71 @@ class PrismaBackgroundSettlementStore: async def pending(self, interaction_id: str) -> PendingBackgroundInteraction | None: row: Final = await self._row(interaction_id) - if row is None or row.claimed_at is not None: + if row is None or self._is_owned(row): return None - return next(iter(_pending_row(row)), None) + return next(iter(_pending_row(row, claim_expires_at=None)), None) async def is_claimed(self, interaction_id: str) -> bool: row: Final = await self._row(interaction_id) - return row is not None and row.claimed_at is not None + return row is not None and self._is_owned(row) async def claim(self, interaction_id: str) -> bool: + if await self._claim_where(_UnclaimedRowKey(interaction_id=interaction_id, claimed_at=None)): + return True + settled_before: Final = await self._outcome_a_dead_settler_left(interaction_id) + took_over: Final = await self._claim_where( + _ExpiredClaimRowKey( + interaction_id=interaction_id, + settled_at=None, + claimed_at=_Before(lt=datetime.now(timezone.utc) - self._lease), + ) + ) + if not took_over or settled_before is None: + return took_over + if settled_before == "unsettled": + verbose_proxy_logger.warning( + "Background interaction %s lost its settler mid-settle, and with spend logs off nothing shows " + "whether it was billed, so it is left unbilled rather than risk billing it twice", + interaction_id, + ) + await self.record_outcome(interaction_id, settled_before) + return False + + async def renew_claim(self, interaction_id: str) -> None: + await self.table.update_many( + data=_Claim(claimed_at=datetime.now(timezone.utc), claimed_by=self.claimed_by), + where=_OwnClaimRowKey(interaction_id=interaction_id, claimed_by=self.claimed_by, settled_at=None), + ) + + async def _claim_where(self, where: _UnclaimedRowKey | _ExpiredClaimRowKey) -> bool: claimed_rows: Final = await _read_from_a_table_that_may_not_exist( self.table.update_many( - data=_Claim(claimed_at=datetime.now(timezone.utc), claimed_by=self.claimed_by), - where=_UnclaimedRowKey(interaction_id=interaction_id, claimed_at=None), + data=_Claim(claimed_at=datetime.now(timezone.utc), claimed_by=self.claimed_by), where=where ), when_missing=0, ) return claimed_rows == 1 + async def _outcome_a_dead_settler_left(self, interaction_id: str) -> Literal["billed", "unsettled"] | None: + if not self.spend_logs_kept(): + return "unsettled" + if await self.spend_logs.find_unique(where=_SpendLogKey(request_id=interaction_id)) is None: + return None + return "billed" + + @property + def _lease(self) -> timedelta: + return timedelta(seconds=self.claim_lease_seconds) + + def _claim_expires_at(self, row: _SettlementRow) -> datetime | None: + return None if row.claimed_at is None else row.claimed_at + self._lease + + def _is_owned(self, row: _SettlementRow) -> bool: + if row.settled_at is not None: + return True + claim_expires_at: Final = self._claim_expires_at(row) + return claim_expires_at is not None and claim_expires_at > datetime.now(timezone.utc) + async def _row(self, interaction_id: str) -> _SettlementRow | None: return await _read_from_a_table_that_may_not_exist( self.table.find_unique(where=_RowKey(interaction_id=interaction_id)), when_missing=None @@ -174,19 +269,78 @@ class PrismaBackgroundSettlementStore: where=_RowKey(interaction_id=interaction_id), ) - async def unclaimed(self) -> Sequence[PendingBackgroundInteraction]: - return _pending_rows(await self.table.find_many(where=_UnclaimedRows(claimed_at=None))) + async def unsettled(self) -> Sequence[PendingBackgroundInteraction]: + rows: Final = await self.table.find_many(where=_UnsettledRows(settled_at=None)) + return tuple(chain.from_iterable(_pending_row(row, self._claim_expires_at(row)) for row in rows)) + + +def _provider_credentials(router: "Router", deployment_id: str, custom_llm_provider: str) -> Mapping[str, object]: + credentials: Final = router.get_deployment_credentials_with_provider(model_id=deployment_id) + if credentials is None or credentials.get("custom_llm_provider") != custom_llm_provider: + return MappingProxyType({}) + return MappingProxyType(credentials) + + +def _group_deployment_credentials(router: "Router", context: BackgroundInteractionPollContext) -> Mapping[str, object]: + if context.model_group is None: + return MappingProxyType({}) + candidates: Final = tuple( + credentials + for deployment_id in router.get_model_ids(model_name=context.model_group) + if (credentials := _provider_credentials(router, deployment_id, context.custom_llm_provider)) + ) + return candidates[0] if len(candidates) == 1 else MappingProxyType({}) + + +def _deployment_credentials( + router: "Router | None", context: BackgroundInteractionPollContext +) -> CredentialLiteLLMParams | None: + if router is None: + return None + serving: Final = ( + _provider_credentials(router, context.deployment_id, context.custom_llm_provider) + if context.deployment_id is not None and router.get_deployment(model_id=context.deployment_id) is not None + else _group_deployment_credentials(router, context) + ) + return CredentialLiteLLMParams.model_validate(serving) if serving else None + + +def fetch_with_deployment_credentials( + router: Callable[[], "Router | None"], fetch_interaction: FetchInteraction = fetch_background_interaction +) -> FetchInteraction: + """ + A resumed poll has no request to take credentials from, and they are never + stored, so it fetches with the credentials of the deployment that served + the create, looked up by its id on the live router at each fetch so a + deployment loaded from the database after startup is found too. Editing a + deployment changes its id, so a missing id falls back to the create's model + group when exactly one of its deployments serves the same provider. Anything + else falls back to the provider's environment credentials, never to a + wildcard route, whose key belongs to whatever provider it routes to. + """ + + async def fetch(context: BackgroundInteractionPollContext) -> InteractionsAPIResponse: + credentials: Final = _deployment_credentials(router(), context) + if credentials is None: + return await fetch_interaction(context) + return await fetch_interaction(replace(context, credentials=credentials)) + + return fetch async def configure_background_interaction_settlement( table: _SettlementTableActions, + spend_logs: _SpendLogLookup, + spend_logs_kept: Callable[[], bool], claimed_by: str, fetch_interaction: FetchInteraction = fetch_background_interaction, schedule: PollSchedule = DEFAULT_POLL_SCHEDULE, ) -> tuple["asyncio.Task[SettlementOutcome | None]", ...]: if not BACKGROUND_INTERACTION_COST_POLLING_ENABLED: return () - store: Final = PrismaBackgroundSettlementStore(table=table, claimed_by=claimed_by) + store: Final = PrismaBackgroundSettlementStore( + table=table, spend_logs=spend_logs, spend_logs_kept=spend_logs_kept, claimed_by=claimed_by + ) configure_background_settlement_store(store) resumed: Final = await resume_unsettled_background_interactions(store, fetch_interaction, schedule) if resumed: @@ -194,11 +348,16 @@ async def configure_background_interaction_settlement( return resumed -async def install_background_interaction_settlement(prisma_client: "PrismaClient") -> None: +async def install_background_interaction_settlement( + prisma_client: "PrismaClient", router: Callable[[], "Router | None"], spend_logs_kept: Callable[[], bool] +) -> None: try: await configure_background_interaction_settlement( table=BackgroundInteractionSettlementRepository(prisma_client).table, + spend_logs=SpendLogsRepository(prisma_client).table, + spend_logs_kept=spend_logs_kept, claimed_by=f"{socket.gethostname()}:{os.getpid()}", + fetch_interaction=fetch_with_deployment_credentials(router), ) except Exception as e: # noqa: BLE001 # a boot step must survive any DB error; billing then settles in-process as before verbose_proxy_logger.warning( diff --git a/schema.prisma b/schema.prisma index d40bac380aa..83a0e6f9e05 100644 --- a/schema.prisma +++ b/schema.prisma @@ -1981,6 +1981,7 @@ model LiteLLM_BackgroundInteractionSettlement { outcome String? @@index([claimed_at], map: "idx_background_interaction_settlement_claimed_at") + @@index([settled_at], map: "idx_background_interaction_settlement_settled_at") } model LiteLLM_Lens { diff --git a/tests/integration/spend/test_background_interaction_settlement.py b/tests/integration/spend/test_background_interaction_settlement.py index 16eb0a3434e..4408bd147f9 100644 --- a/tests/integration/spend/test_background_interaction_settlement.py +++ b/tests/integration/spend/test_background_interaction_settlement.py @@ -1,3 +1,4 @@ +import json import math import socket import time @@ -17,6 +18,8 @@ from integration._support.client import ( Gateway, eventually, gateway_from_environment, + held, + list_value, object_value, string_value, ) @@ -54,6 +57,9 @@ _CUSTOM_INPUT_RATE: Final = 2e-06 _CUSTOM_OUTPUT_RATE: Final = 4e-05 _ENV_KEY: Final = "integration-gemini-env-key" _DEPLOYMENT_KEY: Final = "integration-gemini-deployment-key" +_DATABASE_DEPLOYMENT_KEY: Final = "integration-gemini-database-deployment-key" +_ROTATED_DEPLOYMENT_KEY: Final = "integration-gemini-rotated-deployment-key" +_TAG: Final = "integration-settlement-team" _CREATOR_POLL: Final = {"BACKGROUND_INTERACTION_COST_POLL_INITIAL_INTERVAL_SECONDS": "300"} _SETTLER_POLL: Final = { "BACKGROUND_INTERACTION_COST_POLL_INITIAL_INTERVAL_SECONDS": "1", @@ -65,8 +71,10 @@ _RESUMER_POLL: Final = { "BACKGROUND_INTERACTION_COST_POLL_MAX_INTERVAL_SECONDS": "1", "BACKGROUND_INTERACTION_COST_POLL_TIMEOUT_SECONDS": "120", } +_EAGER_RESUMER_POLL: Final = {**_RESUMER_POLL, "BACKGROUND_INTERACTION_COST_POLL_INITIAL_INTERVAL_SECONDS": "0.01"} +_LATE_RESUMER_POLL: Final = {**_RESUMER_POLL, "BACKGROUND_INTERACTION_COST_POLL_TIMEOUT_SECONDS": "30"} _SPEND_QUERY: Final = ( - "SELECT request_id, spend, call_type, status, model, prompt_tokens, completion_tokens " + "SELECT request_id, spend, call_type, status, model, prompt_tokens, completion_tokens, request_tags " 'FROM "LiteLLM_SpendLogs" WHERE request_id = %s' ) _SETTLEMENT_QUERY: Final = ( @@ -74,6 +82,16 @@ _SETTLEMENT_QUERY: Final = ( 'settled_at IS NOT NULL AS settled, create_context FROM "LiteLLM_BackgroundInteractionSettlement" ' "WHERE interaction_id = %s" ) +_ABANDONED_CLAIM: Final = ( + 'UPDATE "LiteLLM_BackgroundInteractionSettlement" SET claimed_by = %s, ' + "claimed_at = (now() AT TIME ZONE 'UTC') - interval '10 minutes', settled_at = NULL, outcome = NULL, " + "create_context = %s::jsonb WHERE interaction_id = %s" +) +_DEAD_SETTLER: Final = "dead-settler:1" +_CREATED_AN_HOUR_AGO: Final = ( + 'UPDATE "LiteLLM_BackgroundInteractionSettlement" ' + "SET created_at = (now() AT TIME ZONE 'UTC') - interval '1 hour' WHERE interaction_id = %s" +) _SETTLEMENT_TABLE_PRESENT_QUERY: Final = "SELECT to_regclass(%s) IS NOT NULL AS present" _SETTLEMENT_TABLE: Final = '"LiteLLM_BackgroundInteractionSettlement"' _SETTLEMENT_BY_CALL_QUERY: Final = ( @@ -179,6 +197,40 @@ def _write_config(directory: Path, upstream_url: str, models: Deployments) -> Pa return path +def _variant_config(rig: Rig, name: str, document: Mapping[str, JsonValue]) -> Path: + path: Final = rig.directory / name + path.write_text(yaml.safe_dump(dict(document))) + return path + + +def _rig_config(rig: Rig) -> dict[str, JsonValue]: + return JSON_OBJECT.validate_python(yaml.safe_load(rig.config.read_text())) + + +def _config_without_spend_logs(rig: Rig) -> Path: + config: Final = _rig_config(rig) + general_settings: Final = {**object_value(config["general_settings"]), "disable_spend_logs": True} + return _variant_config( + rig, "settlement_config_without_spend_logs.yaml", {**config, "general_settings": general_settings} + ) + + +def _config_with_a_rotated_key(rig: Rig) -> Path: + config: Final = _rig_config(rig) + model_list: Final = [_with_rotated_key(entry, rig.models.in_progress) for entry in list_value(config["model_list"])] + return _variant_config(rig, "settlement_config_rotated_key.yaml", {**config, "model_list": model_list}) + + +def _with_rotated_key(entry: JsonValue, model_name: str) -> JsonValue: + deployment: Final = object_value(entry) + if deployment["model_name"] != model_name: + return deployment + return { + **deployment, + "litellm_params": {**object_value(deployment["litellm_params"]), "api_key": _ROTATED_DEPLOYMENT_KEY}, + } + + def _interaction(status: str, usage: dict[str, JsonValue] | None, http_status: int = 200) -> JsonResponse: return JsonResponse( content_type="application/json", @@ -206,12 +258,14 @@ def _create( path: str = "/v1beta/interactions", background: bool = True, text: str | None = None, + headers: Mapping[str, str] | None = None, ) -> str: response: Final = replica.request( "POST", path, {"model": model, "input": text or f"settle {uuid.uuid4().hex}", "background": background}, key=key, + headers=headers, ) assert response.status_code == 200, response.text return string_value(JSON_OBJECT.validate_json(response.content)["id"]) @@ -252,6 +306,16 @@ def _settlement(interaction_id: str) -> dict[str, JsonValue] | None: return rows[0] if rows else None +def _registered(interaction_id: str) -> dict[str, JsonValue]: + row: Final = eventually(lambda: _settlement(interaction_id), lambda value: value is not None) + assert row is not None + return row + + +def _abandon_claim(interaction_id: str, create_context: JsonValue) -> None: + write_rows(_ABANDONED_CLAIM, (_DEAD_SETTLER, json.dumps(create_context), interaction_id)) + + def _settlement_table_present() -> bool: return read_rows(_SETTLEMENT_TABLE_PRESENT_QUERY, (_SETTLEMENT_TABLE,))[0]["present"] is True @@ -315,6 +379,11 @@ def _assert_billed(row: Mapping[str, JsonValue], rates: tuple[float, float]) -> return spend +def _assert_billed_with_tag(row: Mapping[str, JsonValue], rates: tuple[float, float]) -> None: + _assert_billed(row, rates) + assert _TAG in list_value(row["request_tags"]), row + + def _key_spend(replica: Gateway, key: str) -> float: spend: Final = object_value(replica.get("/key/info", {"key": key})["info"])["spend"] assert isinstance(spend, float | int), spend @@ -333,12 +402,12 @@ def _drain(rig: Rig) -> list[JsonValue]: return requests -def _calls(rig: Rig, interaction_id: str) -> tuple[tuple[str, str], ...]: - suffix: Final = f"/v1beta/interactions/{interaction_id}" +def _calls(rig: Rig, *interaction_ids: str) -> tuple[tuple[str, str], ...]: + suffixes: Final = tuple(f"/v1beta/interactions/{interaction_id}" for interaction_id in interaction_ids) return tuple( (string_value(object_value(entry)["method"]), string_value(object_value(entry)["api_key"])) for entry in _drain(rig) - if string_value(object_value(entry)["path"]).endswith(suffix) + if string_value(object_value(entry)["path"]).endswith(suffixes) ) @@ -447,7 +516,7 @@ def test_a_replica_booting_later_resumes_and_bills_unclaimed_interactions(rig: R model: Final = rig.models.in_progress key: Final = scenario.key() created_after: Final = time.time() - created: Final = tuple(_create(rig.creator, model, key) for _ in range(3)) + created: Final = tuple(_create(rig.creator, model, key, headers={"x-litellm-tags": _TAG}) for _ in range(3)) for item in created: _state(rig, item, _completed()) rates: Final = _rates(rig.creator, model) @@ -457,13 +526,113 @@ def test_a_replica_booting_later_resumes_and_bills_unclaimed_interactions(rig: R pids: Final = _worker_pids(resumer.process.pid) assert len(pids) == 2, pids for item in created: - _assert_billed(_await_spend_row(item, seconds=90), rates) + _assert_billed_with_tag(_await_spend_row(item, seconds=90), rates) claimer: Final = _claimer_pid(_await_outcome(item, "billed")) assert claimer in pids or _booted_after(claimer, created_after), (claimer, pids) + assert {key for _, key in _calls(rig, *created)} == {_DEPLOYMENT_KEY} for item in created: assert len(_spend_rows(item)) == 1 +def test_a_replica_booting_later_resumes_a_database_model_with_its_own_key(rig: Rig) -> None: + with rig.creator.scenario() as scenario: + model: Final = scenario.model( + model=_MODEL, api_key=_DATABASE_DEPLOYMENT_KEY, api_base=f"{rig.upstream.url}/{rig.models.in_progress}" + ) + created: Final = _create(rig.creator, model, scenario.key()) + _state(rig, created, _completed()) + with owned_proxy_process( + rig.gateway, rig.directory, rig.environment(**_EAGER_RESUMER_POLL), config=rig.config, workers=2 + ): + _await_spend_row(created, seconds=90) + _await_outcome(created, "billed") + assert {key for _, key in _calls(rig, created)} == {_DATABASE_DEPLOYMENT_KEY} + + +def test_a_replica_booting_later_takes_over_lapsed_claims_and_bills_only_the_unbilled(rig: Rig) -> None: + with rig.creator.scenario() as scenario: + model: Final = rig.models.in_progress + key: Final = scenario.key() + unbilled: Final = _create(rig.creator, model, key) + billed: Final = _create(rig.settler, model, key) + contexts: Final = {item: _registered(item)["create_context"] for item in (unbilled, billed)} + for item in (unbilled, billed): + _state(rig, item, _completed()) + rates: Final = _rates(rig.creator, model) + spend: Final = _assert_billed(_await_spend_row(billed), rates) + _await_outcome(billed, "billed") + for item, create_context in contexts.items(): + _abandon_claim(item, create_context) + with owned_proxy_process( + rig.gateway, rig.directory, rig.environment(**_EAGER_RESUMER_POLL), config=rig.config, workers=2 + ): + _assert_billed(_await_spend_row(unbilled, seconds=90), rates) + assert _await_outcome(unbilled, "billed")["claimed_by"] != _DEAD_SETTLER + assert _await_outcome(billed, "billed", seconds=90)["claimed_by"] != _DEAD_SETTLER + held( + lambda: _key_spend(rig.creator, key), + lambda observed: math.isclose(observed, 2 * spend, rel_tol=1e-9), + holding=5, + seconds=30, + ) + for item in (unbilled, billed): + assert len(_spend_rows(item)) == 1 + + +def test_a_lapsed_claim_with_spend_logs_off_is_recorded_unsettled_and_never_billed(rig: Rig) -> None: + with rig.creator.scenario() as scenario: + model: Final = rig.models.in_progress + key: Final = scenario.key() + created: Final = _create(rig.creator, model, key) + _abandon_claim(created, _registered(created)["create_context"]) + _state(rig, created, _completed()) + with owned_proxy_process( + rig.gateway, + rig.directory, + rig.environment(**_EAGER_RESUMER_POLL), + config=_config_without_spend_logs(rig), + workers=2, + ): + assert _await_outcome(created, "unsettled", seconds=90)["claimed_by"] != _DEAD_SETTLER + assert _spend_rows(created) == [] + assert _key_spend(rig.creator, key) == 0 + + +def test_a_resumed_interaction_whose_deployment_key_rotated_fetches_with_the_new_key(rig: Rig) -> None: + with rig.creator.scenario() as scenario: + model: Final = rig.models.in_progress + created: Final = _create(rig.creator, model, scenario.key()) + _state(rig, created, _completed()) + with owned_proxy_process( + rig.gateway, + rig.directory, + rig.environment(**_EAGER_RESUMER_POLL), + config=_config_with_a_rotated_key(rig), + workers=2, + ) as resumer: + assert _model_info(resumer.gateway, model)["id"] != _model_info(rig.creator, model)["id"] + _assert_billed(_await_spend_row(created, seconds=90), _rates(rig.creator, model)) + _await_outcome(created, "billed") + assert {key for _, key in _calls(rig, created)} == {_ROTATED_DEPLOYMENT_KEY} + + +def test_a_late_restart_retries_a_failed_fetch_within_a_full_poll_window(rig: Rig) -> None: + with rig.creator.scenario() as scenario: + model: Final = rig.models.in_progress + created: Final = _create(rig.creator, model, scenario.key()) + _registered(created) + write_rows(_CREATED_AN_HOUR_AGO, (created,)) + _state(rig, created, InteractionState(status="completed", usage=_USAGE, get_status=503)) + _drain(rig) + with owned_proxy_process( + rig.gateway, rig.directory, rig.environment(**_LATE_RESUMER_POLL), config=rig.config, workers=2 + ): + eventually(lambda: _calls(rig, created), lambda calls: calls != (), seconds=90) + _state(rig, created, _completed()) + _assert_billed(_await_spend_row(created), _rates(rig.creator, model)) + _await_outcome(created, "billed") + + def test_deletes_on_the_creating_proxy_bill_each_interaction_once(rig: Rig) -> None: with rig.settler.scenario() as scenario: model: Final = rig.models.in_progress diff --git a/tests/unit/interactions/test_background_cost_polling.py b/tests/unit/interactions/test_background_cost_polling.py index 9dc71a4fa1b..d55994dfb19 100644 --- a/tests/unit/interactions/test_background_cost_polling.py +++ b/tests/unit/interactions/test_background_cost_polling.py @@ -1,11 +1,14 @@ import asyncio +import json import time -from datetime import datetime, timezone +from datetime import datetime, timedelta, timezone from itertools import islice from typing import Optional +import httpx import pytest +import litellm from litellm.interactions.background_cost_polling import ( _create_context, _poll_intervals, @@ -21,6 +24,7 @@ from litellm.interactions.background_cost_polling import ( ) from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kwargs from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging +from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload from litellm.types.interactions import InteractionsAPIResponse USAGE_BLOCK = { @@ -74,6 +78,7 @@ def _context( logging_obj: LitellmLogging, timeout_seconds: float = 1.0, store: Optional[InMemoryBackgroundSettlementStore] = None, + resumed: bool = False, ) -> BackgroundInteractionPollContext: return BackgroundInteractionPollContext( interaction_id="interactions/bg-abc", @@ -82,7 +87,9 @@ def _context( initial_interval_seconds=0.001, max_interval_seconds=0.002, timeout_seconds=timeout_seconds, + outcome_write_timeout_seconds=0.01, store=store if store is not None else InMemoryBackgroundSettlementStore(), + resumed=resumed, ) @@ -255,6 +262,180 @@ async def test_poller_retries_after_fetch_error_and_still_bills(): assert logging_obj.model_call_details["response_cost"] > 0 +def _provider_error_response(status_code: int) -> httpx.Response: + return httpx.Response(status_code, request=httpx.Request("GET", "https://generativelanguage.googleapis.com")) + + +@pytest.mark.parametrize( + "error", + [ + litellm.BadRequestError(message="invalid interaction id", model="gemini-2.5-flash", llm_provider="gemini"), + litellm.AuthenticationError(message="API key not valid", model="gemini-2.5-flash", llm_provider="gemini"), + litellm.PermissionDeniedError( + message="permission denied", + model="gemini-2.5-flash", + llm_provider="gemini", + response=_provider_error_response(403), + ), + litellm.NotFoundError(message="interaction not found", model="gemini-2.5-flash", llm_provider="gemini"), + ], +) +@pytest.mark.asyncio +async def test_a_resumed_poll_stops_fetching_at_the_first_error_no_retry_changes_and_releases_the_reservation(error): + """ + The regression: a resumed fetch the provider refuses for good, like an + unknown id or a key rotated away, was retried every interval until the + poll timed out. + """ + reservation = _reservation() + logging_obj = _logging_obj_with_reservation(reservation) + fetch, calls = _fetch_sequence(error) + + outcome = await poll_and_log_background_interaction_cost( + _context(logging_obj, timeout_seconds=0.05, resumed=True), fetch_interaction=fetch + ) + + assert outcome == "unsettled" + assert len(calls) == 1 + assert reservation["finalized"] is True + + +@pytest.mark.asyncio +async def test_a_resumed_poll_the_provider_refuses_leaves_the_bill_to_the_replica_still_polling_it(): + """ + The regression: a replica that booted with credentials unable to read the + interaction settled it as unsettled on its first 404, and the replica that + created it, still polling, then found it settled and never billed it. + """ + store = InMemoryBackgroundSettlementStore() + creator_logging_obj = _logging_obj() + resumed_fetched = asyncio.Event() + + async def creator_fetch(context): + if resumed_fetched.is_set(): + return _response("completed", with_usage=True) + return _response("in_progress", with_usage=False) + + async def resumed_fetch(context): + resumed_fetched.set() + raise litellm.NotFoundError(message="interaction not found", model="gemini-2.5-flash", llm_provider="gemini") + + creator_outcome, resumed_outcome = await asyncio.gather( + poll_and_log_background_interaction_cost( + _context(creator_logging_obj, store=store), fetch_interaction=creator_fetch + ), + poll_and_log_background_interaction_cost( + _context(_logging_obj(), store=store, resumed=True), fetch_interaction=resumed_fetch + ), + ) + + assert (creator_outcome, resumed_outcome) == ("billed", None) + assert creator_logging_obj.model_call_details["response_cost"] > 0 + + +@pytest.mark.parametrize( + "error", + [ + litellm.RateLimitError(message="quota exceeded", model="gemini-2.5-flash", llm_provider="gemini"), + litellm.ServiceUnavailableError(message="overloaded", model="gemini-2.5-flash", llm_provider="gemini"), + litellm.InternalServerError(message="internal error", model="gemini-2.5-flash", llm_provider="gemini"), + litellm.Timeout(message="timed out", model="gemini-2.5-flash", llm_provider="gemini"), + litellm.BadRequestError(message="invalid interaction id", model="gemini-2.5-flash", llm_provider="gemini"), + litellm.AuthenticationError(message="API key not valid", model="gemini-2.5-flash", llm_provider="gemini"), + litellm.NotFoundError(message="interaction not found", model="gemini-2.5-flash", llm_provider="gemini"), + ], +) +@pytest.mark.asyncio +async def test_the_creating_poll_retries_through_any_provider_error(error): + logging_obj = _logging_obj() + fetch, calls = _fetch_sequence(error, _response("completed", with_usage=True)) + + outcome = await poll_and_log_background_interaction_cost(_context(logging_obj), fetch_interaction=fetch) + + assert outcome == "billed" + assert len(calls) == 2 + + +class _OutcomeWritesFail: + def __init__(self): + self.store = InMemoryBackgroundSettlementStore() + + async def register(self, pending): + await self.store.register(pending) + + async def pending(self, interaction_id): + return await self.store.pending(interaction_id) + + async def is_claimed(self, interaction_id): + return await self.store.is_claimed(interaction_id) + + async def claim(self, interaction_id): + return await self.store.claim(interaction_id) + + async def record_outcome(self, interaction_id, outcome): + raise RuntimeError("database unavailable") + + async def unsettled(self): + return await self.store.unsettled() + + +@pytest.mark.parametrize( + "last_status, timeout_seconds", + [("failed", 1.0), ("in_progress", 0.01)], +) +@pytest.mark.asyncio +async def test_a_release_it_could_not_record_leaves_the_reservation_to_the_settler_that_takes_over( + last_status, timeout_seconds +): + """ + The regression: a settler released the reservation and then failed to + record it, so the replica that took over the lapsed claim released the + same reservation again and the key's budget lost that amount twice. + """ + reservation = _reservation() + logging_obj = _logging_obj_with_reservation(reservation) + fetch, _ = _fetch_sequence(_response(last_status, with_usage=False)) + + await poll_and_log_background_interaction_cost( + _context(logging_obj, timeout_seconds=timeout_seconds, store=_OutcomeWritesFail()), + fetch_interaction=fetch, + ) + + assert reservation["finalized"] is False + + +class _FirstOutcomeWriteFails(_OutcomeWritesFail): + def __init__(self): + super().__init__() + self.failed_once = False + + async def record_outcome(self, interaction_id, outcome): + if not self.failed_once: + self.failed_once = True + raise RuntimeError("connection reset") + await self.store.record_outcome(interaction_id, outcome) + + +@pytest.mark.asyncio +async def test_a_release_whose_first_outcome_write_failed_still_frees_the_reservation(): + """ + The regression: one failed outcome write left the key's budget + reservation held until a later boot or delete took the row over, and an + active key's spend counter never expired it in the meantime. + """ + reservation = _reservation() + store = _FirstOutcomeWriteFails() + fetch, _ = _fetch_sequence(_response("failed", with_usage=False)) + + outcome = await poll_and_log_background_interaction_cost( + _context(_logging_obj_with_reservation(reservation), store=store), fetch_interaction=fetch + ) + + assert outcome == "released" + assert reservation["finalized"] is True + assert await store.store.unsettled() == () + + @pytest.mark.asyncio async def test_schedule_creates_poll_task_for_in_progress_create(): logging_obj = _logging_obj() @@ -302,6 +483,59 @@ async def test_schedule_registers_an_agent_only_create_that_names_no_model(): assert _rebuild_logging_obj(pending.create_context).model is None +@pytest.mark.asyncio +async def test_schedule_registers_a_create_whose_tags_are_a_string_and_keeps_them_as_sent(): + """ + The regression: storing the create's context read string tags as a list + and raised, so a create sent with a tags string answered with an error + after the provider had already started the interaction. + """ + logging_obj = _logging_obj(litellm_params={"metadata": _create_metadata(tags="team-a,nightly-batch")}) + store = InMemoryBackgroundSettlementStore() + + task = await maybe_schedule_background_interaction_cost_polling( + response=_response("in_progress", with_usage=False), + create_kwargs={"litellm_logging_obj": logging_obj}, + custom_llm_provider="gemini", + store=store, + ) + + assert isinstance(task, asyncio.Task) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + pending = await store.pending("interactions/bg-abc") + assert pending is not None + rebuilt = _rebuild_logging_obj(pending.create_context) + assert get_litellm_metadata_from_kwargs(kwargs=rebuilt.model_call_details)["tags"] == "team-a,nightly-batch" + + +@pytest.mark.asyncio +async def test_delete_fetches_with_the_vertex_project_location_and_credentials_the_request_used(): + store = InMemoryBackgroundSettlementStore() + await _create_on_a_replica_that_then_dies(_logging_obj(litellm_params={"metadata": _create_metadata()}), store) + fetch, captured = _capturing_fetch(_response("completed", with_usage=True)) + + outcome = await maybe_settle_background_interaction_before_delete( + interaction_id="interactions/bg-abc", + delete_kwargs={ + "vertex_project": "request-project", + "vertex_location": "europe-west4", + "vertex_credentials": {"type": "service_account", "project_id": "request-project"}, + }, + fetch_interaction=fetch, + store=store, + ) + + assert outcome == "billed" + credentials = captured[0].credentials + assert (credentials.vertex_project, credentials.vertex_location, credentials.vertex_credentials) == ( + "request-project", + "europe-west4", + {"type": "service_account", "project_id": "request-project"}, + ) + + @pytest.mark.asyncio @pytest.mark.parametrize( "response,create_kwargs", @@ -781,6 +1015,105 @@ async def test_restart_resumes_only_the_rows_no_replica_claimed(): assert await store.is_claimed("interactions/bg-orphaned") +@pytest.mark.asyncio +async def test_a_restart_long_after_the_create_still_retries_a_failed_fetch_before_giving_up(): + """ + The regression: a resumed poll kept only what was left of the original + window, so a row older than the window got a single fetch, and one + transient error right after the restart ended it unbilled. + """ + store = InMemoryBackgroundSettlementStore() + await store.register( + PendingBackgroundInteraction( + interaction_id="interactions/bg-abc", + custom_llm_provider="gemini", + create_context=_create_context(_logging_obj(litellm_params={"metadata": _create_metadata()}), "gemini"), + created_at=datetime.now(timezone.utc) - timedelta(seconds=FAST_SCHEDULE.timeout_seconds * 100), + ) + ) + fetch, calls = _fetch_sequence( + RuntimeError("transient network error"), + _response("completed", with_usage=True), + ) + + (resumed,) = await resume_unsettled_background_interactions(store, fetch, schedule=FAST_SCHEDULE) + + assert await asyncio.wait_for(resumed, timeout=5) == "billed" + assert len(calls) == 2 + + +def _create_through_the_proxy_router(tags: list) -> LitellmLogging: + logging_obj = LitellmLogging( + model="gemini-2.5-flash", + messages=[], + stream=False, + call_type="acreate_interaction", + start_time=time.time(), + litellm_call_id="bg-interactions-call-id", + function_id="bg-interactions-fn-id", + ) + logging_obj.update_from_kwargs( + kwargs={ + "metadata": _create_metadata( + tags=tags, spend_logs_metadata={"job": "nightly"}, model_group="client-chosen-group" + ), + "litellm_metadata": {"model_info": {"id": "deployment-1"}, "model_group": "gemini-background"}, + }, + model="gemini-2.5-flash", + optional_params={}, + litellm_params={"litellm_call_id": "bg-interactions-call-id"}, + custom_llm_provider="gemini", + ) + return logging_obj + + +def _create_with_header_tags(tags: list) -> LitellmLogging: + return _logging_obj( + litellm_params={ + "metadata": _create_metadata(tags=tags, spend_logs_metadata={"job": "nightly"}), + "proxy_server_request": {"headers": {"user-agent": "claude-cli/2.1.0"}}, + } + ) + + +async def _billed_spend_row(logging_obj: LitellmLogging) -> dict: + response = _response("completed", with_usage=True) + fetch, _ = _fetch_sequence(response) + assert await poll_and_log_background_interaction_cost(_context(logging_obj), fetch_interaction=fetch) == "billed" + settled_at = datetime.now(timezone.utc) + row = get_logging_payload( + kwargs={**logging_obj.model_call_details, "completion_start_time": settled_at}, + response_obj=response, + start_time=settled_at, + end_time=settled_at, + ) + return { + "request_tags": json.loads(row["request_tags"]), + "model_id": row["model_id"], + "model_group": row["model_group"], + "spend_logs_metadata": json.loads(row["metadata"]).get("spend_logs_metadata"), + } + + +@pytest.mark.parametrize("create", [_create_through_the_proxy_router, _create_with_header_tags]) +@pytest.mark.asyncio +async def test_a_bill_settled_off_the_creating_worker_matches_the_creating_worker_spend_row(create): + """ + The regressions: the stored create context kept only the router's metadata + on the proxy path, so a row billed after a restart or by another replica + lost the request tags that tag-based spend reports group by; then it let + the client's request metadata override the router's, so a client-sent + model group relabeled the bill another replica wrote. + """ + creator_row = await _billed_spend_row(create(["team-a", "nightly-batch"])) + stored = _create_context(create(["team-a", "nightly-batch"]), "gemini") + + rebuilt_row = await _billed_spend_row(_rebuild_logging_obj(stored)) + + assert "team-a" in creator_row["request_tags"] + assert rebuilt_row == creator_row + + class _ClaimAnswersOnlyAfterTheLastFetch: def __init__(self): self.store = InMemoryBackgroundSettlementStore() @@ -809,8 +1142,8 @@ class _ClaimAnswersOnlyAfterTheLastFetch: async def record_outcome(self, interaction_id, outcome): return None - async def unclaimed(self): - return await self.store.unclaimed() + async def unsettled(self): + return await self.store.unsettled() @pytest.mark.asyncio @@ -844,7 +1177,7 @@ class _DownStore: async def record_outcome(self, interaction_id, outcome): raise RuntimeError("database unavailable") - async def unclaimed(self): + async def unsettled(self): raise RuntimeError("database unavailable") @@ -868,8 +1201,8 @@ class _RegistersThenRaises: async def record_outcome(self, interaction_id, outcome): return None - async def unclaimed(self): - return await self.store.unclaimed() + async def unsettled(self): + return await self.store.unsettled() @pytest.mark.asyncio @@ -898,7 +1231,7 @@ async def test_create_whose_registration_raised_after_landing_still_claims_the_s assert outcome == "billed" assert await store.is_claimed("interactions/bg-abc") - assert await store.unclaimed() == () + assert await store.unsettled() == () task.cancel() with pytest.raises(asyncio.CancelledError): await task @@ -907,7 +1240,7 @@ async def test_create_whose_registration_raised_after_landing_still_claims_the_s @pytest.mark.asyncio async def test_delete_on_a_worker_that_resumed_the_poll_fails_when_it_cannot_fetch(): """ - After a restart every worker resumes the unclaimed rows, so none of them + After a restart every worker resumes the unsettled rows, so none of them is the creator whose delete may release and delete on a failed fetch. A resumed worker's delete fails like any other replica's, and its own poll still bills the interaction once it completes. @@ -965,9 +1298,9 @@ class _LandsThenGoesDown: async def record_outcome(self, interaction_id, outcome): return None - async def unclaimed(self): + async def unsettled(self): self._answer() - return await self.store.unclaimed() + return await self.store.unsettled() def _answer(self): if self.down: @@ -992,7 +1325,7 @@ class _TableLessStore: async def record_outcome(self, interaction_id, outcome): raise RuntimeError("the settlement table does not exist") - async def unclaimed(self): + async def unsettled(self): raise RuntimeError("the settlement table does not exist") @@ -1027,7 +1360,7 @@ async def test_create_whose_registration_and_read_back_both_failed_bills_once_th assert (while_down, recovered) == (None, "billed") assert len(calls) == 2 assert await store.is_claimed("interactions/bg-abc") - assert await store.unclaimed() == () + assert await store.unsettled() == () assert await resume_unsettled_background_interactions(store, poll_fetch, schedule=FAST_SCHEDULE) == () task.cancel() with pytest.raises(asyncio.CancelledError): diff --git a/tests/unit/proxy/spend_tracking/test_background_interaction_settlement.py b/tests/unit/proxy/spend_tracking/test_background_interaction_settlement.py index 21384279bbb..eea7f8a81aa 100644 --- a/tests/unit/proxy/spend_tracking/test_background_interaction_settlement.py +++ b/tests/unit/proxy/spend_tracking/test_background_interaction_settlement.py @@ -1,22 +1,28 @@ import asyncio import time from dataclasses import dataclass -from datetime import datetime, timezone +from datetime import datetime, timedelta, timezone from typing import Optional import pytest import litellm.interactions.background_cost_polling as bg +from litellm import Router +from litellm.constants import BACKGROUND_INTERACTION_SETTLEMENT_CLAIM_LEASE_SECONDS from litellm.interactions.background_cost_polling import ( _create_context, + BackgroundInteractionPollContext, configure_background_settlement_store, maybe_settle_background_interaction_before_delete, PendingBackgroundInteraction, + poll_and_log_background_interaction_cost, PollSchedule, + resume_unsettled_background_interactions, ) from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging from litellm.proxy.spend_tracking.background_interaction_settlement import ( configure_background_interaction_settlement, + fetch_with_deployment_credentials, install_background_interaction_settlement, PrismaBackgroundSettlementStore, ) @@ -78,7 +84,42 @@ class _FakeSettlementTable: return len(matched) def _matching(self, where) -> list: - return [row for row in self.rows.values() if all(getattr(row, column) == value for column, value in where.items())] + return [row for row in self.rows.values() if _matches(row, where)] + + +def _matches(row: "_Row", where: dict) -> bool: + return all(_satisfies(row, column, condition) for column, condition in where.items()) + + +def _satisfies(row: "_Row", column: str, condition) -> bool: + value = getattr(row, column) + if isinstance(condition, dict): + return value is not None and value < condition["lt"] + return value == condition + + +class _FakeSpendLogs: + """The spend log table, holding only the request ids of the rows written so far.""" + + def __init__(self, request_ids: tuple[str, ...] = ()): + self.request_ids = request_ids + + async def find_unique(self, *, where): + return where["request_id"] if where["request_id"] in self.request_ids else None + + +def _store( + table, + claimed_by: str, + claim_lease_seconds: float = BACKGROUND_INTERACTION_SETTLEMENT_CLAIM_LEASE_SECONDS, +) -> PrismaBackgroundSettlementStore: + return PrismaBackgroundSettlementStore( + table=table, + spend_logs=_FakeSpendLogs(), + spend_logs_kept=lambda: True, + claimed_by=claimed_by, + claim_lease_seconds=claim_lease_seconds, + ) def _logging_obj(metadata: Optional[dict] = None) -> LitellmLogging: @@ -110,7 +151,12 @@ def _pending(interaction_id: str) -> PendingBackgroundInteraction: ) -def _stored_row(interaction_id: str, claimed: bool = False, create_context: Optional[object] = None) -> _Row: +def _stored_row( + interaction_id: str, + claimed_at: Optional[datetime] = None, + settled: bool = False, + create_context: Optional[object] = None, +) -> _Row: return _Row( interaction_id=interaction_id, custom_llm_provider="gemini", @@ -120,11 +166,21 @@ def _stored_row(interaction_id: str, claimed: bool = False, create_context: Opti else _create_context(_logging_obj(), "gemini").model_dump(mode="json") ), created_at=datetime.now(timezone.utc), - claimed_at=datetime.now(timezone.utc) if claimed else None, - claimed_by="replica-a:1" if claimed else None, + claimed_at=claimed_at, + claimed_by="replica-a:1" if claimed_at is not None else None, + settled_at=claimed_at if settled else None, + outcome="billed" if settled else None, ) +def _minutes_ago(minutes: float) -> datetime: + return datetime.now(timezone.utc) - timedelta(minutes=minutes) + + +def _past_the_default_lease() -> datetime: + return datetime.now(timezone.utc) - timedelta(seconds=BACKGROUND_INTERACTION_SETTLEMENT_CLAIM_LEASE_SECONDS + 60) + + def _completed(interaction_id: str) -> InteractionsAPIResponse: return InteractionsAPIResponse( id=interaction_id, model="gemini-2.5-flash", status="completed", steps=[], usage=dict(USAGE_BLOCK) @@ -144,20 +200,20 @@ def _capturing_fetch(): @pytest.mark.asyncio async def test_registered_row_reads_back_as_the_same_pending_interaction(): table = _FakeSettlementTable() - store = PrismaBackgroundSettlementStore(table=table, claimed_by="replica-a:1") + store = _store(table, claimed_by="replica-a:1") pending = _pending("interactions/bg-1") await store.register(pending) assert await store.pending("interactions/bg-1") == pending - assert await store.unclaimed() == (pending,) + assert await store.unsettled() == (pending,) @pytest.mark.asyncio async def test_claim_is_won_by_exactly_one_settler(): table = _FakeSettlementTable() - replica_a = PrismaBackgroundSettlementStore(table=table, claimed_by="replica-a:1") - replica_b = PrismaBackgroundSettlementStore(table=table, claimed_by="replica-b:1") + replica_a = _store(table, claimed_by="replica-a:1") + replica_b = _store(table, claimed_by="replica-b:1") await replica_a.register(_pending("interactions/bg-1")) assert await replica_b.claim("interactions/bg-1") is True @@ -200,7 +256,7 @@ class _MissingSettlementTable: async def test_a_missing_table_holds_no_rows_and_takes_no_registration(): from prisma.errors import TableNotFoundError - store = PrismaBackgroundSettlementStore(table=_MissingSettlementTable(), claimed_by="replica-a:1") + store = _store(_MissingSettlementTable(), claimed_by="replica-a:1") with pytest.raises(TableNotFoundError): await store.register(_pending("interactions/bg-1")) @@ -208,29 +264,62 @@ async def test_a_missing_table_holds_no_rows_and_takes_no_registration(): assert await store.is_claimed("interactions/bg-1") is False assert await store.claim("interactions/bg-1") is False with pytest.raises(TableNotFoundError): - await store.unclaimed() + await store.unsettled() @pytest.mark.asyncio -async def test_unclaimed_skips_claimed_and_unreadable_rows(): +async def test_unsettled_lists_claimed_rows_with_no_outcome_and_when_their_lease_ends(): + claimed_at = _minutes_ago(1) table = _FakeSettlementTable( rows=( _stored_row("interactions/bg-orphaned"), - _stored_row("interactions/bg-settled", claimed=True), + _stored_row("interactions/bg-mid-settlement", claimed_at=claimed_at), + _stored_row("interactions/bg-settled", claimed_at=claimed_at, settled=True), _stored_row("interactions/bg-from-the-future", create_context={"schema": "unknown"}), ) ) - store = PrismaBackgroundSettlementStore(table=table, claimed_by="replica-b:1") + store = _store(table, claimed_by="replica-b:1", claim_lease_seconds=300) - unclaimed = await store.unclaimed() + unsettled = await store.unsettled() - assert [row.interaction_id for row in unclaimed] == ["interactions/bg-orphaned"] + assert [(row.interaction_id, row.claim_expires_at) for row in unsettled] == [ + ("interactions/bg-orphaned", None), + ("interactions/bg-mid-settlement", claimed_at + timedelta(seconds=300)), + ] + + +@pytest.mark.asyncio +async def test_a_claim_whose_settler_died_is_taken_over_once_its_lease_runs_out(): + """ + The regression: a settler that died between its claim and its outcome left + the row claimed forever, so no replica, restart, or delete ever billed it. + """ + table = _FakeSettlementTable( + rows=( + _stored_row("interactions/bg-dead-settler", claimed_at=_minutes_ago(10)), + _stored_row("interactions/bg-settled", claimed_at=_minutes_ago(10), settled=True), + ) + ) + replica_b = _store(table, claimed_by="replica-b:1", claim_lease_seconds=300) + replica_c = _store(table, claimed_by="replica-c:1", claim_lease_seconds=300) + + assert await replica_b.is_claimed("interactions/bg-dead-settler") is False + assert await replica_b.pending("interactions/bg-dead-settler") is not None + assert await replica_b.claim("interactions/bg-dead-settler") is True + assert await replica_c.claim("interactions/bg-dead-settler") is False + assert await replica_c.is_claimed("interactions/bg-dead-settler") is True + assert table.rows["interactions/bg-dead-settler"].claimed_by == "replica-b:1" + + assert await replica_b.is_claimed("interactions/bg-settled") is True + assert await replica_b.pending("interactions/bg-settled") is None + assert await replica_b.claim("interactions/bg-settled") is False + assert table.rows["interactions/bg-settled"].claimed_by == "replica-a:1" @pytest.mark.asyncio async def test_record_outcome_keeps_the_audit_trail_and_drops_the_stored_request_context(): table = _FakeSettlementTable() - store = PrismaBackgroundSettlementStore(table=table, claimed_by="replica-a:1") + store = _store(table, claimed_by="replica-a:1") await store.register(_pending("interactions/bg-1")) assert await store.claim("interactions/bg-1") assert table.rows["interactions/bg-1"].create_context @@ -244,23 +333,198 @@ async def test_record_outcome_keeps_the_audit_trail_and_drops_the_stored_request assert row.create_context == {} +@pytest.mark.asyncio +async def test_a_takeover_records_the_bill_a_settler_wrote_before_dying_instead_of_billing_again(): + """ + The regression: a settler that billed and then died, or failed to record + the outcome, left its claim to lapse with no outcome, and the replica that + took the claim over billed the interaction a second time. + """ + table = _FakeSettlementTable( + rows=( + _stored_row("interactions/bg-billed-then-died", claimed_at=_past_the_default_lease()), + _stored_row("interactions/bg-died-before-billing", claimed_at=_past_the_default_lease()), + ) + ) + fetch, _ = _capturing_fetch() + previous_store = bg._STORE.store + try: + resumed = await configure_background_interaction_settlement( + table=table, + spend_logs=_FakeSpendLogs(request_ids=("interactions/bg-billed-then-died",)), + spend_logs_kept=lambda: True, + claimed_by="replica-b:1", + fetch_interaction=fetch, + schedule=FAST_SCHEDULE, + ) + + assert [await asyncio.wait_for(task, timeout=5) for task in resumed] == [None, "billed"] + assert [(row.claimed_by, row.outcome) for row in table.rows.values()] == [ + ("replica-b:1", "billed"), + ("replica-b:1", "billed"), + ] + assert await bg._STORE.store.unsettled() == () + finally: + configure_background_settlement_store(previous_store) + + +@pytest.mark.asyncio +async def test_with_spend_logs_off_a_takeover_leaves_the_interaction_unbilled_instead_of_billing_it_twice(): + """ + The regression: with spend logs off a settler that billed and then died + leaves no spend row to find, so the replica that took over its lapsed + claim billed the interaction a second time. + """ + table = _FakeSettlementTable( + rows=( + _stored_row("interactions/bg-claimed-then-died", claimed_at=_past_the_default_lease()), + _stored_row("interactions/bg-orphaned"), + ) + ) + fetch, _ = _capturing_fetch() + previous_store = bg._STORE.store + try: + resumed = await configure_background_interaction_settlement( + table=table, + spend_logs=_FakeSpendLogs(), + spend_logs_kept=lambda: False, + claimed_by="replica-b:1", + fetch_interaction=fetch, + schedule=FAST_SCHEDULE, + ) + + assert [await asyncio.wait_for(task, timeout=5) for task in resumed] == [None, "billed"] + assert [(row.claimed_by, row.outcome) for row in table.rows.values()] == [ + ("replica-b:1", "unsettled"), + ("replica-b:1", "billed"), + ] + finally: + configure_background_settlement_store(previous_store) + + +class _SpendLogsWhoseFirstReadFails(_FakeSpendLogs): + def __init__(self): + super().__init__() + self.reads = 0 + + async def find_unique(self, *, where): + self.reads += 1 + if self.reads == 1: + raise RuntimeError("read replica unavailable") + return await super().find_unique(where=where) + + +@pytest.mark.asyncio +async def test_a_takeover_whose_spend_log_read_failed_retries_and_bills_instead_of_stranding_the_row(): + """ + The regression: a takeover wrote its claim before reading the spend log, + so a failed read left the row under a fresh claim of its own, which the + poll then took for another settler's and left unbilled until a reboot. + """ + table = _FakeSettlementTable( + rows=(_stored_row("interactions/bg-dead-settler", claimed_at=_past_the_default_lease()),) + ) + fetch, _ = _capturing_fetch() + previous_store = bg._STORE.store + try: + (resumed,) = await configure_background_interaction_settlement( + table=table, + spend_logs=_SpendLogsWhoseFirstReadFails(), + spend_logs_kept=lambda: True, + claimed_by="replica-b:1", + fetch_interaction=fetch, + schedule=FAST_SCHEDULE, + ) + + assert await asyncio.wait_for(resumed, timeout=5) == "billed" + assert table.rows["interactions/bg-dead-settler"].outcome == "billed" + finally: + configure_background_settlement_store(previous_store) + + +def _logging_obj_whose_bill_takes(seconds: float, bills: list) -> LitellmLogging: + logging_obj = _logging_obj() + + async def bill(result): + bills.append(result.id) + await asyncio.sleep(seconds) + + logging_obj.async_log_background_interaction_completion = bill + return logging_obj + + +@pytest.mark.asyncio +async def test_a_settle_that_outlasts_the_lease_is_not_taken_over_and_billed_twice(): + """ + The regression: a success callback that held the settle past the claim + lease let a booting replica take the live claim over and bill the same + interaction a second time. + """ + lease_seconds = 0.6 + table = _FakeSettlementTable() + creator = _store(table, claimed_by="replica-a:1", claim_lease_seconds=lease_seconds) + resumer = _store(table, claimed_by="replica-b:1", claim_lease_seconds=lease_seconds) + await creator.register(_pending("interactions/bg-slow-settle")) + fetch, _ = _capturing_fetch() + bills = [] + + def poll(store, resumed: bool, start_delay_seconds: float): + return poll_and_log_background_interaction_cost( + BackgroundInteractionPollContext( + interaction_id="interactions/bg-slow-settle", + custom_llm_provider="gemini", + logging_obj=_logging_obj_whose_bill_takes(lease_seconds * 2.5, bills), + initial_interval_seconds=0.001, + max_interval_seconds=0.002, + timeout_seconds=lease_seconds * 3, + start_delay_seconds=start_delay_seconds, + claim_renewal_interval_seconds=lease_seconds / 10, + store=store, + resumed=resumed, + ), + fetch, + ) + + outcomes = await asyncio.gather( + poll(creator, resumed=False, start_delay_seconds=0.0), + poll(resumer, resumed=True, start_delay_seconds=lease_seconds * 1.5), + ) + + assert outcomes == ["billed", None] + assert bills == ["interactions/bg-slow-settle"] + assert table.rows["interactions/bg-slow-settle"].claimed_by == "replica-a:1" + + @pytest.mark.asyncio async def test_configure_installs_the_store_and_resumes_the_orphaned_rows(): table = _FakeSettlementTable( - rows=(_stored_row("interactions/bg-orphaned"), _stored_row("interactions/bg-settled", claimed=True)) + rows=( + _stored_row("interactions/bg-orphaned"), + _stored_row("interactions/bg-dead-settler", claimed_at=_past_the_default_lease()), + _stored_row("interactions/bg-settled", claimed_at=_minutes_ago(1), settled=True), + ) ) fetch, captured = _capturing_fetch() previous_store = bg._STORE.store try: resumed = await configure_background_interaction_settlement( - table=table, claimed_by="replica-b:1", fetch_interaction=fetch, schedule=FAST_SCHEDULE + table=table, + spend_logs=_FakeSpendLogs(), + spend_logs_kept=lambda: True, + claimed_by="replica-b:1", + fetch_interaction=fetch, + schedule=FAST_SCHEDULE, ) - assert len(resumed) == 1 - assert await asyncio.wait_for(resumed[0], timeout=5) == "billed" - assert [context.interaction_id for context in captured] == ["interactions/bg-orphaned"] - assert table.rows["interactions/bg-orphaned"].claimed_by == "replica-b:1" - assert table.rows["interactions/bg-orphaned"].outcome == "billed" + assert [await asyncio.wait_for(task, timeout=5) for task in resumed] == ["billed", "billed"] + assert sorted(context.interaction_id for context in captured) == [ + "interactions/bg-dead-settler", + "interactions/bg-orphaned", + ] + for interaction_id in ("interactions/bg-orphaned", "interactions/bg-dead-settler"): + assert table.rows[interaction_id].claimed_by == "replica-b:1" + assert table.rows[interaction_id].outcome == "billed" + assert table.rows["interactions/bg-settled"].claimed_by == "replica-a:1" await table.create( data={ @@ -288,7 +552,9 @@ class _PrismaClientWithoutSettlementTable: async def test_install_keeps_booting_when_the_settlement_table_is_unreachable(): previous_store = bg._STORE.store - await install_background_interaction_settlement(_PrismaClientWithoutSettlementTable()) + await install_background_interaction_settlement( + _PrismaClientWithoutSettlementTable(), router=lambda: None, spend_logs_kept=lambda: True + ) assert bg._STORE.store is previous_store @@ -296,3 +562,166 @@ async def test_install_keeps_booting_when_the_settlement_table_is_unreachable(): @dataclass(frozen=True) class _JsonLike: data: object + + +@pytest.mark.asyncio +async def test_a_resumed_row_still_under_a_live_claim_waits_for_the_lease_before_settling(): + claimed_at = datetime.now(timezone.utc) + table = _FakeSettlementTable(rows=(_stored_row("interactions/bg-mid-settlement", claimed_at=claimed_at),)) + store = _store(table, claimed_by="replica-b:1", claim_lease_seconds=0.2) + fetched_at = [] + + async def fetch(context): + fetched_at.append(datetime.now(timezone.utc)) + return _completed(context.interaction_id) + + (resumed,) = await resume_unsettled_background_interactions(store, fetch, schedule=FAST_SCHEDULE) + + assert await asyncio.wait_for(resumed, timeout=5) == "billed" + assert fetched_at[0] >= claimed_at + timedelta(seconds=0.2) + assert table.rows["interactions/bg-mid-settlement"].claimed_by == "replica-b:1" + + +def _router_with_a_deployment_scoped_key() -> Router: + return Router( + model_list=[ + { + "model_name": "gemini-background", + "litellm_params": { + "model": "gemini/gemini-2.5-flash", + "api_key": "deployment-scoped-key", + "api_base": "https://deployment-scoped.example", + }, + "model_info": {"id": "deployment-1"}, + } + ] + ) + + +def _router_with_a_catch_all_route() -> Router: + return Router(model_list=[{"model_name": "*", "litellm_params": {"model": "gemini/*", "api_key": "catch-all-key"}}]) + + +def _router_with_two_deployments_in_the_group() -> Router: + return Router( + model_list=[ + { + "model_name": "gemini-background", + "litellm_params": {"model": "gemini/gemini-2.5-flash", "api_key": f"deployment-key-{index}"}, + "model_info": {"id": f"deployment-{index}"}, + } + for index in (2, 3) + ] + ) + + +def _router_with_a_vertex_deployment() -> Router: + return Router( + model_list=[ + { + "model_name": "vertex-background", + "litellm_params": { + "model": "vertex_ai/gemini-2.5-flash", + "vertex_project": "deployment-project", + "vertex_location": "europe-west4", + "vertex_credentials": "/secrets/deployment-service-account.json", + }, + "model_info": {"id": "vertex-deployment-1"}, + } + ] + ) + + +def _create_context_served_by(deployment_id: str) -> dict: + metadata = { + "user_api_key": "0123456789abcdef" * 4, + "model_info": {"id": deployment_id}, + "model_group": "gemini-background", + } + return _create_context(_logging_obj(metadata), "gemini").model_dump(mode="json") + + +@pytest.mark.parametrize("deployment_id", ["deployment-1", "deployment-1-before-an-edit"]) +@pytest.mark.asyncio +async def test_a_resumed_poll_fetches_with_the_credentials_of_the_deployment_that_served_the_create(deployment_id): + """ + The regression: a resumed poll has no request to take credentials from, so + it fetched with the environment's key, which a deployment-scoped key is + not, and every fetch failed until the row was given up unbilled. Editing + the deployment since gives it a new id, and its model group still names it. + """ + table = _FakeSettlementTable( + rows=(_stored_row("interactions/bg-1", create_context=_create_context_served_by(deployment_id)),) + ) + router = _router_with_a_deployment_scoped_key() + fetch, captured = _capturing_fetch() + previous_store = bg._STORE.store + try: + (resumed,) = await configure_background_interaction_settlement( + table=table, + spend_logs=_FakeSpendLogs(), + spend_logs_kept=lambda: True, + claimed_by="replica-b:1", + fetch_interaction=fetch_with_deployment_credentials(lambda: router, fetch), + schedule=FAST_SCHEDULE, + ) + + assert await asyncio.wait_for(resumed, timeout=5) == "billed" + assert (captured[0].credentials.api_key, captured[0].credentials.api_base) == ( + "deployment-scoped-key", + "https://deployment-scoped.example", + ) + finally: + configure_background_settlement_store(previous_store) + + +@pytest.mark.parametrize( + "router, deployment_id, model_group", + [ + (_router_with_a_deployment_scoped_key, "deployment-removed-since", None), + (_router_with_a_deployment_scoped_key, None, None), + (lambda: None, "deployment-1", "gemini-background"), + (_router_with_a_catch_all_route, "deployment-removed-since", None), + (_router_with_two_deployments_in_the_group, "deployment-removed-since", "gemini-background"), + (_router_with_a_vertex_deployment, "deployment-removed-since", "vertex-background"), + (_router_with_a_vertex_deployment, "vertex-deployment-1", None), + ], +) +@pytest.mark.asyncio +async def test_a_resumed_poll_without_a_known_deployment_fetches_with_the_environment_credentials( + router, deployment_id, model_group +): + live_router = router() + fetch, captured = _capturing_fetch() + context = BackgroundInteractionPollContext( + interaction_id="interactions/bg-1", + custom_llm_provider="gemini", + logging_obj=_logging_obj(), + deployment_id=deployment_id, + model_group=model_group, + ) + + await fetch_with_deployment_credentials(lambda: live_router, fetch)(context) + + assert captured == [context] + + +@pytest.mark.asyncio +async def test_a_resumed_vertex_poll_fetches_with_the_project_location_and_credentials_of_its_deployment(): + router = _router_with_a_vertex_deployment() + fetch, captured = _capturing_fetch() + context = BackgroundInteractionPollContext( + interaction_id="interactions/bg-1", + custom_llm_provider="vertex_ai", + logging_obj=_logging_obj(), + deployment_id="vertex-deployment-1", + ) + + await fetch_with_deployment_credentials(lambda: router, fetch)(context) + + credentials = captured[0].credentials + assert (credentials.vertex_project, credentials.vertex_location, credentials.vertex_credentials) == ( + "deployment-project", + "europe-west4", + "/secrets/deployment-service-account.json", + ) From d5ce7dbb8cb8c387493ee5afba79c16f91ad9beb Mon Sep 17 00:00:00 2001 From: Deepanshu Pal <40927968+deepanshupal@users.noreply.github.com> Date: Fri, 16 Oct 2026 18:45:05 +0530 Subject: [PATCH 17/23] fix(proxy)!: honor caller litellm_session_id in spend log session_id (#44199) * fix(responses): scope session history lookup to the referenced key * Implement SQLite-backed DB for session tests Added a SQLite-backed database class for testing and a new test to ensure session isolation between different API keys. * fix(proxy): honor caller litellm_session_id in spend log session_id Updated docstring to clarify session ID resolution logic for spend log rows. * test(proxy): caller litellm_session_id wins over the trace id * style(proxy): fix blank lines around new tests * fix(proxy): read the caller session from the logging payload for spend logs * fix(responses): match user and team too when replaying session history * fix(proxy): keep bridged Responses chains whole and scope session history by owner * fix(proxy): keep a metadata session next to a trace header and let session headers win * fix(proxy): send the trace-header session only to spend logs and let session headers win beside a trace header * fix(proxy): drop a client-sent caller session and keep it out of requester metadata * fix(proxy): keep the trace id as the spend log session when the caller session contains a NUL byte * fix(responses): keep a continued bridged chain's session under request_correlation_in_logs With the flag on, the caller session outranked the chain session the bridge sets for a continued Responses chain, so a chain whose session changed or first appeared on a later turn split. The continuation now logs under the chain's session and drops the turn's caller session key, as on main. Adds the /audit integration cells for the spend log caller session. * fix(responses): log a continued bridged chain under its session through the spend-log key only The previous commit set the chain session as the turn's litellm_session_id, which also moved the OTel conversation id of a continued turn away from its trace header. The bridge now writes the chain session into the internal spend-log session key of the metadata the spend log reads, so only the spend log session changes, as on main. Adds an OTel cell for a continued chain. * test(spend): cover bridged chains that share one caller session * test(spend): cover the caller session on configured pass-through routes * test(spend): type the SQLite fake and split stacked comprehensions in the caller session tests --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/constants.py | 1 + litellm/proxy/litellm_pre_call_utils.py | 52 +- .../spend_tracking/spend_tracking_utils.py | 35 +- .../session_handler.py | 25 +- .../transformation.py | 26 +- .../spend/test_spend_log_caller_session.py | 783 ++++++++++++++++++ .../test_spend_tracking_utils.py | 80 +- .../unit/proxy/test_litellm_pre_call_utils.py | 88 ++ .../test_litellm_completion_responses.py | 67 ++ .../test_session_handler.py | 78 ++ 10 files changed, 1208 insertions(+), 27 deletions(-) create mode 100644 tests/integration/spend/test_spend_log_caller_session.py diff --git a/litellm/constants.py b/litellm/constants.py index 271e42ee661..eb03d034c7e 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1629,6 +1629,7 @@ INTERNAL_CALL_ORIGIN_METADATA_KEY: Final = "internal_call_origin" MANAGED_BACKGROUND_RESPONSE_METADATA_KEY: Final = "_managed_background_response" SESSION_ID_GENERATED_METADATA_KEY: Final = "litellm_session_id_generated" SESSION_ID_OMITTED_METADATA_KEY: Final = "litellm_session_id_omitted" +CALLER_SESSION_ID_METADATA_KEY: Final = "litellm_caller_session_id" LITELLM_TRUNCATED_PAYLOAD_FIELD: Final = "litellm_truncated" LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE: Final = ( "Truncation is a DB storage safeguard. " diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 6bd31a42e9e..675bf0830f9 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -27,6 +27,7 @@ from litellm._logging import verbose_logger, verbose_proxy_logger from litellm._service_logger import ServiceLogging from litellm._uuid import uuid from litellm.constants import ( + CALLER_SESSION_ID_METADATA_KEY, CLIENT_OUTPUT_CEILING_METADATA_KEY, CONSUMED_REQUEST_TAGS_METADATA_KEY, INTERNAL_CALL_ORIGIN_METADATA_KEY, @@ -187,6 +188,30 @@ def _caller_trace_field(data: Mapping[str, object], metadata_variable_name: str, return requester_value if isinstance(requester_value, str) and requester_value else None +def _caller_body_session_id(data: Mapping[str, object], metadata_variable_name: str) -> str | None: + top_level: Final = data.get("litellm_session_id") + if isinstance(top_level, str) and top_level: + return top_level + return _caller_trace_field(data, metadata_variable_name, "session_id") + + +def _caller_session_beside_trace_header( + normalized_headers: Mapping[str, str], data: Mapping[str, object], metadata_variable_name: str +) -> str | None: + session_headers: Final = {key: value for key, value in normalized_headers.items() if key != "x-litellm-trace-id"} + return get_chain_id_from_headers(session_headers) or _caller_body_session_id(data, metadata_variable_name) + + +def _drop_client_caller_session(request_body: Mapping[str, object]) -> None: + for metadata_key in ("metadata", "litellm_metadata"): + if isinstance(metadata := request_body.get(metadata_key), dict): + metadata.pop(CALLER_SESSION_ID_METADATA_KEY, None) + + +def _requester_metadata_snapshot(metadata: Mapping[str, object]) -> Mapping[str, object]: + return copy.deepcopy({key: value for key, value in metadata.items() if key != CALLER_SESSION_ID_METADATA_KEY}) + + def _stampable_key_hash(user_api_key_dict: UserAPIKeyAuth) -> str | None: """Only proxy-validated keys are stamped, proven by the unforgeable via_virtual_key marker AND a known non-secret shape: the sha256 hex digest @@ -1623,6 +1648,8 @@ class LiteLLMProxyRequestSetup: from litellm.proxy._types import LitellmMetadataFromRequestHeaders metadata_from_headers: Final = LitellmMetadataFromRequestHeaders() + request_body: Final = cast(Mapping[str, object], data) # cast-ok: request body is a str-keyed JSON object + _drop_client_caller_session(request_body) spend_logs_metadata: Final = LiteLLMProxyRequestSetup._get_spend_logs_metadata_from_request_headers(headers) if spend_logs_metadata is not None: metadata_from_headers["spend_logs_metadata"] = spend_logs_metadata @@ -1631,6 +1658,9 @@ class LiteLLMProxyRequestSetup: # Finally update the requests metadata with the `metadata_from_headers` ######################################################################################### + normalized_headers: Final = MappingProxyType( + {k.lower(): v for k, v in headers.items() if isinstance(k, str) and isinstance(v, str)} + ) agent_id_from_header: Final = headers.get("x-litellm-agent-id") # Explicit litellm headers take precedence; fall back to any x-*-session-id header. chain_id: Final = get_chain_id_from_headers(dict(headers)) @@ -1639,6 +1669,13 @@ class LiteLLMProxyRequestSetup: metadata_from_headers["agent_id"] = agent_id_from_header verbose_proxy_logger.debug("Extracted agent_id from header: %s", agent_id_from_header) + spend_log_session_id: Final = ( + _caller_session_beside_trace_header(normalized_headers, request_body, _metadata_variable_name) + if litellm.request_correlation_in_logs + and chain_id is not None + and chain_id == normalized_headers.get("x-litellm-trace-id") + else None + ) if chain_id: metadata_from_headers["trace_id"] = chain_id metadata_from_headers["session_id"] = chain_id @@ -1663,7 +1700,6 @@ class LiteLLMProxyRequestSetup: # - but lets a caller's existing traceparent/baggage headers (from # real OTel instrumentation) correlate with litellm's own logs instead # of generating an unrelated trace_id. - normalized_headers: Final = MappingProxyType({k.lower(): v for k, v in headers.items() if isinstance(k, str)}) if "litellm_trace_id" not in data: traceparent: Final = normalized_headers.get("traceparent") if isinstance(traceparent, str): @@ -1689,11 +1725,7 @@ class LiteLLMProxyRequestSetup: baggage: Final = normalized_headers.get("baggage") if isinstance(baggage, str): session_id_from_baggage: Final = _session_id_from_baggage(baggage) - caller_session_id: Final = _caller_trace_field( - cast(Mapping[str, object], data), # cast-ok: request body is a str-keyed JSON object - _metadata_variable_name, - "session_id", - ) + caller_session_id: Final = _caller_trace_field(request_body, _metadata_variable_name, "session_id") if session_id_from_baggage and ( caller_session_id is None or caller_session_id == session_id_from_baggage ): @@ -1702,7 +1734,11 @@ class LiteLLMProxyRequestSetup: verbose_proxy_logger.debug("Extracted session_id from W3C baggage header") if isinstance(data[_metadata_variable_name], dict): - data[_metadata_variable_name].update(metadata_from_headers) + data[_metadata_variable_name].update( + {**metadata_from_headers, CALLER_SESSION_ID_METADATA_KEY: spend_log_session_id} + if spend_log_session_id + else metadata_from_headers + ) return data @staticmethod @@ -2417,7 +2453,7 @@ async def add_litellm_data_to_request( # them — from leaking into requester_metadata where guardrails and audit # paths may read from it. if "metadata" in data and isinstance(data["metadata"], dict): - data[_metadata_variable_name]["requester_metadata"] = copy.deepcopy(data["metadata"]) + data[_metadata_variable_name]["requester_metadata"] = _requester_metadata_snapshot(data["metadata"]) if _metadata_variable_name == "litellm_metadata": data[_metadata_variable_name].update( _promoted_trace_control_fields( diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 49644f1aa12..fe55092591e 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -15,6 +15,7 @@ from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError import litellm from litellm._logging import verbose_proxy_logger from litellm.constants import ( + CALLER_SESSION_ID_METADATA_KEY, CLI_SESSION_KEY_PREFIX, EMPTY_MAPPING, INTERNAL_CALL_ORIGIN_METADATA_KEY, @@ -33,6 +34,8 @@ from litellm.constants import ( from litellm.constants import ( MAX_STRING_LENGTH_PROMPT_IN_DB as DEFAULT_MAX_STRING_LENGTH_PROMPT_IN_DB, ) +from litellm.integrations.otel.model.metadata import caller_session_id +from litellm.integrations.otel.model.trace_controls import TraceControls from litellm.litellm_core_utils.classifier_logging import ( CLASSIFIER_AUDIT_FIELDS, classifier_audit_fields, @@ -930,6 +933,16 @@ def _get_batch_trace_session_id(call_type: str | None, request_id: str | None) - return request_id.removesuffix(BATCH_COST_REQUEST_ID_SUFFIX) +def _caller_session_for_spend_log(kwargs: Mapping[str, object], metadata: Mapping[str, object] | None) -> str | None: + beside_trace_header: Final = metadata.get(CALLER_SESSION_ID_METADATA_KEY) if metadata else None + caller_session: Final = ( + beside_trace_header + if isinstance(beside_trace_header, str) and beside_trace_header + else caller_session_id(kwargs, TraceControls()) + ) + return caller_session if caller_session and "\x00" not in caller_session else None + + def _get_session_id_for_spend_log( kwargs: Mapping[str, object], metadata: Mapping[str, object] | None, @@ -937,9 +950,12 @@ def _get_session_id_for_spend_log( omit_when_missing: bool, batch_trace_session_id: str | None = None, ) -> str | None: - """Under `omit` only `metadata.session_id`, the key Langfuse reads, counts as a session; `litellm_session_id` may - be a copied trace id. Batch call types carry a deterministic session derived from the batch id, which outranks - the per-request trace ids because those differ between the create call and the cost poller's row.""" + """Resolve the session id for the spend log row: `omit` honors only metadata.session_id, + batch sessions outrank everything, then the conversation id the caller sent (also one an x-litellm-trace-id header + displaced from the request's session), then trace ids. + + A session minted by `missing_session_id: generate` or copied from the request's trace id is not a conversation, so + it yields to the trace id, which the Responses bridge points at the previous turn's session to keep a chain whole.""" if omit_when_missing: session_id: Final = metadata.get("session_id") if metadata else None return str(session_id) if session_id else None @@ -948,11 +964,14 @@ def _get_session_id_for_spend_log( if batch_trace_session_id is not None: return batch_trace_session_id - if standard_logging_payload is not None and standard_logging_payload.get("trace_id") is not None: - return str(standard_logging_payload.get("trace_id")) - if kwargs.get("litellm_trace_id") is not None: - return str(kwargs.get("litellm_trace_id")) - return str(uuid.uuid4()) + if litellm.request_correlation_in_logs and (caller_session := _caller_session_for_spend_log(kwargs, metadata)): + return caller_session + match standard_logging_payload: + case {"trace_id": trace_id} if trace_id is not None: + return str(trace_id) + case _: + litellm_trace_id: Final = kwargs.get("litellm_trace_id") + return str(litellm_trace_id) if litellm_trace_id is not None else str(uuid.uuid4()) def _get_request_duration_ms(start_time: datetime, end_time: datetime) -> int | None: diff --git a/litellm/responses/litellm_completion_transformation/session_handler.py b/litellm/responses/litellm_completion_transformation/session_handler.py index 732c32ff2cb..e78ec5e5d3e 100644 --- a/litellm/responses/litellm_completion_transformation/session_handler.py +++ b/litellm/responses/litellm_completion_transformation/session_handler.py @@ -281,7 +281,13 @@ class ResponsesSessionHandler: SQL query - SELECT session_id FROM spend_logs WHERE response_id = previous_response_id, SELECT * FROM spend_logs WHERE session_id = session_id + SELECT session_id, api_key, "user", team_id FROM spend_logs WHERE request_id = previous_response_id, then + SELECT * FROM spend_logs with that session_id and the same owner + + Only rows with the referenced response's owner are returned: the same user and team, and for a key with + neither, the same key. The session id can be caller supplied, so matching on it alone would mix owners. A + user's or team's regenerated or second key keeps the history; a key with neither does not, and turns written + by another owner (an admin continuing the chain) are not replayed. A just-finished turn gets a short second chance: the worker that served it may still be writing its spend log when the follow-up arrives, and an empty result @@ -303,13 +309,24 @@ class ResponsesSessionHandler: query: Final = """ WITH matching_session AS ( - SELECT session_id + SELECT session_id, api_key, "user", team_id FROM "LiteLLM_SpendLogs" WHERE request_id = $1 ) SELECT * - FROM "LiteLLM_SpendLogs" - WHERE session_id IN (SELECT session_id FROM matching_session) + FROM "LiteLLM_SpendLogs" AS logs + WHERE EXISTS ( + SELECT 1 + FROM matching_session + WHERE matching_session.session_id = logs.session_id + AND matching_session."user" IS NOT DISTINCT FROM logs."user" + AND matching_session.team_id IS NOT DISTINCT FROM logs.team_id + AND ( + COALESCE(matching_session."user", '') <> '' + OR COALESCE(matching_session.team_id, '') <> '' + OR matching_session.api_key = logs.api_key + ) + ) ORDER BY "endTime" ASC; """ diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 2e7fb6da194..587cdfcee25 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -39,7 +39,11 @@ from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger from litellm.caching import InMemoryCache -from litellm.constants import REDACTED_BY_LITELLM, REDACTED_TOOL_CALL_ARGUMENTS_PLACEHOLDER +from litellm.constants import ( + CALLER_SESSION_ID_METADATA_KEY, + REDACTED_BY_LITELLM, + REDACTED_TOOL_CALL_ARGUMENTS_PLACEHOLDER, +) from litellm.integrations.anthropic_cache_control_hook import MAX_CACHE_CONTROL_BLOCKS, AnthropicCacheControlHook from litellm.litellm_core_utils.core_helpers import normalize_drop_params from litellm.litellm_core_utils.get_supported_openai_params import ( @@ -563,6 +567,14 @@ class ChatCompletionSession(TypedDict, total=False): ########### End of Initialize Classes used for Responses API ########### +def _chain_session_metadata(chain_session: str, request_kwargs: Mapping[str, object]) -> Mapping[str, object]: + key: Final = "litellm_metadata" if request_kwargs.get("litellm_metadata") else "metadata" + metadata: Final = request_kwargs.get(key) + if not isinstance(metadata, dict): + return MappingProxyType({}) + return {key: {**_STR_KEY_DICT_ADAPTER.validate_python(metadata), CALLER_SESSION_ID_METADATA_KEY: chain_session}} + + class LiteLLMCompletionResponsesConfig: @staticmethod def get_supported_openai_params(model: str) -> list: @@ -1078,8 +1090,6 @@ class LiteLLMCompletionResponsesConfig: else: # Both are empty - this likely means function_call_output had empty/invalid call_id # Provide a helpful error message - import litellm - raise litellm.BadRequestError( message=( f"Unable to create messages for completion request. " @@ -1094,8 +1104,14 @@ class LiteLLMCompletionResponsesConfig: ) litellm_completion_request["messages"] = combined_messages - litellm_completion_request["litellm_trace_id"] = chat_completion_session.get("litellm_session_id") - return litellm_completion_request + chain_session: Final = chat_completion_session.get("litellm_session_id") + litellm_completion_request["litellm_trace_id"] = chain_session + chain_session_metadata: Final[Mapping[str, object]] = ( + _chain_session_metadata(chain_session, request_kwargs) + if chain_session and litellm.request_correlation_in_logs + else MappingProxyType({}) + ) + return {**litellm_completion_request, **chain_session_metadata} @staticmethod def _transform_response_input_param_to_chat_completion_message( diff --git a/tests/integration/spend/test_spend_log_caller_session.py b/tests/integration/spend/test_spend_log_caller_session.py new file mode 100644 index 00000000000..213583d21a1 --- /dev/null +++ b/tests/integration/spend/test_spend_log_caller_session.py @@ -0,0 +1,783 @@ +import itertools +import json +import re +import threading +import uuid +from urllib.parse import urlsplit +from collections import deque +from collections.abc import Iterator, Mapping, Sequence +from concurrent.futures import ThreadPoolExecutor +from contextlib import ExitStack +from dataclasses import dataclass, replace +from datetime import datetime, timedelta, timezone +from hashlib import sha256 +from pathlib import Path +from queue import SimpleQueue +from types import MappingProxyType +from typing import Final + +import anthropic +import httpx +import openai +import psutil +import pytest +import yaml +from integration._support.client import ( + JSON_OBJECT, + Gateway, + Scenario, + eventually, + gateway_from_environment, + list_value, + object_value, + string_value, +) +from integration._support.budget import readiness_seconds +from integration._support.database import read_rows +from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue + +pytestmark: Final = pytest.mark.timeout(2 * readiness_seconds(2) + 2 * graceful_stop_seconds() + 240) + +MARKER: Final = re.compile(rb"callersess-[0-9a-f]{32}") +CALLER_SESSION_KEY: Final = "litellm_caller_session_id" +PATHS: Final = MappingProxyType( + {"chat": "/v1/chat/completions", "messages": "/v1/messages", "responses": "/v1/responses"} +) +ENDPOINT_STREAMS: Final = tuple(itertools.product(PATHS, (False, True))) +Headers = Mapping[str, str] | Sequence[tuple[str, str]] + + +def _marker() -> str: + return "callersess-" + uuid.uuid4().hex + + +def _token(label: str) -> str: + return f"{label}-{uuid.uuid4().hex}" + + +def _frame(payload: Mapping[str, JsonValue]) -> bytes: + return b"data: " + json.dumps(payload).encode() + b"\n\n" + + +def _chat_reply(text: str, stream: bool) -> Reply: + identity: Final = "chatcmpl-" + uuid.uuid4().hex + usage: Final[dict[str, JsonValue]] = {"prompt_tokens": 7, "completion_tokens": 2, "total_tokens": 9} + if not stream: + return Reply( + body=json.dumps( + { + "id": identity, + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": text}, "finish_reason": "stop"} + ], + "usage": usage, + } + ).encode() + ) + chunk: Final[dict[str, JsonValue]] = { + "id": identity, + "object": "chat.completion.chunk", + "created": 1, + "model": "gpt-4o-mini", + } + return Reply( + content_type="text/event-stream", + chunks=( + _frame({**chunk, "choices": [{"index": 0, "delta": {"role": "assistant", "content": text}}]}), + _frame({**chunk, "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], "usage": usage}), + b"data: [DONE]\n\n", + ), + ) + + +def _upstream(request: Request) -> Reply: + markers: Final = tuple(found.group(0).decode() for found in MARKER.finditer(request.body)) + if request.method != "POST" or not request.target.endswith("/chat/completions") or not markers: + return Reply(status=404, body=b'{"error":"unexpected request"}') + return _chat_reply(f"answer {markers[-1]}", JSON_OBJECT.validate_json(request.body).get("stream") is True) + + +def _user_markers(body: Mapping[str, JsonValue]) -> tuple[str, ...]: + user_contents: Final = tuple( + json.dumps(object_value(message)["content"]) + for message in list_value(body.get("messages", [])) + if object_value(message)["role"] == "user" + ) + return tuple(found.group(0).decode() for content in user_contents if (found := MARKER.search(content.encode()))) + + +def _spans(batch: Request) -> Iterator[dict[str, JsonValue]]: + for resource in list_value(JSON_OBJECT.validate_json(batch.body)["resourceSpans"]): + for scope in list_value(object_value(resource)["scopeSpans"]): + yield from (object_value(span) for span in list_value(object_value(scope)["spans"])) + + +def _attributes(span: Mapping[str, JsonValue]) -> dict[str, JsonValue]: + return { + string_value(object_value(attribute)["key"]): object_value(attribute)["value"] + for attribute in list_value(span.get("attributes", [])) + } + + +def _batch_attributes(batches: Sequence[Request]) -> Iterator[dict[str, JsonValue]]: + return (_attributes(span) for span in itertools.chain.from_iterable(_spans(batch) for batch in batches)) + + +@dataclass(frozen=True, slots=True) +class Collector: + wire: Wire + accepted: deque[Request] + guard: threading.Lock + + def span_attributes(self, response_id: str) -> tuple[dict[str, JsonValue], ...]: + with self.guard: + batches: Final = tuple(self.accepted) + attributes: Final = tuple(_batch_attributes(batches)) + return tuple( + found + for found in attributes + if object_value(found.get("gen_ai.response.id", {})).get("stringValue") == response_id + ) + + def conversation_ids(self, conversations: frozenset[str]) -> tuple[str, ...]: + with self.guard: + batches: Final = tuple(self.accepted) + responded: Final = ( + object_value(found.get("gen_ai.conversation.id", {})).get("stringValue") + for found in _batch_attributes(batches) + if "gen_ai.response.id" in found + ) + return tuple( + conversation + for conversation in responded + if isinstance(conversation, str) and conversation in conversations + ) + + +@dataclass(frozen=True, slots=True) +class Rig: + proxy: Gateway + scenario: Scenario + model: str + upstream: Wire + process: OwnedProxy + + def body(self, endpoint: str, marker: str, extra: Mapping[str, JsonValue]) -> dict[str, JsonValue]: + match endpoint: + case "chat": + return {"model": self.model, "messages": [{"role": "user", "content": marker}], **extra} + case "messages": + return { + "model": self.model, + "max_tokens": 16, + "messages": [{"role": "user", "content": marker}], + **extra, + } + case _: + return {"model": self.model, "input": marker, **extra} + + def post( + self, + endpoint: str, + marker: str, + fields: Mapping[str, JsonValue] = MappingProxyType({}), + *, + stream: bool = False, + headers: Headers = (), + key: str | None = None, + ) -> httpx.Response: + payload: Final = self.body(endpoint, marker, fields) | ({"stream": True} if stream else {}) + request_headers: Final = ( + ("Authorization", f"Bearer {self.proxy.key if key is None else key}"), + *(headers.items() if isinstance(headers, Mapping) else headers), + ) + with self.proxy.client.stream("POST", PATHS[endpoint], json=payload, headers=request_headers) as response: + response.read() + return response + + def send( + self, + endpoint: str, + marker: str, + fields: Mapping[str, JsonValue] = MappingProxyType({}), + *, + stream: bool = False, + headers: Headers = (), + key: str | None = None, + ) -> httpx.Response: + response: Final = self.post(endpoint, marker, fields, stream=stream, headers=headers, key=key) + assert response.status_code == 200, response.text + return response + + def upstream_bodies(self, marker: str) -> tuple[dict[str, JsonValue], ...]: + bodies: Final = tuple(JSON_OBJECT.validate_json(request.body) for request in self.upstream.drain()) + return tuple(body for body in bodies if _user_markers(body)[-1:] == (marker,)) + + +@dataclass(frozen=True, slots=True) +class Rigs: + flag_off: Rig + flag_on: Rig + sink: Collector + + +def _spend_rows(call_id: str) -> list[dict[str, JsonValue]]: + return read_rows( + 'SELECT session_id, request_id, metadata::text AS metadata_text FROM "LiteLLM_SpendLogs" ' + "WHERE litellm_call_id=%s", + (call_id,), + ) + + +def _spend_row(response: httpx.Response) -> dict[str, JsonValue]: + call_id: Final = response.headers["x-litellm-call-id"] + rows: Final = eventually(lambda: _spend_rows(call_id), lambda values: len(values) >= 1, seconds=70) + assert len(rows) == 1, rows + return rows[0] + + +def _session(response: httpx.Response) -> JsonValue: + return _spend_row(response)["session_id"] + + +def _model(scenario: Scenario, upstream: Wire) -> str: + return scenario.model( + model="hosted_vllm/gpt-4o-mini", + api_base=f"{upstream.url}/v1", + use_chat_completions_api=True, + input_cost_per_token=0.001, + output_cost_per_token=0.002, + ) + + +def _rig_config( + directory: Path, name: str, litellm_settings: Mapping[str, JsonValue], extra: Mapping[str, JsonValue] +) -> Path: + shared: Final = object_value(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + config: Final = { + **shared, + "general_settings": {**object_value(shared["general_settings"]), "store_prompts_in_spend_logs": True}, + "litellm_settings": {**object_value(shared["litellm_settings"]), **litellm_settings}, + **extra, + } + path: Final = directory / f"{name}.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +@pytest.fixture(scope="module") +def rigs(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rigs]: + accepted: Final[deque[Request]] = deque() # mutable-ok: sink thread appends each accepted OTLP batch + guard: Final = threading.Lock() + + def sink(request: Request) -> Reply: + with guard: + accepted.append(request) + return Reply() + + directory: Final = tmp_path_factory.mktemp("caller-session") + with ExitStack() as stack: + upstream: Final = stack.enter_context(wire_server(_upstream)) + collector: Final = Collector(stack.enter_context(wire_server(sink)), accepted, guard) + shared: Final = stack.enter_context(gateway_from_environment()) + flag_off: Final = stack.enter_context( + owned_proxy_process(shared, directory, {}, config=_rig_config(directory, "flag-off", {}, {}), workers=2) + ) + flag_on: Final = stack.enter_context( + owned_proxy_process( + shared, + directory, + {"LITELLM_OTEL_V2": "1", "OTEL_BSP_SCHEDULE_DELAY": "300"}, + config=_rig_config( + directory, + "flag-on", + {"request_correlation_in_logs": True, "callbacks": ["otel"]}, + { + "callback_settings": { + "otel": {"exporter": "http/json", "endpoint": collector.wire.url, "mapper_names": ["genai"]} + } + }, + ), + workers=2, + ) + ) + off_scenario: Final = stack.enter_context(flag_off.gateway.scenario()) + on_scenario: Final = stack.enter_context(flag_on.gateway.scenario()) + yield Rigs( + Rig(flag_off.gateway, off_scenario, _model(off_scenario, upstream), upstream, flag_off), + Rig(flag_on.gateway, on_scenario, _model(on_scenario, upstream), upstream, flag_on), + collector, + ) + + +@pytest.mark.parametrize(("endpoint", "stream"), ENDPOINT_STREAMS) +def test_spend_log_session_is_the_body_session_beside_a_trace_header(rigs: Rigs, endpoint: str, stream: bool) -> None: + marker: Final = _marker() + session: Final = _token("convo") + response: Final = rigs.flag_on.send( + endpoint, + marker, + stream=stream, + headers={"x-litellm-trace-id": _token("trace")}, + fields={"litellm_session_id": session}, + ) + + row: Final = _spend_row(response) + upstream: Final = rigs.flag_on.upstream_bodies(marker) + + assert row["session_id"] == session, row + assert CALLER_SESSION_KEY not in string_value(row["metadata_text"]), row + assert len(upstream) == 1, upstream + assert CALLER_SESSION_KEY not in json.dumps(upstream[0]), upstream + + +def test_sdk_clients_record_the_body_session_beside_a_trace_header(rigs: Rigs) -> None: + rig: Final = rigs.flag_on + base_url: Final = str(rig.proxy.client.base_url) + traces: Final = tuple(_token("trace") for _ in range(3)) + sessions: Final = tuple(_token("convo") for _ in range(3)) + markers: Final = tuple(_marker() for _ in range(3)) + chat_client: Final = openai.OpenAI(base_url=base_url + "/v1", api_key=rig.proxy.key, max_retries=0) + messages_client: Final = anthropic.Anthropic(base_url=base_url, api_key=rig.proxy.key, max_retries=0) + + with chat_client.chat.completions.with_streaming_response.create( + model=rig.model, + messages=[{"role": "user", "content": markers[0]}], + stream=True, + extra_headers={"x-litellm-trace-id": traces[0]}, + extra_body={"litellm_session_id": sessions[0]}, + ) as chat_stream: + chat_call: Final = chat_stream.headers["x-litellm-call-id"] + chat_text: Final = "".join( + chunk.choices[0].delta.content or "" for chunk in chat_stream.parse() if chunk.choices + ) + message: Final = messages_client.messages.with_raw_response.create( + model=rig.model, + max_tokens=16, + messages=[{"role": "user", "content": markers[1]}], + extra_headers={"x-litellm-trace-id": traces[1]}, + extra_body={"litellm_session_id": sessions[1]}, + ) + with chat_client.responses.with_streaming_response.create( + model=rig.model, + input=markers[2], + stream=True, + extra_headers={"x-litellm-trace-id": traces[2]}, + extra_body={"litellm_session_id": sessions[2]}, + ) as responses_stream: + responses_call: Final = responses_stream.headers["x-litellm-call-id"] + completed: Final = tuple(event for event in responses_stream.parse() if event.type == "response.completed") + + first_block: Final = message.parse().content[0] + assert chat_text == f"answer {markers[0]}", chat_text + assert isinstance(first_block, anthropic.types.TextBlock) and first_block.text == f"answer {markers[1]}", message + assert len(completed) == 1, completed + calls: Final = (chat_call, message.headers["x-litellm-call-id"], responses_call) + landed: Final = tuple( + eventually(lambda call=call: _spend_rows(call), lambda values: len(values) == 1, seconds=70)[0]["session_id"] + for call in calls + ) + assert landed == sessions, (landed, calls) + + +def test_conversation_span_keeps_the_trace_header_and_no_internal_key(rigs: Rigs) -> None: + trace: Final = _token("trace") + response: Final = rigs.flag_on.send( + "chat", _marker(), headers={"x-litellm-trace-id": trace}, fields={"litellm_session_id": _token("convo")} + ) + response_id: Final = string_value(JSON_OBJECT.validate_json(response.content)["id"]) + + spans: Final = eventually(lambda: rigs.sink.span_attributes(response_id), lambda found: len(found) == 1, seconds=30) + + assert spans[0]["gen_ai.conversation.id"] == {"stringValue": trace}, spans + assert CALLER_SESSION_KEY not in json.dumps(spans[0]), spans + + +def test_a_continued_chains_conversation_spans_keep_each_turns_trace_header(rigs: Rigs) -> None: + traces: Final = (_token("trace"), _token("trace")) + first: Final = _turn(rigs.flag_on, _marker(), headers={"x-litellm-trace-id": traces[0]}) + _turn( + rigs.flag_on, + _marker(), + first, + headers={"x-litellm-trace-id": traces[1]}, + fields={"litellm_session_id": _token("convo")}, + ) + + conversations: Final = eventually( + lambda: rigs.sink.conversation_ids(frozenset(traces)), lambda found: len(found) == 2, seconds=30 + ) + + assert sorted(conversations) == sorted(traces), conversations + + +@pytest.mark.parametrize("endpoint", list(PATHS)) +def test_flag_off_keeps_the_trace_id_beside_a_body_session(rigs: Rigs, endpoint: str) -> None: + trace: Final = _token("trace") + + response: Final = rigs.flag_off.send( + endpoint, _marker(), headers={"x-litellm-trace-id": trace}, fields={"litellm_session_id": _token("convo")} + ) + + assert _session(response) == trace + + +@pytest.mark.parametrize("endpoint", list(PATHS)) +def test_response_cache_hit_records_the_body_session_beside_a_trace_header(rigs: Rigs, endpoint: str) -> None: + marker: Final = _marker() + sessions: Final = (_token("convo"), _token("convo")) + responses: Final = tuple( + rigs.flag_on.send( + endpoint, marker, headers={"x-litellm-trace-id": _token("trace")}, fields={"litellm_session_id": session} + ) + for session in sessions + ) + + assert tuple(_session(response) for response in responses) == sessions + assert len(rigs.flag_on.upstream_bodies(marker)) == 1 + + +def test_metadata_session_beside_a_new_trace_per_turn_groups_both_turns(rigs: Rigs) -> None: + session: Final = _token("convo") + responses: Final = tuple( + rigs.flag_on.send( + "chat", + _marker(), + headers={"x-litellm-trace-id": _token("trace")}, + fields={"metadata": {"session_id": session}}, + ) + for _ in range(2) + ) + + assert tuple(_session(response) for response in responses) == (session, session) + now: Final = datetime.now(timezone.utc) + page: Final = rigs.flag_on.proxy.get( + "/spend/logs/ui", + params={ + "session_id": session, + "start_date": (now - timedelta(days=1)).strftime("%Y-%m-%d %H:%M:%S"), + "end_date": (now + timedelta(days=1)).strftime("%Y-%m-%d %H:%M:%S"), + }, + ) + listed: Final = sorted(string_value(object_value(row)["request_id"]) for row in list_value(page["data"])) + assert listed == sorted( + string_value(JSON_OBJECT.validate_json(response.content)["id"]) for response in responses + ), page + + +def _shape(shape: str, trace: str, session: str, header: str) -> tuple[Headers, dict[str, JsonValue], str]: + traced: Final = {"x-litellm-trace-id": trace} + match shape: + case "session-header-wins": + return {"x-litellm-session-id": header}, {"litellm_session_id": session}, header + case "trace-header-only": + return traced, {}, trace + case "body-session-only": + return {}, {"litellm_session_id": session}, session + case "int-session": + return traced, {"litellm_session_id": 7}, trace + case "list-session": + return traced, {"litellm_session_id": [session]}, trace + case "empty-session": + return traced, {"litellm_session_id": ""}, trace + case "null-session": + return traced, {"litellm_session_id": None}, trace + case "nul-session": + return traced, {"litellm_session_id": session + "\x00x"}, trace + case "two-trace-headers": + return ( + (("x-litellm-trace-id", trace), ("x-litellm-trace-id", header)), + {"litellm_session_id": session}, + session, + ) + case _: + return traced, {"litellm_session_id": session + "s" * 5120}, session + "s" * 5120 + + +@pytest.mark.parametrize( + "shape", + [ + "session-header-wins", + "trace-header-only", + "body-session-only", + "int-session", + "list-session", + "empty-session", + "null-session", + "nul-session", + "two-trace-headers", + "5kb-session", + ], +) +def test_spend_log_session_for_each_session_shape(rigs: Rigs, shape: str) -> None: + headers, fields, expected = _shape(shape, _token("trace"), _token("convo"), _token("header")) + neighbor_trace: Final = _token("neighbor") + + response: Final = rigs.flag_on.send("chat", _marker(), fields, headers=headers) + neighbor: Final = rigs.flag_on.send("chat", _marker(), headers={"x-litellm-trace-id": neighbor_trace}) + + assert _session(response) == expected + assert _session(neighbor) == neighbor_trace + + +@pytest.mark.parametrize("endpoint", ["chat", "messages"]) +def test_client_sent_internal_session_key_is_ignored(rigs: Rigs, endpoint: str) -> None: + forged: Final = _token("forged") + trace: Final = _token("trace") + forged_metadata: Final[dict[str, JsonValue]] = { + "metadata" if endpoint == "chat" else "litellm_metadata": {CALLER_SESSION_KEY: forged} + } + + untraced: Final = rigs.flag_on.send(endpoint, _marker(), forged_metadata) + traced: Final = rigs.flag_on.send(endpoint, _marker(), forged_metadata, headers={"x-litellm-trace-id": trace}) + + assert _session(untraced) not in (forged, None) + assert _session(traced) == trace + + +def test_identical_requests_each_log_once_with_the_body_session(rigs: Rigs) -> None: + session: Final = _token("convo") + trace: Final = _token("trace") + marker: Final = _marker() + + responses: Final = tuple( + rigs.flag_on.send("chat", marker, headers={"x-litellm-trace-id": trace}, fields={"litellm_session_id": session}) + for _ in range(3) + ) + + assert len({response.headers["x-litellm-call-id"] for response in responses}) == 3 + assert tuple(_session(response) for response in responses) == (session, session, session) + + +def _pass_through_route(rig: Rig, target: str) -> str: + path: Final = f"/caller-session-{uuid.uuid4().hex}" + created: Final = rig.proxy.post("/config/pass_through_endpoint", {"path": path, "target": target, "auth": True}) + endpoint_id: Final = string_value(object_value(list_value(created["endpoints"])[0])["id"]) + rig.scenario.cleanups.callback( + lambda: rig.proxy.request("DELETE", "/config/pass_through_endpoint", params={"endpoint_id": endpoint_id}) + ) + return path + + +@pytest.mark.parametrize("field", ["litellm_session_id", "metadata"]) +def test_pass_through_spend_log_session_is_the_body_session_beside_a_trace_header(rigs: Rigs, field: str) -> None: + marker: Final = _marker() + session: Final = _token("convo") + rig: Final = rigs.flag_on + path: Final = _pass_through_route(rig, f"{rig.upstream.url}/v1/chat/completions") + key: Final = rig.scenario.key(allowed_passthrough_routes=[path]) + fields: Final[dict[str, JsonValue]] = ( + {"litellm_session_id": session} if field == "litellm_session_id" else {"metadata": {"session_id": session}} + ) + + response: Final = rig.proxy.request( + "POST", + path, + rig.body("chat", marker, fields), + key=key, + headers={"x-litellm-trace-id": _token("trace")}, + ) + + rows: Final = eventually( + lambda: read_rows( + 'SELECT session_id FROM "LiteLLM_SpendLogs" WHERE api_key=%s', (sha256(key.encode()).hexdigest(),) + ), + lambda found: len(found) >= 1, + seconds=70, + ) + assert response.status_code == 200, response.text + assert f"answer {marker}" in response.text + assert len(rig.upstream_bodies(marker)) == 1 + assert rows == [{"session_id": session}] + + +def _turn( + rig: Rig, + marker: str, + previous: str | None = None, + *, + key: str | None = None, + headers: Headers = (), + fields: Mapping[str, JsonValue] = MappingProxyType({}), +) -> str: + chained: Final[dict[str, JsonValue]] = {"previous_response_id": previous} if previous else {} + response: Final = rig.send("responses", marker, {**chained, **fields}, headers=headers, key=key) + _spend_row(response) + return string_value(JSON_OBJECT.validate_json(response.content)["id"]) + + +def _turn_fields(shape: str, session: str) -> tuple[dict[str, JsonValue], dict[str, JsonValue]]: + match shape: + case "session-every-turn": + return {"litellm_session_id": session}, {"litellm_session_id": session} + case "session-from-turn-2": + return {}, {"litellm_session_id": session} + case _: + return {"litellm_session_id": _token("convo")}, {"litellm_session_id": _token("convo")} + + +@pytest.mark.parametrize( + ("correlated", "shape", "replays_turn_1"), + [ + (True, "session-every-turn", True), + (True, "session-from-turn-2", True), + (True, "new-session-per-turn", True), + (False, "session-from-turn-2", False), + ], +) +def test_bridged_chain_history(rigs: Rigs, correlated: bool, shape: str, replays_turn_1: bool) -> None: + rig: Final = rigs.flag_on if correlated else rigs.flag_off + first_fields, second_fields = _turn_fields(shape, _token("convo")) + markers: Final = (_marker(), _marker(), _marker()) + first: Final = _turn(rig, markers[0], headers={"x-litellm-trace-id": _token("trace")}, fields=first_fields) + second: Final = _turn(rig, markers[1], first, headers={"x-litellm-trace-id": _token("trace")}, fields=second_fields) + + _turn(rig, markers[2], second) + + replayed: Final = rig.upstream_bodies(markers[2]) + assert len(replayed) == 1, replayed + assert _user_markers(replayed[0]) == (markers if replays_turn_1 else markers[1:]) + + +def test_bridged_chains_sharing_a_metadata_session_replay_each_others_turns(rigs: Rigs) -> None: + shared: Final[dict[str, JsonValue]] = {"metadata": {"session_id": _token("convo")}} + markers: Final = (_marker(), _marker(), _marker()) + first: Final = _turn(rigs.flag_on, markers[0], fields=shared) + _turn(rigs.flag_on, markers[1], fields=shared) + + _turn(rigs.flag_on, markers[2], first) + + replayed: Final = rigs.flag_on.upstream_bodies(markers[2]) + assert len(replayed) == 1, replayed + assert _user_markers(replayed[0]) == markers + + +@dataclass(frozen=True, slots=True) +class Owners: + first: str + second: str + + +def _owners(scenario: Scenario, model: str, shape: str) -> Owners: + user: Final = _token("user") + match shape: + case "other-user": + return Owners( + scenario.key(user_id=user, models=[model]), scenario.key(user_id=_token("user"), models=[model]) + ) + case "same-user": + return Owners(scenario.key(user_id=user, models=[model]), scenario.key(user_id=user, models=[model])) + case "same-team": + team: Final = scenario.team() + return Owners(scenario.key(team_id=team, models=[model]), scenario.key(team_id=team, models=[model])) + case _: + return Owners(scenario.key(models=[model]), scenario.key(models=[model])) + + +@pytest.mark.parametrize( + ("shape", "replays_first_owner"), + [("other-user", False), ("same-user", True), ("same-team", True), ("ownerless-keys", False)], +) +def test_bridged_history_follows_the_previous_responses_owner( + rigs: Rigs, shape: str, replays_first_owner: bool +) -> None: + rig: Final = rigs.flag_off + owners: Final = _owners(rig.scenario, rig.model, shape) + session: Final = _token("convo") + markers: Final = (_marker(), _marker(), _marker()) + _turn(rig, markers[0], key=owners.first, fields={"litellm_session_id": session}) + second: Final = _turn(rig, markers[1], key=owners.second, fields={"litellm_session_id": session}) + + _turn(rig, markers[2], second, key=owners.second) + + replayed: Final = rig.upstream_bodies(markers[2]) + assert len(replayed) == 1, replayed + assert _user_markers(replayed[0]) == (markers if replays_first_owner else markers[1:]) + + +BURST: Final = ENDPOINT_STREAMS * 4 + + +def test_concurrent_burst_across_endpoints_logs_each_session_once(rigs: Rigs) -> None: + sessions: Final = tuple(_token("convo") for _ in BURST) + + def one(index: int) -> httpx.Response: + endpoint, stream = BURST[index] + return rigs.flag_on.send( + endpoint, + _marker(), + stream=stream, + headers={"x-litellm-trace-id": _token("trace")}, + fields={"litellm_session_id": sessions[index]}, + ) + + with ThreadPoolExecutor(max_workers=12) as pool: + responses: Final = tuple(pool.map(one, range(len(BURST)))) + + assert tuple(_session(response) for response in responses) == sessions + + +def _open_upstream_connections(pid: int, upstream: Wire) -> int: + port: Final = urlsplit(upstream.url).port + return sum( + 1 + for connection in psutil.Process(pid).net_connections(kind="tcp") + if connection.status == psutil.CONN_ESTABLISHED and connection.raddr and connection.raddr.port == port + ) + + +def _post_or_none(rig: Rig, session: str) -> httpx.Response | None: + try: + return rig.post( + "chat", _marker(), headers={"x-litellm-trace-id": _token("trace")}, fields={"litellm_session_id": session} + ) + except httpx.HTTPError: + return None + + +def test_killing_one_worker_mid_burst_keeps_logging_the_body_session(rigs: Rigs) -> None: + release: Final = threading.Event() + arrived: Final[SimpleQueue[bytes]] = SimpleQueue() + + def held(request: Request) -> Reply: + arrived.put(request.body) + assert release.wait(timeout=60), "the burst was never released" + return _upstream(request) + + process: Final = rigs.flag_on.process + sessions: Final = tuple(_token("convo") for _ in range(20)) + after_session: Final = _token("convo") + with wire_server(held) as wire: + rig: Final = replace(rigs.flag_on, model=_model(rigs.flag_on.scenario, wire), upstream=wire) + workers: Final = eventually(process.worker_pids, lambda found: len(found) == 2, seconds=30) + with ThreadPoolExecutor(max_workers=len(sessions)) as pool: + pending: Final = tuple(pool.submit(_post_or_none, rig, session) for session in sessions) + eventually(arrived.qsize, lambda size: size == len(sessions), seconds=60) + held_by: Final = MappingProxyType({pid: _open_upstream_connections(pid, wire) for pid in workers}) + victim, survivor = sorted(workers, key=held_by.__getitem__) + with process.replacing_worker(victim) as kill: + kill.freeze() + kill() + release.set() + results: Final = tuple(future.result() for future in pending) + after: Final = rig.send( + "chat", + _marker(), + headers={"x-litellm-trace-id": _token("trace")}, + fields={"litellm_session_id": after_session}, + ) + + served: Final = tuple( + (response, session) + for response, session in zip(results, sessions, strict=True) + if response is not None and response.status_code == 200 + ) + assert sum(held_by.values()) == len(sessions), held_by + assert len(served) == held_by[survivor] >= len(sessions) // 2, (held_by, len(served)) + assert _session(after) == after_session + assert tuple(_session(response) for response, _ in served) == tuple(session for _, session in served) diff --git a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py index 7291f49af99..de2676bafaa 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py @@ -14,8 +14,8 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm import litellm.constants as litellm_constants import litellm.proxy.spend_tracking.spend_tracking_utils as spend_tracking_utils -from litellm.constants import LITELLM_TRUNCATED_PAYLOAD_FIELD, LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE, LITTELM_CLI_SERVICE_ACCOUNT_NAME, LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME, MAX_SPEND_LOG_MODEL_NAME_LENGTH, REDACTED_BY_LITELM_STRING, SESSION_ID_OMITTED_METADATA_KEY, UNKNOWN_MODEL_SPEND_LOG_MODEL -from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY, MANAGED_BACKGROUND_RESPONSE_METADATA_KEY +from litellm.constants import LITELLM_TRUNCATED_PAYLOAD_FIELD, LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE, LITTELM_CLI_SERVICE_ACCOUNT_NAME, LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME, MAX_SPEND_LOG_MODEL_NAME_LENGTH, REDACTED_BY_LITELM_STRING, SESSION_ID_GENERATED_METADATA_KEY, SESSION_ID_OMITTED_METADATA_KEY, UNKNOWN_MODEL_SPEND_LOG_MODEL +from litellm.constants import CALLER_SESSION_ID_METADATA_KEY, INTERNAL_CALL_ORIGIN_METADATA_KEY, MANAGED_BACKGROUND_RESPONSE_METADATA_KEY from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup from litellm.llms.base_llm.ocr.transformation import OCRResponse, OCRUsageInfo from litellm.litellm_core_utils.safe_json_dumps import safe_dumps @@ -243,6 +243,82 @@ def test_batch_session_outranks_the_per_request_trace_id(): assert session_id == "batch-uid-1" +def _chat_session_id( + monkeypatch: pytest.MonkeyPatch, + litellm_params: Mapping[str, object], + correlation_in_logs: bool = True, + batch_trace_session_id: str | None = None, + metadata: Mapping[str, object] = MappingProxyType({}), +) -> str | None: + monkeypatch.setattr(litellm, "request_correlation_in_logs", correlation_in_logs) + return _get_session_id_for_spend_log( + kwargs={"litellm_trace_id": "trace-abc", "litellm_params": litellm_params}, + metadata=metadata, + standard_logging_payload=_TRACE_ONLY_STANDARD_LOGGING, + omit_when_missing=False, + batch_trace_session_id=batch_trace_session_id, + ) + + +_CALLER_SESSION_PARAMS: Final = MappingProxyType({"litellm_session_id": "sess-1", "metadata": {"trace_id": "otel-1"}}) + + +def test_caller_session_wins_over_the_per_request_trace_id(monkeypatch: pytest.MonkeyPatch): + assert _chat_session_id(monkeypatch, _CALLER_SESSION_PARAMS) == "sess-1" + + +def test_batch_session_outranks_a_caller_session(monkeypatch: pytest.MonkeyPatch): + assert _chat_session_id(monkeypatch, _CALLER_SESSION_PARAMS, batch_trace_session_id="batch-uid-1") == "batch-uid-1" + + +def test_caller_session_keeps_the_legacy_trace_id_with_correlation_logging_off(monkeypatch: pytest.MonkeyPatch): + assert _chat_session_id(monkeypatch, _CALLER_SESSION_PARAMS, correlation_in_logs=False) == "trace-abc" + + +@pytest.mark.parametrize(("correlation_in_logs", "expected_session_id"), [(True, "convo-1"), (False, "trace-abc")]) +def test_session_displaced_by_a_trace_header_wins_only_with_correlation_logging( + monkeypatch: pytest.MonkeyPatch, correlation_in_logs: bool, expected_session_id: str +): + header_chain: Final = {"litellm_session_id": "trace-abc", "metadata": {"trace_id": "trace-abc", "session_id": "trace-abc"}} + session_id: Final = _chat_session_id( + monkeypatch, + header_chain, + correlation_in_logs=correlation_in_logs, + metadata={CALLER_SESSION_ID_METADATA_KEY: "convo-1"}, + ) + assert session_id == expected_session_id + + +@pytest.mark.parametrize( + ("litellm_params", "metadata"), + [ + pytest.param({"litellm_session_id": "convo\x00x"}, {}, id="body-session"), + pytest.param({"litellm_session_id": "trace-abc"}, {CALLER_SESSION_ID_METADATA_KEY: "convo\x00x"}, id="beside-trace-header"), + ], +) +def test_caller_session_postgres_cannot_store_keeps_the_trace_id( + monkeypatch: pytest.MonkeyPatch, litellm_params: Mapping[str, object], metadata: Mapping[str, object] +): + assert _chat_session_id(monkeypatch, litellm_params, metadata=metadata) == "trace-abc" + + +@pytest.mark.parametrize( + "litellm_params", + [ + pytest.param({"metadata": {}}, id="no-session"), + pytest.param( + {"litellm_session_id": "minted-1", "metadata": {"session_id": "minted-1", SESSION_ID_GENERATED_METADATA_KEY: True}}, + id="minted-by-generate-policy", + ), + pytest.param({"litellm_session_id": "job-7", "metadata": {"trace_id": "job-7"}}, id="copied-from-trace-id"), + ], +) +def test_session_that_is_not_a_caller_conversation_keeps_the_trace_id( + monkeypatch: pytest.MonkeyPatch, litellm_params: Mapping[str, object] +): + assert _chat_session_id(monkeypatch, litellm_params) == "trace-abc" + + def test_omit_policy_still_suppresses_batch_sessions(): session_id: Final = _get_session_id_for_spend_log( kwargs={}, diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index 96c11be2650..93202d93566 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -17,6 +17,7 @@ from starlette.datastructures import Headers import litellm from litellm.constants import ( + CALLER_SESSION_ID_METADATA_KEY, ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY, SERVER_STREAMING_CLASSIFICATION_KEY, SERVER_STREAMING_CLASSIFICATION_MARKER, @@ -3308,6 +3309,93 @@ def test_add_litellm_metadata_from_request_headers_x_litellm_trace_id_sets_chain assert data["litellm_trace_id"] == "foo" +@pytest.mark.parametrize(("correlation_in_logs", "expected_spend_log_session"), [(True, "convo-1"), (False, None)]) +@pytest.mark.parametrize("session_in_metadata", [False, True], ids=["top-level", "metadata"]) +def test_x_litellm_trace_id_keeps_the_body_session_for_spend_logs_only( + monkeypatch: pytest.MonkeyPatch, + correlation_in_logs: bool, + expected_spend_log_session: str | None, + session_in_metadata: bool, +): + monkeypatch.setattr(litellm, "request_correlation_in_logs", correlation_in_logs) + metadata: Final[dict[str, str]] = {"session_id": "convo-1"} if session_in_metadata else {} + data: Final[dict[str, object]] = ( + {"metadata": metadata} if session_in_metadata else {"metadata": metadata, "litellm_session_id": "convo-1"} + ) + LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers( + headers={"x-litellm-trace-id": "req-trace-1"}, data=data, _metadata_variable_name="metadata" + ) + assert metadata.get(CALLER_SESSION_ID_METADATA_KEY) == expected_spend_log_session + assert data["litellm_session_id"] == "req-trace-1" + assert metadata["session_id"] == "req-trace-1" + assert data["litellm_trace_id"] == "req-trace-1" + + +@pytest.mark.parametrize( + "session_header", ["x-litellm-session-id", "x-claude-code-session-id", "x-session-id"] +) +@pytest.mark.parametrize("with_trace_header", [False, True], ids=["session-header-only", "beside-trace-header"]) +def test_session_headers_still_beat_the_body_session_with_correlation_logging( + monkeypatch: pytest.MonkeyPatch, session_header: str, with_trace_header: bool +): + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) + metadata: Final[dict[str, str]] = {} + data: Final[dict[str, object]] = {"metadata": metadata, "litellm_session_id": "convo-1"} + trace_header: Final = {"x-litellm-trace-id": "req-trace-1"} if with_trace_header else {} + LiteLLMProxyRequestSetup.add_litellm_metadata_from_request_headers( + headers={**trace_header, session_header: "hdr-session-1"}, data=data, _metadata_variable_name="metadata" + ) + spend_log_session: Final = metadata.get(CALLER_SESSION_ID_METADATA_KEY) or data["litellm_session_id"] + assert spend_log_session == "hdr-session-1" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("forged_in", ["metadata", "litellm_metadata"]) +@pytest.mark.parametrize( + "headers", + [{"x-litellm-session-id": "hdr-session-1"}, {}], + ids=["session-header", "no-headers"], +) +async def test_client_metadata_cannot_set_the_spend_log_caller_session( + monkeypatch: pytest.MonkeyPatch, forged_in: str, headers: dict[str, str] +): + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) + updated: Final = await add_litellm_data_to_request( + data={ + "model": "gpt-5.4-mini", + "messages": [{"role": "user", "content": "hi"}], + forged_in: {CALLER_SESSION_ID_METADATA_KEY: "forged-session"}, + }, + request=_make_request_mock("/v1/chat/completions", headers), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={}, + version="test-version", + ) + assert CALLER_SESSION_ID_METADATA_KEY not in updated["metadata"] + + +@pytest.mark.asyncio +async def test_spend_log_caller_session_stays_out_of_requester_metadata(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) + updated: Final = await add_litellm_data_to_request( + data={ + "model": "gpt-5.4-mini", + "messages": [{"role": "user", "content": "hi"}], + "litellm_session_id": "convo-1", + "metadata": {"team_label": "kept"}, + }, + request=_make_request_mock("/v1/chat/completions", {"x-litellm-trace-id": "req-trace-1"}), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={}, + version="test-version", + ) + assert updated["metadata"][CALLER_SESSION_ID_METADATA_KEY] == "convo-1" + assert updated["metadata"]["requester_metadata"]["team_label"] == "kept" + assert CALLER_SESSION_ID_METADATA_KEY not in updated["metadata"]["requester_metadata"] + + def test_add_litellm_metadata_from_request_headers_x_litellm_session_id_sets_chain_id(): """x-litellm-session-id sets both metadata and top-level litellm_session_id/litellm_trace_id for call chaining.""" headers = {"x-litellm-session-id": "bar"} diff --git a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py index 19bbb6481fb..cc5d5c62739 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -17,6 +17,7 @@ from openai.types.responses.response_function_web_search import ( from pydantic import TypeAdapter import litellm +from litellm.constants import CALLER_SESSION_ID_METADATA_KEY from litellm.integrations.anthropic_cache_control_hook import MAX_CACHE_CONTROL_BLOCKS, AnthropicCacheControlHook from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR from litellm.litellm_core_utils.prompt_templates.factory import ( @@ -6499,3 +6500,69 @@ def test_transform_request_reads_service_tier_from_responses_api_request() -> No ) assert result["service_tier"] == "priority" + + +def _continued_chain_request(request_kwargs: Mapping[str, object]) -> Mapping[str, object]: + async def session_lookup(previous_response_id: str) -> ChatCompletionSession: + assert previous_response_id == "resp_turn_1" + return ChatCompletionSession( + messages=[{"role": "user", "content": "turn 1"}], litellm_session_id="chain-session" + ) + + return asyncio.run( + LiteLLMCompletionResponsesConfig.async_responses_api_session_handler( + previous_response_id="resp_turn_1", + litellm_completion_request={ + "model": _CLAUDE_ON_BEDROCK, + "custom_llm_provider": "bedrock", + "messages": [{"role": "user", "content": "turn 2"}], + "metadata": request_kwargs.get("metadata"), + }, + instructions=None, + session_lookup=session_lookup, + request_kwargs=request_kwargs, + ) + ) + + +@pytest.mark.parametrize( + ("request_kwargs", "logged_metadata_key"), + [ + ( + { + "metadata": {"tag": "client"}, + "litellm_metadata": {CALLER_SESSION_ID_METADATA_KEY: "turn-session", "user_api_key": "hashed"}, + }, + "litellm_metadata", + ), + ({"metadata": {"user_api_key": "hashed"}}, "metadata"), + ], +) +def test_a_continued_chain_logs_its_spend_under_the_chain_session_with_request_correlation( + monkeypatch: pytest.MonkeyPatch, request_kwargs: Mapping[str, object], logged_metadata_key: str +): + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) + sent_kwargs: Final = deepcopy(request_kwargs) + + request: Final = _continued_chain_request(request_kwargs) + + assert request["litellm_trace_id"] == "chain-session" + assert "litellm_session_id" not in request + assert request[logged_metadata_key] == { + **TypeAdapter(dict[str, object]).validate_python(request_kwargs[logged_metadata_key]), + CALLER_SESSION_ID_METADATA_KEY: "chain-session", + } + assert request_kwargs == sent_kwargs + + +def test_a_continued_chain_leaves_the_turns_metadata_alone_without_request_correlation( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr(litellm, "request_correlation_in_logs", False) + caller_metadata: Final = {CALLER_SESSION_ID_METADATA_KEY: "turn-session", "user_api_key": "hashed"} + + request: Final = _continued_chain_request({"metadata": caller_metadata, "litellm_metadata": caller_metadata}) + + assert request["litellm_trace_id"] == "chain-session" + assert "litellm_metadata" not in request + assert request["metadata"] == {CALLER_SESSION_ID_METADATA_KEY: "turn-session", "user_api_key": "hashed"} diff --git a/tests/unit/responses/litellm_completion_transformation/test_session_handler.py b/tests/unit/responses/litellm_completion_transformation/test_session_handler.py index 002a595a5e7..5db563be4b0 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_session_handler.py +++ b/tests/unit/responses/litellm_completion_transformation/test_session_handler.py @@ -1,4 +1,5 @@ import json +import sqlite3 from typing import Final from unittest.mock import AsyncMock, patch @@ -787,3 +788,80 @@ async def test_message_history_replays_real_key_named_tool_payloads() -> None: tool_message: Final = result["messages"][2] assert json.loads(assistant_message["tool_calls"][0]["function"]["arguments"]) == function_arguments assert json.loads(tool_message["content"]) == function_output + + +class _SqliteBackedPrismaDB: + def __init__(self, rows: list[tuple[str, str, str, str, str | None, str | None]]) -> None: + self._connection = sqlite3.connect(":memory:") + self._connection.row_factory = sqlite3.Row + self._connection.execute( + 'CREATE TABLE "LiteLLM_SpendLogs" ' + '(request_id TEXT, api_key TEXT, session_id TEXT, "endTime" TEXT, "user" TEXT, team_id TEXT)' + ) + self._connection.executemany('INSERT INTO "LiteLLM_SpendLogs" VALUES (?, ?, ?, ?, ?, ?)', rows) + + async def query_raw(self, query: str, *args: object) -> list[dict[str, object]]: + return [dict(row) for row in self._connection.execute(query.replace("$1", "?"), args)] + + +@pytest.mark.asyncio +async def test_session_lookup_does_not_return_another_keys_rows_for_a_shared_session_id(): + """A caller-supplied litellm_session_id can collide with another key's session; history must stay per key.""" + sqlite_db: Final = _SqliteBackedPrismaDB( + [ + ("victim-1", "victim-key", "shared-session", "2026-01-01T00:00:01", "", ""), + ("victim-2", "victim-key", "shared-session", "2026-01-01T00:00:02", "", ""), + ("attacker-1", "attacker-key", "shared-session", "2026-01-01T00:00:03", "", ""), + ("attacker-2", "attacker-key", "other-session", "2026-01-01T00:00:04", "", ""), + ] + ) + fake_prisma_client = _FakePrismaClient(results=[]) + fake_prisma_client.db = sqlite_db + + with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma_client): + spend_logs = await ResponsesSessionHandler.get_all_spend_logs_for_previous_response_id("attacker-1") + victim_logs = await ResponsesSessionHandler.get_all_spend_logs_for_previous_response_id("victim-1") + + assert [row["request_id"] for row in spend_logs] == ["attacker-1"] + assert [row["request_id"] for row in victim_logs] == ["victim-1", "victim-2"] + + +async def _history_request_ids(sqlite_db: _SqliteBackedPrismaDB, response_id: str) -> list[str]: + fake_prisma_client: Final = _FakePrismaClient(results=[]) + fake_prisma_client.db = sqlite_db + with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma_client): + spend_logs: Final = await ResponsesSessionHandler.get_all_spend_logs_for_previous_response_id(response_id) + return [row["request_id"] for row in spend_logs] + + +@pytest.mark.asyncio +async def test_session_lookup_keeps_jwt_users_and_teams_apart_when_they_share_an_empty_api_key(): + sqlite_db: Final = _SqliteBackedPrismaDB( + [ + ("alice-1", "", "shared-session", "2026-01-01T00:00:01", "alice", "team-a"), + ("bob-1", "", "shared-session", "2026-01-01T00:00:02", "bob", "team-a"), + ("service-a-1", "", "shared-session", "2026-01-01T00:00:03", "", "team-a"), + ("service-b-1", "", "shared-session", "2026-01-01T00:00:04", "", "team-b"), + ("alice-2", "", "shared-session", "2026-01-01T00:00:05", "alice", "team-a"), + ] + ) + + assert await _history_request_ids(sqlite_db, "alice-1") == ["alice-1", "alice-2"] + assert await _history_request_ids(sqlite_db, "bob-1") == ["bob-1"] + assert await _history_request_ids(sqlite_db, "service-a-1") == ["service-a-1"] + assert await _history_request_ids(sqlite_db, "service-b-1") == ["service-b-1"] + + +@pytest.mark.asyncio +async def test_session_lookup_follows_the_owner_across_keys_and_keeps_ownerless_keys_apart(): + sqlite_db: Final = _SqliteBackedPrismaDB( + [ + ("alice-old-key", "hashed-old", "shared-session", "2026-01-01T00:00:01", "alice", "team-a"), + ("ownerless-1", "hashed-other", "shared-session", "2026-01-01T00:00:02", "", ""), + ("alice-new-key", "hashed-new", "shared-session", "2026-01-01T00:00:03", "alice", "team-a"), + ("ownerless-2", "hashed-another", "shared-session", "2026-01-01T00:00:04", "", ""), + ] + ) + + assert await _history_request_ids(sqlite_db, "alice-new-key") == ["alice-old-key", "alice-new-key"] + assert await _history_request_ids(sqlite_db, "ownerless-2") == ["ownerless-2"] From 7930c0e43407fcb3317f201921cc6b2c20c1a0e0 Mon Sep 17 00:00:00 2001 From: Ayush Shekhar <106994833+ayushh0110@users.noreply.github.com> Date: Fri, 16 Oct 2026 19:21:32 +0530 Subject: [PATCH 18/23] feat(router): mid-stream fallback continuation for chat completions (#41127) * feat(router): continue chat streams on fallback after a mid-stream break A chat-completions stream that breaks after content has been delivered was re-raised and the fallback deployment never ran (#40404), while the Responses-API path already continues via a prefilled assistant turn. Add the same for chat completions behind `enable_mid_stream_fallback_continuation` (opt-in, async-only), so a post-content break re-enters the fallback chain with the partial text as an assistant prefill instead of surfacing the error. Continuation is only attempted when it is safe: the stream emitted plain assistant text (no tool/function calls, thinking blocks, reasoning items, audio or images), the request is not constrained output (response_format or a forced tool_choice) and not merge-reasoning mode, and the fallback target's model supports assistant prefill. The target check runs as a deployment pre-call filter, so a chain with no prefill-capable deployment empties and the original error is surfaced rather than a duplicated or rejected request being sent. Plain reasoning_content is treated as out-of-band and does not block a continuation, matching the Responses-API path; Anthropic thinking is excluded because its signed thinking blocks cannot ride a text-only prefill. CustomStreamWrapper now latches whether a disqualifying delta was streamed and carries it on MidStreamFallbackError so the router can decide without re-scanning chunks. Default behavior is unchanged. * fix(router): harden mid-stream continuation eligibility and prefill folding Address review of the mid-stream continuation: - keep response_format={"type": "text"} eligible; only json_object / json_schema and other structured formats decline, since "text" is the unconstrained default - fold a new partial into an existing trailing assistant prefill instead of appending a second one, so a nested mid-stream break stays a single prefill turn even on providers that do not merge consecutive assistant messages - split the deployment prefill-capability lookup into two readable steps - cover the merge_reasoning decline branch, the text response_format pass, and the prefill-folding path with tests * fix(router): address review of mid-stream continuation - register the continuation deployment filter as a process-global singleton instead of tracking it per router, so discarding one router no longer removes the type-deduplicated filter that other live routers still depend on - disqualify reasoning_content: it reaches the caller as visible reasoning that a text-only prefill cannot carry, so a reasoning-capable fallback would re-derive it and produce an incoherent response - move MID_STREAM_CONTINUATION_KWARG to constants and lazy-import the filter class in the router, breaking the module-level import cycle CodeQL flagged - trim the added docstrings to the repository comment policy * fix(router): keep mid-stream continuation within the type-discipline budget The LIT* type-discipline gate counts mutable-collection use the ruff pass does not see. Annotate the reads with read-only Sequence/Mapping views, compare response_format against a frozen MappingProxyType sentinel, and mark the genuinely mutable litellm message and deployment-list boundaries with mutable-ok reasons, so LIT001 and LIT002 stay under their limits. * fix(router): clear basedpyright operator error and the import-cycle alert - narrow the accumulated stream delta content to str before appending, so the shared helper no longer trips reportOperatorIssue (str += object) - keep the continuation marker in the pre-call-check module and lazy-import it at the two router call sites instead of routing it through constants, so the router has no module-level dependency that closes an import cycle * fix(router): make the mid-stream continuation marker unforgeable The deployment filter is process-global and read the continuation marker as a truthy value from request_kwargs. Since the proxy can forward arbitrary request-body fields into the router, a client could send `"_mid_stream_continuation": true` on a normal request to filter a mixed model group down to prefill-capable deployments and steer their prompt onto that provider. Set a private sentinel object internally and check it by type instead of truthiness: a JSON request body cannot construct one, so a forged flag is ignored. Adds a regression test that forged client values (true, "true", 1, a dict) leave the deployment list untouched. * feat(router): honor per-deployment model_info for continuation prefill capability The deployment filter only read supports_assistant_prefill from the cost map, so a model configured with `model_info: {supports_assistant_prefill: true}` on the deployment (the usual way to declare a capability for a model that is not in the cost map, or to override it) was wrongly skipped. Check the deployment's model_info first (True or False both authoritative), then fall back to the cost map. Adds tests for the override in both directions. * fix(router): bound mid-stream continuation to the caller's output allowance A continuation handed the fallback the full max_tokens / max_completion_tokens even though part of the answer had already streamed, so the total could exceed the caller's limit and repeated interruptions could re-grant the allowance on every hop, amplifying spend. Reduce each configured ceiling by the tokens already emitted before re-entry, and decline the continuation (surface the original error) when the allowance is exhausted. The reductions compound across hops because the trimmed ceiling is what the next hop sees. Also add the missing explanatory comment on the model_info fall-through except. * refactor(router): extract mid-stream decline to stay within the complexity budget The output-ceiling branch pushed stream_with_fallbacks past the C901 limit. Both decline paths (ineligible after content, and exhausted allowance) re-raised the original provider error the same way, so pull that into _raise_original_mid_stream_error. It removes the duplicated branching, drops the function back under the complexity ceiling, and reads clearer. * fix(cost-map): mark Claude Sonnet 4.6 as not supporting assistant prefill The Anthropic API and Bedrock both reject an assistant-last message on claude-sonnet-4-6 with "This model does not support assistant message prefill", so the continuation filter must not route a mid-stream continuation to it. Flip the flag on every claude-sonnet-4-6 entry and on databricks' opus-4-6 entry, whose anthropic sibling already says false. The prefill check tests now register synthetic cost-map entries instead of pinning real models, which broke when claude-3-opus left the map. * refactor(router): type the streaming fallback kwargs and drop narration comments Removes an unused mutable-ok marker (LIT013) and the explanatory comments Greptile flagged, and types initial_kwargs on both streaming fallback iterators as dict[str, Any] like their Responses siblings so the new continuation helpers take a typed mapping and the basedpyright budget stays at or below main. * fix(router): run post_call_rules over the joined text of a continued stream Each provider stream wrapper validates only the fragment it produced, so a rule that trips on the primary's prefix plus the continuation's text never fired. The router now feeds litellm.post_call_rules the emitted text plus every continuation chunk before yielding it. Also drops a narrating test-helper docstring. * refactor(router): read a deployment's declared prefill support without an empty except * test(router): cover _raise_original_mid_stream_error directly for the router coverage gate * fix(router): fold the continuation prefill into a trailing assistant turn A request that already ends in an assistant message took the partial output as a second assistant turn, which prefill providers reject. The partial now extends that turn in place, with prefix set, so the continuation request never ends in two assistant messages. * fix(router): keep content blocks on a trailing assistant turn when extending it with the partial output A trailing assistant message whose content is a list of blocks was stringified before the partial output was appended, so the fallback received the Python repr of the blocks as its prefill. The partial now lands as one more text block, and a content shape that cannot be extended surfaces the original stream error instead of a garbled continuation. * fix(router): continue the trailing text block in place when extending a prefill with content blocks Anthropic strips trailing whitespace per content block, so appending the partial output as a new text block after an assistant turn that ends in "Sure, " would send the prefill as "Sure,here". The continuation now extends the last text block itself and only opens a new text block when the trailing block is not text (an image, say), so the prefill reads exactly as the text the caller already received * test(router): type the prefill check's deployment helper * fix(router): keep a continued stream on the primary's identity and widen what can continue Continuation chunks now keep the primary stream's id and created time and drop the fallback's second role delta. Requests with n > 1 no longer continue. Prefill support reads model_info.base_model and litellm_params.base_model before the deployment model. enable_mid_stream_fallback_continuation is runtime-updatable through update_settings and registers the prefill filter. The bridge rebuild of MidStreamFallbackError keeps the inner stream's disqualifier. A fallback that repeats the prefill is not doubled in the continuation input. * fix(router): decline mid-stream continuation after a content filter, merged reasoning, or a plain trailing assistant turn Also count the partial stream's completion tokens in the continued stream's usage, accept string values for enable_mid_stream_fallback_continuation, and list the new router helpers in the router coverage check * fix(router): drop re-emitted whitespace at the continuation join and sum both legs' usage consistently * fix(router): keep the full prefill so a nested break keeps the separator * fix(router): bill both calls of a continued stream and fix Bedrock Converse, Databricks, and Vertex DeepSeek prefill routes * test(router): integration cells for mid-stream fallback continuation Covers the continued stream the caller gets, the prefill turn and reduced max_tokens the fallback receives, the prefill-capability filter, the runtime toggle, repeated requests through the response cache, and the spend row of a continued stream * fix(router): log a continued stream under the caller's id, price each cut on its own prompt, and keep prefill flags per deployment * fix(router): price and log each cut leg before its continuation reuses the logging object Price the cut stream's partial and let its failure log finish (bounded) before the fallback overwrites a shared proxy logging object, and skip continuation when an output guardrail is registered or a request fallback carries its own messages. * test(router): integration cells for the output-guardrail and request-fallback declines and the cut call's failure log Cover a registered post_call guardrail and a request fallback carrying its own messages keeping the original error, an input-only guardrail leaving continuation on, and the cut call's failure event naming its own deployment and price. * fix(router): skip mid-stream continuation when a request fallback overrides request params A request-body fallback entry that sets anything besides model (stream, messages, max_tokens) changes the request the prefill continuation relies on, so a stream: false entry ended the answer with no error. * fix(router): skip mid-stream continuation for a fallback group target that overrides request params The override check read only top-level fallback entries, so a target inside a group's list, such as {"primary": [{"model": "backup", "stream": false}]} in the request, the config, or a key's router settings, still ended the answer with no error. * fix(router): check only the request's own fallback chain for request param overrides Scanning every group's chain turned continuation off for all requests when one unrelated group's target overrode a param. Resolve the chain the router will use for this request, the same way the fallback path does. * fix(router): keep the original error when request fallbacks are malformed The request-param override check resolved a standard-shape chain outside the router's own error handling, so an empty entry or a non-list chain raised StopIteration or TypeError and replaced the original mid-stream error. Skip empty entries and treat a non-list chain as no chain. * fix(router): read deployment-configured constraints, decline deployment-id fallbacks without prefill, strip only the last Bedrock text block * fix(router): read a deployment's tool_choice ahead of router defaults in the continuation check, and restore the renamed prefill test names * fix(router): filter continuation deployments on every routing strategy, read extra_body and fallback deployments' own params The deployment filter now also runs on the sync selection path that usage-based-routing v1 and custom strategies use, one predicate reads n, response_format, tool_choice and merge_reasoning_content_in_choices with extra_body laid over them, fallback deployments whose own config constrains the output are skipped, and continuation declines when the cut deployment's id is not in the router (per-request credential deployments) * fix(router): decline a cut mid-echo, reduce extra_body output limits, and check post_call_rules once on nested continuations --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/__init__.py | 2 +- litellm/constants.py | 2 + litellm/exceptions.py | 2 + litellm/litellm_core_utils/litellm_logging.py | 49 +- .../prompt_templates/common_utils.py | 5 + .../prompt_templates/factory.py | 30 +- .../litellm_core_utils/streaming_handler.py | 56 +- .../llms/databricks/chat/transformation.py | 4 +- ...odel_prices_and_context_window_backup.json | 32 +- .../spend_tracking/spend_tracking_utils.py | 8 +- litellm/router.py | 294 +++- litellm/router_utils/common_utils.py | 3 +- .../router_utils/fallback_event_handlers.py | 26 + .../continuation_prefill_check.py | 264 ++++ litellm/types/utils.py | 4 +- litellm/utils.py | 14 + model_prices_and_context_window.json | 32 +- .../router_code_coverage.py | 6 + .../test_mid_stream_fallback_continuation.py | 749 ++++++++++ ...llm_core_utils_prompt_templates_factory.py | 58 + .../test_streaming_handler.py | 134 ++ .../test_databricks_chat_transformation.py | 15 + .../test_continuation_prefill_check.py | 288 ++++ tests/unit/test_router/test_router.py | 1218 ++++++++++++++++- .../unit/test_router_model_cost_isolation.py | 13 +- 25 files changed, 3211 insertions(+), 97 deletions(-) create mode 100644 litellm/router_utils/pre_call_checks/continuation_prefill_check.py create mode 100644 tests/integration/streaming/test_mid_stream_fallback_continuation.py create mode 100644 tests/unit/router_utils/pre_call_checks/test_continuation_prefill_check.py diff --git a/litellm/__init__.py b/litellm/__init__.py index af812dade3c..88e56c14187 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -207,7 +207,7 @@ _async_failure_callback: List[Union[str, Callable[..., object], "CustomLogger"]] [] ) # internal variable - async custom callbacks are routed here. pre_call_rules: List[Callable] = [] -post_call_rules: List[Callable] = [] +post_call_rules: List[Callable[..., object]] = [] turn_off_message_logging: Optional[bool] = False standard_logging_payload_excluded_fields: Optional[List[str]] = ( None # Fields to exclude from StandardLoggingPayload before callbacks receive it diff --git a/litellm/constants.py b/litellm/constants.py index eb03d034c7e..39363cc02e4 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -57,6 +57,7 @@ RUNTIME_UPDATABLE_ROUTER_SETTINGS: Final[frozenset[str]] = frozenset( "model_group_retry_policy", "model_group_alias", "enable_weighted_failover", + "enable_mid_stream_fallback_continuation", "enable_tag_filtering", "tag_routing_prefix", "optional_pre_call_checks", @@ -146,6 +147,7 @@ HF_CONFIG_FETCH_TIMEOUT_SECONDS: Final = 10.0 # Streams exceeding this duration are terminated with a Timeout error. # None (default) = no limit. Set env var to a number of seconds to enable globally. LITELLM_MAX_STREAMING_DURATION_SECONDS: Final = get_env_float_or_none("LITELLM_MAX_STREAMING_DURATION_SECONDS") +MID_STREAM_FAILURE_LOG_WAIT_SECONDS: Final = 1.0 # Maximum number of base64 characters to keep in logging payloads. # Data URIs exceeding this are replaced with a size placeholder. diff --git a/litellm/exceptions.py b/litellm/exceptions.py index 31388b00554..b3a0802db13 100644 --- a/litellm/exceptions.py +++ b/litellm/exceptions.py @@ -1217,6 +1217,7 @@ class MidStreamFallbackError(ServiceUnavailableError): num_retries: int | None = None, generated_content: str = "", is_pre_first_chunk: bool = False, + emitted_disqualifying_content: bool = False, ): original_status: Final = getattr(original_exception, "status_code", None) self.status_code = int(original_status) if original_status is not None else 503 @@ -1229,6 +1230,7 @@ class MidStreamFallbackError(ServiceUnavailableError): self.num_retries = num_retries self.generated_content = generated_content self.is_pre_first_chunk = is_pre_first_chunk + self.emitted_disqualifying_content = emitted_disqualifying_content # Create a response if one wasn't provided if response is None: diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index e766a2be9c3..fa009ca6327 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -2295,6 +2295,36 @@ class Logging(LiteLLMLoggingBaseClass): self.model_call_details["combined_usage_object"] = usage self.model_call_details["response_cost"] = response_cost + def record_continued_partial(self, usage: Usage, response_cost: float, response_id: str) -> None: + """Bill what a cut stream already consumed on this continuation's success log, under the id the caller saw.""" + self.model_call_details["continued_partial_usage"] = usage + self.model_call_details["continued_partial_cost"] = response_cost + self.model_call_details["continued_response_id"] = response_id + + def _add_continued_partial_spend(self, response: object) -> None: + partial_usage: Final = self.model_call_details.get("continued_partial_usage") + partial_cost: Final = self.model_call_details.get("continued_partial_cost") + own_usage: Final = getattr(response, "usage", None) + own_cost: Final = self.model_call_details.get("response_cost") + if not (isinstance(partial_usage, Usage) and isinstance(partial_cost, float) and isinstance(own_usage, Usage)): + return + from litellm.cost_calculator import BaseTokenUsageProcessor + + self.model_call_details["combined_usage_object"] = BaseTokenUsageProcessor.combine_usage_objects( + usage_objects=[own_usage, partial_usage] + ) + if not isinstance(own_cost, (int, float)): + return + self.model_call_details["response_cost"] = own_cost + partial_cost + if self.cost_breakdown is not None: + self.cost_breakdown = self.cost_breakdown | CostBreakdown( + total_cost=own_cost + partial_cost, + additional_costs={ + **self.cost_breakdown.get("additional_costs", {}), + "mid_stream_partial_cost": partial_cost, + }, + ) + def record_assembled_response_for_failure(self, assembled: ModelResponse) -> None: """Bill a fully streamed response on the failure log when a post-call hook rejects it.""" usage: Final = getattr(assembled, "usage", None) @@ -2590,6 +2620,15 @@ class Logging(LiteLLMLoggingBaseClass): return response return response.model_copy(update={"model": model}) + def _with_continued_response_id( + self, + response: ModelResponse | TextCompletionResponse | ResponsesAPIResponse | InteractionsAPIResponse, + ) -> ModelResponse | TextCompletionResponse | ResponsesAPIResponse | InteractionsAPIResponse: + response_id: Final = self.model_call_details.get("continued_response_id") + if not isinstance(response_id, str) or getattr(response, "id", None) == response_id: + return response + return response.model_copy(update={"id": response_id}) + def _success_handler_helper_fn( self, result: object = None, @@ -2908,8 +2947,11 @@ class Logging(LiteLLMLoggingBaseClass): self.model_call_details["response_cost"] = self.response_cost_calculator( result=complete_streaming_response ) + self._add_continued_partial_spend(complete_streaming_response) self._merge_hidden_params_from_response_into_metadata(complete_streaming_response) - logged_streaming_response: Final = self._with_client_facing_stream_model(complete_streaming_response) + logged_streaming_response: Final = self._with_continued_response_id( + self._with_client_facing_stream_model(complete_streaming_response) + ) self.model_call_details["complete_streaming_response"] = logged_streaming_response ## STANDARDIZED LOGGING PAYLOAD self.model_call_details["standard_logging_object"] = self._build_standard_logging_payload( @@ -3445,12 +3487,15 @@ class Logging(LiteLLMLoggingBaseClass): self.model, ) self.model_call_details["response_cost"] = None + self._add_continued_partial_spend(complete_streaming_response) self._merge_hidden_params_from_response_into_metadata(complete_streaming_response) await self._prepare_baseline_cache_estimate(complete_streaming_response) - logged_streaming_response: Final = self._with_client_facing_stream_model(complete_streaming_response) + logged_streaming_response: Final = self._with_continued_response_id( + self._with_client_facing_stream_model(complete_streaming_response) + ) self.model_call_details["async_complete_streaming_response"] = logged_streaming_response ## STANDARDIZED LOGGING PAYLOAD diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index 81fc865a5bc..044123d8729 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -1894,6 +1894,11 @@ def strip_litellm_internal_message_fields(message: AllMessageValues) -> AllMessa return _message_without_fields(message, LITELLM_INTERNAL_MESSAGE_FIELDS) +def strip_prefill_prefix_field(message: AllMessageValues) -> AllMessageValues: + """Drop the ``prefix`` flag Mistral and DeepSeek read off a prefill, for endpoints that reject unknown keys.""" + return _message_without_fields(message, frozenset({"prefix"})) + + def filter_value_from_dict(dictionary: dict, key: str, depth: int = 0) -> Any: """ Filters a value from a dictionary diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 07534e70a27..87e76fe39f7 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -4303,6 +4303,24 @@ def _with_text_when_document_only(message: BedrockMessageBlock) -> BedrockMessag return BedrockMessageBlock(role="user", content=[*blocks[:cut], placeholder, *blocks[cut:]]) +def _rstrip_text_block(block: BedrockContentBlock) -> BedrockContentBlock: + return BedrockContentBlock(text=block["text"].rstrip()) if "text" in block else block + + +def _assistant_turn_rstripped(message: BedrockMessageBlock) -> BedrockMessageBlock: + """Claude on Converse rejects a final assistant turn that ends in whitespace, as the Anthropic path strips it.""" + if message["role"] != "assistant": + return message + blocks: Final = message["content"] + last_text_index: Final = max((index for index, block in enumerate(blocks) if "text" in block), default=None) + return BedrockMessageBlock( + role="assistant", + content=[ + _rstrip_text_block(block) if index == last_text_index else block for index, block in enumerate(blocks) + ], + ) + + def _ensure_document_messages_have_text( contents: list[BedrockMessageBlock], ) -> list[BedrockMessageBlock]: @@ -4894,7 +4912,12 @@ class BedrockConverseMessagesProcessor: llm_provider=llm_provider, ) - return _ensure_document_messages_have_text(_rename_duplicate_bedrock_document_names(contents)) + converse_messages: Final = _ensure_document_messages_have_text( + _rename_duplicate_bedrock_document_names(contents) + ) + if not (to_claude and converse_messages): + return converse_messages + return [*converse_messages[:-1], _assistant_turn_rstripped(converse_messages[-1])] @staticmethod def translate_thinking_blocks_to_reasoning_content_blocks( @@ -5303,7 +5326,10 @@ def _bedrock_converse_messages_pt( llm_provider=llm_provider, ) - return _ensure_document_messages_have_text(_rename_duplicate_bedrock_document_names(contents)) + converse_messages: Final = _ensure_document_messages_have_text(_rename_duplicate_bedrock_document_names(contents)) + if not (to_claude and converse_messages): + return converse_messages + return [*converse_messages[:-1], _assistant_turn_rstripped(converse_messages[-1])] def make_valid_bedrock_tool_name(input_tool_name: str) -> str: diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 46ad1b70abe..95d062a190c 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -242,6 +242,7 @@ class CustomStreamWrapper: self.completion_stream = completion_stream self.sent_first_chunk = False self.sent_last_chunk = False + self._failure_logging: asyncio.Task[None] | None = None self._closed = False self._stream_created_time: float = time.time() @@ -269,6 +270,7 @@ class CustomStreamWrapper: self.holding_chunk = "" self.complete_response = "" self.response_uptil_now = "" + self._emitted_disqualifying_content = False _model_info: Final[dict] = litellm_params.model_info or {} _api_base: Final = get_api_base( @@ -1129,6 +1131,7 @@ class CustomStreamWrapper: if self.merge_reasoning_content_in_choices is True: reasoning_content: Final = getattr(model_response.choices[0].delta, "reasoning_content", None) if reasoning_content: + self._emitted_disqualifying_content = True if self.sent_first_thinking_block is False: # Ensure content is not None before concatenation if model_response.choices[0].delta.content is None: @@ -1865,9 +1868,7 @@ class CustomStreamWrapper: if response.choices: choice = response.choices[0] if isinstance(choice, StreamingChoices): - self.response_uptil_now += choice.delta.get("content", "") or "" - else: - self.response_uptil_now += "" + self._accumulate_streamed_delta(choice.delta) self.rules.post_call_rules(input=self.response_uptil_now, model=self.model) # HANDLE STREAM OPTIONS self.chunks.append(response) @@ -2084,9 +2085,7 @@ class CustomStreamWrapper: if processed_chunk.choices: choice = processed_chunk.choices[0] if isinstance(choice, StreamingChoices): - self.response_uptil_now += choice.delta.get("content", "") or "" - else: - self.response_uptil_now += "" + self._accumulate_streamed_delta(choice.delta) self.rules.post_call_rules(input=self.response_uptil_now, model=self.model) # Add mcp_list_tools to first chunk if present if not self.sent_first_chunk and processed_chunk.choices: @@ -2156,9 +2155,7 @@ class CustomStreamWrapper: choice = processed_chunk.choices[0] if isinstance(choice, StreamingChoices): - self.response_uptil_now += choice.delta.get("content", "") or "" - else: - self.response_uptil_now += "" + self._accumulate_streamed_delta(choice.delta) self.rules.post_call_rules(input=self.response_uptil_now, model=self.model) # RETURN RESULT self.chunks.append(processed_chunk) @@ -2172,7 +2169,7 @@ class CustomStreamWrapper: if self.logging_obj is not None: self._record_partial_usage_for_failure() ## LOGGING - asyncio.create_task( + self._failure_logging = asyncio.create_task( self.logging_obj.dispatch_failure_handlers(e, traceback_exception, prefer_async_handlers=True) ) await self.aclose() @@ -2305,11 +2302,17 @@ class CustomStreamWrapper: if self.logging_obj is not None: self._record_partial_usage_for_failure() ## LOGGING - asyncio.create_task( + self._failure_logging = asyncio.create_task( self.logging_obj.dispatch_failure_handlers(e, traceback_exception, prefer_async_handlers=True) ) self._handle_stream_fallback_error(e) + async def wait_for_failure_logging(self, timeout: float) -> None: + """Let this stream's failure log read its own deployment before a fallback reuses the logging object.""" + if self._failure_logging is None: + return + await asyncio.wait((self._failure_logging,), timeout=timeout) + def _record_partial_usage_for_failure(self) -> None: """ A stream that breaks mid-flight still billed the provider for the chunks @@ -2346,6 +2349,35 @@ class CustomStreamWrapper: recover_error, ) + # Delta fields a text-only prefill continuation cannot carry, so a stream + # that emitted any of them is not eligible for mid-stream continuation. + _CONTINUATION_DISQUALIFYING_DELTA_FIELDS: Final = ( + "tool_calls", + "function_call", + "reasoning_content", + "thinking_blocks", + "reasoning_items", + "audio", + "images", + "annotations", + ) + + @classmethod + def _delta_disqualifies_continuation(cls, delta: object) -> bool: + get: Final = getattr(delta, "get", None) + if not callable(get): + return False + return any(get(field) for field in cls._CONTINUATION_DISQUALIFYING_DELTA_FIELDS) + + def _accumulate_streamed_delta(self, delta: object) -> None: + # Shared by the sync, async, and non-aiohttp iteration sites so answer + # text and the disqualifying-content latch stay in step across all three. + get: Final = getattr(delta, "get", None) + content: Final = get("content") if callable(get) else None + self.response_uptil_now += content if isinstance(content, str) else "" + if not self._emitted_disqualifying_content and self._delta_disqualifies_continuation(delta): + self._emitted_disqualifying_content = True + def _handle_stream_fallback_error(self, e: Exception) -> "NoReturn": """ Common error handling for both __next__ and __anext__. @@ -2377,6 +2409,7 @@ class CustomStreamWrapper: original_exception=e.original_exception, generated_content=self.response_uptil_now, is_pre_first_chunk=not self.sent_first_chunk, + emitted_disqualifying_content=self._emitted_disqualifying_content or e.emitted_disqualifying_content, ) # Map to OpenAI exception format. Some providers' mappers (e.g. @@ -2437,6 +2470,7 @@ class CustomStreamWrapper: original_exception=mapped_exception, generated_content=self.response_uptil_now, is_pre_first_chunk=not self.sent_first_chunk, + emitted_disqualifying_content=self._emitted_disqualifying_content, ) @staticmethod diff --git a/litellm/llms/databricks/chat/transformation.py b/litellm/llms/databricks/chat/transformation.py index 2145c19a71a..74bdd2cd733 100644 --- a/litellm/llms/databricks/chat/transformation.py +++ b/litellm/llms/databricks/chat/transformation.py @@ -20,6 +20,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( merge_consecutive_system_messages, strip_litellm_internal_message_fields, strip_name_from_message, + strip_prefill_prefix_field, ) from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator from litellm.llms.base_llm.base_utils import ( @@ -487,6 +488,7 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): Databricks does not support: - 'name' in user message. - litellm's internal `thinking_blocks` / `reasoning_content` on assistant messages. + - the `prefix` prefill flag Mistral and DeepSeek read, which Claude does not need to continue an assistant turn. - non-Claude json_object "json" checks read only string content, so a single text block is sent as a string. """ new_messages = [] @@ -496,7 +498,7 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): else: _message = message _message = strip_name_from_message(_message, allowed_name_roles=["user"]) - _message = strip_litellm_internal_message_fields(_message) + _message = strip_prefill_prefix_field(strip_litellm_internal_message_fields(_message)) # Move message-level cache_control into a content block when content is a string. if "cache_control" in _message and isinstance(_message.get("content"), str): _message = self._move_cache_control_into_string_content_block(_message) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index b744032fb6c..5fff760d85a 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -2773,7 +2773,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_pdf_input": true, @@ -2808,7 +2808,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_pdf_input": true, @@ -2843,7 +2843,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_pdf_input": true, @@ -2878,7 +2878,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_pdf_input": true, @@ -2913,7 +2913,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_pdf_input": true, @@ -2948,7 +2948,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_pdf_input": true, @@ -3743,7 +3743,7 @@ "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 1.5e-05, - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_pdf_input": true, @@ -15494,7 +15494,7 @@ }, "supports_adaptive_thinking": true, "supports_legacy_thinking": true, - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_pdf_input": true, @@ -20832,7 +20832,7 @@ "output_cost_per_token": 2.5000010000000002e-05, "output_dbu_cost_per_token": 0.000357143, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_function_calling": true, "supports_legacy_thinking": true, "supports_prompt_caching": true, @@ -21047,7 +21047,7 @@ "output_cost_per_token": 1.5000020000000002e-05, "output_dbu_cost_per_token": 0.000214286, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_function_calling": true, "supports_legacy_thinking": true, "supports_prompt_caching": true, @@ -49176,7 +49176,7 @@ "max_tokens": 64000, "mode": "chat", "output_cost_per_token": 2.5e-05, - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_prompt_caching": true, @@ -51006,7 +51006,7 @@ "mode": "chat", "output_cost_per_token": 1.5e-05, "output_cost_per_token_batches": 7.5e-06, - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_pdf_input": true, @@ -51179,7 +51179,7 @@ "supported_regions": [ "global" ], - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -51200,7 +51200,7 @@ "supported_regions": [ "us-central1" ], - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -58822,7 +58822,7 @@ "mode": "chat", "output_cost_per_token": 1.5e-05, "output_cost_per_token_batches": 7.5e-06, - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_pdf_input": true, @@ -72334,7 +72334,7 @@ "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_pdf_input": true, diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index fe55092591e..9ebea589cc9 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -573,11 +573,11 @@ def get_logging_payload( elif isinstance(_usage, dict): usage = _usage - # A request that failed mid-stream has no usable response_obj usage, but the - # streaming handler may have recovered the usage from the chunks already - # delivered. Honor that override so the partial usage lands in spend tracking. + # A request that failed mid-stream has no usable response_obj usage, and a stream + # continued on a fallback only has the fallback's. combined_usage_object holds the + # usage of every call the client's answer took, so it wins in both cases. _combined_usage: Final = kwargs.get("combined_usage_object") - if not usage and isinstance(_combined_usage, litellm.Usage): + if (not usage or "continued_partial_usage" in kwargs) and isinstance(_combined_usage, litellm.Usage): usage = _combined_usage.model_dump() id = get_spend_logs_id(call_type or "acompletion", response_obj_dict, kwargs, metadata) diff --git a/litellm/router.py b/litellm/router.py index 151a8a75dc3..3c969610e83 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -38,7 +38,7 @@ from dataclasses import dataclass from datetime import datetime, timezone from functools import lru_cache, partial from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Generic, Literal, Optional, TypeAlias, TypeVar, Union, cast +from typing import TYPE_CHECKING, Any, Final, Generic, Literal, NoReturn, Optional, TypeAlias, TypeVar, Union, cast import anyio import httpx @@ -68,6 +68,7 @@ from litellm.constants import ( DEFAULT_HEALTH_CHECK_STALENESS_MULTIPLIER, DEFAULT_MAX_LRU_CACHE_SIZE, INTERNAL_CALL_ORIGIN_METADATA_KEY, + MID_STREAM_FAILURE_LOG_WAIT_SECONDS, OUTPUT_TOKEN_CEILING_PARAMS, PROMPT_CACHE_PIN_CACHE_MAX_SIZE, ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY, @@ -75,7 +76,7 @@ from litellm.constants import ( RUNTIME_UPDATABLE_ROUTER_SETTINGS, SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, ) -from litellm.integrations.custom_guardrail import is_guardrail_intervention +from litellm.integrations.custom_guardrail import CustomGuardrail, is_guardrail_intervention from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.otel.routing import routing_decision_attributes from litellm.integrations.otel.runtime import phase_attributes, phase_event, phase_span @@ -194,6 +195,7 @@ from litellm.router_utils.clientside_credential_handler import ( is_clientside_credential, ) from litellm.router_utils.common_utils import ( + MID_STREAM_CONTINUATION_KWARG, filter_team_based_models, filter_web_search_deployments, format_fallback_outcome_message, @@ -233,7 +235,10 @@ from litellm.router_utils.fallback_event_handlers import ( check_non_standard_fallback_format, clear_pre_routing_selection, committed_retry_budget_for_request, + continuation_fallback_targets, fallback_lookup_groups, + fallback_target_name, + fallback_target_overrides_request_params, fallbacks_disabled_for_request, get_fallback_model_group_for_lookup_groups, get_pre_routing_selection, @@ -258,6 +263,17 @@ from litellm.router_utils.handle_error import ( send_llm_exception_alert, ) from litellm.router_utils.health_state_cache import DeploymentHealthCache +from litellm.router_utils.pre_call_checks.continuation_prefill_check import ( + MID_STREAM_CONTINUATION_MARKER, + ContinuationOutput, + ContinuationPrefillDeploymentCheck, + ContinuedPartial, + constrains_continuation, + continuation_deployments, + continuation_output_ceilings, + continued_text, + deployment_can_continue, +) from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( DeploymentAffinityCheck, warn_on_unknown_model_group_affinity_flags, @@ -303,6 +319,8 @@ from litellm.router_utils.routing_groups import ( ) from litellm.router_utils.routing_read_batch import ROUTER_USAGE_TARGET, RoutingPrefetch, RoutingReadBatch from litellm.scheduler import FlowItem, Scheduler +from litellm.secret_managers.main import str_to_bool +from litellm.types.guardrails import GuardrailEventHooks from litellm.types.litellm_params import RoutingStrategyName from litellm.types.llms.openai import ( AllMessageValues, @@ -371,6 +389,7 @@ from litellm.utils import ( EmbeddingResponse, ModelResponse, Rules, + TextCompletionResponse, function_setup, get_llm_provider, get_non_default_completion_params, @@ -527,6 +546,15 @@ _ROUTING_KWARGS_ADAPTER: Final[TypeAdapter[Mapping[str, object] | None]] = TypeA _FALLBACK_HOP_ADAPTER: Final = TypeAdapter(Mapping[str, object]) _OPAQUE_VALUE_ADAPTER: Final = TypeAdapter(object) _DEPLOYMENT_SELECTED_EVENT: Final = "litellm.request.deployment_selected" +_CONTENT_BLOCKS_ADAPTER: Final = TypeAdapter(list[dict[str, object]]) + + +def _extend_trailing_text_block(blocks: Sequence[Mapping[str, object]], text: str) -> Sequence[Mapping[str, object]]: + """Continue the last text block in place so the prefill reads as the text the caller already received.""" + last: Final = blocks[-1] if blocks else None + if last is not None and last.get("type") == "text" and isinstance(last_text := last.get("text"), str): + return [*blocks[:-1], {**last, "text": f"{last_text}{text}"}] + return [*blocks, {"type": "text", "text": text}] def _is_fallback_hop(request_kwargs: Mapping[str, object]) -> bool: @@ -1119,6 +1147,7 @@ class Router: health_check_ignore_transient_errors: bool = False, background_health_check_model_groups: Sequence[str] | None = None, enable_weighted_failover: bool = False, + enable_mid_stream_fallback_continuation: bool = False, fallback_access_check: FallbackAccessCheck | None = None, fallback_budget_check: FallbackBudgetCheck | None = None, auto_router_capability_limit: AutoRouterCapabilityLimit | None = None, @@ -1160,6 +1189,7 @@ class Router: deployment_affinity_ttl_seconds (int): TTL for user-key -> deployment affinity mapping. Defaults to 3600. ignore_invalid_deployments (bool): Ignores invalid deployments, and continues with other deployments. Default is to raise an error. enable_weighted_failover (bool): When True and the routing strategy is "simple-shuffle", a retryable failure on one deployment causes the request to re-pick (weighted) across the other deployments in the same model group before any cross-group fallback runs. Bounded by `max_fallbacks`. Async-only: currently honored by `router.acompletion()` and other async entrypoints. The sync `router.completion()` path falls back to the regular fallback flow. Defaults to False. + enable_mid_stream_fallback_continuation (bool): When True, a chat-completions stream that breaks after plain assistant text has been delivered continues on a fallback deployment via assistant prefill instead of surfacing the error. Only deployments whose model supports assistant prefill are eligible, so the partial text is continued, not regenerated; if none is, the original error is surfaced. Streams that emitted tool calls, thinking blocks, audio/images, or constrained (JSON / forced tool_choice) output are never continued. Async-only. Defaults to False. fallback_access_check (Optional[FallbackAccessCheck]): Awaited before each cross-model-group fallback attempt on the async path; a fallback target it rejects is skipped. Defaults to None (every configured fallback is attempted). fallback_budget_check (Optional[FallbackBudgetCheck]): Awaited before each cross-model-group fallback attempt on the async path; a fallback target it rejects as over budget is skipped. Defaults to None (budget is not re-checked on fallback). Returns: @@ -1355,6 +1385,7 @@ class Router: self.fallback_on_unsupported_params = fallback_on_unsupported_params self.enable_health_check_routing = enable_health_check_routing self.enable_weighted_failover = enable_weighted_failover + self.enable_mid_stream_fallback_continuation = self._is_flag_enabled(enable_mid_stream_fallback_continuation) self.health_check_ignore_transient_errors = health_check_ignore_transient_errors self.background_health_check_model_groups: frozenset[str] | None = ( frozenset(background_health_check_model_groups) @@ -1549,6 +1580,14 @@ class Router: default_pre_call_checks: Final[OptionalPreCallChecks] = [] self.add_optional_pre_call_checks(default_pre_call_checks) + if self.enable_mid_stream_fallback_continuation: + self._register_continuation_prefill_check() + + @staticmethod + def _register_continuation_prefill_check() -> None: + # Process-global on purpose: Router.discard() must not drop a filter another router still needs. + litellm.logging_callback_manager.add_litellm_callback(ContinuationPrefillDeploymentCheck()) + def discard(self): """ Pseudo-destructor to be invoked to clean up global data structures when router is no longer used. @@ -3143,15 +3182,53 @@ class Router: usage: Final = cast(Usage | None, getattr(fallback_item, "usage", None)) usage_objects: Final = [usage] if usage is not None else [] - if ( - complete_response_object_usage is not None - and hasattr(complete_response_object_usage, "usage") - and complete_response_object_usage.usage is not None - ): + if complete_response_object_usage is not None: usage_objects.append(complete_response_object_usage) combined_usage: Final = BaseTokenUsageProcessor.combine_usage_objects(usage_objects=usage_objects) setattr(fallback_item, "usage", combined_usage) + @staticmethod + def _price_cut_partial( + cut_stream: CustomStreamWrapper, + partial_response: ModelResponse | TextCompletionResponse | None, + ) -> "ContinuedPartial | None": + """Price the cut stream's partial, plus any partial it inherited, before a fallback can reuse its logging.""" + from litellm.cost_calculator import BaseTokenUsageProcessor + + cut_logging: Final = cut_stream.logging_obj + partial_usage: Final = getattr(partial_response, "usage", None) + if partial_response is None or not isinstance(partial_usage, Usage): + return None + inherited_usage: Final = cut_logging.model_call_details.get("continued_partial_usage") + inherited_cost: Final = cut_logging.model_call_details.get("continued_partial_cost") + inherited_id: Final = cut_logging.model_call_details.get("continued_response_id") + priced_partial: Final = partial_response.model_copy( + update={"model": cut_stream.model} if cut_stream.model else {} + ) + priced_partial._hidden_params = { # pyright: ignore[reportPrivateUsage] # no public accessor + key: value for key, value in get_hidden_params_dict(partial_response).items() if key != "response_cost" + } + return ContinuedPartial( + usage=BaseTokenUsageProcessor.combine_usage_objects( + usage_objects=[partial_usage, inherited_usage] + if isinstance(inherited_usage, Usage) + else [partial_usage] + ), + response_cost=(cut_logging.response_cost_calculator(result=priced_partial) or 0.0) + + (inherited_cost if isinstance(inherited_cost, float) else 0.0), + response_id=inherited_id if isinstance(inherited_id, str) else partial_response.id, + ) + + @staticmethod + def _carry_partial_spend_to_fallback(fallback_response: object, cut_partial: "ContinuedPartial | None") -> None: + """Bill what the cut legs consumed on the continuation's success log.""" + fallback_logging: Final = getattr(fallback_response, "logging_obj", None) + if cut_partial is None or not isinstance(fallback_logging, LiteLLMLogging): + return + fallback_logging.record_continued_partial( + usage=cut_partial.usage, response_cost=cut_partial.response_cost, response_id=cut_partial.response_id + ) + @staticmethod def _prepare_fallback_hidden_params( fallback_response: object, @@ -3201,7 +3278,7 @@ class Router: self, model_response: CustomStreamWrapper, messages: list[dict[str, str]], - initial_kwargs: dict, + initial_kwargs: dict[str, Any], deployment_slot: contextlib.AsyncExitStack | None = None, first_chunk: _FirstChunk[AsyncIterator[object]] | None = None, ) -> CustomStreamWrapper: @@ -3264,16 +3341,22 @@ class Router: with anyio.CancelScope(shield=True): await close_model_response() await held_slot.aclose() - if not e.is_pre_first_chunk and ( - e.generated_content or _stream_chunks_have_generated_content(model_response.chunks) - ): - if e.original_exception is not None: - raise e.original_exception from e - raise + committed: Final = bool( + not e.is_pre_first_chunk + and (e.generated_content or _stream_chunks_have_generated_content(model_response.chunks)) + ) + cut_hidden_params: Final = get_hidden_params(model_response) or {} + continue_after_content: Final = committed and self._mid_stream_continuation_eligible( + e=e, request_kwargs=initial_kwargs, cut_deployment_id=cut_hidden_params.get("model_id") + ) + if committed and not continue_after_content: + self._raise_original_mid_stream_error(e) from litellm.main import stream_chunk_builder - complete_response_object: Final = stream_chunk_builder(chunks=model_response.chunks) + complete_response_object: Final = stream_chunk_builder( + chunks=model_response.chunks, messages=messages if continue_after_content else None + ) complete_response_object_usage: Final = cast( Usage | None, getattr(complete_response_object, "usage", None), @@ -3289,7 +3372,41 @@ class Router: "content_policy_fallbacks", self.content_policy_fallbacks ) initial_kwargs["original_function"] = self._acompletion - initial_kwargs["messages"] = messages + continuation_messages: Final = ( + self._build_completion_continuation_input(messages, e.generated_content) + if continue_after_content + else None + ) + continuation_output: Final = ( + ContinuationOutput( + e.generated_content, + continuation_messages[-1], + model_group, + complete_response_object, + checks_post_call_rules=initial_kwargs.get(MID_STREAM_CONTINUATION_KWARG) + is not MID_STREAM_CONTINUATION_MARKER, + ) + if continuation_messages is not None + else None + ) + cut_partial: Final = ( + Router._price_cut_partial(model_response, complete_response_object) + if continuation_output is not None + else None + ) + if continue_after_content: + emitted_tokens: Final = int( + getattr(complete_response_object_usage, "completion_tokens", 0) or 0 + ) + reduced_ceilings: Final = continuation_output_ceilings(initial_kwargs, emitted_tokens) + if reduced_ceilings is None or continuation_messages is None: + self._raise_original_mid_stream_error(e) + await model_response.wait_for_failure_logging(timeout=MID_STREAM_FAILURE_LOG_WAIT_SECONDS) + initial_kwargs.update(reduced_ceilings) + initial_kwargs["messages"] = continuation_messages + initial_kwargs[MID_STREAM_CONTINUATION_KWARG] = MID_STREAM_CONTINUATION_MARKER + else: + initial_kwargs["messages"] = messages self._update_kwargs_before_fallbacks(model=model_group, kwargs=initial_kwargs) fallback_response = await self.async_function_with_fallbacks_common_utils( e=e, @@ -3305,11 +3422,17 @@ class Router: # If fallback returns a streaming response, iterate over it if hasattr(fallback_response, "__aiter__"): + Router._carry_partial_spend_to_fallback(fallback_response, cut_partial) prepared_fallback_hidden_params = Router._adopt_fallback_response_headers( wrapper_ref, fallback_response ) fallback_headers_are_settled = False - async for fallback_item in fallback_response: + fallback_items: Final = ( + continuation_output.splice(fallback_response) + if continuation_output is not None + else fallback_response + ) + async for fallback_item in fallback_items: if not fallback_headers_are_settled: fallback_headers_are_settled = True # a fallback that failed over again only repoints itself once it yields @@ -3322,7 +3445,10 @@ class Router: and isinstance(fallback_item, ModelResponseStream) and hasattr(fallback_item, "usage") ): - self._combine_fallback_usage(fallback_item, complete_response_object_usage) + self._combine_fallback_usage( + fallback_item, + complete_response_object_usage if continuation_output is not None else None, + ) yield fallback_item else: # If fallback returns a non-streaming response, yield None @@ -3479,6 +3605,97 @@ class Router: total_tokens=(partial_usage.total_tokens or 0) + (fb.total_tokens or 0), ) + @staticmethod + def _raise_original_mid_stream_error(e: "MidStreamFallbackError") -> NoReturn: + """Surface the provider error the stream wrapper carried instead of the internal MidStreamFallbackError.""" + if e.original_exception is not None: + raise e.original_exception from e + raise e + + def _params_sent_to_deployment( + self, request_kwargs: Mapping[str, object], deployment: Deployment | None + ) -> Mapping[str, object]: + """The request as the deployment received it, merged in the order _update_kwargs_with_deployment applies.""" + deployment_params: Final = deployment.litellm_params.model_dump(exclude_none=True) if deployment else {} + defaults: Final = {key: value for key, value in self.default_litellm_params.items() if value is not None} + deployment_tool_choice: Final = ( + {"tool_choice": deployment_params["tool_choice"]} + if deployment_params.get("tool_choice") is not None + else {} + ) + return MappingProxyType({**deployment_params, **defaults, **deployment_tool_choice, **request_kwargs}) + + def _names_deployment_that_cannot_continue(self, target: object) -> bool: + name: Final = fallback_target_name(target) + if name is None or name in self.model_names or not self.has_model_id(name): + return False + deployment: Final = self.get_deployment(model_id=name) + return deployment is not None and not deployment_can_continue(deployment.model_dump(exclude_none=True)) + + def _mid_stream_continuation_eligible( + self, + e: "MidStreamFallbackError", + request_kwargs: Mapping[str, object], + cut_deployment_id: object = None, + ) -> bool: + """Whether a stream that broke after plain assistant text may continue via prefill.""" + if not self.enable_mid_stream_fallback_continuation: + return False + if not e.generated_content or e.emitted_disqualifying_content: + return False + if isinstance(e.original_exception, litellm.ContentPolicyViolationError): + return False + cut_deployment: Final = ( + self.get_deployment(model_id=cut_deployment_id) if isinstance(cut_deployment_id, str) else None + ) + if isinstance(cut_deployment_id, str) and cut_deployment is None: + return False + if constrains_continuation(self._params_sent_to_deployment(request_kwargs, cut_deployment)): + return False + requested_group: Final = request_kwargs.get("model") + lookup_groups: Final = fallback_lookup_groups( + request_kwargs, requested_group if isinstance(requested_group, str) else None + ) + targets: Final = continuation_fallback_targets(request_kwargs.get("fallbacks", self.fallbacks), lookup_groups) + if any(fallback_target_overrides_request_params(target) for target in targets): + return False + if any(self._names_deployment_that_cannot_continue(target) for target in targets): + return False + return not any( + isinstance(callback, CustomGuardrail) + and callback._event_hook_is_event_type( # pyright: ignore[reportPrivateUsage] # the only check of a guardrail's hooks + GuardrailEventHooks.post_call + ) + for callback in litellm.callbacks + ) + + @staticmethod + def _build_completion_continuation_input( + messages: Sequence[Mapping[str, object]], + generated_content: str, + ) -> Sequence[Mapping[str, object]] | None: + """Append the partial output as an assistant prefill, or extend a trailing `prefix: True` assistant turn in + place so the request never ends in two assistant messages. None when the request ends in a plain assistant + turn, in a prefill turn whose content has a shape this cannot extend, or when the cut came mid-echo of the + prefill.""" + last: Final = messages[-1] if messages else None + if last is None or last.get("role") != "assistant": + return [*messages, {"role": "assistant", "content": generated_content, "prefix": True}] + if last.get("prefix") is not True: + return None + new_text: Final = continued_text(last, generated_content) + if new_text is None: + return None + content: Final = last.get("content") + if content is None or isinstance(content, str): + return [*messages[:-1], {**last, "content": f"{content or ''}{new_text}", "prefix": True}] + try: + blocks: Final = _CONTENT_BLOCKS_ADAPTER.validate_python(content) + except ValidationError: + return None + extended: Final = _extend_trailing_text_block(blocks, new_text) + return [*messages[:-1], {**last, "content": extended, "prefix": True}] + @staticmethod def _build_responses_continuation_input( input_val: Union[str, "ResponseInputParam"] | None, @@ -3880,13 +4097,6 @@ class Router: raise e.original_exception from e raise - from litellm.main import stream_chunk_builder - - complete_response_object: Final = stream_chunk_builder(chunks=model_response.chunks) - complete_response_object_usage: Final = cast( - Usage | None, - getattr(complete_response_object, "usage", None), - ) try: model_group: Final = cast(str, initial_kwargs.get("model")) fallbacks: Final[list | None] = initial_kwargs.get("fallbacks", router_self.fallbacks) @@ -3932,7 +4142,7 @@ class Router: and isinstance(fallback_item, ModelResponseStream) and hasattr(fallback_item, "usage") ): - router_self._combine_fallback_usage(fallback_item, complete_response_object_usage) + router_self._combine_fallback_usage(fallback_item, None) yield fallback_item else: yield None @@ -13004,6 +13214,7 @@ class Router: "retry_policy", "model_group_alias", "enable_weighted_failover", + "enable_mid_stream_fallback_continuation", "enable_tag_filtering", "tag_routing_prefix", "fallback_on_unsupported_params", @@ -13060,12 +13271,10 @@ class Router: rebuild_routing_groups = True elif var == "optional_pre_call_checks": self.set_optional_pre_call_checks(kwargs[var]) + elif var == "enable_mid_stream_fallback_continuation": + self._set_mid_stream_fallback_continuation(kwargs[var]) elif var == "retry_policy": - value = kwargs[var] - if isinstance(value, dict): - value = RetryPolicy(**value) - if value is None or isinstance(value, RetryPolicy): - setattr(self, var, value) + self._set_retry_policy(kwargs[var]) else: value = kwargs[var] # only run routing strategy init if it has changed @@ -13092,6 +13301,20 @@ class Router: self._routing_groups_input = routing_groups_input verbose_router_logger.debug("Updated Router settings: %s", self.get_settings()) + @staticmethod + def _is_flag_enabled(value: object) -> bool: + return value is True or (isinstance(value, str) and str_to_bool(value) is True) + + def _set_mid_stream_fallback_continuation(self, enabled: object) -> None: + self.enable_mid_stream_fallback_continuation = self._is_flag_enabled(enabled) + if self.enable_mid_stream_fallback_continuation: + self._register_continuation_prefill_check() + + def _set_retry_policy(self, value: object) -> None: + policy: Final = RetryPolicy.model_validate(value) if isinstance(value, dict) else value + if policy is None or isinstance(policy, RetryPolicy): + self.retry_policy = policy + def _get_client(self, deployment, kwargs, client_type=None): """ Returns the appropriate client based on the given deployment, kwargs, and client_type. @@ -15119,9 +15342,8 @@ class Router: request_kwargs=request_kwargs, ) strategy, strategy_selector = self._get_routing_context(model, request_kwargs) - pick_attributes: Final = _deployment_pick_attributes( - model, _ROUTING_KWARGS_ADAPTER.validate_python(request_kwargs) - ) + routing_kwargs: Final = _ROUTING_KWARGS_ADAPTER.validate_python(request_kwargs) + pick_attributes: Final = _deployment_pick_attributes(model, routing_kwargs) if isinstance(healthy_deployments, dict): if (healthy_deployments.get("model_info") or {}).get("blocked") is True: @@ -15159,7 +15381,9 @@ class Router: ) healthy_deployments = _pre_cooldown_deployments - healthy_deployments = self._filter_blocked_deployments(healthy_deployments) + healthy_deployments = continuation_deployments( + self._filter_blocked_deployments(healthy_deployments), routing_kwargs + ) # filter pre-call checks if self.enable_pre_call_checks and (messages is not None or input is not None): diff --git a/litellm/router_utils/common_utils.py b/litellm/router_utils/common_utils.py index e12304c6ed1..c4e4d9ccc0b 100644 --- a/litellm/router_utils/common_utils.py +++ b/litellm/router_utils/common_utils.py @@ -27,7 +27,8 @@ from litellm.types.utils import LlmProviders, ProviderSpecificModelInfo _V = TypeVar("_V") -ROUTER_ONLY_CALL_KWARGS: Final = frozenset({"silent_model", "include_fallback_errors"}) +MID_STREAM_CONTINUATION_KWARG: Final = "_mid_stream_continuation" +ROUTER_ONLY_CALL_KWARGS: Final = frozenset({"silent_model", "include_fallback_errors", MID_STREAM_CONTINUATION_KWARG}) _OBJECT_MAPPING_ADAPTER: Final[TypeAdapter[Mapping[object, object]]] = TypeAdapter(Mapping[object, object]) diff --git a/litellm/router_utils/fallback_event_handlers.py b/litellm/router_utils/fallback_event_handlers.py index 1f37c6c8cb1..06bd3329cc1 100644 --- a/litellm/router_utils/fallback_event_handlers.py +++ b/litellm/router_utils/fallback_event_handlers.py @@ -615,6 +615,32 @@ def get_fallback_model_group( return fallback_model_group, generic_fallback_idx +def _resolved_fallback_chain(entries: Sequence[object], lookup_groups: tuple[str, ...]) -> Sequence[object]: + group_entries: Final[Sequence[Mapping[str, list[str]]]] = tuple( # mutable-ok: mirrors the fallbacks config shape + entry for entry in entries if isinstance(entry, Mapping) and entry + ) + chain, _ = get_fallback_model_group_for_lookup_groups(fallbacks=group_entries, lookup_groups=lookup_groups) + return chain if isinstance(chain, list) else () + + +def continuation_fallback_targets(fallbacks: object, lookup_groups: tuple[str, ...]) -> Sequence[object]: + entries: Final[Sequence[object]] = fallbacks if isinstance(fallbacks, list) else () + return ( + entries + if check_non_standard_fallback_format(fallbacks=entries) + else _resolved_fallback_chain(entries, lookup_groups) + ) + + +def fallback_target_overrides_request_params(target: object) -> bool: + return isinstance(target, Mapping) and "model" in target and len(target) > 1 + + +def fallback_target_name(target: object) -> str | None: + name: Final = target.get("model") if isinstance(target, Mapping) else target + return name if isinstance(name, str) else None + + PROVIDER_SCOPED_RESOURCE_KEYS: Final = ("input_file_id", "training_file", "batch_id", "file_id", "fine_tuning_job_id") PROVIDER_SCOPED_RESOURCE_FUNCTION_NAMES: Final = frozenset( { diff --git a/litellm/router_utils/pre_call_checks/continuation_prefill_check.py b/litellm/router_utils/pre_call_checks/continuation_prefill_check.py new file mode 100644 index 00000000000..8f211ef67db --- /dev/null +++ b/litellm/router_utils/pre_call_checks/continuation_prefill_check.py @@ -0,0 +1,264 @@ +"""Keeps a mid-stream continuation on deployments whose model supports assistant prefill and whose own params leave +the output unconstrained.""" + +from collections.abc import AsyncIterable, AsyncIterator, Mapping, Sequence +from dataclasses import dataclass +from itertools import accumulate +from os.path import commonprefix +from types import MappingProxyType +from typing import Final + +from pydantic import TypeAdapter, ValidationError + +import litellm +from litellm.integrations.custom_logger import CustomLogger, Span +from litellm.litellm_core_utils.rules import Rules +from litellm.router_utils.common_utils import MID_STREAM_CONTINUATION_KWARG +from litellm.types.llms.openai import AllMessageValues +from litellm.types.utils import ModelResponse, ModelResponseStream, Usage +from litellm.utils import supports_assistant_prefill + + +class _ContinuationMarker: + """Only the router constructs one, so a JSON request body cannot forge the continuation flag.""" + + +MID_STREAM_CONTINUATION_MARKER: Final = _ContinuationMarker() + +_STR_KEYED_DICT_ADAPTER: Final = TypeAdapter(dict[str, object]) +_NO_FIELDS: Final[Mapping[str, object]] = MappingProxyType({}) +_UNCONSTRAINED_RESPONSE_FORMAT: Final[Mapping[str, str]] = MappingProxyType({"type": "text"}) +_OUTPUT_CEILING_KEYS: Final = ("max_tokens", "max_completion_tokens") + + +def _str_keyed(value: object) -> Mapping[str, object]: + try: + return _STR_KEYED_DICT_ADAPTER.validate_python(value) + except ValidationError: + return _NO_FIELDS + + +def deployment_supports_prefill(deployment: object) -> bool: + deployment_map: Final = _str_keyed(deployment) + model_info: Final = _str_keyed(deployment_map.get("model_info")) + declared: Final = model_info.get("supports_assistant_prefill") + if isinstance(declared, bool): + return declared + litellm_params: Final = _str_keyed(deployment_map.get("litellm_params")) + candidates: Final = (model_info.get("base_model"), litellm_params.get("base_model"), litellm_params.get("model")) + model: Final = next((candidate for candidate in candidates if isinstance(candidate, str) and candidate), None) + return model is not None and supports_assistant_prefill(model=model) + + +def constrains_continuation(params: Mapping[str, object]) -> bool: + sent: Final = {**params, **_str_keyed(params.get("extra_body"))} + choice_count: Final = sent.get("n") + response_format: Final = sent.get("response_format") + tool_choice: Final = sent.get("tool_choice") + return ( + (isinstance(choice_count, int) and choice_count > 1) + or (response_format is not None and response_format != _UNCONSTRAINED_RESPONSE_FORMAT) + or tool_choice == "required" + or isinstance(tool_choice, Mapping) + or sent.get("merge_reasoning_content_in_choices") is True + ) + + +def deployment_can_continue(deployment: object) -> bool: + litellm_params: Final = _str_keyed(_str_keyed(deployment).get("litellm_params")) + return deployment_supports_prefill(deployment) and not constrains_continuation(litellm_params) + + +def continuation_deployments( + healthy_deployments: list[dict[str, object]], # mutable-ok: router deployment-list contract + request_kwargs: Mapping[str, object] | None, +) -> list[dict[str, object]]: # mutable-ok: returns a mutable deployment list + marker: Final = (request_kwargs or {}).get(MID_STREAM_CONTINUATION_KWARG) + if not isinstance(marker, _ContinuationMarker): + return healthy_deployments + return [deployment for deployment in healthy_deployments if deployment_can_continue(deployment)] + + +class ContinuationPrefillDeploymentCheck(CustomLogger): + async def async_filter_deployments( + self, + model: str, + healthy_deployments: list[dict[str, object]], # mutable-ok: CustomLogger deployment-list contract + messages: Sequence[AllMessageValues] | None, + request_kwargs: Mapping[str, object] | None = None, + parent_otel_span: Span | None = None, + ) -> list[dict[str, object]]: # mutable-ok: returns a mutable deployment list + return continuation_deployments(healthy_deployments, request_kwargs) + + +def _delta_text(chunk: ModelResponseStream) -> str: + if not chunk.choices: + return "" + content: Final = chunk.choices[0].delta.content + return content if isinstance(content, str) else "" + + +_CONTENT_BLOCKS_ADAPTER: Final = TypeAdapter(list[dict[str, object]]) + + +def _prefill_text(prefill_message: Mapping[str, object]) -> str: + content: Final = prefill_message.get("content") + if isinstance(content, str): + return content + try: + blocks: Final = _CONTENT_BLOCKS_ADAPTER.validate_python(content) + except ValidationError: + return "" + return "".join( + text for block in blocks if block.get("type") == "text" and isinstance(text := block.get("text"), str) + ) + + +def _resendable_whitespace(prefill: str) -> str: + return prefill[len(prefill.rstrip()) :] if prefill.strip() else "" + + +def continued_text(prefill_message: Mapping[str, object], generated_content: str) -> str | None: + prefill: Final = _prefill_text(prefill_message) + if generated_content and len(generated_content) < len(prefill) and prefill.startswith(generated_content): + return None + echo: Final = len(prefill) if prefill and generated_content.startswith(prefill) else 0 + after_echo: Final = generated_content[echo:] + return after_echo[len(commonprefix((after_echo, _resendable_whitespace(prefill)))) :] + + +def _reduced_ceilings(params: Mapping[str, object], emitted_tokens: int) -> Mapping[str, int]: + return MappingProxyType( + {key: value - emitted_tokens for key in _OUTPUT_CEILING_KEYS if isinstance(value := params.get(key), int)} + ) + + +def continuation_output_ceilings( + request_kwargs: Mapping[str, object], emitted_tokens: int +) -> Mapping[str, object] | None: + """The caller's output ceilings, top-level and in extra_body, minus the tokens already emitted, or None once + nothing is left.""" + extra_body: Final = _str_keyed(request_kwargs.get("extra_body")) + top_level: Final = _reduced_ceilings(request_kwargs, emitted_tokens) + in_extra_body: Final = _reduced_ceilings(extra_body, emitted_tokens) + if any(ceiling <= 0 for ceiling in (*top_level.values(), *in_extra_body.values())): + return None + if not in_extra_body: + return top_level + return MappingProxyType({**top_level, "extra_body": {**extra_body, **in_extra_body}}) + + +def _with_first_delta(chunk: ModelResponseStream, update: Mapping[str, object]) -> ModelResponseStream: + choice: Final = chunk.choices[0] + delta: Final = choice.delta.model_copy(update=update) + return chunk.model_copy(update={"choices": [choice.model_copy(update={"delta": delta}), *chunk.choices[1:]]}) + + +def _without_leading_text(chunk: ModelResponseStream, count: int) -> ModelResponseStream: + if count <= 0: + return chunk + return _with_first_delta(chunk, {"content": _delta_text(chunk)[count:]}) + + +def _without_role(chunk: ModelResponseStream) -> ModelResponseStream: + if not chunk.choices or chunk.choices[0].delta.role is None: + return chunk + return _with_first_delta(chunk, {"role": None}) + + +def _stream_identity(primary_response: object) -> Mapping[str, object]: + if not isinstance(primary_response, ModelResponse): + return _NO_FIELDS + return MappingProxyType({"id": primary_response.id, "created": primary_response.created}) + + +def _drop_leading_text(chunks: tuple[ModelResponseStream, ...], count: int) -> tuple[ModelResponseStream, ...]: + lengths: Final = tuple(len(_delta_text(chunk)) for chunk in chunks) + offsets: Final = tuple(accumulate(lengths, initial=0))[:-1] + return tuple( + _without_leading_text(chunk, min(length, count - offset)) + for chunk, offset, length in zip(chunks, offsets, lengths, strict=True) + ) + + +@dataclass(frozen=True, slots=True) +class ContinuedPartial: + usage: Usage + response_cost: float + response_id: str + + +class ContinuationOutput: + """Splices a fallback's continuation onto the stream the client already received. Drops a leading echo of the + prefill, since some providers (Mistral) return the prefix inside the completion, and the prefill's trailing + whitespace sent again by a provider that strips it (Anthropic), keeps the primary's chunk id, created time and single role delta, and + runs litellm.post_call_rules over the joined text, since each stream wrapper only sees its own fragment. A + continuation nested in another leaves the rules to the outermost one, the only one holding the whole text.""" + + def __init__( + self, + emitted_text: str, + prefill_message: Mapping[str, object], + model: str, + primary_response: object, + checks_post_call_rules: bool = True, + ) -> None: + self._text = emitted_text + self._checks_post_call_rules: Final = checks_post_call_rules + self._prefill_text: Final = _prefill_text(prefill_message) + self._model: Final = model + self._identity: Final = _stream_identity(primary_response) + self._rules: Final = Rules() + self._held: tuple[ModelResponseStream, ...] = () + self._echo_settled = not self._prefill_text + self._seam = _resendable_whitespace(self._prefill_text) + + async def splice(self, stream: AsyncIterable[object]) -> AsyncIterator[object]: + async for item in stream: + if not isinstance(item, ModelResponseStream): + yield item + continue + for released in self._observed(self._past_seam(self._released(item))): + yield self._in_primary_stream(released) + for released in self._observed(self._past_seam(self._flushed())): + yield self._in_primary_stream(released) + + def _in_primary_stream(self, chunk: ModelResponseStream) -> ModelResponseStream: + return _without_role(chunk.model_copy(update=self._identity) if self._identity else chunk) + + def _flushed(self) -> tuple[ModelResponseStream, ...]: + held: Final = self._held + self._held = () + self._echo_settled = True + return held + + def _released(self, chunk: ModelResponseStream) -> tuple[ModelResponseStream, ...]: + if self._echo_settled: + return (chunk,) + held: Final = (*self._held, chunk) + text: Final = "".join(_delta_text(held_chunk) for held_chunk in held) + if len(text) < len(self._prefill_text) and self._prefill_text.startswith(text): + self._held = held + return () + self._held = () + self._echo_settled = True + if not text.startswith(self._prefill_text): + return held + return _drop_leading_text(held, len(self._prefill_text)) + + def _past_seam(self, chunks: tuple[ModelResponseStream, ...]) -> tuple[ModelResponseStream, ...]: + return tuple(self._without_seam(chunk) for chunk in chunks) + + def _without_seam(self, chunk: ModelResponseStream) -> ModelResponseStream: + text: Final = _delta_text(chunk) + matched: Final = len(commonprefix((text, self._seam))) + self._seam = self._seam[matched:] if matched == len(text) else "" + return _without_leading_text(chunk, matched) + + def _observed(self, chunks: tuple[ModelResponseStream, ...]) -> tuple[ModelResponseStream, ...]: + if not self._checks_post_call_rules or not litellm.post_call_rules: + return chunks + for chunk in chunks: + self._text += _delta_text(chunk) + self._rules.post_call_rules(input=self._text, model=self._model) + return chunks diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 0bcb732a0f5..7df8cc7452f 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -4045,7 +4045,9 @@ class CustomPricingLiteLLMParams(MirroredPricingParams): DEPLOYMENT_SCOPED_PRICING_FIELDS: Final[frozenset[str]] = frozenset({"off_peak_pricing"}) -DEPLOYMENT_SCOPED_CAPABILITY_FIELDS: Final[frozenset[str]] = frozenset({"supports_regex_lookaround"}) +DEPLOYMENT_SCOPED_CAPABILITY_FIELDS: Final[frozenset[str]] = frozenset( + {"supports_regex_lookaround", "supports_assistant_prefill"} +) SHARED_BACKEND_MODEL_INFO_FIELDS: Final[frozenset[str]] = ( frozenset(ModelInfoBase.__required_keys__ | ModelInfoBase.__optional_keys__) diff --git a/litellm/utils.py b/litellm/utils.py index 65408663f6f..ea36a68ad5a 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -3060,6 +3060,20 @@ def supports_prompt_cache_breakpoint(model: str, custom_llm_provider: str | None ) +def supports_assistant_prefill(model: str, custom_llm_provider: str | None = None) -> bool: + """ + Whether the model can continue a prefilled assistant message (Anthropic's + ``prefix: True`` trick and equivalents). Missing metadata reads as False, + so a mid-stream fallback continuation is only routed to a model known to + support it. + """ + return _supports_factory( + model=model, + custom_llm_provider=custom_llm_provider, + key="supports_assistant_prefill", + ) + + def supports_thinking_cache_preservation(model: str, custom_llm_provider: str | None = None) -> bool: return supports_factory( model=model, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index b744032fb6c..5fff760d85a 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -2773,7 +2773,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_pdf_input": true, @@ -2808,7 +2808,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_pdf_input": true, @@ -2843,7 +2843,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_pdf_input": true, @@ -2878,7 +2878,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_pdf_input": true, @@ -2913,7 +2913,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_pdf_input": true, @@ -2948,7 +2948,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_pdf_input": true, @@ -3743,7 +3743,7 @@ "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 1.5e-05, - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_pdf_input": true, @@ -15494,7 +15494,7 @@ }, "supports_adaptive_thinking": true, "supports_legacy_thinking": true, - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_pdf_input": true, @@ -20832,7 +20832,7 @@ "output_cost_per_token": 2.5000010000000002e-05, "output_dbu_cost_per_token": 0.000357143, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_function_calling": true, "supports_legacy_thinking": true, "supports_prompt_caching": true, @@ -21047,7 +21047,7 @@ "output_cost_per_token": 1.5000020000000002e-05, "output_dbu_cost_per_token": 0.000214286, "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_function_calling": true, "supports_legacy_thinking": true, "supports_prompt_caching": true, @@ -49176,7 +49176,7 @@ "max_tokens": 64000, "mode": "chat", "output_cost_per_token": 2.5e-05, - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_prompt_caching": true, @@ -51006,7 +51006,7 @@ "mode": "chat", "output_cost_per_token": 1.5e-05, "output_cost_per_token_batches": 7.5e-06, - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_pdf_input": true, @@ -51179,7 +51179,7 @@ "supported_regions": [ "global" ], - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -51200,7 +51200,7 @@ "supported_regions": [ "us-central1" ], - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -58822,7 +58822,7 @@ "mode": "chat", "output_cost_per_token": 1.5e-05, "output_cost_per_token_batches": 7.5e-06, - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_pdf_input": true, @@ -72334,7 +72334,7 @@ "max_output_tokens": 64000, "max_tokens": 64000, "mode": "chat", - "supports_assistant_prefill": true, + "supports_assistant_prefill": false, "supports_computer_use": true, "supports_function_calling": true, "supports_pdf_input": true, diff --git a/tests/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py index a3ed70f4a56..6762e4cd35c 100644 --- a/tests/code_coverage_tests/router_code_coverage.py +++ b/tests/code_coverage_tests/router_code_coverage.py @@ -114,6 +114,12 @@ ignored_function_names = [ "_deployment_that_just_failed", # Tested through the same-boundary hop test in test_router_order_fallback.py "_collect_ptu_share_teams", # Cached behind team_ptu_ceiling, tested through the live model list test in test_router.py "_resolve_team_ptu_ceiling", # Cached behind team_ptu_ceiling, tested through the live model list test in test_router.py + "_register_continuation_prefill_check", # Tested through the continuation tests in test_router.py + "_set_mid_stream_fallback_continuation", # Tested through the update_settings continuation test in test_router.py + "_is_flag_enabled", # Tested through the string-flag continuation tests in test_router.py + "_params_sent_to_deployment", # Tested through the configured-constraint continuation tests in test_router.py + "_names_deployment_that_cannot_continue", # Tested through the deployment-id continuation tests in test_router.py + "_set_retry_policy", # Tested through update_settings in test_router_retry_policy_update.py "_resolve_pre_call_check_input_tokens", # Tested through the token-count failure pre-call check tests in test_router.py "_params_neither_dropped_nor_allowed", # Tested through the allowed-or-dropped param pre-call check tests in test_router.py "_deployment_rejects_params", # Tested through the unsupported-param pre-call check tests in test_router.py diff --git a/tests/integration/streaming/test_mid_stream_fallback_continuation.py b/tests/integration/streaming/test_mid_stream_fallback_continuation.py new file mode 100644 index 00000000000..8ec52f26dde --- /dev/null +++ b/tests/integration/streaming/test_mid_stream_fallback_continuation.py @@ -0,0 +1,749 @@ +from __future__ import annotations + +import functools +import hashlib +import itertools +import json +import math +import operator +import threading +import uuid +from collections.abc import Callable, Generator, Iterable, Iterator, Mapping, Sequence +from contextlib import ExitStack, contextmanager +from dataclasses import dataclass +from pathlib import Path +from typing import Final + +import openai +import pytest +import yaml +from integration._support.client import JSON_OBJECT, Gateway, eventually, gateway_from_environment, object_value +from integration._support.database import read_rows, scratch_database +from integration._support.process import graceful_stop_seconds, owned_proxy +from integration._support.wire import Reply, Request, Wire, wire_server +from openai.types.chat import ChatCompletionChunk +from pydantic import JsonValue, TypeAdapter + +pytestmark: Final = pytest.mark.timeout(4 * graceful_stop_seconds() + 120) + +SETTING: Final = "enable_mid_stream_fallback_continuation" +PARTIAL: Final = ("Alpha beta", " gamma") +JOINED_PARTIAL: Final = "".join(PARTIAL) +RELAY_TEXT: Final = " delta" +BACKUP_TEXT: Final = " omega." +PLAIN_TEXT: Final = "plain answer" +TOOL_PREAMBLE: Final = "Let me check" +BACKUP_USAGE: Final = (40, 6) +MAX_TOKENS: Final = 200 +PRIMARY_PRICES: Final = (0.001, 0.002) +BACKUP_PRICES: Final = (0.003, 0.005) +SHARED_BACKEND: Final = "openai/gpt-4o-mini" +SEPARATE_BACKEND: Final = "openai/gpt-4o" +WEATHER_TOOL: Final[Mapping[str, JsonValue]] = { + "type": "function", + "function": { + "name": "get_weather", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}, + }, +} +TOOL_CALL_DELTA: Final[Mapping[str, JsonValue]] = { + "index": 0, + "id": "call_1", + "type": "function", + "function": {"name": "get_weather", "arguments": ""}, +} +NOTHING_SEEN: Final[tuple[ChatCompletionChunk, ...]] = () +ALLOW_EVERYTHING: Final = "def apply_guardrail(inputs, request_data, input_type):\n return allow()\n" +EVENT_BATCH: Final = TypeAdapter(list[dict[str, JsonValue]]) + + +@functools.cache +def cut_gate(marker: str, hop: str, attempt: int) -> threading.Event: + return threading.Event() + + +@functools.cache +def attempts(marker: str, hop: str) -> Iterator[int]: + return itertools.count() + + +def frame(identity: str, delta: Mapping[str, JsonValue], finish: str | None = None) -> bytes: + body: Final = { + "id": identity, + "object": "chat.completion.chunk", + "created": 1700000000, + "model": "gpt-4o-mini", + "choices": [{"index": 0, "delta": dict(delta), "finish_reason": finish}], + } + return b"data: " + json.dumps(body).encode() + b"\n\n" + + +def complete_stream(identity: str, text: str, usage: tuple[int, int]) -> Reply: + usage_frame: Final[dict[str, JsonValue]] = { + "id": identity, + "object": "chat.completion.chunk", + "created": 1700000000, + "model": "gpt-4o-mini", + "choices": [], + "usage": {"prompt_tokens": usage[0], "completion_tokens": usage[1], "total_tokens": sum(usage)}, + } + return Reply( + content_type="text/event-stream", + chunks=( + frame(identity, {"role": "assistant", "content": text}), + frame(identity, {}, finish="stop"), + b"data: " + json.dumps(usage_frame).encode() + b"\n\n", + b"data: [DONE]\n\n", + ), + ) + + +def cut_after_caller_sees(marker: str, hop: str, *frames: bytes) -> Reply: + return Reply( + content_type="text/event-stream", + chunks=(b"".join(frames), b"data: never sent\n\n"), + abort_after=1, + gate_after_first=cut_gate(marker, hop, next(attempts(marker, hop))), + ) + + +def first_user_text(request: Request) -> str: + messages: Final = JSON_OBJECT.validate_json(request.body)["messages"] + assert isinstance(messages, list), messages + first: Final = object_value(messages[0])["content"] + assert isinstance(first, str), first + return first + + +def marker_of(request: Request) -> str: + return first_user_text(request).split()[-1] + + +def is_chat(request: Request) -> bool: + return request.target.endswith("/chat/completions") + + +def model_list() -> Reply: + return Reply(body=b'{"object": "list", "data": []}') + + +def primary(request: Request) -> Reply: + if not is_chat(request): + return model_list() + scenario: Final = first_user_text(request).split()[0] + marker: Final = marker_of(request) + identity: Final = f"chatcmpl-{marker}-primary" + if scenario == "reject": + return Reply(status=500, body=b'{"error": {"message": "primary is down", "type": "server_error"}}') + if scenario == "tool-cut": + return cut_after_caller_sees( + marker, + "primary", + frame(identity, {"role": "assistant", "content": TOOL_PREAMBLE}), + frame(identity, {"tool_calls": [dict(TOOL_CALL_DELTA)]}), + ) + return cut_after_caller_sees( + marker, + "primary", + frame(identity, {"role": "assistant", "content": PARTIAL[0]}), + frame(identity, {"content": PARTIAL[1]}), + ) + + +def relay(request: Request) -> Reply: + if not is_chat(request): + return model_list() + marker: Final = marker_of(request) + return cut_after_caller_sees( + marker, "relay", frame(f"chatcmpl-{marker}-relay", {"role": "assistant", "content": RELAY_TEXT}) + ) + + +def backup(request: Request) -> Reply: + if not is_chat(request): + return model_list() + return complete_stream(f"chatcmpl-{marker_of(request)}-backup", BACKUP_TEXT, BACKUP_USAGE) + + +def no_prefill(request: Request) -> Reply: + if not is_chat(request): + return model_list() + return complete_stream(f"chatcmpl-{marker_of(request)}-plain", PLAIN_TEXT, (9, 2)) + + +def accept_events(request: Request) -> Reply: + return Reply() + + +def registered_guardrail(name: str, mode: str) -> dict[str, JsonValue]: + return { + "guardrail_name": name, + "litellm_params": { + "guardrail": "custom_code", + "mode": mode, + "default_on": False, + "custom_code": ALLOW_EVERYTHING, + }, + } + + +@dataclass(frozen=True, slots=True) +class Rig: + gateway: Gateway + database_url: str + primary: Wire + relay: Wire + backup: Wire + no_prefill: Wire + events: Wire + + +def joined_text(chunks: Iterable[ChatCompletionChunk]) -> str: + return "".join(choice.delta.content or "" for chunk in chunks for choice in chunk.choices) + + +def has_tool_call(chunks: Iterable[ChatCompletionChunk]) -> bool: + return any(choice.delta.tool_calls for chunk in chunks for choice in chunk.choices) + + +@dataclass(frozen=True, slots=True) +class Release: + gate: threading.Event + ready: Callable[[tuple[ChatCompletionChunk, ...]], bool] + + +def after_text(marker: str, hop: str, text: str, attempt: int = 0) -> Release: + return Release(cut_gate(marker, hop, attempt), lambda seen: joined_text(seen) == text) + + +@dataclass(frozen=True, slots=True) +class Streamed: + chunks: tuple[ChatCompletionChunk, ...] + error: openai.APIError | None + + @property + def text(self) -> str: + return joined_text(self.chunks) + + @property + def ids(self) -> frozenset[str]: + return frozenset(chunk.id for chunk in self.chunks) + + @property + def role_deltas(self) -> int: + return sum(1 for chunk in self.chunks for choice in chunk.choices if choice.delta.role is not None) + + @property + def finish_reasons(self) -> tuple[str, ...]: + return tuple(choice.finish_reason for chunk in self.chunks for choice in chunk.choices if choice.finish_reason) + + @property + def usage(self) -> tuple[int, int]: + reported: Final = tuple(chunk.usage for chunk in self.chunks if chunk.usage is not None) + assert reported, self.chunks + return (reported[-1].prompt_tokens, reported[-1].completion_tokens) + + +def deployment( + name: str, + wire: Wire, + prices: tuple[float, float], + *, + backend: str = SHARED_BACKEND, + prefill: bool | None = None, +) -> dict[str, JsonValue]: + return { + "model_name": name, + "litellm_params": { + "model": backend, + "api_base": wire.url + "/v1", + "api_key": "synthetic-continuation-key", + "input_cost_per_token": prices[0], + "output_cost_per_token": prices[1], + }, + **({"model_info": {"supports_assistant_prefill": prefill}} if prefill is not None else {}), + } + + +def write_config( + directory: Path, + model_entries: Sequence[Mapping[str, JsonValue]], + fallbacks: Sequence[Mapping[str, JsonValue]], + router_settings: Mapping[str, JsonValue], + guardrails: Sequence[Mapping[str, JsonValue]], +) -> Path: + base: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + config: Final = { + **base, + "model_list": [dict(entry) for entry in model_entries], + "litellm_settings": { + **object_value(base["litellm_settings"]), + "callbacks": ["generic_api"], + "DEFAULT_FLUSH_INTERVAL_SECONDS": 1, + }, + "guardrails": [dict(entry) for entry in guardrails], + "router_settings": { + **object_value(base["router_settings"]), + "num_retries": 0, + "fallbacks": [dict(entry) for entry in fallbacks], + **router_settings, + }, + } + path: Final = directory / f"continuation-{uuid.uuid4().hex[:8]}.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +@contextmanager +def booted_rig( + directory: Path, + fallbacks: Sequence[Mapping[str, JsonValue]], + router_settings: Mapping[str, JsonValue], + workers: int, + guardrails: Sequence[Mapping[str, JsonValue]] = (), +) -> Generator[Rig]: + with ExitStack() as stack: + shared: Final = stack.enter_context(gateway_from_environment()) + database_url: Final = stack.enter_context(scratch_database()) + wires: Final = tuple( + stack.enter_context(wire_server(respond)) for respond in (primary, relay, backup, no_prefill, accept_events) + ) + primary_wire, relay_wire, backup_wire, no_prefill_wire, events_wire = wires + config: Final = write_config( + directory, + ( + *( + deployment(name, primary_wire, PRIMARY_PRICES) + for name in ("primary", "primary-skip", "primary-none", "primary-relay", "primary-shared") + ), + deployment("relay", relay_wire, PRIMARY_PRICES, prefill=True), + deployment("backup", backup_wire, BACKUP_PRICES, prefill=True), + deployment("no-prefill", no_prefill_wire, PRIMARY_PRICES, backend=SEPARATE_BACKEND), + deployment("no-prefill-shared", no_prefill_wire, PRIMARY_PRICES), + ), + fallbacks, + router_settings, + guardrails, + ) + gateway: Final = stack.enter_context( + owned_proxy( + shared, + directory, + {"DATABASE_URL": database_url, "GENERIC_LOGGER_ENDPOINT": events_wire.url}, + config=config, + remove_environment=("DATABASE_URL_READ_REPLICA",), + workers=workers, + ) + ) + yield Rig(gateway, database_url, primary_wire, relay_wire, backup_wire, no_prefill_wire, events_wire) + + +@pytest.fixture(scope="module") +def configured(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + fallbacks: Final[tuple[Mapping[str, JsonValue], ...]] = ( + {"primary": ["backup"]}, + {"primary-skip": ["no-prefill", "backup"]}, + {"primary-none": ["no-prefill"]}, + {"primary-relay": ["relay"]}, + {"relay": ["backup"]}, + {"primary-shared": ["no-prefill-shared"]}, + ) + with booted_rig( + tmp_path_factory.mktemp("continuation-configured"), + fallbacks, + {SETTING: True}, + 2, + (registered_guardrail("input-check", "pre_call"),), + ) as rig: + yield rig + + +@pytest.fixture(scope="module") +def output_guarded(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + with booted_rig( + tmp_path_factory.mktemp("continuation-output-guarded"), + ({"primary": ["backup"]},), + {SETTING: True}, + 2, + (registered_guardrail("output-check", "post_call"),), + ) as rig: + yield rig + + +@pytest.fixture(scope="module") +def unset(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + with booted_rig( + tmp_path_factory.mktemp("continuation-unset"), ({"primary": ["no-prefill", "backup"]},), {}, 1 + ) as rig: + yield rig + + +def new_key(rig: Rig) -> str: + key: Final = rig.gateway.post("/key/generate", {})["key"] + assert isinstance(key, str), key + return key + + +def appended(seen: tuple[ChatCompletionChunk, ...], chunk: ChatCompletionChunk) -> tuple[ChatCompletionChunk, ...]: + return (*seen, chunk) + + +def observed( + response: Iterable[ChatCompletionChunk], releases: Sequence[Release] +) -> Iterator[ChatCompletionChunk | openai.APIError]: + try: + for seen in itertools.accumulate(response, appended, initial=NOTHING_SEEN): + for release in releases: + if release.ready(seen): + release.gate.set() + yield from seen[-1:] + except openai.APIError as error: + yield error + + +def stream( + rig: Rig, + key: str, + model: str, + prompt: str, + releases: Sequence[Release] = (), + extra: Mapping[str, JsonValue] | None = None, +) -> Streamed: + base_url: Final = f"{rig.gateway.client.base_url}".rstrip("/") + "/v1" + try: + with openai.OpenAI(base_url=base_url, api_key=key, max_retries=0, timeout=60) as client: + response: Final = client.chat.completions.create( + model=model, + messages=[{"role": "user", "content": prompt}], + stream=True, + stream_options={"include_usage": True}, + max_tokens=MAX_TOKENS, + extra_body=dict(extra or {}), + ) + outcome: Final = tuple(observed(response, releases)) + finally: + for release in releases: + release.gate.set() + errors: Final = tuple(item for item in outcome if isinstance(item, openai.APIError)) + return Streamed( + tuple(item for item in outcome if isinstance(item, ChatCompletionChunk)), errors[0] if errors else None + ) + + +def chat_bodies(wire: Wire, marker: str) -> tuple[dict[str, JsonValue], ...]: + return tuple( + JSON_OBJECT.validate_json(request.body) + for request in wire.drain() + if is_chat(request) and marker_of(request) == marker + ) + + +def spend_rows(rig: Rig, key: str) -> list[dict[str, JsonValue]]: + return eventually( + lambda: read_rows( + 'SELECT request_id, status, spend, prompt_tokens, completion_tokens, metadata FROM "LiteLLM_SpendLogs" ' + "WHERE api_key=%s", + (hashlib.sha256(key.encode()).hexdigest(),), + database_url=rig.database_url, + ), + lambda rows: len(rows) >= 1, + seconds=70, + ) + + +def amount(value: JsonValue) -> float: + assert isinstance(value, (int, float)), value + return float(value) + + +def key_spend(rig: Rig, key: str) -> float: + return amount(object_value(rig.gateway.get("/key/info", {"key": key})["info"])["spend"]) + + +def close(observed_value: float, expected: float) -> bool: + return math.isclose(observed_value, expected, rel_tol=1e-9, abs_tol=1e-12) + + +def cost(tokens: tuple[int, int], prices: tuple[float, float]) -> float: + return tokens[0] * prices[0] + tokens[1] * prices[1] + + +def prefill_turn(text: str) -> dict[str, JsonValue]: + return {"role": "assistant", "content": text, "prefix": True} + + +@dataclass(frozen=True, slots=True) +class Billed: + rows: list[dict[str, JsonValue]] + partial: tuple[int, int] + + +def billed_once(rig: Rig, key: str, received: Streamed) -> Billed: + rows: Final = spend_rows(rig, key) + assert [row["status"] for row in rows] == ["success"], rows + row_tokens: Final = (rows[0]["prompt_tokens"], rows[0]["completion_tokens"]) + assert isinstance(row_tokens[0], int) and isinstance(row_tokens[1], int), rows + assert received.usage == row_tokens, (received.usage, row_tokens) + partial: Final = (row_tokens[0] - BACKUP_USAGE[0], row_tokens[1] - BACKUP_USAGE[1]) + assert partial[0] > 0 and partial[1] > 0, (row_tokens, BACKUP_USAGE) + return Billed(rows, partial) + + +def assert_spend_prices_every_token(rig: Rig, key: str, billed: Billed) -> None: + partial_cost: Final = cost(billed.partial, PRIMARY_PRICES) + breakdown: Final = object_value(object_value(billed.rows[0]["metadata"])["cost_breakdown"]) + spend: Final = amount(billed.rows[0]["spend"]) + reported_partial_cost: Final = amount(object_value(breakdown["additional_costs"])["mid_stream_partial_cost"]) + assert close(spend, cost(BACKUP_USAGE, BACKUP_PRICES) + partial_cost) and close( + reported_partial_cost, partial_cost + ), f"partial={billed.partial} spend={spend} breakdown={json.dumps(breakdown)}" + eventually(lambda: key_spend(rig, key), lambda value: close(value, spend), seconds=70) + assert spend_rows(rig, key) == billed.rows + + +def test_text_cut_continues_on_the_prefill_fallback_and_bills_both_calls_once(configured: Rig) -> None: + marker: Final = uuid.uuid4().hex + key: Final = new_key(configured) + prompt: Final = f"text-cut {marker}" + received: Final = stream(configured, key, "primary", prompt, (after_text(marker, "primary", JOINED_PARTIAL),)) + assert received.error is None, received.error + assert received.text == JOINED_PARTIAL + BACKUP_TEXT, received.chunks + assert received.ids == {f"chatcmpl-{marker}-primary"}, received.ids + assert received.role_deltas == 1, received.chunks + assert received.finish_reasons == ("stop",), received.chunks + assert len(chat_bodies(configured.primary, marker)) == 1 + continued: Final = chat_bodies(configured.backup, marker) + assert len(continued) == 1, continued + assert continued[0]["messages"] == [{"role": "user", "content": prompt}, prefill_turn(JOINED_PARTIAL)] + assert "_mid_stream_continuation" not in continued[0], continued[0] + billed: Final = billed_once(configured, key, received) + assert continued[0]["max_tokens"] == MAX_TOKENS - billed.partial[1], (continued[0], billed.partial) + assert_spend_prices_every_token(configured, key, billed) + + +def test_spend_row_of_a_continued_stream_carries_the_id_the_caller_received(configured: Rig) -> None: + marker: Final = uuid.uuid4().hex + key: Final = new_key(configured) + received: Final = stream( + configured, key, "primary", f"text-cut {marker}", (after_text(marker, "primary", JOINED_PARTIAL),) + ) + assert received.error is None, received.error + assert [row["request_id"] for row in spend_rows(configured, key)] == [received.chunks[0].id] + + +def test_tool_call_before_the_cut_keeps_the_original_error(configured: Rig) -> None: + marker: Final = uuid.uuid4().hex + received: Final = stream( + configured, + new_key(configured), + "primary", + f"tool-cut {marker}", + (Release(cut_gate(marker, "primary", 0), has_tool_call),), + {"tools": [dict(WEATHER_TOOL)]}, + ) + assert received.error is not None, received.chunks + assert (received.text, has_tool_call(received.chunks)) == (TOOL_PREAMBLE, True), received.chunks + assert len(chat_bodies(configured.primary, marker)) == 1 + assert chat_bodies(configured.backup, marker) == () + + +def test_continuation_skips_a_fallback_without_prefill_support(configured: Rig) -> None: + marker: Final = uuid.uuid4().hex + prompt: Final = f"text-cut {marker}" + received: Final = stream( + configured, new_key(configured), "primary-skip", prompt, (after_text(marker, "primary", JOINED_PARTIAL),) + ) + assert received.error is None, received.error + assert received.text == JOINED_PARTIAL + BACKUP_TEXT, received.chunks + assert chat_bodies(configured.no_prefill, marker) == () + assert [body["messages"] for body in chat_bodies(configured.backup, marker)] == [ + [{"role": "user", "content": prompt}, prefill_turn(JOINED_PARTIAL)] + ] + + +def test_no_prefill_capable_fallback_keeps_the_original_error(configured: Rig) -> None: + marker: Final = uuid.uuid4().hex + received: Final = stream( + configured, + new_key(configured), + "primary-none", + f"text-cut {marker}", + (after_text(marker, "primary", JOINED_PARTIAL),), + ) + assert received.error is not None, received.chunks + assert "No deployments available" not in str(received.error), received.error + assert received.text == JOINED_PARTIAL, received.chunks + assert len(chat_bodies(configured.primary, marker)) == 1 + assert chat_bodies(configured.no_prefill, marker) == () + + +def test_undeclared_fallback_sharing_a_backend_with_a_prefill_deployment_keeps_the_original_error( + configured: Rig, +) -> None: + marker: Final = uuid.uuid4().hex + received: Final = stream( + configured, + new_key(configured), + "primary-shared", + f"text-cut {marker}", + (after_text(marker, "primary", JOINED_PARTIAL),), + ) + assert (received.error is not None, received.text) == (True, JOINED_PARTIAL), received.chunks + assert chat_bodies(configured.no_prefill, marker) == () + + +def test_failure_before_content_still_falls_back_to_a_fallback_without_prefill_support(configured: Rig) -> None: + marker: Final = uuid.uuid4().hex + prompt: Final = f"reject {marker}" + received: Final = stream(configured, new_key(configured), "primary-none", prompt) + assert received.error is None, received.error + assert received.text == PLAIN_TEXT, received.chunks + assert [body["messages"] for body in chat_bodies(configured.no_prefill, marker)] == [ + [{"role": "user", "content": prompt}] + ] + + +def stream_with_two_cuts(rig: Rig, key: str, marker: str) -> Streamed: + return stream( + rig, + key, + "primary-relay", + f"text-cut {marker}", + (after_text(marker, "primary", JOINED_PARTIAL), after_text(marker, "relay", JOINED_PARTIAL + RELAY_TEXT)), + ) + + +def test_second_cut_continues_with_both_partials_in_one_prefill(configured: Rig) -> None: + marker: Final = uuid.uuid4().hex + key: Final = new_key(configured) + received: Final = stream_with_two_cuts(configured, key, marker) + assert received.error is None, received.error + assert received.text == JOINED_PARTIAL + RELAY_TEXT + BACKUP_TEXT, received.chunks + assert received.role_deltas == 1, received.chunks + user_turn: Final = {"role": "user", "content": f"text-cut {marker}"} + relayed: Final = chat_bodies(configured.relay, marker) + continued: Final = chat_bodies(configured.backup, marker) + assert [body["messages"] for body in relayed] == [[user_turn, prefill_turn(JOINED_PARTIAL)]] + assert [body["messages"] for body in continued] == [[user_turn, prefill_turn(JOINED_PARTIAL + RELAY_TEXT)]] + partial: Final = billed_once(configured, key, received).partial + relay_ceiling: Final = relayed[0]["max_tokens"] + assert isinstance(relay_ceiling, int), relayed[0] + assert 0 < MAX_TOKENS - relay_ceiling < partial[1], (relayed[0], partial) + assert continued[0]["max_tokens"] == MAX_TOKENS - partial[1], (continued[0], partial) + + +def test_second_cut_bills_every_call_on_one_row(configured: Rig) -> None: + key: Final = new_key(configured) + received: Final = stream_with_two_cuts(configured, key, uuid.uuid4().hex) + assert received.error is None, received.error + assert_spend_prices_every_token(configured, key, billed_once(configured, key, received)) + + +def test_repeating_a_continued_request_never_returns_only_the_continuation(configured: Rig) -> None: + marker: Final = uuid.uuid4().hex + key: Final = new_key(configured) + prompt: Final = f"text-cut {marker}" + answers: Final = tuple( + stream(configured, key, "primary", prompt, (after_text(marker, "primary", JOINED_PARTIAL, attempt),)) + for attempt in range(2) + ) + assert tuple((answer.error, answer.text) for answer in answers) == ((None, JOINED_PARTIAL + BACKUP_TEXT),) * 2 + + +def test_without_the_setting_a_text_cut_keeps_the_original_error(unset: Rig) -> None: + marker: Final = uuid.uuid4().hex + received: Final = stream( + unset, new_key(unset), "primary", f"text-cut {marker}", (after_text(marker, "primary", JOINED_PARTIAL),) + ) + assert received.error is not None, received.chunks + assert received.text == JOINED_PARTIAL, received.chunks + assert len(chat_bodies(unset.primary, marker)) == 1 + assert chat_bodies(unset.no_prefill, marker) == () + assert chat_bodies(unset.backup, marker) == () + + +def router_setting(rig: Rig) -> JsonValue: + return object_value(rig.gateway.get("/router/settings")["current_values"]).get(SETTING) + + +def test_turning_the_setting_on_at_runtime_continues_on_a_prefill_fallback(unset: Rig) -> None: + rig: Final = unset + rig.gateway.post("/config/update", {"router_settings": {SETTING: True}}) + try: + assert router_setting(rig) is True + marker: Final = uuid.uuid4().hex + prompt: Final = f"text-cut {marker}" + received: Final = stream(rig, new_key(rig), "primary", prompt, (after_text(marker, "primary", JOINED_PARTIAL),)) + assert received.error is None, received.error + assert received.text == JOINED_PARTIAL + BACKUP_TEXT, received.chunks + assert chat_bodies(rig.no_prefill, marker) == () + assert [body["messages"] for body in chat_bodies(rig.backup, marker)] == [ + [{"role": "user", "content": prompt}, prefill_turn(JOINED_PARTIAL)] + ] + finally: + rig.gateway.post("/config/update", {"router_settings": {SETTING: False}}) + + +def test_a_registered_output_guardrail_keeps_the_original_error(output_guarded: Rig) -> None: + marker: Final = uuid.uuid4().hex + received: Final = stream( + output_guarded, + new_key(output_guarded), + "primary", + f"text-cut {marker}", + (after_text(marker, "primary", JOINED_PARTIAL),), + ) + assert (received.error is not None, received.text) == (True, JOINED_PARTIAL), received.chunks + assert len(chat_bodies(output_guarded.primary, marker)) == 1 + assert chat_bodies(output_guarded.backup, marker) == () + + +def test_a_request_fallback_bringing_its_own_messages_keeps_the_original_error(configured: Rig) -> None: + marker: Final = uuid.uuid4().hex + received: Final = stream( + configured, + new_key(configured), + "primary", + f"text-cut {marker}", + (after_text(marker, "primary", JOINED_PARTIAL),), + {"fallbacks": [{"model": "backup", "messages": [{"role": "user", "content": f"start over {marker}"}]}]}, + ) + assert (received.error is not None, received.text) == (True, JOINED_PARTIAL), received.chunks + assert len(chat_bodies(configured.primary, marker)) == 1 + assert chat_bodies(configured.backup, marker) == () + + +def carries_tag(event: Mapping[str, JsonValue], marker: str) -> bool: + tags: Final = event.get("request_tags") + return isinstance(tags, list) and marker in tags + + +def tagged_events(batches: Iterable[Request], marker: str) -> tuple[dict[str, JsonValue], ...]: + events: Final = itertools.chain.from_iterable(EVENT_BATCH.validate_json(batch.body) for batch in batches) + return tuple(event for event in events if carries_tag(event, marker)) + + +def logged_events(sink: Wire, marker: str) -> Iterator[tuple[dict[str, JsonValue], ...]]: + return itertools.accumulate((tagged_events(sink.drain(), marker) for _ in itertools.count()), operator.add) + + +def test_the_cut_call_logs_its_failure_under_its_own_deployment_and_price(configured: Rig) -> None: + marker: Final = uuid.uuid4().hex + received: Final = stream( + configured, + new_key(configured), + "primary", + f"text-cut {marker}", + (after_text(marker, "primary", JOINED_PARTIAL),), + {"metadata": {"tags": [marker]}}, + ) + assert received.error is None, received.error + polls: Final = logged_events(configured.events, marker) + events: Final = eventually(lambda: next(polls), lambda logged: len(logged) >= 2, seconds=30) + assert sorted(str(event["status"]) for event in events) == ["failure", "success"], events + failure: Final = next(event for event in events if event["status"] == "failure") + assert failure["model_group"] == "primary", failure + assert str(failure["api_base"]).startswith(configured.primary.url), failure + tokens: Final = (failure["prompt_tokens"], failure["completion_tokens"]) + assert isinstance(tokens[0], int) and isinstance(tokens[1], int) and tokens[0] > 0 and tokens[1] > 0, failure + assert close(amount(failure["response_cost"]), cost((tokens[0], tokens[1]), PRIMARY_PRICES)), failure diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py index c878f0e149d..bf977325a85 100644 --- a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py +++ b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py @@ -642,6 +642,64 @@ async def test_bedrock_converse_replays_reasoning_to_claude_only_as_base64_of_cl assert assistant_content == [*sent, {"text": "hello"}] +@pytest.mark.asyncio +@pytest.mark.parametrize("use_async", [False, True]) +@pytest.mark.parametrize( + "model, final_text", + [("us.anthropic.claude-sonnet-4-6", "Leaves fall"), ("us.amazon.nova-pro-v1:0", "Leaves fall \n")], + ids=["claude", "nova"], +) +async def test_bedrock_converse_strips_trailing_whitespace_from_a_final_assistant_turn_for_claude_only( + use_async: bool, model: str, final_text: str +): + """Claude on Converse 400s on a final assistant turn ending in whitespace, which a cut stream's prefill often has.""" + messages: Final = [ + {"role": "user", "content": "Write a haiku."}, + {"role": "assistant", "content": "Sure. "}, + {"role": "user", "content": "About autumn."}, + {"role": "assistant", "content": "Leaves fall \n"}, + ] + + result: Final = ( + await BedrockConverseMessagesProcessor._bedrock_converse_messages_pt_async( + messages=messages, model=model, llm_provider="bedrock" + ) + if use_async + else _bedrock_converse_messages_pt(messages=messages, model=model, llm_provider="bedrock") + ) + + assert [m["content"] for m in result if m["role"] == "assistant"] == [ + [{"text": "Sure. "}], + [{"text": final_text}], + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_async", [False, True]) +async def test_bedrock_converse_keeps_whitespace_between_text_blocks_of_a_final_claude_assistant_turn( + use_async: bool, +): + messages: Final = [ + {"role": "user", "content": "Write a haiku."}, + { + "role": "assistant", + "content": [{"type": "text", "text": "Leaves fall "}, {"type": "text", "text": "softly \n"}], + }, + ] + + result: Final = ( + await BedrockConverseMessagesProcessor._bedrock_converse_messages_pt_async( + messages=messages, model="us.anthropic.claude-sonnet-4-6", llm_provider="bedrock" + ) + if use_async + else _bedrock_converse_messages_pt( + messages=messages, model="us.anthropic.claude-sonnet-4-6", llm_provider="bedrock" + ) + ) + + assert result[-1]["content"] == [{"text": "Leaves fall "}, {"text": "softly"}] + + def test_translate_thinking_blocks_keeps_dropping_redacted_data_when_called_the_old_way(): """Outside callers that only pass the thinking blocks keep the old behavior, with no redactedContent sent.""" assert BedrockConverseMessagesProcessor.translate_thinking_blocks_to_reasoning_content_blocks( diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index 4347a6639bd..fa649b48a51 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -1019,6 +1019,87 @@ async def test_bridged_stream_mid_stream_fallback_error_is_rebuilt_around_the_pr assert excinfo.value.generated_content == "" +@pytest.mark.asyncio +@pytest.mark.parametrize("merge_reasoning", [True, False], ids=["merged_into_content", "kept_as_reasoning"]) +async def test_mid_stream_break_after_reasoning_disqualifies_continuation(logging_obj: Logging, merge_reasoning: bool): + import httpx + + from litellm.exceptions import MidStreamFallbackError + + class _ReasoningThenCut: + def __init__(self) -> None: + self.chunks = iter( + ( + ModelResponseStream(choices=[{"index": 0, "delta": {"reasoning_content": "Let me think"}}]), + ModelResponseStream(choices=[{"index": 0, "delta": {"content": "The answer"}}]), + ) + ) + + def __aiter__(self) -> "_ReasoningThenCut": + return self + + async def __anext__(self) -> ModelResponseStream: + next_chunk: Final = next(self.chunks, None) + if next_chunk is None: + raise httpx.ReadTimeout("connection dropped") + return next_chunk + + response: Final = CustomStreamWrapper( + completion_stream=_ReasoningThenCut(), model="deepseek-reasoner", logging_obj=logging_obj + ) + response.merge_reasoning_content_in_choices = merge_reasoning + + with pytest.raises(MidStreamFallbackError) as excinfo: + async for _ in response: + pass + + assert excinfo.value.is_pre_first_chunk is False + assert excinfo.value.emitted_disqualifying_content is True + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "outer_delta,inner_flag,expected", + [ + ({"tool_calls": [{"index": 0}]}, False, True), + ({"content": "Let me check."}, True, True), + ({"content": "Let me check."}, False, False), + ], + ids=["bridge_wrapper_sent_a_tool_call", "inner_stream_sent_a_tool_call", "plain_text_only"], +) +async def test_bridged_stream_rebuild_keeps_the_continuation_disqualifier( + logging_obj: Logging, outer_delta: dict[str, object], inner_flag: bool, expected: bool +): + from litellm.exceptions import InternalServerError, MidStreamFallbackError + + provider_error: Final = InternalServerError(message="upstream 500", llm_provider="azure", model="gpt-6.1-sol") + inner_error: Final = MidStreamFallbackError( + message=str(provider_error), + model="gpt-6.1-sol", + llm_provider="azure", + original_exception=provider_error, + is_pre_first_chunk=False, + emitted_disqualifying_content=inner_flag, + ) + + async def _raise_inner_error(**kwargs: object) -> NoReturn: + raise inner_error + + response: Final = CustomStreamWrapper( + completion_stream=None, + model="gpt-6.1-sol", + logging_obj=logging_obj, + custom_llm_provider="azure", + make_call=_raise_inner_error, + ) + response._accumulate_streamed_delta(outer_delta) + + with pytest.raises(MidStreamFallbackError) as excinfo: + await response.__anext__() + + assert excinfo.value.emitted_disqualifying_content is expected + + @pytest.mark.asyncio @pytest.mark.parametrize("sync_mode", [True, False]) async def test_streaming_rate_limit_midstream_fallback_keeps_provider_headers(logging_obj: Logging, sync_mode: bool): @@ -6230,6 +6311,59 @@ class TestStableStreamingResponseId: assert wrapper.model_response_creator().id == "chatcmpl-from-provider" +class TestContinuationDisqualifiers: + """The mid-stream continuation eligibility hinges on classifying which deltas + carry output a text-only prefill cannot represent.""" + + @pytest.mark.parametrize( + "field", + [ + "tool_calls", + "function_call", + "reasoning_content", + "thinking_blocks", + "reasoning_items", + "audio", + "images", + "annotations", + ], + ) + def test_disqualifying_fields_flagged(self, field): + # reasoning_content included: it reaches the caller as visible reasoning a + # text-only prefill cannot carry, so a reasoning fallback would re-derive it + assert CustomStreamWrapper._delta_disqualifies_continuation({field: [{"x": 1}]}) is True + + @pytest.mark.parametrize("delta", [{"content": "hi"}, {}, {"role": "assistant"}]) + def test_plain_text_not_flagged(self, delta): + assert CustomStreamWrapper._delta_disqualifies_continuation(delta) is False + + def test_non_mapping_delta_is_safe(self): + assert CustomStreamWrapper._delta_disqualifies_continuation(object()) is False + + def test_accumulate_grows_text_and_leaves_flag_clear_for_plain_text(self): + wrapper = object.__new__(CustomStreamWrapper) + wrapper.response_uptil_now = "" + wrapper._emitted_disqualifying_content = False + + wrapper._accumulate_streamed_delta({"content": "Hel"}) + wrapper._accumulate_streamed_delta({"content": "lo"}) + + assert wrapper.response_uptil_now == "Hello" + assert wrapper._emitted_disqualifying_content is False + + def test_accumulate_latches_flag_on_disqualifying_delta(self): + wrapper = object.__new__(CustomStreamWrapper) + wrapper.response_uptil_now = "" + wrapper._emitted_disqualifying_content = False + + wrapper._accumulate_streamed_delta({"content": "Hi"}) + wrapper._accumulate_streamed_delta({"tool_calls": [{"index": 0}]}) + wrapper._accumulate_streamed_delta({"content": "there"}) + + assert wrapper.response_uptil_now == "Hithere" + assert wrapper._emitted_disqualifying_content is True + + @pytest.mark.asyncio async def test_async_stream_without_usage_counts_tokens_off_the_event_loop(): from tests.large_text import text diff --git a/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py b/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py index 650517c92ca..1e353a7a82a 100644 --- a/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py +++ b/tests/unit/llms/databricks/chat/test_databricks_chat_transformation.py @@ -461,6 +461,21 @@ def test_transform_messages_sanitizes_empty_content(): assert result[1]["content"] == "Hi" +def test_transform_messages_drops_the_prefill_prefix_flag(): + """Databricks rejects unknown message keys, so a continuation prefill reaches Claude without ``prefix``.""" + messages = [ + {"role": "user", "content": "How do transformers work?"}, + {"role": "assistant", "content": "Transformers use", "prefix": True}, + ] + + result = DatabricksConfig().transform_messages( + messages=messages, model="databricks-claude-sonnet-4-6", is_async=False + ) + + assert result[-1] == {"role": "assistant", "content": "Transformers use"} + assert messages[-1]["prefix"] is True + + def test_transform_request_preserves_unity_model_service_name(): config = DatabricksConfig() result = config.transform_request( diff --git a/tests/unit/router_utils/pre_call_checks/test_continuation_prefill_check.py b/tests/unit/router_utils/pre_call_checks/test_continuation_prefill_check.py new file mode 100644 index 00000000000..8fa1c799863 --- /dev/null +++ b/tests/unit/router_utils/pre_call_checks/test_continuation_prefill_check.py @@ -0,0 +1,288 @@ +import pytest + +import litellm +from litellm.router_utils.pre_call_checks.continuation_prefill_check import ( + MID_STREAM_CONTINUATION_KWARG, + MID_STREAM_CONTINUATION_MARKER, + ContinuationOutput, + ContinuationPrefillDeploymentCheck, + continuation_output_ceilings, + deployment_supports_prefill, +) + +PREFILL_MODEL = "anthropic/prefill-capable-test-model" +NON_PREFILL_MODEL = "openai/prefill-unknown-test-model" +UNMAPPED_MODEL = "openai/unmapped-test-model" + + +@pytest.fixture(autouse=True) +def _synthetic_cost_map_entries() -> None: + litellm.register_model( + { + PREFILL_MODEL: {"litellm_provider": "anthropic", "mode": "chat", "supports_assistant_prefill": True}, + NON_PREFILL_MODEL: {"litellm_provider": "openai", "mode": "chat"}, + }, + persist_across_reloads=False, + ) + litellm.get_model_info.cache_clear() + + +def _deployment(model: str, dep_id: str) -> dict[str, object]: + return {"litellm_params": {"model": model}, "model_info": {"id": dep_id}} + + +def test_deployment_supports_prefill_reads_capability(): + assert deployment_supports_prefill(_deployment(PREFILL_MODEL, "a")) is True + assert deployment_supports_prefill(_deployment(NON_PREFILL_MODEL, "b")) is False + assert deployment_supports_prefill({"litellm_params": {"model": PREFILL_MODEL}}) is True + assert deployment_supports_prefill({"litellm_params": {"model": PREFILL_MODEL}, "model_info": "bogus"}) is True + + +def test_deployment_model_info_override_wins_over_cost_map(): + # model_info True opts in a model that is not in the cost map + assert ( + deployment_supports_prefill( + {"litellm_params": {"model": "vendor/custom-model"}, "model_info": {"supports_assistant_prefill": True}} + ) + is True + ) + # model_info False opts out a model the cost map would otherwise allow + assert ( + deployment_supports_prefill( + {"litellm_params": {"model": PREFILL_MODEL}, "model_info": {"supports_assistant_prefill": False}} + ) + is False + ) + # model_info without the key falls through to the cost map + assert deployment_supports_prefill(_deployment(PREFILL_MODEL, "z")) is True + + +@pytest.mark.parametrize( + "deployment,expected", + [ + ({"litellm_params": {"model": "bedrock/arn:custom-import", "base_model": PREFILL_MODEL}}, True), + ({"litellm_params": {"model": "azure/my-deployment"}, "model_info": {"base_model": PREFILL_MODEL}}, True), + ( + { + "litellm_params": {"model": "azure/my-deployment", "base_model": NON_PREFILL_MODEL}, + "model_info": {"base_model": PREFILL_MODEL}, + }, + True, + ), + ( + { + "litellm_params": {"model": PREFILL_MODEL, "base_model": NON_PREFILL_MODEL}, + "model_info": {"base_model": NON_PREFILL_MODEL}, + }, + False, + ), + ], + ids=[ + "litellm_params_base_model", + "model_info_base_model", + "model_info_base_model_wins", + "base_model_wins_over_model", + ], +) +def test_deployment_supports_prefill_reads_the_cost_map_key_of_the_base_model(deployment, expected): + assert deployment_supports_prefill(deployment) is expected + + +def test_deployment_supports_prefill_rejects_malformed_deployments(): + assert deployment_supports_prefill({}) is False + assert deployment_supports_prefill({"litellm_params": {}}) is False + assert deployment_supports_prefill("not-a-dict") is False + + +@pytest.mark.asyncio +async def test_filter_is_noop_without_continuation_marker(): + """A normal (non-continuation) request must be passed through untouched, even + if some deployments cannot prefill.""" + check = ContinuationPrefillDeploymentCheck() + deployments = [_deployment(PREFILL_MODEL, "a"), _deployment(NON_PREFILL_MODEL, "b")] + + for request_kwargs in ({}, None, {MID_STREAM_CONTINUATION_KWARG: False}): + result = await check.async_filter_deployments( + model="group", healthy_deployments=deployments, messages=None, request_kwargs=request_kwargs + ) + assert result == deployments + + +@pytest.mark.asyncio +async def test_filter_ignores_forged_client_flag(): + """A client cannot steer routing: a plain truthy value under the marker key + (which the proxy could forward from the request body) is not the internal + sentinel, so the filter leaves the deployment list untouched.""" + check = ContinuationPrefillDeploymentCheck() + deployments = [_deployment(PREFILL_MODEL, "a"), _deployment(NON_PREFILL_MODEL, "b")] + + for forged in (True, "true", 1, {"any": "json"}): + result = await check.async_filter_deployments( + model="group", + healthy_deployments=deployments, + messages=None, + request_kwargs={MID_STREAM_CONTINUATION_KWARG: forged}, + ) + assert result == deployments + + +@pytest.mark.asyncio +async def test_filter_keeps_only_prefill_capable_on_continuation(): + check = ContinuationPrefillDeploymentCheck() + deployments = [_deployment(PREFILL_MODEL, "a"), _deployment(NON_PREFILL_MODEL, "b")] + + result = await check.async_filter_deployments( + model="group", + healthy_deployments=deployments, + messages=None, + request_kwargs={MID_STREAM_CONTINUATION_KWARG: MID_STREAM_CONTINUATION_MARKER}, + ) + assert [d["model_info"]["id"] for d in result] == ["a"] + + +@pytest.mark.asyncio +async def test_filter_empties_group_when_no_prefill_capable_deployment(): + """No prefill-capable deployment -> empty result, so the router advances the + fallback chain and ultimately surfaces the original error.""" + check = ContinuationPrefillDeploymentCheck() + deployments = [_deployment(NON_PREFILL_MODEL, "b"), _deployment(UNMAPPED_MODEL, "c")] + + result = await check.async_filter_deployments( + model="group", + healthy_deployments=deployments, + messages=None, + request_kwargs={MID_STREAM_CONTINUATION_KWARG: MID_STREAM_CONTINUATION_MARKER}, + ) + assert result == [] + + +@pytest.mark.asyncio +async def test_filter_drops_deployments_whose_own_params_constrain_the_continuation(): + check = ContinuationPrefillDeploymentCheck() + plain = _deployment(PREFILL_MODEL, "plain") + constrained = [ + {"litellm_params": {"model": PREFILL_MODEL, **params}, "model_info": {"id": dep_id}} + for dep_id, params in ( + ("json", {"response_format": {"type": "json_object"}}), + ("two_choices", {"extra_body": {"n": 2}}), + ("forced_tool", {"tool_choice": {"type": "function", "function": {"name": "f"}}}), + ("merged_reasoning", {"merge_reasoning_content_in_choices": True}), + ) + ] + + result = await check.async_filter_deployments( + model="group", + healthy_deployments=[plain, *constrained], + messages=None, + request_kwargs={MID_STREAM_CONTINUATION_KWARG: MID_STREAM_CONTINUATION_MARKER}, + ) + assert [d["model_info"]["id"] for d in result] == ["plain"] + + +def _text_chunk(text: str) -> litellm.ModelResponseStream: + return litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": text}}]) + + +@pytest.mark.parametrize( + "prefill_content,fallback_texts,expected_texts", + [ + ("Hello", ("", "Hel", "lo wor", "ld"), ("", "", " wor", "ld")), + ("Hello", ("Hello",), ("",)), + ("Hello", (" world",), (" world",)), + ("Hello", ("Hel", "p me"), ("Hel", "p me")), + ("Hello", ("Hel",), ("Hel",)), + ([{"type": "text", "text": "Sure: "}, {"type": "text", "text": "Hello"}], ("Sure: Hello", " world"), ("", " world")), + ], + ids=[ + "echo_split_across_chunks_is_dropped", + "exact_echo_then_end_is_dropped", + "continuation_without_echo_passes_through", + "near_echo_is_kept_in_order", + "stream_ending_mid_match_is_kept", + "content_block_prefill_echo_is_dropped", + ], +) +@pytest.mark.asyncio +async def test_continuation_output_drops_only_a_full_echo_of_the_prefill( + prefill_content, fallback_texts, expected_texts +): + output = ContinuationOutput( + "Hello", {"role": "assistant", "content": prefill_content, "prefix": True}, "group", () + ) + + async def fallback_stream(): + for text in fallback_texts: + yield _text_chunk(text) + + released = [chunk async for chunk in output.splice(fallback_stream())] + + assert tuple(chunk.choices[0].delta.content for chunk in released) == expected_texts + + +@pytest.mark.parametrize( + "prefill,fallback_texts,expected_texts", + [ + ("Transformers use ", ("self-attention",), ("self-attention",)), + ("Hello ", ("Hello ", " world"), ("", "world")), + ("Hello \n", ("\nworld",), ("\nworld",)), + ("A.\n\n", ("\n Hi", "\nthere"), (" Hi", "\nthere")), + ("\n\n", ("\nHi",), ("\nHi",)), + ("\nHello\n", ("\n",), ("",)), + ], + ids=[ + "continuation_without_whitespace_passes_through", + "echo_then_re_emitted_space_are_both_dropped", + "different_whitespace_is_kept", + "seam_ends_at_the_first_divergence", + "whitespace_only_prefill_has_no_seam", + "re_emitted_whitespace_at_stream_end_is_dropped", + ], +) +@pytest.mark.asyncio +async def test_continuation_output_drops_only_whitespace_the_client_already_has( + prefill, fallback_texts, expected_texts +): + output = ContinuationOutput(prefill, {"role": "assistant", "content": prefill, "prefix": True}, "group", ()) + + async def fallback_stream(): + for text in fallback_texts: + yield _text_chunk(text) + + released = [chunk async for chunk in output.splice(fallback_stream())] + + assert tuple(chunk.choices[0].delta.content for chunk in released) == expected_texts + + +@pytest.mark.parametrize( + "request_kwargs,expected", + [ + ({"max_tokens": 100}, {"max_tokens": 70}), + ({"max_tokens": 100, "max_completion_tokens": 40}, {"max_tokens": 70, "max_completion_tokens": 10}), + ({}, {}), + ( + {"max_tokens": 100, "extra_body": {"max_tokens": 50, "top_k": 3}}, + {"max_tokens": 70, "extra_body": {"max_tokens": 20, "top_k": 3}}, + ), + ({"extra_body": {"max_completion_tokens": 45}}, {"extra_body": {"max_completion_tokens": 15}}), + ({"extra_body": {"top_k": 3}}, {}), + ], + ids=["max_tokens", "both_ceilings", "no_ceiling", "extra_body_max_tokens", "extra_body_only", "extra_body_no_ceiling"], +) +def test_continuation_output_ceilings_take_the_emitted_tokens_off_every_ceiling_the_request_sends( + request_kwargs, expected +): + assert continuation_output_ceilings(request_kwargs, 30) == expected + + +@pytest.mark.parametrize( + "request_kwargs", + [ + {"max_tokens": 30}, + {"max_tokens": 20}, + {"max_tokens": 100, "max_completion_tokens": 30}, + {"max_tokens": 100, "extra_body": {"max_tokens": 30}}, + ], + ids=["max_tokens_met", "max_tokens_exceeded", "one_of_two_met", "extra_body_met"], +) +def test_continuation_output_ceilings_decline_once_a_ceiling_has_no_tokens_left(request_kwargs): + assert continuation_output_ceilings(request_kwargs, 30) is None diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 3bd051d71d4..6bc52ab3dfc 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -1,4 +1,5 @@ import ast, asyncio, importlib, time +import contextlib import copy import functools import gc @@ -12,7 +13,7 @@ import warnings from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Callable, Mapping from datetime import datetime, timedelta, timezone from pathlib import Path -from types import SimpleNamespace +from types import MappingProxyType, SimpleNamespace from typing import TYPE_CHECKING, Final, Literal, cast from unittest.mock import AsyncMock, MagicMock, Mock, patch @@ -2934,6 +2935,1216 @@ async def test_acompletion_streaming_iterator(): print("\n=== All tests passed! ===") +def _make_midstream_source(error, chunks=None, first_chunk=None): + from unittest.mock import MagicMock + + opening_chunk = ( + first_chunk if first_chunk is not None else MagicMock(choices=[MagicMock(delta=MagicMock(content="Hello"))]) + ) + + class _Source: + def __init__(self): + self.index = 0 + self.chunks = chunks if chunks is not None else [] + self.model = "gpt-4" + self.custom_llm_provider = "openai" + self.logging_obj = MagicMock() + + def __aiter__(self): + return self + + async def __anext__(self): + if self.index == 0: + self.index += 1 + return opening_chunk + raise error + + async def wait_for_failure_logging(self, timeout): + return None + + return _Source() + + +class _FakeFallbackStream: + def __init__(self, item): + self._item = item + self._done = False + + def __aiter__(self): + return self + + async def __anext__(self): + if self._done: + raise StopAsyncIteration + self._done = True + return self._item + + +@pytest.mark.asyncio +async def test_acompletion_streaming_iterator_continues_after_content_when_eligible(): + """Flag on + plain-text break: the router re-enters the fallback chain with an + assistant-prefill continuation and the mid-stream marker, then streams the + fallback's output instead of re-raising.""" + from unittest.mock import AsyncMock, MagicMock, patch + + from litellm.exceptions import MidStreamFallbackError + from litellm.router_utils.pre_call_checks.continuation_prefill_check import ( + MID_STREAM_CONTINUATION_KWARG, + MID_STREAM_CONTINUATION_MARKER, + ) + + router = litellm.Router( + model_list=[ + {"model_name": "gpt-4", "litellm_params": {"model": "gpt-4", "api_key": "k1"}}, + {"model_name": "backup", "litellm_params": {"model": "anthropic/claude-3-opus-20240229", "api_key": "k2"}}, + ], + fallbacks=[{"gpt-4": ["backup"]}], + enable_mid_stream_fallback_continuation=True, + ) + + error = MidStreamFallbackError( + message="Connection lost", + model="gpt-4", + llm_provider="openai", + generated_content="Hello", + is_pre_first_chunk=False, + emitted_disqualifying_content=False, + ) + source = _make_midstream_source(error) + fallback_chunk = litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": "world"}}]) + initial_kwargs = {"model": "gpt-4", "stream": True} + + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + new=AsyncMock(return_value=_FakeFallbackStream(fallback_chunk)), + ) as mock_fallback: + result = await router._acompletion_streaming_iterator( + model_response=source, messages=[{"role": "user", "content": "Hi"}], initial_kwargs=initial_kwargs + ) + collected = [chunk async for chunk in result] + + mock_fallback.assert_awaited_once() + passed_kwargs = mock_fallback.await_args.kwargs["kwargs"] + assert passed_kwargs[MID_STREAM_CONTINUATION_KWARG] is MID_STREAM_CONTINUATION_MARKER + assert passed_kwargs["messages"][-1] == {"role": "assistant", "content": "Hello", "prefix": True} + assert fallback_chunk in collected + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "continuation_text,joined_output_trips_rule", + [("6789", True), ("6788", False)], + ids=["rule_trips_only_on_the_joined_text", "clean_continuation_streams"], +) +async def test_acompletion_streaming_iterator_continuation_runs_post_call_rules_on_the_joined_text( + monkeypatch, continuation_text, joined_output_trips_rule +): + from unittest.mock import AsyncMock, patch + + from litellm.exceptions import MidStreamFallbackError + + monkeypatch.setattr(litellm, "post_call_rules", [lambda output: "123-45-6789" not in output]) + router = litellm.Router( + model_list=[ + {"model_name": "gpt-4", "litellm_params": {"model": "gpt-4", "api_key": "k1"}}, + {"model_name": "backup", "litellm_params": {"model": "anthropic/claude-3-opus-20240229", "api_key": "k2"}}, + ], + fallbacks=[{"gpt-4": ["backup"]}], + enable_mid_stream_fallback_continuation=True, + ) + error = MidStreamFallbackError( + message="Connection lost", model="gpt-4", llm_provider="openai", generated_content="123-45-", + is_pre_first_chunk=False, emitted_disqualifying_content=False, + ) + fallback_chunk = litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": continuation_text}}]) + + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + new=AsyncMock(return_value=_FakeFallbackStream(fallback_chunk)), + ): + result = await router._acompletion_streaming_iterator( + model_response=_make_midstream_source(error), + messages=[{"role": "user", "content": "Hi"}], + initial_kwargs={"model": "gpt-4", "stream": True}, + ) + if joined_output_trips_rule: + with pytest.raises(litellm.APIResponseValidationError): + async for _ in result: + pass + else: + assert fallback_chunk in [chunk async for chunk in result] + + +@pytest.mark.asyncio +async def test_a_second_cut_runs_post_call_rules_only_on_the_text_the_caller_received(monkeypatch): + from unittest.mock import AsyncMock, patch + + from litellm.exceptions import MidStreamFallbackError + from litellm.router_utils.pre_call_checks.continuation_prefill_check import ( + MID_STREAM_CONTINUATION_KWARG, + MID_STREAM_CONTINUATION_MARKER, + ) + + checked_texts = [] + monkeypatch.setattr(litellm, "post_call_rules", [lambda output: checked_texts.append(output) is None]) + router = litellm.Router( + model_list=[ + {"model_name": "gpt-4", "litellm_params": {"model": "gpt-4", "api_key": "k1"}}, + {"model_name": "backup", "litellm_params": {"model": "anthropic/claude-3-opus-20240229", "api_key": "k2"}}, + ], + fallbacks=[{"gpt-4": ["backup"]}], + enable_mid_stream_fallback_continuation=True, + ) + + def cut(generated_content): + return MidStreamFallbackError( + message="Connection lost", model="gpt-4", llm_provider="openai", generated_content=generated_content, + is_pre_first_chunk=False, emitted_disqualifying_content=False, + ) + + messages = [{"role": "user", "content": "Hi"}] + middle_leg = await router._acompletion_streaming_iterator( + model_response=_make_midstream_source( + cut(" there"), + first_chunk=litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": " there"}}]), + ), + messages=[*messages, {"role": "assistant", "content": "Hello", "prefix": True}], + initial_kwargs={"model": "gpt-4", "stream": True, MID_STREAM_CONTINUATION_KWARG: MID_STREAM_CONTINUATION_MARKER}, + ) + backup_leg = _FakeFallbackStream( + litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": " friend"}}]) + ) + + with patch.object( + router, "async_function_with_fallbacks_common_utils", new=AsyncMock(side_effect=[middle_leg, backup_leg]) + ): + result = await router._acompletion_streaming_iterator( + model_response=_make_midstream_source(cut("Hello")), + messages=messages, + initial_kwargs={"model": "gpt-4", "stream": True}, + ) + _ = [chunk async for chunk in result] + + assert checked_texts == ["Hello there", "Hello there friend"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "fallback_texts", + [("Hel", "lo", " world"), (" world",)], + ids=["fallback_echoes_the_prefill", "fallback_continues_after_the_prefill"], +) +async def test_acompletion_streaming_iterator_continuation_never_repeats_the_partial_output(fallback_texts): + from unittest.mock import AsyncMock, patch + + from litellm.exceptions import MidStreamFallbackError + + router = litellm.Router( + model_list=[ + {"model_name": "gpt-4", "litellm_params": {"model": "gpt-4", "api_key": "k1"}}, + {"model_name": "backup", "litellm_params": {"model": "mistral/mistral-small-latest", "api_key": "k2"}}, + ], + fallbacks=[{"gpt-4": ["backup"]}], + enable_mid_stream_fallback_continuation=True, + ) + error = MidStreamFallbackError( + message="Connection lost", model="gpt-4", llm_provider="openai", generated_content="Hello", + is_pre_first_chunk=False, emitted_disqualifying_content=False, + ) + + async def fallback_stream(): + for text in fallback_texts: + yield litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": text}}]) + + with patch.object( + router, "async_function_with_fallbacks_common_utils", new=AsyncMock(return_value=fallback_stream()) + ): + result = await router._acompletion_streaming_iterator( + model_response=_make_midstream_source(error), + messages=[{"role": "user", "content": "Hi"}], + initial_kwargs={"model": "gpt-4", "stream": True}, + ) + collected = [chunk async for chunk in result] + + continuation = [chunk for chunk in collected if isinstance(chunk, litellm.ModelResponseStream)] + assert "".join(chunk.choices[0].delta.content for chunk in continuation) == " world" + + +@pytest.mark.asyncio +async def test_acompletion_streaming_iterator_continuation_stays_in_the_primary_stream(): + from unittest.mock import AsyncMock, patch + + from litellm.exceptions import MidStreamFallbackError + + router = litellm.Router( + model_list=[ + {"model_name": "gpt-4", "litellm_params": {"model": "gpt-4", "api_key": "k1"}}, + {"model_name": "backup", "litellm_params": {"model": "anthropic/claude-3-opus-20240229", "api_key": "k2"}}, + ], + fallbacks=[{"gpt-4": ["backup"]}], + enable_mid_stream_fallback_continuation=True, + ) + error = MidStreamFallbackError( + message="Connection lost", model="gpt-4", llm_provider="openai", generated_content="Hello", + is_pre_first_chunk=False, emitted_disqualifying_content=False, + ) + primary_chunk = litellm.ModelResponseStream( + id="chatcmpl-primary", created=111, choices=[{"index": 0, "delta": {"role": "assistant", "content": "Hello"}}] + ) + + async def fallback_stream(): + yield litellm.ModelResponseStream( + id="msg-fallback", created=222, choices=[{"index": 0, "delta": {"role": "assistant", "content": " wor"}}] + ) + yield litellm.ModelResponseStream( + id="msg-fallback", created=222, choices=[{"index": 0, "delta": {"content": "ld"}}] + ) + + with patch.object( + router, "async_function_with_fallbacks_common_utils", new=AsyncMock(return_value=fallback_stream()) + ): + result = await router._acompletion_streaming_iterator( + model_response=_make_midstream_source(error, chunks=[primary_chunk]), + messages=[{"role": "user", "content": "Hi"}], + initial_kwargs={"model": "gpt-4", "stream": True}, + ) + collected = [chunk async for chunk in result] + + continuation = [chunk for chunk in collected if isinstance(chunk, litellm.ModelResponseStream)] + assert [(chunk.id, chunk.created, chunk.choices[0].delta.role) for chunk in continuation] == [ + ("chatcmpl-primary", 111, None), + ("chatcmpl-primary", 111, None), + ] + assert "".join(chunk.choices[0].delta.content for chunk in continuation) == " world" + + +def _continuation_router(): + return litellm.Router( + model_list=[ + {"model_name": "gpt-4", "litellm_params": {"model": "gpt-4", "api_key": "k1"}}, + {"model_name": "backup", "litellm_params": {"model": "anthropic/claude-3-opus-20240229", "api_key": "k2"}}, + ], + fallbacks=[{"gpt-4": ["backup"]}], + enable_mid_stream_fallback_continuation=True, + ) + + +def _usage_chunk(prompt_tokens, completion_tokens): + from litellm.types.utils import Usage + + return litellm.ModelResponseStream( + model="gpt-4", + choices=[], + usage=Usage( + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=prompt_tokens + completion_tokens, + ), + ) + + +async def _drain_with_fallback(router, error, primary_chunks, messages, fallback_items): + from unittest.mock import AsyncMock, patch + + async def fallback_stream(): + for item in fallback_items: + yield item + + with patch.object( + router, "async_function_with_fallbacks_common_utils", new=AsyncMock(return_value=fallback_stream()) + ) as mock_fallback: + result = await router._acompletion_streaming_iterator( + model_response=_make_midstream_source(error, chunks=primary_chunks), + messages=messages, + initial_kwargs={"model": "gpt-4", "stream": True}, + ) + collected = [chunk async for chunk in result] + streamed = [chunk for chunk in collected if isinstance(chunk, litellm.ModelResponseStream)] + return mock_fallback.await_args.kwargs["kwargs"], streamed + + +@pytest.mark.asyncio +@pytest.mark.parametrize("primary_reports_usage", [True, False], ids=["usage_reported", "usage_not_reported"]) +async def test_acompletion_streaming_iterator_continuation_usage_sums_both_legs(primary_reports_usage): + from litellm.exceptions import MidStreamFallbackError + + messages = [{"role": "user", "content": "Explain transformers in one paragraph."}] + partial = "Transformers use " + error = MidStreamFallbackError( + message="Connection lost", model="gpt-4", llm_provider="openai", generated_content=partial, + is_pre_first_chunk=False, emitted_disqualifying_content=False, + ) + content_chunk = litellm.ModelResponseStream( + model="gpt-4", choices=[{"index": 0, "delta": {"role": "assistant", "content": partial}}] + ) + primary_chunks = [content_chunk, _usage_chunk(21, 3)] if primary_reports_usage else [content_chunk] + primary_usage = litellm.stream_chunk_builder(chunks=list(primary_chunks), messages=messages).usage + + _, streamed = await _drain_with_fallback( + _continuation_router(), error, primary_chunks, messages, [_usage_chunk(91, 5)] + ) + + assert primary_usage.prompt_tokens > 0 + assert streamed[-1].usage.prompt_tokens == primary_usage.prompt_tokens + 91 + assert streamed[-1].usage.completion_tokens == primary_usage.completion_tokens + 5 + + +class _CutSSE(httpx.AsyncByteStream): + def __init__(self, content: str) -> None: + self._content: Final = content + + async def __aiter__(self) -> AsyncIterator[bytes]: + yield f'data: {json.dumps({"id": "cut", "object": "chat.completion.chunk", "created": 1700000000, "model": "m", "choices": [{"index": 0, "delta": {"role": "assistant", "content": self._content}, "finish_reason": None}]})}\n\n'.encode() + raise httpx.RemoteProtocolError("peer closed connection without sending complete message body") + + +_BACKUP_SSE: Final = "".join( + f"data: {json.dumps(chunk)}\n\n" + for chunk in ( + {"id": "b", "object": "chat.completion.chunk", "created": 1700000000, "model": "mistral-large-latest", + "choices": [{"index": 0, "delta": {"role": "assistant", "content": " self-attention."}, "finish_reason": None}]}, + {"id": "b", "object": "chat.completion.chunk", "created": 1700000000, "model": "mistral-large-latest", + "choices": [{"index": 0, "delta": {"content": ""}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 91, "completion_tokens": 5, "total_tokens": 96}}, + ) +) + "data: [DONE]\n\n" + + +class _SpendRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.successes: Final[list[dict[str, object]]] = [] + self.failures: Final[list[dict[str, object]]] = [] + self.spend_rows: Final[list[dict[str, object]]] = [] + self.succeeded: Final = asyncio.Event() + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time) -> None: + from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload + + self.successes.append(kwargs["standard_logging_object"]) + self.spend_rows.append(get_logging_payload(kwargs, response_obj, start_time, end_time)) + self.succeeded.set() + + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time) -> None: + self.failures.append(kwargs["standard_logging_object"]) + + +class _SlowFailureRecorder(_SpendRecorder): + def __init__(self) -> None: + super().__init__() + self.failed_deployment_ids: Final[list[str]] = [] + + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time) -> None: + await asyncio.sleep(0.05) + self.failed_deployment_ids.append(kwargs["litellm_params"]["metadata"]["model_info"]["id"]) + await super().async_log_failure_event(kwargs, response_obj, start_time, end_time) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("cut_groups", [("primary",), ("primary", "middle")], ids=["one_break", "two_breaks"]) +async def test_continued_stream_success_log_bills_every_cut_leg( + monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter, cut_groups: tuple[str, ...] +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + recorder: Final = _SpendRecorder() + monkeypatch.setattr(litellm, "callbacks", [recorder]) + respx_mock.post("https://api.openai.com/v1/chat/completions").mock( + return_value=httpx.Response(200, headers={"content-type": "text/event-stream"}, stream=_CutSSE("Transformers use")) + ) + respx_mock.post(host="api.deepseek.com").mock( + return_value=httpx.Response(200, headers={"content-type": "text/event-stream"}, stream=_CutSSE(" learned")) + ) + respx_mock.post(host="api.mistral.ai").mock( + return_value=httpx.Response(200, headers={"content-type": "text/event-stream"}, text=_BACKUP_SSE) + ) + router: Final = Router( + model_list=[ + { + "model_name": "primary", + "litellm_params": { + "model": "openai/gpt-4o", + "api_key": "k1", + "input_cost_per_token": 0.001, + "output_cost_per_token": 0.002, + }, + }, + { + "model_name": "middle", + "litellm_params": { + "model": "deepseek/deepseek-chat", + "api_key": "k2", + "input_cost_per_token": 0.003, + "output_cost_per_token": 0.004, + }, + }, + {"model_name": "backup", "litellm_params": {"model": "mistral/mistral-large-latest", "api_key": "k3"}}, + ], + fallbacks=[{"primary": ["middle" if "middle" in cut_groups else "backup"]}, {"middle": ["backup"]}], + enable_mid_stream_fallback_continuation=True, + ) + + response: Final = await router.acompletion( + model="primary", + messages=[{"role": "user", "content": "How do transformers work?"}], + stream=True, + stream_options={"include_usage": True}, + ) + chunks: Final = [chunk async for chunk in response] + await asyncio.wait_for(recorder.succeeded.wait(), timeout=10) + await asyncio.sleep(0.5) + + backup_prompt_cost, backup_completion_cost = litellm.cost_per_token( + model="mistral/mistral-large-latest", prompt_tokens=91, completion_tokens=5 + ) + success: Final = recorder.successes[-1] + assert len(recorder.failures) == len(cut_groups) + assert [failure["response_cost"] for failure in recorder.failures] == pytest.approx( + [ + failure["prompt_tokens"] * input_price + failure["completion_tokens"] * output_price + for failure, (input_price, output_price) in zip(recorder.failures, ((0.001, 0.002), (0.003, 0.004))) + ] + ) + assert success["response_cost"] == pytest.approx( + sum(failure["response_cost"] for failure in recorder.failures) + backup_prompt_cost + backup_completion_cost + ) + assert success["prompt_tokens"] == sum(failure["prompt_tokens"] for failure in recorder.failures) + 91 + assert success["completion_tokens"] == sum(failure["completion_tokens"] for failure in recorder.failures) + 5 + assert success["cost_breakdown"]["total_cost"] == pytest.approx(success["response_cost"]) + spend_row: Final = recorder.spend_rows[-1] + assert (spend_row["spend"], spend_row["prompt_tokens"], spend_row["completion_tokens"]) == ( + success["response_cost"], + success["prompt_tokens"], + success["completion_tokens"], + ) + client_usage: Final = next(chunk.usage for chunk in reversed(chunks) if getattr(chunk, "usage", None)) + assert (client_usage.prompt_tokens, client_usage.completion_tokens) == ( + success["prompt_tokens"], + success["completion_tokens"], + ) + assert {chunk.id for chunk in chunks} == {spend_row["request_id"]} + + +@pytest.mark.asyncio +async def test_continued_stream_sharing_the_proxy_logging_object_bills_the_cut_to_its_own_deployment( + monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + recorder: Final = _SlowFailureRecorder() + monkeypatch.setattr(litellm, "callbacks", [recorder]) + respx_mock.post("https://api.openai.com/v1/chat/completions").mock( + return_value=httpx.Response(200, headers={"content-type": "text/event-stream"}, stream=_CutSSE("Transformers use")) + ) + respx_mock.post(host="api.mistral.ai").mock( + return_value=httpx.Response(200, headers={"content-type": "text/event-stream"}, text=_BACKUP_SSE) + ) + router: Final = Router( + model_list=[ + { + "model_name": "primary", + "litellm_params": { + "model": "openai/gpt-4o", + "api_key": "k1", + "input_cost_per_token": 0.001, + "output_cost_per_token": 0.002, + }, + "model_info": {"id": "primary-deployment"}, + }, + {"model_name": "backup", "litellm_params": {"model": "mistral/mistral-large-latest", "api_key": "k3"}}, + ], + fallbacks=[{"primary": ["backup"]}], + enable_mid_stream_fallback_continuation=True, + ) + messages: Final = [{"role": "user", "content": "How do transformers work?"}] + shared_logging, _ = litellm.utils.function_setup( + original_function="acompletion", + rules_obj=litellm.utils.Rules(), + start_time=datetime.now(), + model="primary", + messages=messages, + stream=True, + litellm_call_id="shared-call", + ) + + response: Final = await router.acompletion( + model="primary", + messages=messages, + stream=True, + stream_options={"include_usage": True}, + litellm_call_id="shared-call", + litellm_logging_obj=shared_logging, + ) + _ = [chunk async for chunk in response] + await asyncio.wait_for(recorder.succeeded.wait(), timeout=10) + await asyncio.sleep(0.5) + + backup_prompt_cost, backup_completion_cost = litellm.cost_per_token( + model="mistral/mistral-large-latest", prompt_tokens=91, completion_tokens=5 + ) + (failure,) = recorder.failures + assert (failure["model_group"], failure["model_id"]) == ("primary", "primary-deployment") + assert recorder.failed_deployment_ids == ["primary-deployment"] + assert failure["response_cost"] == pytest.approx( + failure["prompt_tokens"] * 0.001 + failure["completion_tokens"] * 0.002 + ) + assert recorder.successes[-1]["response_cost"] == pytest.approx( + failure["response_cost"] + backup_prompt_cost + backup_completion_cost + ) + + +def test_carry_partial_spend_to_fallback_adds_the_partial_the_cut_stream_inherited(): + from litellm.types.utils import Usage + + cut_logging: Final = LiteLLMLogging( + model="deepseek/deepseek-chat", messages=[], stream=True, call_type="acompletion", + start_time=datetime.now(), litellm_call_id="cut", function_id="cut", + ) + cut_logging.model_call_details["response_cost"] = 0.5 + cut_logging.update_environment_variables( + litellm_params={"custom_llm_provider": "deepseek"}, optional_params={}, model="deepseek/deepseek-chat" + ) + cut_logging.record_continued_partial( + usage=Usage(prompt_tokens=40, completion_tokens=7, total_tokens=47), + response_cost=0.001, + response_id="chatcmpl-primary", + ) + fallback_logging: Final = LiteLLMLogging( + model="mistral-large-latest", messages=[], stream=True, call_type="acompletion", + start_time=datetime.now(), litellm_call_id="fallback", function_id="fallback", + ) + + Router._carry_partial_spend_to_fallback( + SimpleNamespace(logging_obj=fallback_logging), + Router._price_cut_partial( + cast(CustomStreamWrapper, SimpleNamespace(logging_obj=cut_logging, model="deepseek/deepseek-chat")), + litellm.ModelResponse( + id="chatcmpl-relay", usage=Usage(prompt_tokens=60, completion_tokens=3, total_tokens=63) + ), + ), + ) + + relay_prompt_cost, relay_completion_cost = litellm.cost_per_token( + model="deepseek/deepseek-chat", prompt_tokens=60, completion_tokens=3 + ) + carried_usage: Final = fallback_logging.model_call_details["continued_partial_usage"] + assert (carried_usage.prompt_tokens, carried_usage.completion_tokens) == (100, 10) + assert fallback_logging.model_call_details["continued_partial_cost"] == pytest.approx( + relay_prompt_cost + relay_completion_cost + 0.001 + ) + assert fallback_logging.model_call_details["continued_response_id"] == "chatcmpl-primary" + + +@pytest.mark.asyncio +async def test_acompletion_streaming_iterator_plain_fallback_keeps_the_fallback_usage(): + from litellm.exceptions import MidStreamFallbackError + + error = MidStreamFallbackError( + message="Connection lost", model="gpt-4", llm_provider="anthropic", generated_content="", + is_pre_first_chunk=True, + ) + role_chunk = litellm.ModelResponseStream( + model="gpt-4", choices=[{"index": 0, "delta": {"role": "assistant", "content": ""}}] + ) + + _, streamed = await _drain_with_fallback( + _continuation_router(), + error, + [role_chunk, _usage_chunk(21, 0)], + [{"role": "user", "content": "Hi"}], + [_usage_chunk(91, 5)], + ) + + assert (streamed[-1].usage.prompt_tokens, streamed[-1].usage.completion_tokens) == (91, 5) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "partial,fallback_texts,expected_answer", + [ + ("Transformers use ", (" self", "-attention"), "Transformers use self-attention"), + ("def f():\n ", ("\n", " return 1"), "def f():\n return 1"), + ("First.\n", ("\n\nSecond",), "First.\n\nSecond"), + ], + ids=["space", "indentation", "paragraph_break"], +) +async def test_acompletion_streaming_iterator_continuation_does_not_double_trailing_whitespace( + partial, fallback_texts, expected_answer +): + from litellm.exceptions import MidStreamFallbackError + + error = MidStreamFallbackError( + message="Connection lost", model="gpt-4", llm_provider="openai", generated_content=partial, + is_pre_first_chunk=False, emitted_disqualifying_content=False, + ) + fallback_items = [ + litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": text}}]) for text in fallback_texts + ] + + passed_kwargs, streamed = await _drain_with_fallback( + _continuation_router(), error, [], [{"role": "user", "content": "Hi"}], fallback_items + ) + + assert passed_kwargs["messages"][-1]["content"] == partial + assert partial + "".join(chunk.choices[0].delta.content for chunk in streamed) == expected_answer + + +@pytest.mark.asyncio +@pytest.mark.parametrize("enabled", [True, "true", " True "]) +async def test_update_settings_turns_on_mid_stream_continuation(monkeypatch, enabled): + from unittest.mock import AsyncMock, patch + + from litellm.exceptions import MidStreamFallbackError + from litellm.router_utils.pre_call_checks.continuation_prefill_check import ( + MID_STREAM_CONTINUATION_KWARG, + MID_STREAM_CONTINUATION_MARKER, + ContinuationPrefillDeploymentCheck, + ) + + monkeypatch.setattr(litellm, "callbacks", []) + router = litellm.Router( + model_list=[ + {"model_name": "gpt-4", "litellm_params": {"model": "gpt-4", "api_key": "k1"}}, + {"model_name": "backup", "litellm_params": {"model": "anthropic/claude-3-opus-20240229", "api_key": "k2"}}, + ], + fallbacks=[{"gpt-4": ["backup"]}], + ) + router.update_settings(enable_mid_stream_fallback_continuation=enabled) + + error = MidStreamFallbackError( + message="Connection lost", model="gpt-4", llm_provider="openai", generated_content="Hello", + is_pre_first_chunk=False, emitted_disqualifying_content=False, + ) + fallback_chunk = litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": " world"}}]) + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + new=AsyncMock(return_value=_FakeFallbackStream(fallback_chunk)), + ) as mock_fallback: + result = await router._acompletion_streaming_iterator( + model_response=_make_midstream_source(error), + messages=[{"role": "user", "content": "Hi"}], + initial_kwargs={"model": "gpt-4", "stream": True}, + ) + collected = [chunk async for chunk in result] + + assert mock_fallback.await_args.kwargs["kwargs"][MID_STREAM_CONTINUATION_KWARG] is MID_STREAM_CONTINUATION_MARKER + assert fallback_chunk in collected + assert any(isinstance(callback, ContinuationPrefillDeploymentCheck) for callback in litellm.callbacks) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "error_kwargs,request_kwargs", + [ + ({"emitted_disqualifying_content": True}, {}), + ({"emitted_disqualifying_content": False}, {"response_format": {"type": "json_object"}}), + ({"emitted_disqualifying_content": False}, {"tool_choice": "required"}), + ({"emitted_disqualifying_content": False}, {"merge_reasoning_content_in_choices": True}), + ({"emitted_disqualifying_content": False}, {"n": 2}), + ], + ids=["tool_or_thinking_emitted", "json_mode", "forced_tool_choice", "merged_reasoning", "multiple_choices"], +) +async def test_acompletion_streaming_iterator_declines_ineligible_after_content(error_kwargs, request_kwargs): + """Flag on but the break is not continuation-safe: the router re-raises and + never enters the fallback chain, so no duplicated/rejected request is sent.""" + from unittest.mock import AsyncMock, patch + + from litellm.exceptions import MidStreamFallbackError + + router = litellm.Router( + model_list=[ + {"model_name": "gpt-4", "litellm_params": {"model": "gpt-4", "api_key": "k1"}}, + {"model_name": "backup", "litellm_params": {"model": "anthropic/claude-3-opus-20240229", "api_key": "k2"}}, + ], + fallbacks=[{"gpt-4": ["backup"]}], + enable_mid_stream_fallback_continuation=True, + ) + error = MidStreamFallbackError( + message="boom", model="gpt-4", llm_provider="openai", generated_content="Hello", + is_pre_first_chunk=False, **error_kwargs, + ) + source = _make_midstream_source(error) + initial_kwargs = {"model": "gpt-4", "stream": True, **request_kwargs} + + with patch.object(router, "async_function_with_fallbacks_common_utils", new=AsyncMock()) as mock_fallback: + result = await router._acompletion_streaming_iterator( + model_response=source, messages=[{"role": "user", "content": "Hi"}], initial_kwargs=initial_kwargs + ) + with pytest.raises(MidStreamFallbackError): + async for _ in result: + pass + + mock_fallback.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_acompletion_streaming_iterator_raises_content_filter_instead_of_continuing(): + """A content filter that stops a stream mid-answer reaches the client as ContentPolicyViolationError, + never as a continuation that lets a backup finish the filtered answer.""" + from unittest.mock import AsyncMock, patch + + from litellm.exceptions import MidStreamFallbackError + + router = litellm.Router( + model_list=[ + {"model_name": "gpt-4", "litellm_params": {"model": "azure/gpt-4", "api_key": "k1"}}, + {"model_name": "backup", "litellm_params": {"model": "anthropic/claude-3-opus-20240229", "api_key": "k2"}}, + ], + fallbacks=[{"gpt-4": ["backup"]}], + enable_mid_stream_fallback_continuation=True, + ) + content_filter = litellm.ContentPolicyViolationError(message="filtered", model="gpt-4", llm_provider="azure") + error = MidStreamFallbackError( + message="filtered", model="gpt-4", llm_provider="azure", generated_content="Step one is to obtain the", + is_pre_first_chunk=False, emitted_disqualifying_content=False, original_exception=content_filter, + ) + + with patch.object(router, "async_function_with_fallbacks_common_utils", new=AsyncMock()) as mock_fallback: + result = await router._acompletion_streaming_iterator( + model_response=_make_midstream_source(error), + messages=[{"role": "user", "content": "Hi"}], + initial_kwargs={"model": "gpt-4", "stream": True}, + ) + with pytest.raises(litellm.ContentPolicyViolationError): + async for _ in result: + pass + + mock_fallback.assert_not_awaited() + + +async def _continues_after_content( + initial_kwargs: dict[str, object], + primary_params: Mapping[str, object] = MappingProxyType({}), + default_litellm_params: dict[str, object] | None = None, + cut_model_id: str = "primary-deployment", +) -> bool: + from unittest.mock import AsyncMock, patch + + from litellm.exceptions import MidStreamFallbackError + + router: Final = litellm.Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4", "api_key": "k1", **primary_params}, + "model_info": {"id": "primary-deployment"}, + }, + {"model_name": "backup", "litellm_params": {"model": "anthropic/claude-3-opus-20240229", "api_key": "k2"}}, + { + "model_name": "openai-backup", + "litellm_params": {"model": "gpt-4o", "api_key": "k3"}, + "model_info": {"id": "no-prefill-deployment", "supports_assistant_prefill": False}, + }, + { + "model_name": "claude-backup", + "litellm_params": {"model": "anthropic/claude-3-opus-20240229", "api_key": "k4"}, + "model_info": {"id": "prefill-deployment", "supports_assistant_prefill": True}, + }, + { + "model_name": "json-claude-backup", + "litellm_params": { + "model": "anthropic/claude-3-opus-20240229", + "api_key": "k5", + "response_format": {"type": "json_object"}, + }, + "model_info": {"id": "json-prefill-deployment", "supports_assistant_prefill": True}, + }, + ], + fallbacks=[{"gpt-4": ["backup"]}], + default_litellm_params=default_litellm_params, + enable_mid_stream_fallback_continuation=True, + ) + error: Final = MidStreamFallbackError( + message="boom", model="gpt-4", llm_provider="openai", generated_content="Hello", + is_pre_first_chunk=False, emitted_disqualifying_content=False, + ) + fallback_chunk: Final = litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": " world"}}]) + cut_stream: Final = _make_midstream_source(error) + cut_stream._hidden_params = {"model_id": cut_model_id} + with patch.object( + router, + "async_function_with_fallbacks_common_utils", + new=AsyncMock(return_value=_FakeFallbackStream(fallback_chunk)), + ) as mock_fallback: + result: Final = await router._acompletion_streaming_iterator( + model_response=cut_stream, + messages=[{"role": "user", "content": "Hi"}], + initial_kwargs={"model": "gpt-4", "stream": True, **initial_kwargs}, + ) + with contextlib.suppress(MidStreamFallbackError): + _ = [chunk async for chunk in result] + return mock_fallback.await_count == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "event_hook,continues", + [("post_call", False), (None, False), ("pre_call", True)], + ids=["post_call_guardrail", "every_hook_guardrail", "input_only_guardrail"], +) +async def test_mid_stream_continuation_declines_while_an_output_guardrail_is_registered( + monkeypatch: pytest.MonkeyPatch, event_hook: str | None, continues: bool +) -> None: + from litellm.types.guardrails import GuardrailEventHooks + + guardrail: Final = CustomGuardrail( + guardrail_name="output-check", event_hook=GuardrailEventHooks(event_hook) if event_hook else None + ) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + + assert await _continues_after_content({}) is continues + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "fallbacks,continues", + [ + ([{"model": "backup", "messages": [{"role": "user", "content": "Say hi"}]}], False), + ([{"model": "backup", "stream": False}], False), + ([{"model": "backup", "max_tokens": 300}], False), + ([{"model": "backup"}], True), + ([{"gpt-4": [{"model": "backup", "stream": False}]}], False), + ([{"*": [{"model": "backup", "max_tokens": 300}]}], False), + ([{"gpt-4": [{"model": "backup"}, "backup"]}], True), + ([{"gpt-4": ["backup"]}, {"other-group": [{"model": "backup", "stream": False}]}], True), + ([{}], True), + ([{"gpt-4": 5}], True), + ], + ids=[ + "entry_with_its_own_messages", + "entry_turning_off_streaming", + "entry_with_its_own_max_tokens", + "entry_with_model_only", + "group_target_turning_off_streaming", + "generic_group_target_with_its_own_max_tokens", + "group_targets_with_model_only", + "override_only_in_another_groups_chain", + "empty_entry", + "non_list_chain", + ], +) +async def test_mid_stream_continuation_declines_a_request_fallback_that_overrides_request_params( + fallbacks: list[object], continues: bool +) -> None: + assert await _continues_after_content({"fallbacks": fallbacks}) is continues + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "fallbacks,continues", + [ + ([{"gpt-4": ["no-prefill-deployment"]}], False), + ([{"gpt-4": ["backup", "no-prefill-deployment"]}], False), + ([{"model": "no-prefill-deployment"}], False), + ([{"gpt-4": ["prefill-deployment"]}], True), + ([{"gpt-4": ["openai-backup"]}], True), + ([{"gpt-4": ["json-prefill-deployment"]}], False), + ], + ids=[ + "deployment_id_without_prefill", + "deployment_id_without_prefill_later_in_chain", + "flat_entry_naming_a_deployment_id_without_prefill", + "deployment_id_with_prefill", + "group_without_prefill_left_to_the_deployment_filter", + "deployment_id_with_its_own_response_format", + ], +) +async def test_mid_stream_continuation_declines_a_fallback_naming_a_deployment_that_cannot_continue( + fallbacks: list[object], continues: bool +) -> None: + assert await _continues_after_content({"fallbacks": fallbacks}) is continues + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "primary_params,default_litellm_params,continues", + [ + ({"response_format": {"type": "json_object"}}, None, False), + ({"n": 2}, None, False), + ({"tool_choice": "required"}, None, False), + ({"merge_reasoning_content_in_choices": True}, None, False), + ({}, {"response_format": {"type": "json_object"}}, False), + ({"response_format": {"type": "text"}}, None, True), + ({"temperature": 0.2}, None, True), + ({"tool_choice": "required"}, {"tool_choice": "auto"}, False), + ({"response_format": {"type": "json_object"}}, {"response_format": {"type": "text"}}, True), + ({"extra_body": {"n": 2}}, None, False), + ], + ids=[ + "deployment_response_format", + "deployment_n", + "deployment_tool_choice", + "deployment_merged_reasoning", + "router_default_response_format", + "deployment_plain_text_format", + "deployment_unrelated_param", + "deployment_tool_choice_over_router_default", + "router_default_format_over_deployment", + "deployment_extra_body_n", + ], +) +async def test_mid_stream_continuation_reads_constraints_configured_on_the_cut_deployment( + primary_params: Mapping[str, object], default_litellm_params: dict[str, object] | None, continues: bool +) -> None: + assert await _continues_after_content({}, primary_params, default_litellm_params) is continues + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "extra_body,continues", + [({"response_format": {"type": "json_object"}}, False), ({"n": 2}, False), ({"top_k": 5}, True)], + ids=["json_format_in_extra_body", "n_in_extra_body", "unrelated_extra_body_key"], +) +async def test_mid_stream_continuation_reads_constraints_the_request_sends_in_extra_body( + extra_body: dict[str, object], continues: bool +) -> None: + assert await _continues_after_content({"extra_body": extra_body}) is continues + + +@pytest.mark.asyncio +async def test_mid_stream_continuation_declines_when_the_cut_deployment_is_not_in_the_router() -> None: + assert await _continues_after_content({}, cut_model_id="per-request-credential-deployment") is False + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "routing_strategy", + [ + "simple-shuffle", + "usage-based-routing", + "usage-based-routing-v2", + "least-busy", + "latency-based-routing", + "cost-based-routing", + ], +) +async def test_continuation_routes_only_to_deployments_that_can_continue_under_every_routing_strategy( + routing_strategy: str, +) -> None: + from litellm.router_utils.pre_call_checks.continuation_prefill_check import ( + MID_STREAM_CONTINUATION_KWARG, + MID_STREAM_CONTINUATION_MARKER, + ) + from litellm.types.router import RouterRateLimitError + + router: Final = litellm.Router( + model_list=[ + { + "model_name": "backup", + "litellm_params": {"model": "gpt-4o", "api_key": "k1"}, + "model_info": {"id": "no-prefill-deployment", "supports_assistant_prefill": False}, + }, + ], + routing_strategy=routing_strategy, + enable_mid_stream_fallback_continuation=True, + ) + messages: Final = [{"role": "user", "content": "Hi"}] + plain_pick: Final = await router.async_get_available_deployment( + model="backup", messages=messages, request_kwargs={} + ) + assert plain_pick["model_info"]["id"] == "no-prefill-deployment" + with pytest.raises(RouterRateLimitError): + await router.async_get_available_deployment( + model="backup", + messages=messages, + request_kwargs={MID_STREAM_CONTINUATION_KWARG: MID_STREAM_CONTINUATION_MARKER}, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "flag_kwargs", [{}, {"enable_mid_stream_fallback_continuation": "false"}], ids=["default", "string_false"] +) +async def test_acompletion_streaming_iterator_flag_off_declines_after_content(flag_kwargs): + """Regression guard for the opt-in default: with the flag off, a post-content + break re-raises and never falls back, exactly as before this feature.""" + from unittest.mock import AsyncMock, patch + + from litellm.exceptions import MidStreamFallbackError + + router = litellm.Router( + model_list=[ + {"model_name": "gpt-4", "litellm_params": {"model": "gpt-4", "api_key": "k1"}}, + {"model_name": "backup", "litellm_params": {"model": "anthropic/claude-3-opus-20240229", "api_key": "k2"}}, + ], + fallbacks=[{"gpt-4": ["backup"]}], + **flag_kwargs, + ) + error = MidStreamFallbackError( + message="boom", model="gpt-4", llm_provider="openai", generated_content="Hello", + is_pre_first_chunk=False, emitted_disqualifying_content=False, + ) + source = _make_midstream_source(error) + + with patch.object(router, "async_function_with_fallbacks_common_utils", new=AsyncMock()) as mock_fallback: + result = await router._acompletion_streaming_iterator( + model_response=source, messages=[{"role": "user", "content": "Hi"}], + initial_kwargs={"model": "gpt-4", "stream": True}, + ) + with pytest.raises(MidStreamFallbackError): + async for _ in result: + pass + + mock_fallback.assert_not_awaited() + + +def test_build_completion_continuation_input_appends_assistant_prefill(): + messages = [{"role": "user", "content": "hi"}] + built = litellm.Router._build_completion_continuation_input(messages, "partial answer") + assert built[:-1] == messages + assert built[-1] == {"role": "assistant", "content": "partial answer", "prefix": True} + + +@pytest.mark.parametrize( + "prefill_content,second_partial,expected_content", + [ + ("part one ", " part two ", "part one part two "), + ("part one ", "part two ", "part one part two "), + ("First.\n\n", "Second", "First.\n\nSecond"), + ([{"type": "text", "text": "part one "}], " part two ", [{"type": "text", "text": "part one part two "}]), + ], + ids=["separator_sent_again", "separator_not_sent_again", "paragraph_break_not_sent_again", "content_blocks"], +) +def test_build_completion_continuation_input_folds_into_existing_prefill( + prefill_content, second_partial, expected_content +): + """A nested break must not leave two trailing assistant turns: the new partial folds into the prior prefill, + which keeps reading as the text the caller already received.""" + messages = [{"role": "user", "content": "hi"}, {"role": "assistant", "content": prefill_content, "prefix": True}] + built = litellm.Router._build_completion_continuation_input(messages, second_partial) + assert built == [messages[0], {"role": "assistant", "content": expected_content, "prefix": True}] + + +@pytest.mark.parametrize( + "prefill_content,expected_content", + [ + ("Hello", "Hello world"), + ([{"type": "text", "text": "Hello"}], [{"type": "text", "text": "Hello world"}]), + ], + ids=["string_prefill", "content_block_prefill"], +) +def test_build_completion_continuation_input_does_not_double_an_echoed_prefill(prefill_content, expected_content): + messages = [{"role": "user", "content": "hi"}, {"role": "assistant", "content": prefill_content, "prefix": True}] + built = litellm.Router._build_completion_continuation_input(messages, "Hello world") + assert built is not None + assert built[-1] == {"role": "assistant", "content": expected_content, "prefix": True} + + +@pytest.mark.parametrize( + "cut_partial,expected_content", + [("Hel", None), ("Hello", "Hello"), ("Hello wor", "Hello wor")], + ids=["cut_mid_echo", "cut_after_a_full_echo", "cut_after_new_text"], +) +def test_build_completion_continuation_input_declines_a_cut_in_the_middle_of_a_prefill_echo( + cut_partial, expected_content +): + messages = [{"role": "user", "content": "hi"}, {"role": "assistant", "content": "Hello", "prefix": True}] + built = litellm.Router._build_completion_continuation_input(messages, cut_partial) + expected = None if expected_content is None else [messages[0], {**messages[1], "content": expected_content}] + assert built == expected + + +@pytest.mark.parametrize( + "assistant_content", ["Earlier answer.", [{"type": "text", "text": "Earlier answer."}]], ids=["string", "blocks"] +) +def test_build_completion_continuation_input_declines_a_trailing_plain_assistant_turn(assistant_content): + """The caller's own assistant turn is history, not a prefill, so the partial is never fused into it.""" + messages = [{"role": "user", "content": "hi"}, {"role": "assistant", "content": assistant_content}] + assert litellm.Router._build_completion_continuation_input(messages, "The capital of France is") is None + + +@pytest.mark.asyncio +async def test_acompletion_streaming_iterator_declines_when_request_ends_in_plain_assistant_turn(): + from unittest.mock import AsyncMock, patch + + from litellm.exceptions import MidStreamFallbackError + + router = litellm.Router( + model_list=[ + {"model_name": "gpt-4", "litellm_params": {"model": "gpt-4", "api_key": "k1"}}, + {"model_name": "backup", "litellm_params": {"model": "anthropic/claude-3-opus-20240229", "api_key": "k2"}}, + ], + fallbacks=[{"gpt-4": ["backup"]}], + enable_mid_stream_fallback_continuation=True, + ) + error = MidStreamFallbackError( + message="boom", model="gpt-4", llm_provider="openai", generated_content="Hello", + is_pre_first_chunk=False, emitted_disqualifying_content=False, + ) + + with patch.object(router, "async_function_with_fallbacks_common_utils", new=AsyncMock()) as mock_fallback: + result = await router._acompletion_streaming_iterator( + model_response=_make_midstream_source(error), + messages=[{"role": "user", "content": "Hi"}, {"role": "assistant", "content": "Earlier answer."}], + initial_kwargs={"model": "gpt-4", "stream": True}, + ) + with pytest.raises(MidStreamFallbackError): + async for _ in result: + pass + + mock_fallback.assert_not_awaited() + + +def test_build_completion_continuation_input_keeps_structured_assistant_content(): + """Content blocks on a trailing assistant turn stay blocks: the partial continues the last text + block in place (a block boundary would let a provider drop the space between "Sure, " and "here"), + lands as a new text block only after a non-text block, and a content shape that cannot be extended declines.""" + from litellm.router import _extend_trailing_text_block + + messages = [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": [{"type": "text", "text": "Sure, "}], "prefix": True}, + ] + built = litellm.Router._build_completion_continuation_input(messages, "here it is") + assert built is not None + assert [m["role"] for m in built] == ["user", "assistant"] + assert built[-1] == {"role": "assistant", "content": [{"type": "text", "text": "Sure, here it is"}], "prefix": True} + image_block = {"type": "image_url", "image_url": {"url": "https://example.test/a.png"}} + assert _extend_trailing_text_block([image_block], "here it is") == [image_block, {"type": "text", "text": "here it is"}] + assert _extend_trailing_text_block([], "here it is") == [{"type": "text", "text": "here it is"}] + assert ( + litellm.Router._build_completion_continuation_input([{"role": "assistant", "content": 42, "prefix": True}], "x") + is None + ) + + +def test_mid_stream_continuation_eligible_allows_text_response_format(): + """response_format={"type": "text"} is the unconstrained default and must stay + eligible, unlike json_object / json_schema.""" + from litellm.exceptions import MidStreamFallbackError + + router = litellm.Router( + model_list=[{"model_name": "gpt-4", "litellm_params": {"model": "gpt-4", "api_key": "k"}}], + enable_mid_stream_fallback_continuation=True, + ) + e = MidStreamFallbackError( + message="boom", model="gpt-4", llm_provider="openai", generated_content="Hello", + is_pre_first_chunk=False, emitted_disqualifying_content=False, + ) + assert router._mid_stream_continuation_eligible(e=e, request_kwargs={"response_format": {"type": "text"}}) is True + assert router._mid_stream_continuation_eligible(e=e, request_kwargs={"response_format": {"type": "json_object"}}) is False + + +def test_raise_original_mid_stream_error_surfaces_the_provider_exception(): + from litellm.exceptions import MidStreamFallbackError, RateLimitError + + provider_error = RateLimitError(message="rate limited", llm_provider="openai", model="gpt-4") + wrapped = MidStreamFallbackError( + message="rate limited", model="gpt-4", llm_provider="openai", + original_exception=provider_error, generated_content="Hello", + ) + with pytest.raises(RateLimitError) as raised: + litellm.Router._raise_original_mid_stream_error(wrapped) + assert raised.value is provider_error + assert raised.value.__cause__ is wrapped + + bare = MidStreamFallbackError(message="boom", model="gpt-4", llm_provider="openai", generated_content="Hello") + with pytest.raises(MidStreamFallbackError) as bare_raised: + litellm.Router._raise_original_mid_stream_error(bare) + assert bare_raised.value is bare + + @pytest.mark.asyncio async def test_acompletion_streaming_iterator_reraises_original_exception_when_available(): """Async: when the mid-stream MidStreamFallbackError wraps a real provider @@ -8694,6 +9905,11 @@ def test_combine_fallback_usage(): assert chunk.usage.completion_tokens == 5 assert chunk.usage.total_tokens == 15 + Router._combine_fallback_usage(chunk, Usage(prompt_tokens=0, completion_tokens=4, total_tokens=4)) + assert chunk.usage.prompt_tokens == 10 + assert chunk.usage.completion_tokens == 9 + assert chunk.usage.total_tokens == 19 + @pytest.mark.asyncio async def test_acompletion_streaming_iterator_does_not_log_success_on_terminal_failure(): diff --git a/tests/unit/test_router_model_cost_isolation.py b/tests/unit/test_router_model_cost_isolation.py index d343f464a77..d8625ff9da3 100644 --- a/tests/unit/test_router_model_cost_isolation.py +++ b/tests/unit/test_router_model_cost_isolation.py @@ -544,13 +544,14 @@ def test_should_not_pollute_shared_key_with_custom_nonzero_pricing(): ) -def test_regex_lookaround_flag_stays_on_the_deployment_that_set_it() -> None: - """A deployment's ``supports_regex_lookaround`` override must not land on the shared +@pytest.mark.parametrize("flag", ["supports_regex_lookaround", "supports_assistant_prefill"]) +def test_capability_flag_stays_on_the_deployment_that_set_it(flag: str) -> None: + """A deployment's capability override must not land on the shared ``{provider}/{model}`` key, or every sibling deployment of that model would inherit it.""" backend_model = "bedrock/us.xai.grok-4.6" deploy_id = "grok-deploy-keep-regex" - builtin_flag = litellm.get_model_info(model=backend_model).get("supports_regex_lookaround") + builtin_flag = litellm.get_model_info(model=backend_model).get(flag) model_keys = { deploy_id: litellm.model_cost.get(deploy_id), backend_model: copy.deepcopy(litellm.model_cost.get(backend_model)), @@ -561,13 +562,13 @@ def test_regex_lookaround_flag_stays_on_the_deployment_that_set_it() -> None: { "model_name": "grok-keep-regex", "litellm_params": {"model": backend_model}, - "model_info": {"id": deploy_id, "supports_regex_lookaround": not builtin_flag}, + "model_info": {"id": deploy_id, flag: not builtin_flag}, } ], ) - assert litellm.model_cost[deploy_id]["supports_regex_lookaround"] is (not builtin_flag) - assert litellm.get_model_info(model=backend_model).get("supports_regex_lookaround") is builtin_flag + assert litellm.model_cost[deploy_id][flag] is (not builtin_flag) + assert litellm.get_model_info(model=backend_model).get(flag) is builtin_flag finally: _restore_model_cost_entries(model_keys) From f81e657903c6ee1969e8173353f875b8785f973f Mon Sep 17 00:00:00 2001 From: Rick <26716961+Bytechoreographer@users.noreply.github.com> Date: Fri, 16 Oct 2026 21:52:15 +0800 Subject: [PATCH 19/23] feat(volcengine): add Doubao-Seedance video generation via the OpenAI /v1/videos surface (#46156) * feat(volcengine): add Doubao-Seedance video generation via the OpenAI /v1/videos surface New VolcEngineVideoConfig maps litellm's OpenAI Videos API onto Volcengine Ark's contents/generations task API: create hits /api/v3/contents/generations/tasks and returns a cgt-... task id, status and content poll /api/v3/contents/generations/tasks/{id}, list maps the same collection (limit -> page_size, Ark's filter.* through extra_query), delete maps DELETE on the task URL, and edit/extension reuse create with a reference_video content item since Ark models both as ordinary generation tasks. Ark's own edit and extension require ratio adaptive (edit also duration -1), so those are pinned unless the caller overrides them. Remix has no Ark equivalent and raises. The body is flat apart from a typed content array: the prompt is a text item and every image/video/audio input is a url item carrying a role, which is the only way to express reference-to-video, edit and extension jobs, so a caller-supplied content array wins outright. input_reference becomes a first_frame item, base64 encoded into a data URI for file inputs. OpenAI params translate per Ark's contract: seconds becomes duration and size yields both a ratio and a named resolution tier, with the ratio snapped onto the set Ark accepts (854x480 reduces to 427:240, which Ark rejects) and the tier taken from the shortest side. Emitting only the ratio leaves the request on the model default and bills at that default's rate. Billing reads the delivered duration and resolution rather than the requested ones, because duration -1 and ratio adaptive both let Ark deliver something else. Ark bills per output token (tokens = width * height * 24 * seconds / 1024), so the cost map carries per-second rates by resolution tier, derived from the console's CNY per million token list rates (2.5 70, 2.0 46, 2.0-fast 37, 2.0-mini 23 without a video input) at the repo's 7.0 CNY/USD convention. Requests that do carry a video input bill at the lower rate instead, which litellm's per-second tiers cannot express, so those are overestimated. usage.total_tokens rides along for invoice reconciliation. Ark echoes the delivered geometry only as a resolution tier plus a ratio, and the 2.0 family's rendered frame does not follow the tier name: its 16:9 "480p" canvas is 864x496, not the 854x480 the tier implies, and its other ratios differ from the tier's nominal shortest side too. A size reconstructed from the tier alone therefore reports pixels the model never produces and scales the 480p rate against the wrong pixel count, so both the reported size and the 2.0 480p rates come from the measured frame (a per-family table, sourced from the Seedance 2.0 Mini model card and re-verified against a live gateway). The create response reports the requested duration and resolution, since only the create, edit and remix calls are costed and Ark's create response carries no usage block, so without it every video logs zero spend. A caller that leaves seconds off is not free either: Ark runs its own default length, which is 5s for the 2.0 family, so that is what the create response bills, while the 2.5 family's default is -1 (Ark picks) and stays unbilled until the status response. The create response also echoes the requested seconds and the frame the job will be rendered at. Editing and extending need the source video as a URL: Ark's task query is a GET while litellm's edit prefetch hook only POSTs, so a task id cannot be resolved into one. The model comes from an explicit extra_body.model, else the encoded video id, else the model the router already resolved, with litellm's provider prefix dropped. The credential's api_base is the Anthropic-compat surface (/api/compatible), so the base is cut at the first /api/ marker rather than by a /api/v3 suffix before the video paths append their own version. The cost-map schema test's exception list gains the 2.0 4K tier, whose $1.28/second rate exceeds its $1 sanity bound, alongside the existing runwayml/seedance2 entry. * fix(volcengine): answer 4xx when a Seedance video is downloaded before it finishes Downloading a task that is still queued or running raised a bare ValueError, which the video content handler synthesizes into a 500. The proxy then retried the same poll twice and reported the caller's own early download as a gateway fault, and the router went looking for a model group fallback for a job that was simply unfinished The download path now raises VolcEngineVideoError carrying the status the OpenAI videos API documents for content that is not ready (404), 400 for a task the provider reported as failed, 404 for an expired task record, and 502 for a succeeded task with no artifact. None of them are 5xx, so the router stops retrying them Regression tests drive avideo_content and assert the caller sees NotFoundError or BadRequestError instead of an internal server error * fix(volcengine): clear the type-discipline and override findings on the Seedance video path The module annotated four mutable collections and carried five `# mutable-ok` markers on lines that suppressed nothing, so the type-discipline gate counted LIT001 x4 and LIT013 x5 against the PR. The content-item and task-dump helpers hand back read-only views of what they build, so they now say Mapping. VideoObject.error is declared as a dict by the type it feeds, so that one marker moves onto the annotation it actually suppresses transform_video_status_retrieve_response also dropped the client keyword the base class declares, which basedpyright reported as an incompatible override * fix(volcengine): bill every Seedance video at create and keep its polls on the deployment that made it Ark bills a task when it finishes while litellm bills it at create, so a request whose cost is not known then is rejected: Ark-picked length (duration -1), frames, a video input, and with it edit and extension. An omitted length is sent as 5s, since the 2.5 family otherwise lets Ark pick it and logged $0. Video ids now carry the requested and the served model group, read from litellm_metadata, so polls and downloads go back to the deployment that created the task. The routing and input_reference helpers move out of the DashScope config into a shared module both providers use. A download before the render finishes answers 400, unreadable or unmapped task bodies answer 502, malformed seconds, size, parameters and Ark fields raise unless drop_params is set, an unpriced resolution tier is refused, and the rendered frame table moves into the cost map. Co-Authored-By: AaronHowell <111272993+AaronHowell@users.noreply.github.com> Co-Authored-By: Claude Opus 5.5 (1M context) * fix(volcengine): reject a text item in content so the guarded prompt is the only text sent Input guardrails scan data["prompt"], so a caller could pass a harmless prompt for scanning and a different, unscanned text item in content, which this branch kept and the guarded prompt was dropped. The prompt now always leads the content array and a second text item answers 400. Co-Authored-By: AaronHowell <111272993+AaronHowell@users.noreply.github.com> Co-Authored-By: Claude Opus 5.5 (1M context) * fix(volcengine): route video to VOLCENGINE_API_BASE when the deployment sets no api_base * fix(volcengine): let a model with no tier prices keep its resolution instead of rejecting every tier * fix(volcengine): answer 400 for malformed or unknown create fields, send a 720p default tier, and stop the list cursor * fix(volcengine): map size to the closest tier by pixel count, reject draft, remix and list cursors, map user to safety_identifier, read ARK_API_BASE * fix(volcengine): drop unsupported create values under drop_params, take OpenAI's object input_reference, default to 4 seconds, price endpoint ids by base_model * fix(videos): unwrap nested file tuples without recursion * fix(volcengine): price Seedance tiers at Ark's own per-resolution token rates Ark prices 1080p and 4k per token above and below the 480p/720p rate, so the single 46 CNY per million tokens the rows were derived from billed 2.0 1080p and 2.5 1080p about 10% low and 2.0 4k 77% high. Ark's per-resolution table and its own 16:9 examples agree: 1080p 2.48 CNY/s, 4k 5.05 CNY/s for 2.0 and 3.74 CNY/s for 2.5 1080p, which is 51, 26 and 77 CNY per million tokens. Rewrites the tier test to check each rate against the rendered pixels at that tier's own price instead of assuming every tier scales from the base rate, and drops the 4k cost exception that only existed for the old rate. Co-Authored-By: Bytechoreographer Co-Authored-By: AaronHowell <111272993+AaronHowell@users.noreply.github.com> * fix(volcengine): unpack Ark parameters sent through extra_body * fix(volcengine): ignore the extra_body keys litellm maps itself The shared request builder merges the caller's raw extra_body back over the mapped params, so a key this adapter maps itself comes back a second time as a field of its own, and the unknown field check turned the supported request into a 400. 2aefc18b0a covered parameters; size, seconds and user still failed that way. All three are already mapped from the top-level copy the same builder writes, so the leftover is ignored rather than rejected. characters and image stay rejected, since Ark has no equivalent and accepting them would hide a parameter that never took effect. Co-Authored-By: Bytechoreographer Co-Authored-By: AaronHowell <111272993+AaronHowell@users.noreply.github.com> * fix(volcengine): reject litellm request fields inside parameters and accept litellm's own controls A seconds, size, user, prompt or nested parameters inside an Ark parameters object, and an extra_body copy of prompt, model or extra_headers, answer 400 again unless drop_params is on, instead of being ignored. A user that is not a string answers 400 the same way. drop_params, additional_drop_params and a forwarded organization are no longer rejected as unknown Ark fields, and additional_drop_params removes the fields it names. * fix(volcengine): apply additional_drop_params before mapping request fields A drop list naming user, seconds, size, parameters or an unsupported field such as image had no effect, because the fields were already mapped to Ark names before the list was applied. The list now removes them first, so the request sends what it would send if those fields were never set --------- Co-authored-by: AaronHowell <111272993+AaronHowell@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 (1M context) Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Co-authored-by: zijie.zhou Co-authored-by: Bytechoreographer --- .../llms/base_llm/videos/async_task_utils.py | 123 ++ .../llms/dashscope/videos/transformation.py | 102 +- litellm/llms/volcengine/videos/__init__.py | 0 .../llms/volcengine/videos/transformation.py | 904 +++++++++++++ ...odel_prices_and_context_window_backup.json | 55 + litellm/utils.py | 4 + model_prices_and_context_window.json | 55 + provider_endpoints_support.json | 3 +- .../test_dashscope_video_transformation.py | 3 +- .../test_volcengine_video_transformation.py | 1131 +++++++++++++++++ 10 files changed, 2284 insertions(+), 96 deletions(-) create mode 100644 litellm/llms/base_llm/videos/async_task_utils.py create mode 100644 litellm/llms/volcengine/videos/__init__.py create mode 100644 litellm/llms/volcengine/videos/transformation.py create mode 100644 tests/unit/llms/volcengine/test_volcengine_video_transformation.py diff --git a/litellm/llms/base_llm/videos/async_task_utils.py b/litellm/llms/base_llm/videos/async_task_utils.py new file mode 100644 index 00000000000..d2fb3ed7606 --- /dev/null +++ b/litellm/llms/base_llm/videos/async_task_utils.py @@ -0,0 +1,123 @@ +import base64 +from collections.abc import Mapping +from io import BufferedReader, BytesIO +from itertools import accumulate, repeat +from pathlib import Path, PurePath +from typing import Final + +from pydantic import TypeAdapter, ValidationError + +from litellm.images.utils import ImageEditRequestUtils +from litellm.types.videos.utils import decode_video_id_with_provider + +_FIELDS_ADAPTER: Final = TypeAdapter(dict[str, object]) +_FILE_TUPLE_ADAPTER: Final = TypeAdapter(tuple[object, ...]) + + +def _fields(value: object) -> Mapping[str, object] | None: + try: + return _FIELDS_ADAPTER.validate_python(value) + except ValidationError: + return None + + +def _logged_litellm_params(logging_obj: object) -> Mapping[str, object]: + return _fields(getattr(logging_obj, "litellm_params", None)) or {} + + +def _id_safe_name(name: object) -> str | None: + return name if isinstance(name, str) and name and ";" not in name else None + + +def _stamped_models(metadata: object) -> tuple[str, str | None] | None: + fields: Final = _fields(metadata) + if fields is None: + return None + requested: Final = _id_safe_name(fields.get("model_group")) + served: Final = _id_safe_name(fields.get("deployment_model_name")) + if requested is None: + return None + return requested, None if served is None or served == requested or "*" in served else served + + +def routed_video_models(logging_obj: object) -> tuple[str, str | None] | None: + """ + A later status or content request checks the key's model access against the model encoded in the id, and + routes by it. One upstream model name can be served by several deployments (regions, accounts), so only the + group that served the task can send its polls back to the deployment that created it, while a key may be + scoped to the alias, routing group or wildcard name it asked for instead. The id carries both names: the one + the router was asked for on this hop, and the group that served it when that differs and is not a wildcard. + The router stamps both into litellm_metadata, and metadata carries a copy only when the caller sent none. + """ + params: Final = _logged_litellm_params(logging_obj) + stamped: Final = (_stamped_models(params.get(key)) for key in ("litellm_metadata", "metadata")) + return next((models for models in stamped if models), None) + + +def polled_video_models(logging_obj: object) -> tuple[str | None, str | None]: + """ + The served group survives into the returned id only when the router sent this poll to it, so an edited id + cannot write any name it likes back. + """ + video_id: Final = _logged_litellm_params(logging_obj).get("video_id") + if not isinstance(video_id, str): + return None, None + decoded: Final = decode_video_id_with_provider(video_id) + served: Final = decoded.get("served_model") + routed: Final = routed_video_models(logging_obj) + return decoded.get("model_id") or None, served if routed is None or served in routed else None + + +def _tuple_payload(file_obj: object) -> object: + if not isinstance(file_obj, tuple): + return file_obj + file_tuple: Final = _FILE_TUPLE_ADAPTER.validate_python(file_obj) + return file_tuple[1] if len(file_tuple) >= 2 else None + + +def _file_content(file_obj: object) -> object: + unwrapped: Final = accumulate(repeat(None), lambda current, _: _tuple_payload(current), initial=file_obj) + return next(item for item in unwrapped if not isinstance(item, tuple)) + + +def input_reference_bytes(file_obj: object) -> bytes: + content: Final = _file_content(file_obj) + if isinstance(content, PurePath): + return Path(content).read_bytes() + if isinstance(content, (BytesIO, BufferedReader)): + current_position: Final = content.tell() + _ = content.seek(0) + data: Final = content.read() + _ = content.seek(current_position) + return data + if isinstance(content, bytes): + return content + if isinstance(content, bytearray): + return bytes(content) + read: Final = getattr(content, "read", None) + if callable(read): + read_data: Final[object] = read() + if isinstance(read_data, bytes): + return read_data + raise ValueError("input_reference must be a URL string, bytes, a path, a file object, or a file tuple") + + +def _reference_object_url(reference: Mapping[str, object]) -> str: + image_url: Final = reference.get("image_url") + if isinstance(image_url, str) and image_url: + return image_url + raise ValueError("input_reference as an object must carry an image_url, a file_id is not supported") + + +def input_reference_url(image: object) -> str: + """A string is a URL the provider fetches itself and is never opened as a local path. OpenAI's object form + carries that URL, or a data URL, in ``image_url``.""" + if isinstance(image, str): + return image + reference: Final = _fields(image) + if reference is not None: + return _reference_object_url(reference) + content: Final = input_reference_bytes(image) + content_type: Final = ImageEditRequestUtils.get_image_content_type(content) + encoded: Final = base64.b64encode(content).decode("utf-8") + return f"data:{content_type};base64,{encoded}" diff --git a/litellm/llms/dashscope/videos/transformation.py b/litellm/llms/dashscope/videos/transformation.py index 17b685a0e1a..1db91cbb127 100644 --- a/litellm/llms/dashscope/videos/transformation.py +++ b/litellm/llms/dashscope/videos/transformation.py @@ -3,13 +3,10 @@ DashScope (Alibaba Model Studio) async video task API: create, poll and download delete or remix endpoint, so those raise. """ -import base64 import json from collections.abc import Mapping from datetime import datetime -from io import BufferedReader, BytesIO from math import log -from pathlib import Path, PurePath from types import MappingProxyType from typing import TYPE_CHECKING, Annotated, Final @@ -19,9 +16,13 @@ from pydantic import BaseModel, Field, StrictBool, StrictInt, StrictStr, TypeAda import litellm from litellm.exceptions import UnsupportedParamsError -from litellm.images.utils import ImageEditRequestUtils from litellm.litellm_core_utils.url_utils import encode_url_path_segment from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.base_llm.videos.async_task_utils import ( + input_reference_url, + polled_video_models, + routed_video_models, +) from litellm.llms.base_llm.videos.transformation import BaseVideoConfig from litellm.llms.custom_httpx.http_handler import ( _get_httpx_client, # pyright: ignore[reportPrivateUsage, reportUnknownVariableType] # house cached-client factory has no public alias and its stub leaves params untyped @@ -32,7 +33,6 @@ from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import ModelInfo from litellm.types.videos.main import VideoCreateOptionalRequestParams, VideoObject from litellm.types.videos.utils import ( - decode_video_id_with_provider, encode_video_id_with_provider, extract_original_video_id, ) @@ -86,7 +86,6 @@ _EMPTY_PARAMS: Final[dict[str, object]] = {} # mutable-ok: BaseVideoConfig cont _TASK_RESPONSE_ADAPTER: Final = TypeAdapter(_TaskResponse) _PARAMETERS_ADAPTER: Final = TypeAdapter(dict[str, object]) -_FILE_TUPLE_ADAPTER: Final = TypeAdapter(tuple[object, ...]) _STATUS_MAP: Final = MappingProxyType( { @@ -383,39 +382,6 @@ def _is_priced_tier(resolution: object, priced_tiers: tuple[str, ...]) -> bool: return isinstance(resolution, str) and resolution.strip().upper() in priced_tiers -def _read_all_bytes(file_obj: object) -> bytes: - if isinstance(file_obj, tuple): - file_tuple: Final = _FILE_TUPLE_ADAPTER.validate_python(file_obj) - return _read_all_bytes(file_tuple[1] if len(file_tuple) >= 2 else None) - if isinstance(file_obj, PurePath): - return Path(file_obj).read_bytes() - if isinstance(file_obj, (BytesIO, BufferedReader)): - current_position: Final = file_obj.tell() - file_obj.seek(0) - content: Final = file_obj.read() - file_obj.seek(current_position) - return content - if isinstance(file_obj, bytes): - return file_obj - if isinstance(file_obj, bytearray): - return bytes(file_obj) - read: Final = getattr(file_obj, "read", None) - if callable(read): - data: Final = read() - if isinstance(data, bytes): - return data - raise ValueError("input_reference must be a URL string, bytes, a path, a file object, or a file tuple") - - -def _image_url(image: object) -> str: - if isinstance(image, str): - return image - content: Final = _read_all_bytes(image) - content_type: Final = ImageEditRequestUtils.get_image_content_type(content) - encoded: Final = base64.b64encode(content).decode("utf-8") - return f"data:{content_type};base64,{encoded}" - - def _media_reference_type(model: str) -> str | None: """Models without the field predate the media array and take the image on the flat ``img_url`` field.""" info: Final = _dashscope_model_info(model) @@ -558,58 +524,6 @@ def _video_url_from_task(task: _TaskResponse, raw_response: httpx.Response) -> s raise _video_error(raw_response, 502, "Video URL not found in task response. The task may not have succeeded yet.") -def _logged_litellm_params(logging_obj: object) -> Mapping[str, object]: - litellm_params: Final[object] = getattr(logging_obj, "litellm_params", None) - try: - return _PARAMETERS_ADAPTER.validate_python(litellm_params) - except ValidationError: - return _EMPTY_PARAMS - - -def _polled_models(logging_obj: object) -> tuple[str | None, str | None]: - """ - The proxy routes by the models encoded in the polled id, so they must survive into the returned id. The served - group survives only when the router sent this poll to it, so an edited id cannot write any name it likes back. - """ - video_id: Final = _logged_litellm_params(logging_obj).get("video_id") - if not isinstance(video_id, str): - return None, None - decoded: Final = decode_video_id_with_provider(video_id) - served: Final = decoded.get("served_model") - routed: Final = _poll_models(logging_obj) - return decoded.get("model_id") or None, served if routed is None or served in routed else None - - -def _poll_models(logging_obj: object) -> tuple[str, str | None] | None: - """ - A later status or content request checks the key's model access against the model encoded in the id, and - routes by it. DashScope serves the same model name from several regional hosts, so only the group that - served the task can send its polls back to the region that created it, while a key may be scoped to the - alias, routing group or wildcard name it asked for instead. The id carries both names: the one the router - was asked for on this hop, and the group that served it when that differs and is not a wildcard. The router - stamps both into litellm_metadata, and metadata carries a copy only when the caller sent none. - """ - params: Final = _logged_litellm_params(logging_obj) - stamped: Final = (_stamped_poll_models(params.get(key)) for key in ("litellm_metadata", "metadata")) - return next((models for models in stamped if models), None) - - -def _stamped_poll_models(metadata: object) -> tuple[str, str | None] | None: - try: - fields: Final = _PARAMETERS_ADAPTER.validate_python(metadata) - except ValidationError: - return None - requested: Final = _id_safe_name(fields.get("model_group")) - served: Final = _id_safe_name(fields.get("deployment_model_name")) - if requested is None: - return None - return requested, None if served is None or served == requested or "*" in served else served - - -def _id_safe_name(name: object) -> str | None: - return name if isinstance(name, str) and name and ";" not in name else None - - class DashScopeVideoConfig(BaseVideoConfig): def get_supported_openai_params(self, model: str) -> list[str]: # mutable-ok: BaseVideoConfig contract returns list return [ @@ -804,7 +718,7 @@ class DashScopeVideoConfig(BaseVideoConfig): input_reference: Final = video_create_optional_request_params.get("input_reference") if input_reference is None: return None, None - image_url: Final = _image_url(input_reference) + image_url: Final = input_reference_url(input_reference) media_type: Final = _media_reference_type(model) if media_type is None: return "img_url", image_url @@ -823,7 +737,7 @@ class DashScopeVideoConfig(BaseVideoConfig): requested: Final = self._requested_parameters(request_data) video_obj: Final = _video_object_from_task(task, raw_response, model=model, requested=requested) if custom_llm_provider and video_obj.id: - requested_model, served_model = _poll_models(logging_obj) or (model, None) + requested_model, served_model = routed_video_models(logging_obj) or (model, None) video_obj.id = encode_video_id_with_provider( video_obj.id, custom_llm_provider, requested_model, served_model ) @@ -869,7 +783,7 @@ class DashScopeVideoConfig(BaseVideoConfig): client: "HTTPHandler | None" = None, ) -> VideoObject: task: Final = _parse_task_response(raw_response) - model, served_model = _polled_models(logging_obj) + model, served_model = polled_video_models(logging_obj) video_obj: Final = _video_object_from_task(task, raw_response, model=model) if custom_llm_provider and video_obj.id: video_obj.id = encode_video_id_with_provider(video_obj.id, custom_llm_provider, model, served_model) diff --git a/litellm/llms/volcengine/videos/__init__.py b/litellm/llms/volcengine/videos/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/llms/volcengine/videos/transformation.py b/litellm/llms/volcengine/videos/transformation.py new file mode 100644 index 00000000000..a75a4185186 --- /dev/null +++ b/litellm/llms/volcengine/videos/transformation.py @@ -0,0 +1,904 @@ +import json +from collections.abc import Mapping, Sequence +from math import log +from types import MappingProxyType +from typing import TYPE_CHECKING, Final, TypeVar + +import httpx +from httpx._types import RequestFiles +from pydantic import BaseModel, StrictBool, StrictInt, StrictStr, TypeAdapter, ValidationError + +import litellm +from litellm.exceptions import UnsupportedParamsError +from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.base_llm.videos.async_task_utils import ( + input_reference_url, + polled_video_models, + routed_video_models, +) +from litellm.llms.base_llm.videos.transformation import BaseVideoConfig +from litellm.llms.custom_httpx.http_handler import ( + AsyncHTTPHandler, + HTTPHandler, + _get_httpx_client, # pyright: ignore[reportPrivateUsage, reportUnknownVariableType] # house cached-client factory has no public alias and its stub leaves params untyped + get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] # factory stub leaves params untyped +) +from litellm.llms.volcengine.common_utils import get_volcengine_base_url +from litellm.secret_managers.main import get_secret_str +from litellm.types.router import GenericLiteLLMParams +from litellm.types.utils import ModelInfo +from litellm.types.videos.main import VideoCreateOptionalRequestParams, VideoObject +from litellm.types.videos.utils import encode_video_id_with_provider, extract_original_video_id +from litellm.utils import get_model_info + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging + + +_ParsedT: Final = TypeVar("_ParsedT") + + +class VolcEngineVideoError(BaseLLMException): + pass + + +class _TaskError(BaseModel, frozen=True): + code: str | None = None + message: str | None = None + + +class _TaskContent(BaseModel, frozen=True): + video_url: str | None = None + + +class _TaskUsage(BaseModel, frozen=True): + total_tokens: int | None = None + + +class _Task(BaseModel, frozen=True): + id: str + status: str + model: str | None = None + error: _TaskError | None = None + content: _TaskContent | None = None + usage: _TaskUsage | None = None + created_at: int | None = None + updated_at: int | None = None + resolution: str | None = None + ratio: str | None = None + duration: int | float | None = None + + +class _TaskCreateResponse(BaseModel, frozen=True): + id: str | None = None + error: _TaskError | None = None + + +class _TaskList(BaseModel, frozen=True): + items: Sequence[_Task] = () + total: int | None = None + + +class _ContentItem(BaseModel, frozen=True): + type: str | None = None + role: str | None = None + + +class _KnownFields(BaseModel, frozen=True): + resolution: StrictStr | None = None + ratio: StrictStr | None = None + seed: StrictInt | None = None + camera_fixed: StrictBool | None = None + watermark: StrictBool | None = None + generate_audio: StrictBool | None = None + + +VOLCENGINE_VIDEO_TASKS_PATH: Final = "/api/v3/contents/generations/tasks" +VOLCENGINE_VIDEO_DEFAULT_DURATION_SECONDS: Final = 4 +VOLCENGINE_VIDEO_DEFAULT_RESOLUTION: Final = "720p" +_VOLCENGINE_API_MARKER: Final = "/api/" +_VOLCENGINE_COMPAT_SUFFIXES: Final = ("/api", "/v1") +_FRAMES_FIELD: Final = "volcengine_video_frames" + +_EMPTY_PARAMS: Final[dict[str, object]] = {} # mutable-ok: BaseVideoConfig contract empty params + +_TASK_ADAPTER: Final = TypeAdapter(_Task) +_TASK_CREATE_ADAPTER: Final = TypeAdapter(_TaskCreateResponse) +_TASK_LIST_ADAPTER: Final = TypeAdapter(_TaskList) +_FIELDS_ADAPTER: Final = TypeAdapter(dict[str, object]) +_CONTENT_ITEMS_ADAPTER: Final = TypeAdapter(tuple[dict[str, object], ...]) +_FRAMES_ADAPTER: Final = TypeAdapter(dict[str, dict[str, str]]) +_EMPTY_FRAMES: Final[Mapping[str, Mapping[str, str]]] = MappingProxyType({}) + +_STATUS_MAP: Final = MappingProxyType( + { + "queued": "queued", + "running": "in_progress", + "succeeded": "completed", + "failed": "failed", + "cancelled": "cancelled", + "expired": "failed", + } +) +_PENDING_STATUSES: Final = frozenset({"queued", "running"}) +_TERMINAL_STATUSES: Final = frozenset({"completed", "failed", "cancelled"}) + +_CREATE_PASSTHROUGH_KEYS: Final = ( + "resolution", + "ratio", + "duration", + "seed", + "camera_fixed", + "watermark", + "generate_audio", + "service_tier", + "priority", + "output_format", + "execution_expires_after", + "callback_url", + "safety_identifier", + "tools", + "omni_reference_task_type", +) + +_REQUEST_METADATA: Final = frozenset({"model", "prompt", "user", "extra_headers", "extra_query", "extra_body"}) +_UNSUPPORTED_BY_VOLCENGINE: Final = frozenset({"characters", "image"}) +_MAPPED_ELSEWHERE: Final = frozenset({*_REQUEST_METADATA, *_UNSUPPORTED_BY_VOLCENGINE, "seconds", "size", "parameters"}) +_READ_FROM_THE_REQUEST: Final = _MAPPED_ELSEWHERE - _UNSUPPORTED_BY_VOLCENGINE +_LITELLM_CONTROLS: Final = frozenset({"drop_params", "additional_drop_params", "organization"}) +_DROP_KEYS_ADAPTER: Final = TypeAdapter(frozenset[str]) + +# litellm/videos/utils.py merges the caller's raw ``extra_body`` back over the mapped params, which returns the keys +# this adapter maps itself a second time, as fields of their own. They are already accounted for, so they count as +# handled rather than unknown. +_MAPPED_EXTRA_BODY_COPIES: Final = frozenset({"size", "seconds", "user"}) +_CREATE_HANDLED_KEYS: Final = frozenset( + { + *_CREATE_PASSTHROUGH_KEYS, + "content", + "input_reference", + "frames", + *_MAPPED_EXTRA_BODY_COPIES, + *_LITELLM_CONTROLS, + } +) + +_VIDEO_INPUT_UNSUPPORTED: Final = ( + "{operation} is not supported through litellm: Ark bills a task with a video input at a different rate per " + "token, and litellm bills the video when the task is created, before that cost is known" +) +_INPUT_VIDEO_ROLE: Final = "reference_video" +_INPUT_VIDEO_TYPE: Final = "video_url" + +_RATIO_SET: Final = ("16:9", "9:16", "4:3", "3:4", "1:1", "21:9") +_RESOLUTION_SHORTEST_SIDE: Final = MappingProxyType({"480p": 480, "720p": 720, "1080p": 1080, "4k": 2160}) +_TIER_FRAME_ASPECT: Final = 16 / 9 + + +def _video_error(raw_response: httpx.Response, status_code: int, message: str) -> VolcEngineVideoError: + """The provider's own error status wins, and a body that failed on a 2xx answers the caller's status instead, + since a 2xx error has the proxy retry a call the provider already rejected.""" + return VolcEngineVideoError( + status_code=raw_response.status_code if raw_response.status_code >= 400 else status_code, + message=message, + headers=raw_response.headers, + ) + + +def _parse(adapter: TypeAdapter[_ParsedT], raw_response: httpx.Response) -> _ParsedT: + try: + return adapter.validate_python(raw_response.json()) + except (json.JSONDecodeError, ValidationError) as e: + raise _video_error(raw_response, 502, f"Unexpected Volcengine video response: {e}") from e + + +def _unsupported(model: str, message: str) -> UnsupportedParamsError: + return UnsupportedParamsError(message=message, model=model, llm_provider="volcengine") + + +def _droppable(model: str, message: str) -> UnsupportedParamsError: + return _unsupported(model, f"{message}, set drop_params=True to ignore it") + + +def _normalized_api_base(api_base: str) -> str: + """The Ark credential is shared with the chat, responses and Anthropic surfaces, so its api_base may carry any of + their paths, and the video paths carry their own version segment.""" + trimmed: Final = api_base.rstrip("/") + marker_index: Final = trimmed.rfind(_VOLCENGINE_API_MARKER) + if marker_index != -1: + return trimmed[:marker_index] + suffix: Final = next((suffix for suffix in _VOLCENGINE_COMPAT_SUFFIXES if trimmed.endswith(suffix)), None) + return trimmed[: -len(suffix)] if suffix is not None else trimmed + + +def _size_dimensions(size: str) -> tuple[int, int] | None: + width_str, separator, height_str = size.partition("x") + if not separator or not (width_str.isdigit() and height_str.isdigit()): + return None + if not int(width_str) or not int(height_str): + return None + return int(width_str), int(height_str) + + +def _ratio_value(ratio: str) -> float: + width_str, _, height_str = ratio.partition(":") + return int(width_str) / int(height_str) + + +def _ratio_from_size(size: str) -> str | None: + """Ark only takes its named ratios, and its own tiers do not reduce onto them (854x480 is 427:240), so the + closest ratio wins.""" + dimensions: Final = _size_dimensions(size) + if dimensions is None: + return None + width, height = dimensions + return min(_RATIO_SET, key=lambda candidate: abs(log(width / height / _ratio_value(candidate)))) + + +def _resolution_from_size(size: str) -> str | None: + """Ark keeps a tier's pixel count across its ratios (960x960 and 1280x720 are both 720p), so the tier whose 16:9 + frame has the closest pixel count wins.""" + dimensions: Final = _size_dimensions(size) + if dimensions is None: + return None + width, height = dimensions + return min( + _RESOLUTION_SHORTEST_SIDE, + key=lambda tier: abs(log(width * height / (_RESOLUTION_SHORTEST_SIDE[tier] ** 2 * _TIER_FRAME_ASPECT))), + ) + + +def _duration_param(seconds: object) -> int | None: + if isinstance(seconds, bool): + return None + if isinstance(seconds, int): + return seconds + if not isinstance(seconds, (str, float)): + return None + try: + parsed: Final = float(seconds) + except ValueError: + return None + return int(parsed) if parsed.is_integer() else None + + +def _duration_fields(seconds: object, model: str, drop_params: bool) -> Mapping[str, object]: + if seconds is None: + return _EMPTY_PARAMS + duration: Final = _duration_param(seconds) + if duration is not None: + return MappingProxyType({"duration": duration}) + if drop_params: + return _EMPTY_PARAMS + raise _droppable(model, f"seconds must be a whole number of seconds, got {seconds!r}") + + +def _size_fields(size: object, model: str, drop_params: bool) -> Mapping[str, object]: + if size is None: + return _EMPTY_PARAMS + ratio: Final = _ratio_from_size(size) if isinstance(size, str) else None + resolution: Final = _resolution_from_size(size) if isinstance(size, str) else None + if ratio is not None and resolution is not None: + return MappingProxyType({"ratio": ratio, "resolution": resolution}) + if drop_params: + return _EMPTY_PARAMS + raise _droppable(model, f"size must be WIDTHxHEIGHT, got {size!r}") + + +def _parameters_fields(parameters: object, model: str, drop_params: bool) -> Mapping[str, object]: + if parameters is None: + return _EMPTY_PARAMS + try: + fields: Final = _FIELDS_ADAPTER.validate_python(parameters) + except ValidationError: + if drop_params: + return _EMPTY_PARAMS + raise _droppable( + model, f"parameters must be an object of Ark request fields, got {type(parameters).__name__}" + ) from None + misplaced: Final = tuple(sorted(key for key in fields if key in _READ_FROM_THE_REQUEST)) + if misplaced and not drop_params: + raise _droppable(model, f"send {', '.join(misplaced)} as top-level request fields, not inside parameters") + return MappingProxyType({key: value for key, value in fields.items() if key not in _READ_FROM_THE_REQUEST}) + + +def _user_fields(user: object, model: str, drop_params: bool) -> Mapping[str, object]: + if isinstance(user, str) and user: + return MappingProxyType({"safety_identifier": user}) + if user is None or isinstance(user, str) or drop_params: + return _EMPTY_PARAMS + raise _droppable(model, f"user must be a string, got {type(user).__name__}") + + +def _additional_drop_keys(params: Mapping[str, object]) -> frozenset[str]: + try: + return _DROP_KEYS_ADAPTER.validate_python(params.get("additional_drop_params") or ()) + except ValidationError: + return frozenset() + + +def _without(params: Mapping[str, object], keys: frozenset[str]) -> Mapping[str, object]: + return MappingProxyType({key: value for key, value in params.items() if key not in keys}) + + +def _malformed_fields(fields: Mapping[str, object]) -> tuple[tuple[str, str], ...]: + try: + _ = _KnownFields.model_validate(dict(fields)) + except ValidationError as e: + return tuple((str(error["loc"][0]), error["msg"]) for error in e.errors() if error["loc"]) + return () + + +def _drops_params(litellm_params: GenericLiteLLMParams) -> bool: + return litellm_params.drop_params is True or litellm.drop_params is True + + +def _model_info(model: str) -> ModelInfo | None: + try: + return get_model_info(model=model, custom_llm_provider="volcengine") + except Exception: # noqa: BLE001 # get_model_info raises bare Exception for unmapped models + return None + + +def _pricing_model(model: str, litellm_params: GenericLiteLLMParams) -> str: + """A deployment priced through ``model_info.base_model``, like an ``ep-`` endpoint id, bills at that model's + tiers.""" + model_info: Final = _FIELDS_ADAPTER.validate_python(litellm_params["model_info"] or _EMPTY_PARAMS) + base_model: Final = model_info.get("base_model") + return base_model if isinstance(base_model, str) and base_model else model + + +def _priced_tiers(model: str) -> tuple[str, ...] | None: + """A tier the model has no price for is one Ark does not offer it at, and litellm would bill it at the base rate. + A model with no tier price listed, missing from the cost map or registered unpriced by the router, has nothing + to check against.""" + info: Final = _model_info(model) + if info is None: + return None + return ( + tuple(tier for tier in _RESOLUTION_SHORTEST_SIDE if info.get(f"output_cost_per_second_{tier}") is not None) + or None + ) + + +def _delivered_frames(model: str | None) -> Mapping[str, Mapping[str, str]]: + info: Final = _model_info(model) if model else None + provider_specific: Final[object] = info.get("provider_specific_entry") if info is not None else None + frames: Final[object] = provider_specific.get(_FRAMES_FIELD) if isinstance(provider_specific, Mapping) else None + try: + return _FRAMES_ADAPTER.validate_python(frames) if frames is not None else _EMPTY_FRAMES + except ValidationError: + return _EMPTY_FRAMES + + +def _size_from_geometry(resolution: str | None, ratio: str | None, model: str | None) -> str | None: + """Ark echoes a tier and a ratio rather than pixels, and its rendered frame does not always follow the tier's + nominal size, so the frame listed for the model wins over scaling the tier's shortest side up the ratio.""" + if not resolution or not ratio or ":" not in ratio: + return None + tier: Final = resolution.strip().lower() + normalized_ratio: Final = ratio.strip() + listed: Final = _delivered_frames(model).get(tier, MappingProxyType({})).get(normalized_ratio) + if listed is not None: + return listed + shortest_side: Final = _RESOLUTION_SHORTEST_SIDE.get(tier) + width_str, _, height_str = normalized_ratio.partition(":") + if shortest_side is None or not (width_str.isdigit() and height_str.isdigit()): + return None + ratio_width: Final = int(width_str) + ratio_height: Final = int(height_str) + if not ratio_width or not ratio_height: + return None + shortest_ratio_side: Final = min(ratio_width, ratio_height) + return f"{-(-shortest_side * ratio_width // shortest_ratio_side)}x{-(-shortest_side * ratio_height // shortest_ratio_side)}" + + +def _is_input_video(item: _ContentItem) -> bool: + return item.type == _INPUT_VIDEO_TYPE or item.role == _INPUT_VIDEO_ROLE + + +def _supplied_items(supplied: object, model: str, drop_params: bool) -> tuple[Mapping[str, object], ...]: + try: + items: Final = _CONTENT_ITEMS_ADAPTER.validate_python(supplied) + entries: Final = tuple(_ContentItem.model_validate(item) for item in items) + except ValidationError: + if drop_params: + return () + raise _droppable(model, "content must be an array of Ark content items") from None + if any(entry.type == "text" for entry in entries) and not drop_params: + raise _droppable( + model, + "A text item in content is not supported, put the text in prompt since it is the only text input " + "guardrails scan before the request reaches Ark", + ) + if any(_is_input_video(entry) for entry in entries) and not drop_params: + raise _droppable(model, _VIDEO_INPUT_UNSUPPORTED.format(operation="Video input (a reference_video item)")) + kept: Final = (item for item, entry in zip(items, entries) if entry.type != "text" and not _is_input_video(entry)) + return tuple(MappingProxyType(item) for item in kept) + + +def _reference_items(params: Mapping[str, object], model: str, drop_params: bool) -> tuple[Mapping[str, object], ...]: + """A caller-supplied ``content`` array wins over ``input_reference``, since its roles are the only way to express + reference-to-video.""" + supplied: Final = params.get("content") + if supplied is not None: + return _supplied_items(supplied, model, drop_params) + reference: Final = params.get("input_reference") + if reference is None: + return () + try: + url: Final = input_reference_url(reference) + except ValueError as e: + if drop_params: + return () + raise _droppable(model, str(e)) from None + return (MappingProxyType({"type": "image_url", "image_url": {"url": url}, "role": "first_frame"}),) + + +def _sent_content( + prompt: str, params: Mapping[str, object], model: str, drop_params: bool +) -> tuple[Mapping[str, object], ...]: + """The prompt always leads, as it is the only text the input guardrails scan.""" + return (MappingProxyType({"type": "text", "text": prompt}), *_reference_items(params, model, drop_params)) + + +def _seconds_label(seconds: object) -> str | None: + if isinstance(seconds, float) and seconds.is_integer(): + return str(int(seconds)) + return None if seconds is None else str(seconds) + + +def _requested_usage(request_data: Mapping[str, object]) -> Mapping[str, object]: + """The create call is the only billed one and Ark's create response carries no usage, so the requested duration + and tier are what the video bills at.""" + duration: Final = _duration_param(request_data.get("duration")) + resolution: Final = request_data.get("resolution") + return MappingProxyType( + { + key: value + for key, value in ( + ("duration_seconds", float(duration) if duration is not None and duration > 0 else None), + ("video_resolution", resolution.strip().lower() if isinstance(resolution, str) else None), + ) + if value is not None + } + ) + + +def _task_usage(task: _Task) -> Mapping[str, object]: + resolution: Final = task.resolution.strip().lower() if task.resolution else None + return MappingProxyType( + { + key: value + for key, value in ( + ("duration_seconds", float(task.duration) if task.duration is not None else None), + ("video_resolution", resolution), + ("total_tokens", task.usage.total_tokens if task.usage is not None else None), + ) + if value is not None + } + ) + + +def _task_status(task: _Task, raw_response: httpx.Response) -> str: + status: Final = _STATUS_MAP.get(task.status.strip().lower()) + if status is None: # defaulting it would hand back a job the caller polls forever + raise _video_error(raw_response, 502, f"Unknown Volcengine task status: {task.status}") + return status + + +def _video_object_from_task(task: _Task, raw_response: httpx.Response, model: str | None) -> VideoObject: + status: Final = _task_status(task, raw_response) + usage: Final = _task_usage(task) + seconds: Final = usage.get("duration_seconds") + error: Final[dict[str, object] | None] = ( # mutable-ok: VideoObject.error is a dict field + {key: value for key, value in (("code", task.error.code), ("message", task.error.message)) if value is not None} + if task.error is not None + else None + ) + return VideoObject( + id=task.id, + object="video", + status=status, + created_at=task.created_at, + completed_at=task.updated_at if status in _TERMINAL_STATUSES else None, + error=error, + seconds=_seconds_label(seconds), + size=_size_from_geometry(task.resolution, task.ratio, task.model), + model=model or task.model, + usage=dict(usage) if usage else None, + ) + + +def _video_url_from_task(task: _Task, raw_response: httpx.Response) -> str: + """A download before the video is ready is the caller polling too early, not a fault of the deployment, so it + answers 400, which litellm neither retries nor counts toward a deployment cooldown. A 404 would cool the + deployment down, and a 409 or a 5xx would be retried against a job that is simply not finished.""" + if task.content is not None and task.content.video_url: + return task.content.video_url + status: Final = task.status.strip().lower() + if status in _PENDING_STATUSES: + raise _video_error( + raw_response, 400, f"Video is still processing (status: {status}). Please wait and try again." + ) + if task.error is not None and (task.error.message or task.error.code): + raise _video_error(raw_response, 400, f"Video generation failed: {task.error.message or task.error.code}") + if status == "expired": + raise _video_error(raw_response, 400, "Task expired. Ark keeps task records for 7 days after creation.") + if status in ("failed", "cancelled"): + raise _video_error(raw_response, 400, f"Video generation ended as {status} with no video URL.") + raise _video_error(raw_response, 502, "Video URL not found in task response.") + + +class VolcEngineVideoConfig(BaseVideoConfig): + """Volcengine Ark (Doubao-Seedance) on Ark's async task API: create, poll, download from the signed + ``content.video_url``, and list.""" + + def get_supported_openai_params(self, model: str) -> list[str]: # mutable-ok: BaseVideoConfig contract returns list + return [ + "model", + "prompt", + "input_reference", + "seconds", + "size", + "user", + "extra_headers", + "content", + "parameters", + *_CREATE_PASSTHROUGH_KEYS, + ] + + def map_openai_params( + self, + video_create_optional_params: VideoCreateOptionalRequestParams, + model: str, + drop_params: bool, + ) -> dict[str, object]: # mutable-ok: BaseVideoConfig contract + requested: Final = _without(video_create_optional_params, _additional_drop_keys(video_create_optional_params)) + unsupported: Final = next( + (key for key, value in requested.items() if key in _UNSUPPORTED_BY_VOLCENGINE and value is not None), + None, + ) + if unsupported is not None and not drop_params: + raise _droppable(model, f"'{unsupported}' is not supported by Volcengine video generation") + passthrough: Final = MappingProxyType( + {key: value for key, value in requested.items() if value is not None and key not in _MAPPED_ELSEWHERE} + ) + return { + **_user_fields(requested.get("user"), model, drop_params), + **_size_fields(requested.get("size"), model, drop_params), + **_duration_fields(requested.get("seconds"), model, drop_params), + **passthrough, + **_parameters_fields(requested.get("parameters"), model, drop_params), + } + + def _resolve_api_key(self, api_key: str | None) -> str: + resolved_api_key: Final = api_key or get_secret_str("ARK_API_KEY") or get_secret_str("VOLCENGINE_API_KEY") + if resolved_api_key is None: + raise ValueError( + "Volcengine API key is required. Set ARK_API_KEY or VOLCENGINE_API_KEY environment variable or pass " + "api_key parameter." + ) + return resolved_api_key + + def validate_environment( + self, + headers: dict[str, str], # mutable-ok: BaseVideoConfig contract; handler expects a mutable headers dict + model: str, + api_key: str | None = None, + litellm_params: GenericLiteLLMParams | None = None, + ) -> dict[str, str]: # mutable-ok: BaseVideoConfig contract + resolved_api_key: Final = self._resolve_api_key( + api_key + or (litellm_params.api_key if litellm_params is not None and litellm_params.api_key else None) + or litellm.api_key + ) + headers.update({"Authorization": f"Bearer {resolved_api_key}", "Content-Type": "application/json"}) + return headers + + def get_complete_url( + self, + model: str, + api_base: str | None, + litellm_params: dict[str, object], # mutable-ok: BaseVideoConfig contract + ) -> str: + return _normalized_api_base( + get_volcengine_base_url(api_base or get_secret_str("VOLCENGINE_API_BASE") or get_secret_str("ARK_API_BASE")) + ) + + def transform_video_create_request( + self, + model: str, + prompt: str, + api_base: str, + video_create_optional_request_params: dict[str, object], # mutable-ok: BaseVideoConfig contract + litellm_params: GenericLiteLLMParams, + headers: dict[str, str], # mutable-ok: BaseVideoConfig contract + ) -> tuple[dict[str, object], RequestFiles, str]: # mutable-ok: BaseVideoConfig contract + """ + Ark bills a video when the task finishes, while litellm bills it here at create, so every request has to + carry a length and a tier whose cost is known now. Ark-picked length (``duration: -1``) and a video input + (edits, extensions and reference videos, billed at a different rate per token) are rejected for that reason, + and an omitted length is sent as OpenAI's documented default of 4 seconds, since the 2.5 family's own default + is ``-1``. + """ + drop_params: Final = _drops_params(litellm_params) + dropped_by_config: Final = _additional_drop_keys(video_create_optional_request_params) + kept: Final = _without(video_create_optional_request_params, dropped_by_config) + params: Final = _without( + { + **{key: value for key, value in kept.items() if key != "parameters"}, + **_parameters_fields(kept.get("parameters"), model, drop_params), + }, + dropped_by_config, + ) + content: Final = _sent_content(prompt, params, model, drop_params) + if params.get("frames") is not None and not drop_params: + raise _droppable(model, "frames is not supported through litellm, pass seconds instead") + requested_duration: Final = params.get("duration") + parsed_duration: Final = None if requested_duration is None else _duration_param(requested_duration) + if requested_duration is not None and (parsed_duration is None or parsed_duration <= 0) and not drop_params: + raise _droppable( + model, + f"duration must be a positive whole number of seconds, got {requested_duration!r}. Ark-picked length " + "(-1) is not supported through litellm, since the video is billed when the task is created, before " + "Ark has picked its length", + ) + duration: Final = ( + parsed_duration + if parsed_duration is not None and parsed_duration > 0 + else VOLCENGINE_VIDEO_DEFAULT_DURATION_SECONDS + ) + + unknown_fields: Final = tuple( + sorted(key for key, value in params.items() if value is not None and key not in _CREATE_HANDLED_KEYS) + ) + if unknown_fields and not drop_params: + raise _droppable(model, f"litellm does not send {', '.join(unknown_fields)} to Ark video generation") + requested_fields: Final = MappingProxyType( + {key: params[key] for key in _CREATE_PASSTHROUGH_KEYS if params.get(key) is not None} + ) + malformed: Final = _malformed_fields(requested_fields) + if malformed and not drop_params: + raise _droppable(model, "; ".join(f"invalid {field}: {message}" for field, message in malformed)) + resolution: Final = requested_fields.get("resolution") + priced_tiers: Final = _priced_tiers(_pricing_model(model, litellm_params)) + drops_resolution: Final = ( + priced_tiers is not None and isinstance(resolution, str) and resolution.strip().lower() not in priced_tiers + ) + if drops_resolution and not drop_params: + raise _droppable( + model, f"{model} is not offered at resolution {resolution!r}, only at {', '.join(priced_tiers or ())}" + ) + dropped: Final = frozenset({*(field for field, _ in malformed), *(("resolution",) if drops_resolution else ())}) + kept_fields: Final = MappingProxyType( + {key: value for key, value in requested_fields.items() if key not in dropped} + ) + default_tier: Final = ( + MappingProxyType({"resolution": VOLCENGINE_VIDEO_DEFAULT_RESOLUTION}) + if "resolution" not in kept_fields + and priced_tiers is not None + and VOLCENGINE_VIDEO_DEFAULT_RESOLUTION in priced_tiers + else _EMPTY_PARAMS + ) + + request_data: Final[dict[str, object]] = { # mutable-ok: request body dict, JSON-serialized by the handler + "model": model, + "content": [dict(item) for item in content], + **default_tier, + **kept_fields, + "duration": duration, + } + return request_data, (), f"{api_base}{VOLCENGINE_VIDEO_TASKS_PATH}" + + def transform_video_create_response( + self, + model: str, + raw_response: httpx.Response, + logging_obj: "Logging", + custom_llm_provider: str | None = None, + request_data: dict[str, object] | None = None, # mutable-ok: BaseVideoConfig contract + ) -> VideoObject: + task: Final = _parse(_TASK_CREATE_ADAPTER, raw_response) + if not task.id: + message: Final = task.error.message if task.error is not None and task.error.message else "no task id" + raise _video_error(raw_response, 400, f"Volcengine did not return a video task id: {message}") + requested: Final = request_data or _EMPTY_PARAMS + usage: Final = _requested_usage(requested) + seconds: Final = usage.get("duration_seconds") + requested_model, served_model = routed_video_models(logging_obj) or (model, None) + resolution: Final = requested.get("resolution") + ratio: Final = requested.get("ratio") + return VideoObject( + id=encode_video_id_with_provider(task.id, custom_llm_provider, requested_model, served_model) + if custom_llm_provider + else task.id, + object="video", + status="queued", + model=model, + seconds=_seconds_label(seconds), + size=_size_from_geometry( + resolution if isinstance(resolution, str) else None, ratio if isinstance(ratio, str) else None, model + ), + usage=dict(usage) if usage else None, + ) + + @staticmethod + def _task_url(video_id: str, api_base: str) -> str: + original_task_id: Final = extract_original_video_id(video_id) + encoded_task_id: Final = encode_url_path_segment(original_task_id, field_name="video_id") + return f"{api_base}{VOLCENGINE_VIDEO_TASKS_PATH}/{encoded_task_id}" + + def transform_video_status_retrieve_request( + self, + video_id: str, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: dict[str, str], # mutable-ok: BaseVideoConfig contract + ) -> tuple[str, dict[str, object]]: # mutable-ok: BaseVideoConfig contract + return self._task_url(video_id, api_base), _EMPTY_PARAMS + + def transform_video_status_retrieve_response( + self, + raw_response: httpx.Response, + logging_obj: "Logging", + custom_llm_provider: str | None = None, + client: "HTTPHandler | None" = None, + ) -> VideoObject: + task: Final = _parse(_TASK_ADAPTER, raw_response) + model, served_model = polled_video_models(logging_obj) + video_obj: Final = _video_object_from_task(task, raw_response, model) + if custom_llm_provider: + video_obj.id = encode_video_id_with_provider(task.id, custom_llm_provider, model, served_model) + return video_obj + + def transform_video_content_request( + self, + video_id: str, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: dict[str, str], # mutable-ok: BaseVideoConfig contract + variant: str | None = None, + ) -> tuple[str, dict[str, object]]: # mutable-ok: BaseVideoConfig contract + return self._task_url(video_id, api_base), _EMPTY_PARAMS + + def transform_video_content_response( + self, + raw_response: httpx.Response, + logging_obj: "Logging", + client: "HTTPHandler | None" = None, + ) -> bytes: + video_url: Final = _video_url_from_task(_parse(_TASK_ADAPTER, raw_response), raw_response) + httpx_client: Final[HTTPHandler] = client if client is not None else _get_httpx_client() + video_response: Final = httpx_client.get(video_url) # pyright: ignore[reportUnknownMemberType] # HTTPHandler.get stub leaves params/headers untyped + video_response.raise_for_status() + return video_response.content + + async def async_transform_video_content_response( + self, + raw_response: httpx.Response, + logging_obj: "Logging", + client: "AsyncHTTPHandler | None" = None, + ) -> bytes: + video_url: Final = _video_url_from_task(_parse(_TASK_ADAPTER, raw_response), raw_response) + async_httpx_client: Final[AsyncHTTPHandler] = ( + client if client is not None else get_async_httpx_client(llm_provider=litellm.LlmProviders.VOLCENGINE) + ) + video_response: Final = await async_httpx_client.get(video_url) # pyright: ignore[reportUnknownMemberType] # HTTPHandler.get stub leaves params/headers untyped + video_response.raise_for_status() + return video_response.content + + def transform_video_remix_request( + self, + video_id: str, + prompt: str, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: dict[str, str], # mutable-ok: BaseVideoConfig contract + extra_body: dict[str, object] | None = None, # mutable-ok: BaseVideoConfig contract + ) -> tuple[str, dict[str, object]]: # mutable-ok: BaseVideoConfig contract + raise _unsupported("", _VIDEO_INPUT_UNSUPPORTED.format(operation="Video remix")) + + def transform_video_remix_response( + self, + raw_response: httpx.Response, + logging_obj: "Logging", + custom_llm_provider: str | None = None, + ) -> VideoObject: + raise NotImplementedError("Video remix is not supported by Volcengine.") + + def transform_video_edit_request( + self, + prompt: str, + video_id: str, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: dict[str, str], # mutable-ok: BaseVideoConfig contract + video_file: object | None = None, + extra_body: dict[str, object] | None = None, # mutable-ok: BaseVideoConfig contract + prefetched_source_data: dict[str, object] | None = None, # mutable-ok: BaseVideoConfig contract + ) -> tuple[str, Mapping[str, object], RequestFiles | None]: + raise _unsupported("", _VIDEO_INPUT_UNSUPPORTED.format(operation="Video edit")) + + def transform_video_extension_request( + self, + prompt: str, + video_id: str, + seconds: str, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: dict[str, str], # mutable-ok: BaseVideoConfig contract + extra_body: dict[str, object] | None = None, # mutable-ok: BaseVideoConfig contract + ) -> tuple[str, dict[str, object]]: # mutable-ok: BaseVideoConfig contract + raise _unsupported("", _VIDEO_INPUT_UNSUPPORTED.format(operation="Video extension")) + + def transform_video_list_request( + self, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: dict[str, str], # mutable-ok: BaseVideoConfig contract + after: str | None = None, + limit: int | None = None, + order: str | None = None, + extra_query: dict[str, object] | None = None, # mutable-ok: BaseVideoConfig contract + ) -> tuple[str, dict[str, object]]: # mutable-ok: BaseVideoConfig contract + """Ark pages by page_num/page_size with filters and has no cursor or order.""" + unsupported: Final = tuple(name for name, value in (("after", after), ("order", order)) if value is not None) + if unsupported and not _drops_params(litellm_params): + raise _droppable( + "", f"Ark's task list pages by page_num in extra_query and has no {' or '.join(unsupported)}" + ) + return f"{api_base}{VOLCENGINE_VIDEO_TASKS_PATH}", { + **({"page_size": str(limit)} if limit is not None else {}), + **(extra_query or {}), + } + + def transform_video_list_response( # pyright: ignore[reportIncompatibleMethodOverride] # base declares dict[str, str] but the payload is a heterogeneous list body + self, + raw_response: httpx.Response, + logging_obj: "Logging", + custom_llm_provider: str | None = None, + ) -> dict[str, object]: # mutable-ok: OpenAI list body served as JSON by the proxy + list_payload: Final = _parse(_TASK_LIST_ADAPTER, raw_response) + videos: Final = tuple(_video_object_from_task(task, raw_response, None) for task in list_payload.items) + data: Final = [ + video.model_dump() + | ( + {"id": encode_video_id_with_provider(video.id, custom_llm_provider, video.model)} + if custom_llm_provider + else {} + ) + for video in videos + ] + return { + "object": "list", + "data": data, + "total": list_payload.total if list_payload.total is not None else len(data), + "has_more": False, + **({"first_id": data[0]["id"], "last_id": data[-1]["id"]} if data else {}), + } + + def transform_video_delete_request( + self, + video_id: str, + api_base: str, + litellm_params: GenericLiteLLMParams, + headers: dict[str, str], # mutable-ok: BaseVideoConfig contract + ) -> tuple[str, dict[str, object]]: # mutable-ok: BaseVideoConfig contract + raise NotImplementedError("Video delete is not supported by Volcengine.") + + def transform_video_delete_response( + self, + raw_response: httpx.Response, + logging_obj: "Logging", + ) -> VideoObject: + raise NotImplementedError("Video delete is not supported by Volcengine.") + + def get_error_class( + self, + error_message: str, + status_code: int, + headers: dict[str, str] | httpx.Headers, # mutable-ok: BaseVideoConfig contract + ) -> BaseLLMException: + return VolcEngineVideoError(status_code=status_code, message=error_message, headers=headers) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 5fff760d85a..49d2ce90af4 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -60467,6 +60467,61 @@ } ] }, + "volcengine/doubao-seedance-2-5-260628": { + "litellm_provider": "volcengine", + "mode": "video_generation", + "source": "https://www.volcengine.com/docs/82379/1544106", + "output_cost_per_second": 0.216, + "output_cost_per_second_480p": 0.096075, + "output_cost_per_second_720p": 0.216, + "output_cost_per_second_1080p": 0.5346, + "supported_modalities": ["text", "image", "video", "audio"], + "supported_output_modalities": ["video"], + "supported_endpoints": ["/v1/videos"], + "provider_specific_entry": {"volcengine_video_frames": {"480p": {"16:9": "854x480", "4:3": "752x560", "1:1": "640x640", "3:4": "560x752", "9:16": "480x854", "21:9": "992x432"}, "720p": {"16:9": "1280x720", "4:3": "1112x834", "1:1": "960x960", "3:4": "834x1112", "9:16": "720x1280", "21:9": "1470x630"}}}, + "metadata": {"comment": "Ark bills per output token, not per second: tokens = width * height * 24 * seconds / 1024, so a per-second rate per resolution tier is exact for the 16:9 geometry Ark publishes (1280x720 is 21600 tokens/s). List 70 CNY per million tokens at 480p and 720p and 77 at 1080p without a video input, matching Ark's per-resolution table and its 1080p 16:9 example of 3.74 CNY/s, converted at 7.0 CNY/USD like the other volcengine entries. Accepts up to 50 reference assets (30 images, 10 videos, 10 audios) and clips of 4-30s; per-asset billing beyond the free allowance is not encoded. Requests with a video input bill at a lower rate per token and are rejected, since their cost is not known at create time. provider_specific_entry.volcengine_video_frames is the frame Ark renders per resolution tier and ratio, verified with ffprobe at 480p for every ratio and at 720p for 16:9 and 4:3; the other 720p ratios are from the Seedance 2.0 model card."} + }, + "volcengine/doubao-seedance-2-0-260128": { + "litellm_provider": "volcengine", + "mode": "video_generation", + "source": "https://www.volcengine.com/docs/82379/1544106", + "output_cost_per_second": 0.141942857, + "output_cost_per_second_480p": 0.066003429, + "output_cost_per_second_720p": 0.141942857, + "output_cost_per_second_1080p": 0.354085714, + "output_cost_per_second_4k": 0.722057143, + "supported_modalities": ["text", "image", "video", "audio"], + "supported_output_modalities": ["video"], + "supported_endpoints": ["/v1/videos"], + "provider_specific_entry": {"volcengine_video_frames": {"480p": {"16:9": "864x496", "4:3": "752x560", "1:1": "640x640", "3:4": "560x752", "9:16": "496x864", "21:9": "992x432"}, "720p": {"16:9": "1280x720", "4:3": "1112x834", "1:1": "960x960", "3:4": "834x1112", "9:16": "720x1280", "21:9": "1470x630"}}}, + "metadata": {"comment": "Ark bills per output token, not per second: tokens = width * height * 24 * seconds / 1024, so a per-second rate per resolution tier is exact for the 16:9 geometry Ark publishes (1280x720 is 21600 tokens/s; the 480p frame is 864x496, 10044 tokens/s, not the tier's nominal 854x480). List 46 CNY per million tokens at 480p and 720p, 51 at 1080p, and 26 at 4k without a video input, matching Ark's per-resolution table and its 1080p and 4k 16:9 examples of 2.48 and 5.05 CNY/s, converted at 7.0 CNY/USD like the other volcengine entries. Clips of 4-15s; 4k is exclusive to this tier. Requests with a video input bill at a lower rate per token and are rejected, since their cost is not known at create time. provider_specific_entry.volcengine_video_frames is the frame Ark renders per resolution tier and ratio, verified with ffprobe at 480p for every ratio and at 720p for 16:9 and 4:3; the other 720p ratios are from the Seedance 2.0 model card."} + }, + "volcengine/doubao-seedance-2-0-fast-260128": { + "litellm_provider": "volcengine", + "mode": "video_generation", + "source": "https://www.volcengine.com/docs/82379/1544106", + "output_cost_per_second": 0.114171429, + "output_cost_per_second_480p": 0.053089714, + "output_cost_per_second_720p": 0.114171429, + "supported_modalities": ["text", "image", "video", "audio"], + "supported_output_modalities": ["video"], + "supported_endpoints": ["/v1/videos"], + "provider_specific_entry": {"volcengine_video_frames": {"480p": {"16:9": "864x496", "4:3": "752x560", "1:1": "640x640", "3:4": "560x752", "9:16": "496x864", "21:9": "992x432"}, "720p": {"16:9": "1280x720", "4:3": "1112x834", "1:1": "960x960", "3:4": "834x1112", "9:16": "720x1280", "21:9": "1470x630"}}}, + "metadata": {"comment": "Ark bills per output token, not per second: tokens = width * height * 24 * seconds / 1024, so a per-second rate per resolution tier is exact for the 16:9 geometry Ark publishes (1280x720 is 21600 tokens/s; the 480p frame is 864x496, 10044 tokens/s, not the tier's nominal 854x480). List 37 CNY per million tokens without a video input, converted at 7.0 CNY/USD like the other volcengine entries. Clips of 4-15s at 480p or 720p. Requests with a video input bill at a lower rate per token and are rejected, since their cost is not known at create time. provider_specific_entry.volcengine_video_frames is the frame Ark renders per resolution tier and ratio, verified with ffprobe at 480p for every ratio and at 720p for 16:9 and 4:3; the other 720p ratios are from the Seedance 2.0 model card."} + }, + "volcengine/doubao-seedance-2-0-mini-260615": { + "litellm_provider": "volcengine", + "mode": "video_generation", + "source": "https://www.volcengine.com/docs/82379/1544106", + "output_cost_per_second": 0.070971429, + "output_cost_per_second_480p": 0.033001714, + "output_cost_per_second_720p": 0.070971429, + "supported_modalities": ["text", "image"], + "supported_output_modalities": ["video"], + "supported_endpoints": ["/v1/videos"], + "provider_specific_entry": {"volcengine_video_frames": {"480p": {"16:9": "864x496", "4:3": "752x560", "1:1": "640x640", "3:4": "560x752", "9:16": "496x864", "21:9": "992x432"}, "720p": {"16:9": "1280x720", "4:3": "1112x834", "1:1": "960x960", "3:4": "834x1112", "9:16": "720x1280", "21:9": "1470x630"}}}, + "metadata": {"comment": "Ark bills per output token, not per second: tokens = width * height * 24 * seconds / 1024, so a per-second rate per resolution tier is exact for the 16:9 geometry Ark publishes (1280x720 is 21600 tokens/s; the 480p frame is 864x496, 10044 tokens/s, not the tier's nominal 854x480). List 23 CNY per million tokens without a video input, converted at 7.0 CNY/USD like the other volcengine entries. Clips of 4-15s at 480p or 720p. Only text and image inputs are declared; the mini tier's reference video and audio caps were not verified. Requests with a video input bill at a lower rate per token and are rejected, since their cost is not known at create time. provider_specific_entry.volcengine_video_frames is the frame Ark renders per resolution tier and ratio, verified with ffprobe at 480p for every ratio and at 720p for 16:9 and 4:3; the other 720p ratios are from the Seedance 2.0 model card."} + }, "bedrock/us-east-1/zai.glm-5": { "input_cost_per_token": 1e-06, "output_cost_per_token": 3.2e-06, diff --git a/litellm/utils.py b/litellm/utils.py index ea36a68ad5a..d2ca9faef41 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -9947,6 +9947,10 @@ class ProviderConfigManager: return MinimaxVideoConfig() elif LlmProviders.EDENAI == provider: return litellm.EdenAIVideoConfig() + elif LlmProviders.VOLCENGINE == provider: + from litellm.llms.volcengine.videos.transformation import VolcEngineVideoConfig + + return VolcEngineVideoConfig() return None @staticmethod diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 5fff760d85a..49d2ce90af4 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -60467,6 +60467,61 @@ } ] }, + "volcengine/doubao-seedance-2-5-260628": { + "litellm_provider": "volcengine", + "mode": "video_generation", + "source": "https://www.volcengine.com/docs/82379/1544106", + "output_cost_per_second": 0.216, + "output_cost_per_second_480p": 0.096075, + "output_cost_per_second_720p": 0.216, + "output_cost_per_second_1080p": 0.5346, + "supported_modalities": ["text", "image", "video", "audio"], + "supported_output_modalities": ["video"], + "supported_endpoints": ["/v1/videos"], + "provider_specific_entry": {"volcengine_video_frames": {"480p": {"16:9": "854x480", "4:3": "752x560", "1:1": "640x640", "3:4": "560x752", "9:16": "480x854", "21:9": "992x432"}, "720p": {"16:9": "1280x720", "4:3": "1112x834", "1:1": "960x960", "3:4": "834x1112", "9:16": "720x1280", "21:9": "1470x630"}}}, + "metadata": {"comment": "Ark bills per output token, not per second: tokens = width * height * 24 * seconds / 1024, so a per-second rate per resolution tier is exact for the 16:9 geometry Ark publishes (1280x720 is 21600 tokens/s). List 70 CNY per million tokens at 480p and 720p and 77 at 1080p without a video input, matching Ark's per-resolution table and its 1080p 16:9 example of 3.74 CNY/s, converted at 7.0 CNY/USD like the other volcengine entries. Accepts up to 50 reference assets (30 images, 10 videos, 10 audios) and clips of 4-30s; per-asset billing beyond the free allowance is not encoded. Requests with a video input bill at a lower rate per token and are rejected, since their cost is not known at create time. provider_specific_entry.volcengine_video_frames is the frame Ark renders per resolution tier and ratio, verified with ffprobe at 480p for every ratio and at 720p for 16:9 and 4:3; the other 720p ratios are from the Seedance 2.0 model card."} + }, + "volcengine/doubao-seedance-2-0-260128": { + "litellm_provider": "volcengine", + "mode": "video_generation", + "source": "https://www.volcengine.com/docs/82379/1544106", + "output_cost_per_second": 0.141942857, + "output_cost_per_second_480p": 0.066003429, + "output_cost_per_second_720p": 0.141942857, + "output_cost_per_second_1080p": 0.354085714, + "output_cost_per_second_4k": 0.722057143, + "supported_modalities": ["text", "image", "video", "audio"], + "supported_output_modalities": ["video"], + "supported_endpoints": ["/v1/videos"], + "provider_specific_entry": {"volcengine_video_frames": {"480p": {"16:9": "864x496", "4:3": "752x560", "1:1": "640x640", "3:4": "560x752", "9:16": "496x864", "21:9": "992x432"}, "720p": {"16:9": "1280x720", "4:3": "1112x834", "1:1": "960x960", "3:4": "834x1112", "9:16": "720x1280", "21:9": "1470x630"}}}, + "metadata": {"comment": "Ark bills per output token, not per second: tokens = width * height * 24 * seconds / 1024, so a per-second rate per resolution tier is exact for the 16:9 geometry Ark publishes (1280x720 is 21600 tokens/s; the 480p frame is 864x496, 10044 tokens/s, not the tier's nominal 854x480). List 46 CNY per million tokens at 480p and 720p, 51 at 1080p, and 26 at 4k without a video input, matching Ark's per-resolution table and its 1080p and 4k 16:9 examples of 2.48 and 5.05 CNY/s, converted at 7.0 CNY/USD like the other volcengine entries. Clips of 4-15s; 4k is exclusive to this tier. Requests with a video input bill at a lower rate per token and are rejected, since their cost is not known at create time. provider_specific_entry.volcengine_video_frames is the frame Ark renders per resolution tier and ratio, verified with ffprobe at 480p for every ratio and at 720p for 16:9 and 4:3; the other 720p ratios are from the Seedance 2.0 model card."} + }, + "volcengine/doubao-seedance-2-0-fast-260128": { + "litellm_provider": "volcengine", + "mode": "video_generation", + "source": "https://www.volcengine.com/docs/82379/1544106", + "output_cost_per_second": 0.114171429, + "output_cost_per_second_480p": 0.053089714, + "output_cost_per_second_720p": 0.114171429, + "supported_modalities": ["text", "image", "video", "audio"], + "supported_output_modalities": ["video"], + "supported_endpoints": ["/v1/videos"], + "provider_specific_entry": {"volcengine_video_frames": {"480p": {"16:9": "864x496", "4:3": "752x560", "1:1": "640x640", "3:4": "560x752", "9:16": "496x864", "21:9": "992x432"}, "720p": {"16:9": "1280x720", "4:3": "1112x834", "1:1": "960x960", "3:4": "834x1112", "9:16": "720x1280", "21:9": "1470x630"}}}, + "metadata": {"comment": "Ark bills per output token, not per second: tokens = width * height * 24 * seconds / 1024, so a per-second rate per resolution tier is exact for the 16:9 geometry Ark publishes (1280x720 is 21600 tokens/s; the 480p frame is 864x496, 10044 tokens/s, not the tier's nominal 854x480). List 37 CNY per million tokens without a video input, converted at 7.0 CNY/USD like the other volcengine entries. Clips of 4-15s at 480p or 720p. Requests with a video input bill at a lower rate per token and are rejected, since their cost is not known at create time. provider_specific_entry.volcengine_video_frames is the frame Ark renders per resolution tier and ratio, verified with ffprobe at 480p for every ratio and at 720p for 16:9 and 4:3; the other 720p ratios are from the Seedance 2.0 model card."} + }, + "volcengine/doubao-seedance-2-0-mini-260615": { + "litellm_provider": "volcengine", + "mode": "video_generation", + "source": "https://www.volcengine.com/docs/82379/1544106", + "output_cost_per_second": 0.070971429, + "output_cost_per_second_480p": 0.033001714, + "output_cost_per_second_720p": 0.070971429, + "supported_modalities": ["text", "image"], + "supported_output_modalities": ["video"], + "supported_endpoints": ["/v1/videos"], + "provider_specific_entry": {"volcengine_video_frames": {"480p": {"16:9": "864x496", "4:3": "752x560", "1:1": "640x640", "3:4": "560x752", "9:16": "496x864", "21:9": "992x432"}, "720p": {"16:9": "1280x720", "4:3": "1112x834", "1:1": "960x960", "3:4": "834x1112", "9:16": "720x1280", "21:9": "1470x630"}}}, + "metadata": {"comment": "Ark bills per output token, not per second: tokens = width * height * 24 * seconds / 1024, so a per-second rate per resolution tier is exact for the 16:9 geometry Ark publishes (1280x720 is 21600 tokens/s; the 480p frame is 864x496, 10044 tokens/s, not the tier's nominal 854x480). List 23 CNY per million tokens without a video input, converted at 7.0 CNY/USD like the other volcengine entries. Clips of 4-15s at 480p or 720p. Only text and image inputs are declared; the mini tier's reference video and audio caps were not verified. Requests with a video input bill at a lower rate per token and are rejected, since their cost is not known at create time. provider_specific_entry.volcengine_video_frames is the frame Ark renders per resolution tier and ratio, verified with ffprobe at 480p for every ratio and at 720p for 16:9 and 4:3; the other 720p ratios are from the Seedance 2.0 model card."} + }, "bedrock/us-east-1/zai.glm-5": { "input_cost_per_token": 1e-06, "output_cost_per_token": 3.2e-06, diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json index 45104d40018..0325fe1f796 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -2766,7 +2766,8 @@ "batches": false, "rerank": false, "a2a": true, - "interactions": true + "interactions": true, + "video_generations": true } }, "voyage": { diff --git a/tests/unit/llms/dashscope/test_dashscope_video_transformation.py b/tests/unit/llms/dashscope/test_dashscope_video_transformation.py index 65eefcab1af..15b77aaee3d 100644 --- a/tests/unit/llms/dashscope/test_dashscope_video_transformation.py +++ b/tests/unit/llms/dashscope/test_dashscope_video_transformation.py @@ -194,8 +194,9 @@ class TestDashScopeVideoInputReference: lambda tmp_path: ("first.png", PNG_BYTES, "image/png"), lambda tmp_path: ("first.png", io.BytesIO(PNG_BYTES)), lambda tmp_path: _write_png(tmp_path), + lambda tmp_path: ("first.png", ("inner.png", PNG_BYTES)), ], - ids=["filename-bytes-content-type", "filename-file", "path"], + ids=["filename-bytes-content-type", "filename-file", "path", "nested-tuple"], ) def test_file_tuple_and_path_references_are_encoded(self, make_reference, tmp_path): """Regression: tuples and paths are valid FileTypes but raised before reaching DashScope.""" diff --git a/tests/unit/llms/volcengine/test_volcengine_video_transformation.py b/tests/unit/llms/volcengine/test_volcengine_video_transformation.py new file mode 100644 index 00000000000..1fcf49eeb13 --- /dev/null +++ b/tests/unit/llms/volcengine/test_volcengine_video_transformation.py @@ -0,0 +1,1131 @@ +import base64 +import io +import json +from pathlib import Path +from typing import Final +from unittest.mock import Mock + +import httpx +import pytest +import respx + +import litellm +from litellm import Router +from litellm.exceptions import UnsupportedParamsError +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +from litellm.llms.openai.cost_calculation import video_generation_cost +from litellm.llms.volcengine.videos.transformation import VolcEngineVideoConfig, VolcEngineVideoError +from litellm.types.router import GenericLiteLLMParams +from litellm.types.videos.utils import decode_video_id_with_provider, encode_video_id_with_provider +from litellm.utils import ProviderConfigManager +from litellm.videos.main import avideo_content, avideo_generation + +PNG_BYTES: Final = base64.b64decode( + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8BQDwAEhQGAhKmMIQAAAABJRU5ErkJggg==" +) +PNG_DATA_URL: Final = f"data:image/png;base64,{base64.b64encode(PNG_BYTES).decode()}" + +API_BASE: Final = "https://ark.cn-beijing.volces.com" +TASKS_URL: Final = f"{API_BASE}/api/v3/contents/generations/tasks" +SEEDANCE_2_5: Final = "doubao-seedance-2-5-260628" +SEEDANCE_2_0: Final = "doubao-seedance-2-0-260128" +SEEDANCE_2_0_FAST: Final = "doubao-seedance-2-0-fast-260128" +TASK_ID: Final = "cgt-2026-abc" + + +def _response(payload: object, status_code: int = 200) -> httpx.Response: + return httpx.Response(status_code, json=payload, request=httpx.Request("GET", TASKS_URL)) + + +def _text_response(text: str, status_code: int = 200) -> httpx.Response: + return httpx.Response(status_code, text=text, request=httpx.Request("GET", TASKS_URL)) + + +def _create_body( + params: dict[str, object], model: str = SEEDANCE_2_5, drop_params: bool | None = None +) -> dict[str, object]: + body, files, url = VolcEngineVideoConfig().transform_video_create_request( + model=model, + prompt="a cat on a roof", + api_base=API_BASE, + video_create_optional_request_params=params, + litellm_params=GenericLiteLLMParams(drop_params=drop_params), + headers={}, + ) + assert (files, url) == ((), TASKS_URL) + return body + + +def _created(request_data: dict[str, object], model: str = SEEDANCE_2_5, logging_obj: object = None): + return VolcEngineVideoConfig().transform_video_create_response( + model=model, + raw_response=_response({"id": TASK_ID}), + logging_obj=logging_obj, # pyright: ignore[reportArgumentType] # the transform only reads litellm_params off it + custom_llm_provider="volcengine", + request_data=request_data, + ) + + +def _polled(payload: object, logging_obj: object = None): + return VolcEngineVideoConfig().transform_video_status_retrieve_response( + raw_response=_response(payload), + logging_obj=logging_obj, # pyright: ignore[reportArgumentType] # the transform only reads litellm_params off it + custom_llm_provider="volcengine", + ) + + +def _task(**fields: object) -> dict[str, object]: + return {"id": TASK_ID, "model": SEEDANCE_2_5, "status": "succeeded", **fields} + + +@pytest.mark.parametrize( + "input_reference", + [ + PNG_BYTES, + io.BytesIO(PNG_BYTES), + ("frame.png", PNG_BYTES, "image/png"), + ("frame.png", io.BytesIO(PNG_BYTES)), + ], + ids=["bytes", "file", "tuple", "tuple-with-file"], +) +def test_a_file_input_reference_becomes_a_data_uri_first_frame(input_reference: object): + body = _create_body({"input_reference": input_reference, "duration": 4}) + + assert body == { + "model": SEEDANCE_2_5, + "content": [ + {"type": "text", "text": "a cat on a roof"}, + {"type": "image_url", "image_url": {"url": PNG_DATA_URL}, "role": "first_frame"}, + ], + "resolution": "720p", + "duration": 4, + } + + +def test_a_path_input_reference_is_read_from_disk(tmp_path: Path): + frame = tmp_path / "frame.png" + frame.write_bytes(PNG_BYTES) + + body = _create_body({"input_reference": frame, "duration": 4}) + + assert body["content"] == [ + {"type": "text", "text": "a cat on a roof"}, + {"type": "image_url", "image_url": {"url": PNG_DATA_URL}, "role": "first_frame"}, + ] + + +def test_a_url_input_reference_is_passed_as_is_and_never_opened_as_a_path(): + body = _create_body({"input_reference": "https://x/first.png", "duration": 4}) + + assert body["content"] == [ + {"type": "text", "text": "a cat on a roof"}, + {"type": "image_url", "image_url": {"url": "https://x/first.png"}, "role": "first_frame"}, + ] + + +def test_openais_object_input_reference_sends_its_image_url(): + body = _create_body({"input_reference": {"image_url": "https://x/first.png"}, "duration": 4}) + + assert body["content"] == [ + {"type": "text", "text": "a cat on a roof"}, + {"type": "image_url", "image_url": {"url": "https://x/first.png"}, "role": "first_frame"}, + ] + + +def test_a_file_id_input_reference_is_a_bad_request_unless_drop_params(): + with pytest.raises(UnsupportedParamsError, match="file_id is not supported"): + _ = _create_body({"input_reference": {"file_id": "file-abc"}, "duration": 4}) + + body = _create_body({"input_reference": {"file_id": "file-abc"}, "duration": 4}, drop_params=True) + + assert body["content"] == [{"type": "text", "text": "a cat on a roof"}] + + +def test_a_caller_content_array_wins_over_input_reference_and_the_prompt_leads(): + body = _create_body( + { + "content": [{"type": "image_url", "image_url": {"url": "https://x/a.png"}, "role": "reference_image"}], + "input_reference": "https://x/ignored.png", + "duration": 4, + } + ) + + assert body["content"] == [ + {"type": "text", "text": "a cat on a roof"}, + {"type": "image_url", "image_url": {"url": "https://x/a.png"}, "role": "reference_image"}, + ] + + +IMAGE_ITEM: Final = {"type": "image_url", "image_url": {"url": "https://x/a.png"}, "role": "reference_image"} + + +def test_a_text_item_in_content_is_rejected_because_input_guardrails_only_scan_the_prompt(): + with pytest.raises(UnsupportedParamsError, match="put the text in prompt"): + _ = _create_body({"content": [{"type": "text", "text": "text the guardrail never saw"}], "duration": 4}) + + +def test_a_text_item_in_content_is_dropped_under_drop_params_so_ark_never_sees_unscanned_text(): + body = _create_body( + {"content": [{"type": "text", "text": "text the guardrail never saw"}, IMAGE_ITEM], "duration": 4}, + drop_params=True, + ) + + assert body["content"] == [{"type": "text", "text": "a cat on a roof"}, IMAGE_ITEM] + + +@pytest.mark.parametrize( + "item", + [ + {"type": "video_url", "video_url": {"url": "https://x/v.mp4"}, "role": "reference_video"}, + {"type": "video_url", "video_url": {"url": "https://x/v.mp4"}}, + {"type": "image_url", "image_url": {"url": "https://x/v.png"}, "role": "reference_video"}, + ], + ids=["reference-video", "untagged-video", "mistagged-role"], +) +def test_a_video_input_is_rejected_because_its_cost_is_not_known_at_create(item: dict[str, object]): + with pytest.raises(UnsupportedParamsError, match="Video input"): + _ = _create_body({"content": [item], "duration": 4}) + + body = _create_body({"content": [item, IMAGE_ITEM], "duration": 4}, drop_params=True) + + assert body["content"] == [{"type": "text", "text": "a cat on a roof"}, IMAGE_ITEM] + + +@pytest.mark.parametrize( + "content", + ["a string", [{"type": "image_url", "role": 7}], [{"type": 7}]], + ids=["not-an-array", "non-string-role", "non-string-type"], +) +def test_a_content_value_that_is_not_an_array_of_items_is_rejected_unless_drop_params(content: object): + with pytest.raises(UnsupportedParamsError, match="content must be an array"): + _ = _create_body({"content": content, "duration": 4}) + + body = _create_body({"content": content, "duration": 4}, drop_params=True) + + assert body["content"] == [{"type": "text", "text": "a cat on a roof"}] + + +def test_an_omitted_duration_and_tier_are_sent_as_openais_four_seconds_at_720p_so_ark_cannot_pick_what_is_billed(): + body = _create_body({}) + + assert body == { + "model": SEEDANCE_2_5, + "content": [{"type": "text", "text": "a cat on a roof"}], + "resolution": "720p", + "duration": 4, + } + + +@pytest.mark.usefixtures("local_model_cost_map") +def test_an_endpoint_id_priced_through_base_model_gets_the_base_models_default_tier(): + body, _, _ = VolcEngineVideoConfig().transform_video_create_request( + model="ep-20261015-abc", + prompt="a cat on a roof", + api_base=API_BASE, + video_create_optional_request_params={"duration": 4}, + litellm_params=GenericLiteLLMParams(model_info={"base_model": f"volcengine/{SEEDANCE_2_5}"}), + headers={}, + ) + + assert (body["model"], body["resolution"]) == ("ep-20261015-abc", "720p") + + +@pytest.mark.parametrize("field", ["return_last_frame", "draft", "frame_count"]) +def test_a_field_litellm_does_not_send_to_ark_is_rejected(field: str): + with pytest.raises(UnsupportedParamsError, match=f"does not send {field} to Ark"): + _ = _create_body({"duration": 4, field: True}) + + +def test_a_field_litellm_does_not_send_to_ark_is_dropped_under_drop_params(): + body = _create_body({"duration": 4, "return_last_frame": True, "seed": 7}, drop_params=True) + + assert "return_last_frame" not in body + assert body["seed"] == 7 + + +@pytest.mark.parametrize("duration", [-1, "-1", 0, "abc", 4.5]) +def test_a_duration_ark_would_pick_or_that_is_not_whole_seconds_is_rejected(duration: object): + with pytest.raises(UnsupportedParamsError, match="duration must be a positive whole number"): + _ = _create_body({"duration": duration}) + + assert _create_body({"duration": duration}, drop_params=True)["duration"] == 4 + + +def test_frames_is_rejected_because_it_bypasses_the_billed_duration(): + with pytest.raises(UnsupportedParamsError, match="frames"): + _ = _create_body({"frames": 121, "duration": 6}) + + body = _create_body({"frames": 121, "duration": 6}, drop_params=True) + + assert ("frames" not in body, body["duration"]) == (True, 6) + + +def test_ark_fields_reach_the_body_and_false_flags_are_kept(): + body = _create_body( + { + "duration": 4, + "resolution": "480p", + "ratio": "9:16", + "seed": 7, + "watermark": False, + "generate_audio": False, + "service_tier": "flex", + } + ) + + assert body == { + "model": SEEDANCE_2_5, + "content": [{"type": "text", "text": "a cat on a roof"}], + "resolution": "480p", + "ratio": "9:16", + "seed": 7, + "watermark": False, + "generate_audio": False, + "service_tier": "flex", + "duration": 4, + } + + +@pytest.mark.usefixtures("local_model_cost_map") +def test_a_resolution_the_model_has_no_price_for_is_rejected(): + with pytest.raises(UnsupportedParamsError, match="not offered at resolution '1080p', only at 480p, 720p"): + _ = _create_body({"duration": 4, "resolution": "1080p"}, model=SEEDANCE_2_0_FAST) + + +@pytest.mark.usefixtures("local_model_cost_map") +def test_an_unpriced_resolution_is_dropped_under_drop_params(): + body = _create_body( + {"duration": 4, "resolution": "1080p", "ratio": "16:9"}, model=SEEDANCE_2_0_FAST, drop_params=True + ) + + assert body == { + "model": SEEDANCE_2_0_FAST, + "content": [{"type": "text", "text": "a cat on a roof"}], + "resolution": "720p", + "ratio": "16:9", + "duration": 4, + } + + +@pytest.mark.usefixtures("local_model_cost_map") +def test_a_model_the_router_registered_without_prices_keeps_its_resolution(): + endpoint_id: Final = "ep-20261015-unpriced" + _ = Router( + model_list=[ + {"model_name": "ark", "litellm_params": {"model": f"volcengine/{endpoint_id}", "api_key": "sk-test"}} + ] + ) + + body = _create_body({"duration": 4, "resolution": "480p"}, model=endpoint_id) + body_without_a_tier = _create_body({"duration": 4}, model=endpoint_id) + + assert body["resolution"] == "480p" + assert "resolution" not in body_without_a_tier + + +def test_a_mistyped_ark_field_is_rejected(): + with pytest.raises(UnsupportedParamsError, match="invalid seed"): + _ = _create_body({"duration": 4, "seed": "abc", "watermark": False}) + + +def test_a_mistyped_ark_field_is_dropped_under_drop_params(): + body = _create_body({"duration": 4, "seed": "abc", "watermark": False}, drop_params=True) + + assert body == { + "model": SEEDANCE_2_5, + "content": [{"type": "text", "text": "a cat on a roof"}], + "resolution": "720p", + "watermark": False, + "duration": 4, + } + + +@pytest.mark.parametrize( + "size,expected", + [ + ("1920x1080", {"ratio": "16:9", "resolution": "1080p"}), + ("854x480", {"ratio": "16:9", "resolution": "480p"}), + ("864x496", {"ratio": "16:9", "resolution": "480p"}), + ("720x1280", {"ratio": "9:16", "resolution": "720p"}), + ("1024x768", {"ratio": "4:3", "resolution": "720p"}), + ("640x640", {"ratio": "1:1", "resolution": "480p"}), + ("960x960", {"ratio": "1:1", "resolution": "720p"}), + ("1440x1440", {"ratio": "1:1", "resolution": "1080p"}), + ("1112x834", {"ratio": "4:3", "resolution": "720p"}), + ("1470x630", {"ratio": "21:9", "resolution": "720p"}), + ("3840x2160", {"ratio": "16:9", "resolution": "4k"}), + ], +) +def test_size_maps_onto_an_ark_ratio_and_tier(size: str, expected: dict[str, str]): + assert VolcEngineVideoConfig().map_openai_params({"size": size}, SEEDANCE_2_5, False) == expected + + +def test_explicit_ratio_and_resolution_beat_the_size_derived_ones(): + mapped = VolcEngineVideoConfig().map_openai_params( + {"size": "1280x720", "seconds": "4", "ratio": "1:1", "resolution": "480p", "parameters": {"seed": 3}}, + SEEDANCE_2_5, + False, + ) + + assert mapped == {"ratio": "1:1", "resolution": "480p", "duration": 4, "seed": 3} + + +@pytest.mark.parametrize( + "params,message", + [ + ({"seconds": "4.5"}, "seconds must be a whole number"), + ({"size": "big"}, "size must be WIDTHxHEIGHT"), + ({"size": "16:9"}, "size must be WIDTHxHEIGHT"), + ({"parameters": "seed=3"}, "parameters must be an object"), + ({"characters": [{"id": "c"}]}, "'characters' is not supported"), + ({"image": "https://x/a.png"}, "'image' is not supported"), + ], +) +def test_a_value_ark_cannot_take_raises_unless_drop_params(params: dict[str, object], message: str): + with pytest.raises(UnsupportedParamsError, match=message): + _ = VolcEngineVideoConfig().map_openai_params(params, SEEDANCE_2_5, False) + + assert VolcEngineVideoConfig().map_openai_params(params, SEEDANCE_2_5, True) == {} + + +def test_request_metadata_is_not_mapped_into_the_body_and_user_becomes_the_safety_identifier(): + mapped = VolcEngineVideoConfig().map_openai_params( + {"user": "u", "extra_headers": {"x": "y"}, "seconds": 4}, SEEDANCE_2_5, False + ) + + assert mapped == {"safety_identifier": "u", "duration": 4} + + +def test_an_explicit_safety_identifier_wins_over_user(): + mapped = VolcEngineVideoConfig().map_openai_params( + {"user": "u", "safety_identifier": "s", "seconds": 4}, SEEDANCE_2_5, False + ) + + assert mapped["safety_identifier"] == "s" + + +@pytest.mark.usefixtures("local_model_cost_map") +@pytest.mark.parametrize( + "model,request_data,seconds,size,usage", + [ + ( + SEEDANCE_2_5, + {"duration": 4, "resolution": "480p", "ratio": "16:9"}, + "4", + "854x480", + {"duration_seconds": 4.0, "video_resolution": "480p"}, + ), + ( + SEEDANCE_2_0, + {"duration": 5, "resolution": "480p", "ratio": "16:9"}, + "5", + "864x496", + {"duration_seconds": 5.0, "video_resolution": "480p"}, + ), + (SEEDANCE_2_5, {"duration": 5}, "5", None, {"duration_seconds": 5.0}), + ], + ids=["2.5-480p", "2.0-480p-renders-its-own-frame", "no-geometry"], +) +def test_the_create_response_reports_what_the_task_bills( + model: str, request_data: dict[str, object], seconds: str, size: str | None, usage: dict[str, object] +): + video = _created(request_data, model=model) + + assert video.model_dump(include={"status", "seconds", "size", "usage", "model"}) == { + "status": "queued", + "seconds": seconds, + "size": size, + "usage": usage, + "model": model, + } + + +@pytest.mark.parametrize( + "usage,expected_cost", + [ + ({"duration_seconds": 4.0, "video_resolution": "480p"}, 4 * 0.1), + ({"duration_seconds": 5.0, "video_resolution": "720p"}, 5 * 0.2), + ({"duration_seconds": 5.0}, 5 * 0.2), + ], +) +def test_the_create_usage_bills_its_duration_at_its_tier_rate(usage: dict[str, object], expected_cost: float): + model_info = { + "output_cost_per_second": 0.2, + "output_cost_per_second_480p": 0.1, + "output_cost_per_second_720p": 0.2, + } + video = _created({"duration": int(usage["duration_seconds"]), "resolution": usage.get("video_resolution")}) + + cost = video_generation_cost( + model=SEEDANCE_2_5, + duration_seconds=video.usage["duration_seconds"], + model_info=model_info, # pyright: ignore[reportArgumentType] # a hand-written rate card stands in for ModelInfo + video_resolution=video.usage.get("video_resolution"), + ) + + assert cost == pytest.approx(expected_cost) + + +def test_the_created_id_carries_the_requested_and_served_model_groups(): + logging_obj = Mock( + litellm_params={"litellm_metadata": {"model_group": "seedance", "deployment_model_name": "seedance-account-b"}} + ) + + video = _created({"duration": 5}, logging_obj=logging_obj) + + assert decode_video_id_with_provider(video.id) == { + "custom_llm_provider": "volcengine", + "model_id": "seedance", + "video_id": TASK_ID, + "served_model": "seedance-account-b", + } + + +def test_a_create_without_router_metadata_carries_the_model(): + video = _created({"duration": 5}, logging_obj=Mock(litellm_params={})) + + assert decode_video_id_with_provider(video.id) == { + "custom_llm_provider": "volcengine", + "model_id": SEEDANCE_2_5, + "video_id": TASK_ID, + } + + +@pytest.mark.parametrize( + "raw_response,status_code,message", + [ + (_response({"error": {"code": "InvalidParameter", "message": "bad ratio"}}), 400, "bad ratio"), + (_response({"error": {"code": "Forbidden", "message": "no access"}}, 403), 403, "no access"), + (_text_response("bad gateway"), 502, "Unexpected Volcengine video response"), + ], + ids=["error-on-200", "error-status-kept", "not-json"], +) +def test_a_create_response_without_a_task_id_raises_instead_of_queueing_an_empty_video( + raw_response: httpx.Response, status_code: int, message: str +): + with pytest.raises(VolcEngineVideoError, match=message) as raised: + _ = VolcEngineVideoConfig().transform_video_create_response( + model=SEEDANCE_2_5, + raw_response=raw_response, + logging_obj=None, # pyright: ignore[reportArgumentType] # the transform only reads litellm_params off it + custom_llm_provider="volcengine", + request_data={"duration": 5}, + ) + + assert raised.value.status_code == status_code + + +@pytest.mark.parametrize( + "ark_status,expected", + [ + ("queued", "queued"), + ("running", "in_progress"), + ("succeeded", "completed"), + ("failed", "failed"), + ("cancelled", "cancelled"), + ("expired", "failed"), + ], +) +def test_the_ark_status_maps_onto_the_openai_vocabulary(ark_status: str, expected: str): + assert _polled(_task(status=ark_status)).status == expected + + +@pytest.mark.usefixtures("local_model_cost_map") +def test_a_completed_task_reports_its_delivered_geometry_duration_and_tokens(): + video = _polled( + _task( + created_at=1791989938, + updated_at=1791990031, + resolution="480p", + ratio="16:9", + duration=4, + usage={"completion_tokens": 40594, "total_tokens": 40594}, + model=SEEDANCE_2_0, + content={"video_url": "https://tos/v.mp4"}, + ) + ) + + assert video.model_dump(exclude={"id"}) | {"id": decode_video_id_with_provider(video.id)["video_id"]} == { + "id": TASK_ID, + "object": "video", + "status": "completed", + "created_at": 1791989938, + "completed_at": 1791990031, + "expires_at": None, + "error": None, + "progress": None, + "remixed_from_video_id": None, + "seconds": "4", + "size": "864x496", + "model": SEEDANCE_2_0, + "usage": {"duration_seconds": 4.0, "video_resolution": "480p", "total_tokens": 40594}, + } + + +@pytest.mark.parametrize( + "resolution,ratio,expected", + [ + ("720p", "9:16", "720x1280"), + ("1080p", "16:9", "1920x1080"), + ("480p", "adaptive", None), + ], +) +def test_a_frame_the_cost_map_does_not_list_is_scaled_from_the_tier(resolution: str, ratio: str, expected: str | None): + assert _polled(_task(resolution=resolution, ratio=ratio, model="doubao-seedance-unlisted")).size == expected + + +def test_a_failed_task_carries_its_error_block(): + video = _polled(_task(status="failed", error={"code": "OutputVideoSensitiveContentDetected", "message": "flagged"})) + + assert (video.status, video.error, video.completed_at) == ( + "failed", + {"code": "OutputVideoSensitiveContentDetected", "message": "flagged"}, + None, + ) + + +@pytest.mark.parametrize( + "payload,message", + [ + (_task(status="paused"), "Unknown Volcengine task status: paused"), + ({"status": "running"}, "Unexpected Volcengine video response"), + ({"id": TASK_ID}, "Unexpected Volcengine video response"), + ], + ids=["unmapped-status", "no-id", "no-status"], +) +def test_a_poll_litellm_cannot_read_raises_instead_of_returning_a_job_to_poll_forever(payload: object, message: str): + with pytest.raises(VolcEngineVideoError, match=message) as raised: + _ = _polled(payload) + + assert raised.value.status_code == 502 + + +def test_a_poll_keeps_the_ids_requested_and_served_model_when_the_router_served_it(): + video_id = encode_video_id_with_provider(TASK_ID, "volcengine", "seedance", served_model="seedance-account-b") + logging_obj = Mock( + litellm_params={ + "video_id": video_id, + "litellm_metadata": {"model_group": "seedance", "deployment_model_name": "seedance-account-b"}, + } + ) + + video = _polled(_task(status="running"), logging_obj=logging_obj) + + assert video.id == video_id + assert video.model == "seedance" + + +def test_a_poll_drops_a_served_model_the_router_did_not_route_to(): + edited = encode_video_id_with_provider( + TASK_ID, "volcengine", "seedance-account-a", served_model="seedance-account-b" + ) + logging_obj = Mock(litellm_params={"video_id": edited, "litellm_metadata": {"model_group": "seedance-account-a"}}) + + video = _polled(_task(status="running"), logging_obj=logging_obj) + + assert video.id == encode_video_id_with_provider(TASK_ID, "volcengine", "seedance-account-a") + + +async def _download(task_payload: dict[str, object], video_bytes: bytes = b"mp4") -> bytes: + def respond(request: httpx.Request) -> httpx.Response: + if request.url.host == "tos.example": + return httpx.Response(200, content=video_bytes) + return httpx.Response(200, json=task_payload) + + async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as http_client: + handler: Final = AsyncHTTPHandler() + await handler.close() + handler.client = http_client + return await avideo_content( + video_id=TASK_ID, + custom_llm_provider="volcengine", + model=f"volcengine/{SEEDANCE_2_5}", + api_key="sk-test", + api_base=API_BASE, + client=handler, + ) + + +@pytest.mark.asyncio +async def test_the_download_fetches_the_signed_url_through_the_callers_client(): + content = await _download(_task(content={"video_url": "https://tos.example/v.mp4"}), video_bytes=b"the-video") + + assert content == b"the-video" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "payload,message", + [ + (_task(status="queued"), "still processing"), + (_task(status="running"), "still processing"), + (_task(status="failed", error={"message": "flagged"}), "flagged"), + (_task(status="failed"), "ended as failed"), + (_task(status="cancelled"), "ended as cancelled"), + (_task(status="expired"), "Task expired"), + ], +) +async def test_a_download_with_no_video_yet_is_a_bad_request_so_it_is_neither_retried_nor_cools_the_deployment( + payload: dict[str, object], message: str +): + with pytest.raises(litellm.BadRequestError, match=message) as raised: + _ = await _download(payload) + + assert raised.value.status_code == 400 + + +def test_the_list_maps_ark_paging_and_filters(): + url, params = VolcEngineVideoConfig().transform_video_list_request( + api_base=API_BASE, + litellm_params=GenericLiteLLMParams(), + headers={}, + limit=20, + extra_query={"filter.status": "succeeded"}, + ) + + assert (url, params) == (TASKS_URL, {"page_size": "20", "filter.status": "succeeded"}) + + +@pytest.mark.parametrize("paging", [{"after": "video_abc"}, {"order": "asc"}]) +def test_a_list_cursor_or_order_is_rejected_unless_drop_params(paging: dict[str, str]): + config = VolcEngineVideoConfig() + + with pytest.raises(UnsupportedParamsError, match="pages by page_num"): + _ = config.transform_video_list_request( + api_base=API_BASE, litellm_params=GenericLiteLLMParams(), headers={}, **paging + ) + _, params = config.transform_video_list_request( + api_base=API_BASE, litellm_params=GenericLiteLLMParams(drop_params=True), headers={}, limit=5, **paging + ) + + assert params == {"page_size": "5"} + + +def test_the_list_response_is_an_openai_list_body(): + response = VolcEngineVideoConfig().transform_video_list_response( + raw_response=_response({"total": 1, "items": [_task(status="running")]}), + logging_obj=Mock(), # pyright: ignore[reportArgumentType] # unused by the transform + custom_llm_provider="volcengine", + ) + + encoded = encode_video_id_with_provider(TASK_ID, "volcengine", SEEDANCE_2_5) + assert {key: value for key, value in response.items() if key != "data"} == { + "object": "list", + "total": 1, + "has_more": False, + "first_id": encoded, + "last_id": encoded, + } + assert [(item["id"], item["status"]) for item in response["data"]] == [(encoded, "in_progress")] + + +def test_an_empty_list_still_carries_the_object_type(): + response = VolcEngineVideoConfig().transform_video_list_response( + raw_response=_response({"total": 0, "items": []}), + logging_obj=Mock(), # pyright: ignore[reportArgumentType] # unused by the transform + ) + + assert response == {"object": "list", "data": [], "total": 0, "has_more": False} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("call", ["avideo_edit", "avideo_extension", "avideo_remix"]) +async def test_edit_extension_and_remix_are_a_bad_request_because_a_video_input_cannot_be_billed_at_create(call: str): + with pytest.raises(litellm.BadRequestError, match="Ark bills a task with a video input") as raised: + _ = await getattr(litellm, call)( + video_id="https://x/source.mp4", + prompt="make it brighter", + seconds="5", + model=f"volcengine/{SEEDANCE_2_5}", + custom_llm_provider="volcengine", + api_key="sk-test", + api_base=API_BASE, + ) + + assert raised.value.status_code == 400 + + +@pytest.mark.parametrize( + "api_base", + [ + None, + API_BASE, + f"{API_BASE}/", + f"{API_BASE}/api/v3", + f"{API_BASE}/api/compatible", + f"{API_BASE}/api/v3/responses", + f"{API_BASE}/v1", + ], +) +def test_any_ark_surface_on_the_shared_credential_normalizes_to_the_host_root( + api_base: str | None, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.delenv("VOLCENGINE_API_BASE", raising=False) + monkeypatch.delenv("ARK_API_BASE", raising=False) + + assert VolcEngineVideoConfig().get_complete_url("m", api_base, {}) == API_BASE + + +@pytest.mark.parametrize("env_var", ["VOLCENGINE_API_BASE", "ARK_API_BASE"]) +def test_an_api_base_env_var_routes_a_deployment_without_api_base(env_var: str, monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv("VOLCENGINE_API_BASE", raising=False) + monkeypatch.delenv("ARK_API_BASE", raising=False) + monkeypatch.setenv(env_var, "https://ark.ap-southeast.bytepluses.com/api/v3") + + assert VolcEngineVideoConfig().get_complete_url("m", None, {}) == "https://ark.ap-southeast.bytepluses.com" + assert VolcEngineVideoConfig().get_complete_url("m", f"{API_BASE}/api/v3", {}) == API_BASE + + +@pytest.mark.parametrize( + "api_base,expected", + [ + ("https://gw.internal/ark/api/v3", "https://gw.internal/ark"), + ("https://gw.internal/api/ark/api/v3", "https://gw.internal/api/ark"), + ], +) +def test_a_path_prefix_before_the_api_marker_survives(api_base: str, expected: str): + assert VolcEngineVideoConfig().get_complete_url("m", api_base, {}) == expected + + +def test_a_missing_api_key_raises(monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv("ARK_API_KEY", raising=False) + monkeypatch.delenv("VOLCENGINE_API_KEY", raising=False) + + with pytest.raises(ValueError, match="Volcengine API key is required"): + _ = VolcEngineVideoConfig().validate_environment({}, SEEDANCE_2_5) + + +def test_volcengine_resolves_to_the_video_config(): + config = ProviderConfigManager.get_provider_video_config( + model=SEEDANCE_2_5, provider=litellm.LlmProviders.VOLCENGINE + ) + + assert type(config) is VolcEngineVideoConfig + + +ARK_VIDEO_TIERS: Final = { + SEEDANCE_2_5: {"480p": ((854, 480), 70), "720p": ((1280, 720), 70), "1080p": ((1920, 1080), 77)}, + SEEDANCE_2_0: {"480p": ((864, 496), 46), "720p": ((1280, 720), 46), "1080p": ((1920, 1080), 51), "4k": ((3840, 2160), 26)}, + SEEDANCE_2_0_FAST: {"480p": ((864, 496), 37), "720p": ((1280, 720), 37)}, + "doubao-seedance-2-0-mini-260615": {"480p": ((864, 496), 23), "720p": ((1280, 720), 23)}, +} +CNY_PER_USD: Final = 7.0 + + +@pytest.mark.usefixtures("local_model_cost_map") +@pytest.mark.parametrize("model", sorted(ARK_VIDEO_TIERS)) +def test_each_tier_rate_is_its_rendered_pixels_at_arks_price_for_that_tier(model: str): + """Ark prices 1080p and 4k per token above and below the 480p/720p rate, so a tier is not the base + rate scaled by pixels. Prices and frames are Ark's published per-resolution table, read 2026-10-16 + from https://www.volcengine.com/docs/82379/1544106, whose own 2.0 examples agree: 1080p 16:9 5s bills + 2.48 CNY/s and 4k bills 5.05 CNY/s.""" + info = litellm.get_model_info(model=model, custom_llm_provider="volcengine") + frames = info["provider_specific_entry"]["volcengine_video_frames"] + + assert {tier: info[f"output_cost_per_second_{tier}"] for tier in ARK_VIDEO_TIERS[model]} == { + tier: pytest.approx(width * height * 24 / 1024 * cny_per_million / 1e6 / CNY_PER_USD, rel=1e-5) + for tier, ((width, height), cny_per_million) in ARK_VIDEO_TIERS[model].items() + } + + assert { + tier: tuple(int(side) for side in frames[tier]["16:9"].split("x")) + for tier in ARK_VIDEO_TIERS[model] + if tier in frames + } == {tier: frame for tier, (frame, _) in ARK_VIDEO_TIERS[model].items() if tier in frames} + + +async def _wire_body(**kwargs: object) -> dict[str, object]: + sent: Final[list[httpx.Request]] = [] + + def respond(request: httpx.Request) -> httpx.Response: + sent.append(request) + return httpx.Response(200, json={"id": TASK_ID}) + + async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as http_client: + handler: Final = AsyncHTTPHandler() + await handler.close() + handler.client = http_client + _ = await avideo_generation(api_key="sk-test", api_base=API_BASE, client=handler, **kwargs) + + (request,) = sent + assert (str(request.url), request.headers["authorization"]) == (TASKS_URL, "Bearer sk-test") + return json.loads(request.content) + + +@pytest.mark.asyncio +async def test_seconds_size_and_extra_body_content_reach_the_wire(): + body = await _wire_body( + model=f"volcengine/{SEEDANCE_2_5}", + prompt="use these", + seconds="4", + size="854x480", + extra_body={ + "content": [{"type": "image_url", "image_url": {"url": "https://x/a.png"}, "role": "reference_image"}] + }, + ) + + assert body == { + "model": SEEDANCE_2_5, + "content": [ + {"type": "text", "text": "use these"}, + {"type": "image_url", "image_url": {"url": "https://x/a.png"}, "role": "reference_image"}, + ], + "ratio": "16:9", + "resolution": "480p", + "duration": 4, + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "fields", + [{"parameters": {"seed": 3}}, {"extra_body": {"parameters": {"seed": 3}}}], + ids=["top_level", "extra_body"], +) +async def test_parameters_reach_the_wire_unpacked_whether_sent_top_level_or_in_extra_body(fields: dict[str, object]): + body = await _wire_body(model=f"volcengine/{SEEDANCE_2_5}", prompt="p", seconds="4", **fields) + + assert body == { + "model": SEEDANCE_2_5, + "content": [{"type": "text", "text": "p"}], + "resolution": "720p", + "duration": 4, + "seed": 3, + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "extra_body,expected", + [ + ({"size": "1920x1080"}, {"resolution": "1080p"}), + ({"seconds": "8"}, {"duration": 8}), + ({"user": "someone"}, {"safety_identifier": "someone"}), + ], + ids=["size", "seconds", "user"], +) +async def test_an_extra_body_field_litellm_maps_itself_is_honoured_and_its_leftover_copy_ignored( + extra_body: dict[str, object], expected: dict[str, object] +): + """The shared request builder copies extra_body onto the top-level params, where this adapter maps it, and then + merges the raw extra_body back over the mapped params, returning the same key a second time as a field of its + own. The second copy has to be ignored rather than rejected as a field litellm never sends.""" + body = await _wire_body( + model=f"volcengine/{SEEDANCE_2_5}", + prompt="p", + seconds="4", + size="1280x720", + extra_body=extra_body, + ) + + assert body == { + "model": SEEDANCE_2_5, + "content": [{"type": "text", "text": "p"}], + "ratio": "16:9", + "resolution": "720p", + "duration": 4, + **expected, + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "fields,misplaced", + [ + ({"parameters": {"seconds": "8"}}, "seconds"), + ({"extra_body": {"parameters": {"seconds": "8"}}}, "seconds"), + ({"extra_body": {"parameters": {"size": "1920x1080"}}}, "size"), + ({"extra_body": {"parameters": {"user": "someone"}}}, "user"), + ({"extra_body": {"parameters": {"prompt": "a dog"}}}, "prompt"), + ({"extra_body": {"parameters": {"parameters": {"seed": 3}}}}, "parameters"), + ], + ids=["top_level_seconds", "seconds", "size", "user", "prompt", "nested_parameters"], +) +async def test_a_request_field_litellm_maps_itself_is_rejected_inside_parameters( + fields: dict[str, object], misplaced: str +): + with pytest.raises(UnsupportedParamsError, match=f"send {misplaced} as top-level request fields"): + _ = await _wire_body(model=f"volcengine/{SEEDANCE_2_5}", prompt="p", **fields) + + +@pytest.mark.asyncio +async def test_a_request_field_inside_parameters_is_dropped_under_drop_params_and_the_rest_is_sent(): + body = await _wire_body( + model=f"volcengine/{SEEDANCE_2_5}", + prompt="p", + drop_params=True, + extra_body={"parameters": {"seconds": "8", "seed": 3}}, + ) + + assert body == { + "model": SEEDANCE_2_5, + "content": [{"type": "text", "text": "p"}], + "resolution": "720p", + "duration": 4, + "seed": 3, + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "extra_body,field", + [ + ({"prompt": "a dog"}, "prompt"), + ({"model": SEEDANCE_2_0_FAST}, "model"), + ({"extra_headers": {"x-a": "1"}}, "extra_headers"), + ], + ids=["prompt", "model", "extra_headers"], +) +async def test_an_extra_body_copy_of_a_request_field_litellm_does_not_map_is_rejected( + extra_body: dict[str, object], field: str +): + with pytest.raises(UnsupportedParamsError, match=f"does not send {field} to Ark"): + _ = await _wire_body(model=f"volcengine/{SEEDANCE_2_5}", prompt="p", extra_body=extra_body) + + +@pytest.mark.asyncio +async def test_a_user_that_is_not_a_string_is_rejected_rather_than_ignored(): + with pytest.raises(UnsupportedParamsError, match="user must be a string, got int"): + _ = await _wire_body(model=f"volcengine/{SEEDANCE_2_5}", prompt="p", extra_body={"user": 5}) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "control", + [{"drop_params": False}, {"organization": "org-1"}], + ids=["drop_params", "organization"], +) +async def test_a_litellm_control_field_is_not_mistaken_for_an_ark_field(control: dict[str, object]): + body = await _wire_body(model=f"volcengine/{SEEDANCE_2_5}", prompt="p", seconds="4", **control) + + assert body == { + "model": SEEDANCE_2_5, + "content": [{"type": "text", "text": "p"}], + "resolution": "720p", + "duration": 4, + } + + +@pytest.mark.asyncio +async def test_additional_drop_params_removes_the_named_ark_fields_from_the_wire(): + body = await _wire_body( + model=f"volcengine/{SEEDANCE_2_5}", + prompt="p", + seconds="4", + additional_drop_params=["seed"], + extra_body={"seed": 3, "watermark": False}, + ) + + assert body == { + "model": SEEDANCE_2_5, + "content": [{"type": "text", "text": "p"}], + "resolution": "720p", + "duration": 4, + "watermark": False, + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("dropped", "request_fields"), + [ + pytest.param("user", {"user": "u-1"}, id="user"), + pytest.param("seconds", {"seconds": "8"}, id="seconds"), + pytest.param("size", {"size": "1920x1080"}, id="size"), + pytest.param("parameters", {"parameters": {"seed": 3}}, id="parameters"), + pytest.param("parameters", {"extra_body": {"parameters": {"seed": 3}}}, id="extra_body_parameters"), + pytest.param("image", {"image": {"gcsUri": "gs://b/a.png"}}, id="unsupported_image"), + ], +) +async def test_additional_drop_params_naming_a_request_field_sends_what_omitting_it_sends( + dropped: str, request_fields: dict[str, object] +): + request: Final = {"model": f"volcengine/{SEEDANCE_2_5}", "prompt": "p"} + + with_drop = await _wire_body(**request, **request_fields, additional_drop_params=[dropped]) + + assert with_drop == await _wire_body(**request) + + +@pytest.mark.asyncio +async def test_an_extra_body_duration_of_minus_one_is_rejected_before_anything_is_sent(): + with pytest.raises(UnsupportedParamsError, match="Ark-picked length"): + _ = await _wire_body(model=f"volcengine/{SEEDANCE_2_5}", prompt="p", extra_body={"duration": -1}) + + +ACCOUNT_A_HOST: Final = "https://ark-a.example" +ACCOUNT_B_HOST: Final = "https://ark-b.example" + + +def _two_account_router(**router_kwargs: object) -> Router: + return Router( + model_list=[ + { + "model_name": "seedance-a", + "litellm_params": { + "model": f"volcengine/{SEEDANCE_2_5}", + "api_key": "key-a", + "api_base": ACCOUNT_A_HOST, + }, + }, + { + "model_name": "seedance-b", + "litellm_params": { + "model": f"volcengine/{SEEDANCE_2_5}", + "api_key": "key-b", + "api_base": ACCOUNT_B_HOST, + }, + }, + ], + **router_kwargs, # pyright: ignore[reportArgumentType] # forwarded Router options + ) + + +def _proxy_poll_model(router: Router, video_id: str) -> str | None: + decoded: Final = decode_video_id_with_provider(video_id) + requested: Final = router.resolve_model_name_from_model_id( + decoded.get("model_id"), decoded.get("custom_llm_provider") + ) + served: Final = decoded.get("served_model") + return router.resolve_served_model_name(requested, served) if requested and served else requested + + +@pytest.mark.parametrize("caller_metadata", [None, {"tags": ["proxy-request"]}]) +@respx.mock +def test_a_video_is_polled_on_the_ark_account_that_created_it(caller_metadata: dict[str, object] | None): + router = _two_account_router() + _ = respx.post(f"{ACCOUNT_B_HOST}/api/v3/contents/generations/tasks").mock( + return_value=httpx.Response(200, json={"id": TASK_ID}) + ) + task_get = respx.get(f"{ACCOUNT_B_HOST}/api/v3/contents/generations/tasks/{TASK_ID}").mock( + return_value=httpx.Response(200, json=_task(status="running")) + ) + + created = router.video_generation(model="seedance-b", prompt="a cat", seconds="4", metadata=caller_metadata) + _ = router.video_status(model=_proxy_poll_model(router, created.id), video_id=created.id) + + assert [call.request.headers["authorization"] for call in task_get.calls] == ["Bearer key-b"] + + +@respx.mock +def test_a_video_created_through_a_routing_group_is_polled_on_the_member_that_created_it(): + router = _two_account_router( + routing_groups=[ + {"group_name": "seedance", "models": ["seedance-a", "seedance-b"], "routing_strategy": "simple-shuffle"} + ] + ) + for host in (ACCOUNT_A_HOST, ACCOUNT_B_HOST): + _ = respx.post(f"{host}/api/v3/contents/generations/tasks").mock( + return_value=httpx.Response(200, json={"id": TASK_ID}) + ) + _ = respx.get(f"{host}/api/v3/contents/generations/tasks/{TASK_ID}").mock( + return_value=httpx.Response(200, json=_task(status="running")) + ) + + for _ in range(8): + created = router.video_generation(model="seedance", prompt="a cat", seconds="4") + polled = router.video_status(model=_proxy_poll_model(router, created.id), video_id=created.id) + assert polled.id == created.id + + hosts = [call.request.url.host for call in respx.calls] + assert hosts[1::2] == hosts[0::2] From 859529ccbc1f3c0bf9daace4a9c92a84e60f439f Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 16 Oct 2026 15:24:40 +0000 Subject: [PATCH 20/23] fix(realtime)!: keep turn_detection.type when a transcript guardrail disables auto-response (#44939) * fix(realtime): keep turn_detection.type when a transcript guardrail disables auto-response * test(realtime): type the GA session.update helper and mark its locals Final * test(realtime): expect typed turn_detection on the voice guardrail session updates * test(realtime): accept either gate path for the client's first voice session.update * fix(gemini): drop VAD tuning when the guardrail turns auto-response off * fix(gemini): check create_response without a redundant isinstance * fix(vertex_ai): keep the client's VAD tuning on Vertex when auto-response is off * fix(realtime): keep the transcript-guardrail gate typeless on xAI so it does not auto-answer * fix(realtime): replace any client turn_detection with the untyped gate on xAI xAI answers on its own under a typed server_vad even with create_response off, so a client's typed turn_detection must not reach xAI while the transcript guardrail gates the turn. * fix(realtime): keep a client's early turn_detection and hide the guardrail cancel's error A session.update sent before session.created reached OpenAI ahead of the proxy's own gate, so the gate replaced the client's turn_detection. The proxy now re-sends the client's gated turn_detection after the gate when the two differ. The guardrail's response.cancel now carries a litellm_guardrail_ event_id, and the error OpenAI returns for it (response_cancel_not_active when nothing is answering) no longer reaches the client. * fix(realtime): send the guardrail gate before a client's early session.update and let Muse and Chirp fall back to server_vad under drop_params * fix(realtime): hide only the cancel errors that echo the guardrail's own response.cancel ids Track the event ids of the guardrail cancels this session sent and drop a backend error only when its error.event_id is one of them, so a client's own event id with the same prefix still gets its error. Cover the scripted response.cancel reply in the integration upstream, update the stale list-mode test to the typed gate, and mark the xAI cancel error cell as an expected failure, since xAI replies with its own event id. * test(integration support): deliver the disconnect a real server sends after the client goes away The scripted upstream now keeps reading frames after a reply hits a closed socket, so the unit test's fake ASGI receive has to end with the websocket.disconnect message a real server delivers. * fix(realtime): delete the blocked transcript item before the guardrail's refusal turn OpenAI GA answered the blocked question in some refusal turns because the blocked user item stayed in the conversation. The raw OpenAI-protocol path now sends conversation.item.delete for it right after the cancel. Provider config backends have no such event, so that path no longer passes an item id. Also bound the guardrail integration file by the shared graceful-stop timeout and make the xAI cancel-error xfail require an AssertionError. * fix(realtime): keep the blocked item on xAI, which cannot delete an audio input item * test(realtime): assert the xAI voice block sends no conversation.item.delete upstream --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../litellm_core_utils/realtime_streaming.py | 228 ++++++--- .../realtime/transcription_protocol.py | 10 + .../llms/gemini/realtime/transformation.py | 11 +- litellm/llms/meta/realtime/transformation.py | 17 +- .../realtime_transformation.py | 18 +- .../llms/vertex_ai/realtime/transformation.py | 8 +- litellm/llms/xai/realtime/transformation.py | 13 + litellm/realtime_api/main.py | 4 +- litellm/utils.py | 3 +- tests/integration/_support/upstream.py | 198 +++++++- .../cost_calculation/cost_tracking_case.py | 2 + ..._transcription_guardrail_session_update.py | 473 ++++++++++++++++-- .../test_realtime_unguarded_voice_turns.py | 14 +- .../unit/integration_support/test_upstream.py | 6 +- .../test_realtime_streaming.py | 320 +++++++++++- .../test_gemini_realtime_transformation.py | 36 +- .../test_meta_realtime_transformation.py | 18 + .../test_vertex_ai_realtime_transformation.py | 21 + .../test_vertex_ai_realtime_transformation.py | 30 ++ 19 files changed, 1252 insertions(+), 178 deletions(-) diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index c9b96c0a4c2..e4eb9acc37d 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -1,16 +1,18 @@ import asyncio import json import traceback -from collections.abc import Coroutine, Mapping, Sequence +from collections.abc import Callable, Coroutine, Mapping, Sequence from dataclasses import dataclass from enum import Enum, auto -from typing import TYPE_CHECKING, Any, Final, NoReturn, Protocol, TypedDict, cast +from types import MappingProxyType +from typing import TYPE_CHECKING, Any, Final, NoReturn, Protocol, TypeAlias, TypedDict, cast from pydantic import TypeAdapter from typing_extensions import ReadOnly import litellm from litellm._logging import redact_internal_details_from_client_message, verbose_logger +from litellm._uuid import uuid from litellm.constants import ( REALTIME_METERED_USAGE_KEY, REALTIME_SESSION_FAILURE_LOGGED_KEY, @@ -116,10 +118,87 @@ def _decode_json_object(payload: str) -> Mapping[str, object]: return json.loads(payload) +_GUARDRAIL_TURN_DETECTION_GATE: Final[Mapping[str, object]] = MappingProxyType( + {"type": "server_vad", "create_response": False} +) +_EMPTY_SESSION: Final[Mapping[str, object]] = MappingProxyType({}) +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object]) +_GUARDRAIL_EVENT_ID_PREFIX: Final = "litellm_guardrail_" + + +def _json_object(value: object) -> Mapping[str, object] | None: + return _JSON_OBJECT.validate_python(value) if isinstance(value, Mapping) else None + + +def _mapping_field(container: Mapping[str, object], key: str) -> Mapping[str, object] | None: + return _json_object(container.get(key)) + + +def _guardrail_event_id() -> str: + return f"{_GUARDRAIL_EVENT_ID_PREFIX}{uuid.uuid4().hex}" + + +def _error_event_id(event: Mapping[str, object]) -> str | None: + error: Final = _mapping_field(event, "error") + caused_by: Final = None if error is None else error.get("event_id") + return caused_by if isinstance(caused_by, str) else None + + +_TurnDetectionGate: TypeAlias = Callable[[object], Mapping[str, object]] + + +def _turn_detection_with_auto_response_off(turn_detection: object) -> Mapping[str, object]: + client_turn_detection: Final = _json_object(turn_detection) + if client_turn_detection is None: + return dict(_GUARDRAIL_TURN_DETECTION_GATE) + return {**client_turn_detection, "create_response": False} + + +def _nested_turn_detection_part(session: Mapping[str, object], gate: _TurnDetectionGate) -> Mapping[str, object]: + audio: Final = _mapping_field(session, "audio") + if audio is None: + return {} + audio_input: Final = _mapping_field(audio, "input") + if audio_input is None or "turn_detection" not in audio_input: + return {} + gated_input: Final = { + **audio_input, + "turn_detection": gate(audio_input["turn_detection"]), + } + return {"audio": {**audio, "input": gated_input}} + + +def _with_create_response_disabled( + session: Mapping[str, object], gate: _TurnDetectionGate +) -> Mapping[str, object] | None: + """``session`` with every turn_detection the client sent unable to auto-respond, or None when it sent none.""" + flat_part: Final = {"turn_detection": gate(session["turn_detection"])} if "turn_detection" in session else {} + nested_part: Final = _nested_turn_detection_part(session, gate) + if not flat_part and not nested_part: + return None + return {**session, **flat_part, **nested_part} + + +def _gate_session_update( + session: object, gate: _TurnDetectionGate, inject_gate_when_absent: bool +) -> Mapping[str, object] | None: + """The session to forward under a transcript guardrail, or None when the frame needs no change.""" + client_session: Final = _json_object(session) + if client_session is None: + return None + gated: Final = _with_create_response_disabled(client_session, gate) + if gated is not None or not inject_gate_when_absent: + return gated + return {**client_session, "turn_detection": gate(None)} + + class RealtimeEventNormalizer(Protocol): def should_drop(self, event: object) -> bool: ... def normalize(self, event: dict) -> dict: ... def patch_outgoing_session(self, session: dict) -> dict: ... + def gate_turn_detection(self, turn_detection: object, /) -> Mapping[str, object]: ... + @property + def deletes_conversation_items(self) -> bool: ... class RealtimeUsageMeter(Protocol): @@ -197,6 +276,8 @@ class RealTimeStreaming: # Track whether we have already sent the guardrail turn-detection update # that disables provider auto-response for transcription guardrails. self._guardrail_turn_detection_update_sent: bool = False + self._session_gate_sent: bool = False + self._guardrail_cancel_event_ids: frozenset[str] = frozenset() # Deferred Gemini Live setup: Pipecat may stream audio before session.update. # Buffer client audio until the backend acknowledges setup (setupComplete). self._backend_setup_complete: bool = provider_config is None or provider_config.requires_session_configuration() @@ -213,6 +294,12 @@ class RealTimeStreaming: self._is_transcription_session: bool = force_transcription_model is not None # Optional per-provider GA event normalizer (e.g. XAIRealtimeNormalizer). self._event_normalizer = event_normalizer + self._gate_turn_detection: Final[_TurnDetectionGate] = ( + _turn_detection_with_auto_response_off if event_normalizer is None else event_normalizer.gate_turn_detection + ) + self._deletes_conversation_items: Final = ( + event_normalizer is None or event_normalizer.deletes_conversation_items + ) self._usage_meter = usage_meter @property @@ -734,20 +821,37 @@ class RealTimeStreaming: except (json.JSONDecodeError, TypeError): return + def _make_turn_detection_update(self, turn_detection: Mapping[str, object]) -> dict[str, object]: + session: Final[dict[str, object]] = ( + {"turn_detection": turn_detection} + if self._backend_uses_beta_protocol + else {"type": "realtime", "audio": {"input": {"turn_detection": turn_detection}}} + ) + return {"type": "session.update", "session": session} + def _make_disable_auto_response_message(self) -> str: """Return a session.update that disables VAD auto-response.""" - turn_detection: Final[dict[str, str | bool]] = { - "type": "server_vad", - "create_response": False, - } - if self._backend_uses_beta_protocol: - session: dict[str, object] = {"turn_detection": turn_detection} - else: - session = { - "type": "realtime", - "audio": {"input": {"turn_detection": turn_detection}}, - } - return json.dumps({"type": "session.update", "session": session}) + return json.dumps(self._make_turn_detection_update(self._gate_turn_detection(None))) + + async def _send_guardrail_cancel(self) -> None: + event_id: Final = _guardrail_event_id() + self._guardrail_cancel_event_ids = self._guardrail_cancel_event_ids | {event_id} + await self._send_to_backend(json.dumps({"type": "response.cancel", "event_id": event_id})) + + async def _withdraw_blocked_turn(self, item_id: str | None) -> None: + await self._send_guardrail_cancel() + if item_id is not None and self._deletes_conversation_items: + await self._send_to_backend(json.dumps({"type": "conversation.item.delete", "item_id": item_id})) + + def _is_guardrail_cancel_error(self, event: Mapping[str, object]) -> bool: + return _error_event_id(event) in self._guardrail_cancel_event_ids + + async def _send_session_gate_once(self) -> None: + """Gate the backend's default turn_detection ahead of the first client session.update, which overrides it.""" + if self._session_gate_sent or not self._should_disable_vad_auto_response(): + return + self._session_gate_sent = True + await self._send_to_backend(self._make_disable_auto_response_message()) async def _maybe_send_guardrail_turn_detection_update(self) -> None: """Disable provider auto-response once when transcription guardrails are enabled.""" @@ -921,7 +1025,7 @@ class RealTimeStreaming: if pre_block_backend_message is not None: await self._send_to_backend(pre_block_backend_message) if not self._is_transcription_session: - await self._send_to_backend(json.dumps({"type": "response.cancel"})) + await self._withdraw_blocked_turn(item_id) # Send the policy violation hint (shows as small gray status text in UI). await self.websocket.send_text( json.dumps( @@ -1051,10 +1155,7 @@ class RealTimeStreaming: self.store_message(event_str) self._capture_transcription_usage(event) await self._send_event_to_client(event, event_str) - blocked = await self.run_realtime_guardrails( - cast(str, transcript), - item_id=cast(str | None, event.get("item_id")), - ) + blocked = await self.run_realtime_guardrails(cast(str, transcript)) if not blocked and self._should_disable_vad_auto_response(): await self._send_to_backend(json.dumps({"type": "response.create"})) continue @@ -1086,7 +1187,7 @@ class RealTimeStreaming: if event_type == "session.created" and self._should_disable_vad_auto_response(): self.store_message(event_obj) await self.websocket.send_text(self._event_to_client_json(event_obj)) - await self._send_to_backend(self._make_disable_auto_response_message()) + await self._send_session_gate_once() return True if event_type == "conversation.item.input_audio_transcription.completed": @@ -1140,6 +1241,10 @@ class RealTimeStreaming: if self._usage_meter is not None: self._usage_meter.observe_backend_event(event) + if event.get("type") == "error" and self._is_guardrail_cancel_error(event): + verbose_logger.debug("Realtime: not forwarding an error caused by a guardrail event: %s", event) + continue + if await self._handle_raw_backend_message(event, raw_response): continue @@ -1472,19 +1577,21 @@ class RealTimeStreaming: and self.session_configuration_request is None and not self._guardrail_turn_detection_update_sent and self._should_disable_vad_auto_response() - ): - session: Mapping[str, object] | None = msg_obj.setdefault("session", {}) - if isinstance(session, dict): - existing_td = session.get("turn_detection") - if not isinstance(existing_td, dict): - existing_td = {} - existing_td["create_response"] = False - session["turn_detection"] = existing_td - message = json.dumps(msg_obj) - guardrail_turn_detection_injected = True - verbose_logger.debug( - "Injected turn_detection into first session.update for audio transcription guardrails" + and ( + gated_first := _gate_session_update( + client_event.get("session", _EMPTY_SESSION), + self._gate_turn_detection, + inject_gate_when_absent=True, ) + ) + is not None + ): + msg_obj["session"] = gated_first + message = json.dumps(msg_obj) + guardrail_turn_detection_injected = True + verbose_logger.debug( + "Injected turn_detection into first session.update for audio transcription guardrails" + ) ## GUARDRAIL: Force ``create_response`` to False in any # client-provided ``turn_detection`` so a later @@ -1492,52 +1599,25 @@ class RealTimeStreaming: # and bypass the transcription guardrail after the # initial disable. Covers both the flat beta key and the # nested GA ``audio.input.turn_detection`` shape, since - # the GA remap below also accepts either form. Skipped - # when the injection block above already ran for this - # message, to avoid redundant double-serialization. + # the GA remap below also accepts either form. A later + # update with no turn_detection is forwarded as is: OpenAI + # keeps the earlier disable for an unset field, and xAI + # falls back to its manual default, which never auto-responds. if ( msg_type == "session.update" and not guardrail_turn_detection_injected and self._should_disable_vad_auto_response() + and ( + gated_later := _gate_session_update( + client_event.get("session"), + self._gate_turn_detection, + inject_gate_when_absent=False, + ) + ) + is not None ): - session = client_event.get("session") - if isinstance(session, dict): - td_overridden = False - flat_td = session.get("turn_detection") - flat_td_present = flat_td is not None - if flat_td_present: - if not isinstance(flat_td, dict): - flat_td = {} - if flat_td.get("create_response") is not False: - flat_td["create_response"] = False - session["turn_detection"] = flat_td - td_overridden = True - nested_td_present = False - audio = session.get("audio") - if isinstance(audio, dict): - audio_input = audio.get("input") - if isinstance(audio_input, dict): - nested_td = audio_input.get("turn_detection") - if nested_td is not None: - nested_td_present = True - if not isinstance(nested_td, dict): - nested_td = {} - if nested_td.get("create_response") is not False: - nested_td["create_response"] = False - audio_input["turn_detection"] = nested_td - td_overridden = True - # Symmetric with the first-update injection block: - # if the client omitted turn_detection entirely on - # a subsequent session.update, still inject the - # ``create_response: False`` override so the - # transcription guardrail cannot be re-enabled by - # any downstream merge that drops the original - # disable. - if not flat_td_present and not nested_td_present: - session["turn_detection"] = {"create_response": False} - td_overridden = True - if td_overridden: - message = json.dumps(msg_obj) + msg_obj["session"] = gated_later + message = json.dumps(msg_obj) # GA compatibility: remap beta-style session fields only when # the upstream is in GA mode. Beta upstreams expect the flat @@ -1590,6 +1670,8 @@ class RealTimeStreaming: # would permanently disable the injection if ``_send_to_backend`` # raised — neither this loop nor # ``_maybe_send_guardrail_turn_detection_update`` would retry. + if msg_type == "session.update" and self.provider_config is None: + await self._send_session_gate_once() sent = await self._send_to_backend(message) if guardrail_turn_detection_injected and sent: self._guardrail_turn_detection_update_sent = True diff --git a/litellm/llms/base_llm/realtime/transcription_protocol.py b/litellm/llms/base_llm/realtime/transcription_protocol.py index c3264911d46..b4514fd1da2 100644 --- a/litellm/llms/base_llm/realtime/transcription_protocol.py +++ b/litellm/llms/base_llm/realtime/transcription_protocol.py @@ -102,6 +102,16 @@ def json_integer( return value +def uses_server_vad( + update: TranscriptionSessionUpdate, backend: str, drop_params: bool, error: ProtocolErrorType +) -> bool: + if update.turn_detection_disabled: + return False + if update.turn_detection_type in (None, "server_vad") or drop_params: + return True + raise error(f"{backend} supports server_vad turn detection or null. Set drop_params to use server_vad instead") + + def new_event_id() -> str: return f"event_{uuid.uuid4().hex}" diff --git a/litellm/llms/gemini/realtime/transformation.py b/litellm/llms/gemini/realtime/transformation.py index 7c04ebf311f..f4e2941670a 100644 --- a/litellm/llms/gemini/realtime/transformation.py +++ b/litellm/llms/gemini/realtime/transformation.py @@ -161,6 +161,10 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): """Google AI Studio Gemini 3.5+ accepts ``id`` on functionResponses; Vertex AI rejects it.""" return True + def _rejects_vad_tuning_with_detection_off(self) -> bool: + """Google AI Studio closes the socket (1007) on padding or silence tuning next to ``disabled: true``.""" + return True + @staticmethod def _usage_detail_alias(details: Mapping[str, int | None] | None, defaults: dict[str, int]) -> dict[str, int]: if not isinstance(details, dict): @@ -279,10 +283,15 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): OpenAI ``semantic_vad`` has no Gemini Live equivalent — return an empty dict so callers omit ``realtimeInputConfig`` (mapping it with - ``disabled: true`` breaks native-audio sessions). + ``disabled: true`` breaks native-audio sessions). An explicit + ``create_response: false`` maps to detection off alone where the backend + rejects tuning next to ``disabled: true``. """ if isinstance(value, dict) and value.get("type") == "semantic_vad" and "create_response" not in value: return AutomaticActivityDetection() + auto_response_off: Final = "create_response" in value and value["create_response"] is False + if auto_response_off and self._rejects_vad_tuning_with_detection_off(): + return AutomaticActivityDetection(disabled=True) automatic_activity_dection: Final = AutomaticActivityDetection() if "create_response" in value and isinstance(value["create_response"], bool): diff --git a/litellm/llms/meta/realtime/transformation.py b/litellm/llms/meta/realtime/transformation.py index 46231e5e980..892d7ff66de 100644 --- a/litellm/llms/meta/realtime/transformation.py +++ b/litellm/llms/meta/realtime/transformation.py @@ -26,6 +26,7 @@ from litellm.llms.base_llm.realtime.transcription_protocol import ( speech_event, transcription_session, transcription_session_created_event, + uses_server_vad, ) from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig from litellm.secret_managers.main import get_secret_str @@ -230,15 +231,7 @@ def _parse_sample_rate(update: TranscriptionSessionUpdate) -> MuseSampleRate: return 16_000 if rate == 16_000 else 24_000 -def _parse_mode(update: TranscriptionSessionUpdate) -> MuseMode: - if update.turn_detection_disabled: - return "PUSH_TO_TALK" - if update.turn_detection_type not in (None, "server_vad"): - raise MuseProtocolError("Muse Voice supports server_vad turn detection or null") - return "ENDPOINTING" - - -def parse_session_update(payload: str, expected_model: str) -> MuseSessionConfig: +def parse_session_update(payload: str, expected_model: str, *, drop_params: bool = False) -> MuseSessionConfig: update: Final = parse_transcription_session_update(payload, MuseProtocolError) if update.session_type not in (None, "transcription", "realtime"): raise MuseProtocolError("Muse Voice supports transcription sessions only") @@ -253,7 +246,7 @@ def parse_session_update(payload: str, expected_model: str) -> MuseSessionConfig raise MuseProtocolError("realtime session model cannot be changed") return MuseSessionConfig( model=normalized_model, - mode=_parse_mode(update), + mode="ENDPOINTING" if uses_server_vad(update, "Muse Voice", drop_params, MuseProtocolError) else "PUSH_TO_TALK", sample_rate=_parse_sample_rate(update), language_bias=() if update.language is None else (normalize_language(update.language),), ) @@ -441,9 +434,11 @@ class MetaRealtimeConfig(BaseRealtimeConfig): *, monotonic: Callable[[], float] = time.monotonic, sleep: Callable[[float], Awaitable[None]] = asyncio.sleep, + drop_params: bool = False, ) -> None: self._monotonic: Final = monotonic self._sleep: Final = sleep + self._drop_params: Final = drop_params self._transformer: Final = MuseEventTransformer() self._access_token: str | None = None self._config: MuseSessionConfig | None = None @@ -553,7 +548,7 @@ class MetaRealtimeConfig(BaseRealtimeConfig): access_token: Final = self._access_token if access_token is None: raise MuseProtocolError("Meta API key was not validated before the session was configured") - config: Final = parse_session_update(message, model) + config: Final = parse_session_update(message, model, drop_params=self._drop_params) self._config = config self._transformer.configure(config) return (json.dumps(config.handshake(access_token), separators=(",", ":")),) diff --git a/litellm/llms/vertex_ai/audio_transcription/realtime_transformation.py b/litellm/llms/vertex_ai/audio_transcription/realtime_transformation.py index 1d4a974fc54..5015678c467 100644 --- a/litellm/llms/vertex_ai/audio_transcription/realtime_transformation.py +++ b/litellm/llms/vertex_ai/audio_transcription/realtime_transformation.py @@ -13,7 +13,6 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.llms.base_llm.realtime.transcription_protocol import ( RealtimeTranscriptionProtocolError, TranscriptionAudioFormat, - TranscriptionSessionUpdate, completed_event, decode_pcm16_append, delta_event, @@ -23,6 +22,7 @@ from litellm.llms.base_llm.realtime.transcription_protocol import ( speech_event, transcription_session, transcription_session_created_event, + uses_server_vad, ) from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig, RealtimeBackend from litellm.llms.vertex_ai.audio_transcription.transformation import ( @@ -125,7 +125,7 @@ def default_session_config(model: str) -> ChirpSessionConfig: ) -def parse_chirp_session_update(payload: str, expected_model: str) -> ChirpSessionConfig: +def parse_chirp_session_update(payload: str, expected_model: str, *, drop_params: bool = False) -> ChirpSessionConfig: update: Final = parse_transcription_session_update(payload, ChirpProtocolError) if update.session_type not in (None, "transcription", "realtime"): raise ChirpProtocolError("Speech-to-Text streaming supports transcription sessions only") @@ -141,7 +141,7 @@ def parse_chirp_session_update(payload: str, expected_model: str) -> ChirpSessio model=model, language=None if update.language is None else normalize_transcription_language_to_bcp47(update.language), sample_rate=_parse_sample_rate(update.audio_format), - server_vad=_parse_server_vad(update), + server_vad=uses_server_vad(update, "Speech-to-Text streaming", drop_params, ChirpProtocolError), ) @@ -161,14 +161,6 @@ def _parse_sample_rate(audio_format: TranscriptionAudioFormat | None) -> int: return rate -def _parse_server_vad(update: TranscriptionSessionUpdate) -> bool: - if update.turn_detection_disabled: - return False - if update.turn_detection_type not in (None, "server_vad"): - raise ChirpProtocolError("Speech-to-Text streaming supports server_vad turn detection or null") - return True - - def session_created_event(config: ChirpSessionConfig, session_id: str) -> OpenAIRealtimeTranscriptionSessionCreated: return transcription_session_created_event(config.openai_session(session_id)) @@ -331,11 +323,13 @@ class VertexChirpRealtimeConfig(BaseRealtimeConfig): project: str, location: str | None, backend_factory: Callable[[SpeechStreamingTarget], RealtimeBackend] = _default_backend_factory, + drop_params: bool = False, ) -> None: self._resolve_access_token: Final = resolve_access_token self._project: Final = validate_vertex_transcription_project_id(project) self._location: Final = validate_vertex_transcription_location(location, DEFAULT_SPEECH_TO_TEXT_LOCATION) self._backend_factory: Final = backend_factory + self._drop_params: Final = drop_params self._transformer: Final = ChirpEventTransformer() self._config: ChirpSessionConfig | None = None self._session_id: str | None = None @@ -423,7 +417,7 @@ class VertexChirpRealtimeConfig(BaseRealtimeConfig): if self._config is not None: verbose_logger.debug("Speech-to-Text streaming: ignoring session.update after the stream was configured") return () - config: Final = parse_chirp_session_update(message, model) + config: Final = parse_chirp_session_update(message, model, drop_params=self._drop_params) self._config = config self._transformer.configure(config, self._session_id or f"sess_{uuid.uuid4().hex}") return (config.configure_command(),) diff --git a/litellm/llms/vertex_ai/realtime/transformation.py b/litellm/llms/vertex_ai/realtime/transformation.py index c14dd535145..55f1cf58cbe 100644 --- a/litellm/llms/vertex_ai/realtime/transformation.py +++ b/litellm/llms/vertex_ai/realtime/transformation.py @@ -47,6 +47,9 @@ class VertexAIRealtimeConfig(GeminiRealtimeConfig): def _include_function_response_id(self) -> bool: return False + def _rejects_vad_tuning_with_detection_off(self) -> bool: + return False + # ------------------------------------------------------------------ # URL # ------------------------------------------------------------------ @@ -248,9 +251,12 @@ def vertex_realtime_config( resolve_access_token: Callable[[], Awaitable[str]], project: str, location: str | None, + drop_params: bool, ) -> VertexAIRealtimeConfig | VertexChirpRealtimeConfig: if is_vertex_speech_to_text_model(model): - return VertexChirpRealtimeConfig(resolve_access_token=resolve_access_token, project=project, location=location) + return VertexChirpRealtimeConfig( + resolve_access_token=resolve_access_token, project=project, location=location, drop_params=drop_params + ) return VertexAIRealtimeConfig( access_token=access_token, project=project, diff --git a/litellm/llms/xai/realtime/transformation.py b/litellm/llms/xai/realtime/transformation.py index 5efe125ee60..ffbf814788b 100644 --- a/litellm/llms/xai/realtime/transformation.py +++ b/litellm/llms/xai/realtime/transformation.py @@ -16,6 +16,7 @@ construction time (see ``handler.py``) so all normalization is isolated here and ``RealTimeStreaming`` stays provider-agnostic. """ +from collections.abc import Mapping from typing import Final @@ -77,6 +78,18 @@ class XAIRealtimeNormalizer: event = self._normalize_response_usage_event(event, event_type) return event + def gate_turn_detection(self, _turn_detection: object, /) -> Mapping[str, object]: + """xAI auto-answers under ``server_vad`` even with ``create_response`` off, so the gate is always untyped.""" + return {"create_response": False} + + @property + def deletes_conversation_items(self) -> bool: + """xAI answers ``conversation.item.delete`` with ``Item not found`` until it commits the audio item. + + The guardrail blocks on an interim transcript, before that commit. + """ + return False + def patch_outgoing_session(self, session: dict) -> dict: """Patch a client ``session.update`` payload before forwarding to xAI. diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py index 367d5be7bfc..aeec0bf67c0 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -500,11 +500,13 @@ async def _arealtime( ) return + drop_params: Final = litellm.drop_params is True or normalize_drop_params(kwargs.get("drop_params")) is True provider_config: BaseRealtimeConfig | None = None if _custom_llm_provider in LlmProviders._member_map_.values(): provider_config = ProviderConfigManager.get_provider_realtime_config( model=model, provider=LlmProviders(_custom_llm_provider), + drop_params=drop_params, ) if provider_config is not None: await base_llm_http_handler.async_realtime( @@ -585,7 +587,6 @@ async def _arealtime( aws_bedrock_runtime_endpoint: Final = kwargs.get("aws_bedrock_runtime_endpoint") aws_external_id: Final = kwargs.get("aws_external_id") aws_session_tags: Final = kwargs.get("aws_session_tags") - drop_params: Final = litellm.drop_params is True or normalize_drop_params(kwargs.get("drop_params")) is True await bedrock_realtime.async_realtime( model=model, @@ -671,6 +672,7 @@ async def _arealtime( resolve_access_token=resolve_vertex_access_token, project=resolved_project, location=vertex_location, + drop_params=drop_params, ) await base_llm_http_handler.async_realtime( diff --git a/litellm/utils.py b/litellm/utils.py index d2ca9faef41..27669d8fd4a 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -9975,6 +9975,7 @@ class ProviderConfigManager: def get_provider_realtime_config( model: str, provider: LlmProviders, + drop_params: bool = False, ) -> BaseRealtimeConfig | None: if LlmProviders.GEMINI == provider: from litellm.llms.gemini.realtime.transformation import GeminiRealtimeConfig @@ -9983,7 +9984,7 @@ class ProviderConfigManager: if LlmProviders.META == provider: from litellm.llms.meta.realtime.transformation import MetaRealtimeConfig - return MetaRealtimeConfig() + return MetaRealtimeConfig(drop_params=drop_params) return None @staticmethod diff --git a/tests/integration/_support/upstream.py b/tests/integration/_support/upstream.py index d2b67c5438e..186bc345fab 100644 --- a/tests/integration/_support/upstream.py +++ b/tests/integration/_support/upstream.py @@ -9,9 +9,10 @@ import struct import uuid import zlib from collections import deque -from collections.abc import AsyncIterator, Mapping +from collections.abc import AsyncIterator, Callable, Iterator, Mapping from dataclasses import dataclass, field from queue import SimpleQueue +from types import MappingProxyType from typing import Final, cast import httpx @@ -42,7 +43,7 @@ from starlette.datastructures import UploadFile as StarletteUploadFile from starlette.requests import Request from starlette.responses import JSONResponse, Response, StreamingResponse from starlette.routing import Route, WebSocketRoute -from starlette.websockets import WebSocket, WebSocketDisconnect +from starlette.websockets import WebSocket, WebSocketDisconnect, WebSocketState JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) INTERNAL_FIELDS: Final = frozenset( @@ -81,6 +82,25 @@ class Observation: api_key: str = "" +@dataclass(frozen=True, slots=True) +class _RealtimeConnection: + websocket: WebSocket + response: RealtimeResponse + session: dict[str, JsonValue] + scenario_id: str + authorization: str + api_key: str + events: Iterator[dict[str, JsonValue]] + typed_turn_detection: deque[JsonValue] + + def auto_answers(self) -> bool: + return ( + self.response.dialect == "xai" + and len(self.typed_turn_detection) > 0 + and _is_typed_server_vad(self.typed_turn_detection[-1]) + ) + + class InteractionState(BaseModel): """What the scripted Interactions API answers for one interaction id until a DELETE drops it.""" @@ -92,6 +112,12 @@ class InteractionState(BaseModel): delay_seconds: float = 0 +class _ReceivedFrame(BaseModel): + model_config = ConfigDict(frozen=True) + type: str + text: str | None = None + + class _ScenarioRegistration(BaseModel): scenario_id: str response: StoredResponse @@ -415,23 +441,47 @@ class Provider: ) -> None: await websocket.accept() session: Final = _realtime_session(response, scenario_id, websocket.query_params.get("model", "")) + frames: Final = _received_texts(websocket) + held: Final = await anext(frames, None) if response.created_after_first_frame else None for _ in range(response.created_repeats): await websocket.send_json({"type": response.created_event, "event_id": _event_id(), "session": session}) - event_index: Final = iter(response.events) - async for message in websocket.iter_json(): - payload: Final = JSON_OBJECT.validate_python(message) - self.observations.put(Observation(websocket.url.path, authorization, payload, "WEBSOCKET_FRAME", api_key)) - if payload.get("type") == "session.update": - await websocket.send_json( - _realtime_update_reply(payload.get("session"), session, response.session_type) + connection: Final = _RealtimeConnection( + websocket, response, session, scenario_id, authorization, api_key, iter(response.events), deque(maxlen=1) + ) + if held is not None: + await self._realtime_frame(connection, held) + async for message in frames: + await self._realtime_frame(connection, message) + + async def _realtime_frame(self, connection: _RealtimeConnection, message: str) -> None: + websocket: Final = connection.websocket + response: Final = connection.response + payload: Final = JSON_OBJECT.validate_json(message) + self.observations.put( + Observation(websocket.url.path, connection.authorization, payload, "WEBSOCKET_FRAME", connection.api_key) + ) + frame_type: Final = payload.get("type") + if frame_type == "session.update": + reply: Final = _realtime_update_reply( + payload.get("session"), connection.session, response.session_type, response.dialect + ) + if reply["type"] == "session.updated": + connection.typed_turn_detection.extend( + filter(_declares_a_type, _declared_turn_detection(payload.get("session"))) ) - continue - if payload.get("type") not in _REALTIME_TRIGGERS: - continue - event: Final = next(event_index, None) - if event is None: - continue - await websocket.send_json(_rendered_realtime_event(event, scenario_id)) + await _reply_while_open(websocket, reply) + return + if frame_type == "response.cancel": + await _reply_while_open(websocket, _cancel_not_active(payload.get("event_id"), response.dialect)) + return + if frame_type not in _REALTIME_TRIGGERS: + return + event: Final = next(connection.events, None) + if event is None: + return + await _reply_while_open(websocket, _rendered_realtime_event(event, connection.scenario_id)) + if frame_type == "input_audio_buffer.commit" and connection.auto_answers(): + await _reply_while_open(websocket, _rendered_realtime_event(dict(_AUTO_ANSWER), connection.scenario_id)) async def muse_realtime(self, websocket: WebSocket) -> None: await websocket.accept() @@ -454,8 +504,23 @@ class Provider: await websocket.send_json(_rendered_realtime_event(pending.popleft(), scenario_id)) async def gemini_live(self, websocket: WebSocket) -> None: - authorization: Final = websocket.headers.get("authorization", "") - scenario_id: Final = websocket.headers.get("x-goog-user-project", "") + await self._gemini_live_session( + websocket, + websocket.headers.get("x-goog-user-project", ""), + websocket.headers.get("authorization", ""), + _accepts_every_setup, + ) + + async def gemini_ai_studio_live(self, websocket: WebSocket) -> None: + await self._gemini_live_session(websocket, websocket.query_params.get("key", ""), "", _ai_studio_setup_refusal) + + async def _gemini_live_session( + self, + websocket: WebSocket, + scenario_id: str, + authorization: str, + setup_refusal: Callable[[Mapping[str, JsonValue]], str | None], + ) -> None: self.observations.put( Observation( websocket.url.path, @@ -477,6 +542,9 @@ class Provider: Observation(websocket.url.path, authorization, payload, "WEBSOCKET_FRAME", scenario_id) ) if "setup" in payload: + if (refusal := setup_refusal(payload)) is not None: + await websocket.close(code=_INVALID_ARGUMENT_CLOSE, reason=refusal[:_CLOSE_REASON_BYTES]) + return await websocket.send_json({"setupComplete": {}}) continue if not _GEMINI_LIVE_TRIGGERS.intersection(payload): @@ -587,6 +655,7 @@ class Provider: WebSocketRoute("/openai/realtime", self.realtime), WebSocketRoute("/v1/asr/realtime", self.muse_realtime), WebSocketRoute(GEMINI_LIVE_PATH, self.gemini_live), + WebSocketRoute(GEMINI_AI_STUDIO_LIVE_PATH, self.gemini_ai_studio_live), ] ) @@ -612,6 +681,19 @@ def _interaction_body(interaction_id: str, state: InteractionState) -> dict[str, _REALTIME_TRIGGERS: Final = frozenset({"response.create", "input_audio_buffer.commit"}) _GEMINI_LIVE_TRIGGERS: Final = frozenset({"clientContent", "realtimeInput"}) GEMINI_LIVE_PATH: Final = "/ws/google.cloud.aiplatform.v1.LlmBidiService/BidiGenerateContent" +GEMINI_AI_STUDIO_LIVE_PATH: Final = "/ws/google.ai.generativelanguage.v1beta.GenerativeService.BidiGenerateContent" +_INVALID_ARGUMENT_CLOSE: Final = 1007 +_CLOSE_REASON_BYTES: Final = 123 +_AI_STUDIO_TUNING_FIELDS: Final = MappingProxyType( + {"prefixPaddingMs": "prefix_padding_ms", "silenceDurationMs": "silence_duration_ms"} +) +_AUTO_ANSWER: Final[Mapping[str, JsonValue]] = MappingProxyType( + { + "type": "response.done", + "event_id": "evt_$UNIQUE_ID", + "response": {"id": "resp_auto_$UNIQUE_ID", "object": "realtime.response", "status": "completed", "output": []}, + } +) _TRANSCRIPTION_UPDATE_REFUSED: Final = "Passing a realtime session update to a transcription session is not allowed." _REALTIME_UPDATE_REFUSED: Final = "Passing a transcription session update to a realtime session is not allowed." _NESTED_TURN_DETECTION_TYPE: Final = "session.audio.input.turn_detection.type" @@ -651,11 +733,11 @@ def _missing_turn_detection_type(session: Mapping[str, JsonValue]) -> str | None def _realtime_update_reply( - session: JsonValue, created: Mapping[str, JsonValue], session_type: str | None + session: JsonValue, created: Mapping[str, JsonValue], session_type: str | None, dialect: str ) -> dict[str, JsonValue]: if not isinstance(session, dict): return _realtime_error("invalid_value", "Invalid value for 'session': expected an object.", "session") - missing: Final = _missing_turn_detection_type(session) + missing: Final = None if dialect == "xai" else _missing_turn_detection_type(session) if missing is not None: return _realtime_error("missing_required_parameter", f"Missing required parameter: '{missing}'.", missing) declared: Final = session.get("type") @@ -666,6 +748,82 @@ def _realtime_update_reply( return {"type": "session.updated", "event_id": _event_id(), "session": {**created, **session}} +async def _received_texts(websocket: WebSocket) -> AsyncIterator[str]: + while (frame := _ReceivedFrame.model_validate(await websocket.receive())).type == "websocket.receive": + if frame.text is not None: + yield frame.text + + +async def _reply_while_open(websocket: WebSocket, frame: Mapping[str, JsonValue]) -> None: + if websocket.application_state is not WebSocketState.CONNECTED: + return + try: + await websocket.send_json(frame) + except WebSocketDisconnect: + return + + +def _cancel_not_active(cancel_event_id: JsonValue, dialect: str) -> dict[str, JsonValue]: + message: Final = "Cancellation failed: no active response found" + if dialect == "xai": + own_id: Final = _event_id() + return { + "type": "error", + "event_id": own_id, + "error": { + "message": message, + "type": "invalid_request_error", + "code": "invalid_request_error", + "event_id": own_id, + }, + } + return { + "type": "error", + "event_id": _event_id(), + "error": { + "type": "invalid_request_error", + "code": "response_cancel_not_active", + "message": message, + "param": None, + "event_id": cancel_event_id, + }, + } + + +def _declared_turn_detection(session: JsonValue) -> tuple[JsonValue, ...]: + if not isinstance(session, dict): + return () + audio: Final = session.get("audio") + audio_input: Final = audio.get("input") if isinstance(audio, dict) else None + if isinstance(audio_input, dict) and "turn_detection" in audio_input: + return (audio_input["turn_detection"],) + return (session["turn_detection"],) if "turn_detection" in session else () + + +def _declares_a_type(turn_detection: JsonValue) -> bool: + return isinstance(turn_detection, dict) and "type" in turn_detection + + +def _is_typed_server_vad(turn_detection: JsonValue) -> bool: + return isinstance(turn_detection, dict) and turn_detection.get("type") == "server_vad" + + +def _accepts_every_setup(_payload: Mapping[str, JsonValue]) -> str | None: + return None + + +def _ai_studio_setup_refusal(payload: Mapping[str, JsonValue]) -> str | None: + setup: Final = payload.get("setup") + config: Final = setup.get("realtimeInputConfig") if isinstance(setup, dict) else None + detection: Final = config.get("automaticActivityDetection") if isinstance(config, dict) else None + if not isinstance(detection, dict) or detection.get("disabled") is not True: + return None + tuning: Final = tuple(name for key, name in _AI_STUDIO_TUNING_FIELDS.items() if key in detection) + if not tuning: + return None + return f"automatic_activity_detection.disabled is true, but the following fields were also set: {', '.join(tuning)}" + + def _rendered_realtime_event(event: Mapping[str, JsonValue], scenario_id: str) -> dict[str, JsonValue]: return JSON_OBJECT.validate_json( json.dumps(event, separators=(",", ":")) diff --git a/tests/integration/cost_calculation/cost_tracking_case.py b/tests/integration/cost_calculation/cost_tracking_case.py index 038c8d6b9e3..208d95e6a14 100644 --- a/tests/integration/cost_calculation/cost_tracking_case.py +++ b/tests/integration/cost_calculation/cost_tracking_case.py @@ -235,6 +235,8 @@ class RealtimeResponse(BaseModel): session_type: str | None = None created_event: Literal["session.created", "transcription_session.created"] = "session.created" created_repeats: int = 1 + created_after_first_frame: bool = False + dialect: Literal["openai", "xai"] = "openai" StoredResponse: TypeAlias = Annotated[ diff --git a/tests/integration/providers/test_realtime_transcription_guardrail_session_update.py b/tests/integration/providers/test_realtime_transcription_guardrail_session_update.py index f12254b8d1d..ba2261e5a1f 100644 --- a/tests/integration/providers/test_realtime_transcription_guardrail_session_update.py +++ b/tests/integration/providers/test_realtime_transcription_guardrail_session_update.py @@ -8,10 +8,10 @@ while a client frame alone never flips a voice session into one. Every row runs which answers the injected and rewritten updates the way the vendors do. A blocked transcript on a transcription session reaches the client as a ``guardrail_violation`` error and -nothing else: the ``response.cancel``, the voiced block prompt and the ``response.create`` a voice session gets -never go to a backend that cannot speak, while ``on_violation: end_session`` and ``end_session_after_n_fails`` -still close the session. A guardrail that raises anything but a block closes the session the way it does on a -voice session. +nothing else: the ``response.cancel``, the blocked item's deletion, the voiced block prompt and the +``response.create`` a voice session gets never go to a backend that cannot speak, while +``on_violation: end_session`` and ``end_session_after_n_fails`` still close the session. A guardrail that raises +anything but a block closes the session the way it does on a voice session. """ from __future__ import annotations @@ -51,6 +51,7 @@ from tests.integration._support.database import read_rows from tests.integration._support.process import ( OwnedProxy, UpstreamSlot, + graceful_stop_seconds, owned_proxy_process, owned_upstream, stop_root_process, @@ -60,7 +61,7 @@ from tests.integration.cost_calculation.cost_tracking_case import RealtimeRespon RecordProperty = Callable[[str, object], None] -pytestmark: Final = pytest.mark.timeout(180) +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) WORKERS: Final = 2 BURST: Final = 20 @@ -83,7 +84,7 @@ CREATED_TYPES: Final = frozenset({"session.created", "transcription_session.crea TRANSCRIPT_COMPLETED: Final = "conversation.item.input_audio_transcription.completed" RESPONSE_DONE: Final = "response.done" GUARDRAIL_VIOLATION: Final = ("guardrail_violation", "content_policy_violation") -MISSING_TURN_DETECTION_TYPE: Final = ("invalid_request_error", "missing_required_parameter") +SESSION_TYPE_MISMATCH: Final = ("invalid_request_error", "invalid_parameter") INVALID_SESSION_VALUE: Final = ("invalid_request_error", "invalid_value") FIVE_KB: Final = "x" * 5120 UNAUTHENTICATED_STATUS: Final = 403 @@ -104,10 +105,9 @@ BETA_TRANSCRIPTION_UPDATE: Final[dict[str, JsonValue]] = { "input_audio_transcription": {"model": TRANSCRIBE_MODEL}, } VOICE_UPDATE: Final[dict[str, JsonValue]] = {"type": "realtime", "instructions": "answer briefly"} -VOICE_UPDATE_FORWARDED: Final[dict[str, JsonValue]] = { - "type": "realtime", - "instructions": "answer briefly", - "audio": {"input": {"turn_detection": {"create_response": False}}}, +VOICE_UPDATE_GATED: Final[dict[str, JsonValue]] = { + **VOICE_UPDATE, + "audio": {"input": {"turn_detection": {"type": "server_vad", "create_response": False}}}, } RE_ENABLE_UPDATE: Final[dict[str, JsonValue]] = { "type": "realtime", @@ -123,14 +123,32 @@ CLIENT_DECLARED_TRANSCRIPTION: Final[dict[str, JsonValue]] = { } CLIENT_DECLARED_TRANSCRIPTION_FORWARDED: Final[dict[str, JsonValue]] = { "type": "transcription", - "audio": {"input": {"turn_detection": {"create_response": False}}}, + "audio": {"input": {"turn_detection": {"type": "server_vad", "create_response": False}}}, } -TURN_DETECTION_TYPE_MISSING: Final = "Missing required parameter: 'session.audio.input.turn_detection.type'." +REALTIME_UPDATE_REFUSED: Final = "Passing a transcription session update to a realtime session is not allowed." GA_INJECTED_UPDATE: Final[dict[str, JsonValue]] = { "type": "realtime", "audio": {"input": {"turn_detection": {"type": "server_vad", "create_response": False}}}, } BETA_INJECTED_UPDATE: Final[dict[str, JsonValue]] = {"turn_detection": {"type": "server_vad", "create_response": False}} +XAI_VOICE_MODEL: Final = "grok-voice" +GEMINI_LIVE_MODEL: Final = "gemini-3.8-live" +VERTEX_LIVE_LOCATION: Final = "us-central1" +UNTYPED_GATE: Final[dict[str, JsonValue]] = {"create_response": False} +SEMANTIC_VAD: Final[dict[str, JsonValue]] = {"type": "semantic_vad", "eagerness": "low"} +TUNED_SERVER_VAD: Final[dict[str, JsonValue]] = { + "type": "server_vad", + "prefix_padding_ms": 300, + "silence_duration_ms": 500, +} +GEMINI_TUNED_DETECTION: Final[dict[str, JsonValue]] = { + "disabled": True, + "prefixPaddingMs": 300, + "silenceDurationMs": 500, +} +CLIENT_CANCEL_ID: Final = "client_cancel_1" +CANCEL_NOT_ACTIVE: Final = ("invalid_request_error", "response_cancel_not_active") +AUTO_ANSWER_PREFIX: Final = "resp_auto_" COMMIT: Final[dict[str, JsonValue]] = {"type": "input_audio_buffer.commit"} APPEND: Final[dict[str, JsonValue]] = {"type": "input_audio_buffer.append", "audio": PCM_200MS_24KHZ} PUSH_TO_TALK: Final = "PUSH_TO_TALK" @@ -172,7 +190,13 @@ GUARDRAIL_END_CLOSE: Final = 1000 PROXY_FAILURE_CLOSE: Final = 1011 PROXY_FAILURE_REASON: Final = "proxy failed while relaying the upstream websocket" TOOL_OUTPUT_BLOCKED: Final = json.dumps({"error": "Tool output blocked by content policy"}) -VOICE_BLOCK_FRAMES: Final = ("response.cancel", "conversation.item.create", "response.create") +VOICE_BLOCK_FRAMES: Final = ( + "response.cancel", + "conversation.item.delete", + "conversation.item.create", + "response.create", +) +XAI_VOICE_BLOCK_FRAMES: Final = ("response.cancel", "conversation.item.create", "response.create") CREATED_SECONDS: Final = 60.0 STEP_SECONDS: Final = 15.0 SDK_SECONDS: Final = 45.0 @@ -318,6 +342,30 @@ def _voice_scenario(*transcripts: JsonValue) -> RealtimeResponse: return RealtimeResponse(content_type="application/x-realtime", events=tuple(_voice_turns(transcripts))) +def _xai_voice_scenario(transcript: str) -> RealtimeResponse: + return RealtimeResponse( + content_type="application/x-realtime", events=tuple(_voice_turns((transcript,))), dialect="xai" + ) + + +def _voice_scenario_created_after_the_first_frame(transcript: str) -> RealtimeResponse: + return RealtimeResponse( + content_type="application/x-realtime", events=tuple(_voice_turns((transcript,))), created_after_first_frame=True + ) + + +def _live_scenario() -> RealtimeResponse: + return RealtimeResponse(content_type="application/x-realtime", events=()) + + +def _realtime_update(turn_detection: JsonValue) -> dict[str, JsonValue]: + return {"type": "realtime", "audio": {"input": {"turn_detection": turn_detection}}} + + +def _without_auto_response(turn_detection: dict[str, JsonValue]) -> dict[str, JsonValue]: + return {**turn_detection, "create_response": False} + + def _muse_scenario(transcript: str) -> RealtimeResponse: return RealtimeResponse( content_type="application/x-realtime", @@ -358,8 +406,41 @@ def _with_transcription_model(update: dict[str, JsonValue], model: str) -> dict[ return {**update, "audio": {"input": {**audio_input, "transcription": {"model": model}}}} -def _muse_deployment(scenario: Scenario, scenario_id: str, api_base: str) -> str: - return scenario.model(model=f"meta/{MUSE_MODEL}", api_key=scenario_id, api_base=api_base) +def _muse_deployment(scenario: Scenario, scenario_id: str, api_base: str, *, drop_params: bool) -> str: + return scenario.model(model=f"meta/{MUSE_MODEL}", api_key=scenario_id, api_base=api_base, drop_params=drop_params) + + +def _gemini_studio_deployment(scenario: Scenario, scenario_id: str, upstream_url: str) -> str: + return scenario.model( + model=f"gemini/{GEMINI_LIVE_MODEL}", + api_key=scenario_id, + api_base=upstream_url.rstrip("/"), + model_info={"mode": "realtime"}, + ) + + +def _vertex_credentials(upstream_url: str) -> str: + return json.dumps( + { + "type": "external_account", + "audience": "synthetic-vertex-audience", + "subject_token_type": "urn:ietf:params:oauth:token-type:jwt", + "token_url": f"{upstream_url}/_oauth/token", + "credential_source": {"url": f"{upstream_url}/health"}, + } + ) + + +def _vertex_live_deployment(scenario: Scenario, project: str, upstream_url: str) -> str: + upstream: Final = upstream_url.rstrip("/") + return scenario.model( + model=f"vertex_ai/{GEMINI_LIVE_MODEL}", + vertex_project=project, + vertex_credentials=_vertex_credentials(upstream), + vertex_location=VERTEX_LIVE_LOCATION, + api_base=upstream, + model_info={"mode": "realtime"}, + ) def _named_deployment_id(creator: Gateway, scenario: Scenario, name: str, model: str, scenario_id: str) -> str: @@ -549,6 +630,55 @@ def _transcribe( return asyncio.run(_session(ws_base, path, query, key, frames, until=until, headers=headers)) +async def _received_after_sending( + socket: ClientConnection, frames: tuple[dict[str, JsonValue], ...], until: str +) -> AsyncIterator[dict[str, JsonValue]]: + try: + for frame in frames: + await socket.send(json.dumps(frame)) + async for event in _until(socket, until, STEP_SECONDS): + yield event + except ConnectionClosed as closed: + yield {"type": "closed", "code": _close_code(closed)} + + +def _is_proxy_answer(event: dict[str, JsonValue]) -> bool: + if event.get("type") != RESPONSE_DONE: + return False + return not string_value(object_value(event["response"])["id"]).startswith(AUTO_ANSWER_PREFIX) + + +async def _until_the_proxy_answer( + socket: ClientConnection, frames: tuple[dict[str, JsonValue], ...] +) -> AsyncIterator[dict[str, JsonValue]]: + try: + yield JSON_OBJECT.validate_json(await asyncio.wait_for(socket.recv(), CREATED_SECONDS)) + for frame in frames: + await socket.send(json.dumps(frame)) + deadline: Final = asyncio.get_running_loop().time() + STEP_SECONDS + while (event := await _next_event(socket, deadline)) is not None: + yield event + if _is_proxy_answer(event): + return + yield {"type": "timeout"} + except ConnectionClosed as closed: + yield {"type": "closed", "code": _close_code(closed)} + + +async def _talked_until_the_proxy_answer( + ws_base: str, query: str, key: str, frames: tuple[dict[str, JsonValue], ...] +) -> Session: + headers: Final = {"Authorization": f"Bearer {key}"} + async with websockets.connect(f"{ws_base}/v1/realtime?{query}", additional_headers=headers) as socket: + return Session(tuple([event async for event in _until_the_proxy_answer(socket, frames)]), None) + + +async def _talked_at_connect(ws_base: str, query: str, key: str, frames: tuple[dict[str, JsonValue], ...]) -> Session: + headers: Final = {"Authorization": f"Bearer {key}"} + async with websockets.connect(f"{ws_base}/v1/realtime?{query}", additional_headers=headers) as socket: + return Session(tuple([event async for event in _received_after_sending(socket, frames, RESPONSE_DONE)]), None) + + def _talk( ws_base: str, query: str, @@ -617,10 +747,43 @@ def _session_updates(observed: tuple[dict[str, JsonValue], ...]) -> tuple[JsonVa return tuple(frame.get("session") for frame in _sent(observed) if frame.get("type") == "session.update") +def _assert_voice_updates(observed: tuple[dict[str, JsonValue], ...], *later: dict[str, JsonValue]) -> None: + updates: Final = _session_updates(observed) + assert updates[0] == GA_INJECTED_UPDATE, updates + assert updates[1] == VOICE_UPDATE_GATED, updates + assert updates[2:] == later, updates + + def _query(upgrade: dict[str, JsonValue]) -> JsonValue: return upgrade["query"] +def _declared_turn_detection(update: JsonValue) -> JsonValue: + session: Final = object_value(update) + audio_input: Final = object_value(object_value(session.get("audio", {})).get("input", {})) + return audio_input["turn_detection"] if "turn_detection" in audio_input else session.get("turn_detection") + + +def _upstream_turn_detections(observed: tuple[dict[str, JsonValue], ...]) -> tuple[JsonValue, ...]: + return tuple(_declared_turn_detection(update) for update in _session_updates(observed)) + + +def _response_ids(session: Session) -> tuple[str, ...]: + done: Final = tuple(object_value(event["response"]) for event in session.events if event["type"] == RESPONSE_DONE) + return tuple(string_value(response["id"]) for response in done) + + +def _error_event_ids(session: Session) -> tuple[JsonValue, ...]: + errors: Final = tuple(object_value(event["error"]) for event in session.events if event["type"] == "error") + return tuple(error.get("event_id") for error in errors) + + +def _setup_detection(observed: tuple[dict[str, JsonValue], ...]) -> JsonValue: + setups: Final = tuple(object_value(frame["setup"]) for frame in _sent(observed) if "setup" in frame) + assert len(setups) == 1, _sent(observed) + return object_value(setups[0].get("realtimeInputConfig", {})).get("automaticActivityDetection") + + def _spend_rows(key: str, count: int) -> list[dict[str, JsonValue]]: return eventually( lambda: read_rows(SPEND_SQL, (sha256(key.encode()).hexdigest(),)), @@ -704,6 +867,7 @@ def _write_config(directory: Path, ca_bundle: Path) -> Path: "litellm_settings": { "ssl_verify": str(ca_bundle), "enable_redis_auth_cache": True, + "gemini_live_defer_setup": True, "cache": True, "cache_params": {"type": "redis", "host": "os.environ/REDIS_HOST", "port": "os.environ/REDIS_PORT"}, }, @@ -918,19 +1082,18 @@ def test_voice_session_keeps_the_injected_update_and_the_clean_transcript_flows( assert session.types == ( "session.created", "session.updated", - "error", + "session.updated", TRANSCRIPT_COMPLETED, RESPONSE_DONE, ), session - assert session.errors == (MISSING_TURN_DETECTION_TYPE,), session - assert session.error_messages == (TURN_DETECTION_TYPE_MISSING,), session + assert session.errors == (), session assert _sent_types(observed) == ( "session.update", "session.update", "input_audio_buffer.commit", "response.create", ), _sent(observed) - assert _session_updates(observed) == (GA_INJECTED_UPDATE, VOICE_UPDATE_FORWARDED), _session_updates(observed) + _assert_voice_updates(observed) rows: Final = _spend_rows(key, 1) assert rows[0]["call_type"] == "_arealtime", rows @@ -945,17 +1108,15 @@ def test_voice_session_blocked_transcript_is_refused_through_the_backend(guardra ) observed: Final = _observed(guardrail_proxy.gateway.upstream_url, handle.scenario_id) assert session.transcripts == (BLOCKED_TRANSCRIPT,), session - assert session.errors == (MISSING_TURN_DETECTION_TYPE, GUARDRAIL_VIOLATION), session + assert session.errors == (GUARDRAIL_VIOLATION,), session assert session.types[-1] == RESPONSE_DONE, session assert _sent_types(observed) == ( "session.update", "session.update", "input_audio_buffer.commit", - "response.cancel", - "conversation.item.create", - "response.create", + *VOICE_BLOCK_FRAMES, ), _sent(observed) - assert _session_updates(observed) == (GA_INJECTED_UPDATE, VOICE_UPDATE_FORWARDED), _session_updates(observed) + _assert_voice_updates(observed) def test_voice_session_cannot_re_enable_auto_response_with_a_later_update(guardrail_proxy: OwnedProxy) -> None: @@ -967,12 +1128,8 @@ def test_voice_session_cannot_re_enable_auto_response_with_a_later_update(guardr session: Final = _talk(_ws_base(_owned_url(guardrail_proxy)), f"model={model}", key, frames) observed: Final = _observed(guardrail_proxy.gateway.upstream_url, handle.scenario_id) assert session.types[-1] == RESPONSE_DONE, session - assert session.errors == (MISSING_TURN_DETECTION_TYPE,), session - assert _session_updates(observed) == ( - GA_INJECTED_UPDATE, - VOICE_UPDATE_FORWARDED, - RE_ENABLE_FORWARDED, - ), _session_updates(observed) + assert session.errors == (), session + _assert_voice_updates(observed, RE_ENABLE_FORWARDED) def test_client_declared_transcription_type_does_not_bypass_the_voice_guardrail(guardrail_proxy: OwnedProxy) -> None: @@ -984,16 +1141,14 @@ def test_client_declared_transcription_type_does_not_bypass_the_voice_guardrail( session: Final = _talk(_ws_base(_owned_url(guardrail_proxy)), f"model={model}", key, frames, seconds=20) observed: Final = _observed(guardrail_proxy.gateway.upstream_url, handle.scenario_id) assert session.transcripts == (BLOCKED_TRANSCRIPT,), session - assert session.errors == (MISSING_TURN_DETECTION_TYPE, GUARDRAIL_VIOLATION), session - assert session.error_messages[0] == TURN_DETECTION_TYPE_MISSING, session + assert session.errors == (SESSION_TYPE_MISMATCH, GUARDRAIL_VIOLATION), session + assert session.error_messages[0] == REALTIME_UPDATE_REFUSED, session assert session.types[-1] == RESPONSE_DONE, session assert _sent_types(observed) == ( "session.update", "session.update", "input_audio_buffer.commit", - "response.cancel", - "conversation.item.create", - "response.create", + *VOICE_BLOCK_FRAMES, ), _sent(observed) assert _session_updates(observed) == ( GA_INJECTED_UPDATE, @@ -1014,6 +1169,202 @@ def test_backend_session_created_typed_transcription_skips_the_injection_on_the_ assert [_query(upgrade) for upgrade in _upgrades(observed)] == [[["model", TRANSCRIBE_MODEL]]] +def test_voice_session_typed_turn_detection_keeps_its_type_and_tuning(guardrail_proxy: OwnedProxy) -> None: + with guardrail_proxy.gateway.scenario() as scenario: + handle: Final = _scripted(scenario, _voice_scenario(CLEAN_TRANSCRIPT)) + key: Final = scenario.key() + model: Final = _openai_deployment(scenario, handle.scenario_id, model=VOICE_MODEL) + frames: Final = _voice_frames(_realtime_update(SEMANTIC_VAD)) + session: Final = _talk(_ws_base(_owned_url(guardrail_proxy)), f"model={model}", key, frames) + observed: Final = _observed(guardrail_proxy.gateway.upstream_url, handle.scenario_id) + assert session.errors == (), session + assert session.types[-1] == RESPONSE_DONE, session + assert _session_updates(observed) == ( + GA_INJECTED_UPDATE, + _realtime_update(_without_auto_response(SEMANTIC_VAD)), + ), _session_updates(observed) + + +@dataclass(frozen=True, slots=True) +class LaterUpdateCase: + headers: Mapping[str, str] + first: dict[str, JsonValue] + later: dict[str, JsonValue] + forwarded: tuple[dict[str, JsonValue], ...] + + +LATER_UPDATE_CASES: Final = ( + pytest.param( + LaterUpdateCase( + {}, + VOICE_UPDATE, + {"type": "realtime", "instructions": "later"}, + (GA_INJECTED_UPDATE, VOICE_UPDATE_GATED), + ), + id="ga", + ), + pytest.param( + LaterUpdateCase( + BETA_HEADERS, + {"turn_detection": {"type": "server_vad"}}, + {"instructions": "later"}, + (BETA_INJECTED_UPDATE, BETA_INJECTED_UPDATE), + ), + id="beta", + ), +) + + +@pytest.mark.parametrize("case", LATER_UPDATE_CASES) +def test_later_update_without_turn_detection_is_forwarded_verbatim( + guardrail_proxy: OwnedProxy, case: LaterUpdateCase +) -> None: + with guardrail_proxy.gateway.scenario() as scenario: + handle: Final = _scripted(scenario, _voice_scenario(CLEAN_TRANSCRIPT)) + key: Final = scenario.key() + model: Final = _openai_deployment(scenario, handle.scenario_id, model=VOICE_MODEL) + session: Final = _drive( + _ws_base(_owned_url(guardrail_proxy)), + f"model={model}", + key, + (_step(_update_frame(case.first), _update_frame(case.later), COMMIT, until=RESPONSE_DONE),), + headers=case.headers, + ) + observed: Final = _observed(guardrail_proxy.gateway.upstream_url, handle.scenario_id) + assert session.errors == (), session + assert session.types[-1] == RESPONSE_DONE, session + assert _session_updates(observed) == (*case.forwarded, case.later), _session_updates(observed) + + +def test_session_update_sent_before_session_created_follows_the_gate(guardrail_proxy: OwnedProxy) -> None: + with guardrail_proxy.gateway.scenario() as scenario: + handle: Final = _scripted(scenario, _voice_scenario_created_after_the_first_frame(CLEAN_TRANSCRIPT)) + key: Final = scenario.key() + model: Final = _openai_deployment(scenario, handle.scenario_id, model=VOICE_MODEL) + frames: Final = _voice_frames(_realtime_update(TUNED_SERVER_VAD)) + session: Final = asyncio.run( + _talked_at_connect(_ws_base(_owned_url(guardrail_proxy)), f"model={model}", key, frames) + ) + observed: Final = _observed(guardrail_proxy.gateway.upstream_url, handle.scenario_id) + assert session.types == ( + "session.created", + "session.updated", + "session.updated", + TRANSCRIPT_COMPLETED, + RESPONSE_DONE, + ), session + assert session.errors == (), session + assert _session_updates(observed) == ( + GA_INJECTED_UPDATE, + _realtime_update(_without_auto_response(TUNED_SERVER_VAD)), + ), _session_updates(observed) + + +def test_client_response_cancel_error_still_reaches_the_client(guardrail_proxy: OwnedProxy) -> None: + with guardrail_proxy.gateway.scenario() as scenario: + handle: Final = _scripted(scenario, _voice_scenario(CLEAN_TRANSCRIPT)) + key: Final = scenario.key() + model: Final = _openai_deployment(scenario, handle.scenario_id, model=VOICE_MODEL) + cancel: Final[dict[str, JsonValue]] = {"type": "response.cancel", "event_id": CLIENT_CANCEL_ID} + session: Final = _drive( + _ws_base(_owned_url(guardrail_proxy)), f"model={model}", key, (_step(cancel, until="error"),) + ) + assert session.errors == (CANCEL_NOT_ACTIVE,), session + assert _error_event_ids(session) == (CLIENT_CANCEL_ID,), session + + +def _xai_blocked_session( + guardrail_proxy: OwnedProxy, scenario: Scenario +) -> tuple[Session, tuple[dict[str, JsonValue], ...]]: + handle: Final = _scripted(scenario, _xai_voice_scenario(BLOCKED_TRANSCRIPT)) + key: Final = scenario.key() + model: Final = scenario.model(model=f"xai/{XAI_VOICE_MODEL}", api_key=handle.scenario_id) + frames: Final = _voice_frames(_realtime_update(TUNED_SERVER_VAD)) + session: Final = asyncio.run( + _talked_until_the_proxy_answer(_ws_base(_owned_url(guardrail_proxy)), f"model={model}", key, frames) + ) + return session, _observed(guardrail_proxy.gateway.upstream_url, handle.scenario_id) + + +def test_xai_voice_session_gets_untyped_gates_and_no_unguarded_answer(guardrail_proxy: OwnedProxy) -> None: + with guardrail_proxy.gateway.scenario() as scenario: + session, observed = _xai_blocked_session(guardrail_proxy, scenario) + turn_detections: Final = _upstream_turn_detections(observed) + assert turn_detections == (UNTYPED_GATE, UNTYPED_GATE), turn_detections + assert session.transcripts == (BLOCKED_TRANSCRIPT,), session + assert GUARDRAIL_VIOLATION in session.errors, session + assert not any(response.startswith(AUTO_ANSWER_PREFIX) for response in _response_ids(session)), session + + +def test_xai_voice_session_block_sends_xai_no_item_delete(guardrail_proxy: OwnedProxy) -> None: + with guardrail_proxy.gateway.scenario() as scenario: + session, observed = _xai_blocked_session(guardrail_proxy, scenario) + assert session.transcripts == (BLOCKED_TRANSCRIPT,), session + assert _sent_types(observed) == ( + "session.update", + "session.update", + "input_audio_buffer.commit", + *XAI_VOICE_BLOCK_FRAMES, + ), _sent(observed) + + +@pytest.mark.xfail( + strict=True, + raises=AssertionError, + reason=( + "xAI answers the guardrail's response.cancel with an error carrying its own event_id, so the proxy cannot " + "tell it apart from an error for a client event and forwards it, as the merge base does" + ), +) +def test_xai_guardrail_cancel_error_does_not_reach_the_client(guardrail_proxy: OwnedProxy) -> None: + with guardrail_proxy.gateway.scenario() as scenario: + session, _ = _xai_blocked_session(guardrail_proxy, scenario) + assert session.errors == (GUARDRAIL_VIOLATION,), session + + +TUNED_VOICE_UPDATES: Final = ( + pytest.param(_realtime_update(TUNED_SERVER_VAD), id="nested"), + pytest.param({"turn_detection": TUNED_SERVER_VAD}, id="flat"), +) + + +@pytest.mark.parametrize("update", TUNED_VOICE_UPDATES) +def test_gemini_ai_studio_voice_session_with_vad_tuning_connects_under_the_guardrail( + guardrail_proxy: OwnedProxy, update: dict[str, JsonValue] +) -> None: + with guardrail_proxy.gateway.scenario() as scenario: + handle: Final = _scripted(scenario, _live_scenario()) + key: Final = scenario.key() + upstream_url: Final = guardrail_proxy.gateway.upstream_url + model: Final = _gemini_studio_deployment(scenario, handle.scenario_id, upstream_url) + session: Final = _drive( + _ws_base(_owned_url(guardrail_proxy)), + f"model={model}", + key, + (_step(_update_frame(update), until="session.updated"),), + ) + observed: Final = _observed(upstream_url, handle.scenario_id) + assert session.types == ("session.created", "session.updated"), session + assert _setup_detection(observed) == {"disabled": True}, _sent(observed) + + +def test_vertex_live_voice_session_keeps_vad_tuning_under_the_guardrail(guardrail_proxy: OwnedProxy) -> None: + with guardrail_proxy.gateway.scenario() as scenario: + handle: Final = _scripted(scenario, _live_scenario()) + key: Final = scenario.key() + upstream_url: Final = guardrail_proxy.gateway.upstream_url + model: Final = _vertex_live_deployment(scenario, handle.scenario_id, upstream_url) + session: Final = _drive( + _ws_base(_owned_url(guardrail_proxy)), + f"model={model}", + key, + (_step(_update_frame(_realtime_update(TUNED_SERVER_VAD)), until="session.updated"),), + ) + observed: Final = _observed(upstream_url, handle.scenario_id) + assert session.types == ("session.created", "session.updated"), session + assert _setup_detection(observed) == GEMINI_TUNED_DETECTION, _sent(observed) + + MUSE_PUSH_TO_TALK_FRAMES: Final[tuple[dict[str, JsonValue], ...]] = ( {"binary_bytes": MUSE_PACKET_BYTES}, {"binary_bytes": MUSE_PACKET_BYTES}, @@ -1041,10 +1392,11 @@ def _muse_session( turn_detection: JsonValue, *, until: str | None = None, + drop_params: bool = False, ) -> tuple[Session, tuple[dict[str, JsonValue], ...]]: handle: Final = _scripted(scenario, _muse_scenario(transcript), control_url=tls_upstream) key: Final = scenario.key() - model: Final = _muse_deployment(scenario, handle.scenario_id, tls_upstream) + model: Final = _muse_deployment(scenario, handle.scenario_id, tls_upstream, drop_params=drop_params) verdict: Final = TRANSCRIPT_COMPLETED if BLOCKED_WORD not in transcript else "error" session: Final = _talk( _ws_base(_owned_url(guardrail_proxy)), @@ -1120,6 +1472,40 @@ def test_muse_server_vad_transcription_session_keeps_endpointing( ), _muse_audio_frames(observed) +def test_muse_semantic_vad_under_drop_params_falls_back_to_endpointing( + guardrail_proxy: OwnedProxy, tls_upstream: str +) -> None: + with guardrail_proxy.gateway.scenario() as scenario: + session, observed = _muse_session( + guardrail_proxy, + scenario, + tls_upstream, + CLEAN_TRANSCRIPT, + f"&{TRANSCRIPTION_QUERY}", + {"type": "semantic_vad"}, + drop_params=True, + ) + assert session.transcripts == (CLEAN_TRANSCRIPT,), session + assert session.errors == (), session + assert _muse_handshake(observed)["mode"] == ENDPOINTING, _muse_handshake(observed) + + +def test_muse_semantic_vad_without_drop_params_is_refused(guardrail_proxy: OwnedProxy, tls_upstream: str) -> None: + with guardrail_proxy.gateway.scenario() as scenario: + session, observed = _muse_session( + guardrail_proxy, + scenario, + tls_upstream, + CLEAN_TRANSCRIPT, + "", + {"type": "semantic_vad"}, + until=TRANSCRIPT_COMPLETED, + ) + assert session.types[-1] == "closed", session + assert session.transcripts == (), session + assert _upgrades(observed) == (), observed + + @pytest.mark.xfail( strict=True, reason=( @@ -1229,8 +1615,8 @@ def test_opt_in_guardrail_gates_a_voice_session_on_an_opted_out_key(guardrail_pr ) observed: Final = _observed(guardrail_proxy.gateway.upstream_url, handle.scenario_id) assert session.transcripts == (BLOCKED_TRANSCRIPT,), session - assert session.errors == (MISSING_TURN_DETECTION_TYPE, GUARDRAIL_VIOLATION), session - assert _session_updates(observed) == (GA_INJECTED_UPDATE, VOICE_UPDATE_FORWARDED), _session_updates(observed) + assert session.errors == (GUARDRAIL_VIOLATION,), session + _assert_voice_updates(observed) def test_opt_in_guardrail_leaves_a_transcription_session_alone_on_an_opted_out_key(guardrail_proxy: OwnedProxy) -> None: @@ -1717,12 +2103,11 @@ def test_voice_session_guardrail_raising_runtime_error_closes_with_the_same_prox assert session.types == ( "session.created", "session.updated", - "error", + "session.updated", TRANSCRIPT_COMPLETED, "error", "closed", ), session - assert session.errors[0] == MISSING_TURN_DETECTION_TYPE, session _assert_closed_by_a_proxy_failure(session) assert _sent_types(observed) == ("session.update", "session.update", "input_audio_buffer.commit"), _sent( observed @@ -1740,8 +2125,8 @@ def test_voice_session_guardrail_raising_value_error_is_voiced_through_the_backe (_step(VOICE_UPDATE_FRAME, COMMIT, until=RESPONSE_DONE),), ) assert session.transcripts == (blocked,), session - assert session.errors == (MISSING_TURN_DETECTION_TYPE, GUARDRAIL_VIOLATION), session - assert session.error_messages[1] == f"{VALUE_ERROR_WORD} is not allowed", session + assert session.errors == (GUARDRAIL_VIOLATION,), session + assert session.error_messages == (f"{VALUE_ERROR_WORD} is not allowed",), session assert session.types[-1] == RESPONSE_DONE, session assert _sent_types(observed) == ( "session.update", @@ -1764,7 +2149,7 @@ def test_voice_session_second_violation_under_end_session_after_n_fails_closes_a (_step(VOICE_UPDATE_FRAME, COMMIT, until=RESPONSE_DONE), _step(COMMIT, until="closed")), ) assert session.transcripts == (blocked, blocked), session - assert session.errors == (MISSING_TURN_DETECTION_TYPE, GUARDRAIL_VIOLATION, GUARDRAIL_VIOLATION), session + assert session.errors == (GUARDRAIL_VIOLATION, GUARDRAIL_VIOLATION), session assert session.types[-1] == "closed", session assert session.close_code == GUARDRAIL_END_CLOSE, session assert _sent_types(observed) == ( @@ -1948,7 +2333,7 @@ def test_non_string_transcript_field_closes_a_voice_session_with_the_same_proxy_ assert session.types == ( "session.created", "session.updated", - "error", + "session.updated", TRANSCRIPT_COMPLETED, "error", "closed", diff --git a/tests/integration/providers/test_realtime_unguarded_voice_turns.py b/tests/integration/providers/test_realtime_unguarded_voice_turns.py index 4dd7fd9aeef..c7489685273 100644 --- a/tests/integration/providers/test_realtime_unguarded_voice_turns.py +++ b/tests/integration/providers/test_realtime_unguarded_voice_turns.py @@ -114,12 +114,10 @@ TRANSCRIPT_FIELDS: Final = ( pytest.param({"transcript": {"text": "a"}}, id="object"), ) VOICE_UPDATE: Final[dict[str, JsonValue]] = {"type": "realtime", "instructions": "answer briefly"} -VOICE_UPDATE_FORWARDED: Final[dict[str, JsonValue]] = { - "type": "realtime", - "instructions": "answer briefly", - "audio": {"input": {"turn_detection": {"create_response": False}}}, +VOICE_UPDATE_GATED: Final[dict[str, JsonValue]] = { + **VOICE_UPDATE, + "audio": {"input": {"turn_detection": {"type": "server_vad", "create_response": False}}}, } -TURN_DETECTION_TYPE_MISSING: Final = "Missing required parameter: 'session.audio.input.turn_detection.type'." INJECTED_UPDATE: Final[dict[str, JsonValue]] = { "type": "realtime", "audio": {"input": {"turn_detection": {"type": "server_vad", "create_response": False}}}, @@ -982,15 +980,15 @@ def test_list_mode_transcript_guardrail_still_turns_auto_response_off_and_answer assert session.types == ( "session.created", "session.updated", - "error", + "session.updated", TRANSCRIPT_COMPLETED, RESPONSE_DONE, ), session assert session.transcript_fields == ({"transcript": TRANSCRIPT},), session - assert session.error_messages == (TURN_DETECTION_TYPE_MISSING,), session + assert session.error_messages == (), session assert _sent(observed) == ( _update_frame(INJECTED_UPDATE), - _update_frame(VOICE_UPDATE_FORWARDED), + _update_frame(VOICE_UPDATE_GATED), COMMIT, PROXY_RESPONSE_CREATE, ), _sent(observed) diff --git a/tests/unit/integration_support/test_upstream.py b/tests/unit/integration_support/test_upstream.py index 0cbcf1b7ee1..9e114e3b0b7 100644 --- a/tests/unit/integration_support/test_upstream.py +++ b/tests/unit/integration_support/test_upstream.py @@ -21,7 +21,11 @@ async def test_realtime_disconnect_after_session_creation_keeps_unexpected_error response: Final = RealtimeResponse(content_type="application/x-realtime", events=({"type": "response.done"},)) provider.scenario_store.put("owned-scenario", response) incoming: Final[deque[Message]] = deque( - ({"type": "websocket.connect"}, {"type": "websocket.receive", "text": '{"type":"response.create"}'}) + ( + {"type": "websocket.connect"}, + {"type": "websocket.receive", "text": '{"type":"response.create"}'}, + {"type": "websocket.disconnect", "code": 1006}, + ) ) outgoing: Final[deque[Message]] = deque() diff --git a/tests/unit/litellm_core_utils/test_realtime_streaming.py b/tests/unit/litellm_core_utils/test_realtime_streaming.py index 8de5e27d99c..9b2c42450ae 100644 --- a/tests/unit/litellm_core_utils/test_realtime_streaming.py +++ b/tests/unit/litellm_core_utils/test_realtime_streaming.py @@ -2351,8 +2351,7 @@ async def test_guardrail_turn_detection_injection_tolerates_non_dict_value( injected_turn_detection = session_obj.get("turn_detection") or session_obj.get("audio", {}).get("input", {}).get( "turn_detection" ) - assert isinstance(injected_turn_detection, dict) - assert injected_turn_detection["create_response"] is False + assert injected_turn_detection == {"type": "server_vad", "create_response": False} assert streaming._guardrail_turn_detection_update_sent is True @@ -2417,8 +2416,7 @@ async def test_subsequent_session_update_cannot_reenable_vad_when_guardrails_act forwarded_turn_detection = session_obj.get("turn_detection") or session_obj.get("audio", {}).get("input", {}).get( "turn_detection" ) - assert isinstance(forwarded_turn_detection, dict) - assert forwarded_turn_detection["create_response"] is False + assert forwarded_turn_detection == {"type": "server_vad", "create_response": False} @pytest.mark.asyncio @@ -3927,3 +3925,317 @@ async def test_transcription_session_guardrail_block_only_reports_violation( sent_to_backend: Final = [call.args[0] for call in backend_ws.send.await_args_list] assert sent_to_backend == [], sent_to_backend assert backend_ws.close.await_count == (1 if expect_session_closed else 0) + + +_GA_GUARDRAIL_GATE: Final = { + "type": "realtime", + "audio": {"input": {"turn_detection": {"type": "server_vad", "create_response": False}}}, +} + + +async def _forwarded_ga_session_updates( + monkeypatch: pytest.MonkeyPatch, + *client_sessions: dict[str, object], + event_normalizer: XAIRealtimeNormalizer | None = None, +) -> tuple[dict[str, object], ...]: + """Run a GA client's session.updates through a transcript-guardrailed proxy and return the forwarded sessions.""" + monkeypatch.setattr(litellm, "callbacks", [_transcription_guardrail()]) + client_ws: Final = AsyncMock() + client_ws.scope = {"headers": []} + client_ws.receive_text = AsyncMock( + side_effect=[ + *(json.dumps({"type": "session.update", "session": session}) for session in client_sessions), + ConnectionClosed(None, None), + ] + ) + backend_ws: Final = MagicMock() + backend_ws.send = AsyncMock() + logging_obj: Final = MagicMock() + logging_obj.async_success_handler = AsyncMock() + logging_obj.success_handler = MagicMock() + await RealTimeStreaming(client_ws, backend_ws, logging_obj, event_normalizer=event_normalizer).client_ack_messages() + return tuple(json.loads(call.args[0])["session"] for call in backend_ws.send.await_args_list) + + +@pytest.mark.asyncio +async def test_first_ga_session_update_keeps_client_turn_detection_when_guardrail_disables_auto_response( + monkeypatch: pytest.MonkeyPatch, +): + """A transcript guardrail only flips create_response: the client's nested GA + turn_detection, its type included, must reach the backend intact, or OpenAI + rejects the whole update and the client's session config is lost.""" + forwarded: Final = await _forwarded_ga_session_updates( + monkeypatch, + { + "type": "realtime", + "instructions": "hi", + "audio": { + "input": { + "turn_detection": {"type": "server_vad", "silence_duration_ms": 700, "create_response": True}, + "transcription": {"model": "gpt-4o-transcribe"}, + } + }, + }, + ) + + gate, client_update = forwarded + assert gate == _GA_GUARDRAIL_GATE + assert "turn_detection" not in client_update + assert client_update["instructions"] == "hi" + assert client_update["audio"]["input"] == { + "turn_detection": {"type": "server_vad", "silence_duration_ms": 700, "create_response": False}, + "transcription": {"model": "gpt-4o-transcribe"}, + } + + +@pytest.mark.asyncio +async def test_first_ga_session_update_without_turn_detection_gets_a_typed_gate(monkeypatch: pytest.MonkeyPatch): + forwarded: Final = await _forwarded_ga_session_updates(monkeypatch, {"type": "realtime", "instructions": "hi"}) + + gate, client_update = forwarded + assert gate == _GA_GUARDRAIL_GATE + assert "turn_detection" not in client_update + assert client_update["audio"]["input"]["turn_detection"] == {"type": "server_vad", "create_response": False} + + +@pytest.mark.asyncio +async def test_later_ga_session_update_without_turn_detection_is_forwarded_without_one( + monkeypatch: pytest.MonkeyPatch, +): + """OpenAI keeps the earlier create_response=false for a field left unset, while a + bare create_response carries no type and gets the whole update rejected.""" + forwarded: Final = await _forwarded_ga_session_updates( + monkeypatch, + {"type": "realtime", "audio": {"input": {"turn_detection": {"type": "server_vad", "create_response": True}}}}, + {"type": "realtime", "instructions": "later"}, + ) + + gate, first, later = forwarded + assert gate == _GA_GUARDRAIL_GATE + assert first["audio"]["input"]["turn_detection"] == {"type": "server_vad", "create_response": False} + assert later == {"type": "realtime", "instructions": "later"} + + +_XAI_TYPED_VAD_SESSION: Final = { + "type": "realtime", + "audio": {"input": {"turn_detection": {"type": "server_vad", "silence_duration_ms": 700}}}, +} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "client_sessions", + [ + pytest.param(({"type": "realtime", "instructions": "hi"},), id="no-turn-detection"), + pytest.param(({"type": "realtime", "audio": {"input": {"turn_detection": None}}},), id="push-to-talk"), + pytest.param((_XAI_TYPED_VAD_SESSION,), id="typed-server-vad"), + pytest.param( + ({"type": "realtime", "instructions": "hi"}, _XAI_TYPED_VAD_SESSION), id="later-typed-server-vad" + ), + ], +) +async def test_xai_session_update_gate_leaves_out_the_vad_type( + monkeypatch: pytest.MonkeyPatch, client_sessions: tuple[dict[str, object], ...] +): + """xAI answers on its own under a typed turn_detection even with create_response off, + so every turn_detection the proxy forwards to xAI under the guardrail must carry no type.""" + forwarded: Final = await _forwarded_ga_session_updates( + monkeypatch, *client_sessions, event_normalizer=XAIRealtimeNormalizer() + ) + + assert [session["audio"]["input"]["turn_detection"] for session in forwarded] == [ + {"create_response": False} + ] * (len(client_sessions) + 1) + + +@pytest.mark.asyncio +async def test_xai_session_created_gate_leaves_out_the_vad_type(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "callbacks", [_transcription_guardrail()]) + client_ws: Final = MagicMock() + client_ws.send_text = AsyncMock() + backend_ws: Final = MagicMock() + backend_ws.send = AsyncMock() + backend_ws.recv = AsyncMock( + side_effect=[ + json.dumps({"type": "session.created", "session": {"turn_detection": {"type": None}}}).encode(), + ConnectionClosed(None, None), + ] + ) + logging_obj: Final = _mock_logging() + logging_obj.async_success_handler = AsyncMock() + logging_obj.success_handler = MagicMock() + streaming: Final = RealTimeStreaming(client_ws, backend_ws, logging_obj, event_normalizer=XAIRealtimeNormalizer()) + + await streaming.backend_to_client_send_messages() + + sent: Final = [json.loads(call.args[0]) for call in backend_ws.send.await_args_list] + assert sent == [ + { + "type": "session.update", + "session": {"type": "realtime", "audio": {"input": {"turn_detection": {"create_response": False}}}}, + } + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "client_turn_detection", + [ + pytest.param({"type": "semantic_vad", "eagerness": "low"}, id="semantic-vad"), + pytest.param({"type": "server_vad", "silence_duration_ms": 700}, id="tuned-server-vad"), + ], +) +async def test_guardrail_gate_goes_ahead_of_a_session_update_sent_before_session_created( + monkeypatch: pytest.MonkeyPatch, client_turn_detection: dict[str, object] +): + """OpenAI applies events in order, so the gate must precede a client update sent on connect: + the client's turn_detection then wins, and when OpenAI refuses that update the gate still holds.""" + monkeypatch.setattr(litellm, "callbacks", [_transcription_guardrail()]) + client_ws: Final = AsyncMock() + client_ws.scope = {"headers": []} + client_ws.receive_text = AsyncMock( + side_effect=[ + json.dumps( + { + "type": "session.update", + "session": {"type": "realtime", "audio": {"input": {"turn_detection": client_turn_detection}}}, + } + ), + ConnectionClosed(None, None), + ] + ) + backend_ws: Final = MagicMock() + backend_ws.send = AsyncMock() + backend_ws.recv = AsyncMock( + side_effect=[json.dumps({"type": "session.created", "session": {}}).encode(), ConnectionClosed(None, None)] + ) + logging_obj: Final = _mock_logging() + logging_obj.async_success_handler = AsyncMock() + logging_obj.success_handler = MagicMock() + streaming: Final = RealTimeStreaming(client_ws, backend_ws, logging_obj) + + await streaming.client_ack_messages() + await streaming.backend_to_client_send_messages() + + sent: Final = [json.loads(call.args[0]) for call in backend_ws.send.await_args_list] + assert [event["session"]["audio"]["input"]["turn_detection"] for event in sent] == [ + {"type": "server_vad", "create_response": False}, + {**client_turn_detection, "create_response": False}, + ] + + +@pytest.mark.asyncio +async def test_client_never_sees_the_error_caused_by_the_guardrail_cancel(monkeypatch: pytest.MonkeyPatch): + """OpenAI answers the guardrail's response.cancel with response_cancel_not_active when nothing + is answering yet, an error about an event the client never sent.""" + + class _BlockingGuardrail(CustomGuardrail): + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + raise ValueError("blocked") + + monkeypatch.setattr( + litellm, + "callbacks", + [ + _BlockingGuardrail( + guardrail_name="blocking", event_hook=GuardrailEventHooks.realtime_input_transcription, default_on=True + ) + ], + ) + client_ws: Final = MagicMock() + client_ws.send_text = AsyncMock() + backend_ws: Final = MagicMock() + backend_ws.send = AsyncMock() + + def cancel_not_active(event_id: str) -> bytes: + return json.dumps( + {"type": "error", "error": {"code": "response_cancel_not_active", "event_id": event_id}} + ).encode() + + def proxy_cancel_event_id() -> str: + sent = [json.loads(call.args[0]) for call in backend_ws.send.await_args_list] + return next(event["event_id"] for event in sent if event["type"] == "response.cancel") + + backend_frames: Final = iter( + ( + lambda: _make_transcript_event("blocked words"), + lambda: cancel_not_active(proxy_cancel_event_id()), + lambda: cancel_not_active("evt_from_client"), + lambda: cancel_not_active("litellm_guardrail_from_client"), + ) + ) + + async def recv(**_kwargs: object) -> bytes: + frame = next(backend_frames, None) + if frame is None: + raise ConnectionClosed(None, None) + return frame() + + backend_ws.recv = recv + logging_obj: Final = _mock_logging() + logging_obj.async_success_handler = AsyncMock() + logging_obj.success_handler = MagicMock() + + await RealTimeStreaming(client_ws, backend_ws, logging_obj).backend_to_client_send_messages() + + client_errors: Final = [ + event["error"] for call in client_ws.send_text.await_args_list if (event := json.loads(call.args[0]))["type"] == "error" + ] + assert [(error["code"], error.get("event_id")) for error in client_errors] == [ + ("content_policy_violation", None), + ("response_cancel_not_active", "evt_from_client"), + ("response_cancel_not_active", "litellm_guardrail_from_client"), + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("event_normalizer", "expected_frames"), + [ + (None, ["response.cancel", "conversation.item.delete", "conversation.item.create", "response.create"]), + (XAIRealtimeNormalizer(), ["response.cancel", "conversation.item.create", "response.create"]), + ], + ids=["deletes_the_blocked_item", "xai_keeps_the_item_it_cannot_delete"], +) +async def test_blocked_transcript_is_removed_from_the_conversation_before_the_refusal_turn( + monkeypatch: pytest.MonkeyPatch, + event_normalizer: XAIRealtimeNormalizer | None, + expected_frames: list[str], +): + class _BlockingGuardrail(CustomGuardrail): + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + raise ValueError("blocked") + + monkeypatch.setattr( + litellm, + "callbacks", + [ + _BlockingGuardrail( + guardrail_name="blocking", event_hook=GuardrailEventHooks.realtime_input_transcription, default_on=True + ) + ], + ) + client_ws: Final = MagicMock() + client_ws.send_text = AsyncMock() + backend_ws: Final = MagicMock() + backend_ws.send = AsyncMock() + backend_frames: Final = iter((_make_transcript_event("blocked words", item_id="item_blocked"),)) + + async def recv(**_kwargs: object) -> bytes: + frame = next(backend_frames, None) + if frame is None: + raise ConnectionClosed(None, None) + return frame + + backend_ws.recv = recv + logging_obj: Final = _mock_logging() + logging_obj.async_success_handler = AsyncMock() + logging_obj.success_handler = MagicMock() + + await RealTimeStreaming( + client_ws, backend_ws, logging_obj, event_normalizer=event_normalizer + ).backend_to_client_send_messages() + + sent: Final = [json.loads(call.args[0]) for call in backend_ws.send.await_args_list] + assert [event["type"] for event in sent] == expected_frames + assert all(event["item_id"] == "item_blocked" for event in sent if event["type"] == "conversation.item.delete") diff --git a/tests/unit/llms/gemini/realtime/test_gemini_realtime_transformation.py b/tests/unit/llms/gemini/realtime/test_gemini_realtime_transformation.py index bcd5f3d8d19..017f279d6c4 100644 --- a/tests/unit/llms/gemini/realtime/test_gemini_realtime_transformation.py +++ b/tests/unit/llms/gemini/realtime/test_gemini_realtime_transformation.py @@ -1,6 +1,6 @@ import json from collections.abc import Mapping -from typing import cast +from typing import Final, cast from unittest.mock import MagicMock import pytest @@ -1323,6 +1323,40 @@ def test_gemini_realtime_pipecat_semantic_vad_omits_realtime_input_config(): assert setup["tools"][0]["function_declarations"][0]["name"] == "terminate_call" +@pytest.mark.parametrize( + ("create_response", "expected_detection"), + [ + (False, {"disabled": True}), + (True, {"disabled": False, "prefixPaddingMs": 200, "silenceDurationMs": 700}), + ], +) +def test_gemini_turn_detection_tuning_is_dropped_when_auto_response_is_off( + create_response: bool, expected_detection: dict[str, object] +): + session_update: Final = { + "type": "session.update", + "session": { + "audio": { + "input": { + "turn_detection": { + "type": "server_vad", + "create_response": create_response, + "prefix_padding_ms": 200, + "silence_duration_ms": 700, + } + }, + }, + }, + } + messages: Final = GeminiRealtimeConfig().transform_realtime_request( + json.dumps(session_update), + "gemini-live-2.5-flash-native-audio", + session_configuration_request=None, + ) + setup: Final = json.loads(messages[0])["setup"] + assert setup["realtimeInputConfig"]["automaticActivityDetection"] == expected_detection + + def test_gemini_input_audio_buffer_commit_maps_to_audio_stream_end(): config = GeminiRealtimeConfig() setup = { diff --git a/tests/unit/llms/meta/realtime/test_meta_realtime_transformation.py b/tests/unit/llms/meta/realtime/test_meta_realtime_transformation.py index a5d7e47fb65..8b46ea53713 100644 --- a/tests/unit/llms/meta/realtime/test_meta_realtime_transformation.py +++ b/tests/unit/llms/meta/realtime/test_meta_realtime_transformation.py @@ -20,6 +20,8 @@ from litellm.llms.meta.realtime.transformation import ( session_created_event, ) from litellm.types.llms.meta import MuseMode +from litellm.types.utils import LlmProviders +from litellm.utils import ProviderConfigManager from litellm.types.realtime import RealtimeResponseTransformInput EMPTY_TRANSFORM_INPUT: Final[RealtimeResponseTransformInput] = { @@ -164,6 +166,22 @@ def test_session_rejects_unsupported_audio_model_and_hints(session: dict[str, ob parse_session_update(_event("session.update", session={"type": "transcription", **session}), MUSE_MODEL) +def _handshake(turn_detection: str, drop_params: bool) -> tuple[str | bytes, ...]: + config = ProviderConfigManager.get_provider_realtime_config(MUSE_MODEL, LlmProviders.META, drop_params=drop_params) + assert config is not None + config.validate_environment({}, MUSE_MODEL, api_key="secret-token") + return tuple(config.transform_realtime_request(_ga_session_update(turn_detection=turn_detection), MUSE_MODEL)) + + +def test_semantic_vad_falls_back_to_server_vad_under_drop_params(): + assert _handshake("semantic_vad", drop_params=True) == _handshake("server_vad", drop_params=False) + + +def test_semantic_vad_without_drop_params_names_the_drop_params_fallback(): + with pytest.raises(MuseProtocolError, match="Set drop_params to use server_vad instead"): + _handshake("semantic_vad", drop_params=False) + + def test_session_created_event_exposes_openai_transcription_shape(): config = parse_session_update( _event( diff --git a/tests/unit/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_transformation.py b/tests/unit/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_transformation.py index b50d6e19458..5c7099832ba 100644 --- a/tests/unit/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_transformation.py +++ b/tests/unit/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_transformation.py @@ -17,6 +17,7 @@ from litellm.llms.vertex_ai.audio_transcription.realtime_transformation import ( parse_chirp_session_update, ) from litellm.llms.vertex_ai.common_utils import VertexAIError +from litellm.llms.vertex_ai.realtime.transformation import vertex_realtime_config from litellm.types.llms.vertex_ai_speech_to_text import ( VertexSpeechStreamingConfigured, VertexSpeechStreamingResponse, @@ -173,6 +174,26 @@ def test_unsupported_session_settings_are_rejected(payload: str, message: str): parse_chirp_session_update(payload, MODEL) +def _configure_command(turn_detection: str, drop_params: bool) -> list[object]: + config = vertex_realtime_config( + MODEL, access_token="token", resolve_access_token=_token, project="proj-1", location="us", drop_params=drop_params + ) + config.transform_session_created_event(MODEL, "sess_1") + return [ + json.loads(command) if isinstance(command, str) else command + for command in config.transform_realtime_request(_ga_session_update(turn_detection=turn_detection), MODEL) + ] + + +def test_semantic_vad_falls_back_to_server_vad_under_drop_params(): + assert _configure_command("semantic_vad", drop_params=True) == _configure_command("server_vad", drop_params=False) + + +def test_semantic_vad_without_drop_params_names_the_drop_params_fallback(): + with pytest.raises(ChirpProtocolError, match="Set drop_params to use server_vad instead"): + _configure_command("semantic_vad", drop_params=False) + + def test_session_update_configures_once_and_later_updates_are_ignored(): config = _config() config.transform_session_created_event(MODEL, "sess_1") diff --git a/tests/unit/llms/vertex_ai/realtime/test_vertex_ai_realtime_transformation.py b/tests/unit/llms/vertex_ai/realtime/test_vertex_ai_realtime_transformation.py index a705c169782..0bfff870d15 100644 --- a/tests/unit/llms/vertex_ai/realtime/test_vertex_ai_realtime_transformation.py +++ b/tests/unit/llms/vertex_ai/realtime/test_vertex_ai_realtime_transformation.py @@ -223,6 +223,36 @@ def test_vertex_session_update_normalizes_ga_remapped_fields( ) +def test_vertex_keeps_turn_detection_tuning_when_auto_response_is_off(): + session_update: Final = { + "type": "session.update", + "session": { + "audio": { + "input": { + "turn_detection": { + "type": "server_vad", + "create_response": False, + "prefix_padding_ms": 200, + "silence_duration_ms": 700, + } + }, + }, + }, + } + cfg: Final = VertexAIRealtimeConfig(access_token="tok", project="my-proj", location="us-central1") + messages: Final = cfg.transform_realtime_request( + json.dumps(session_update), + _NATIVE_AUDIO_MODEL, + session_configuration_request=None, + ) + setup: Final = json.loads(messages[0])["setup"] + assert setup["realtimeInputConfig"]["automaticActivityDetection"] == { + "disabled": True, + "prefixPaddingMs": 200, + "silenceDurationMs": 700, + } + + # --------------------------------------------------------------------------- # Round-trip test: text-in / text-out via RealTimeStreaming # --------------------------------------------------------------------------- From 977a58121372025f63e74657a589e18e066f59a0 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Fri, 16 Oct 2026 08:52:42 -0700 Subject: [PATCH 21/23] fix(router): stop queueing budget increments for Redis when no Redis is configured (#46831) * fix(router): stop queueing budget increments for Redis when no Redis is configured Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): assert Redis-less model budget limiter tracks spend locally without queueing Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/router_strategy/budget_limiter.py | 5 +- .../sdk/test_router_budget_limiter_wire.py | 94 +++++++++++++++++++ ...test_unit_test_max_model_budget_limiter.py | 8 +- .../router_strategy/test_budget_limiter.py | 38 ++++---- 4 files changed, 119 insertions(+), 26 deletions(-) diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index b20f71cd51f..1622a30b95f 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/litellm/router_strategy/budget_limiter.py @@ -405,7 +405,8 @@ class RouterBudgetLimiting(CustomLogger): Runs once the budget start time exists in Redis Cache (on the 2nd and subsequent requests to the same provider) - Increments the spend in memory cache (so spend instantly updated in memory) - - Queues the increment operation to Redis Pipeline (using batched pipeline to optimize performance. Using Redis for multi instance environment of LiteLLM) + - Queues the increment for the batched Redis pipeline when Redis is configured, since only a multi instance + deployment has anything to sync """ increment_op: Final = RedisPipelineIncrementOperation( key=spend_key, @@ -418,6 +419,8 @@ class RouterBudgetLimiting(CustomLogger): value=response_cost, ttl=ttl, ) + if self.dual_cache.redis_cache is None: + return self.redis_increment_operation_queue.append(increment_op) def _get_redis_increment_queue_lock(self) -> asyncio.Lock: diff --git a/tests/integration/sdk/test_router_budget_limiter_wire.py b/tests/integration/sdk/test_router_budget_limiter_wire.py index cf4fda41891..0863ee6d6cc 100644 --- a/tests/integration/sdk/test_router_budget_limiter_wire.py +++ b/tests/integration/sdk/test_router_budget_limiter_wire.py @@ -1,7 +1,9 @@ from __future__ import annotations +import asyncio import json import os +import time import uuid from collections.abc import Iterator from typing import Final @@ -222,3 +224,95 @@ async def test_tag_budget_blocks_exhausted_tag_but_allows_another( assert response.choices[0].message.content == "openai" assert len(requests) == 1 assert json.loads(requests[0].body)["model"] == _MODEL + + + +_ISOLATED_SPEND_KEY: Final = "provider_spend:hosted_vllm:1d" +_INPUT_COST_PER_TOKEN: Final = 0.001 +_OUTPUT_COST_PER_TOKEN: Final = 0.002 + + +def _isolated_provider_deployment(api_base: str) -> dict[str, JsonValue]: + """Only these two tests route to hosted_vllm, so success callbacks still draining from earlier tests cannot land here.""" + return { + "model_name": "isolated-budget", + "litellm_params": { + "model": "hosted_vllm/llama-3", + "api_base": api_base, + "api_key": _API_KEY, + "input_cost_per_token": _INPUT_COST_PER_TOKEN, + "output_cost_per_token": _OUTPUT_COST_PER_TOKEN, + }, + "model_info": {"id": "vllm-deployment"}, + } + + +def _isolated_budget_router(api_base: str, redis_client: Redis | None) -> Router: + redis_params: Final = ( + {} + if redis_client is None + else {"redis_host": os.environ["REDIS_HOST"], "redis_port": int(os.environ["REDIS_PORT"])} + ) + return Router( + model_list=[_isolated_provider_deployment(api_base)], + provider_budget_config={"hosted_vllm": BudgetConfig(budget_duration="1d", max_budget=10.0)}, + num_retries=0, + disable_cooldowns=True, + **redis_params, + ) + + +async def _complete_and_wait_for_spend(router: Router, previous: float) -> float: + response: Final = await router.acompletion( + model="isolated-budget", + messages=[{"role": "user", "content": "budget tracking"}], + ) + cost: Final = ( + response.usage.prompt_tokens * _INPUT_COST_PER_TOKEN + response.usage.completion_tokens * _OUTPUT_COST_PER_TOKEN + ) + deadline: Final = time.monotonic() + 5.0 + while time.monotonic() < deadline: + spend: float | None = router.cache.in_memory_cache.get_cache(_ISOLATED_SPEND_KEY) + if spend is not None and spend > previous: + assert spend == pytest.approx(previous + cost) + return spend + await asyncio.sleep(0.05) + raise AssertionError(f"{_ISOLATED_SPEND_KEY} never rose above {previous}") + + +@pytest.mark.asyncio +async def test_router_without_redis_tracks_spend_but_queues_no_redis_increments() -> None: + """A Router with no Redis never flushes, so anything it queues for Redis is retained for the life of the process.""" + with wire_server(_peer) as wire: + router: Final = _isolated_budget_router(f"{wire.url}/openai/v1", redis_client=None) + first: Final = await _complete_and_wait_for_spend(router, previous=0.0) + second: Final = await _complete_and_wait_for_spend(router, previous=first) + third: Final = await _complete_and_wait_for_spend(router, previous=second) + wire.drain() + + assert third > second > first > 0.0 + assert router.router_budget_logger is not None + assert router.router_budget_logger.redis_increment_operation_queue == [] + + +async def _wait_for_redis_spend(redis_client: Redis, at_least: float) -> float: + deadline: Final = time.monotonic() + 15.0 + while time.monotonic() < deadline: + raw: str | None = redis_client.get(_ISOLATED_SPEND_KEY) + if raw is not None and float(raw) >= at_least: + return float(raw) + await asyncio.sleep(0.1) + raise AssertionError(f"{_ISOLATED_SPEND_KEY} in Redis never reached {at_least}") + + +@pytest.mark.asyncio +async def test_router_with_redis_flushes_queued_spend_increments(redis_client: Redis) -> None: + redis_client.delete(_ISOLATED_SPEND_KEY) + with wire_server(_peer) as wire: + router: Final = _isolated_budget_router(f"{wire.url}/openai/v1", redis_client=redis_client) + first: Final = await _complete_and_wait_for_spend(router, previous=0.0) + second: Final = await _complete_and_wait_for_spend(router, previous=first) + redis_spend: Final = await _wait_for_redis_spend(redis_client, at_least=second) + wire.drain() + + assert redis_spend == pytest.approx(second) diff --git a/tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py b/tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py index 33904637366..c6369a1f3d7 100644 --- a/tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py +++ b/tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py @@ -20,7 +20,6 @@ from litellm.proxy.hooks.model_max_budget_limiter import ( model_budget_start_time_cache_key, resolve_model_budget, ) -from litellm.types.caching import RedisPipelineIncrementOperation async def _reset_at_of_open_window( @@ -500,16 +499,15 @@ async def test_async_log_success_event_pushes_redis_increments_when_redis_config @pytest.mark.asyncio -async def test_model_budget_limiter_initializes_redis_increment_queue_lock(): +async def test_model_budget_limiter_without_redis_tracks_spend_locally_and_queues_nothing(): dual_cache = DualCache() limiter = PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache) spend_key = "virtual_key_spend:test-key:gpt-4:1d" await limiter._increment_spend_in_current_window(spend_key=spend_key, response_cost=0.01, ttl=86400) - assert limiter.redis_increment_operation_queue == [ - RedisPipelineIncrementOperation(key=spend_key, increment_value=0.01, ttl=86400) - ] + assert await dual_cache.async_get_cache(spend_key) == 0.01 + assert limiter.redis_increment_operation_queue == [] @pytest.mark.asyncio diff --git a/tests/unit/router_strategy/test_budget_limiter.py b/tests/unit/router_strategy/test_budget_limiter.py index 685ea46547e..380aa573b3c 100644 --- a/tests/unit/router_strategy/test_budget_limiter.py +++ b/tests/unit/router_strategy/test_budget_limiter.py @@ -14,6 +14,7 @@ import pytest from litellm import Router from litellm.caching.caching import DualCache +from litellm.caching.in_memory_cache import InMemoryCache from litellm.router_strategy.budget_limiter import RouterBudgetLimiting from litellm.types.utils import BudgetConfig import sys, os, asyncio, time, random @@ -232,7 +233,11 @@ async def test_a_monthly_router_budget_ends_a_calendar_month_after_it_opens( """ monkeypatch.setattr(litellm, "tag_budget_config", tag_budget_config) monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) - limiter = RouterBudgetLimiting(dual_cache=DualCache(), clock=charged_at.timestamp, **limiter_config) + limiter = RouterBudgetLimiting( + dual_cache=DualCache(in_memory_cache=InMemoryCache(clock=charged_at.timestamp)), + clock=charged_at.timestamp, + **limiter_config, + ) await limiter.dual_cache.async_set_cache(key=start_time_key, value=opened_at.timestamp()) await _log_success( @@ -243,9 +248,8 @@ async def test_a_monthly_router_budget_ends_a_calendar_month_after_it_opens( }, ) - assert limiter.redis_increment_operation_queue == [ - {"key": spend_key, "increment_value": 0.25, "ttl": int((closes_at - charged_at).total_seconds())} - ] + assert await limiter.dual_cache.async_get_cache(spend_key) == 0.25 + assert limiter.dual_cache.in_memory_cache.ttl_dict[spend_key] == closes_at.timestamp() @pytest.mark.asyncio @@ -362,24 +366,18 @@ async def test_new_budget_window_sets_spend_and_start_time() -> None: @pytest.mark.asyncio -async def test_spend_increment_updates_memory_and_queues_redis_operation() -> None: - provider_budget: Final = RouterBudgetLimiting(dual_cache=DualCache(), provider_budget_config={}) - await provider_budget.dual_cache.async_set_cache(key="provider_spend:openai:1d", value=1.0, ttl=86400) - - await provider_budget._increment_spend_in_current_window( - spend_key="provider_spend:openai:1d", - response_cost=0.5, - ttl=86400, +async def test_redis_less_limiter_tracks_spend_without_queueing_redis_increments(disable_budget_sync) -> None: + """With no Redis there is nothing to flush to, so queued increments would be retained for the life of the process.""" + limiter: Final = RouterBudgetLimiting( + dual_cache=DualCache(), + provider_budget_config={"openai": {"budget_limit": 10.0, "time_period": "1d"}}, ) - assert await provider_budget.dual_cache.async_get_cache("provider_spend:openai:1d") == 1.5 - assert provider_budget.redis_increment_operation_queue == [ - { - "key": "provider_spend:openai:1d", - "increment_value": 0.5, - "ttl": 86400, - } - ] + for _ in range(3): + await _log_success(limiter, _success_kwargs(provider_in_litellm_params="openai", provider_in_payload="openai")) + + assert await limiter.dual_cache.async_get_cache("provider_spend:openai:1d") == pytest.approx(0.75) + assert limiter.redis_increment_operation_queue == [] @pytest.mark.asyncio From e45f2737d33772efd72df232cdc7a6e56154dfd7 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 16 Oct 2026 09:22:22 -0700 Subject: [PATCH 22/23] fix(proxy)!: delete callbacks configured under callbacks or failure_callback (#42740) * fix(proxy): delete callbacks configured under callbacks or failure_callback Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(e2e): assert on configured callback rows and use a callback absent from CI configs Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(e2e): follow logging suite conventions for callback delete tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): flatten remaining callbacks with a helper generator Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): type callback delete helpers to satisfy basedpyright budget Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): save callback config without mutating the loaded dict Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): delete callbacks from every litellm_settings list GET shows POST /config/callback/delete only searched success_callback while GET /get/config/callbacks also lists failure_callback and callbacks, so the Admin UI Delete returned 404 for those rows Remove the name from every list that holds it, matching entries case-insensitively, and skip lists the config file owns when the database owns another copy so the save is not refused Replace the e2e test with an integration test: the shared e2e stack config owns litellm_settings.callbacks, so it cannot exercise that list * test(proxy): cover config-owned success_callback on callback delete and wait on runtime state Add the mirror mixed-ownership case where the config file sets success_callback and the database owns failure_callback, which main refused with a 400, plus a string-valued list case Poll GET listings and /active/callbacks instead of reading them once, so a request served by a worker that has not synced yet cannot fail the test, and reuse the existing callback list key map in the route * test(proxy): annotate callback delete test helper locals with Final * fix(proxy): save callback delete through the config it read so concurrent settings edits are kept --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: yucheng-berri <293926326+yucheng-berri@users.noreply.github.com> --- litellm/proxy/proxy_server.py | 62 +++- .../configuration/test_callback_delete.py | 351 ++++++++++++++++++ .../proxy/proxy_server/test_routes_config.py | 227 +++++++++++ 3 files changed, 628 insertions(+), 12 deletions(-) create mode 100644 tests/integration/configuration/test_callback_delete.py diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 9eb2bc93391..b56f0e68c4f 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -19698,6 +19698,26 @@ async def delete_config_general_settings( return response +def _configured_callback_names(value: object) -> tuple[JsonValue, ...]: + if isinstance(value, str): + return (value,) + if isinstance(value, list): + return tuple(cast(Iterable[JsonValue], value)) # cast-ok: config.yaml values are untyped + return () + + +def _is_callback_name(entry: object, callback_name: str) -> bool: + return isinstance(entry, str) and entry.lower() == callback_name + + +def _callback_lists_to_rewrite( + matching_keys: tuple[str, ...], owned_by_config: Callable[[str], bool] +) -> tuple[str, ...]: + """The matching lists the database owns, or all of them when the config file owns every one.""" + database_owned: Final = tuple(key for key in matching_keys if not owned_by_config(key)) + return database_owned or matching_keys + + @router.post( "/config/callback/delete", tags=["config.yaml"], @@ -19733,24 +19753,42 @@ async def delete_callback( try: # Get current configuration - config: Final = await proxy_config.get_config() + config: Final = cast( # cast-ok: get_config returns an untyped config dict + dict[str, object], + await proxy_config.get_config(), + ) callback_name: Final = data.callback_name.lower() # Check if callback exists in current configuration - litellm_settings: Final = config.get("litellm_settings", {}) - success_callbacks: Final = litellm_settings.get("success_callback", []) + litellm_settings: Final[Mapping[str, object]] = cast( # cast-ok: the config dict's values are untyped + Mapping[str, object], + config.get("litellm_settings", {}), + ) + configured_lists: Final = { + key: _configured_callback_names(litellm_settings.get(key)) for key in _DB_CONFIG_CALLBACK_EVENT_TYPES + } + matching_keys: Final = tuple( + key + for key, names in configured_lists.items() + if any(_is_callback_name(entry, callback_name) for entry in names) + ) - if callback_name not in success_callbacks: + if not matching_keys: raise HTTPException( status_code=404, detail={"error": f"Callback '{callback_name}' not found in active configuration"}, ) - before_success_callbacks: Final = list(success_callbacks) - - # Remove callback from success_callback list - success_callbacks.remove(callback_name) - config.setdefault("litellm_settings", {})["success_callback"] = success_callbacks + rewritten_keys: Final = _callback_lists_to_rewrite(matching_keys, proxy_config.litellm_settings.owned_by_config) + remaining_lists: Final = MappingProxyType( + { + key: tuple(entry for entry in configured_lists[key] if not _is_callback_name(entry, callback_name)) + for key in rewritten_keys + } + ) + before_callbacks: Final[JsonValue] = {key: list(configured_lists[key]) for key in rewritten_keys} + after_callbacks: Final[JsonValue] = {key: list(names) for key, names in remaining_lists.items()} + config["litellm_settings"] = {**litellm_settings, **after_callbacks} # Save the updated configuration await proxy_config.save_config(new_config=config) @@ -19759,8 +19797,8 @@ async def delete_callback( create_config_audit_log( "litellm_settings", "deleted", - {"success_callback": before_success_callbacks}, - {"success_callback": success_callbacks}, + before_callbacks, + after_callbacks, user_api_key_dict, ) ) @@ -19771,7 +19809,7 @@ async def delete_callback( return { "message": f"Successfully deleted callback: {callback_name}", "removed_callback": callback_name, - "remaining_callbacks": success_callbacks, + "remaining_callbacks": tuple(chain.from_iterable(remaining_lists.values())), "deleted_at": datetime.now().isoformat(), } diff --git a/tests/integration/configuration/test_callback_delete.py b/tests/integration/configuration/test_callback_delete.py new file mode 100644 index 00000000000..d58dd252593 --- /dev/null +++ b/tests/integration/configuration/test_callback_delete.py @@ -0,0 +1,351 @@ +import json +import os +import uuid +from collections.abc import Generator, Iterator, Mapping +from contextlib import contextmanager +from dataclasses import dataclass +from pathlib import Path +from types import MappingProxyType +from typing import Final + +import pytest +import yaml +from integration._support.client import JSON_OBJECT, Gateway, JsonValue, eventually, list_value, object_value +from integration._support.database import read_rows, scratch_database +from integration._support.process import owned_proxy +from integration._support.wire import Reply, Request, Wire, wire_server + +CALLBACK: Final = "generic_api" +SIBLING: Final = "langsmith" +LIST_TYPES: Final = MappingProxyType( + { + "success_callback": "success", + "failure_callback": "failure", + "callbacks": "success_and_failure", + } +) +SINK_ENVIRONMENT: Final = MappingProxyType( + { + "GENERIC_LOGGER_ENDPOINT": "http://127.0.0.1:9/unused", + "LANGSMITH_API_KEY": "synthetic-langsmith-key", + "LANGSMITH_BASE_URL": "http://127.0.0.1:9", + } +) +NOT_FOUND: Final = f"Callback '{CALLBACK}' not found in active configuration" +WORKER_SYNC_SECONDS: Final = 40 + + +@dataclass(frozen=True, slots=True) +class Rig: + proxy: Gateway + database_url: str + + +@contextmanager +def _rig(gateway: Gateway, tmp_path: Path, file_settings: Mapping[str, JsonValue]) -> Generator[Rig]: + config: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + path: Final = tmp_path / f"callback_delete_{uuid.uuid4().hex}.yaml" + path.write_text( + yaml.safe_dump({**config, "litellm_settings": {**object_value(config["litellm_settings"]), **file_settings}}) + ) + with ( + scratch_database() as database_url, + owned_proxy( + gateway, + tmp_path, + {"DATABASE_URL": database_url, **SINK_ENVIRONMENT}, + config=path, + remove_environment=("DATABASE_URL_READ_REPLICA",), + workers=2, + ) as proxy, + ): + yield Rig(proxy, database_url) + + +@pytest.fixture +def rig(gateway: Gateway, tmp_path: Path) -> Iterator[Rig]: + with _rig(gateway, tmp_path, {}) as created: + yield created + + +def _configured(proxy: Gateway) -> tuple[tuple[str, str], ...]: + rows: Final = proxy.get("/get/config/callbacks")["callbacks"] + assert isinstance(rows, list) + return tuple( + sorted( + (str(row["name"]), str(row["type"])) + for row in (object_value(entry) for entry in rows) + if not row.get("read_only") + ) + ) + + +def _listed(proxy: Gateway, expected: tuple[tuple[str, str], ...]) -> None: + eventually(lambda: _configured(proxy), lambda rows: rows == expected, seconds=WORKER_SYNC_SECONDS) + + +def _stored_settings(database_url: str) -> Mapping[str, JsonValue]: + rows: Final = read_rows( + 'SELECT param_value FROM "LiteLLM_Config" WHERE param_name = %s', + ("litellm_settings",), + database_url=database_url, + ) + assert len(rows) == 1, rows + value: Final = rows[0]["param_value"] + return object_value(json.loads(value) if isinstance(value, str) else value) + + +def _configure( + proxy: Gateway, settings: Mapping[str, JsonValue], environment: Mapping[str, JsonValue] | None = None +) -> None: + body: Final[Mapping[str, JsonValue]] = { + "litellm_settings": dict(settings), + "environment_variables": None if environment is None else dict(environment), + } + assert proxy.post("/config/update", body) == {"message": "Config updated successfully"} + + +def _delete(proxy: Gateway, name: str) -> tuple[int, Mapping[str, JsonValue]]: + response: Final = proxy.request("POST", "/config/callback/delete", {"callback_name": name}) + return response.status_code, JSON_OBJECT.validate_json(response.content) + + +@pytest.mark.parametrize("list_key", tuple(LIST_TYPES)) +def test_callback_listed_under_any_settings_list_is_deleted_from_the_listing_and_the_store( + rig: Rig, list_key: str +) -> None: + _configure(rig.proxy, {list_key: [CALLBACK]}) + _listed(rig.proxy, ((CALLBACK, LIST_TYPES[list_key]),)) + + status, body = _delete(rig.proxy, CALLBACK) + + assert status == 200, body + assert (body["removed_callback"], body["remaining_callbacks"]) == (CALLBACK, []) + assert _stored_settings(rig.database_url)[list_key] == [] + _listed(rig.proxy, ()) + + +def test_callback_in_both_success_and_failure_lists_is_removed_from_both_in_one_delete(rig: Rig) -> None: + _configure(rig.proxy, {"success_callback": [CALLBACK], "failure_callback": [CALLBACK]}) + _listed(rig.proxy, ((CALLBACK, "failure"), (CALLBACK, "success"))) + + status, body = _delete(rig.proxy, CALLBACK) + + assert status == 200, body + stored: Final = _stored_settings(rig.database_url) + assert (stored["success_callback"], stored["failure_callback"]) == ([], []) + _listed(rig.proxy, ()) + + +def test_delete_keeps_other_callbacks_and_reports_what_is_left_in_the_list_it_edited(rig: Rig) -> None: + _configure(rig.proxy, {"callbacks": [SIBLING, CALLBACK], "failure_callback": [SIBLING]}) + _listed(rig.proxy, ((CALLBACK, "success_and_failure"), (SIBLING, "failure"), (SIBLING, "success_and_failure"))) + + status, body = _delete(rig.proxy, CALLBACK) + + assert status == 200, body + assert body["remaining_callbacks"] == [SIBLING] + stored: Final = _stored_settings(rig.database_url) + assert (stored["callbacks"], stored["failure_callback"]) == ([SIBLING], [SIBLING]) + _listed(rig.proxy, ((SIBLING, "failure"), (SIBLING, "success_and_failure"))) + + +def test_delete_matches_the_callback_name_case_insensitively(rig: Rig) -> None: + _configure(rig.proxy, {"failure_callback": [CALLBACK.upper()], "callbacks": [CALLBACK]}) + _listed(rig.proxy, ((CALLBACK.upper(), "failure"), (CALLBACK, "success_and_failure"))) + + status, body = _delete(rig.proxy, CALLBACK.upper()) + + assert status == 200, body + assert body["removed_callback"] == CALLBACK + _listed(rig.proxy, ()) + + +def test_delete_of_an_unconfigured_callback_is_a_404_and_changes_nothing(rig: Rig) -> None: + _configure(rig.proxy, {"callbacks": [SIBLING]}) + _listed(rig.proxy, ((SIBLING, "success_and_failure"),)) + before: Final = _stored_settings(rig.database_url) + + status, body = _delete(rig.proxy, CALLBACK) + + assert (status, body) == (404, {"detail": {"error": NOT_FOUND}}) + assert _stored_settings(rig.database_url) == before + _listed(rig.proxy, ((SIBLING, "success_and_failure"),)) + + +@pytest.mark.parametrize( + ("file_key", "database_key"), (("failure_callback", "success_callback"), ("success_callback", "failure_callback")) +) +def test_delete_removes_the_database_copy_and_keeps_the_copy_the_config_file_sets( + gateway: Gateway, tmp_path: Path, file_key: str, database_key: str +) -> None: + with _rig(gateway, tmp_path, {file_key: [CALLBACK]}) as file_owned: + _configure(file_owned.proxy, {database_key: [CALLBACK]}) + _listed(file_owned.proxy, ((CALLBACK, "failure"), (CALLBACK, "success"))) + + status, body = _delete(file_owned.proxy, CALLBACK) + + assert status == 200, body + assert body["remaining_callbacks"] == [] + assert _stored_settings(file_owned.database_url)[database_key] == [] + _listed(file_owned.proxy, ((CALLBACK, LIST_TYPES[file_key]),)) + + +def test_delete_of_a_callback_only_the_config_file_lists_is_refused_and_names_the_file( + gateway: Gateway, tmp_path: Path +) -> None: + with _rig(gateway, tmp_path, {"callbacks": [CALLBACK]}) as file_owned: + _listed(file_owned.proxy, ((CALLBACK, "success_and_failure"),)) + + status, body = _delete(file_owned.proxy, CALLBACK) + + assert status == 400, body + detail: Final = object_value(body["detail"]) + assert (detail["keys"], detail["section"]) == (["callbacks"], "litellm_settings") + assert "set in the config file" in str(detail["error"]) + _listed(file_owned.proxy, ((CALLBACK, "success_and_failure"),)) + + +def _completion(request: Request) -> Reply: + message: Final = object_value(list_value(JSON_OBJECT.validate_json(request.body)["messages"])[0]) + if str(message["content"]).endswith("-failure"): + return Reply( + status=400, + body=json.dumps( + {"error": {"type": "invalid_request_error", "code": "synthetic", "message": "synthetic failure"}} + ).encode(), + ) + return Reply( + body=json.dumps( + { + "id": "chatcmpl-" + uuid.uuid4().hex, + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 3, "completion_tokens": 1, "total_tokens": 4}, + } + ).encode() + ) + + +@dataclass(slots=True) +class Sink: + wire: Wire + bodies: list[bytes] # mutable-ok: drain() consumes the queue, later reads keep earlier batches + + def received(self, markers: tuple[str, ...]) -> bool: + self.bodies.extend(request.body for request in self.wire.drain()) + return all(any(marker.encode() in body for body in self.bodies) for marker in markers) + + +@dataclass(frozen=True, slots=True) +class Pair: + first: Gateway + peer: Gateway + database_url: str + provider: Wire + generic: Sink + langsmith: Sink + + def running(self, name: str) -> tuple[bool, bool]: + return (_runs(self.first, name), _runs(self.peer, name)) + + +def _runs(proxy: Gateway, name: str) -> bool: + active: Final = list_value(proxy.get("/active/callbacks")["all_litellm_callbacks"]) + return any(name in str(entry).lower() for entry in active) + + +@contextmanager +def _pair(gateway: Gateway, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Generator[Pair]: + with ( + scratch_database() as database_url, + wire_server(_completion) as provider, + wire_server(lambda _: Reply()) as generic, + wire_server(lambda _: Reply()) as langsmith, + ): + monkeypatch.setenv("DATABASE_URL", database_url) + environment: Final = { + "DATABASE_URL": database_url, + "DEFAULT_FLUSH_INTERVAL_SECONDS": "1", + "LITELLM_STORE_AUDIT_LOGS": "True", + } + removed: Final = ( + "DATABASE_URL_READ_REPLICA", + "GENERIC_LOGGER_ENDPOINT", + "GENERIC_LOGGER_HEADERS", + "LANGSMITH_API_KEY", + "LANGSMITH_BASE_URL", + ) + with ( + owned_proxy(gateway, tmp_path, environment, remove_environment=removed) as first, + owned_proxy(gateway, tmp_path, environment, remove_environment=removed) as peer, + ): + yield Pair(first, peer, database_url, provider, Sink(generic, []), Sink(langsmith, [])) + + +def _send_on_both(pair: Pair, model: str, key: str, marker: str, outcome: str) -> tuple[str, str]: + markers: Final = (f"{marker}-first-{outcome}", f"{marker}-peer-{outcome}") + for proxy, text in zip((pair.first, pair.peer), markers, strict=True): + response = proxy.request( + "POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": text}]}, key=key + ) + assert response.status_code == (400 if outcome == "failure" else 200), response.text + return markers + + +@pytest.mark.parametrize(("list_key", "outcome"), (("failure_callback", "failure"), ("callbacks", "success"))) +def test_deleted_callback_stops_receiving_events_on_both_processes( + gateway: Gateway, tmp_path: Path, monkeypatch: pytest.MonkeyPatch, list_key: str, outcome: str +) -> None: + assert os.environ.get("LITELLM_LICENSE"), "LITELLM_LICENSE must be set: audit logs are an enterprise feature" + marker: Final = "callback-delete-" + uuid.uuid4().hex + with _pair(gateway, tmp_path, monkeypatch) as pair, pair.first.scenario() as scenario: + _configure( + pair.first, + {list_key: [SIBLING, CALLBACK]}, + { + "GENERIC_LOGGER_ENDPOINT": pair.generic.wire.url, + "LANGSMITH_API_KEY": "langsmith-key-" + marker, + "LANGSMITH_BASE_URL": pair.langsmith.wire.url, + }, + ) + eventually(lambda: pair.running(SIBLING), lambda running: running == (True, True), seconds=WORKER_SYNC_SECONDS) + model: Final = scenario.model(api_base=pair.provider.url + "/v1") + key: Final = scenario.key(models=[model]) + before: Final = _send_on_both(pair, model, key, marker + "-before", outcome) + eventually(lambda: pair.generic.received(before), bool, seconds=25) + eventually(lambda: pair.langsmith.received(before), bool, seconds=25) + + status, body = _delete(pair.first, SIBLING) + + assert status == 200, body + assert (body["removed_callback"], body["remaining_callbacks"]) == (SIBLING, [CALLBACK]) + assert _stored_settings(pair.database_url)[list_key] == [CALLBACK] + audit: Final = eventually( + lambda: read_rows( + 'SELECT before_value, updated_values FROM "LiteLLM_AuditLog" WHERE object_id = %s AND action = %s', + ("litellm_settings", "deleted"), + database_url=pair.database_url, + ), + bool, + seconds=10, + ) + assert len(audit) == 1, audit + before_value: Final = list_value(object_value(audit[0]["before_value"])[list_key]) + assert sorted(str(name) for name in before_value) == sorted([SIBLING, CALLBACK]), audit + assert audit[0]["updated_values"] == {list_key: [CALLBACK]}, audit + eventually( + lambda: pair.running(SIBLING), lambda running: running == (False, False), seconds=WORKER_SYNC_SECONDS + ) + + after: Final = _send_on_both(pair, model, key, marker + "-after", outcome) + eventually(lambda: pair.generic.received(after), bool, seconds=25) + _configure(pair.first, {list_key: [CALLBACK, SIBLING]}) + eventually(lambda: pair.running(SIBLING), lambda running: running == (True, True), seconds=WORKER_SYNC_SECONDS) + barrier: Final = _send_on_both(pair, model, key, marker + "-barrier", outcome) + eventually(lambda: pair.langsmith.received(barrier), bool, seconds=25) + + leaked: Final = tuple(text for text in after if any(text.encode() in body for body in pair.langsmith.bodies)) + assert leaked == (), leaked diff --git a/tests/unit/proxy/proxy_server/test_routes_config.py b/tests/unit/proxy/proxy_server/test_routes_config.py index 4c8265b82c3..2d86ff6a7c2 100644 --- a/tests/unit/proxy/proxy_server/test_routes_config.py +++ b/tests/unit/proxy/proxy_server/test_routes_config.py @@ -14,10 +14,15 @@ Routes covered: from __future__ import annotations import asyncio +import contextlib import json +from collections.abc import Callable +from typing import Final from unittest.mock import AsyncMock, MagicMock +import httpx import pytest +from fastapi.testclient import TestClient from .conftest import VOLATILE_KEYS, normalize @@ -944,6 +949,228 @@ def test_config_callback_delete_not_found(client, auth_as, mock_prisma, monkeypa assert "langfuse" in str(response.json()).lower() or "not found" in str(response.json()).lower() +def _delete_callback_roundtrip( + client: TestClient, + auth_as: Callable[..., contextlib.AbstractContextManager[None]], + mock_prisma: MagicMock, + monkeypatch: pytest.MonkeyPatch, + litellm_settings: dict[str, object], + callback_name: str, + owned_by_config: frozenset[str] = frozenset(), +) -> tuple[httpx.Response, dict[str, object]]: + """POST /config/callback/delete as admin against a stubbed config, and + return (response, saved litellm_settings) so each case pins the full body + plus the exact config persisted.""" + from litellm.proxy import proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + _install_litellm_config(mock_prisma) + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + monkeypatch.setattr(ps, "store_model_in_db", True) + + fake_proxy_config: Final = MagicMock() + fake_proxy_config.get_config = AsyncMock(return_value={"litellm_settings": litellm_settings}) + fake_proxy_config.save_config = AsyncMock() + fake_proxy_config.add_deployment = AsyncMock() + fake_proxy_config.litellm_settings.owned_by_config = owned_by_config.__contains__ + monkeypatch.setattr(ps, "proxy_config", fake_proxy_config) + + with auth_as(LitellmUserRoles.PROXY_ADMIN): + response: Final = client.post("/config/callback/delete", json={"callback_name": callback_name}) + saved: Final = fake_proxy_config.save_config.await_args.kwargs["new_config"]["litellm_settings"] + return response, saved + + +def test_config_callback_delete_from_failure_callback( + client: TestClient, + auth_as: Callable[..., contextlib.AbstractContextManager[None]], + mock_prisma: MagicMock, + monkeypatch: pytest.MonkeyPatch, +): + """A callback configured only under litellm_settings.failure_callback is + listed by GET /get/config/callbacks, so the delete route must remove it.""" + response, saved = _delete_callback_roundtrip( + client, + auth_as, + mock_prisma, + monkeypatch, + {"success_callback": ["slack"], "failure_callback": ["datadog", "sentry"]}, + "datadog", + ) + assert response.status_code == 200 + assert response.json()["removed_callback"] == "datadog" + assert saved == {"success_callback": ["slack"], "failure_callback": ["sentry"]} + assert response.json()["remaining_callbacks"] == ["sentry"] + + +def test_config_callback_delete_from_callbacks_list( + client: TestClient, + auth_as: Callable[..., contextlib.AbstractContextManager[None]], + mock_prisma: MagicMock, + monkeypatch: pytest.MonkeyPatch, +): + """A callback configured only under litellm_settings.callbacks is deleted + while unrelated lists pass through untouched.""" + response, saved = _delete_callback_roundtrip( + client, + auth_as, + mock_prisma, + monkeypatch, + {"success_callback": ["langfuse"], "callbacks": ["datadog"], "drop_params": True}, + "datadog", + ) + assert response.status_code == 200 + assert response.json()["removed_callback"] == "datadog" + assert saved == {"success_callback": ["langfuse"], "callbacks": [], "drop_params": True} + assert response.json()["remaining_callbacks"] == [] + + +def test_config_callback_delete_from_multiple_lists( + client: TestClient, + auth_as: Callable[..., contextlib.AbstractContextManager[None]], + mock_prisma: MagicMock, + monkeypatch: pytest.MonkeyPatch, +): + """A name present in both success_callback and callbacks is removed from + both, and remaining_callbacks reports the leftovers of the lists it left.""" + response, saved = _delete_callback_roundtrip( + client, + auth_as, + mock_prisma, + monkeypatch, + { + "success_callback": ["datadog", "slack"], + "failure_callback": ["sentry"], + "callbacks": ["langfuse", "datadog"], + }, + "datadog", + ) + assert response.status_code == 200 + assert saved == { + "success_callback": ["slack"], + "failure_callback": ["sentry"], + "callbacks": ["langfuse"], + } + assert response.json()["remaining_callbacks"] == ["slack", "langfuse"] + + +def test_config_callback_delete_skips_lists_the_config_file_owns( + client: TestClient, + auth_as: Callable[..., contextlib.AbstractContextManager[None]], + mock_prisma: MagicMock, + monkeypatch: pytest.MonkeyPatch, +): + """A name in a database-owned list and a config-file-owned list is removed + from the database-owned one only, so the save is not refused.""" + response, saved = _delete_callback_roundtrip( + client, + auth_as, + mock_prisma, + monkeypatch, + {"success_callback": ["datadog", "slack"], "failure_callback": ["datadog"]}, + "datadog", + owned_by_config=frozenset({"failure_callback"}), + ) + assert response.status_code == 200 + assert saved == {"success_callback": ["slack"], "failure_callback": ["datadog"]} + assert response.json()["remaining_callbacks"] == ["slack"] + + +def test_config_callback_delete_removes_database_copy_when_config_file_owns_success_callback( + client: TestClient, + auth_as: Callable[..., contextlib.AbstractContextManager[None]], + mock_prisma: MagicMock, + monkeypatch: pytest.MonkeyPatch, +): + """A name the config file sets in success_callback is left there while its + database-owned failure_callback copy is removed.""" + response, saved = _delete_callback_roundtrip( + client, + auth_as, + mock_prisma, + monkeypatch, + {"success_callback": ["datadog"], "failure_callback": ["datadog", "sentry"]}, + "datadog", + owned_by_config=frozenset({"success_callback"}), + ) + assert response.status_code == 200 + assert saved == {"success_callback": ["datadog"], "failure_callback": ["sentry"]} + assert response.json()["remaining_callbacks"] == ["sentry"] + + +def test_config_callback_delete_from_string_valued_list( + client: TestClient, + auth_as: Callable[..., contextlib.AbstractContextManager[None]], + mock_prisma: MagicMock, + monkeypatch: pytest.MonkeyPatch, +): + """A list written as a single string is matched and saved back as an empty list.""" + response, saved = _delete_callback_roundtrip( + client, + auth_as, + mock_prisma, + monkeypatch, + {"failure_callback": "datadog"}, + "datadog", + ) + assert response.status_code == 200 + assert saved == {"failure_callback": []} + + +def test_config_callback_delete_case_insensitive( + client: TestClient, + auth_as: Callable[..., contextlib.AbstractContextManager[None]], + mock_prisma: MagicMock, + monkeypatch: pytest.MonkeyPatch, +): + """A mixed-case entry is deleted by the name GET /get/config/callbacks shows for it.""" + response, saved = _delete_callback_roundtrip( + client, + auth_as, + mock_prisma, + monkeypatch, + {"callbacks": ["DataDog", "otel"]}, + "DataDog", + ) + assert response.status_code == 200 + assert response.json()["removed_callback"] == "datadog" + assert saved == {"callbacks": ["otel"]} + + +def test_config_callback_delete_saves_against_the_baseline_it_read( + client: TestClient, + auth_as: Callable[..., contextlib.AbstractContextManager[None]], + mock_prisma: MagicMock, + monkeypatch: pytest.MonkeyPatch, +): + """save_config gets the object get_config returned, so it diffs against the + settings as they were read and writes only the edited callback list.""" + from litellm.proxy import proxy_server as ps + from litellm.proxy._types import LitellmUserRoles + + _install_litellm_config(mock_prisma) + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + monkeypatch.setattr(ps, "store_model_in_db", True) + + read_time: Final = {"litellm_settings": {"failure_callback": ["datadog", "sentry"], "max_budget": 10}} + loaded: Final = ps._ConfigWithBaseline(read_time) + fake_proxy_config: Final = MagicMock() + fake_proxy_config.get_config = AsyncMock(return_value=loaded) + fake_proxy_config.save_config = AsyncMock() + fake_proxy_config.add_deployment = AsyncMock() + fake_proxy_config.litellm_settings.owned_by_config = frozenset[str]().__contains__ + monkeypatch.setattr(ps, "proxy_config", fake_proxy_config) + + with auth_as(LitellmUserRoles.PROXY_ADMIN): + response: Final = client.post("/config/callback/delete", json={"callback_name": "datadog"}) + + assert response.status_code == 200 + saved: Final = fake_proxy_config.save_config.await_args.kwargs["new_config"] + assert saved is loaded + assert saved.baseline == read_time + assert saved["litellm_settings"] == {"failure_callback": ["sentry"], "max_budget": 10} + + # --------------------------------------------------------------------------- # GET /get/config/callbacks # --------------------------------------------------------------------------- From e4dc696af85d2034ffd5824460c596f6dd9b2fc1 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Fri, 16 Oct 2026 09:33:49 -0700 Subject: [PATCH 23/23] feat(ui): turn the Chat UI on by default (#46918) Co-authored-by: Krrish Dholakia Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy_setting_endpoints.py | 2 +- .../test_proxy_setting_endpoints.py | 18 ++++++++++++++++++ 2 files changed, 19 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index b327129793d..4ece88f757d 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -343,7 +343,7 @@ class UISettings(LiteLLMBaseModel): ) enable_chat_ui: bool = Field( - default=False, + default=True, description="If true, shows the Chat page in the UI sidebar, letting users chat with an LLM and connect their own MCP server credentials via OAuth.", ) diff --git a/tests/unit/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py b/tests/unit/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py index 6705e5ac8a5..da799a307a6 100644 --- a/tests/unit/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py +++ b/tests/unit/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py @@ -3771,6 +3771,24 @@ class TestPtuCostAttributionUISetting: assert not mock_prisma.db.litellm_uisettings.upsert.called +class TestEnableChatUISetting: + def test_chat_ui_is_on_when_nothing_is_stored(self, mock_auth, monkeypatch): + TestPtuCostAttributionUISetting._mock_prisma(monkeypatch) + + response = client.get("/get/ui_settings") + + assert response.status_code == 200 + assert response.json()["values"]["enable_chat_ui"] is True + + def test_admin_can_still_turn_chat_ui_off(self, mock_auth, monkeypatch): + TestPtuCostAttributionUISetting._mock_prisma(monkeypatch, stored={"enable_chat_ui": False}) + + response = client.get("/get/ui_settings") + + assert response.status_code == 200 + assert response.json()["values"]["enable_chat_ui"] is False + + class TestApplyUserBudgetToTeamKeysUISetting: """``apply_user_budget_to_team_keys`` mirrors general_settings on every GET.