From 045fb7ba84803e3265462c7988ecfaab35351860 Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Tue, 20 Oct 2026 23:18:42 -0500 Subject: [PATCH 01/14] feat(vertex-ai)!: support gemini 3.1 flash tts (#31915) * feat(vertex-ai): support gemini 3.1 flash tts * fix(vertex-ai): keep TTS metadata compatible with the cost map guard * fix(vertex-ai): preserve existing Gemini TTS models and voice mappings * fix(vertex-ai): recognize provider-stripped TTS model names * test(vertex-ai): match current speech HTTP error response * fix(vertex-ai): preserve nested speech locale and voice * fix(vertex-ai): use provider speech media type in proxy * fix(gemini): preserve TTS dispatch and nested speech settings * test(gemini): run TTS regressions in CI shards * fix(vertex-ai): account for Gemini Cloud TTS audio cost * fix(vertex-ai): bill Gemini TTS without optional audio decoder * test(gemini): cover TTS mapping and billing edge cases * test(vertex-ai): cover Gemini voice normalization in provider shard * fix(vertex-ai): keep Gemini TTS routing with generic provider metadata * chore(vertex-ai): remove obsolete lint suppressions * fix(gemini): strip audio format from speech configuration * chore(vertex-ai): remove obsolete mutation suppressions * fix(vertex-ai): preserve FLAC billing and wrapped audio headers * fix(vertex-ai): reject unmeasurable FLAC before billing * fix(vertex-ai): reconcile speech support with current provider mappings * fix(vertex-ai): validate speech mappings and usage types * fix(vertex-ai): preserve speech mapper override forwarding * fix(gemini): map voice aliases across speech request forms * fix(vertex-ai): send only formats the speech bridge cannot serve to Cloud TTS Gemini TTS detection now goes by model name, so newer Vertex Gemini TTS models keep working, and only mp3, opus, alaw, mulaw and other container formats move to Cloud Text-to-Speech. wav and pcm requests stay on generateContent as on main. A Cloud TTS response whose duration cannot be measured now returns the audio unbilled with a warning instead of raising after a paid call. FLAC handling is removed because Cloud TTS has no FLAC encoding. The cost map, the google_genai mappers and the proxy error path go back to main, since those changes were outside this feature. * fix(vertex_ai): keep Gemini TTS formats, region and instructions on routes that serve them Only mp3, opus, alaw and mulaw go to Cloud Text-to-Speech. Every other format gets a local 400 that lists the supported ones, before Google is called. Cloud TTS uses the deployment's us or eu regional host. Other regions and a custom api_base get a 400 that says to request wav or pcm. instructions is sent as the Cloud TTS style prompt, and a request with no voice keeps the Gemini model. get_usage_object is back to main's code, so dict usages still go through Usage.__init__, and the proxy no longer labels pcm16 or linear16 output as WAV. * fix(vertex_ai): send Gemini TTS to every regional Cloud TTS host and keep speak tags as text Cloud Text-to-Speech serves Gemini TTS on a host named after each Vertex location, so a non-global vertex_location now uses {location}-texttospeech.googleapis.com instead of only the us and eu hosts. Google answers for regions it does not serve. Gemini TTS voices reject SSML, so input that contains is sent as text for Gemini TTS models. Chirp keeps the SSML auto-detect. Removes the proxy alaw and mulaw media-type entries. Cloud TTS returns those encodings with a WAV header, so the response is labeled audio/wav and the entries never applied. Adds a regression test that custom per-token pricing bills the usage a speech response reports. * fix(vertex_ai): map Gemini TTS regions to Google's Cloud TTS locations and keep a bad voice from cooling down the deployment US and EU Vertex regions go to the us and eu Cloud TTS hosts, northamerica-northeast1 to its own, global or unset to global, and any other region gets a local 400. A Cloud TTS 404 on the Gemini path, such as a misspelled voice, comes back as a 400 so the router keeps the deployment serving wav. Usage reads the sent request from the response instead of the logging object. * fix(vertex_ai): default Gemini TTS Cloud formats to a Gemini voice and validate speakers Parse only usage fields from the sent Cloud TTS request, so a numeric speakingRate no longer fails after Google bills it. Send the Gemini voice Kore when no voice is named, and reject Cloud TTS voice names unless drop_params is set. Reject malformed speaker entries unless drop_params is set, and translate Cloud TTS speaker fields on wav and chat audio. Use the Vertex default region when no location is set, so mp3 and wav share a region. * test(vertex_ai): cover two-speaker Gemini TTS mp3 through aspeech with fake credentials Replaces two tests that patched the Vertex TTS config's own token and URL methods. * fix(vertex_ai): reject malformed Gemini TTS voices and keep beta Chirp and Lyria params unmapped * fix(vertex_ai): drop unknown Gemini TTS voice keys on wav and pcm instead of rejecting them * fix(vertex_ai): reject unknown Gemini TTS voice keys on wav, pcm, and chat audio unless drop_params is set * fix(vertex_ai): honor every drop_params setting in the Gemini TTS voice key check Global drop_params, batch row and deployment drop_params on Vertex batch file uploads now drop unknown voice keys. A voice object with both speechConfig and speech_config gets a 400 for the ignored speech_config. * fix(gemini): drop unknown nested voice fields under drop_params and detect TTS by served model * fix(gemini): filter unknown speech config fields with fixed-depth helpers and pass audio params through unchanged * fix(gemini): camel-case snake replicated voice keys so drop_params keeps them * fix(vertex-ai): reject voice fields Cloud TTS cannot send unless drop_params is set * fix(vertex-ai): reject wrong-typed Cloud TTS voice values and check speakers as sent Gemini TTS on mp3, opus, alaw, or mulaw now answers 400 for a voice field whose value Cloud TTS cannot send, such as a list languageCode or a string voiceConfig, and for nested voice fields under a speaker that also names speakerId. drop_params still drops them. * fix(vertex-ai): reject wrong-typed speech config wrappers and speakerId overriding a prebuilt voice on Cloud TTS * fix(vertex-ai): keep speakers and voice_config beside a speech_config wrapper on Cloud TTS * fix(vertex-ai): pick Cloud TTS voices from the same speech config the WAV path sends to Gemini * fix(files): carry the deployment's drop_params on model= file uploads so Vertex batch rows drop unknown voice keys * fix(gemini): keep a structured speech voice when the model is a model_alias_map alias for a Gemini TTS model --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/cost_calculator.py | 6 +- .../speech_to_completion_bridge/handler.py | 1 + .../transformation.py | 67 +- .../llms/vertex_ai/files/transformation.py | 8 +- .../vertex_and_google_ai_studio_gemini.py | 334 ++++- .../text_to_speech/transformation.py | 567 +++++++- litellm/main.py | 7 +- .../openai_files_endpoints/common_utils.py | 12 + .../openai_files_endpoints/files_endpoints.py | 3 +- litellm/types/llms/vertex_ai.py | 10 + .../types/llms/vertex_ai_text_to_speech.py | 17 +- litellm/utils.py | 5 + tests/audio_tests/gemini_tts_speech.mp3 | Bin 0 -> 4992 bytes tests/audio_tests/gemini_tts_speech.ogg | Bin 0 -> 4179 bytes .../test_transformation.py | 121 +- tests/unit/llms/gemini/test_gemini_tts.py | 262 +++- .../files/test_vertex_ai_files_streaming.py | 40 + ...test_vertex_and_google_ai_studio_gemini.py | 239 ++++ .../text_to_speech/test_transformation.py | 1233 ++++++++++++++++- .../test_files_endpoint.py | 43 + tests/unit/test_cost_calculator.py | 29 +- 21 files changed, 2878 insertions(+), 126 deletions(-) create mode 100644 tests/audio_tests/gemini_tts_speech.mp3 create mode 100644 tests/audio_tests/gemini_tts_speech.ogg diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index f952883bc87..550bbfcde98 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -1673,8 +1673,10 @@ def completion_cost( verbose_logger.debug("selected model name for cost calculation: %s", model) if completion_response is not None and ( - isinstance(completion_response, BaseModel) or isinstance(completion_response, dict) - ): # tts returns a custom class + isinstance(completion_response, (BaseModel, dict)) + or isinstance(completion_response, HttpxBinaryResponseContent) + and completion_response.usage is not None + ): if isinstance(completion_response, dict): usage_obj: dict | Usage | None = completion_response.get("usage", {}) else: diff --git a/litellm/endpoints/speech/speech_to_completion_bridge/handler.py b/litellm/endpoints/speech/speech_to_completion_bridge/handler.py index 00d112773f1..faf2daecf5f 100644 --- a/litellm/endpoints/speech/speech_to_completion_bridge/handler.py +++ b/litellm/endpoints/speech/speech_to_completion_bridge/handler.py @@ -122,6 +122,7 @@ class SpeechToCompletionBridgeHandler: if isinstance(result, ModelResponse): return self.transformation_handler.transform_response( model_response=result, + model=model, response_format=requested_response_format if isinstance(requested_response_format, str) else None, ) else: diff --git a/litellm/endpoints/speech/speech_to_completion_bridge/transformation.py b/litellm/endpoints/speech/speech_to_completion_bridge/transformation.py index 9c1e5ae7a29..ad6ed269adf 100644 --- a/litellm/endpoints/speech/speech_to_completion_bridge/transformation.py +++ b/litellm/endpoints/speech/speech_to_completion_bridge/transformation.py @@ -5,6 +5,7 @@ from typing import TYPE_CHECKING, Final, cast from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm.constants import OPENAI_CHAT_COMPLETION_PARAMS +from litellm.litellm_core_utils.core_helpers import drops_unsupported_params from litellm.litellm_core_utils.hidden_params import get_hidden_params, served_from_cache, set_hidden_param if TYPE_CHECKING: @@ -38,6 +39,11 @@ class ChatAudioParam(TypedDict): format: ReadOnly[NotRequired[str]] +class ChatSpeechConfigAudioParam(TypedDict): + speech_config: ReadOnly[Mapping[str, object]] + format: ReadOnly[str] + + class SpeechToCompletionBridgeTransformationHandler: def _validate_response_format( self, model: str, custom_llm_provider: str, optional_params: Mapping[str, object] @@ -49,6 +55,20 @@ class SpeechToCompletionBridgeTransformationHandler: return from litellm.exceptions import BadRequestError + if custom_llm_provider.startswith("vertex_ai"): + from litellm.llms.vertex_ai.text_to_speech.transformation import VertexAITextToSpeechConfig + + vertex_supported: Final = ", ".join( + sorted({*GEMINI_TTS_SUPPORTED_RESPONSE_FORMATS, *VertexAITextToSpeechConfig.GEMINI_FORMAT_MAPPINGS}) + ) + raise BadRequestError( + message=( + f"Vertex AI Gemini TTS does not support response_format='{response_format}'." + f" Supported response formats: {vertex_supported}." + ), + model=model, + llm_provider=custom_llm_provider, + ) supported: Final = ", ".join(sorted(GEMINI_TTS_SUPPORTED_RESPONSE_FORMATS)) raise BadRequestError( message=( @@ -75,8 +95,27 @@ class SpeechToCompletionBridgeTransformationHandler: return response_format if isinstance(response_format, str) else None def _chat_audio_param( - self, model: str, voice: str | Mapping[str, object] | None, optional_params: Mapping[str, object] - ) -> ChatAudioParam | None: + self, + model: str, + voice: str | Mapping[str, object] | None, + optional_params: Mapping[str, object], + drop_params: bool, + custom_llm_provider: str, + ) -> ChatAudioParam | ChatSpeechConfigAudioParam | None: + import litellm + + if isinstance(voice, Mapping) and self._is_gemini_tts_model(litellm.model_alias_map.get(model, model)): + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + normalize_gemini_speech_config, + reject_unknown_gemini_tts_voice_keys, + ) + + reject_unknown_gemini_tts_voice_keys(model, voice, drop_params, custom_llm_provider) + speech_audio: Final[ChatSpeechConfigAudioParam] = { + "speech_config": normalize_gemini_speech_config(voice, drops_unsupported_params(drop_params)), + "format": GEMINI_TTS_CHAT_AUDIO_FORMAT, + } + return speech_audio if not isinstance(voice, str): return None audio_format: Final = self._chat_audio_format(model, optional_params) @@ -90,9 +129,9 @@ class SpeechToCompletionBridgeTransformationHandler: self, model: str, input: str, - voice: str | dict | None, - optional_params: dict, - litellm_params: dict, + voice: str | Mapping[str, object] | None, + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], headers: dict, litellm_logging_obj: "LiteLLMLoggingObj", custom_llm_provider: str, @@ -104,7 +143,13 @@ class SpeechToCompletionBridgeTransformationHandler: "messages": [user_message], "modalities": ["audio"], **self._chat_completion_params(optional_params), - "audio": self._chat_audio_param(model, voice, optional_params), + "audio": self._chat_audio_param( + model, + voice, + optional_params, + drops_unsupported_params(litellm_params.get("drop_params")), + custom_llm_provider, + ), **litellm_params, "headers": headers, "litellm_logging_obj": litellm_logging_obj, @@ -153,8 +198,9 @@ class SpeechToCompletionBridgeTransformationHandler: return wav_header + pcm_data def _is_gemini_tts_model(self, model: str) -> bool: - """Check if the model is a Gemini TTS model that returns PCM16 data.""" - return "gemini" in model.lower() and ("tts" in model.lower() or "preview-tts" in model.lower()) + from litellm.utils import is_gemini_tts_model + + return is_gemini_tts_model(model) def _gemini_tts_response_body(self, decoded_audio: bytes, response_format: str | None) -> tuple[bytes, str]: if response_format == GEMINI_TTS_RAW_RESPONSE_FORMAT: @@ -162,7 +208,7 @@ class SpeechToCompletionBridgeTransformationHandler: return self._convert_pcm16_to_wav(decoded_audio), "audio/wav" def transform_response( - self, model_response: "ModelResponse", response_format: str | None + self, model_response: "ModelResponse", model: str, response_format: str | None ) -> "HttpxBinaryResponseContent": import base64 @@ -176,10 +222,9 @@ class SpeechToCompletionBridgeTransformationHandler: raise ValueError("No audio part found in the response") decoded_audio: Final = base64.b64decode(audio_part.data) - model: Final = getattr(model_response, "model", "") content, content_type = ( self._gemini_tts_response_body(decoded_audio, response_format) - if self._is_gemini_tts_model(model) + if self._is_gemini_tts_model(model) or self._is_gemini_tts_model(model_response.model or "") else (decoded_audio, "audio/mpeg") ) response: Final = httpx.Response( diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index 176569f5d8d..d1e3f10fdd4 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -8,6 +8,7 @@ import time from collections.abc import AsyncGenerator, Callable, Iterable, Iterator, Mapping from contextlib import aclosing from dataclasses import dataclass +from functools import partial from types import MappingProxyType from typing import IO, Any, Final, TypedDict from urllib.parse import quote, unquote @@ -30,6 +31,7 @@ from litellm.litellm_core_utils.cloud_storage_security import ( split_configured_cloud_bucket_name, validate_managed_cloud_file_id, ) +from litellm.litellm_core_utils.core_helpers import drops_unsupported_params from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.prompt_templates.common_utils import ( extract_file_data, @@ -1103,6 +1105,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): def _map_openai_to_vertex_params( self, openai_request_body: dict[str, Any], + deployment_drop_params: object = None, ) -> dict[str, object]: """ wrapper to call VertexGeminiConfig.map_openai_params @@ -1121,7 +1124,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): model=model, non_default_params=openai_request_body, optional_params={}, - drop_params=False, + drop_params=drops_unsupported_params(openai_request_body.get("drop_params", deployment_drop_params)), allow_integer_enums=False, ) return vertex_params @@ -1157,11 +1160,12 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): create_file_data=create_file_data, content_type=content_type, ): + deployment_drop_params: Final[object] = litellm_params.get("drop_params") return { "streaming_media_upload": StreamingMediaUploadConfig( body_stream=_OpenAIToVertexBatchUploadStream( file_data, - self._map_openai_to_vertex_params, + partial(self._map_openai_to_vertex_params, deployment_drop_params=deployment_drop_params), ), content_type="application/json", ) 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 241b2d2a3a0..ce632bc0eef 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 @@ -12,7 +12,7 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, Union, cast, get_args import httpx -from pydantic import JsonValue +from pydantic import JsonValue, TypeAdapter import litellm from litellm import verbose_logger @@ -28,7 +28,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.core_helpers import as_str_mapping, drops_unsupported_params, normalize_drop_params from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap from litellm.litellm_core_utils.json_fragment_accumulator import JSONFragmentAccumulator from litellm.litellm_core_utils.prompt_templates.factory import ( @@ -232,6 +232,274 @@ def _served_model_name(model_version: object) -> str | None: return model_version.split("@", 1)[0] +_GEMINI_SPEECH_CONFIG_KEY_MAP: Final[Mapping[str, str]] = MappingProxyType( + { + "multi_speaker_voice_config": "multiSpeakerVoiceConfig", + "speaker_voice_configs": "speakerVoiceConfigs", + "voice_config": "voiceConfig", + "prebuilt_voice_config": "prebuiltVoiceConfig", + "voice_name": "voiceName", + "replicated_voice_config": "replicatedVoiceConfig", + "mime_type": "mimeType", + "voice_sample_audio": "voiceSampleAudio", + "language_code": "languageCode", + "speaker_alias": "speakerAlias", + "speaker_id": "speakerId", + } +) + + +_SPEECH_CONFIG_ITEMS: Final = TypeAdapter(list[object]) + + +def _normalize_gemini_speech_config_item(value: Mapping[str, object]) -> Mapping[str, object]: + normalized: Final[dict[str, object]] = {} # mutable-ok: iterative traversal fills nested output without recursion + + pending: Final[ # mutable-ok: the explicit worklist must grow while avoiding recursive traversal + list[tuple[dict[str, object] | list[object], Mapping[str, object] | list[object]]] + ] = [(normalized, value)] + while pending: + target, source = pending.pop() + items = source.items() if isinstance(source, Mapping) else enumerate(source) + for key, item in items: + if isinstance(source, Mapping): + if not isinstance(key, str): + continue + normalized_key: str | int = _GEMINI_SPEECH_CONFIG_KEY_MAP.get(key, key) + else: + normalized_key = len(target) + + child: object + if (mapping_item := as_str_mapping(item)) is not None: + pending.append((child_mapping := dict[str, object](), mapping_item)) + child = child_mapping + elif isinstance(item, list): + pending.append((child_items := list[object](), _SPEECH_CONFIG_ITEMS.validate_python(item))) + child = child_items + else: + child = ( + OPENAI_TO_GEMINI_TTS_VOICES.get(item.lower(), item) + if normalized_key == "voiceName" and isinstance(item, str) + else item + ) + + if isinstance(target, dict): + assert isinstance(normalized_key, str) + target[normalized_key] = child + else: + target.append(child) + + return normalized + + +def _gemini_speaker_voice_config(item: object) -> object: + speaker_config: Final = as_str_mapping(item) + if speaker_config is None: + return item + speaker: Final = speaker_config.get("speakerAlias") + voice_name: Final = speaker_config.get("speakerId") + if not isinstance(speaker, str) or not isinstance(voice_name, str): + return item + return { + **{key: value for key, value in speaker_config.items() if key not in ("speakerAlias", "speakerId")}, + "speaker": speaker, + "voiceConfig": { + "prebuiltVoiceConfig": {"voiceName": OPENAI_TO_GEMINI_TTS_VOICES.get(voice_name.lower(), voice_name)} + }, + } + + +def _with_gemini_speaker_voice_configs(config: Mapping[str, object]) -> Mapping[str, object]: + multi_speaker_config: Final = as_str_mapping(config.get("multiSpeakerVoiceConfig")) + if multi_speaker_config is None: + return config + speaker_voice_configs: Final = multi_speaker_config.get("speakerVoiceConfigs") + if not isinstance(speaker_voice_configs, list): + return config + return { + **config, + "multiSpeakerVoiceConfig": { + **multi_speaker_config, + "speakerVoiceConfigs": [ + _gemini_speaker_voice_config(item) + for item in _SPEECH_CONFIG_ITEMS.validate_python(speaker_voice_configs) + ], + }, + } + + +_GEMINI_SPEECH_CONFIG_FIELDS: Final = frozenset({"voiceConfig", "multiSpeakerVoiceConfig", "languageCode"}) + +_GEMINI_VOICE_CONFIG_FIELDS: Final = MappingProxyType( + { + "prebuiltVoiceConfig": frozenset({"voiceName"}), + "replicatedVoiceConfig": frozenset({"mimeType", "voiceSampleAudio"}), + "voice": None, + } +) + +_GEMINI_SPEAKER_VOICE_CONFIG_FIELDS: Final = frozenset({"speaker", "voiceConfig"}) + + +def _known_leaf_fields(value: object, fields: frozenset[str] | None) -> object: + mapping: Final = as_str_mapping(value) + if fields is None or mapping is None: + return value + return {key: item for key, item in mapping.items() if key in fields} + + +def _known_voice_config_fields(value: object) -> object: + mapping: Final = as_str_mapping(value) + if mapping is None: + return value + return { + key: _known_leaf_fields(item, _GEMINI_VOICE_CONFIG_FIELDS[key]) + for key, item in mapping.items() + if key in _GEMINI_VOICE_CONFIG_FIELDS + } + + +def _known_speaker_voice_config_fields(value: object) -> object: + mapping: Final = as_str_mapping(value) + if mapping is None: + return value + return { + key: _known_voice_config_fields(item) if key == "voiceConfig" else item + for key, item in mapping.items() + if key in _GEMINI_SPEAKER_VOICE_CONFIG_FIELDS + } + + +def _known_speaker_voice_configs(value: object) -> object: + if not isinstance(value, list): + return value + return [_known_speaker_voice_config_fields(item) for item in _SPEECH_CONFIG_ITEMS.validate_python(value)] + + +def _known_multi_speaker_voice_config_fields(value: object) -> object: + mapping: Final = as_str_mapping(value) + if mapping is None: + return value + return {key: _known_speaker_voice_configs(item) for key, item in mapping.items() if key == "speakerVoiceConfigs"} + + +def _known_speech_config_field(key: str, value: object) -> object: + match key: + case "voiceConfig": + return _known_voice_config_fields(value) + case "multiSpeakerVoiceConfig": + return _known_multi_speaker_voice_config_fields(value) + case _: + return value + + +def _known_gemini_speech_fields( + config: Mapping[str, object], +) -> dict[str, object]: # mutable-ok: provider request serialization requires a concrete dict + return { + key: _known_speech_config_field(key, item) + for key, item in config.items() + if key in _GEMINI_SPEECH_CONFIG_FIELDS + } + + +_GEMINI_TTS_VOICE_KEYS: Final = frozenset( + { + "format", + "name", + "voice", + "voiceName", + "voice_name", + "languageCode", + "language_code", + "modelName", + "model_name", + "speechConfig", + "speech_config", + "voiceConfig", + "voice_config", + "multiSpeakerVoiceConfig", + "multi_speaker_voice_config", + } +) + + +_NAMED_UNKNOWN_VOICE_KEYS: Final = 10 + +_NAMED_VOICE_KEY_CHARS: Final = 100 + + +def _gemini_speech_config_wrapper(value: Mapping[str, object]) -> Mapping[str, object] | None: + camel_config: Final = as_str_mapping(value.get("speechConfig")) + return camel_config if camel_config is not None else as_str_mapping(value.get("speech_config")) + + +def named_voice_keys(keys: Sequence[str]) -> str: + named: Final = ", ".join(key[:_NAMED_VOICE_KEY_CHARS] for key in keys[:_NAMED_UNKNOWN_VOICE_KEYS]) + unnamed_count: Final = len(keys) - _NAMED_UNKNOWN_VOICE_KEYS + return f"{named}, and {unnamed_count} more" if unnamed_count > 0 else named + + +def reject_unknown_gemini_tts_voice_keys( + model: str, voice: Mapping[str, object], drop_params: bool, llm_provider: str | None +) -> None: + nested_config: Final = _gemini_speech_config_wrapper(voice) or {} + ignored_wrapper: Final = ( + ("speech_config",) + if as_str_mapping(voice.get("speechConfig")) is not None and voice.get("speech_config") is not None + else () + ) + unknown_keys: Final = sorted( + {*(key for key in (*voice, *nested_config) if key not in _GEMINI_TTS_VOICE_KEYS), *ignored_wrapper} + ) + if not unknown_keys or drops_unsupported_params(drop_params): + return + raise UnsupportedParamsError( + message=( + f"{model} does not support the voice keys {named_voice_keys(unknown_keys)}. Remove them, or set " + "drop_params to ignore them." + ), + model=model, + llm_provider=llm_provider, + ) + + +def normalize_gemini_speech_config( + value: Mapping[str, object], drop_unknown_fields: bool = False +) -> dict[str, object]: # mutable-ok: provider request serialization requires a concrete dict + speech_config: Final = {**_with_gemini_speaker_voice_configs(mapped_gemini_speech_config(value))} + return _known_gemini_speech_fields(speech_config) if drop_unknown_fields else speech_config + + +def _merge_gemini_speech_config_wrapper(value: Mapping[str, object]) -> Mapping[str, object]: + nested_config: Final = _gemini_speech_config_wrapper(value) + if nested_config is None: + return value + return { + **{key: item for key, item in value.items() if key not in ("speechConfig", "speech_config")}, + **nested_config, + } + + +def mapped_gemini_speech_config( + value: Mapping[str, object], +) -> dict[str, object]: # mutable-ok: provider request serialization requires a concrete dict + normalized: Final = _normalize_gemini_speech_config_item(_merge_gemini_speech_config_wrapper(value)) + speech_config: Final = {key: item for key, item in normalized.items() if key in _GEMINI_SPEECH_CONFIG_FIELDS} + voice_name: Final = next( + (normalized[key] for key in ("name", "voiceName", "voice") if isinstance(normalized.get(key), str)), + None, + ) + if "voiceConfig" in speech_config or "multiSpeakerVoiceConfig" in speech_config or not isinstance(voice_name, str): + return speech_config + return { + **speech_config, + "voiceConfig": { + "prebuiltVoiceConfig": {"voiceName": OPENAI_TO_GEMINI_TTS_VOICES.get(voice_name.lower(), voice_name)} + }, + } + + class VertexAIBaseConfig: def get_mapped_special_auth_params(self) -> dict: """ @@ -1237,7 +1505,9 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): status_code=400, ) - def _map_audio_params(self, value: dict) -> dict: + def _map_audio_params( + self, value: Mapping[str, object], drop_unknown_fields: bool = False + ) -> dict[str, object]: # mutable-ok: the public forwarding interface returns a concrete request dictionary """ Expected input: { @@ -1256,12 +1526,6 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): languageCode: "en-US", } """ - from litellm.types.llms.vertex_ai import ( - PrebuiltVoiceConfig, - SpeechConfig, - VoiceConfig, - ) - # Validate audio format - Gemini TTS only supports pcm16 audio_format: Final = value.get("format") if audio_format is not None and audio_format != "pcm16": @@ -1271,22 +1535,46 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): f"Please set audio format to 'pcm16'." ) - # Map OpenAI audio parameter to Gemini speech config - speech_config: Final[SpeechConfig] = {} + camel_speech_config: Final = value.get("speechConfig") + snake_speech_config: Final = value.get("speech_config") + speech_config_value: Final = ( + camel_speech_config if isinstance(camel_speech_config, Mapping) else snake_speech_config + ) + language_code: Final = value.get("language_code", value.get("languageCode")) + if isinstance(speech_config_value, Mapping): + return normalize_gemini_speech_config(value, drop_unknown_fields) - if "voice" in value: - voice: Final = value["voice"] - voice_name: Final = ( - OPENAI_TO_GEMINI_TTS_VOICES.get(voice.lower(), voice) if isinstance(voice, str) else voice - ) - prebuilt_voice_config: Final[PrebuiltVoiceConfig] = {"voiceName": voice_name} - voice_config: Final[VoiceConfig] = {"prebuiltVoiceConfig": prebuilt_voice_config} - speech_config["voiceConfig"] = voice_config + voice: Final = value.get("voice") + speech_config_without_language: Final[ # mutable-ok: provider request serialization requires a concrete dict + dict[str, object] + ] = ( + { + "voiceConfig": { + "prebuiltVoiceConfig": { + "voiceName": OPENAI_TO_GEMINI_TTS_VOICES.get(voice.lower(), voice) + if isinstance(voice, str) + else voice, + } + } + } + if "voice" in value + else {} + ) - if "language_code" in value: - speech_config["languageCode"] = value["language_code"] + if language_code is not None: + return { + **speech_config_without_language, + "languageCode": language_code, + } - return cast(dict, speech_config) + return speech_config_without_language + + def _map_chat_audio_params( + self, model: str, value: Mapping[str, object], drop_params: bool + ) -> dict[str, object]: # mutable-ok: the public forwarding interface returns a concrete request dictionary + if _gemini_speech_config_wrapper(value) is not None: + reject_unknown_gemini_tts_voice_keys(model, value, drop_params, self.custom_llm_provider) + return self._map_audio_params(value, drop_unknown_fields=drops_unsupported_params(drop_params)) def map_audio_params( self, @@ -1351,7 +1639,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): elif param == "n": optional_params["candidate_count"] = value elif param == "audio" and isinstance(value, dict): - optional_params["speechConfig"] = self._map_audio_params(value) + optional_params["speechConfig"] = self._map_chat_audio_params(model, value, drop_params) elif param == "stop": if isinstance(value, str): optional_params["stop_sequences"] = [value] diff --git a/litellm/llms/vertex_ai/text_to_speech/transformation.py b/litellm/llms/vertex_ai/text_to_speech/transformation.py index 2fae2a9de1c..1441d01e41e 100644 --- a/litellm/llms/vertex_ai/text_to_speech/transformation.py +++ b/litellm/llms/vertex_ai/text_to_speech/transformation.py @@ -6,20 +6,25 @@ Reference: https://cloud.google.com/text-to-speech/docs/reference/rest/v1/text/s """ import base64 -from collections.abc import Coroutine, Mapping +import math +from collections.abc import Coroutine, Iterator, Mapping, Sequence from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, TypeAlias, Union +from typing import TYPE_CHECKING, Any, ClassVar, Final, TypeAlias, Union import httpx -from pydantic import ConfigDict, TypeAdapter +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError import litellm +from litellm._logging import verbose_logger from litellm.exceptions import UnsupportedParamsError from litellm.litellm_core_utils.audio_utils.utils import ( DEFAULT_SPEECH_MEDIA_TYPE, + calculate_request_duration, speech_media_type_from_audio_bytes, ) +from litellm.litellm_core_utils.core_helpers import as_str_mapping, drops_unsupported_params 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.text_to_speech.transformation import ( BaseTextToSpeechConfig, TextToSpeechRequestData, @@ -33,6 +38,7 @@ from litellm.types.llms.vertex_ai import VERTEX_CREDENTIALS_TYPES from litellm.types.llms.vertex_ai_text_to_speech import ( VertexTextToSpeechAudioConfig, VertexTextToSpeechInput, + VertexTextToSpeechSpeakerVoiceConfig, VertexTextToSpeechVoice, ) @@ -49,6 +55,166 @@ _LyriaVoice: TypeAlias = str | dict | None _JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True)) _JSON_ARRAY: Final = TypeAdapter(tuple[object, ...], config=ConfigDict(hide_input_in_errors=True)) _STR: Final = TypeAdapter(str, config=ConfigDict(hide_input_in_errors=True)) +_GEMINI_CLOUD_TTS_LOCATIONS: Final = MappingProxyType( + { + **dict.fromkeys( + ("us", "us-central1", "us-east1", "us-east4", "us-east5", "us-south1", "us-west1", "us-west4"), "us" + ), + **dict.fromkeys( + ("eu", "europe-central2", "europe-north1", "europe-southwest1", "europe-west1", "europe-west4"), "eu" + ), + "northamerica-northeast1": "northamerica-northeast1", + } +) + + +_DEFAULT_GEMINI_TTS_VOICE: Final = "Kore" + + +class _GeminiCloudTtsInput(BaseModel): + model_config = ConfigDict(frozen=True) + + text: str | None = None + ssml: str | None = None + prompt: str | None = None + + +class _GeminiCloudTtsAudioConfig(BaseModel): + model_config = ConfigDict(frozen=True) + + audioEncoding: str = "" + sampleRateHertz: int | None = None + + +class _GeminiCloudTtsRequest(BaseModel): + model_config = ConfigDict(frozen=True) + + input: _GeminiCloudTtsInput + audioConfig: _GeminiCloudTtsAudioConfig + + +def _fallback_gemini_tts_audio_duration(audio: bytes, encoding: str, sample_rate: int) -> float | None: + if encoding == "PCM": + return len(audio) / (2 * sample_rate) if sample_rate > 0 else None + if encoding in ("ALAW", "MULAW") and not (audio[:4] == b"RIFF" and audio[8:12] == b"WAVE"): + return len(audio) / sample_rate if sample_rate > 0 else None + + if encoding in ("LINEAR16", "ALAW", "MULAW") and audio[:4] == b"RIFF" and audio[8:12] == b"WAVE": + format_offset: Final = audio.find(b"fmt ", 12) + data_offset: Final = audio.find(b"data", 12) + if format_offset >= 0 and data_offset >= 0 and format_offset + 20 <= len(audio): + byte_rate: Final = int.from_bytes(audio[format_offset + 16 : format_offset + 20], "little") + data_size: Final = int.from_bytes(audio[data_offset + 4 : data_offset + 8], "little") + if byte_rate > 0 and data_size <= len(audio) - data_offset - 8: + return data_size / byte_rate + + if encoding == "MP3" and speech_media_type_from_audio_bytes(audio) == "audio/mpeg": + return len(audio) / 4000 + + if encoding == "OGG_OPUS" and audio[:4] == b"OggS": + opus_header: Final = audio.find(b"OpusHead") + last_page: Final = audio.rfind(b"OggS") + if opus_header >= 0 and opus_header + 12 <= len(audio) and last_page + 14 <= len(audio): + pre_skip: Final = int.from_bytes(audio[opus_header + 10 : opus_header + 12], "little") + granule: Final = int.from_bytes(audio[last_page + 6 : last_page + 14], "little") + if pre_skip <= granule < 2**64 - 1: + return (granule - pre_skip) / 48000 + + return None + + +_SPEAKER_CONFIG_ITEMS: Final = TypeAdapter(list[object]) + +_CloudTtsVoiceFields: TypeAlias = "Mapping[str, _CloudTtsVoiceFields]" + +_CLOUD_TTS_VOICE_LEAF: Final[_CloudTtsVoiceFields] = MappingProxyType({}) + +_CLOUD_TTS_PREBUILT_VOICE_FIELDS: Final[_CloudTtsVoiceFields] = MappingProxyType( + {"prebuiltVoiceConfig": MappingProxyType({"voiceName": _CLOUD_TTS_VOICE_LEAF})} +) + +_CLOUD_TTS_SPEECH_CONFIG_FIELDS: Final[_CloudTtsVoiceFields] = MappingProxyType( + { + "languageCode": _CLOUD_TTS_VOICE_LEAF, + "voiceConfig": _CLOUD_TTS_PREBUILT_VOICE_FIELDS, + "multiSpeakerVoiceConfig": MappingProxyType( + { + "speakerVoiceConfigs": MappingProxyType( + { + "speaker": _CLOUD_TTS_VOICE_LEAF, + "speakerAlias": _CLOUD_TTS_VOICE_LEAF, + "speakerId": _CLOUD_TTS_VOICE_LEAF, + "voiceConfig": _CLOUD_TTS_PREBUILT_VOICE_FIELDS, + } + ) + } + ), + } +) + + +_CLOUD_TTS_LIST_FIELDS: Final = frozenset({"speakerVoiceConfigs"}) + +_GEMINI_TTS_VOICE_NAME_KEYS: Final = ("name", "voiceName", "voice_name", "voice") + + +def _fields_cloud_tts_cannot_send( + config: Mapping[str, object], fields: _CloudTtsVoiceFields, path: str +) -> Iterator[str]: + for key, value in config.items(): + yield from _field_cloud_tts_cannot_send(f"{path}{key}", value, fields.get(key), key in _CLOUD_TTS_LIST_FIELDS) + + +def _field_cloud_tts_cannot_send( + path: str, value: object, fields: "_CloudTtsVoiceFields | None", is_list_field: bool +) -> Iterator[str]: + if fields is None: + yield path + return + if value is None: + return + if not fields: + yield from () if isinstance(value, str) else (path,) + return + items: Final = ( + _SPEAKER_CONFIG_ITEMS.validate_python(value) if is_list_field and isinstance(value, list) else (value,) + ) + for item in items: + mapping = as_str_mapping(item) + yield from (path,) if mapping is None else _fields_cloud_tts_cannot_send(mapping, fields, f"{path}.") + + +def _non_string_voice_names(config: Mapping[str, object], path: str) -> Iterator[str]: + return ( + f"{path}{key}" + for key in _GEMINI_TTS_VOICE_NAME_KEYS + if config.get(key) is not None and not isinstance(config.get(key), str) + ) + + +def _non_mapping_speech_config_wrappers(voice: Mapping[str, object]) -> Iterator[str]: + return ( + key + for key in ("speechConfig", "speech_config") + if voice.get(key) is not None and as_str_mapping(voice.get(key)) is None + ) + + +def _speaker_entries(config: Mapping[str, object]) -> Iterator[Mapping[str, object]]: + multi_speaker_config: Final = as_str_mapping(config.get("multiSpeakerVoiceConfig")) + speakers: Final = multi_speaker_config.get("speakerVoiceConfigs") if multi_speaker_config is not None else None + items: Final = _SPEAKER_CONFIG_ITEMS.validate_python(speakers) if isinstance(speakers, list) else () + return (entry for item in items if (entry := as_str_mapping(item)) is not None) + + +def _prebuilt_voices_replaced_by_speaker_id(config: Mapping[str, object]) -> Iterator[str]: + return ( + "multiSpeakerVoiceConfig.speakerVoiceConfigs.voiceConfig.prebuiltVoiceConfig" + for entry in _speaker_entries(config) + if entry.get("speakerId") is not None + and (voice_config := as_str_mapping(entry.get("voiceConfig"))) is not None + and voice_config.get("prebuiltVoiceConfig") is not None + ) class VertexAITextToSpeechConfig(BaseTextToSpeechConfig, VertexBase): @@ -81,14 +247,19 @@ class VertexAITextToSpeechConfig(BaseTextToSpeechConfig, VertexBase): } # Response format mappings from OpenAI to Google Cloud audio encoding - FORMAT_MAPPINGS = { - "mp3": "MP3", - "opus": "OGG_OPUS", - "aac": "MP3", # Google doesn't have AAC, use MP3 - "flac": "FLAC", - "wav": "LINEAR16", - "pcm": "LINEAR16", - } + FORMAT_MAPPINGS: ClassVar[Mapping[str, str]] = MappingProxyType( + { + "mp3": "MP3", + "opus": "OGG_OPUS", + "aac": "MP3", # Google doesn't have AAC, use MP3 + "flac": "FLAC", + "wav": "LINEAR16", + "pcm": "LINEAR16", + } + ) + GEMINI_FORMAT_MAPPINGS: ClassVar[Mapping[str, str]] = MappingProxyType( + {"mp3": "MP3", "opus": "OGG_OPUS", "alaw": "ALAW", "mulaw": "MULAW"} + ) def __init__(self) -> None: BaseTextToSpeechConfig.__init__(self) @@ -96,7 +267,9 @@ class VertexAITextToSpeechConfig(BaseTextToSpeechConfig, VertexBase): def _map_voice_to_vertex_format( self, - voice: str | dict | None, + voice: str | dict[str, object] | None, + model: str | None = None, + drop_params: bool = False, ) -> tuple[str | None, dict | None]: """ Map voice to Vertex AI format. @@ -111,6 +284,9 @@ class VertexAITextToSpeechConfig(BaseTextToSpeechConfig, VertexBase): - voice_str: Original string voice (for interface compatibility) - voice_dict: Vertex AI format dict with languageCode and name """ + if model is not None and self._is_gemini_tts_model(model): + return self._map_gemini_tts_voice_to_vertex_format(model=model, voice=voice, drop_params=drop_params) + if voice is None: return None, None @@ -142,11 +318,225 @@ class VertexAITextToSpeechConfig(BaseTextToSpeechConfig, VertexBase): return voice_str, voice_dict + @staticmethod + def _is_gemini_tts_model(model: str) -> bool: + from litellm.utils import is_gemini_tts_model + + return is_gemini_tts_model(model) + + @staticmethod + def _get_str_value( + source: Mapping[str, object], + *keys: str, + ) -> str | None: + for key in keys: + value = source.get(key) + if isinstance(value, str): + return value + return None + + @staticmethod + def _get_dict_value( + source: Mapping[str, object], + *keys: str, + ) -> Mapping[str, object] | None: + for key in keys: + value = source.get(key) + if (mapping := as_str_mapping(value)) is not None: + return mapping + return None + + @staticmethod + def _resolve_gemini_tts_voice_name(model: str, voice_name: str, drop_params: bool) -> str: + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import OPENAI_TO_GEMINI_TTS_VOICES + + gemini_voice_name: Final = OPENAI_TO_GEMINI_TTS_VOICES.get(voice_name.lower(), voice_name) + if "-" not in gemini_voice_name: + return gemini_voice_name + if drop_params: + return _DEFAULT_GEMINI_TTS_VOICE + raise UnsupportedParamsError( + message=( + f"{model} speaks Gemini voices such as {_DEFAULT_GEMINI_TTS_VOICE}, and {voice_name} is a Cloud " + "Text-to-Speech voice. Use a Gemini voice, or set drop_params to use the default Gemini voice." + ), + model=model, + llm_provider="vertex_ai", + ) + + @staticmethod + def _reject_malformed_gemini_tts_voice(model: str, drop_params: bool) -> None: + if drop_params: + return + raise UnsupportedParamsError( + message=( + f"{model} needs voice to be a Gemini voice name such as {_DEFAULT_GEMINI_TTS_VOICE} or a speech " + "config object. Send one, or set drop_params to use the default Gemini voice." + ), + model=model, + llm_provider="vertex_ai", + ) + + @staticmethod + def _reject_voice_fields_cloud_tts_cannot_send(model: str, voice: Mapping[str, object], drop_params: bool) -> None: + if drop_params: + return + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + mapped_gemini_speech_config, + named_voice_keys, + ) + + speech_config: Final = VertexAITextToSpeechConfig._get_dict_value(voice, "speechConfig", "speech_config") + mapped_config: Final = mapped_gemini_speech_config(voice) + unsendable: Final = sorted( + { + *_fields_cloud_tts_cannot_send(mapped_config, _CLOUD_TTS_SPEECH_CONFIG_FIELDS, ""), + *_prebuilt_voices_replaced_by_speaker_id(mapped_config), + *_non_mapping_speech_config_wrappers(voice), + *_non_string_voice_names(voice, ""), + *_non_string_voice_names(speech_config or {}, "speechConfig."), + } + ) + if not unsendable: + return + raise UnsupportedParamsError( + message=( + f"{model} cannot send the voice fields {named_voice_keys(unsendable)} as mp3, opus, alaw, or mulaw. " + "Request wav or pcm, remove them, or set drop_params to ignore them." + ), + model=model, + llm_provider="vertex_ai", + ) + + @staticmethod + def _reject_malformed_gemini_tts_speaker(model: str, drop_params: bool) -> None: + if drop_params: + return + raise UnsupportedParamsError( + message=( + f"{model} needs each entry of speaker_voice_configs to name a speaker and a voice, as speaker with " + "voice_config.prebuilt_voice_config.voice_name, or speakerAlias with speakerId." + ), + model=model, + llm_provider="vertex_ai", + ) + + def _extract_gemini_tts_speaker_configs( + self, + model: str, + speech_config: Mapping[str, object], + drop_params: bool, + ) -> Sequence[VertexTextToSpeechSpeakerVoiceConfig]: + multi_speaker_config: Final = self._get_dict_value(speech_config, "multiSpeakerVoiceConfig") + if multi_speaker_config is None: + return [] + raw_speaker_configs: Final = multi_speaker_config.get("speakerVoiceConfigs", []) + try: + validated_speaker_configs: Final = _SPEAKER_CONFIG_ITEMS.validate_python(raw_speaker_configs, strict=True) + except ValidationError: + self._reject_malformed_gemini_tts_speaker(model, drop_params) + return [] + speaker_configs: Final[ # mutable-ok: validated payloads are accumulated for serialization + list[VertexTextToSpeechSpeakerVoiceConfig] + ] = [] + for raw_item in validated_speaker_configs: + if (raw_config := as_str_mapping(raw_item)) is None: + self._reject_malformed_gemini_tts_speaker(model, drop_params) + continue + speaker_alias = self._get_str_value(raw_config, "speakerAlias", "speaker") + speaker_id = self._get_str_value(raw_config, "speakerId") or self._prebuilt_voice_name(raw_config) + if speaker_alias is None or speaker_id is None: + self._reject_malformed_gemini_tts_speaker(model, drop_params) + continue + speaker_configs.append( + { + "speakerAlias": speaker_alias, + "speakerId": self._resolve_gemini_tts_voice_name(model, speaker_id, drop_params), + } + ) + return speaker_configs + + def _prebuilt_voice_name(self, config: Mapping[str, object]) -> str | None: + voice_config: Final = self._get_dict_value(config, "voiceConfig") + prebuilt_voice_config: Final = ( + self._get_dict_value(voice_config, "prebuiltVoiceConfig") if voice_config is not None else None + ) + if prebuilt_voice_config is None: + return None + return self._get_str_value(prebuilt_voice_config, "voiceName") + + def _map_gemini_tts_voice_to_vertex_format( + self, + model: str, + voice: object, + drop_params: bool = False, + ) -> tuple[str | None, dict[str, object]]: # mutable-ok: provider request serialization requires a concrete dict + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + mapped_gemini_speech_config, + reject_unknown_gemini_tts_voice_keys, + ) + + default_voice: Final[dict[str, object]] = { + "languageCode": self.DEFAULT_LANGUAGE_CODE, + "modelName": model, + "name": _DEFAULT_GEMINI_TTS_VOICE, + } + if voice is None: + return None, default_voice + + if isinstance(voice, str) and voice: + return voice, { + "languageCode": self.DEFAULT_LANGUAGE_CODE, + "modelName": model, + "name": self._resolve_gemini_tts_voice_name(model, voice, drop_params), + } + + voice_config: Final = as_str_mapping(voice) + if voice_config is None: + self._reject_malformed_gemini_tts_voice(model, drop_params) + return None, default_voice + reject_unknown_gemini_tts_voice_keys(model, voice_config, drop_params, "vertex_ai") + + speech_config: Final = mapped_gemini_speech_config(voice_config) + language_code: Final = self._get_str_value(speech_config, "languageCode") or self.DEFAULT_LANGUAGE_CODE + model_name: Final = model + speaker_configs: Final = self._extract_gemini_tts_speaker_configs(model, speech_config, drop_params) + self._reject_voice_fields_cloud_tts_cannot_send(model, voice_config, drop_params) + if speaker_configs: + return None, { + "languageCode": language_code, + "modelName": model_name, + "multiSpeakerVoiceConfig": { + "speakerVoiceConfigs": speaker_configs, + }, + } + voice_name: Final = self._prebuilt_voice_name(speech_config) + if voice_name is not None: + return None, { + "languageCode": language_code, + "modelName": model_name, + "name": self._resolve_gemini_tts_voice_name(model, voice_name, drop_params), + } + return None, { + "languageCode": language_code, + "modelName": model_name, + "name": _DEFAULT_GEMINI_TTS_VOICE, + } + + @staticmethod + def _dispatch_voice_name(voice: str | Mapping[str, object] | None) -> str | None: + if isinstance(voice, str): + return voice + if not isinstance(voice, Mapping): + return None + name: Final = voice.get("name") + return name if isinstance(name, str) else None + def dispatch_text_to_speech( self, model: str, input: str, - voice: str | dict | None, + voice: str | dict[str, object] | None, optional_params: dict[str, object], litellm_params_dict: dict[str, object], logging_obj: "LiteLLMLoggingObj", @@ -175,14 +565,17 @@ class VertexAITextToSpeechConfig(BaseTextToSpeechConfig, VertexBase): vertex_project: Final = self.safe_get_vertex_ai_project(litellm_params_dict) vertex_location: Final = self.safe_get_vertex_ai_location(litellm_params_dict) - # Convert voice to string if it's a dict (extract name) - # Actual voice mapping happens in map_openai_params - voice_str: str | None = None - if isinstance(voice, str): - voice_str = voice - elif isinstance(voice, dict): - # Extract voice name from dict if needed - voice_str = voice.get("name") if voice else None + mapped_voice, mapped_params = ( + self.map_openai_params( + model=model, + voice=voice, + optional_params=optional_params, + drop_params=drops_unsupported_params(kwargs.get("drop_params")), + kwargs=kwargs, + ) + if self._is_gemini_tts_model(model) + else (self._dispatch_voice_name(voice), optional_params) + ) # Store credentials in litellm_params for use in transform methods litellm_params_dict.update( @@ -198,9 +591,11 @@ class VertexAITextToSpeechConfig(BaseTextToSpeechConfig, VertexBase): response: Final = base_llm_http_handler.text_to_speech_handler( model=model, input=input, - voice=voice_str, - text_to_speech_provider_config=self, - text_to_speech_optional_params=optional_params, + voice=mapped_voice, + text_to_speech_provider_config=( + VertexAIGeminiCloudTextToSpeechConfig() if self._is_gemini_tts_model(model) else self + ), + text_to_speech_optional_params=mapped_params, custom_llm_provider="vertex_ai", litellm_params=litellm_params_dict, logging_obj=logging_obj, @@ -219,13 +614,13 @@ class VertexAITextToSpeechConfig(BaseTextToSpeechConfig, VertexBase): Note: Vertex AI also supports additional parameters like audioConfig which can be passed but are not part of the OpenAI spec """ - return ["voice", "response_format", "speed"] + return ["voice", "response_format", "speed", *(["instructions"] if self._is_gemini_tts_model(model) else [])] def map_openai_params( self, model: str, optional_params: dict, - voice: str | dict | None = None, + voice: str | dict[str, object] | None = None, drop_params: bool = False, kwargs: dict = {}, ) -> tuple[str | None, dict]: @@ -244,19 +639,28 @@ class VertexAITextToSpeechConfig(BaseTextToSpeechConfig, VertexBase): Tuple of (mapped_voice_str, mapped_params) """ mapped_params: Final[dict[str, object]] = {} + openai_params: Final = _JSON_OBJECT.validate_python(optional_params) + drops_params: Final = drop_params or litellm.drop_params is True ########################################################## # Map voice using helper ########################################################## - mapped_voice_str, voice_dict = self._map_voice_to_vertex_format(voice) + mapped_voice_str, voice_dict = self._map_voice_to_vertex_format( + voice=voice, + model=model, + drop_params=drops_params, + ) if voice_dict is not None: mapped_params["vertex_voice_dict"] = voice_dict # Map response format if "response_format" in optional_params: format_name: Final = optional_params["response_format"] - if format_name in self.FORMAT_MAPPINGS: - mapped_params["audioEncoding"] = self.FORMAT_MAPPINGS[format_name] + format_mappings: Final = ( + self.GEMINI_FORMAT_MAPPINGS if self._is_gemini_tts_model(model) else self.FORMAT_MAPPINGS + ) + if format_name in format_mappings: + mapped_params["audioEncoding"] = format_mappings[format_name] else: # Try to use it directly as Google Cloud format mapped_params["audioEncoding"] = format_name @@ -265,10 +669,19 @@ class VertexAITextToSpeechConfig(BaseTextToSpeechConfig, VertexBase): mapped_params["audioEncoding"] = self.DEFAULT_AUDIO_ENCODING # Map speed (OpenAI: 0.25-4.0, Vertex AI: speakingRate 0.25-4.0) - if "speed" in optional_params: - speed: Final = optional_params["speed"] - if speed is not None: - mapped_params["speakingRate"] = str(speed) + speed: Final = openai_params.get("speed") + if speed is not None: + mapped_params["speakingRate"] = str(speed) + + instructions: Final = openai_params.get("instructions") + if isinstance(instructions, str) and self._is_gemini_tts_model(model): + mapped_params["prompt"] = instructions + elif instructions is not None and self._is_gemini_tts_model(model) and not drops_params: + raise UnsupportedParamsError( + message=f"{model} needs instructions to be a string. Send one, or set drop_params to ignore it.", + model=model, + llm_provider="vertex_ai", + ) # Pass through Vertex AI-specific parameters from kwargs if "audioConfig" in kwargs: @@ -311,15 +724,48 @@ class VertexAITextToSpeechConfig(BaseTextToSpeechConfig, VertexBase): Google Cloud TTS endpoint: https://texttospeech.googleapis.com/v1/text:synthesize """ + if self._is_gemini_tts_model(model): + return self._get_gemini_tts_url( + model=model, api_base=api_base, litellm_params=_JSON_OBJECT.validate_python(litellm_params) + ) + if api_base: return api_base return self.TTS_API_URL + def _get_gemini_tts_url(self, model: str, api_base: str | None, litellm_params: Mapping[str, object]) -> str: + if api_base: + raise UnsupportedParamsError( + message=( + f"{model} serves mp3, opus, alaw and mulaw through Cloud Text-to-Speech, which does not use " + "this deployment's api_base. Request wav or pcm, or use a deployment without api_base." + ), + model=model, + llm_provider="vertex_ai", + ) + location: Final = litellm_params.get("vertex_location") + vertex_location: Final = location if isinstance(location, str) else self.get_default_vertex_location() + if vertex_location == "global": + return self.TTS_API_URL + cloud_tts_location: Final = _GEMINI_CLOUD_TTS_LOCATIONS.get(vertex_location) + if cloud_tts_location is None: + raise UnsupportedParamsError( + message=( + f"{model} serves mp3, opus, alaw and mulaw through Cloud Text-to-Speech, which has no Gemini TTS " + f"location for vertex_location {vertex_location}. Use a deployment in global, us, eu, " + "northamerica-northeast1, or a US or EU Vertex AI region." + ), + model=model, + llm_provider="vertex_ai", + ) + return f"https://{cloud_tts_location}-texttospeech.googleapis.com/v1/text:synthesize" + def _validate_vertex_input( self, input_data: VertexTextToSpeechInput, optional_params: dict, + detect_ssml: bool, ) -> VertexTextToSpeechInput: """ Validate and transform input for Vertex AI TTS @@ -343,7 +789,7 @@ class VertexAITextToSpeechConfig(BaseTextToSpeechConfig, VertexBase): else: # LiteLLM will auto-detect if text is in ssml format # check if "text" is an ssml - in this case we should pass it as ssml instead of text - if input_data: + if detect_ssml and input_data: _text: Final = input_data.get("text", None) or "" if "" in _text: input_data["ssml"] = input_data.pop("text") @@ -403,8 +849,13 @@ class VertexAITextToSpeechConfig(BaseTextToSpeechConfig, VertexBase): headers["x-goog-user-project"] = vertex_project ####### Build the request ################ - vertex_input = VertexTextToSpeechInput(text=input) - vertex_input = self._validate_vertex_input(vertex_input, optional_params) + prompt: Final = _JSON_OBJECT.validate_python(optional_params).get("prompt") + validated_input: Final = self._validate_vertex_input( + VertexTextToSpeechInput(text=input), optional_params, detect_ssml=not self._is_gemini_tts_model(model) + ) + vertex_input: Final[VertexTextToSpeechInput] = ( + {**validated_input, "prompt": prompt} if isinstance(prompt, str) else validated_input + ) # Build voice configuration # Check for voice dict stored in: @@ -412,8 +863,8 @@ class VertexAITextToSpeechConfig(BaseTextToSpeechConfig, VertexBase): # 2. optional_params by map_openai_params voice_dict: Final = litellm_params.get("vertex_voice_dict") or optional_params.get("vertex_voice_dict") if voice_dict is not None and isinstance(voice_dict, dict): - vertex_voice = VertexTextToSpeechVoice(**voice_dict) - elif voice is not None and isinstance(voice, str): + vertex_voice = _JSON_OBJECT.validate_python(voice_dict) + elif voice is not None: # Handle string voice (shouldn't normally happen if dispatch was called) parts: Final = voice.split("-") if len(parts) >= 2: @@ -485,8 +936,44 @@ class VertexAITextToSpeechConfig(BaseTextToSpeechConfig, VertexBase): content=binary_data, ) - # Initialize the HttpxBinaryResponseContent instance - return HttpxBinaryResponseContent(response) + binary_response: Final = HttpxBinaryResponseContent(response) + if self._is_gemini_tts_model(model): + from litellm.types.utils import CompletionTokensDetailsWrapper, Usage + + request: Final = _GeminiCloudTtsRequest.model_validate_json(raw_response.request.content) + input_data: Final = request.input + audio_config: Final = request.audioConfig + container_duration: Final = calculate_request_duration(binary_data) + sample_rate: Final = audio_config.sampleRateHertz or 24000 + duration: Final = ( + container_duration + if container_duration is not None + else _fallback_gemini_tts_audio_duration(binary_data, audio_config.audioEncoding, sample_rate) + ) + if duration is None: + verbose_logger.warning( + "Could not measure Gemini TTS audio duration for %s, so its audio tokens are not billed", model + ) + return binary_response + input_text: Final = " ".join( + value for value in (input_data.text, input_data.ssml, input_data.prompt) if value + ) + prompt_tokens: Final = litellm.token_counter(model=model, text=input_text) + audio_tokens: Final = math.ceil(duration * 25) + binary_response.usage = Usage( + prompt_tokens=prompt_tokens, + completion_tokens=audio_tokens, + total_tokens=prompt_tokens + audio_tokens, + completion_tokens_details=CompletionTokensDetailsWrapper(audio_tokens=audio_tokens), + ) + return binary_response + + +class VertexAIGeminiCloudTextToSpeechConfig(VertexAITextToSpeechConfig): + def get_error_class(self, error_message: str, status_code: int, headers: dict[str, str]) -> BaseLLMException: + return super().get_error_class( + error_message=error_message, status_code=400 if status_code == 404 else status_code, headers=headers + ) class VertexAILyriaTextToSpeechConfig(VertexAITextToSpeechConfig): diff --git a/litellm/main.py b/litellm/main.py index d647cab24e3..079b0bcbd6c 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -8846,8 +8846,11 @@ def speech( generic_optional_params: Final = GenericLiteLLMParams.model_validate(kwargs) - # Handle Gemini models separately (they use speech_to_completion_bridge) - if "gemini" in model: + needs_cloud_text_to_speech: Final = ( + litellm.utils.is_gemini_tts_model(model) + and response_format in VertexAITextToSpeechConfig.GEMINI_FORMAT_MAPPINGS + ) + if "gemini" in model and not needs_cloud_text_to_speech: from .endpoints.speech.speech_to_completion_bridge.handler import ( speech_to_completion_bridge_handler, ) diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index 2e07694876b..f7289445d3a 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -477,6 +477,18 @@ async def get_authorized_credentials_for_model( ) +def deployment_drop_params( + llm_router: Optional["Router"], model_id: str, request: Mapping[str, object] +) -> Mapping[str, bool | str]: + deployment: Final = ( + llm_router.get_credential_deployment(model_id=model_id) + if llm_router is not None and "drop_params" not in request + else None + ) + drop_params: Final = deployment.litellm_params.drop_params if deployment is not None else None + return MappingProxyType({}) if drop_params is None else MappingProxyType({"drop_params": drop_params}) + + def _key_allowlist_covers(public_model_name: str, key_model_allowlist: Sequence[str]) -> bool: from litellm.proxy.auth.auth_checks import model_matches_any_wildcard_pattern_in_list diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index 8823f01d781..d2ceefffe47 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -76,6 +76,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import ( # noqa: F401 # add_internal_model_credentials, apply_team_provider_credentials, authorize_model_for_key, + deployment_drop_params, encode_file_id_with_model, extract_file_creation_params, get_authorized_credentials_for_model, @@ -457,7 +458,7 @@ async def route_create_file( # Merge credentials into the request prepare_data_with_credentials( data=_create_file_request, - credentials=credentials, + credentials={**credentials, **deployment_drop_params(llm_router, model, _create_file_request)}, ) # Create the file with model credentials diff --git a/litellm/types/llms/vertex_ai.py b/litellm/types/llms/vertex_ai.py index 0480d6b1f3e..f1176f777c7 100644 --- a/litellm/types/llms/vertex_ai.py +++ b/litellm/types/llms/vertex_ai.py @@ -239,8 +239,18 @@ class VoiceConfig(TypedDict): prebuiltVoiceConfig: PrebuiltVoiceConfig +class SpeakerVoiceConfig(TypedDict): + speaker: ReadOnly[str] + voiceConfig: ReadOnly[VoiceConfig] + + +class MultiSpeakerVoiceConfig(TypedDict): + speakerVoiceConfigs: ReadOnly[Sequence[SpeakerVoiceConfig]] + + class SpeechConfig(TypedDict, total=False): voiceConfig: VoiceConfig + multiSpeakerVoiceConfig: ReadOnly[MultiSpeakerVoiceConfig] languageCode: str diff --git a/litellm/types/llms/vertex_ai_text_to_speech.py b/litellm/types/llms/vertex_ai_text_to_speech.py index e8f9fbe4a24..86b131dce59 100644 --- a/litellm/types/llms/vertex_ai_text_to_speech.py +++ b/litellm/types/llms/vertex_ai_text_to_speech.py @@ -4,7 +4,9 @@ Type definitions for Vertex AI Text-to-Speech API Reference: https://cloud.google.com/text-to-speech/docs/reference/rest/v1/text/synthesize """ -from typing_extensions import TypedDict +from collections.abc import Sequence + +from typing_extensions import ReadOnly, TypedDict class VertexTextToSpeechInput(TypedDict, total=False): @@ -16,6 +18,16 @@ class VertexTextToSpeechInput(TypedDict, total=False): text: str | None ssml: str | None + prompt: ReadOnly[str | None] + + +class VertexTextToSpeechSpeakerVoiceConfig(TypedDict): + speakerAlias: ReadOnly[str] + speakerId: ReadOnly[str] + + +class VertexTextToSpeechMultiSpeakerVoiceConfig(TypedDict): + speakerVoiceConfigs: ReadOnly[Sequence[VertexTextToSpeechSpeakerVoiceConfig]] class VertexTextToSpeechVoice(TypedDict, total=False): @@ -29,6 +41,8 @@ class VertexTextToSpeechVoice(TypedDict, total=False): languageCode: str name: str + modelName: ReadOnly[str] + multiSpeakerVoiceConfig: ReadOnly[VertexTextToSpeechMultiSpeakerVoiceConfig] class VertexTextToSpeechAudioConfig(TypedDict, total=False): @@ -42,6 +56,7 @@ class VertexTextToSpeechAudioConfig(TypedDict, total=False): audioEncoding: str speakingRate: str + sampleRateHertz: ReadOnly[int] class VertexTextToSpeechRequest(TypedDict, total=False): diff --git a/litellm/utils.py b/litellm/utils.py index 1a114ffe9f0..7de7ce6188d 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -860,6 +860,11 @@ def _is_gemini_model(model: str | None, custom_llm_provider: str | None) -> bool return model is not None and "gemini" in model.lower() +def is_gemini_tts_model(model: str) -> bool: + model_name: Final = model.lower() + return "gemini" in model_name and "tts" in model_name + + def _remove_thought_signature_from_id(tool_call_id: str, separator: str) -> str: """ Remove thought signature from a tool call ID. diff --git a/tests/audio_tests/gemini_tts_speech.mp3 b/tests/audio_tests/gemini_tts_speech.mp3 new file mode 100644 index 0000000000000000000000000000000000000000..71a916ad6c0803a578e56162e8d97996d999b228 GIT binary patch literal 4992 zcmeI#cTiK?x&ZLK6TlE!0s$fdhN6b1p$P~G0U}K$Vi2&=5wL|S7ETbUL8T+oq>B+1 z1Tlh05fB0CHUt4hsRtB6kRw?y@ywhvZ|?i+zI*<=-^`xbv-a%Y`sQ2Tti2cbX3-2` zrozHua}2zV0b;Xi8?;SfoPMQ2Ph|GD@s|Mn7)DulJAx&{{9yzZ|7CepQAy#aHW+_e zIEq)g2LM=+6N^qnYCru6bsC-fX6AYpIbyvbcfn)bYkx1m^~EIp><`27L{yRG9phay1Y`NEPBEGIaQ2#u1rsCn%xv5{`QpN46tG03i! z`c2h$W-N?$(|}q3xRjettV&5bE$igU@_}7}vBbEMf*7jd%<-!!}WthPrLz?uWIbs2F<<>GvG* z-PcHujIWM>19PMhjn~KR&fPd~qjO!o@Mueu?sxD500yZb>W~)=*$l6eM_Cz0L|1o+ z=ZoTc5IXQBJd9kHg|xrfSm%+T2G}IgN?i+1d=}1Fp#{LHiHXKr8uiMV%_Wi)H*W@Eiwv&@`CVp)N*ii}G2!T!* z*4&lAVBWG!^+pN%^Q+Kd*uJyzhPRgB_DI~X{3UOrE#L?uoR-qqLp@te8Nm~`&YNd? ztWk!1vYIK>mhW$TvMOvWf+bbbQq~5>TI~&!dTQ=J@*Y`Byh^Y}b+&&LN|V;SH4Yd% zX1A^ys7ggXZ-V_c%V1S_J{o7wBEd2GP@Hme&tVAry^PFdBOD7aqjf?Uq zaFaa)zLXCI2@1L=k;?W15EUP(WL8)#D4OBO?NLYF!+wR{;&Fn?z;F7^m_^-aL~h_^ zAwQlkgl!wNf+H{x+deuaXUqowF!xN34ZaNhG5##6r?jJc#D}f5r^|qmAC9m!@$Zj5vQi;F~_NR%N0Ey51uX@k|4okw2A|^242wDGOyH_{mLKL$g+Z?@leQ(Tf$Cy_^<=5 zC9e5c+?wvso^YJlpq}G*1%fx{o6FJDjte@n1cFeSlOTzN5P! z6iU&P&>UcV{d#in!L&MYBU95R5U6&)ddA>p>U@BhLP_&yhC{A-?NsEsixHy^CrEAn zu2PCQaHQ^p+IGFq$&Q!@t$LE_Pg00Yje0abp|?kPz9yz2`2u_x5BVuCP$Co^IQPyh za!Ifvu*sR++m86O!vgnWV2?Q--rq7-@v2sd9~xGfiFLh7G(ria%#|8Zjgva!U%s`J zF!Agir=#p*cE%ja=qB*BRt@revTT!L6@1+Zx{|ftZe%fD%%FeUF}==rs~<8dCJf4J zRKIN$KJGoj02B6wFOrQk00Ui)a4OHx{XL=Mq5VMr+0`fIgiYqo2P@JoFrr4uqI_@v zSZq5uS9^B`=J^U(!R=M>TOx!{DjGPQRe5F#d>?XtlaE!#2d_X9{)~z$^+N!x7G20} zRqa6QyK*#2?a0toq(;!)l3v}fN?Jyu%|AW{$ZtHTrnEu5FvsNG?K|B4ZuV{aL7u-2 ztHLaUIwMoBznxE^NOciBOiTzz%qQvrm_p@JhJBBx(kYMoYWB-$a_+$>Xr0VWlGsXa zY+SxPa@IcDGxYxRL2vJ_$lj$OW9~aZyD%+D!wg-)7-CKGd}DFn)-q@yGSvek;{slT zwY|RQCF~!(3H45C-F!yW`}jts+om3sOWMD;Zppc}7eOcof5Z1yI^F#uf)tTD?j!)Yvro4J&+~U--V_u;U4-2isrwcJ`soJ;XBz_8^V=R~ zHJ-m4&F=di$2@K_bo$2kSyRQ?Oz$nV>1}BK#?J~paEEmS-XnVF{$a(04(^9E_~6y> z*P_SO&l+1lMCulT%f% zS|H1)oc#~#>dFJ}_itl8tW#L+77!bFEt9?T+mth<8|qEB7pVszr`QDmIfjsz;B&k@ zpTZx>Jl{(^-=+=nu}?i@_)^E(W5CBhe4lu9@RLIo*s?&C5u*?5BQ?Id6%ntZa!-s^ z5yl#hFgY=FXR*$vPA@OZ)xvXOa|-gcpd`dWA$7Nayia0zXkcomZC+n?U@Xr+iG{U# zAbl#G-m|wuhQ%)5*W(RH69AkXKJw)$_v zPMtsM+%&7@M6szCjD{D;HAwO z(jxL;i2=Z}-@-^i$A&alMQT@jcj`WQCHyI~am)wNsBS3-KYC7Ux&yKLW>?$q{_ZiB zat;AmX>xLEwXZu_ZOZD3SORO%1aO1vGsXQ?tDycqIsl$tj0IgVY&$Q`(KXgw5foX- z?ke7(V8@^A=!>aCph_hA>~vFFlUdGB&eT8t$OAYnk81u7vp+0KS(` zPl6sQX{w+*sL@Lx(Ze_XsRq$8P zakwLU2WZ9dVaNzaw*wI ztoU1e^U&cdz35RjmSSjb;(KpR#Dxv{Jii9}s34J-|9&+4l$ch0$4rRl9$K;wZ9aIu zc37Ibuice+q7ik)7|+`1g-SB}h)itrOW1KY)amSlP@m|#v7Jd9WH(_V`B+)9JJz4K z0V`4Wa%(?dOS-oAD9>+|ZfY%pY?|3=vxh74Pd(ppd9X8?kJu~i z9^f65pxVX!erC1)cL3G`y}p?nP-%FAJ8|#o=%#RMR1Dj7&&k`1OLk>JSaf6ymDTQE zxbiE1Sv7Piq(*(@dQqTSygUL0lt7F{HUkC2wz-wz1>PouJP*h{DMOLEi3AM-_or2hT6_NKO3u%$%ce1 zQp18W4_8VlES!?kZc*8J8jU;cvYfEgs|UcMW{GBTzlN8FB$JAFAJB^m12uBt53xIvA$y2@WCDic)!QkC^Ph&-=u4Usi~ zo5f_RlMDm=ELBBca#9sSn!emk-b2yes!1X_I5_-y&+|VjZcx>SFO#4@A8?=jCxX#- zB3#s*=ik9jYLVb32=at>X$8h+Jx1z4d)~t}5Guc=L63LXV+*JgAg1hLBt0BK#LFUs zAgpP_8{r;^BC7s06WS!IKGZADt_+NI7@C*3I{}use*;lRK&itB&o6@@q!q*s3+{kL z9^dHc0xBO}lTLuGS~oTb>{Ul(o0f|Wg}cHg7la`U{YKaqa+837#PUeA018EP-3Gpt zAYc*(HR|x%1;oixliZ`F^cy^09|``oKmX-dzb6Pgg=>eyF%VJ!V;oFr-x~Lfc4j+|FJXBQnPGzW?c2*f5hYQ3V+ zSlp8VH-8{gyWR>W<6-b)S=Jpn{+x*t;&)%#AD(~Y>F9Q=4DV-joW4j)TT?iR9i%k< zVX`TIOu;pw?e2MEkjJ`cj$NlBVV_bKShJ!Tq(*T|qmmkBe%(q5{860F^CxgJ)5^T_ z#bmQ%Z}U}si-;!kW($f%8pyw*Vrv$jcA3xD97dD!6hqJm(F%aXp=gjyLtW;PDv#aj?#xjB>8O-H_3UW>VGvF z>)(#Q`*TK8d|C&NBJ%q4-{$`1#lQV23H>D;+=Tz9{J;9=AIkr~{=Y0Pci`{!{b%a) T=bZm*|M|bt_t2m^|QgoM(m zln4j{FP`&wp7Y_k-gVu3U-#a>|9{Xx3iVgh`EVtj%E0>359 z*O$#>q>Hu96X|P+yoSPp*HDPf$==H4+WB|pptft>b^Ol!yDC+*21W=6$4pH@T0lVP z-d!m^5fNTJOJk(5HUeR2W}%~}g+Rdd;6{)1byd}MG!^AUc{$nGX09%;7eKP9oSI5z zmXRP2JJCH*ygkViZtp8Mt0Sgo{uPxqd8WtBUBj}?(pdGvdsFh|6+guQbnsgV{OYeB z{8h*Qa|42qWRLK2ipLcrC?2c z$K*k~wy`ZmfWL8;#8kUxkak}^&^kiHhBiP$4hV90AN(NYprEptMQDmeF8Yc`aq9M* z&*jTbqCxiyWBI1}8I1I4>0HzHVLhDe?9+T2(3s*-jE3`pc(^-{jjdi%Et3QT(}R*AF{hxudg40IT)G zd-PVx`Z0IMcB!#7lV zzhGO5gOtML;MN1plP$$AG{6jVAn-h3D(Y9b(TR(B;JEU&=567!5D36}@rozy{PlE# zz_EsU$;k;~(8i0$_{!UTK2S^hW=i0)$~>(SsS-P?N<`oJC_Z{8*{%ie_F6KguM0^{Mw+VCZ2ojr4*etO4Amr;EgGd>Wu`)jfKtfEY z%|@D+weJpZW%FVcwycF>?&u5Z)&%x5>2c>Jjdvl5AAZik>wmQbf&ie$wEah!X+|!g z?4K6$nc!UlS*h(LBlFqVM_xXG)nGH*Wtbe9ER7WN;woR`cm~xZ830;ge`mrVd16hi zId*H{#SLWS%hbZ$#tWGxMZx`uJJD)+kHa?@RI4Lk(Cmpy9LhTI3_!3D!+hxYn0+kW z5&VE7(`T8GeeEV~3Mxij*3M!AQD}^S9J^IM$->aF!cuR*CSv7@w9L6fHuAu*o~tun zD1_@kmMfVMTOba|MO@Z36xAY=N}s^dMQ#kWd~B!^S;N7a za~*H`v7cV1(J0!yhY8)*jj1`El)--0*8VFD|ILU9*w(}-NY@G|SZs3Odje*UT69ig zdRBJ|NG6ej;QzopKxT59Gtfw^;?>(mK=lHW1t`ayN2GC2h*0-wfS~+69mL zs+$MA>2CA8t?cZ@Dm{nl?+NeVZASlOxg5x*w7p{}O{5plexUdr$Qsb^4SQ(tb&V`V zV*Z59wj)}j+#~pm>xUqR&GZlCbSX?6H^lmxJm+lDj`+^zw8@OU##NfqR|Ko0q`^kh zF6Xa(maSP_sw7!XYTungcKu7^*puKOR7Us=e<_J!lDj(4n7!oDH)n{QGDo8wtx~S; z1IMb~D z^x*4+pi#LHS9mtKs4iFT&67N!`kmwf{_dd4(Kq2L{xPUDJH&5s+PwMfMKBHYGX;L+ z`5kknS(>g+8~8!ET}kVR^{|$c8`;!r9M(>Z0XWX8cV9!BIr}yp0Sx@)nf?|~uI&8- z=9?e?HN1Y85H0KA)GU6}(PtQoZ)ulxf>WzwX4iyGp1XC!k!nx%|CA4|Ue+(2L~{ganH$ ziR=*<7xrICiVq2Y#sxfkUR{__PKY3+Nh82voC#@swo}Cop&&$%v>b3I?8^(H4oD#Ge=X{3qzo%68 zM89QKyVy37BUX69~XDGfKFlC8{aYr?s5=xgr1;s4Hi-~BOC>_sg>!+&q6CCK@F?; zf^N#`dwc!eFRa%WiuuUoS?0z)qEha%YWYcEOpgS%OSR}?ECetL(yetsgUzG4M|X6B zpZ2I2QXQ1<^#vok;TXpQ(|e?`w8kFR$gPFa@1O+$D5#nzG{nSs_>#pBIFhL0JzB&I zx2fQF&!C2guelEUjP4$~(L-UuaQNb4tLLOeXM=-ira6$h60gV4_yD|g7@nX!xWusG z8Ld3A>bKW}tq!3hw^`syhvD?{;cJ#)ch7eORvzM~MN_^JhZ)acUMSE`c#cHcdWqmc zQYL4>k*&JeD^5U~=!$=};*+raC^}H3$CUBT_PHsEXXVE?OwPE7Vl+}LHLdK?VECm= z4?-Fgn0gBF>T!uGI(6GktkZUCH2;=>u9DhRBeFZ^tP!|89c6sA1Egh1&;JsjpYQkn z9N{HIlrZ#?bK`NT>Pz!Q^qh??g$H169a4h3yF{`m&{<~mc}z`^krk@n zHW+WMtQRM_z<_GaR^6ny$~0L=2I6%5bbsRY-WJkrBvFDGt)M~^dlF{R>uVn81v%(9 z0xHaKomKemK7A_~TefU&5fi4~F45(D7`f%>J+^4f%oEd31_6JU{xD}vEgb#JBwybT zE}s1*r#u&HPy>n7+8-0-%N;Kt%TB@}M=~Z%I9UGDUB)Mf+86~&bT2viQcsyl(4((* zDO!*TVP?(na*nEE8aEUJNt>x%nv~)ppUXvx0!c;JHj2vM)({`pt{{Q2W4lXEYWaTg zAz9eHot`LsQZH44fZwL9L$h4|Dc0n zhU!Y<@zSgdhh6#P%yvtbT^Jd|Ns8 zb$tJ&;^jw^{Hfp-RQRyRb zeq9=t80j|jkc=qZ@S;$6&d?c1&YTQrb|ox1>O=bwEyJZn_Z{Y45Nq@JZ&}2nXEK|3 zCbl~mGm{qe_jCkC$3@?5TI9|2#POe8!f8KYnyp=|4ywZ)M8a6U58dDp6fw7i-nGV{ux0JgIH)Fk|V zAw7AWYua+zWf=rWlst0Bt1R_C_p=ts=YBe593X7QiaLyw#qmBA;r$wn-xhAQ$C#-o zQlfY&!8>;{!x19)O=JHfiHWf?6NbE4R+eI0`QZa0Um=@)C)gJR8 ze0$I;>%ux@M|k&r8)?I?I-q!~Ea&HDK3g}o@N)n+(qc8;O=qj3o~b{eUFW5X;7Ig7 zndKO@gek$WZTh(x>hdjOY(1PB#Fw37z6+UtO-h zX#)G86P^ou2|Xds{%W~(MwpOuAi2gGpg-}r+Er4Kl+C+QQ@*?vMaoB7EVo40Ub-z!GjeWfB@q&h%xHUr@9$5_1)bak^a#>==o9@|AtbcVo@Lc zKb#a2^PQQpIs21ui=485n%K|$Il4%Y)?XtXuEfsYL)7n)4(EF7PY#bJy~WGV&-Nz( zg87Kl_JWzrP@cFwAzXVJWZW0ApH_;!~fAN3&|8OG5JA;4ntNhDf@AF@O3j5Mu=L7$P6CwY%zvJcK{@?lk_$~eq@QA+< literal 0 HcmV?d00001 diff --git a/tests/unit/endpoints/speech/speech_to_completion_bridge/test_transformation.py b/tests/unit/endpoints/speech/speech_to_completion_bridge/test_transformation.py index 953f028af3c..cf53fe90216 100644 --- a/tests/unit/endpoints/speech/speech_to_completion_bridge/test_transformation.py +++ b/tests/unit/endpoints/speech/speech_to_completion_bridge/test_transformation.py @@ -11,14 +11,14 @@ from litellm.endpoints.speech.speech_to_completion_bridge.transformation import ) from litellm.types.utils import ChatCompletionAudioResponse, Choices, Message, ModelResponse +pytestmark = pytest.mark.usefixtures("local_model_cost_map") + GEMINI_TTS_MODEL: Final = "gemini-3.1-flash-tts-preview" PCM_BYTES: Final = b"\x01\x02\x03\x04" * 6 def _model_response(model: str, pcm: bytes) -> ModelResponse: - audio: Final = ChatCompletionAudioResponse( - data=base64.b64encode(pcm).decode(), expires_at=0, transcript="hello" - ) + audio: Final = ChatCompletionAudioResponse(data=base64.b64encode(pcm).decode(), expires_at=0, transcript="hello") return ModelResponse(model=model, choices=[Choices(message=Message(content=None, audio=audio))]) @@ -72,6 +72,103 @@ def test_non_gemini_request_forwards_speech_response_format_as_audio_format() -> assert request["audio"] == {"voice": "alloy", "format": "wav"} +def test_gemini_request_keeps_voice_beside_nested_speech_config() -> None: + request: Final = SpeechToCompletionBridgeTransformationHandler().transform_request( + model=GEMINI_TTS_MODEL, + input="Bonjour", + voice={"speech_config": {"language_code": "fr-FR"}, "voice": "Kore"}, + optional_params={}, + litellm_params={}, + headers={}, + litellm_logging_obj=MagicMock(), + custom_llm_provider="gemini", + ) + + assert request["audio"]["speech_config"] == { + "languageCode": "fr-FR", + "voiceConfig": {"prebuiltVoiceConfig": {"voiceName": "Kore"}}, + } + + +def test_gemini_request_keeps_voice_object_for_a_model_alias_map_alias(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "model_alias_map", {"voice-alias": GEMINI_TTS_MODEL}) + + request: Final = SpeechToCompletionBridgeTransformationHandler().transform_request( + model="voice-alias", + input="Hello there", + voice={"voice": "Puck"}, + optional_params={}, + litellm_params={}, + headers={}, + litellm_logging_obj=MagicMock(), + custom_llm_provider="gemini", + ) + + assert request["audio"] == { + "speech_config": {"voiceConfig": {"prebuiltVoiceConfig": {"voiceName": "Puck"}}}, + "format": "pcm16", + } + + +@pytest.mark.parametrize("wrapper", (None, "speech_config", "speechConfig")) +def test_gemini_request_keeps_voice_and_language_without_audio_format(wrapper: str | None) -> None: + voice: Final = {"language_code": "fr-FR", "voice": "Kore", "format": "pcm16"} + request: Final = SpeechToCompletionBridgeTransformationHandler().transform_request( + model=GEMINI_TTS_MODEL, + input="Bonjour", + voice=voice if wrapper is None else {wrapper: voice}, + optional_params={}, + litellm_params={}, + headers={}, + litellm_logging_obj=MagicMock(), + custom_llm_provider="gemini", + ) + + assert request["audio"]["speech_config"] == { + "languageCode": "fr-FR", + "voiceConfig": {"prebuiltVoiceConfig": {"voiceName": "Kore"}}, + } + + +@pytest.mark.parametrize( + ("drop_params", "voice_config"), + [ + (True, {"prebuiltVoiceConfig": {"voiceName": "Kore"}}), + (False, {"prebuiltVoiceConfig": {"voiceName": "Kore", "style": "calm"}, "tone": "warm"}), + ], +) +def test_gemini_request_drops_unknown_nested_voice_fields_only_with_drop_params( + drop_params: bool, voice_config: dict[str, object] +) -> None: + request: Final = SpeechToCompletionBridgeTransformationHandler().transform_request( + model=GEMINI_TTS_MODEL, + input="Bonjour", + voice={"voice_config": {"prebuilt_voice_config": {"voice_name": "Kore", "style": "calm"}, "tone": "warm"}}, + optional_params={}, + litellm_params={"drop_params": drop_params}, + headers={}, + litellm_logging_obj=MagicMock(), + custom_llm_provider="gemini", + ) + + assert request["audio"]["speech_config"] == {"voiceConfig": voice_config} + + +def test_gemini_request_reads_snake_case_wrapper_beside_null_camel_case_wrapper() -> None: + request: Final = SpeechToCompletionBridgeTransformationHandler().transform_request( + model=GEMINI_TTS_MODEL, + input="Bonjour", + voice={"speechConfig": None, "speech_config": {"voice_name": "Kore"}}, + optional_params={}, + litellm_params={}, + headers={}, + litellm_logging_obj=MagicMock(), + custom_llm_provider="gemini", + ) + + assert request["audio"]["speech_config"] == {"voiceConfig": {"prebuiltVoiceConfig": {"voiceName": "Kore"}}} + + @pytest.mark.parametrize("response_format", ["mp3", "flac", "opus", "aac"]) def test_gemini_tts_request_rejects_formats_gemini_cannot_produce(response_format: str) -> None: with pytest.raises(litellm.BadRequestError) as excinfo: @@ -86,6 +183,7 @@ def test_gemini_tts_request_rejects_formats_gemini_cannot_produce(response_forma def test_gemini_tts_pcm_response_returns_raw_pcm_bytes() -> None: response: Final = SpeechToCompletionBridgeTransformationHandler().transform_response( model_response=_model_response(GEMINI_TTS_MODEL, PCM_BYTES), + model=GEMINI_TTS_MODEL, response_format="pcm", ) @@ -97,6 +195,7 @@ def test_gemini_tts_pcm_response_returns_raw_pcm_bytes() -> None: def test_gemini_tts_wav_and_default_responses_wrap_pcm_in_wav(response_format: str | None) -> None: response: Final = SpeechToCompletionBridgeTransformationHandler().transform_response( model_response=_model_response(GEMINI_TTS_MODEL, PCM_BYTES), + model=GEMINI_TTS_MODEL, response_format=response_format, ) @@ -110,8 +209,24 @@ def test_gemini_tts_wav_and_default_responses_wrap_pcm_in_wav(response_format: s def test_non_gemini_response_keeps_original_bytes_and_mpeg_content_type() -> None: response: Final = SpeechToCompletionBridgeTransformationHandler().transform_response( model_response=_model_response("gpt-4o-audio-preview", PCM_BYTES), + model="gpt-4o-audio-preview", response_format="mp3", ) assert response.response.content == PCM_BYTES assert response.response.headers["content-type"] == "audio/mpeg" + + +@pytest.mark.parametrize( + ("requested_model", "served_model"), + [(GEMINI_TTS_MODEL, "served-speech-alias"), ("requested-speech-alias", GEMINI_TTS_MODEL)], +) +def test_gemini_response_wraps_wav_when_either_model_name_is_gemini_tts(requested_model: str, served_model: str) -> None: + response: Final = SpeechToCompletionBridgeTransformationHandler().transform_response( + model_response=_model_response(served_model, PCM_BYTES), + model=requested_model, + response_format=None, + ) + + assert response.response.content[:4] == b"RIFF" + assert response.response.headers["content-type"] == "audio/wav" diff --git a/tests/unit/llms/gemini/test_gemini_tts.py b/tests/unit/llms/gemini/test_gemini_tts.py index a7522f6b315..d0ef04e6372 100644 --- a/tests/unit/llms/gemini/test_gemini_tts.py +++ b/tests/unit/llms/gemini/test_gemini_tts.py @@ -2,14 +2,70 @@ Test Gemini TTS (Text-to-Speech) functionality """ -import pytest +from typing import Final from unittest.mock import patch, MagicMock +import pytest import litellm +from litellm.endpoints.speech.speech_to_completion_bridge.transformation import ( + SpeechToCompletionBridgeTransformationHandler, +) from litellm.llms.gemini.chat.transformation import GoogleAIStudioGeminiConfig from litellm.utils import get_supported_openai_params +pytestmark = pytest.mark.usefixtures("local_model_cost_map") + +GEMINI_3_1_FLASH_TTS_MODEL = "gemini-3.1-flash-tts-preview" + +MULTI_SPEAKER_SPEECH_CONFIG = { + "multi_speaker_voice_config": { + "speaker_voice_configs": [ + { + "speaker": "Ryan", + "voice_config": { + "prebuilt_voice_config": { + "voice_name": "Umbriel", + }, + }, + }, + { + "speaker": "Katie", + "voice_config": { + "prebuilt_voice_config": { + "voice_name": "Leda", + }, + }, + }, + ], + }, + "language_code": "en-US", +} + +NORMALIZED_MULTI_SPEAKER_SPEECH_CONFIG = { + "multiSpeakerVoiceConfig": { + "speakerVoiceConfigs": [ + { + "speaker": "Ryan", + "voiceConfig": { + "prebuiltVoiceConfig": { + "voiceName": "Umbriel", + }, + }, + }, + { + "speaker": "Katie", + "voiceConfig": { + "prebuiltVoiceConfig": { + "voiceName": "Leda", + }, + }, + }, + ], + }, + "languageCode": "en-US", +} + class TestGeminiTTSTransformation: """Test Gemini TTS transformation functionality""" @@ -19,17 +75,16 @@ class TestGeminiTTSTransformation: config = GoogleAIStudioGeminiConfig() # Test TTS models (both preview and non-preview versions) - assert ( - config.is_model_gemini_audio_model("gemini-2.5-flash-preview-tts") == True - ) - assert config.is_model_gemini_audio_model("gemini-2.5-pro-preview-tts") == True - assert config.is_model_gemini_audio_model("gemini-2.5-flash-tts") == True - assert config.is_model_gemini_audio_model("gemini-2.5-pro-tts") == True + assert config.is_model_gemini_audio_model("gemini-2.5-flash-preview-tts") + assert config.is_model_gemini_audio_model("gemini-2.5-pro-preview-tts") + assert config.is_model_gemini_audio_model("gemini-2.5-flash-tts") + assert config.is_model_gemini_audio_model("gemini-2.5-pro-tts") + assert config.is_model_gemini_audio_model(GEMINI_3_1_FLASH_TTS_MODEL) # Test non-TTS models - assert config.is_model_gemini_audio_model("gemini-2.5-flash") == False - assert config.is_model_gemini_audio_model("gemini-2.5-pro") == False - assert config.is_model_gemini_audio_model("gpt-4o-audio-preview") == False + assert not config.is_model_gemini_audio_model("gemini-2.5-flash") + assert not config.is_model_gemini_audio_model("gemini-2.5-pro") + assert not config.is_model_gemini_audio_model("gpt-4o-audio-preview") def test_gemini_tts_supported_params(self): """Test that audio parameter is included for TTS models""" @@ -66,10 +121,7 @@ class TestGeminiTTSTransformation: assert "speechConfig" in result assert "voiceConfig" in result["speechConfig"] assert "prebuiltVoiceConfig" in result["speechConfig"]["voiceConfig"] - assert ( - result["speechConfig"]["voiceConfig"]["prebuiltVoiceConfig"]["voiceName"] - == "Kore" - ) + assert result["speechConfig"]["voiceConfig"]["prebuiltVoiceConfig"]["voiceName"] == "Kore" # Check response modalities assert "responseModalities" in result @@ -78,9 +130,7 @@ class TestGeminiTTSTransformation: def test_gemini_tts_audio_parameter_mapping_with_language_code(self): config = GoogleAIStudioGeminiConfig() - non_default_params = { - "audio": {"voice": "Kore", "format": "pcm16", "language_code": "en-US"} - } + non_default_params = {"audio": {"voice": "Kore", "format": "pcm16", "language_code": "en-US"}} optional_params = {} result = config.map_openai_params( @@ -102,6 +152,16 @@ class TestGeminiTTSTransformation: assert result["languageCode"] == "de-DE" assert result["voiceConfig"]["prebuiltVoiceConfig"]["voiceName"] == "Kore" + def test_map_audio_params_keeps_voice_beside_nested_speech_config(self): + config = GoogleAIStudioGeminiConfig() + + result = config._map_audio_params( + {"speech_config": {"language_code": "fr-FR"}, "voice": "Kore", "format": "pcm16"} + ) + + assert result["languageCode"] == "fr-FR" + assert result["voiceConfig"]["prebuiltVoiceConfig"]["voiceName"] == "Kore" + def test_map_audio_params_no_language_code(self): config = GoogleAIStudioGeminiConfig() @@ -110,6 +170,55 @@ class TestGeminiTTSTransformation: assert "languageCode" not in result assert result["voiceConfig"]["prebuiltVoiceConfig"]["voiceName"] == "Kore" + def test_map_audio_params_keeps_nested_speaker_list_and_drops_unknown_fields(self): + config: Final = GoogleAIStudioGeminiConfig() + + result: Final = config._map_audio_params( + { + "speech_config": { + "language_code": "fr-FR", + "multi_speaker_voice_config": { + "speaker_voice_configs": [ + {"speaker": "Ryan", "voice_config": {"prebuilt_voice_config": {"voice_name": "Umbriel"}}} + ] + }, + "metadata": ["studio", 2], + }, + "format": "pcm16", + } + ) + + assert result == { + "languageCode": "fr-FR", + "multiSpeakerVoiceConfig": { + "speakerVoiceConfigs": [ + {"speaker": "Ryan", "voiceConfig": {"prebuiltVoiceConfig": {"voiceName": "Umbriel"}}} + ] + }, + } + + def test_gemini_tts_multi_speaker_audio_parameter_mapping(self): + """Test multi-speaker audio parameter mapping for Gemini 3.1 TTS models""" + config = GoogleAIStudioGeminiConfig() + + non_default_params = { + "audio": { + "speech_config": MULTI_SPEAKER_SPEECH_CONFIG, + "format": "pcm16", + } + } + optional_params = {} + + result = config.map_openai_params( + non_default_params=non_default_params, + optional_params=optional_params, + model=GEMINI_3_1_FLASH_TTS_MODEL, + drop_params=False, + ) + + assert result["speechConfig"] == NORMALIZED_MULTI_SPEAKER_SPEECH_CONFIG + assert result["responseModalities"] == ["AUDIO"] + def test_gemini_tts_audio_parameter_with_existing_modalities(self): """Test audio parameter mapping when modalities already exist""" config = GoogleAIStudioGeminiConfig() @@ -193,9 +302,7 @@ class TestGeminiTTSTransformation: } optional_params = {} - with pytest.raises( - ValueError, match="Unsupported audio format for Gemini TTS models" - ): + with pytest.raises(ValueError, match="Unsupported audio format for Gemini TTS models"): config.map_openai_params( non_default_params=non_default_params, optional_params=optional_params, @@ -234,6 +341,28 @@ def test_gemini_tts_completion_mock(): assert response.choices[0].message.content is not None +def test_gemini_tts_speech_bridge_accepts_multi_speaker_voice_dict(): + handler = SpeechToCompletionBridgeTransformationHandler() + + result = handler.transform_request( + model=f"vertex_ai/{GEMINI_3_1_FLASH_TTS_MODEL}", + input="Ryan: How are you doing today Katie?\nKatie: Not too bad.", + voice=MULTI_SPEAKER_SPEECH_CONFIG, + optional_params={"response_format": "pcm"}, + litellm_params={}, + headers={}, + litellm_logging_obj=MagicMock(), + custom_llm_provider="vertex_ai", + ) + + assert result["modalities"] == ["audio"] + assert result["audio"] == { + "speech_config": NORMALIZED_MULTI_SPEAKER_SPEECH_CONFIG, + "format": "pcm16", + } + assert "response_format" not in result + + class TestGeminiTTSSpeechConfigInRequestBody: """Test that speechConfig is properly included in the final request body. @@ -250,6 +379,8 @@ class TestGeminiTTSSpeechConfigInRequestBody: ("gemini-2.5-flash-preview-tts", "vertex_ai"), ("gemini-2.5-flash-preview-tts", "gemini"), ("gemini-2.5-pro-tts", "vertex_ai"), + (GEMINI_3_1_FLASH_TTS_MODEL, "vertex_ai"), + (GEMINI_3_1_FLASH_TTS_MODEL, "gemini"), ], ) def test_speechconfig_in_generation_config_transform_request_body(self, model, custom_llm_provider): @@ -283,12 +414,7 @@ class TestGeminiTTSSpeechConfigInRequestBody: f"speechConfig was filtered out of generationConfig for model={model}, provider={custom_llm_provider}. " "Ensure speechConfig is in the GenerationConfig TypedDict." ) - assert ( - generation_config["speechConfig"]["voiceConfig"]["prebuiltVoiceConfig"][ - "voiceName" - ] - == "Kore" - ) + assert generation_config["speechConfig"]["voiceConfig"]["prebuiltVoiceConfig"]["voiceName"] == "Kore" @pytest.mark.parametrize( "model,custom_llm_provider", @@ -296,6 +422,8 @@ class TestGeminiTTSSpeechConfigInRequestBody: ("gemini-2.5-flash-tts", "vertex_ai"), ("gemini-2.5-flash-tts", "gemini"), ("gemini-2.5-flash-preview-tts", "vertex_ai"), + (GEMINI_3_1_FLASH_TTS_MODEL, "vertex_ai"), + (GEMINI_3_1_FLASH_TTS_MODEL, "gemini"), ], ) def test_speechconfig_end_to_end_mapping(self, model, custom_llm_provider): @@ -342,17 +470,54 @@ class TestGeminiTTSSpeechConfigInRequestBody: f"speechConfig was filtered out during transform_request_body() for model={model}, provider={custom_llm_provider}. " "This breaks Gemini TTS - speechConfig must be in GenerationConfig TypedDict." ) - assert ( - generation_config["speechConfig"]["voiceConfig"]["prebuiltVoiceConfig"][ - "voiceName" - ] - == "Puck" - ) + assert generation_config["speechConfig"]["voiceConfig"]["prebuiltVoiceConfig"]["voiceName"] == "Puck" # Also verify responseModalities is present assert "responseModalities" in generation_config assert "AUDIO" in generation_config["responseModalities"] + @pytest.mark.parametrize( + "custom_llm_provider", + [ + "vertex_ai", + "gemini", + ], + ) + def test_multi_speaker_speechconfig_end_to_end_mapping(self, custom_llm_provider): + """Test full pipeline for Gemini 3.1 multi-speaker TTS speechConfig""" + from litellm.llms.vertex_ai.gemini.transformation import ( + _transform_request_body, + ) + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + VertexGeminiConfig, + ) + + config = VertexGeminiConfig() + + mapped_params = config.map_openai_params( + non_default_params={ + "audio": { + "speech_config": MULTI_SPEAKER_SPEECH_CONFIG, + "format": "pcm16", + } + }, + optional_params={}, + model=GEMINI_3_1_FLASH_TTS_MODEL, + drop_params=False, + ) + + request_body = _transform_request_body( + messages=[{"role": "user", "content": "Ryan: Hi.\nKatie: Hello."}], + model=GEMINI_3_1_FLASH_TTS_MODEL, + optional_params=mapped_params, + custom_llm_provider=custom_llm_provider, + litellm_params={}, + cached_content=None, + ) + + generation_config = request_body["generationConfig"] + assert generation_config["speechConfig"] == NORMALIZED_MULTI_SPEAKER_SPEECH_CONFIG + assert generation_config["responseModalities"] == ["AUDIO"] @pytest.mark.parametrize( "model,custom_llm_provider", @@ -395,14 +560,35 @@ class TestGeminiTTSSpeechConfigInRequestBody: generation_config = request_body["generationConfig"] assert generation_config["speechConfig"]["languageCode"] == "pt-BR" - assert ( - generation_config["speechConfig"]["voiceConfig"]["prebuiltVoiceConfig"][ - "voiceName" - ] - == "Puck" - ) + assert generation_config["speechConfig"]["voiceConfig"]["prebuiltVoiceConfig"]["voiceName"] == "Puck" assert "AUDIO" in generation_config["responseModalities"] if __name__ == "__main__": pytest.main([__file__]) + + +@pytest.mark.parametrize("voice_key", ["name", "voiceName", "voice_name", "voice"]) +@pytest.mark.parametrize("wrapper", [None, "speechConfig", "speech_config"]) +@pytest.mark.parametrize("model", ["gemini-2.5-flash-tts", "gemini-2.5-pro-tts", GEMINI_3_1_FLASH_TTS_MODEL]) +def test_speech_bridge_preserves_single_speaker_voice_mapping(voice_key: str, wrapper: str | None, model: str): + handler = SpeechToCompletionBridgeTransformationHandler() + config = {voice_key: "Umbriel", "language_code": "en-US"} + voice = {wrapper: config} if wrapper else config + result = handler.transform_request( + model=model, + input="Hello", + voice=voice, + optional_params={}, + litellm_params={}, + headers={}, + litellm_logging_obj=MagicMock(), + custom_llm_provider="vertex_ai", + ) + assert result["audio"] == { + "speech_config": { + "voiceConfig": {"prebuiltVoiceConfig": {"voiceName": "Umbriel"}}, + "languageCode": "en-US", + }, + "format": "pcm16", + } diff --git a/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_streaming.py b/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_streaming.py index f59c584eb90..2a550b4a701 100644 --- a/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_streaming.py +++ b/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_streaming.py @@ -278,6 +278,46 @@ class TestGetObjectNameLazyParse: assert exc_info.value.message.startswith("Batch input file line 1 could not be converted") +def _tts_batch_row(**body_extra: object) -> bytes: + body: Final = { + "model": "gemini-2.5-flash-preview-tts", + "messages": [{"role": "user", "content": "Say hi"}], + "audio": {"format": "pcm16", "speech_config": {"voice_name": "Kore", "style": "calm"}}, + **body_extra, + } + return json.dumps({"custom_id": "r-0", "method": "POST", "url": "/v1/chat/completions", "body": body}).encode() + + +def _upload_tts_batch_row(row: bytes, litellm_params: dict[str, object]) -> bytes: + request: Final[CreateFileRequest] = {"file": ("batch.jsonl", row + b"\n", "application/jsonl"), "purpose": "batch"} + return _join_upload_body( + VertexAIFilesConfig().transform_create_file_request( + model="", create_file_data=request, optional_params={}, litellm_params=litellm_params + ) + ) + + +class TestBatchRowUnknownVoiceKeys: + def test_unknown_voice_keys_are_a_400_without_drop_params(self) -> None: + with pytest.raises(VertexAIError, match="voice keys style") as exc_info: + _upload_tts_batch_row(_tts_batch_row(), {}) + assert exc_info.value.status_code == 400 + + @pytest.mark.parametrize( + "row, litellm_params", + [(_tts_batch_row(drop_params=True), {}), (_tts_batch_row(), {"drop_params": True})], + ids=["row", "deployment"], + ) + def test_drop_params_on_the_row_or_deployment_drops_unknown_voice_keys( + self, row: bytes, litellm_params: dict[str, object] + ) -> None: + uploaded: Final = json.loads(_upload_tts_batch_row(row, litellm_params)) + + assert uploaded["request"]["generationConfig"]["speechConfig"] == { + "voiceConfig": {"prebuiltVoiceConfig": {"voiceName": "Kore"}} + } + + class TestStreamingLaziness: def test_upload_stream_transforms_rows_only_as_they_are_pulled(self) -> None: cfg: Final = VertexAIFilesConfig() 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 6ea3039edd8..d2b5c8d5bd9 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 @@ -24,8 +24,11 @@ from litellm.llms.gemini.chat.transformation import GoogleAIStudioGeminiConfig from litellm.llms.vertex_ai.common_utils import GEMINI_MAX_ENUM_CHARACTERS, VertexAIError from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + OPENAI_TO_GEMINI_TTS_VOICES, VertexGeminiConfig, _log_gemini_sampling_params_deprecated, + normalize_gemini_speech_config, + reject_unknown_gemini_tts_voice_keys, ) from litellm.types.llms.vertex_ai import GeminiFinishReason, RequestBody, UsageMetadata from litellm.types.utils import ChoiceLogprobs, Usage @@ -104,6 +107,242 @@ GEMINI_3_IMAGE_SIZE_MAPPINGS: Final[tuple[tuple[str, str, str], ...]] = ( ) +@pytest.mark.parametrize( + "speech_config", + ( + {"name": "aLlOy"}, + {"speech_config": {"voice_name": "aLlOy"}}, + {"voice_config": {"prebuilt_voice_config": {"voice_name": "aLlOy"}}}, + ), +) +def test_gemini_structured_speech_config_maps_openai_voice_names(speech_config: Mapping[str, object]) -> None: + assert normalize_gemini_speech_config(speech_config) == { + "voiceConfig": {"prebuiltVoiceConfig": {"voiceName": OPENAI_TO_GEMINI_TTS_VOICES["alloy"]}} + } + + +@pytest.mark.parametrize( + "speaker", + ( + {"speakerAlias": "Ryan", "speakerId": "alloy"}, + {"speaker_alias": "Ryan", "speaker_id": "alloy"}, + ), +) +def test_gemini_speech_config_accepts_cloud_tts_speaker_entries(speaker: Mapping[str, object]) -> None: + speech_config: Final = { + "multi_speaker_voice_config": { + "speaker_voice_configs": [speaker, {"speaker": "Katie", "voice_config": {"voice_name": "Leda"}}] + } + } + + assert normalize_gemini_speech_config(speech_config) == { + "multiSpeakerVoiceConfig": { + "speakerVoiceConfigs": [ + { + "speaker": "Ryan", + "voiceConfig": {"prebuiltVoiceConfig": {"voiceName": OPENAI_TO_GEMINI_TTS_VOICES["alloy"]}}, + }, + {"speaker": "Katie", "voiceConfig": {"voiceName": "Leda"}}, + ] + } + } + + +@pytest.mark.parametrize( + ("voice", "expected"), + ( + ( + { + "languageCode": "en-US", + "modelName": "gemini-2.5-flash-tts", + "multiSpeakerVoiceConfig": {"speakerVoiceConfigs": [{"speakerAlias": "Ryan", "speakerId": "Umbriel"}]}, + }, + { + "languageCode": "en-US", + "multiSpeakerVoiceConfig": { + "speakerVoiceConfigs": [ + {"speaker": "Ryan", "voiceConfig": {"prebuiltVoiceConfig": {"voiceName": "Umbriel"}}} + ] + }, + }, + ), + ( + {"model_name": "gemini-2.5-flash-tts", "voice_config": {"prebuilt_voice_config": {"voice_name": "Leda"}}}, + {"voiceConfig": {"prebuiltVoiceConfig": {"voiceName": "Leda"}}}, + ), + ({"speech_config": {"modelName": "gemini-2.5-flash-tts", "language_code": "fr-FR"}}, {"languageCode": "fr-FR"}), + ), +) +def test_gemini_speech_config_drops_cloud_tts_model_name( + voice: Mapping[str, object], expected: Mapping[str, object] +) -> None: + assert normalize_gemini_speech_config(voice) == expected + + +def test_gemini_speech_config_drops_non_string_keys_without_adding_voice(): + speech_config: Final = { + "speech_config": { + "language_code": "fr-FR", + "voice_config": {"prebuilt_voice_config": {"voice_name": "Leda", 1: "discard"}}, + }, + "format": "pcm16", + } + + assert normalize_gemini_speech_config(speech_config) == { + "languageCode": "fr-FR", + "voiceConfig": {"prebuiltVoiceConfig": {"voiceName": "Leda"}}, + } + + +def test_gemini_speech_config_keeps_only_speech_config_fields(): + speech_config: Final = { + "voice": "Kore", + "metadata": {"source": "studio"}, + "speech_config": {"language_code": "fr-FR", "style": "calm"}, + } + + assert normalize_gemini_speech_config(speech_config) == { + "languageCode": "fr-FR", + "voiceConfig": {"prebuiltVoiceConfig": {"voiceName": "Kore"}}, + } + + +@pytest.mark.parametrize( + "voice, unknown_keys", + [ + ({"name": "Kore", "style": "calm"}, "style"), + ({"speech_config": {"voice": "Kore", "pitch": 2}}, "pitch"), + ({"speechConfig": {"voiceConfig": {}}, "tone": "x", "pitch": 2}, "pitch, tone"), + ({"speechConfig": {"voiceName": "Kore"}, "speech_config": {"style": "calm"}}, "speech_config"), + ], +) +def test_unknown_gemini_tts_voice_keys_are_rejected_without_drop_params(voice, unknown_keys): + with pytest.raises(litellm.UnsupportedParamsError, match=f"voice keys {unknown_keys}\\."): + reject_unknown_gemini_tts_voice_keys("gemini-3.1-flash-tts-preview", voice, False, "vertex_ai") + + reject_unknown_gemini_tts_voice_keys("gemini-3.1-flash-tts-preview", voice, True, "vertex_ai") + + +def test_chat_audio_speech_config_unknown_keys_are_rejected_unless_dropped(): + audio: Final = {"format": "pcm16", "speech_config": {"voice_name": "Kore", "style": "calm"}} + + with pytest.raises(litellm.UnsupportedParamsError, match="voice keys style\\."): + VertexGeminiConfig().map_openai_params( + non_default_params={"audio": audio}, + optional_params={}, + model="gemini-2.5-flash-preview-tts", + drop_params=False, + ) + + dropped: Final = VertexGeminiConfig().map_openai_params( + non_default_params={"audio": audio}, + optional_params={}, + model="gemini-2.5-flash-preview-tts", + drop_params=True, + ) + assert dropped["speechConfig"] == {"voiceConfig": {"prebuiltVoiceConfig": {"voiceName": "Kore"}}} + + voice_only: Final = VertexGeminiConfig().map_openai_params( + non_default_params={"audio": {"voice": "Kore", "format": "pcm16", "style": "calm"}}, + optional_params={}, + model="gemini-2.5-flash-preview-tts", + drop_params=False, + ) + assert voice_only["speechConfig"] == {"voiceConfig": {"prebuiltVoiceConfig": {"voiceName": "Kore"}}} + + +def test_chat_audio_speech_config_unknown_keys_follow_global_drop_params(monkeypatch): + monkeypatch.setattr(litellm, "drop_params", True) + + dropped: Final = VertexGeminiConfig().map_openai_params( + non_default_params={"audio": {"format": "pcm16", "speech_config": {"voice_name": "Kore", "style": "calm"}}}, + optional_params={}, + model="gemini-2.5-flash-preview-tts", + drop_params=False, + ) + + assert dropped["speechConfig"] == {"voiceConfig": {"prebuiltVoiceConfig": {"voiceName": "Kore"}}} + + +def _chat_speech_config(speech_config: Mapping[str, object], drop_params: bool) -> object: + return VertexGeminiConfig().map_openai_params( + non_default_params={"audio": {"format": "pcm16", "speech_config": speech_config}}, + optional_params={}, + model="gemini-2.5-flash-preview-tts", + drop_params=drop_params, + )["speechConfig"] + + +def test_chat_audio_drop_params_strips_unknown_fields_inside_speaker_voices(): + speech_config: Final = { + "multi_speaker_voice_config": { + "speaker_voice_configs": [ + {"speaker": "Joe", "voice_config": {"prebuilt_voice_config": {"voice_name": "Puck", "pitch": 2}}}, + {"speaker_alias": "Jane", "speaker_id": "Kore", "style": "calm"}, + ], + "turn_order": "alternate", + } + } + + assert _chat_speech_config(speech_config, drop_params=True) == { + "multiSpeakerVoiceConfig": { + "speakerVoiceConfigs": [ + {"speaker": "Joe", "voiceConfig": {"prebuiltVoiceConfig": {"voiceName": "Puck"}}}, + {"speaker": "Jane", "voiceConfig": {"prebuiltVoiceConfig": {"voiceName": "Kore"}}}, + ] + } + } + + +GEMINI_VOICE_CONFIG: Final = { + "voice": "Kore", + "prebuiltVoiceConfig": {"voiceName": "Kore"}, + "replicatedVoiceConfig": {"mimeType": "audio/wav", "voiceSampleAudio": "UklGRg=="}, +} + + +@pytest.mark.parametrize( + "speech_config", + [ + {"voiceConfig": GEMINI_VOICE_CONFIG, "languageCode": "fr-FR"}, + { + "voice_config": { + "voice": "Kore", + "prebuilt_voice_config": {"voice_name": "Kore"}, + "replicated_voice_config": {"mime_type": "audio/wav", "voice_sample_audio": "UklGRg=="}, + }, + "language_code": "fr-FR", + }, + ], +) +def test_chat_audio_drop_params_keeps_every_voice_config_field_gemini_accepts(speech_config: dict[str, object]): + assert _chat_speech_config(speech_config, drop_params=True) == { + "voiceConfig": GEMINI_VOICE_CONFIG, + "languageCode": "fr-FR", + } + + +def test_chat_audio_reads_snake_case_wrapper_beside_null_camel_case_wrapper(): + mapped: Final = VertexGeminiConfig().map_openai_params( + non_default_params={"audio": {"format": "pcm16", "speechConfig": None, "speech_config": {"voice_name": "Kore"}}}, + optional_params={}, + model="gemini-2.5-flash-preview-tts", + drop_params=False, + ) + + assert mapped["speechConfig"] == {"voiceConfig": {"prebuiltVoiceConfig": {"voiceName": "Kore"}}} + + +def test_unknown_voice_key_error_stays_bounded_for_large_voice_objects(): + voice: Final = {"name": "Kore", **{f"{index:04d}{'x' * 1000}": index for index in range(2000)}} + + with pytest.raises(litellm.UnsupportedParamsError) as excinfo: + reject_unknown_gemini_tts_voice_keys("gemini-3.1-flash-tts-preview", voice, False, "vertex_ai") + + assert "and 1990 more" in excinfo.value.message + assert len(excinfo.value.message) < 1500 + + def test_top_logprobs(): non_default_params = { "top_logprobs": 2, diff --git a/tests/unit/llms/vertex_ai/text_to_speech/test_transformation.py b/tests/unit/llms/vertex_ai/text_to_speech/test_transformation.py index dc6f35fcfe5..b3d6d5b952f 100644 --- a/tests/unit/llms/vertex_ai/text_to_speech/test_transformation.py +++ b/tests/unit/llms/vertex_ai/text_to_speech/test_transformation.py @@ -1,4 +1,10 @@ import base64 +import io +import math +import struct +import sys +import wave +from pathlib import Path from typing import Final import json from collections.abc import Iterator, Mapping @@ -12,13 +18,61 @@ from pydantic import ValidationError import litellm from litellm.litellm_core_utils.exception_mapping_utils import exception_type +from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import OPENAI_TO_GEMINI_TTS_VOICES from litellm.llms.vertex_ai.text_to_speech.transformation import ( VertexAILyriaTextToSpeechConfig, VertexAITextToSpeechConfig, + _fallback_gemini_tts_audio_duration, ) +from litellm.cost_calculator import response_cost_calculator +from litellm.litellm_core_utils.audio_utils.utils import calculate_request_duration from litellm.types.utils import LlmProviders from litellm.utils import ProviderConfigManager +pytestmark = pytest.mark.usefixtures("local_model_cost_map") + + +@pytest.mark.parametrize( + ("voice", "multi_speaker"), + ( + ("aLlOy", False), + ({"name": "aLlOy"}, False), + ({"voice_config": {"prebuilt_voice_config": {"voice_name": "aLlOy"}}}, False), + ( + { + "multi_speaker_voice_config": { + "speaker_voice_configs": [ + {"speaker": "reader", "voice_config": {"prebuilt_voice_config": {"voice_name": "aLlOy"}}} + ] + } + }, + True, + ), + ), +) +def test_gemini_cloud_tts_maps_openai_voice_names_in_all_voice_forms( + voice: str | dict[str, object], multi_speaker: bool +) -> None: + config: Final = VertexAITextToSpeechConfig() + model: Final = "gemini-3.1-flash-tts-preview" + _, params = config.map_openai_params(model=model, voice=voice, optional_params={"response_format": "mp3"}) + mapped_voice: Final = OPENAI_TO_GEMINI_TTS_VOICES["alloy"] + expected_voice: Final = { + "languageCode": config.DEFAULT_LANGUAGE_CODE, + "modelName": model, + **( + { + "multiSpeakerVoiceConfig": { + "speakerVoiceConfigs": [{"speakerAlias": "reader", "speakerId": mapped_voice}] + } + } + if multi_speaker + else {"name": mapped_voice} + ), + } + + assert params["vertex_voice_dict"] == expected_voice + class TestVertexAITextToSpeechConfig: """Tests for VertexAITextToSpeechConfig transformation""" @@ -48,8 +102,12 @@ class TestVertexAITextToSpeechConfig: assert url == custom_url - @patch.object(VertexAITextToSpeechConfig, "_ensure_access_token") - @patch.object(VertexAITextToSpeechConfig, "_get_token_and_url") + @patch.object( # test-quality-ok: isolates provider credentials while testing request serialization + VertexAITextToSpeechConfig, "_ensure_access_token" + ) + @patch.object( # test-quality-ok: fixes the provider URL at the authentication boundary + VertexAITextToSpeechConfig, "_get_token_and_url" + ) def test_transform_text_to_speech_request_body(self, mock_get_token, mock_ensure_token): """Test that transform_text_to_speech_request generates correct request body""" # Mock authentication @@ -94,6 +152,20 @@ class TestVertexAITextToSpeechConfig: assert "headers" in result assert "Authorization" in result["headers"] + raw_voice_request: Final = config.transform_text_to_speech_request( + model="vertex_ai/chirp", + input="Hello", + voice="en-US-Chirp3-HD-Charon", + optional_params={}, + litellm_params={ + "vertex_credentials": None, + "vertex_project": "test-project", + "vertex_location": "us-central1", + }, + headers={}, + ) + assert raw_voice_request["dict_body"]["voice"]["name"] == "en-US-Chirp3-HD-Charon" + def test_voice_mapping_openai_to_vertex(self): """Test that OpenAI voice names are correctly mapped to Vertex AI voices""" config = VertexAITextToSpeechConfig() @@ -131,6 +203,506 @@ class TestVertexAITextToSpeechConfig: assert voice_str is None assert voice_dict == voice_input + def test_gemini_tts_multi_speaker_voice_mapping(self): + config = VertexAITextToSpeechConfig() + + voice = { + "multi_speaker_voice_config": { + "speaker_voice_configs": [ + { + "speaker": "Ryan", + "voice_config": { + "prebuilt_voice_config": { + "voice_name": "Umbriel", + }, + }, + }, + { + "speaker": "Katie", + "voice_config": { + "prebuilt_voice_config": { + "voice_name": "Leda", + }, + }, + }, + ], + }, + } + + expected_voice = { + "languageCode": "en-US", + "modelName": "gemini-3.1-flash-tts-preview", + "multiSpeakerVoiceConfig": { + "speakerVoiceConfigs": [ + { + "speakerAlias": "Ryan", + "speakerId": "Umbriel", + }, + { + "speakerAlias": "Katie", + "speakerId": "Leda", + }, + ], + }, + } + + voice_str, optional_params = config.map_openai_params( + model="gemini-3.1-flash-tts-preview", + optional_params={"response_format": "mp3"}, + voice=voice, + ) + assert voice_str is None + assert optional_params["audioEncoding"] == "MP3" + assert optional_params["vertex_voice_dict"] == expected_voice + + voice_str, optional_params = config.map_openai_params( + model="gemini-3.1-flash-tts-preview", + optional_params={"response_format": "opus"}, + voice=voice, + ) + assert voice_str is None + assert optional_params["audioEncoding"] == "OGG_OPUS" + assert optional_params["vertex_voice_dict"] == expected_voice + + def test_gemini_tts_nested_language_and_voice_alias_reach_cloud_tts(self): + config: Final = VertexAITextToSpeechConfig() + voice: Final = {"speech_config": {"language_code": "fr-FR"}, "voice": "Kore"} + + voice_name, optional_params = config.map_openai_params( + model="gemini-3.1-flash-tts-preview", + optional_params={"response_format": "mp3"}, + voice=voice, + ) + + assert voice_name is None + assert optional_params["vertex_voice_dict"] == { + "languageCode": "fr-FR", + "modelName": "gemini-3.1-flash-tts-preview", + "name": "Kore", + } + + @pytest.mark.parametrize( + ("voice", "expected_voice"), + [ + ( + { + "multi_speaker_voice_config": { + "speaker_voice_configs": [ + {"speaker": "Ryan", "voice_config": {"prebuilt_voice_config": {"voice_name": "Umbriel"}}}, + {"speaker": "Katie", "voice_config": {"prebuilt_voice_config": {"voice_name": "Leda"}}}, + ], + }, + "speech_config": {"language_code": "fr-FR"}, + }, + { + "languageCode": "fr-FR", + "modelName": "gemini-3.1-flash-tts-preview", + "multiSpeakerVoiceConfig": { + "speakerVoiceConfigs": [ + {"speakerAlias": "Ryan", "speakerId": "Umbriel"}, + {"speakerAlias": "Katie", "speakerId": "Leda"}, + ], + }, + }, + ), + ( + { + "voice_config": {"prebuilt_voice_config": {"voice_name": "Puck"}}, + "speechConfig": {"languageCode": "en-GB"}, + }, + {"languageCode": "en-GB", "modelName": "gemini-3.1-flash-tts-preview", "name": "Puck"}, + ), + ], + ) + def test_gemini_tts_voice_beside_speech_config_wrapper_reaches_cloud_tts( + self, voice: dict[str, object], expected_voice: dict[str, object] + ): + config: Final = VertexAITextToSpeechConfig() + + voice_name, optional_params = config.map_openai_params( + model="gemini-3.1-flash-tts-preview", + optional_params={"response_format": "mp3"}, + voice=voice, + ) + + assert voice_name is None + assert optional_params["vertex_voice_dict"] == expected_voice + + @pytest.mark.parametrize( + ("voice", "expected_language", "expected_voices"), + [ + ( + { + "multiSpeakerVoiceConfig": { + "speakerVoiceConfigs": [ + {"speaker": "Ryan", "voiceConfig": {"prebuiltVoiceConfig": {"voiceName": "Umbriel"}}}, + {"speaker": "Katie", "voiceConfig": {"prebuiltVoiceConfig": {"voiceName": "Leda"}}}, + ], + }, + "speech_config": { + "multi_speaker_voice_config": { + "speaker_voice_configs": [ + {"speaker": "Ana", "voice_config": {"prebuilt_voice_config": {"voice_name": "Kore"}}}, + {"speaker": "Ben", "voice_config": {"prebuilt_voice_config": {"voice_name": "Puck"}}}, + ], + }, + }, + }, + None, + (("Ana", "Kore"), ("Ben", "Puck")), + ), + ( + {"name": "Puck", "languageCode": "fr-FR", "speech_config": {"language_code": "de-DE"}}, + "de-DE", + ((None, "Puck"),), + ), + ( + {"name": "Puck", "voice_config": {"prebuilt_voice_config": {"voice_name": "Charon"}}}, + None, + ((None, "Charon"),), + ), + ], + ) + def test_gemini_tts_cloud_tts_speaks_the_voice_wav_sends_to_gemini( + self, + voice: dict[str, object], + expected_language: str | None, + expected_voices: tuple[tuple[str | None, str], ...], + ): + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import normalize_gemini_speech_config + + _, optional_params = VertexAITextToSpeechConfig().map_openai_params( + model="gemini-3.1-flash-tts-preview", + optional_params={"response_format": "mp3"}, + voice=voice, + ) + cloud_tts_voice: Final = optional_params["vertex_voice_dict"] + cloud_tts_speakers: Final = cloud_tts_voice.get("multiSpeakerVoiceConfig", {}).get("speakerVoiceConfigs") + cloud_tts_voices: Final = ( + tuple((speaker["speakerAlias"], speaker["speakerId"]) for speaker in cloud_tts_speakers) + if cloud_tts_speakers is not None + else ((None, cloud_tts_voice["name"]),) + ) + wav_config: Final = normalize_gemini_speech_config(voice, drop_unknown_fields=True) + wav_speakers: Final = wav_config.get("multiSpeakerVoiceConfig", {}).get("speakerVoiceConfigs") + wav_voices: Final = ( + tuple( + (speaker["speaker"], speaker["voiceConfig"]["prebuiltVoiceConfig"]["voiceName"]) + for speaker in wav_speakers + ) + if wav_speakers is not None + else ((None, wav_config["voiceConfig"]["prebuiltVoiceConfig"]["voiceName"]),) + ) + + assert cloud_tts_voices == wav_voices == expected_voices + assert cloud_tts_voice["languageCode"] == (expected_language or "en-US") + assert wav_config.get("languageCode") == expected_language + + def test_gemini_tts_string_voice_reaches_cloud_tts(self): + config: Final = VertexAITextToSpeechConfig() + + voice_name, optional_params = config.map_openai_params( + model="gemini-3.1-flash-tts-preview", + optional_params={"response_format": "mp3"}, + voice="Kore", + ) + + assert voice_name == "Kore" + assert optional_params["vertex_voice_dict"] == { + "languageCode": "en-US", + "modelName": "gemini-3.1-flash-tts-preview", + "name": "Kore", + } + + @pytest.mark.parametrize( + "speaker_voice_configs", + ( + [None, {"speakerAlias": "Katie", "speakerId": "Leda"}], + [{"speakerAlias": "Ryan"}, {"speakerAlias": "Katie", "speakerId": "Leda"}], + [ + {"speaker": "Ryan", "voice_config": {"prebuilt_voice_config": {"name": "Umbriel"}}}, + {"speakerAlias": "Katie", "speakerId": "Leda"}, + ], + "invalid", + ), + ) + def test_gemini_tts_malformed_speaker_entry_is_a_bad_request(self, speaker_voice_configs: object): + config: Final = VertexAITextToSpeechConfig() + + with pytest.raises(litellm.UnsupportedParamsError, match="speaker_voice_configs"): + config.map_openai_params( + model="gemini-3.1-flash-tts-preview", + optional_params={"response_format": "mp3"}, + voice={"multiSpeakerVoiceConfig": {"speakerVoiceConfigs": speaker_voice_configs}}, + ) + + def test_gemini_tts_malformed_speaker_entry_is_dropped_with_drop_params(self): + config: Final = VertexAITextToSpeechConfig() + + _, optional_params = config.map_openai_params( + model="gemini-3.1-flash-tts-preview", + optional_params={"response_format": "mp3"}, + voice={ + "multiSpeakerVoiceConfig": { + "speakerVoiceConfigs": [ + {"speakerAlias": "Ryan", "speakerId": "Umbriel"}, + {"speakerAlias": "Katie"}, + ] + } + }, + drop_params=True, + ) + + assert optional_params["vertex_voice_dict"]["multiSpeakerVoiceConfig"] == { + "speakerVoiceConfigs": [{"speakerAlias": "Ryan", "speakerId": "Umbriel"}] + } + + @pytest.mark.parametrize( + "voice", + [ + {"speechConfig": {"voice": "Kore"}}, + {"speechConfig": {"voiceConfig": {"prebuiltVoiceConfig": {"voiceName": "Kore"}}}}, + {"speechConfig": {"voiceConfig": {}}}, + {"speech_config": {}}, + {"multi_speaker_voice_config": {"speaker_voice_configs": []}}, + {"languageCode": "en-US"}, + ], + ) + def test_gemini_tts_nested_voice_mapping_sends_only_cloud_tts_voice_fields(self, voice: dict[str, object]): + config: Final = VertexAITextToSpeechConfig() + + _, optional_params = config.map_openai_params( + model="gemini-3.1-flash-tts-preview", + optional_params={"response_format": "mp3"}, + voice=voice, + ) + + assert optional_params["vertex_voice_dict"] == { + "languageCode": "en-US", + "modelName": "gemini-3.1-flash-tts-preview", + "name": "Kore", + } + + @pytest.mark.parametrize( + ("voice", "unsendable_field"), + [ + ( + {"voice_config": {"replicated_voice_config": {"mime_type": "audio/wav", "voice_sample_audio": "UklG"}}}, + "voiceConfig.replicatedVoiceConfig", + ), + ( + {"voiceConfig": {"prebuiltVoiceConfig": {"voiceName": "Puck", "style": "calm"}}}, + "voiceConfig.prebuiltVoiceConfig.style", + ), + ({"speechConfig": {"voiceConfig": {"voice": "Puck"}}}, "voiceConfig.voice"), + ( + { + "multiSpeakerVoiceConfig": { + "speakerVoiceConfigs": [ + {"speaker": "Ryan", "voiceConfig": {"prebuiltVoiceConfig": {"voiceName": "Puck"}}}, + {"speaker": "Katie", "voiceConfig": {"prebuiltVoiceConfig": {"voiceName": "Leda"}}, "tone": "x"}, + ] + } + }, + "multiSpeakerVoiceConfig.speakerVoiceConfigs.tone", + ), + ( + { + "multiSpeakerVoiceConfig": { + "speakerVoiceConfigs": [ + {"speaker": "Ryan", "voiceConfig": {"prebuiltVoiceConfig": {"voiceName": "Puck"}}, "tone": "x"}, + {"speaker": "Katie", "voiceConfig": {"prebuiltVoiceConfig": {"voiceName": "Leda"}}, "tone": "y"}, + ] + } + }, + "multiSpeakerVoiceConfig.speakerVoiceConfigs.tone", + ), + ( + { + "multiSpeakerVoiceConfig": { + "speakerVoiceConfigs": [ + { + "speakerAlias": "Ryan", + "speakerId": "Puck", + "voiceConfig": { + "replicatedVoiceConfig": {"mimeType": "audio/wav", "voiceSampleAudio": "AAAA"} + }, + }, + {"speakerAlias": "Katie", "speakerId": "Kore"}, + ] + } + }, + "multiSpeakerVoiceConfig.speakerVoiceConfigs.voiceConfig.replicatedVoiceConfig", + ), + ({"languageCode": ["en-US"], "voiceConfig": {"prebuiltVoiceConfig": {"voiceName": "Puck"}}}, "languageCode"), + ({"voiceConfig": {"prebuiltVoiceConfig": {"voiceName": 5}}}, "voiceConfig.prebuiltVoiceConfig.voiceName"), + ({"voiceConfig": "Puck"}, "voiceConfig"), + ({"multiSpeakerVoiceConfig": [{"speaker": "Ryan"}]}, "multiSpeakerVoiceConfig"), + ({"name": 5}, "name"), + ({"speechConfig": {"voice": ["Puck"]}}, "speechConfig.voice"), + ({"speechConfig": "Puck"}, "speechConfig"), + ({"speech_config": [{"voice_config": {"prebuilt_voice_config": {"voice_name": "Puck"}}}]}, "speech_config"), + ( + { + "multiSpeakerVoiceConfig": { + "speakerVoiceConfigs": [ + { + "speakerAlias": "Ryan", + "speakerId": "Umbriel", + "voiceConfig": {"prebuiltVoiceConfig": {"voiceName": "Puck"}}, + }, + {"speakerAlias": "Katie", "speakerId": "Kore"}, + ] + } + }, + "multiSpeakerVoiceConfig.speakerVoiceConfigs.voiceConfig.prebuiltVoiceConfig", + ), + ], + ) + def test_gemini_tts_cloud_tts_rejects_voice_fields_it_cannot_send( + self, voice: dict[str, object], unsendable_field: str + ): + config: Final = VertexAITextToSpeechConfig() + + with pytest.raises(litellm.UnsupportedParamsError) as exc_info: + config.map_openai_params( + model="gemini-3.1-flash-tts-preview", + optional_params={"response_format": "mp3"}, + voice=voice, + ) + + assert f"voice fields {unsendable_field} as mp3" in str(exc_info.value) + + @pytest.mark.parametrize( + ("voice", "voice_name"), + [ + ({"voiceConfig": {"replicatedVoiceConfig": {"mimeType": "audio/wav", "voiceSampleAudio": "UklG"}}}, "Kore"), + ({"voiceConfig": {"prebuiltVoiceConfig": {"voiceName": "Puck", "style": "calm"}}}, "Puck"), + ({"voiceConfig": "Puck"}, "Kore"), + ({"speechConfig": "Puck"}, "Kore"), + ], + ) + def test_gemini_tts_cloud_tts_drops_voice_fields_it_cannot_send_with_drop_params( + self, voice: dict[str, object], voice_name: str + ): + config: Final = VertexAITextToSpeechConfig() + + _, optional_params = config.map_openai_params( + model="gemini-3.1-flash-tts-preview", + optional_params={"response_format": "mp3"}, + voice=voice, + drop_params=True, + ) + + assert optional_params["vertex_voice_dict"] == { + "languageCode": "en-US", + "modelName": "gemini-3.1-flash-tts-preview", + "name": voice_name, + } + + @pytest.mark.parametrize( + "voice", + [ + {"name": "Umbriel", "modelName": "gemini-2.5-flash-tts"}, + {"name": "Umbriel", "model_name": "gemini-2.5-flash-tts"}, + {"modelName": "gemini-2.5-flash-tts", "model_name": "chirp-3"}, + { + "modelName": "gemini-2.5-flash-tts", + "model_name": "chirp-3", + "multi_speaker_voice_config": { + "speaker_voice_configs": [ + { + "speaker": "Ryan", + "voice_config": { + "prebuilt_voice_config": { + "voice_name": "Umbriel", + }, + }, + }, + ], + }, + }, + ], + ) + def test_gemini_tts_ignores_voice_model_name_override(self, voice): + config = VertexAITextToSpeechConfig() + routed_model = "gemini-3.1-flash-tts-preview" + + _, optional_params = config.map_openai_params( + model=routed_model, + optional_params={"response_format": "mp3"}, + voice=voice, + ) + assert optional_params["vertex_voice_dict"]["modelName"] == routed_model + + def test_dispatch_maps_gemini_cloud_tts_params_when_provider_config_skipped(self): + config = VertexAITextToSpeechConfig() + handler = MagicMock() + handler.text_to_speech_handler.return_value = "ok" + + result = config.dispatch_text_to_speech( + model="gemini-3.1-flash-tts-preview", + input="Hi", + voice={ + "name": "Umbriel", + "modelName": "gemini-2.5-flash-tts", + }, + optional_params={"response_format": "mp3"}, + litellm_params_dict={}, + logging_obj=MagicMock(), + timeout=10, + extra_headers=None, + base_llm_http_handler=handler, + aspeech=False, + api_base=None, + api_key=None, + ) + + assert result == "ok" + mapped_params = handler.text_to_speech_handler.call_args.kwargs["text_to_speech_optional_params"] + assert mapped_params["audioEncoding"] == "MP3" + assert mapped_params["vertex_voice_dict"]["name"] == "Umbriel" + assert mapped_params["vertex_voice_dict"]["modelName"] == "gemini-3.1-flash-tts-preview" + + @pytest.mark.parametrize( + ("voice", "expected_name"), + [ + ("en-US-Chirp3-HD-Charon", "en-US-Chirp3-HD-Charon"), + ({"name": "en-US-Chirp3-HD-Charon"}, "en-US-Chirp3-HD-Charon"), + (None, None), + ], + ) + def test_dispatch_keeps_pre_mapped_cloud_tts_params( + self, voice: str | dict[str, str] | None, expected_name: str | None + ): + config = VertexAITextToSpeechConfig() + handler = MagicMock() + handler.text_to_speech_handler.return_value = "ok" + optional_params = { + "audioEncoding": "OGG_OPUS", + "vertex_voice_dict": {"languageCode": "en-US", "name": "en-US-Chirp3-HD-Charon"}, + } + + config.dispatch_text_to_speech( + model="chirp", + input="Hi", + voice=voice, + optional_params=optional_params, + litellm_params_dict={}, + logging_obj=MagicMock(), + timeout=10, + extra_headers=None, + base_llm_http_handler=handler, + aspeech=False, + api_base=None, + api_key=None, + ) + + call_kwargs = handler.text_to_speech_handler.call_args.kwargs + assert call_kwargs["voice"] == expected_name + assert call_kwargs["text_to_speech_optional_params"] is optional_params + assert call_kwargs["text_to_speech_optional_params"]["audioEncoding"] == "OGG_OPUS" + @pytest.mark.parametrize( ("audio", "expected_content_type"), @@ -174,6 +746,150 @@ def test_transform_text_to_speech_response_leaves_unknown_bytes_unlabeled(): assert result.response.content == raw_pcm +@pytest.mark.parametrize("encoding", ["LINEAR16", "PCM", "MP3", "OGG_OPUS", "ALAW", "MULAW"]) +def test_gemini_cloud_tts_response_bills_text_and_audio(encoding: str, monkeypatch: pytest.MonkeyPatch) -> None: + model: Final = "gemini-3.1-flash-tts-preview" + input_text: Final = "Hello from Gemini text to speech" + pcm_audio: Final = b"\x00\x00" * 24000 + wav_buffer: Final = io.BytesIO() + with wave.open(wav_buffer, "wb") as wav_file: + wav_file.setnchannels(1) + wav_file.setsampwidth(2) + wav_file.setframerate(24000) + wav_file.writeframes(pcm_audio) + g711_audio: Final = ( + struct.pack( + "<4sI4s4sIHHIIHH4sI", + b"RIFF", + 36 + 24000, + b"WAVE", + b"fmt ", + 16, + 6 if encoding == "ALAW" else 7, + 1, + 24000, + 24000, + 1, + 8, + b"data", + 24000, + ) + + b"\x00" * 24000 + ) + fixture_dir: Final = Path(__file__).resolve().parents[4] / "audio_tests" + audio_by_encoding: Final = { + "PCM": pcm_audio, + "LINEAR16": wav_buffer.getvalue(), + "MP3": (fixture_dir / "gemini_tts_speech.mp3").read_bytes(), + "OGG_OPUS": (fixture_dir / "gemini_tts_speech.ogg").read_bytes(), + "ALAW": g711_audio, + "MULAW": g711_audio, + } + audio_bytes: Final = audio_by_encoding[encoding] + logger: Final = MagicMock() + sent_request: Final = httpx.Request( + "POST", + _SYNTHESIZE_URL, + json={ + "input": {"text": input_text}, + "audioConfig": {"audioEncoding": encoding, "sampleRateHertz": 24000}, + }, + ) + raw_response: Final = httpx.Response( + 200, json={"audioContent": base64.b64encode(audio_bytes).decode()}, request=sent_request + ) + + result: Final = VertexAITextToSpeechConfig().transform_text_to_speech_response(model, raw_response, logger) + usage: Final = result.usage + assert usage is not None + duration: Final = len(audio_bytes) / 48000 if encoding == "PCM" else calculate_request_duration(audio_bytes) + assert duration is not None + assert _fallback_gemini_tts_audio_duration(audio_bytes, encoding, 24000) == pytest.approx(duration), ( + "Google Cloud AudioEncoding and RFC 7845, checked 2026-09-23: " + "https://cloud.google.com/text-to-speech/docs/reference/rest/v1/AudioEncoding " + "https://www.rfc-editor.org/rfc/rfc7845.html" + ) + assert usage.completion_tokens == math.ceil(duration * 25), ( + "Google Cloud TTS pricing, 2026-09-23: https://cloud.google.com/text-to-speech/pricing" + ) + assert usage.prompt_tokens == litellm.token_counter(model=model, text=input_text) + assert usage.total_tokens == usage.prompt_tokens + usage.completion_tokens + assert usage.completion_tokens_details is not None + assert usage.completion_tokens_details.audio_tokens == usage.completion_tokens + with monkeypatch.context() as dependencies: + dependencies.setitem(sys.modules, "soundfile", None) + sdk_result: Final = VertexAITextToSpeechConfig().transform_text_to_speech_response(model, raw_response, logger) + assert sdk_result.usage == usage + + model_info: Final = litellm.get_model_info(model, custom_llm_provider="vertex_ai") + cost: Final = response_cost_calculator( + response_object=result, + model=model, + custom_llm_provider="vertex_ai", + call_type="speech", + optional_params={}, + prompt=input_text, + ) + expected_cost: Final = ( + usage.prompt_tokens * model_info["input_cost_per_token"] + + usage.completion_tokens * model_info["output_cost_per_audio_token"] + ) + assert cost == pytest.approx(expected_cost) + assert cost > usage.prompt_tokens * model_info["input_cost_per_token"] + + +@pytest.mark.parametrize("encoding", ["ALAW", "MULAW"]) +@pytest.mark.parametrize("audio", [b"\x12" * 24000, b"RIFF" + b"\x12" * 23996]) +def test_gemini_cloud_tts_returns_raw_g711_audio_when_duration_decoder_is_unavailable(encoding: str, audio: bytes): + logger: Final = MagicMock() + sent_request: Final = httpx.Request( + "POST", + _SYNTHESIZE_URL, + json={ + "input": {"text": "Hello"}, + "audioConfig": {"audioEncoding": encoding, "sampleRateHertz": 24000}, + }, + ) + raw_response: Final = httpx.Response( + 200, json={"audioContent": base64.b64encode(audio).decode()}, request=sent_request + ) + + with patch("litellm.llms.vertex_ai.text_to_speech.transformation.calculate_request_duration", return_value=None): + result: Final = VertexAITextToSpeechConfig().transform_text_to_speech_response( + "gemini-3.1-flash-tts-preview", raw_response, logger + ) + + assert result.response.content == audio + assert result.usage is not None + assert result.usage.prompt_tokens > 0 + assert result.usage.completion_tokens == 25, ( + "Google Cloud TTS pricing, 2026-09-23: https://cloud.google.com/text-to-speech/pricing" + ) + + +def test_gemini_cloud_tts_returns_audio_unbilled_when_duration_is_unmeasurable(): + logger: Final = MagicMock() + sent_request: Final = httpx.Request( + "POST", + _SYNTHESIZE_URL, + json={ + "input": {"text": "Hello"}, + "audioConfig": {"audioEncoding": "MP3", "sampleRateHertz": 24000}, + }, + ) + raw_response: Final = httpx.Response( + 200, json={"audioContent": base64.b64encode(b"invalid").decode()}, request=sent_request + ) + + with patch("litellm.llms.vertex_ai.text_to_speech.transformation.calculate_request_duration", return_value=None): + result: Final = VertexAITextToSpeechConfig().transform_text_to_speech_response( + "gemini-3.1-flash-tts-preview", raw_response, logger + ) + + assert result.response.content == b"invalid" + assert result.usage is None + + class TestVertexAILyriaTextToSpeechConfig: @pytest.mark.parametrize( "model", @@ -929,3 +1645,516 @@ def test_lyria_malformed_item_position_does_not_change_the_mapped_error(payload: custom_llm_provider="vertex_ai", original_exception=exc_info.value, ) + + +_PCM_AUDIO: Final = b"\x00\x00\x01\x00" * 2400 +_GENERATE_CONTENT_RESPONSE: Final = { + "candidates": [ + { + "content": { + "role": "model", + "parts": [ + { + "inlineData": { + "mimeType": "audio/L16;codec=pcm;rate=24000", + "data": base64.b64encode(_PCM_AUDIO).decode(), + } + } + ], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": {"promptTokenCount": 5, "candidatesTokenCount": 10, "totalTokenCount": 15}, +} + + +@pytest.mark.parametrize( + ("model", "response_format"), + ( + ("gemini-3.8-flash-tts", None), + ("gemini-3.1-flash-tts-preview", None), + ("gemini-3.1-flash-tts-preview", "wav"), + ("gemini-3.1-flash-tts-preview", "pcm"), + ), +) +@pytest.mark.asyncio +async def test_aspeech_vertex_ai_gemini_tts_formats_the_bridge_serves_stay_on_generate_content( + respx_mock: respx.MockRouter, + google_token_endpoint: responses.RequestsMock, + model: str, + response_format: str | None, +) -> None: + route: Final = respx_mock.post( + url__regex=rf"https://aiplatform\.googleapis\.com/.*/models/{model}:generateContent$" + ).mock(return_value=httpx.Response(200, json=_GENERATE_CONTENT_RESPONSE)) + + response: Final = await litellm.aspeech( + model=f"vertex_ai/{model}", + input="Say hello", + voice="Kore", + vertex_credentials=_AUTHORIZED_USER, + vertex_location="global", + **({} if response_format is None else {"response_format": response_format}), + ) + + assert route.call_count == 1 + assert _PCM_AUDIO in response.content + + +@pytest.mark.asyncio +async def test_aspeech_vertex_ai_gemini_tts_mp3_reaches_cloud_tts_for_any_gemini_tts_model( + respx_mock: respx.MockRouter, google_token_endpoint: responses.RequestsMock +) -> None: + audio_bytes: Final = (Path(__file__).resolve().parents[4] / "audio_tests/gemini_tts_speech.mp3").read_bytes() + route: Final = respx_mock.post(_SYNTHESIZE_URL).mock( + return_value=httpx.Response(200, json={"audioContent": base64.b64encode(audio_bytes).decode()}) + ) + + response: Final = await litellm.aspeech( + model="vertex_ai/gemini-3.8-flash-tts", + input="Say hello", + voice="Kore", + response_format="mp3", + vertex_credentials=_AUTHORIZED_USER, + vertex_location="global", + ) + + assert response.content == audio_bytes + sent: Final = json.loads(route.calls[0].request.content) + assert sent["voice"]["modelName"] == "gemini-3.8-flash-tts" + assert sent["voice"]["name"] == "Kore" + assert sent["audioConfig"]["audioEncoding"] == "MP3" + + +@pytest.mark.asyncio +async def test_aspeech_vertex_ai_gemini_tts_mp3_without_voice_keeps_the_gemini_model( + respx_mock: respx.MockRouter, google_token_endpoint: responses.RequestsMock +) -> None: + audio_bytes: Final = (Path(__file__).resolve().parents[4] / "audio_tests/gemini_tts_speech.mp3").read_bytes() + route: Final = respx_mock.post(_SYNTHESIZE_URL).mock( + return_value=httpx.Response(200, json={"audioContent": base64.b64encode(audio_bytes).decode()}) + ) + + await litellm.aspeech( + model="vertex_ai/gemini-3.1-flash-tts-preview", + input="Say hello", + response_format="mp3", + vertex_credentials=_AUTHORIZED_USER, + vertex_location="global", + ) + + sent: Final = json.loads(route.calls[0].request.content) + assert sent["voice"] == {"languageCode": "en-US", "modelName": "gemini-3.1-flash-tts-preview", "name": "Kore"} + + +@pytest.mark.parametrize("voice", ("en-US-Studio-O", "en-US-Neural2-F")) +@pytest.mark.asyncio +async def test_aspeech_vertex_ai_gemini_tts_mp3_rejects_a_cloud_tts_voice( + voice: str, respx_mock: respx.MockRouter, google_token_endpoint: responses.RequestsMock +) -> None: + with pytest.raises(litellm.BadRequestError, match=f"{voice} is a Cloud Text-to-Speech voice"): + await litellm.aspeech( + model="vertex_ai/gemini-3.1-flash-tts-preview", + input="Say hello", + voice=voice, + response_format="mp3", + vertex_credentials=_AUTHORIZED_USER, + vertex_location="global", + ) + + assert not respx_mock.calls + + +@pytest.mark.asyncio +async def test_aspeech_vertex_ai_gemini_tts_mp3_drop_params_swaps_a_cloud_tts_voice_for_a_gemini_voice( + respx_mock: respx.MockRouter, google_token_endpoint: responses.RequestsMock +) -> None: + audio_bytes: Final = (Path(__file__).resolve().parents[4] / "audio_tests/gemini_tts_speech.mp3").read_bytes() + route: Final = respx_mock.post(_SYNTHESIZE_URL).mock( + return_value=httpx.Response(200, json={"audioContent": base64.b64encode(audio_bytes).decode()}) + ) + + await litellm.aspeech( + model="vertex_ai/gemini-3.1-flash-tts-preview", + input="Say hello", + voice="en-US-Studio-O", + response_format="mp3", + drop_params=True, + vertex_credentials=_AUTHORIZED_USER, + vertex_location="global", + ) + + assert json.loads(route.calls[0].request.content)["voice"]["name"] == "Kore" + + +@pytest.mark.asyncio +async def test_aspeech_vertex_ai_gemini_tts_mp3_without_a_location_uses_the_vertex_default_region( + respx_mock: respx.MockRouter, google_token_endpoint: responses.RequestsMock, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.delenv("VERTEXAI_LOCATION", raising=False) + monkeypatch.delenv("VERTEX_LOCATION", raising=False) + audio_bytes: Final = (Path(__file__).resolve().parents[4] / "audio_tests/gemini_tts_speech.mp3").read_bytes() + route: Final = respx_mock.post("https://us-texttospeech.googleapis.com/v1/text:synthesize").mock( + return_value=httpx.Response(200, json={"audioContent": base64.b64encode(audio_bytes).decode()}) + ) + + await litellm.aspeech( + model="vertex_ai/gemini-3.1-flash-tts-preview", + input="Say hello", + voice="Kore", + response_format="mp3", + vertex_credentials=_AUTHORIZED_USER, + ) + + assert route.called + + +@pytest.mark.asyncio +async def test_aspeech_vertex_ai_gemini_tts_mp3_sends_instructions_as_the_style_prompt( + respx_mock: respx.MockRouter, google_token_endpoint: responses.RequestsMock +) -> None: + audio_bytes: Final = (Path(__file__).resolve().parents[4] / "audio_tests/gemini_tts_speech.mp3").read_bytes() + route: Final = respx_mock.post(_SYNTHESIZE_URL).mock( + return_value=httpx.Response(200, json={"audioContent": base64.b64encode(audio_bytes).decode()}) + ) + + await litellm.aspeech( + model="vertex_ai/gemini-3.1-flash-tts-preview", + input="Have a wonderful day", + voice="Kore", + response_format="mp3", + instructions="Read this slowly and calmly", + vertex_credentials=_AUTHORIZED_USER, + vertex_location="global", + ) + + sent: Final = json.loads(route.calls[0].request.content) + assert sent["input"] == {"text": "Have a wonderful day", "prompt": "Read this slowly and calmly"} + + +@pytest.mark.asyncio +async def test_aspeech_vertex_ai_gemini_tts_mp3_with_a_numeric_speaking_rate_returns_billed_audio( + respx_mock: respx.MockRouter, google_token_endpoint: responses.RequestsMock +) -> None: + audio_bytes: Final = (Path(__file__).resolve().parents[4] / "audio_tests/gemini_tts_speech.mp3").read_bytes() + route: Final = respx_mock.post(_SYNTHESIZE_URL).mock( + return_value=httpx.Response(200, json={"audioContent": base64.b64encode(audio_bytes).decode()}) + ) + + result: Final = await litellm.aspeech( + model="vertex_ai/gemini-3.1-flash-tts-preview", + input="Hello there", + voice="Kore", + response_format="mp3", + audioConfig={"audioEncoding": "MP3", "speakingRate": 1.25}, + vertex_credentials=_AUTHORIZED_USER, + vertex_location="global", + ) + + assert json.loads(route.calls[0].request.content)["audioConfig"]["speakingRate"] == 1.25 + assert result.content == audio_bytes + assert result.usage is not None + assert result.usage.completion_tokens > 0 + + +@pytest.mark.parametrize( + ("vertex_location", "synthesize_url"), + ( + ("us-central1", "https://us-texttospeech.googleapis.com/v1/text:synthesize"), + ("us-east5", "https://us-texttospeech.googleapis.com/v1/text:synthesize"), + ("us", "https://us-texttospeech.googleapis.com/v1/text:synthesize"), + ("europe-west4", "https://eu-texttospeech.googleapis.com/v1/text:synthesize"), + ("europe-north1", "https://eu-texttospeech.googleapis.com/v1/text:synthesize"), + ("eu", "https://eu-texttospeech.googleapis.com/v1/text:synthesize"), + ("northamerica-northeast1", "https://northamerica-northeast1-texttospeech.googleapis.com/v1/text:synthesize"), + ("global", _SYNTHESIZE_URL), + ), +) +@pytest.mark.asyncio +async def test_aspeech_vertex_ai_gemini_tts_mp3_keeps_the_deployment_region( + respx_mock: respx.MockRouter, + google_token_endpoint: responses.RequestsMock, + vertex_location: str, + synthesize_url: str, +) -> None: + audio_bytes: Final = (Path(__file__).resolve().parents[4] / "audio_tests/gemini_tts_speech.mp3").read_bytes() + route: Final = respx_mock.post(synthesize_url).mock( + return_value=httpx.Response(200, json={"audioContent": base64.b64encode(audio_bytes).decode()}) + ) + + response: Final = await litellm.aspeech( + model="vertex_ai/gemini-3.1-flash-tts-preview", + input="Say hello", + voice="Kore", + response_format="mp3", + vertex_credentials=_AUTHORIZED_USER, + vertex_location=vertex_location, + ) + + assert route.call_count == 1 + assert response.content == audio_bytes + + +@pytest.mark.parametrize("vertex_location", ["asia-northeast1", "europe-west2", "australia-southeast1"]) +@pytest.mark.asyncio +async def test_aspeech_vertex_ai_gemini_tts_mp3_rejects_a_region_cloud_tts_does_not_serve( + respx_mock: respx.MockRouter, google_token_endpoint: responses.RequestsMock, vertex_location: str +) -> None: + with pytest.raises(litellm.BadRequestError, match=f"no Gemini TTS location for vertex_location {vertex_location}"): + await litellm.aspeech( + model="vertex_ai/gemini-3.1-flash-tts-preview", + input="Say hello", + voice="Kore", + response_format="mp3", + vertex_credentials=_AUTHORIZED_USER, + vertex_location=vertex_location, + ) + + assert not respx_mock.calls + + +@pytest.mark.asyncio +async def test_router_gemini_tts_mp3_unknown_voice_is_a_bad_request_that_keeps_the_deployment_serving( + respx_mock: respx.MockRouter, google_token_endpoint: responses.RequestsMock +) -> None: + audio_bytes: Final = (Path(__file__).resolve().parents[4] / "audio_tests/gemini_tts_speech.mp3").read_bytes() + respx_mock.post(_SYNTHESIZE_URL).mock( + side_effect=lambda request: ( + httpx.Response( + 404, json={"error": {"code": 404, "message": "Voice Koree not found.", "status": "NOT_FOUND"}} + ) + if json.loads(request.content)["voice"]["name"] == "Koree" + else httpx.Response(200, json={"audioContent": base64.b64encode(audio_bytes).decode()}) + ) + ) + router: Final = litellm.Router( + model_list=[ + { + "model_name": "gemini-tts", + "litellm_params": { + "model": "vertex_ai/gemini-3.1-flash-tts-preview", + "vertex_credentials": _AUTHORIZED_USER, + "vertex_location": "global", + }, + } + ], + num_retries=0, + ) + + with pytest.raises(litellm.BadRequestError, match="Voice Koree not found"): + await router.aspeech(model="gemini-tts", input="Say hello", voice="Koree", response_format="mp3") + response: Final = await router.aspeech(model="gemini-tts", input="Say hello", voice="Kore", response_format="mp3") + + assert response.content == audio_bytes + + +@pytest.mark.asyncio +async def test_aspeech_vertex_ai_chirp_unknown_voice_stays_not_found( + respx_mock: respx.MockRouter, google_token_endpoint: responses.RequestsMock +) -> None: + respx_mock.post(_SYNTHESIZE_URL).mock( + return_value=httpx.Response( + 404, json={"error": {"code": 404, "message": "Voice en-US-Nope not found.", "status": "NOT_FOUND"}} + ) + ) + + with pytest.raises(litellm.NotFoundError, match="Voice en-US-Nope not found"): + await litellm.aspeech( + model="vertex_ai/chirp", + input="Say hello", + voice={"languageCode": "en-US", "name": "en-US-Nope"}, + vertex_credentials=_AUTHORIZED_USER, + ) + + +@pytest.mark.asyncio +async def test_aspeech_vertex_ai_gemini_tts_mp3_rejects_a_deployment_api_base( + respx_mock: respx.MockRouter, google_token_endpoint: responses.RequestsMock +) -> None: + with pytest.raises(litellm.BadRequestError, match="Request wav or pcm"): + await litellm.aspeech( + model="vertex_ai/gemini-3.1-flash-tts-preview", + input="Say hello", + voice="Kore", + response_format="mp3", + vertex_credentials=_AUTHORIZED_USER, + vertex_location="us-central1", + api_base="https://vertex.gateway.internal", + ) + + +@pytest.mark.asyncio +async def test_aspeech_vertex_ai_gemini_tts_mp3_sends_speak_tags_as_text( + respx_mock: respx.MockRouter, google_token_endpoint: responses.RequestsMock +) -> None: + audio_bytes: Final = (Path(__file__).resolve().parents[4] / "audio_tests/gemini_tts_speech.mp3").read_bytes() + route: Final = respx_mock.post(_SYNTHESIZE_URL).mock( + return_value=httpx.Response(200, json={"audioContent": base64.b64encode(audio_bytes).decode()}) + ) + speech_input: Final = "Have a wonderful day" + + await litellm.aspeech( + model="vertex_ai/gemini-3.1-flash-tts-preview", + input=speech_input, + voice="Kore", + response_format="mp3", + vertex_credentials=_AUTHORIZED_USER, + vertex_location="global", + ) + + sent: Final = json.loads(route.calls[0].request.content) + assert sent["input"] == {"text": speech_input} + + +@pytest.mark.parametrize("response_format", ("aac", "flac", "pcm16", "WAV")) +@pytest.mark.asyncio +async def test_aspeech_vertex_ai_gemini_tts_unsupported_format_is_rejected_before_calling_google( + respx_mock: respx.MockRouter, google_token_endpoint: responses.RequestsMock, response_format: str +) -> None: + with pytest.raises(litellm.BadRequestError, match="Supported response formats: alaw, mp3, mulaw, opus, pcm, wav"): + await litellm.aspeech( + model="vertex_ai/gemini-3.1-flash-tts-preview", + input="Say hello", + voice="Kore", + response_format=response_format, + vertex_credentials=_AUTHORIZED_USER, + vertex_location="global", + ) + + +@pytest.mark.asyncio +async def test_aspeech_vertex_ai_gemini_tts_mp3_sends_both_speakers_to_cloud_tts( + respx_mock: respx.MockRouter, google_token_endpoint: responses.RequestsMock +) -> None: + audio_bytes: Final = (Path(__file__).resolve().parents[4] / "audio_tests/gemini_tts_speech.mp3").read_bytes() + route: Final = respx_mock.post(_SYNTHESIZE_URL).mock( + return_value=httpx.Response(200, json={"audioContent": base64.b64encode(audio_bytes).decode()}) + ) + + result: Final = await litellm.aspeech( + model="vertex_ai/gemini-3.1-flash-tts-preview", + input="Ryan: Hi.\nKatie: Hello.", + voice={ + "multi_speaker_voice_config": { + "speaker_voice_configs": [ + {"speaker": "Ryan", "voice_config": {"prebuilt_voice_config": {"voice_name": "Umbriel"}}}, + {"speaker": "Katie", "voice_config": {"prebuilt_voice_config": {"voice_name": "Leda"}}}, + ], + }, + }, + response_format="mp3", + vertex_credentials=_AUTHORIZED_USER, + vertex_location="global", + ) + + assert result.content == audio_bytes + assert result.usage is not None + assert result.usage.prompt_tokens > 0 + assert result.usage.completion_tokens > 0 + assert route.call_count == 1 + sent: Final = json.loads(route.calls[0].request.content) + assert sent["audioConfig"]["audioEncoding"] == "MP3" + assert sent["voice"] == { + "languageCode": "en-US", + "modelName": "gemini-3.1-flash-tts-preview", + "multiSpeakerVoiceConfig": { + "speakerVoiceConfigs": [ + {"speakerAlias": "Ryan", "speakerId": "Umbriel"}, + {"speakerAlias": "Katie", "speakerId": "Leda"}, + ], + }, + } + + +@pytest.mark.parametrize( + ("speech_params", "error"), + ( + ({"voice": ["Kore"]}, "needs voice to be a Gemini voice name"), + ({"voice": ""}, "needs voice to be a Gemini voice name"), + ({"voice": {"name": "Kore", "style": "calm"}}, "does not support the voice keys style"), + ({"voice": {"speechConfig": {"voice": "Kore", "pitch": 2}}}, "does not support the voice keys pitch"), + ({"voice": "Kore", "instructions": ["Read slowly"]}, "needs instructions to be a string"), + ), +) +@pytest.mark.asyncio +async def test_aspeech_vertex_ai_gemini_tts_mp3_malformed_params_are_bad_requests_unless_dropped( + respx_mock: respx.MockRouter, + google_token_endpoint: responses.RequestsMock, + speech_params: dict[str, object], + error: str, +) -> None: + audio_bytes: Final = (Path(__file__).resolve().parents[4] / "audio_tests/gemini_tts_speech.mp3").read_bytes() + route: Final = respx_mock.post(_SYNTHESIZE_URL).mock( + return_value=httpx.Response(200, json={"audioContent": base64.b64encode(audio_bytes).decode()}) + ) + request: Final = { + "model": "vertex_ai/gemini-3.1-flash-tts-preview", + "input": "Say hello", + "response_format": "mp3", + "vertex_credentials": _AUTHORIZED_USER, + "vertex_location": "global", + **speech_params, + } + + with pytest.raises(litellm.BadRequestError, match=error): + await litellm.aspeech(**request) + assert not route.called + + response: Final = await litellm.aspeech(**request, drop_params=True) + + sent: Final = json.loads(route.calls[0].request.content) + assert response.content == audio_bytes + assert sent["input"] == {"text": "Say hello"} + assert sent["voice"] == {"languageCode": "en-US", "modelName": "gemini-3.1-flash-tts-preview", "name": "Kore"} + + +@pytest.mark.asyncio +async def test_aspeech_vertex_ai_beta_lyria_2_serves_the_openai_params_it_ignores( + respx_mock: respx.MockRouter, google_token_endpoint: responses.RequestsMock +) -> None: + route: Final = respx_mock.post(url__regex=r"https://us-central1-aiplatform\.googleapis\.com/.*/lyria-002:predict$").mock( + return_value=httpx.Response(200, json={"predictions": [{"audioContent": "bHlyaWEtMi1hdWRpbw==", "mimeType": "audio/wav"}]}) + ) + + response: Final = await litellm.aspeech( + model="vertex_ai_beta/lyria-002", + input="A bright synth track", + voice="alloy", + response_format="mp3", + speed=1.5, + vertex_credentials=_AUTHORIZED_USER, + vertex_project="test-project", + vertex_location="us-central1", + ) + + assert route.call_count == 1 + assert response.content == b"lyria-2-audio" + + +@pytest.mark.asyncio +async def test_aspeech_vertex_ai_gemini_tts_wav_unknown_voice_keys_are_bad_requests_unless_dropped( + respx_mock: respx.MockRouter, google_token_endpoint: responses.RequestsMock +) -> None: + route: Final = respx_mock.post( + url__regex=r"https://aiplatform\.googleapis\.com/.*/models/gemini-3.1-flash-tts-preview:generateContent$" + ).mock(return_value=httpx.Response(200, json=_GENERATE_CONTENT_RESPONSE)) + request: Final = { + "model": "vertex_ai/gemini-3.1-flash-tts-preview", + "input": "Say hello", + "voice": {"speech_config": {"voice": "Kore", "style": "calm"}, "modelName": "gemini-3.1-flash-tts-preview"}, + "response_format": "wav", + "vertex_credentials": _AUTHORIZED_USER, + "vertex_location": "global", + } + + with pytest.raises(litellm.BadRequestError, match="does not support the voice keys style"): + await litellm.aspeech(**request) + assert not route.called + + response: Final = await litellm.aspeech(**request, drop_params=True) + + assert _PCM_AUDIO in response.content + assert json.loads(route.calls[0].request.content)["generationConfig"]["speechConfig"] == { + "voiceConfig": {"prebuiltVoiceConfig": {"voiceName": "Kore"}} + } diff --git a/tests/unit/proxy/openai_files_endpoint/test_files_endpoint.py b/tests/unit/proxy/openai_files_endpoint/test_files_endpoint.py index eb222366812..d47db69257b 100644 --- a/tests/unit/proxy/openai_files_endpoint/test_files_endpoint.py +++ b/tests/unit/proxy/openai_files_endpoint/test_files_endpoint.py @@ -722,6 +722,16 @@ def _router_with_executed_batch_model() -> Router: "litellm_params": {"model": "gemini/gemini-2.0-flash"}, "model_info": {"id": "gemini-2.0-flash-id"}, }, + { + "model_name": "vertex-batch-drop", + "litellm_params": {"model": "vertex_ai/gemini-2.5-flash", "drop_params": True}, + "model_info": {"id": "vertex-batch-drop-id"}, + }, + { + "model_name": "vertex-batch-keep", + "litellm_params": {"model": "vertex_ai/gemini-2.5-flash", "drop_params": False}, + "model_info": {"id": "vertex-batch-keep-id"}, + }, ] ) @@ -939,6 +949,39 @@ def test_batch_upload_for_a_provider_model_still_goes_to_the_provider(batch_uplo assert provider_upload.call_args.kwargs["custom_llm_provider"] == "gemini" +@pytest.mark.parametrize( + "model, expected_drop_params", + [("vertex-batch-drop", True), ("vertex-batch-keep", False)], +) +def test_batch_upload_for_a_provider_model_carries_the_deployment_drop_params( + batch_upload_seams, model: str, expected_drop_params: bool +): + _, provider_upload, _ = batch_upload_seams + + response = _upload_batch_file({"x-litellm-model": model}, {}) + + assert response.status_code == 200, response.text + assert provider_upload.call_args.kwargs["drop_params"] is expected_drop_params + + +def test_batch_upload_for_a_provider_model_without_deployment_drop_params_sends_no_drop_params(batch_upload_seams): + _, provider_upload, _ = batch_upload_seams + + response = _upload_batch_file({"x-litellm-model": "gemini-2.0-flash"}, {}) + + assert response.status_code == 200, response.text + assert "drop_params" not in provider_upload.call_args.kwargs + + +def test_batch_upload_keeps_the_callers_drop_params_over_the_deployment(batch_upload_seams): + _, provider_upload, _ = batch_upload_seams + + response = _upload_batch_file({"x-litellm-model": "vertex-batch-keep", "user-agent": "claude-cli/2.1.0"}, {}) + + assert response.status_code == 200, response.text + assert provider_upload.call_args.kwargs["drop_params"] is True + + @pytest.mark.skip(reason="mock respx fails on ci/cd - unclear why") def test_create_file_and_call_chat_completion_e2e( mocker: MockerFixture, monkeypatch, llm_router: Router diff --git a/tests/unit/test_cost_calculator.py b/tests/unit/test_cost_calculator.py index 2854f9010b2..5d4eb5fae12 100644 --- a/tests/unit/test_cost_calculator.py +++ b/tests/unit/test_cost_calculator.py @@ -31,7 +31,7 @@ from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.llms.fireworks_ai.cost_calculator import get_base_model_for_pricing from litellm.llms.together_ai.cost_calculator import get_model_params_and_category from litellm.types.llms.base import CachedTokensDetails -from litellm.types.llms.openai import OpenAIRealtimeStreamList, ResponseAPIUsage, ResponsesAPIResponse +from litellm.types.llms.openai import HttpxBinaryResponseContent, OpenAIRealtimeStreamList, ResponseAPIUsage, ResponsesAPIResponse from litellm.types.llms.xai import RealtimeMeteredUsage from litellm.types.rerank import RerankResponse from litellm.types.utils import CallTypes, Choices, CompletionTokensDetailsWrapper, EmbeddingResponse, ImageObject, ImageResponse, ImageUsage, ImageUsageInputTokensDetails, LiteLLMRealtimeStreamLoggingObject, Message, ModelInfo, ModelResponse, PromptTokensDetails, PromptTokensDetailsWrapper, Usage @@ -9134,6 +9134,33 @@ def test_azure_tts_cost_uses_test_owned_character_rate(monkeypatch: pytest.Monke assert cost == pytest.approx(0.05) +def test_speech_cost_with_custom_pricing_bills_the_reported_usage() -> None: + usage: Final = Usage( + prompt_tokens=10, + completion_tokens=1000, + total_tokens=1010, + completion_tokens_details=CompletionTokensDetailsWrapper(audio_tokens=1000), + ) + response: Final = HttpxBinaryResponseContent( + httpx.Response(200, content=b"RIFF....WAVE", headers={"content-type": "audio/wav"}) + ) + response.usage = usage + rates: Final = MappingProxyType({"input_cost_per_token": 1e-6, "output_cost_per_token": 2e-5}) + + cost: Final = completion_cost( + completion_response=response, + model="vertex_ai/gemini-3.1-flash-tts-preview", + call_type="speech", + prompt="Say hello", + custom_cost_per_token=dict(rates), + ) + + assert cost == pytest.approx( + usage.prompt_tokens * rates["input_cost_per_token"] + + usage.completion_tokens * rates["output_cost_per_token"] + ) + + def test_moderation_cost_uses_test_owned_token_rate() -> None: prompt: Final = "moderation input" rates: Final = MappingProxyType( From 1fcbb0355e76065c5ed0e79aba298c8ece25033a Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Tue, 20 Oct 2026 21:31:19 -0700 Subject: [PATCH 02/14] fix(pricing): update baseten reflection/Beam-1.0 prices from the models API (#47784) Price-Sync: litellm-providers Co-authored-by: berriai-litellm-provider-info-sync[bot] <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 6 +++--- model_prices_and_context_window.json | 6 +++--- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index fd72d209ffb..764cf397158 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -76709,14 +76709,14 @@ "supports_response_schema": true }, "baseten/reflection/Beam-1.0": { - "cache_read_input_token_cost": 1.3e-07, - "input_cost_per_token": 1.3e-06, + "cache_read_input_token_cost": 4e-08, + "input_cost_per_token": 4e-07, "litellm_provider": "baseten", "max_input_tokens": 1048576, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 4.3e-06, + "output_cost_per_token": 1.2e-06, "source": "https://inference.baseten.co/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index fd72d209ffb..764cf397158 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -76709,14 +76709,14 @@ "supports_response_schema": true }, "baseten/reflection/Beam-1.0": { - "cache_read_input_token_cost": 1.3e-07, - "input_cost_per_token": 1.3e-06, + "cache_read_input_token_cost": 4e-08, + "input_cost_per_token": 4e-07, "litellm_provider": "baseten", "max_input_tokens": 1048576, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 4.3e-06, + "output_cost_per_token": 1.2e-06, "source": "https://inference.baseten.co/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, From 820725661e865983859d73bca9b9b712706c0371 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 21 Oct 2026 04:46:45 +0000 Subject: [PATCH 03/14] fix(router): log each dropped /v1/messages attempt through its own logging copy and open the retry without its usage (#45194) * fix(proxy): keep a retried /v1/messages stream's spend log by logging the superseded attempt as a failure A /v1/messages stream the router retries before any content shares one logging object across both attempts. The superseded attempt's end-of-stream logging ran as a success (a usage-only or placeholder result) and claimed the request's one-shot success log, so the recovered stream's spend row was either missing or written under the superseded attempt's message id. The passthrough logging path now logs an Anthropic stream whose frames carry an `event: error` as a failure with whatever usage the provider reported, the way a dropped stream already does, and the router closes the superseded stream before it opens the retry so that logging is queued first. * fix(proxy): log a provider error frame through an attempt-scoped logging copy The error-frame failure stashed the superseded attempt's usage on the logging object the retry shares, so the retry's success log carried the failed attempt's token counts. The failure now logs through a copy of the logging object with its own model_call_details, and the helper moves into the Anthropic passthrough logging handler so the proxy streaming handler only routes by endpoint type. * fix(proxy): scope the error-frame failure log to streams closed at the frame A stream that ran to its end at the provider's error frame kept logging through the success route before this PR, and the attempt-scoped failure log is now taken only for a stream the router closed at that frame, with the copy made at the close so the retry's deployment never reaches it * fix(router): log each dropped /v1/messages attempt through its own logging copy and open the retry without its usage A /v1/messages stream dropped after message_start and retried in its model group shares one Logging object across attempts, so the retry's success log inherited the dropped attempt's stashed usage and exception, and a second drop never reached the failure handlers because the first had already marked the object as logged. The dropped attempt now dispatches its failure handlers through an attempt-scoped copy, and the router discards the stashed usage right before it opens the next attempt, on same-group retries and on every fallback hop, so a request that runs out of attempts still bills the last attempt's partial usage. * fix(router): dispatch chat stream failures through the attempt-scoped logging copy The router now clears a dropped attempt's stashed usage before every fallback hop, so the chat stream wrapper's queued failure log must read that usage from its own copy of the logging object, like the /v1/messages path already does. The stash test asserts through a real success log, the cooldown test reads a real Router's cooldown list, and the integration helpers carry full types. * test(vertex): give the stream wrapper's logging stub an attempt_scoped_copy * test: route the pass-through logging mock's attempt copy to itself and read the callback sink through a reader The pass-through failure test stubbed dispatch_failure_handlers on a spec'd MagicMock, which the attempt-scoped copy no longer reached, so the logging worker never finished. The integration tests kept earlier sink batches in a caller-owned list the helper mutated; a small reader keeps them instead. * test(router): use real Logging in fallback cooldown test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(router): preserve superseded retry usage for billing Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): bill retry input after prior stream usage Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): track retry dispatch for terminal billing Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(router): use deterministic retry timestamp Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(router): avoid real clock in fallback tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(tests): remove duplicate streaming logger import Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): retain error-frame usage for exhausted retries * test(proxy): fix retry billing expectations and helper annotations * fix(proxy): preserve failure estimate without carried usage Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(proxy): audit retry helpers and fixtures Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): fix retry fixture timestamps Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): observe retry stash at attempt boundary Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(logging): keep terminal failure attribution when dropped attempts log on copies Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(logging): publish a dropped attempt's payload only while it is still the current attempt Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): cover spend-log encoding error propagation Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(anthropic): skip adapter failure dispatch when the wrapped stream already logged it Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(router): cool down the final fallback attempt when its failure callbacks were blocked Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(anthropic): import Final in adapter regression Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): drop spend-log encoding expectation superseded by main's poison isolation Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(anthropic): finish the wrapped stream's failure log before surfacing the error Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(anthropic): wait for the wrapped stream's failure log to finish like the inline dispatch did Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(anthropic): wait for the executor-run sync failure callback separately from the async caller boundary Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): pin pass-through stream fixture start times Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(logging): publish failed attempt payload before awaiting failure callbacks Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(logging): publish failed attempt payload from the failure handler after baseline invalidation Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(logging): drop unused discard_partial_usage_for_failure 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/litellm_core_utils/litellm_logging.py | 46 +- .../litellm_core_utils/streaming_handler.py | 33 +- .../adapters/streaming_iterator.py | 26 +- .../anthropic_passthrough_logging_handler.py | 22 - .../streaming_handler.py | 50 +- litellm/proxy/utils.py | 30 +- litellm/router.py | 2 + .../router_utils/fallback_event_handlers.py | 31 +- ...thropic_messages_pre_content_retry_wire.py | 163 +++- .../test_litellm_logging.py | 6 + .../test_max_streaming_duration.py | 11 +- .../test_streaming_handler.py | 66 +- ...est_streaming_iterator_mid_stream_error.py | 81 +- ...test_vertex_and_google_ai_studio_gemini.py | 8 + ...t_anthropic_passthrough_logging_handler.py | 51 -- .../test_streaming_handler.py | 4 + .../test_streaming_handler_interrupt.py | 145 +++- ..._utils_model_creation_and_error_logging.py | 726 +++++++++++++++++- .../mcp/test_aresponses_api_with_mcp.py | 19 +- .../test_fallback_event_handlers.py | 233 +++++- tests/unit/test_router/test_router.py | 52 +- 21 files changed, 1631 insertions(+), 174 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index a59672b5cb2..d75d35e4531 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -604,6 +604,10 @@ def _timestamp_seconds(moment: object) -> float | None: return None +def same_attempt_marker(first: object, second: object) -> bool: + return first is second or first == second + + class Logging(LiteLLMLoggingBaseClass): global \ supabaseClient, \ @@ -625,6 +629,7 @@ class Logging(LiteLLMLoggingBaseClass): classifier_input: Mapping[str, JsonValue] | None = None baseline_cache_context: "BaselineCacheContext | None" = None baseline_observation: "CapturedBaselineObservation | None" = None + _on_failure_payload_built: "Callable[[], None] | None" = None @property def _defer_async_logging(self) -> bool: @@ -852,6 +857,7 @@ class Logging(LiteLLMLoggingBaseClass): "litellm_call_id": litellm_call_id, "start_time": start_time, "first_api_call_start_time": None, + "current_attempt_api_call_start_time": None, "litellm_params": litellm_params, } new_call.streaming_chunks = [] @@ -1602,7 +1608,9 @@ class Logging(LiteLLMLoggingBaseClass): # A proxy-gate failure (auth, rate limit, client disconnect) fires pre_call for its callbacks # without any provider handoff, so it must not stamp the handoff instants. if LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL not in self.model_call_details: - self.model_call_details["api_call_start_time"] = datetime.datetime.now() + api_call_start_time: Final = datetime.datetime.now() + self.model_call_details["api_call_start_time"] = api_call_start_time + self.model_call_details["current_attempt_api_call_start_time"] = api_call_start_time # Set-once first provider-handoff instant. api_call_start_time # is overwritten on every retry, so it can't measure one-time # preprocessing; pinning the first attempt excludes retry loops @@ -1610,9 +1618,7 @@ class Logging(LiteLLMLoggingBaseClass): # litellm_params["metadata"] (caller request metadata, typed # Dict[str, str], echoed downstream; a datetime breaks it). if self.model_call_details.get("first_api_call_start_time") is None: - self.model_call_details["first_api_call_start_time"] = self.model_call_details[ - "api_call_start_time" - ] + self.model_call_details["first_api_call_start_time"] = api_call_start_time # Input Integration Logging -> If you want to log the fact that an attempt to call the model was made callbacks: Final = ( *self._without_replaced_global_callbacks(litellm.input_callback), @@ -2474,12 +2480,37 @@ class Logging(LiteLLMLoggingBaseClass): attempt.model_call_details = dict(self.model_call_details) return attempt + def dispatch_attempt_failure_handlers( + self, exception: Exception, traceback_exception: str + ) -> Coroutine[None, None, None]: + shared_payload: Final = self.model_call_details.get("standard_logging_object") + attempt_marker: Final = self.model_call_details.get("current_attempt_api_call_start_time") + attempt: Final = self.attempt_scoped_copy() + + def publish() -> None: + attempt_payload: Final = attempt.model_call_details.get("standard_logging_object") + current_marker: Final = self.model_call_details.get("current_attempt_api_call_start_time") + if ( + attempt_payload is not None + and self.model_call_details.get("standard_logging_object") is shared_payload + and same_attempt_marker(current_marker, attempt_marker) + ): + self.model_call_details["standard_logging_object"] = attempt_payload + + attempt._on_failure_payload_built = publish + + async def dispatch() -> None: + await attempt.dispatch_failure_handlers(exception, traceback_exception, prefer_async_handlers=True) + + return dispatch() + def carry_failed_attempt_usage(self) -> None: """Move what a failed earlier attempt stashed off the retry sharing this logging object, so the retry logs its own usage while the request's failure log still bills every attempt.""" usage: Final = self.model_call_details.pop("combined_usage_object", None) response_cost: Final = self.model_call_details.pop("response_cost", None) if isinstance(usage, Usage): + self.model_call_details["current_attempt_api_call_start_time"] = None self.model_call_details["failed_attempts_usage"] = ( *self.model_call_details.get("failed_attempts_usage", ()), (usage, float(response_cost) if isinstance(response_cost, (int, float)) else 0.0), @@ -2593,6 +2624,10 @@ class Logging(LiteLLMLoggingBaseClass): """ return self.model_call_details[f"has_logged_{event_type}"] = True + if event_type == "async_failure": + self.model_call_details["async_failure_logged_attempt"] = self.model_call_details.get( + "current_attempt_api_call_start_time" + ) return def has_run_logging( @@ -4410,7 +4445,8 @@ class Logging(LiteLLMLoggingBaseClass): start_time=start_time, end_time=end_time, ) - + if self._on_failure_payload_built is not None: + self._on_failure_payload_built() callbacks: Final = self.get_combined_callback_list( dynamic_success_callbacks=self.dynamic_async_failure_callbacks, global_callbacks=litellm._async_failure_callback, diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 5e43157a55d..f994baa52b1 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -2283,12 +2283,7 @@ class CustomStreamWrapper: traceback_exception = traceback.format_exc() ## ADD DEBUG INFORMATION - E.G. LITELLM REQUEST TIMEOUT traceback_exception += f"\nLiteLLM Default Request Timeout - {litellm.request_timeout}" - if self.logging_obj is not None: - self._record_partial_usage_for_failure() - ## LOGGING - self._failure_logging = asyncio.create_task( - self.logging_obj.dispatch_failure_handlers(e, traceback_exception, prefer_async_handlers=True) - ) + self._dispatch_stream_failure(e, traceback_exception) await self.aclose() self._handle_stream_fallback_error(e) except (httpx.ReadError, httpx.RemoteProtocolError) as e: @@ -2389,16 +2384,20 @@ class CustomStreamWrapper: return processed_chunk def _log_stream_failure_and_raise(self, e: Exception) -> NoReturn: - traceback_exception: Final = traceback.format_exc() - if self.logging_obj is not None: - self._record_partial_usage_for_failure() - ## LOGGING - self._failure_logging = asyncio.create_task( - self.logging_obj.dispatch_failure_handlers(e, traceback_exception, prefer_async_handlers=True) - ) + self._dispatch_stream_failure(e, traceback.format_exc()) self._handle_stream_fallback_error(e) - async def wait_for_failure_logging(self, timeout: float) -> None: + def _dispatch_stream_failure(self, e: Exception, traceback_exception: str) -> None: + self._record_partial_usage_for_failure() + self._failure_logging = asyncio.create_task( + self.logging_obj.dispatch_attempt_failure_handlers(e, traceback_exception) + ) + + @property + def has_dispatched_failure_logging(self) -> bool: + return self._failure_logging is not None + + async def wait_for_failure_logging(self, timeout: float | None) -> 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 @@ -2412,9 +2411,9 @@ class CustomStreamWrapper: A stream that breaks mid-flight still billed the provider for the chunks already delivered. Recover that partial usage from the chunks seen so far and stash it, with its cost, on the logging object so the failure - handler records the real partial spend instead of zero. A request that - later recovers via a router fallback overwrites this with the combined - success log on the same request id, so this never double counts. + handler records the real partial spend instead of zero. The router + moves it into terminal failure usage before opening the next attempt, + so a recovered attempt logs its own usage. """ if self.logging_obj is None or not self.chunks: return diff --git a/litellm/llms/anthropic/pass_through/adapters/streaming_iterator.py b/litellm/llms/anthropic/pass_through/adapters/streaming_iterator.py index 52b836c789a..e9fd77ef441 100644 --- a/litellm/llms/anthropic/pass_through/adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/pass_through/adapters/streaming_iterator.py @@ -16,6 +16,7 @@ from typing import ( Protocol, cast, get_args, + runtime_checkable, ) from pydantic import TypeAdapter @@ -65,6 +66,14 @@ class _ChunkStream(Protocol): def __aiter__(self) -> "AsyncIterator[ModelResponseStream]": ... +@runtime_checkable +class _FailureLoggingWaiter(Protocol): + @property + def has_dispatched_failure_logging(self) -> bool: ... + + async def wait_for_failure_logging(self, timeout: float | None) -> None: ... + + def _optional_attr(obj: object, name: str) -> object: return getattr(obj, name, None) @@ -189,6 +198,14 @@ class _CombinedChunkSplitter: self._held_calls: Mapping[int, tuple[ModelResponseStream, _ToolCallPart]] = MappingProxyType({}) self._live_call_index: int | None = None + @property + def has_dispatched_failure_logging(self) -> bool: + return isinstance(self._stream, _FailureLoggingWaiter) and self._stream.has_dispatched_failure_logging is True + + async def wait_for_failure_logging(self, timeout: float | None) -> None: + if isinstance(self._stream, _FailureLoggingWaiter): + await self._stream.wait_for_failure_logging(timeout=timeout) + @property def chunks(self) -> "list[ModelResponseStream] | None": return cast( # cast-ok: chunks is a list of ModelResponseStream on the inner stream @@ -522,7 +539,9 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): # Wrap the upstream stream so chunks that carry both content and a # finish_reason (fake-streamed providers) are split into two — see # _CombinedChunkSplitter. - super().__init__(_CombinedChunkSplitter(completion_stream)) + combined_chunk_splitter: Final = _CombinedChunkSplitter(completion_stream) + super().__init__(combined_chunk_splitter) + self._combined_chunk_splitter: Final = combined_chunk_splitter self.model = model self._message_id: str = f"msg_{uuid.uuid4()}" if litellm_logging_obj is not None: @@ -1244,7 +1263,10 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): if provider_error is e: raise raise provider_error from e - if logging_obj is not None: + stream_has_dispatched_failure_logging: Final = self._combined_chunk_splitter.has_dispatched_failure_logging + if stream_has_dispatched_failure_logging: + await self._combined_chunk_splitter.wait_for_failure_logging(timeout=None) + elif logging_obj is not None: try: await logging_obj.dispatch_failure_handlers( exception=provider_error, diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py index a7f3a4f84a7..23ef89b3c5c 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py @@ -10,14 +10,12 @@ from pydantic import ConfigDict, TypeAdapter import litellm from litellm._logging import verbose_proxy_logger from litellm.constants import ANTHROPIC_BATCHES_ROUTE -from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.litellm_logging import use_custom_pricing_for_model from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_content_from_model_response, ) from litellm.llms.anthropic import get_anthropic_config -from litellm.llms.anthropic.common_utils import anthropic_error_frame_exception from litellm.llms.anthropic.pass_through.stream_assembly import ( build_complete_streaming_response, build_complete_streaming_response_legacy, @@ -256,26 +254,6 @@ class AnthropicPassthroughLoggingHandler: ), ) - @staticmethod - async def log_error_frame_as_failure( - attempt_logging_obj: LiteLLMLoggingObj, - request_body: Mapping[str, object], - all_chunks: Sequence[str | bytes], - error_event: tuple[str, str, int], - ) -> None: - """A stream closed at the provider's `event: error` frame is one failed attempt of the request: it logs - through the attempt-scoped copy its caller took at the close, so the usage it consumed never reaches - the success log of the attempt the router opens in its place.""" - error_type, message, status_code = error_event - await asyncify(AnthropicPassthroughLoggingHandler.record_partial_usage_for_failure)( - litellm_logging_obj=attempt_logging_obj, request_body=request_body, all_chunks=all_chunks - ) - await attempt_logging_obj.dispatch_failure_handlers( - anthropic_error_frame_exception(error_type, message, status_code, str(attempt_logging_obj.model)), - "", - prefer_async_handlers=True, - ) - @staticmethod def _cost_partial_stream_or_zero( partial_response: ModelResponse | TextCompletionResponse, model: str, logging_obj: LiteLLMLoggingObj diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py index 1fcfe2e279a..b6a067b7666 100644 --- a/litellm/proxy/pass_through_endpoints/streaming_handler.py +++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py @@ -1,5 +1,5 @@ import traceback -from collections.abc import Coroutine, Sequence +from collections.abc import Callable, Coroutine, Mapping, Sequence from dataclasses import dataclass from datetime import datetime, timezone from typing import Final, Protocol @@ -12,6 +12,7 @@ from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.core_helpers import bind_budget_reservation_to_callbacks from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER +from litellm.llms.anthropic.common_utils import anthropic_error_frame_exception from litellm.proxy._types import PassThroughEndpointLoggingResultValues from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing from litellm.proxy.common_utils.sse_keepalive import split_complete_sse_frames @@ -87,13 +88,32 @@ class PassThroughStreamingHandler: ) try: GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue( - async_coroutine=litellm_logging_obj.dispatch_failure_handlers( - exception, traceback.format_exc(), prefer_async_handlers=True - ) + async_coroutine=litellm_logging_obj.dispatch_attempt_failure_handlers(exception, traceback.format_exc()) ) except Exception as e: verbose_proxy_logger.error("Error scheduling stream failure logging: %s", e) + @staticmethod + async def _schedule_stream_end_logging( + litellm_logging_obj: LiteLLMLoggingObj, + endpoint_type: EndpointType, + request_body: Mapping[str, object], + raw_bytes: Sequence[bytes], + error_event: tuple[str, str, int] | None, + success_logging: Callable[[], Coroutine[None, None, None]], + ) -> None: + if error_event is None: + GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=success_logging()) + return + error_type, message, status_code = error_event + await PassThroughStreamingHandler.schedule_stream_failure_logging( + litellm_logging_obj=litellm_logging_obj, + endpoint_type=endpoint_type, + request_body=dict(request_body), + raw_bytes=raw_bytes, + exception=anthropic_error_frame_exception(error_type, message, status_code, str(litellm_logging_obj.model)), + ) + @staticmethod def _record_partial_usage_for_failure( litellm_logging_obj: LiteLLMLoggingObj, @@ -276,17 +296,13 @@ class PassThroughStreamingHandler: else PassThroughStreamingHandler._error_frame_closed_at(endpoint_type, raw_bytes) ) try: - GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue( - async_coroutine=( - _build_logging_coroutine() - if superseded_error_event is None - else AnthropicPassthroughLoggingHandler.log_error_frame_as_failure( - attempt_logging_obj=litellm_logging_obj.attempt_scoped_copy(), - request_body=resolved_request_body, - all_chunks=tuple(raw_bytes), - error_event=superseded_error_event, - ) - ) + await PassThroughStreamingHandler._schedule_stream_end_logging( + litellm_logging_obj=litellm_logging_obj, + endpoint_type=endpoint_type, + request_body=resolved_request_body, + raw_bytes=tuple(raw_bytes), + error_event=superseded_error_event, + success_logging=_build_logging_coroutine, ) except Exception as e: verbose_proxy_logger.error("Error scheduling chunk_processor logging: %s", e) @@ -295,10 +311,6 @@ class PassThroughStreamingHandler: @staticmethod def _error_frame_closed_at(endpoint_type: EndpointType, raw_bytes: Sequence[bytes]) -> tuple[str, str, int] | None: - """The provider `event: error` frame of a stream the router closed because a retry or fallback took the - request over. That attempt is logged as its own failure, on a copy taken now so the retry's rebinding of - the shared logging object never reaches it. A stream that ran to its end or that the client closed keeps - logging through the success route, so its usage is still billed.""" from litellm.llms.anthropic.pass_through.messages.streaming_iterator import ( closing_superseded_attempt, parse_anthropic_error_event, diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 6c9b23860da..2a328971817 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1015,7 +1015,7 @@ def _billed_across_attempts( ) -def _failure_usage_to_lift( +def _current_failure_usage_to_lift( model_call_details: Mapping[str, object], request_body: Mapping[str, object], dispatched: bool, @@ -1037,9 +1037,6 @@ def _failure_usage_to_lift( failure stays at zero cost. Returns the (combined_usage_object, response_cost) pair to lift, or None.""" recovered_usage: Final = model_call_details.get("combined_usage_object") - failed_attempts: Final = _FAILED_ATTEMPTS_USAGE.validate_python(model_call_details.get("failed_attempts_usage", ())) - if failed_attempts: - return _billed_across_attempts(failed_attempts, recovered_usage, model_call_details.get("response_cost")) if recovered_usage is not None: return recovered_usage, model_call_details.get("response_cost") if not dispatched: @@ -1281,6 +1278,29 @@ def _estimated_input_cost( return next((cost for cost in candidate_costs if cost is not None), 0.0) +def _failure_usage_to_lift( + model_call_details: Mapping[str, object], + request_body: Mapping[str, object], + dispatched: bool, + provider_billed: bool = False, + deployment_id: str | None = None, + base_model: str | None = None, +) -> tuple[object, object] | None: + current_attempt: Final = _current_failure_usage_to_lift( + model_call_details, + request_body, + dispatched, + provider_billed, + deployment_id, + base_model, + ) + failed_attempts: Final = _FAILED_ATTEMPTS_USAGE.validate_python(model_call_details.get("failed_attempts_usage", ())) + if not failed_attempts: + return current_attempt + current_usage, current_cost = current_attempt if current_attempt is not None else (None, None) + return _billed_across_attempts(failed_attempts, current_usage, current_cost) + + _EMPTY_LIFT: Final = MappingProxyType({}) @@ -1444,7 +1464,7 @@ def _failure_fields_to_lift(request_data: Mapping[str, object], provider_billed: _usage_to_lift: Final = _failure_usage_to_lift( model_call_details=_model_call_details, request_body=request_data, - dispatched=_first_handoff is not None, + dispatched=_model_call_details.get("current_attempt_api_call_start_time", _first_handoff) is not None, provider_billed=_billed, deployment_id=_deployment_id, base_model=_base_model, diff --git a/litellm/router.py b/litellm/router.py index 0ac21960be1..fe418b5830b 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -260,6 +260,7 @@ from litellm.router_utils.fallback_event_handlers import ( attempted_retries_for_request, carry_over_pre_routing_selection, carry_over_routed_deployment, + carry_superseded_attempt_usage, check_non_standard_fallback_format, clear_pre_routing_selection, committed_retry_budget_for_request, @@ -6721,6 +6722,7 @@ class Router: verbose_router_logger.debug( "Retrying stream that failed before content, attempt %s of %s", attempt + 1, budget ) + carry_superseded_attempt_usage(retry_kwargs) try: response = await attempt_function(**retry_kwargs) except Exception as retry_error: # noqa: BLE001 # the fallback chain judges the failed retry diff --git a/litellm/router_utils/fallback_event_handlers.py b/litellm/router_utils/fallback_event_handlers.py index e1201df1a13..6c93fe6d239 100644 --- a/litellm/router_utils/fallback_event_handlers.py +++ b/litellm/router_utils/fallback_event_handlers.py @@ -8,6 +8,7 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final import httpx +from pydantic import TypeAdapter import litellm from litellm._logging import verbose_router_logger @@ -15,6 +16,7 @@ from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs, safe_deep_copy from litellm.litellm_core_utils.credential_accessor import reject_api_key_secret_reference from litellm.litellm_core_utils.get_llm_provider_logic import inferred_provider +from litellm.litellm_core_utils.litellm_logging import Logging, same_attempt_marker from litellm.litellm_core_utils.sensitive_data_masker import mask_sensitive_structure from litellm.router_utils.add_retry_fallback_headers import ( add_fallback_headers_to_response, @@ -51,6 +53,7 @@ else: # (e.g. a nonexistent file/batch/thread id), independent of the selected deployment's health. _REQUEST_SCOPED_STATUS_CODES: Final = frozenset((404,)) _NO_MODEL_CALL_DETAILS: Final[Mapping[str, object]] = MappingProxyType({}) +_FALLBACK_KWARGS_ADAPTER: Final = TypeAdapter(dict[str, object]) def _trigger_cooldown_for_failed_deployment( @@ -447,6 +450,24 @@ def mid_stream_retry_kwargs( return {key: value for key, value in hop_kwargs.items() if key not in _MID_STREAM_RETRY_STRIPPED_KEYS} +def carry_superseded_attempt_usage(kwargs: Mapping[str, object]) -> None: + """Carry the superseded stream's partial usage into billing if the request ultimately fails.""" + logging_obj: Final = kwargs.get("litellm_logging_obj") + if isinstance(logging_obj, Logging): + logging_obj.carry_failed_attempt_usage() + + +def _failure_callbacks_blocked_logging_object(kwargs: Mapping[str, object]) -> Logging | None: + logging_obj: Final = kwargs.get("litellm_logging_obj") + if not isinstance(logging_obj, Logging): + return None + if logging_obj.model_call_details.get("has_logged_async_failure", False) is not True: + return None + current_marker: Final = logging_obj.model_call_details.get("current_attempt_api_call_start_time") + logged_marker: Final = logging_obj.model_call_details.get("async_failure_logged_attempt") + return logging_obj if not same_attempt_marker(current_marker, logged_marker) else None + + def _request_metadata_bucket(kwargs: Mapping[str, object]) -> Mapping[str, object] | None: bucket: Final = kwargs.get(get_metadata_variable_name_from_kwargs(kwargs)) return bucket if isinstance(bucket, Mapping) else None @@ -869,6 +890,7 @@ async def run_async_fallback( kwargs["attempted_targets"] = attempted if include_fallback_errors: kwargs["include_fallback_errors"] = include_fallback_errors + carry_superseded_attempt_usage(_FALLBACK_KWARGS_ADAPTER.validate_python(kwargs)) response = await litellm_router.async_function_with_fallbacks(*args, **kwargs) verbose_router_logger.info("Successful fallback b/w models.") response = add_fallback_headers_to_response( @@ -891,13 +913,16 @@ async def run_async_fallback( kwargs=kwargs, original_exception=original_exception, ) - logging_obj = kwargs.get("litellm_logging_obj") - if logging_obj is not None and logging_obj.model_call_details.get("has_logged_async_failure", False): + if ( + failure_callbacks_blocked_logging_obj := _failure_callbacks_blocked_logging_object( + _FALLBACK_KWARGS_ADAPTER.validate_python(kwargs) + ) + ) is not None: _trigger_cooldown_for_failed_deployment( litellm_router=litellm_router, kwargs=kwargs, exception=e, - model_call_details=logging_obj.model_call_details, + model_call_details=failure_callbacks_blocked_logging_obj.model_call_details, ) raise error_from_fallbacks diff --git a/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_pre_content_retry_wire.py b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_pre_content_retry_wire.py index e0b2e4a336f..041a0a72fdd 100644 --- a/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_pre_content_retry_wire.py +++ b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_messages_pre_content_retry_wire.py @@ -1,13 +1,17 @@ import json import os import uuid -from collections.abc import Callable, Mapping +from collections.abc import Callable, Iterator, Mapping, Sequence +from contextlib import AbstractContextManager from dataclasses import dataclass +from pathlib import Path from typing import Final, Literal import anthropic import httpx +from pydantic import JsonValue, TypeAdapter import pytest +import yaml from integration._support.anthropic_sse import ( LIFECYCLE, Attempts, @@ -27,13 +31,14 @@ from integration._support.anthropic_sse import ( ) from integration._support.client import Gateway, eventually, object_value from integration._support.database import read_rows +from integration._support.process import owned_proxy from integration._support.wire import Reply, Request, Wire, wire_server -from pydantic import JsonValue from redis import Redis _MODEL: Final = "claude-sonnet-4-5-20250929" _API_KEY: Final = "synthetic-anthropic-key" _TEXT: Final = "Hello" +_SINK_EVENT_BATCH_ADAPTER: Final = TypeAdapter(list[dict[str, JsonValue]]) FirstAttempt = Literal[ "drop_after_headers", @@ -127,24 +132,24 @@ def _stream(gateway: Gateway, body: Mapping[str, JsonValue]) -> tuple[int, tuple return response.status_code, parse_sse(response.text) -def _rows(request_id: str) -> list[dict[str, JsonValue]]: - return read_rows('SELECT status, model_group FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (request_id,)) +def _rows(gateway: Gateway, request_id: str) -> list[dict[str, JsonValue]]: + return gateway.rows('SELECT status, model_group FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (request_id,)) -def _success_rows(request_id: str) -> list[dict[str, JsonValue]]: - return eventually(lambda: _rows(request_id), lambda rows: len(rows) >= 1, seconds=70) +def _success_rows(gateway: Gateway, request_id: str) -> list[dict[str, JsonValue]]: + return eventually(lambda: _rows(gateway, request_id), lambda rows: len(rows) >= 1, seconds=70) def _assert_completed_by_retry( - status: int, events: tuple[SseEvent, ...], prompt: str, model: str, wire: Wire, *, attempts: int = 2 + gateway: Gateway, status: int, events: tuple[SseEvent, ...], prompt: str, model: str, wire: Wire, *, attempts: int = 2 ) -> None: assert status == 200, events assert event_types(events) == LIFECYCLE, events assert message_id(events) == _served_id(prompt, attempts), events assert delta_text(events) == _TEXT, events assert [request.target for request in wire.drain()] == ["/v1/messages"] * attempts - assert _success_rows(_served_id(prompt, attempts)) == [{"status": "success", "model_group": model}] - superseded_rows: Final = {attempt: _rows(_served_id(prompt, attempt)) for attempt in range(1, attempts)} + assert _success_rows(gateway, _served_id(prompt, attempts)) == [{"status": "success", "model_group": model}] + superseded_rows: Final = {attempt: _rows(gateway, _served_id(prompt, attempt)) for attempt in range(1, attempts)} assert not any(superseded_rows.values()), superseded_rows @@ -183,7 +188,7 @@ def test_stream_failing_before_content_is_retried_on_the_same_group_and_complete **({"num_retries": 1} if source == "deployment" else {}), ) status, events = _stream(gateway, _body(model, prompt, source)) - _assert_completed_by_retry(status, events, prompt, model, wire) + _assert_completed_by_retry(gateway, status, events, prompt, model, wire) @pytest.mark.parametrize("source", ["deployment", "request"]) @@ -204,7 +209,7 @@ def test_stream_rejected_before_it_opens_is_retried_and_completes( **({"num_retries": 1} if source == "deployment" else {}), ) status, events = _stream(gateway, _body(model, prompt, source)) - _assert_completed_by_retry(status, events, prompt, model, wire) + _assert_completed_by_retry(gateway, status, events, prompt, model, wire) def _sdk(gateway: Gateway) -> anthropic.Anthropic: @@ -232,7 +237,7 @@ def test_anthropic_sdk_sync_stream_completes_after_a_drop_following_message_star assert final.id == _served_id(prompt, 2), final assert [block.text for block in final.content if block.type == "text"] == [_TEXT], final assert [request.target for request in wire.drain()] == ["/v1/messages"] * 2 - assert _success_rows(final.id) == [{"status": "success", "model_group": model}] + assert _success_rows(gateway, final.id) == [{"status": "success", "model_group": model}] async def test_anthropic_sdk_async_stream_completes_after_a_drop_following_message_start(gateway: Gateway) -> None: @@ -321,7 +326,7 @@ def test_stream_error_frame_after_content_is_not_retried_and_keeps_its_spend_row assert message_id(events) == _served_id(prompt, 1), events assert delta_text(events) == _TEXT, events assert [request.target for request in wire.drain()] == ["/v1/messages"] - assert [row["model_group"] for row in _success_rows(_served_id(prompt, 1))] == [model] + assert [row["model_group"] for row in _success_rows(gateway, _served_id(prompt, 1))] == [model] def test_stream_rejected_with_401_before_it_opens_is_not_retried(gateway: Gateway) -> None: @@ -393,6 +398,22 @@ def test_retry_rejected_before_it_opens_counts_against_the_same_budget(gateway: model: Final = scenario.model(model=f"anthropic/{_MODEL}", api_base=wire.url, api_key=_API_KEY, num_retries=1) response: Final = gateway.request("POST", "/v1/messages", _body(model, prompt, "deployment")) _assert_rejected_before_content(response, wire, status=500, attempts=2, error="error") + rows: Final = eventually( + lambda: _request_rows(gateway, response.headers["x-litellm-call-id"]), + lambda values: len(values) >= 1, + seconds=70, + ) + rejected: Final = gateway.request("POST", "/v1/messages", _body(model, prompt, "request", budget=0)) + _assert_rejected_before_content(rejected, wire, status=500, attempts=1, error="error") + rejection_rows: Final = eventually( + lambda: _request_rows(gateway, rejected.headers["x-litellm-call-id"]), + lambda values: len(values) >= 1, + seconds=70, + ) + estimated_input: Final = rejection_rows[0]["prompt_tokens"] + assert isinstance(estimated_input, int) + assert rejection_rows == [{"status": "failure", "prompt_tokens": estimated_input, "completion_tokens": 0}] + assert rows == [{"status": "failure", "prompt_tokens": 5 + estimated_input, "completion_tokens": 1}], rows def test_non_stream_messages_rejected_with_500_is_retried_per_the_deployment_budget(gateway: Gateway) -> None: @@ -409,7 +430,7 @@ def test_non_stream_messages_rejected_with_500_is_retried_per_the_deployment_bud assert payload["id"] == _served_id(prompt, 2), response.text assert payload["content"] == [{"type": "text", "text": _TEXT}], response.text assert [request.target for request in wire.drain()] == ["/v1/messages"] * 2 - assert _success_rows(_served_id(prompt, 2)) == [{"status": "success", "model_group": model}] + assert _success_rows(gateway, _served_id(prompt, 2)) == [{"status": "success", "model_group": model}] def test_retried_stream_response_headers_name_the_attempt(gateway: Gateway) -> None: @@ -422,7 +443,7 @@ def test_retried_stream_response_headers_name_the_attempt(gateway: Gateway) -> N model: Final = scenario.model(model=f"anthropic/{_MODEL}", api_base=wire.url, api_key=_API_KEY, num_retries=1) response: Final = gateway.request("POST", "/v1/messages", _body(model, prompt, "deployment")) events: Final = parse_sse(response.text) - _assert_completed_by_retry(response.status_code, events, prompt, model, wire) + _assert_completed_by_retry(gateway, response.status_code, events, prompt, model, wire) assert response.headers.get("x-litellm-attempted-retries") == "1", dict(response.headers) assert response.headers.get("x-litellm-max-retries") == "1", dict(response.headers) @@ -454,7 +475,7 @@ def test_retried_stream_spend_row_records_the_attempt_count(gateway: Gateway) -> ): model: Final = scenario.model(model=f"anthropic/{_MODEL}", api_base=wire.url, api_key=_API_KEY, num_retries=1) status, events = _stream(gateway, _body(model, prompt, "deployment")) - _assert_completed_by_retry(status, events, prompt, model, wire) + _assert_completed_by_retry(gateway, status, events, prompt, model, wire) rows: Final = read_rows( "SELECT metadata->>'attempted_retries' AS attempted, metadata->>'max_retries' AS budget " 'FROM "LiteLLM_SpendLogs" WHERE request_id=%s', @@ -463,6 +484,114 @@ def test_retried_stream_spend_row_records_the_attempt_count(gateway: Gateway) -> assert rows == [{"attempted": "1", "budget": "1"}], rows +def _proxy_with_callback_sink(gateway: Gateway, tmp_path: Path, sink: Wire) -> AbstractContextManager[Gateway]: + loaded_config: Final = TypeAdapter(dict[str, object]).validate_python( + yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + ) + litellm_settings: Final = TypeAdapter(dict[str, object]).validate_python(loaded_config["litellm_settings"]) + config: Final = { + **loaded_config, + "litellm_settings": { + **litellm_settings, + "callbacks": ["generic_api"], + "DEFAULT_FLUSH_INTERVAL_SECONDS": 1, + }, + } + path: Final = tmp_path / "callbacks.yaml" + path.write_text(yaml.safe_dump(config)) + return owned_proxy( + gateway, + tmp_path, + {"GENERIC_LOGGER_ENDPOINT": sink.url, "GENERIC_LOGGER_HEADERS": "Authorization=Bearer synthetic-sink-secret"}, + config=path, + ) + + +def _sink_events(batches: Sequence[Request]) -> Iterator[dict[str, JsonValue]]: + for batch in batches: + yield from _SINK_EVENT_BATCH_ADAPTER.validate_python(json.loads(batch.body)) + + +class _SinkReader: + """Reads the callback sink across polls: drain() consumes the sink's queue, so earlier batches are kept here.""" + + def __init__(self, sink: Wire, model: str) -> None: + self._sink: Final = sink + self._model: Final = model + self._batches: tuple[Request, ...] = () + + def delivered_usage(self) -> tuple[tuple[str, int, int], ...]: + self._batches = (*self._batches, *self._sink.drain()) + return tuple( + (str(event["status"]), int(event["prompt_tokens"]), int(event["completion_tokens"])) + for event in _sink_events(self._batches) + if event.get("model_group") == self._model + ) + + +def _request_rows(gateway: Gateway, call_id: str) -> list[dict[str, JsonValue]]: + return gateway.rows( + 'SELECT status, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" WHERE litellm_call_id=%s', (call_id,) + ) + + +@pytest.mark.parametrize( + "failure", + (_Failure("drop_after_message_start"), _Failure("error_frame_after_message_start", 500)), + ids=("transport-drop", "recoverable-error-frame"), +) +def test_superseded_attempt_and_its_retry_each_log_their_own_usage( + gateway: Gateway, tmp_path: Path, failure: _Failure +) -> None: + prompt: Final = _prompt() + attempts: Final = Attempts() + with ( + wire_server(_upstream(prompt, failure, attempts)) as wire, + wire_server(lambda request: Reply()) as sink, + _proxy_with_callback_sink(gateway, tmp_path, sink) as candidate, + candidate.scenario() as scenario, + ): + model: Final = scenario.model(model=f"anthropic/{_MODEL}", api_base=wire.url, api_key=_API_KEY, num_retries=1) + status, events = _stream(candidate, _body(model, prompt, "deployment")) + _assert_completed_by_retry(candidate, status, events, prompt, model, wire) + delivered: Final = eventually( + _SinkReader(sink, model).delivered_usage, lambda values: len(values) >= 2, seconds=20 + ) + assert sorted(delivered) == [("failure", 5, 1), ("success", 5, 3)], delivered + + +def test_every_dropped_attempt_logs_its_own_failure_and_the_request_row_still_bills_the_dropped_usage( + gateway: Gateway, tmp_path: Path +) -> None: + prompt: Final = _prompt() + attempts: Final = Attempts() + with ( + wire_server(_upstream(prompt, _Failure("drop_after_message_start"), attempts, failing_attempts=99)) as wire, + wire_server(lambda request: Reply()) as sink, + _proxy_with_callback_sink(gateway, tmp_path, sink) as candidate, + candidate.scenario() as scenario, + ): + model: Final = scenario.model(model=f"anthropic/{_MODEL}", api_base=wire.url, api_key=_API_KEY, num_retries=1) + response: Final = candidate.request("POST", "/v1/messages", _body(model, prompt, "deployment")) + _assert_stream_failed_after_message_start(response.status_code, parse_sse(response.text), wire, attempts=2) + delivered: Final = eventually( + _SinkReader(sink, model).delivered_usage, lambda values: len(values) >= 2, seconds=20 + ) + assert delivered == (("failure", 5, 1), ("failure", 5, 1)), delivered + rows: Final = eventually( + lambda: _request_rows(candidate, response.headers["x-litellm-call-id"]), + lambda values: len(values) >= 1, + seconds=70, + ) + assert rows == [ + { + "status": "failure", + "prompt_tokens": sum(input_tokens for _, input_tokens, _ in delivered), + "completion_tokens": sum(output_tokens for _, _, output_tokens in delivered), + } + ], rows + + def _string_values(cache: Redis) -> tuple[bytes, ...]: keys: Final = tuple(key for key in cache.scan_iter(count=1000) if cache.type(key) == b"string") return tuple(value for value in cache.mget(keys) if value is not None) if keys else () @@ -485,7 +614,7 @@ def test_retried_stream_is_cached_and_the_identical_request_is_served_without_th model: Final = scenario.model(model=f"anthropic/{_MODEL}", api_base=wire.url, api_key=_API_KEY, num_retries=1) body: Final = _body(model, prompt, "deployment") status, events = _stream(gateway, body) - _assert_completed_by_retry(status, events, prompt, model, wire) + _assert_completed_by_retry(gateway, status, events, prompt, model, wire) eventually(lambda: _cached_somewhere(_served_id(prompt, 2)), bool, seconds=30) replay_status, replay = _stream(gateway, body) assert replay_status == 200, replay diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index 43190852c60..bbfae2229ce 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -1849,13 +1849,19 @@ def test_mark_logging_complete_flags_by_stream_and_event( "async_failure", "sync_failure", ) + attempt_marker: Final = datetime_unit_test(2026, 1, 1) logging_obj.stream = stream + logging_obj.model_call_details["current_attempt_api_call_start_time"] = attempt_marker logging_obj.mark_logging_complete(event_type=event) expected_to_run: Final = stream and event in ("async_success", "sync_success") assert logging_obj.should_run_logging(event_type=event) is expected_to_run assert all(logging_obj.should_run_logging(event_type=other_event) for other_event in events if other_event != event) + if event == "async_failure": + assert logging_obj.model_call_details["async_failure_logged_attempt"] is attempt_marker + else: + assert "async_failure_logged_attempt" not in logging_obj.model_call_details @pytest.mark.parametrize("stream", [False, True]) diff --git a/tests/unit/litellm_core_utils/test_max_streaming_duration.py b/tests/unit/litellm_core_utils/test_max_streaming_duration.py index fc8daab3899..234bc83e3c3 100644 --- a/tests/unit/litellm_core_utils/test_max_streaming_duration.py +++ b/tests/unit/litellm_core_utils/test_max_streaming_duration.py @@ -67,12 +67,13 @@ class TestCustomStreamWrapperMaxDuration: @pytest.mark.asyncio async def test_should_raise_on_async_anext_when_exceeded(self): - """__anext__ should check the limit before iterating, dispatching the - same failure-callback/logging path every other stream failure goes - through (dispatch_failure_handlers is async on the real Logging class, - so the mock needs to be awaitable too).""" + """__anext__ checks the limit before iterating and schedules failure logging.""" wrapper = _make_custom_stream_wrapper() - wrapper.logging_obj.dispatch_failure_handlers = AsyncMock() + + async def _noop() -> None: + return None + + wrapper.logging_obj.dispatch_attempt_failure_handlers = MagicMock(side_effect=lambda *_: _noop()) wrapper._stream_created_time = time.time() - 20 with patch("litellm.constants.LITELLM_MAX_STREAMING_DURATION_SECONDS", 10.0): with pytest.raises(litellm.Timeout): diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index 550ecd3c0f7..c13cea51183 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -13,8 +13,7 @@ from unittest.mock import AsyncMock, MagicMock, Mock, patch import httpx import pytest from openai.types.chat import ChatCompletionChunk -from pydantic import BaseModel - +from pydantic import BaseModel, TypeAdapter import litellm from litellm import verbose_logger from litellm._logging import session_id_var, trace_id_var @@ -36,6 +35,7 @@ from litellm.types.utils import ( ModelResponseStream, PromptTokensDetailsWrapper, StreamingChoices, + StandardLoggingPayload, Usage, ) from litellm.utils import ModelResponseListIterator @@ -4265,6 +4265,68 @@ def test_record_partial_usage_for_failure_stashes_usage_and_cost(): assert isinstance(logging_obj.model_call_details["response_cost"], float) +class _FailureRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.failures: tuple[tuple[int, int], ...] = () + self.logged: asyncio.Event = asyncio.Event() + + async def async_log_failure_event( + self, + kwargs: Mapping[str, object], + response_obj: object, + start_time: datetime.datetime, + end_time: datetime.datetime, + ) -> None: + payload: Final = TypeAdapter(StandardLoggingPayload).validate_python(kwargs["standard_logging_object"]) + self.failures = (*self.failures, (payload["prompt_tokens"], payload["completion_tokens"])) + self.logged.set() + + +@pytest.mark.asyncio +async def test_log_stream_failure_bills_the_dropped_attempt_even_after_the_router_discards_its_usage() -> None: + """The dropped attempt's failure log is queued, and the router discards the stashed usage as it opens the + fallback hop before that log runs. The log must still bill what the dropped attempt consumed, and the shared + logging object must not be marked as already failed, since the hop's own failure still has to reach the + integrations.""" + recorder: Final = _FailureRecorder() + logging_obj: Final = Logging( + model="gpt-5.6", + messages=[{"role": "user", "content": "Hey"}], + stream=True, + call_type="completion", + start_time=datetime.datetime(2026, 1, 1), + litellm_call_id="partial-usage-copy", + function_id="1245", + dynamic_async_failure_callbacks=[recorder], + ) + logging_obj.model_call_details["custom_llm_provider"] = "openai" + wrapper: Final = CustomStreamWrapper( + completion_stream=None, + model="gpt-5.6", + logging_obj=logging_obj, + custom_llm_provider="openai", + ) + wrapper.chunks = [ + ModelResponseStream( + id="chatcmpl-partial-copy", + created=1742056047, + model="gpt-5.6", + object="chat.completion.chunk", + choices=[StreamingChoices(finish_reason=None, index=0, delta=Delta(content="The", role="assistant"))], + usage=Usage(prompt_tokens=30, completion_tokens=1, total_tokens=31), + ) + ] + + with pytest.raises(MidStreamFallbackError): + wrapper._log_stream_failure_and_raise(httpx.ReadError("connection dropped")) + logging_obj.carry_failed_attempt_usage() + await asyncio.wait_for(recorder.logged.wait(), timeout=10) + + assert recorder.failures == ((30, 1),) + assert logging_obj.model_call_details.get("has_logged_async_failure") is not True + + def test_record_partial_usage_for_failure_noop_without_chunks(): """With no chunks delivered there is nothing billed to recover, so the failure stash must stay absent and not force a zero-usage row. diff --git a/tests/unit/llms/anthropic/pass_through/adapters/test_streaming_iterator_mid_stream_error.py b/tests/unit/llms/anthropic/pass_through/adapters/test_streaming_iterator_mid_stream_error.py index 785e1f1096b..6951253349e 100644 --- a/tests/unit/llms/anthropic/pass_through/adapters/test_streaming_iterator_mid_stream_error.py +++ b/tests/unit/llms/anthropic/pass_through/adapters/test_streaming_iterator_mid_stream_error.py @@ -19,12 +19,14 @@ import json import os import sys import threading +from collections.abc import Callable, Mapping, Sequence from datetime import datetime -from collections.abc import Callable -from typing import List, Optional +from typing import Final, Iterator, List, Optional from unittest.mock import MagicMock +import httpx import pytest +from pydantic import TypeAdapter import litellm @@ -33,12 +35,13 @@ sys.path.insert(0, os.path.abspath("../../../../..")) from litellm.exceptions import MidStreamFallbackError from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.llms.anthropic.pass_through.adapters.streaming_iterator import ( AnthropicStreamWrapper, _mid_stream_error_sse_event, ) from litellm.llms.bedrock.common_utils import BedrockError -from litellm.types.utils import Delta, StreamingChoices +from litellm.types.utils import Delta, ModelResponseStream, StandardLoggingPayload, StreamingChoices, Usage def _make_chunk(delta: Delta, finish_reason: Optional[str] = None) -> MagicMock: @@ -55,14 +58,14 @@ class _AsyncStreamThenRaise: """Yields the given chunks, then raises ``exc`` (mimics a provider stream that terminates mid-response).""" - def __init__(self, items: List[MagicMock], exc: BaseException): + def __init__(self, items: Sequence[object], exc: BaseException) -> None: self._it = iter(items) self._exc = exc - def __aiter__(self): + def __aiter__(self) -> "_AsyncStreamThenRaise": return self - async def __anext__(self): + async def __anext__(self) -> object: try: return next(self._it) except StopIteration: @@ -182,9 +185,73 @@ class _AsyncFailureRecorder(CustomLogger): def __init__(self): super().__init__() self.exceptions: list[BaseException] = [] + self.usages: tuple[tuple[int, int], ...] = () - async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + async def async_log_failure_event( + self, + kwargs: Mapping[str, object], + response_obj: object, + start_time: datetime, + end_time: datetime, + ) -> None: self.exceptions.append(kwargs["exception"]) + payload: Final = TypeAdapter(StandardLoggingPayload).validate_python(kwargs["standard_logging_object"]) + self.usages = (*self.usages, (payload["prompt_tokens"], payload["completion_tokens"])) + + +class _BoundaryAsyncFailureRecorder(_AsyncFailureRecorder): + def __init__(self) -> None: + super().__init__() + self.started = 0 + self.completed = 0 + + async def async_log_failure_event( + self, + kwargs: Mapping[str, object], + response_obj: object, + start_time: datetime, + end_time: datetime, + ) -> None: + self.started += 1 + await asyncio.sleep(0.05) + await super().async_log_failure_event(kwargs, response_obj, start_time, end_time) + self.completed += 1 + + +def test_anthropic_adapter_does_not_repeat_custom_stream_failure_dispatch() -> None: + async_recorder: Final = _BoundaryAsyncFailureRecorder() + sync_recorder: Final = _SyncFailureRecorder() + logging_obj: Final = _make_logging_obj("wrapped-custom-stream", async_recorder, sync_recorder) + chunk: Final = ModelResponseStream( + id="chatcmpl-dropped", + object="chat.completion.chunk", + created=1, + model="gpt-5.6", + choices=[StreamingChoices(index=0, delta=Delta(content="partial"), finish_reason=None)], + usage=Usage(prompt_tokens=30, completion_tokens=1, total_tokens=31), + ) + inner_stream: Final = CustomStreamWrapper( + completion_stream=_AsyncStreamThenRaise([chunk], httpx.ReadError("stream dropped")), + model="gpt-5.6", + logging_obj=logging_obj, + custom_llm_provider="openai", + ) + wrapper: Final = AnthropicStreamWrapper( + completion_stream=inner_stream, + model="openai/gpt-5.6", + litellm_logging_obj=logging_obj, + ) + + events: Final = asyncio.run(_drain_sse(wrapper)) + + assert async_recorder.started == 1 + assert async_recorder.completed == 1 + assert sync_recorder.called.wait(timeout=5), "sync failure callback never ran" + assert _parse_sse(events[-1])[0] == "error" + assert len(async_recorder.exceptions) == 1 + assert len(sync_recorder.exceptions) == 1 + assert async_recorder.usages == ((30, 1),) + assert sync_recorder.exceptions[0] == async_recorder.exceptions[0] class _SyncFailureRecorder: 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 d2b5c8d5bd9..4c3ee9e7ca5 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 @@ -3278,6 +3278,14 @@ async def test_vertex_ai_streaming_bad_request_is_not_wrapped(): async def dispatch_failure_handlers(self, *args, **kwargs): return None + async def dispatch_attempt_failure_handlers( + self, exception: Exception, traceback_exception: str + ) -> None: + return None + + def attempt_scoped_copy(self) -> "DummyLogging": + return self + async def failing_make_call(client=None, **kwargs): raise VertexAIError(status_code=400, message="bad input", headers={}) diff --git a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py index 4770ecd5485..cc701e1ebd4 100644 --- a/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py +++ b/tests/unit/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py @@ -2643,57 +2643,6 @@ class TestRecordPartialUsageForFailure: assert "response_cost" not in logging_obj.model_call_details -class _FailureRecorder(CustomLogger): - def __init__(self): - super().__init__() - self.failure_kwargs = [] - - async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): - self.failure_kwargs.append(kwargs) - - -class TestLogErrorFrameAsFailure: - """A stream closed at the provider's `event: error` frame is one failed attempt: its failure log bills the - usage streamed before the frame on the attempt-scoped copy its caller hands it, and the request's shared - logging state stays as it was, so the attempt the router opens next logs with nothing left behind.""" - - @pytest.mark.asyncio - async def test_logs_a_failure_with_the_usage_streamed_before_the_frame_and_leaves_the_request_state_alone(self): - recorder = _FailureRecorder() - logging_obj = LiteLLMLoggingObj( - model="claude-sonnet-5", - messages=[{"role": "user", "content": "hello"}], - stream=True, - call_type="anthropic_messages", - start_time=datetime.now(), - litellm_call_id="test-error-frame-failure", - function_id="test-error-frame-failure", - dynamic_async_failure_callbacks=[recorder], - ) - chunks = [ - *TestRecordPartialUsageForFailure()._interrupted_chunks()[:1], - TestRecordPartialUsageForFailure._sse( - "error", {"type": "error", "error": {"type": "overloaded_error", "message": "Overloaded"}} - ), - ] - - await AnthropicPassthroughLoggingHandler.log_error_frame_as_failure( - attempt_logging_obj=logging_obj.attempt_scoped_copy(), - request_body={"model": "claude-sonnet-5", "stream": True}, - all_chunks=chunks, - error_event=("overloaded_error", "Overloaded", 529), - ) - - assert len(recorder.failure_kwargs) == 1 - failure_payload = recorder.failure_kwargs[0]["standard_logging_object"] - assert failure_payload["status"] == "failure" - assert failure_payload["prompt_tokens"] == 52 - assert failure_payload["response_cost"] > 0 - assert isinstance(recorder.failure_kwargs[0]["exception"], litellm.InternalServerError) - assert "combined_usage_object" not in logging_obj.model_call_details - assert "response_cost" not in logging_obj.model_call_details - assert logging_obj.should_run_logging("async_failure") - @pytest.mark.parametrize( ("all_chunks", "interrupted"), diff --git a/tests/unit/proxy/pass_through_endpoints/test_streaming_handler.py b/tests/unit/proxy/pass_through_endpoints/test_streaming_handler.py index da084146894..7d56487a844 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_streaming_handler.py +++ b/tests/unit/proxy/pass_through_endpoints/test_streaming_handler.py @@ -220,7 +220,11 @@ async def test_failed_anthropic_stream_records_partial_usage_off_the_event_loop( logging_obj.model_call_details = {"model": model, "stream": True} logging_obj.litellm_params = {} logging_obj.get_router_model_id.return_value = None + logging_obj.attempt_scoped_copy.return_value = logging_obj logging_obj.dispatch_failure_handlers = AsyncMock() + logging_obj.dispatch_attempt_failure_handlers = AsyncMock( + side_effect=logging_obj.dispatch_failure_handlers + ) _, took, lags = await timed_with_loop_lags( lambda: PassThroughStreamingHandler.schedule_stream_failure_logging( diff --git a/tests/unit/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py b/tests/unit/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py index 8ce41129aec..4aa7e9cfbae 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py +++ b/tests/unit/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py @@ -2,7 +2,9 @@ import asyncio import json +from collections.abc import AsyncGenerator, AsyncIterator, Mapping, Sequence from datetime import datetime +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -22,6 +24,8 @@ from litellm.llms.anthropic.pass_through.messages.streaming_iterator import aclo from litellm.proxy.pass_through_endpoints.streaming_handler import ( PassThroughStreamingHandler, ) +from litellm.proxy.utils import _failure_fields_to_lift +from litellm.router_utils.fallback_event_handlers import carry_superseded_attempt_usage from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType @@ -407,7 +411,7 @@ async def _collect_openai_passthrough_chunks(chunks, endpoint_type): request_body={"model": "gpt-4o-mini", "stream": True}, litellm_logging_obj=_unarmed_logging_obj(), endpoint_type=endpoint_type, - start_time=datetime.now(), + start_time=datetime(2026, 1, 1), passthrough_success_handler_obj=MagicMock(), url_route="/openai/v1/chat/completions", route_streaming_logging=AsyncMock(), @@ -705,7 +709,7 @@ async def test_chunk_processor_logs_failure_not_success_on_mid_stream_exception( messages=[{"role": "user", "content": "hi"}], stream=True, call_type="anthropic_messages", - start_time=datetime.now(), + start_time=datetime(2026, 1, 1), litellm_call_id="test-mid-stream-timeout", function_id="test-mid-stream-timeout", dynamic_async_success_callbacks=[recorder], @@ -750,6 +754,109 @@ async def test_chunk_processor_logs_failure_not_success_on_mid_stream_exception( assert isinstance(recorder.failure_kwargs[0]["exception"], httpx.ReadTimeout) +def _anthropic_stream_dropping_after(frames: tuple[bytes, ...]) -> httpx.Response: + mock: Final = MagicMock(spec=httpx.Response) + mock.status_code = 200 + + async def _aiter_bytes() -> AsyncIterator[bytes]: + for frame in frames: + yield frame + raise httpx.ReadError("upstream closed the connection") + + mock.aiter_bytes = _aiter_bytes + return mock + + +def _anthropic_messages_attempt( + logging_obj: LiteLLMLoggingObj, response: httpx.Response +) -> AsyncGenerator[bytes, None]: + return PassThroughStreamingHandler.chunk_processor( + response=response, + request_body={"model": "claude-sonnet-5", "stream": True}, + litellm_logging_obj=logging_obj, + endpoint_type=EndpointType.ANTHROPIC, + start_time=datetime(2026, 1, 1), + passthrough_success_handler_obj=MagicMock(), + url_route="/v1/messages", + ) + + +def _anthropic_messages_logging_obj(recorder: _EventRecorder, call_id: str) -> LiteLLMLoggingObj: + return LiteLLMLoggingObj( + model="claude-sonnet-5", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="anthropic_messages", + start_time=datetime(2026, 1, 1), + litellm_call_id=call_id, + function_id=call_id, + dynamic_async_success_callbacks=[recorder], + dynamic_async_failure_callbacks=[recorder], + ) + + +async def _drop_after_message_start(logging_obj: LiteLLMLoggingObj, message_id: str, input_tokens: int) -> None: + dropped: Final = _anthropic_messages_attempt( + logging_obj, _anthropic_stream_dropping_after((_anthropic_message_start_frame(message_id, input_tokens),)) + ) + with pytest.raises(httpx.ReadError): + async for _ in dropped: + pass + + +@pytest.mark.asyncio +async def test_dropped_attempt_logs_its_failure_without_claiming_the_shared_logging_object() -> None: + """The router retries a /v1/messages stream the provider drops after message_start on the same logging + object. The dropped attempt's failure log bills its own usage and leaves the shared object free: the + retry's success log carries the retry's usage and no stale exception, while the stash stays for the + failure row of a request that runs out of retries until the router opens the next attempt.""" + recorder: Final = _EventRecorder() + logging_obj: Final = _anthropic_messages_logging_obj(recorder, "test-dropped-attempt") + + await _drop_after_message_start(logging_obj, "msg_a1", 52) + await _settle(recorder.failure_kwargs) + + assert recorder.failure_kwargs[0]["standard_logging_object"]["prompt_tokens"] == 52 + assert isinstance(recorder.failure_kwargs[0]["exception"], httpx.ReadError) + assert logging_obj.model_call_details.get("has_logged_async_failure") is not True + assert "exception" not in logging_obj.model_call_details + stashed: Final = logging_obj.model_call_details.get("combined_usage_object") + assert isinstance(stashed, litellm.Usage) and stashed.prompt_tokens == 52 + + logging_obj.carry_failed_attempt_usage() + retry: Final = _anthropic_messages_attempt( + logging_obj, _anthropic_stream_of((_anthropic_message_start_frame("msg_a2", 7), *_anthropic_completed_stream_tail())) + ) + async for _ in retry: + pass + await _settle(recorder.success_kwargs) + + success_payload: Final = recorder.success_kwargs[0]["standard_logging_object"] + assert success_payload["status"] == "success" + assert (success_payload["prompt_tokens"], success_payload["completion_tokens"]) == (7, 3) + assert recorder.success_kwargs[0].get("exception") is None + assert len(recorder.failure_kwargs) == 1 + + +@pytest.mark.asyncio +async def test_every_dropped_attempt_logs_its_own_failure() -> None: + """Two attempts of one /v1/messages request dropped in a row each log a failure billing their own usage; + the first attempt's log must not swallow the second's, and the shared object ends up holding the last + attempt's usage for the request's failure row.""" + recorder: Final = _EventRecorder() + logging_obj: Final = _anthropic_messages_logging_obj(recorder, "test-every-dropped-attempt") + + await _drop_after_message_start(logging_obj, "msg_a1", 52) + await _settle(recorder.failure_kwargs) + logging_obj.carry_failed_attempt_usage() + await _drop_after_message_start(logging_obj, "msg_a2", 9) + await _settle(recorder.failure_kwargs, count=2) + + assert [kwargs["standard_logging_object"]["prompt_tokens"] for kwargs in recorder.failure_kwargs] == [52, 9] + assert all(isinstance(kwargs["exception"], httpx.ReadError) for kwargs in recorder.failure_kwargs) + stashed: Final = logging_obj.model_call_details.get("combined_usage_object") + assert isinstance(stashed, litellm.Usage) and stashed.prompt_tokens == 9 + def _anthropic_message_start_frame(message_id: str = "msg_1", input_tokens: int = 52) -> bytes: return _anthropic_sse( "message_start", @@ -788,9 +895,9 @@ def _anthropic_completed_stream_tail() -> tuple[bytes, ...]: ) -async def _settle(recorded: list[dict[str, object]]) -> None: +async def _settle(recorded: Sequence[Mapping[str, object]], count: int = 1) -> None: for _ in range(300): - if recorded: + if len(recorded) >= count: return await asyncio.sleep(0.01) @@ -857,8 +964,6 @@ async def test_chunk_processor_logs_a_provider_error_frame_as_a_failure_when_clo assert received == list(frames) assert recorder.success_kwargs == [] assert logging_obj.model_call_details.get("has_dispatched_final_stream_success") is not True - assert "combined_usage_object" not in logging_obj.model_call_details - assert "response_cost" not in logging_obj.model_call_details assert logging_obj.should_run_logging("async_failure") assert len(recorder.failure_kwargs) == 1 failure_payload = recorder.failure_kwargs[0]["standard_logging_object"] @@ -991,6 +1096,31 @@ async def test_superseded_attempt_failure_log_keeps_its_own_deployment_when_the_ assert logging_obj.model_call_details["litellm_params"]["model_info"]["id"] == "deployment-2" +@pytest.mark.asyncio +async def test_error_frame_then_exhausted_retry_bills_both_attempts() -> None: + recorder: Final = _EventRecorder() + logging_obj: Final = _anthropic_messages_logging_obj(recorder, "test-error-frame-terminal-billing") + superseded: Final = _anthropic_messages_attempt( + logging_obj, _anthropic_stream_of((_anthropic_message_start_frame("msg_a1", 52), _anthropic_overloaded_error_frame())) + ) + async for chunk in superseded: + if b"event: error" in chunk: + break + await aclose_superseded_attempt(superseded) + carry_superseded_attempt_usage({"litellm_logging_obj": logging_obj}) + await _drop_after_message_start(logging_obj, "msg_a2", 9) + await _settle(recorder.failure_kwargs, count=2) + + failures: Final = tuple(event["standard_logging_object"] for event in recorder.failure_kwargs) + assert tuple(event["prompt_tokens"] for event in failures) == (52, 9) + terminal: Final = _failure_fields_to_lift({"litellm_logging_obj": logging_obj}) + usage: Final = terminal["combined_usage_object"] + assert isinstance(usage, litellm.Usage) + assert usage.prompt_tokens == sum(event["prompt_tokens"] for event in failures) + assert usage.completion_tokens == sum(event["completion_tokens"] for event in failures) + assert terminal["response_cost"] == pytest.approx(sum(event["response_cost"] for event in failures)) + + @pytest.mark.asyncio async def test_retry_success_log_carries_its_own_usage_after_a_provider_error_frame(): """Both attempts of a retried /v1/messages stream log through the request's one logging object. The @@ -1025,6 +1155,7 @@ async def test_retry_success_log_carries_its_own_usage_after_a_provider_error_fr if b"event: error" in chunk: break await aclose_superseded_attempt(superseded) + carry_superseded_attempt_usage({"litellm_logging_obj": logging_obj}) await _settle(recorder.failure_kwargs) retry = attempt((_anthropic_message_start_frame("msg_a2", 7), *_anthropic_completed_stream_tail())) async for _ in retry: @@ -1036,8 +1167,6 @@ async def test_retry_success_log_carries_its_own_usage_after_a_provider_error_fr success_payload = recorder.success_kwargs[0]["standard_logging_object"] assert success_payload["status"] == "success" assert (success_payload["prompt_tokens"], success_payload["completion_tokens"]) == (7, 3) - - def _google_sse(prompt_tokens: int, completion_tokens: int, text: str) -> bytes: payload = { "candidates": [{"content": {"parts": [{"text": text}], "role": "model"}, "index": 0}], diff --git a/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py b/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py index 7b18cb3350c..ab6a874a010 100644 --- a/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py +++ b/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py @@ -1,5 +1,8 @@ +import asyncio import datetime as real_datetime import smtplib +import traceback +from collections.abc import Mapping from typing import Final import httpx @@ -11,8 +14,15 @@ from litellm.exceptions import InternalServerError from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.bug_report import ISSUE_URL_BASE +from litellm.litellm_core_utils.litellm_logging import Logging from litellm.proxy._types import ProxyErrorTypes, UserAPIKeyAuth -from litellm.proxy.utils import PrismaClient, ProxyLogging, handle_exception_on_proxy +from litellm.proxy.utils import ( + PrismaClient, + ProxyLogging, + _failure_fields_to_lift, + handle_exception_on_proxy, +) +from litellm.router_utils.fallback_event_handlers import carry_superseded_attempt_usage from litellm.types.guardrails import GuardrailEventHooks @@ -2472,10 +2482,681 @@ def test_attempts_already_summed_into_the_final_row_are_not_lifted_again(): lifted = _failure_fields_to_lift({"litellm_logging_obj": logging_obj}, provider_billed=True) - assert lifted["combined_usage_object"].prompt_tokens == 24 + current_only_logging_obj: Final = _dispatched_failure_logging_obj( + FAILURE_USAGE_MODEL, + get_litellm_params( + metadata={ + "attempted_retries": 1, + "model_info": {"id": "dep-lit145-final"}, + "previous_models": ( + {"deployment_id": "dep-lit145-final", "exception_type": "Timeout", "attempted_retries": 0}, + ), + } + ), + ) + current_only: Final = _failure_fields_to_lift( + {"litellm_logging_obj": current_only_logging_obj}, provider_billed=True + ) + current_only_usage: Final = current_only["combined_usage_object"] + + assert current_only_usage.prompt_tokens > 0 + assert lifted["combined_usage_object"].prompt_tokens == 24 + current_only_usage.prompt_tokens assert "billed_failed_attempts" not in lifted +def test_carry_superseded_attempt_usage_without_stashed_usage_preserves_failure_estimate() -> None: + request_start: Final = real_datetime.datetime(2026, 1, 1) + logging_obj: Final = Logging( + model=FAILURE_USAGE_MODEL, + messages=ONE_USER_MESSAGE, + stream=False, + call_type="acompletion", + start_time=request_start, + litellm_call_id="call-lit145-no-stashed-usage", + function_id="fn-lit145-no-stashed-usage", + ) + logging_obj.pre_call(input=ONE_USER_MESSAGE, api_key="sk-test", model=FAILURE_USAGE_MODEL) + logging_obj.model_call_details["call_type"] = "acompletion" + logging_obj.model_call_details["messages"] = ONE_USER_MESSAGE + + try: + try: + raise httpx.ReadTimeout("no bytes in time") + except httpx.ReadTimeout as cause: + raise litellm.Timeout("Request timed out.", model=FAILURE_USAGE_MODEL, llm_provider="openai") from cause + except litellm.Timeout as timeout: + logging_obj.failure_handler( + timeout, + traceback.format_exc(), + request_start, + request_start + real_datetime.timedelta(seconds=1), + ) + + assert "combined_usage_object" not in logging_obj.model_call_details + request_data: Final = {"litellm_logging_obj": logging_obj} + before: Final = _failure_fields_to_lift(request_data, provider_billed=True) + + carry_superseded_attempt_usage(request_data) + + after: Final = _failure_fields_to_lift(request_data, provider_billed=True) + before_usage: Final = before.get("combined_usage_object") + after_usage: Final = after.get("combined_usage_object") + assert before_usage is not None + assert after_usage is not None + assert after_usage == before_usage + assert before["response_cost"] > 0 + assert after["response_cost"] == before["response_cost"] + + +@pytest.mark.asyncio +async def test_dropped_stream_failure_attribution_reaches_terminal_spend_row() -> None: + from litellm.exceptions import MidStreamFallbackError + from litellm.litellm_core_utils.get_litellm_params import get_litellm_params + from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices, Usage + + request_start: Final = real_datetime.datetime(2026, 1, 1) + model: Final = "deployment-model" + logging_obj: Final = Logging( + model=model, + messages=ONE_USER_MESSAGE, + stream=True, + call_type="acompletion", + start_time=request_start, + litellm_call_id="call-lit145-attribution", + function_id="fn-lit145-attribution", + ) + logging_obj.update_environment_variables( + model=model, + user="", + optional_params={}, + litellm_params=get_litellm_params(metadata={"model_info": {"id": "deployment-id"}}), + ) + logging_obj.model_call_details = { + **logging_obj.model_call_details, + "custom_llm_provider": "openai", + "first_api_call_start_time": request_start, + "current_attempt_api_call_start_time": request_start, + } + wrapper: Final = CustomStreamWrapper( + completion_stream=None, + model=model, + logging_obj=logging_obj, + custom_llm_provider="openai", + ) + wrapper.chunks = [ + ModelResponseStream( + id="chatcmpl-attribution", + created=1742056047, + model=model, + object="chat.completion.chunk", + choices=[StreamingChoices(finish_reason=None, index=0, delta=Delta(content="The"))], + usage=Usage(prompt_tokens=30, completion_tokens=1, total_tokens=31), + ) + ] + exception: Final = httpx.ReadError("connection dropped") + + with pytest.raises(MidStreamFallbackError): + wrapper._log_stream_failure_and_raise(exception) + await wrapper.wait_for_failure_logging(timeout=10) + + request_data: Final = {"litellm_logging_obj": logging_obj} + lifted: Final = _failure_fields_to_lift(request_data, provider_billed=True) + terminal_kwargs: Final = { + key: value + for key, value in {**request_data, **lifted}.items() + if key != "litellm_logging_obj" + } + row: Final = get_logging_payload( + kwargs=terminal_kwargs, + response_obj=exception, + start_time=request_start, + end_time=request_start + real_datetime.timedelta(seconds=1), + ) + + assert row["model_id"] == "deployment-id" + assert row["custom_llm_provider"] == "openai" + assert row["model"] == model + + +class _ReplaceSharedFailurePayload(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.logging_obj: Logging | None = None + self.replaced: asyncio.Event = asyncio.Event() + self.started: asyncio.Event = asyncio.Event() + self.release: asyncio.Event = asyncio.Event() + self.pause_until_released: bool = False + self.replacement_payload: object | None = None + + async def async_log_failure_event( + self, + kwargs: Mapping[str, object], + response_obj: object, + start_time: real_datetime.datetime, + end_time: real_datetime.datetime, + ) -> None: + from pydantic import TypeAdapter + + from litellm.types.utils import StandardLoggingPayload + + payload: Final = TypeAdapter(StandardLoggingPayload).validate_python(kwargs["standard_logging_object"]) + self.started.set() + if self.pause_until_released: + await self.release.wait() + logging_obj: Final = self.logging_obj + if logging_obj is None: + return + replacement: Final = self.replacement_payload or TypeAdapter(StandardLoggingPayload).validate_python( + { + **payload, + "model_id": "newer-deployment-id", + "custom_llm_provider": "newer-provider", + "model": "newer-model", + } + ) + logging_obj.model_call_details = { + **logging_obj.model_call_details, + "standard_logging_object": replacement, + } + self.replaced.set() + + +class _OrderedAttemptFailureRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.attempt_a_started: asyncio.Event = asyncio.Event() + self.attempt_b_started: asyncio.Event = asyncio.Event() + self.release_attempt_a: asyncio.Event = asyncio.Event() + self.release_attempt_b: asyncio.Event = asyncio.Event() + + async def async_log_failure_event( + self, + kwargs: Mapping[str, object], + response_obj: object, + start_time: real_datetime.datetime, + end_time: real_datetime.datetime, + ) -> None: + from pydantic import TypeAdapter + + from litellm.types.utils import StandardLoggingPayload + + payload: Final = TypeAdapter(StandardLoggingPayload).validate_python(kwargs["standard_logging_object"]) + if payload["model_id"] == "deployment-a": + self.attempt_a_started.set() + await self.release_attempt_a.wait() + return + if payload["model_id"] == "deployment-b": + self.attempt_b_started.set() + await self.release_attempt_b.wait() + + +class _PausedBaselineFailureRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.started: asyncio.Event = asyncio.Event() + self.release: asyncio.Event = asyncio.Event() + self.standard_logging_object: object | None = None + + async def async_log_failure_event( + self, + kwargs: Mapping[str, object], + response_obj: object, + start_time: real_datetime.datetime, + end_time: real_datetime.datetime, + ) -> None: + self.standard_logging_object = kwargs.get("standard_logging_object") + self.started.set() + await self.release.wait() + + +@pytest.mark.asyncio +async def test_latest_dropped_attempt_publishes_its_payload_after_earlier_dispatch_finishes() -> None: + from litellm.litellm_core_utils.get_litellm_params import get_litellm_params + from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload + from litellm.types.utils import Usage + + request_start: Final = real_datetime.datetime(2026, 1, 1) + marker_a: Final = request_start + marker_b: Final = request_start + real_datetime.timedelta(seconds=1) + recorder: Final = _OrderedAttemptFailureRecorder() + logging_obj: Final = Logging( + model="deployment-model-a", + messages=ONE_USER_MESSAGE, + stream=True, + call_type="acompletion", + start_time=request_start, + litellm_call_id="call-lit145-attempt-order", + function_id="fn-lit145-attempt-order", + dynamic_async_failure_callbacks=[recorder], + ) + logging_obj.update_environment_variables( + model="deployment-model-a", + user="", + optional_params={}, + litellm_params=get_litellm_params(metadata={"model_info": {"id": "deployment-a"}}), + ) + logging_obj.model_call_details = { + **logging_obj.model_call_details, + "custom_llm_provider": "openai", + "first_api_call_start_time": request_start, + "current_attempt_api_call_start_time": marker_a, + "combined_usage_object": Usage(prompt_tokens=4, completion_tokens=1, total_tokens=5), + "response_cost": 0.001, + } + + attempt_a: Final = asyncio.create_task( + logging_obj.dispatch_attempt_failure_handlers(httpx.ReadError("attempt A dropped"), "") + ) + await asyncio.wait_for(recorder.attempt_a_started.wait(), timeout=10) + + logging_obj.carry_failed_attempt_usage() + logging_obj.update_environment_variables( + model="deployment-model-b", + user="", + optional_params={}, + litellm_params=get_litellm_params(metadata={"model_info": {"id": "deployment-b"}}), + ) + logging_obj.model_call_details = { + **logging_obj.model_call_details, + "model": "deployment-model-b", + "custom_llm_provider": "anthropic", + "current_attempt_api_call_start_time": marker_b, + "combined_usage_object": Usage(prompt_tokens=7, completion_tokens=2, total_tokens=9), + "response_cost": 0.002, + } + + attempt_b: Final = asyncio.create_task( + logging_obj.dispatch_attempt_failure_handlers(httpx.ReadError("attempt B dropped"), "") + ) + await asyncio.wait_for(recorder.attempt_b_started.wait(), timeout=10) + recorder.release_attempt_a.set() + await attempt_a + recorder.release_attempt_b.set() + await attempt_b + + request_data: Final = {"litellm_logging_obj": logging_obj} + lifted: Final = _failure_fields_to_lift(request_data, provider_billed=True) + terminal_kwargs: Final = { + key: value + for key, value in {**request_data, **lifted}.items() + if key != "litellm_logging_obj" + } + row: Final = get_logging_payload( + kwargs=terminal_kwargs, + response_obj=httpx.ReadError("attempt B dropped"), + start_time=request_start, + end_time=request_start + real_datetime.timedelta(seconds=1), + ) + + assert row["model_id"] == "deployment-b" + assert row["custom_llm_provider"] == "anthropic" + assert row["model"] == "deployment-model-b" + + +@pytest.mark.asyncio +async def test_dropped_attempt_attribution_is_published_while_failure_callback_is_paused() -> None: + from litellm.litellm_core_utils.get_litellm_params import get_litellm_params + from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload + from litellm.types.utils import Usage + + request_start: Final = real_datetime.datetime(2026, 1, 1) + marker_a: Final = request_start + failure: Final = httpx.ReadError("attempt A dropped") + recorder: Final = _OrderedAttemptFailureRecorder() + logging_obj: Final = Logging( + model="deployment-model-a", + messages=ONE_USER_MESSAGE, + stream=True, + call_type="acompletion", + start_time=request_start, + litellm_call_id="call-lit145-attempt-paused", + function_id="fn-lit145-attempt-paused", + dynamic_async_failure_callbacks=[recorder], + ) + logging_obj.update_environment_variables( + model="deployment-model-a", + user="", + optional_params={}, + litellm_params=get_litellm_params(metadata={"model_info": {"id": "deployment-a"}}), + ) + logging_obj.model_call_details = { + **logging_obj.model_call_details, + "custom_llm_provider": "openai", + "first_api_call_start_time": request_start, + "current_attempt_api_call_start_time": marker_a, + "combined_usage_object": Usage(prompt_tokens=4, completion_tokens=1, total_tokens=5), + "response_cost": 0.001, + } + + attempt: Final = asyncio.create_task(logging_obj.dispatch_attempt_failure_handlers(failure, "")) + await asyncio.wait_for(recorder.attempt_a_started.wait(), timeout=10) + + request_data: Final = {"litellm_logging_obj": logging_obj} + lifted: Final = _failure_fields_to_lift(request_data, provider_billed=True) + terminal_kwargs: Final = { + key: value + for key, value in {**request_data, **lifted}.items() + if key != "litellm_logging_obj" + } + row: Final = get_logging_payload( + kwargs=terminal_kwargs, + response_obj=failure, + start_time=request_start, + end_time=request_start + real_datetime.timedelta(seconds=1), + ) + + try: + assert row["model_id"] == "deployment-a" + assert row["custom_llm_provider"] == "openai" + assert row["model"] == "deployment-model-a" + assert row["prompt_tokens"] == 4 + assert row["completion_tokens"] == 1 + finally: + recorder.release_attempt_a.set() + await attempt + + +@pytest.mark.asyncio +async def test_failed_attempt_publishes_baseline_reason_after_invalidation( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from pydantic import TypeAdapter + + from litellm.litellm_core_utils.get_litellm_params import get_litellm_params + from litellm.proxy.hooks.autorouter_baseline_cache import AutoRouterBaselineCache, CapturedBaselineObservation + from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload + from litellm.router import Router + from litellm.types.utils import CallTypes, StandardLoggingPayload, Usage + from tests.unit.proxy.hooks.test_autorouter_baseline_cache import _MODELS, _kwargs + + request_start: Final = real_datetime.datetime(2026, 1, 1) + marker: Final = request_start + model: Final = "deployment-model-a" + router: Final = Router(model_list=_MODELS, disable_cooldowns=True) + baseline_cache: Final = AutoRouterBaselineCache( + None, + router=lambda: router, + clock=lambda: request_start.timestamp() + 10, + ) + recorder: Final = _PausedBaselineFailureRecorder() + logging_obj: Final = Logging( + model=model, + messages=ONE_USER_MESSAGE, + stream=True, + call_type=CallTypes.anthropic_messages.value, + start_time=request_start, + litellm_call_id="call-lit145-baseline-failure", + function_id="fn-lit145-baseline-failure", + dynamic_async_failure_callbacks=[recorder], + ) + logging_obj.update_environment_variables( + model=model, + user="", + optional_params={}, + litellm_params=get_litellm_params(metadata={"model_info": {"id": "deployment-a"}}), + ) + logging_obj.model_call_details = { + **logging_obj.model_call_details, + "custom_llm_provider": "openai", + "first_api_call_start_time": request_start, + "current_attempt_api_call_start_time": marker, + "combined_usage_object": Usage(prompt_tokens=4, completion_tokens=1, total_tokens=5), + "response_cost": 0.001, + } + await baseline_cache.async_pre_call_deployment_hook(_kwargs(logging_obj), CallTypes.anthropic_messages) + assert logging_obj.baseline_cache_context is not None + + failure: Final = httpx.ReadError("attempt A dropped") + attempt: Final = asyncio.create_task(logging_obj.dispatch_attempt_failure_handlers(failure, "")) + try: + await asyncio.wait_for(recorder.started.wait(), timeout=10) + request_data: Final = {"litellm_logging_obj": logging_obj} + lifted: Final = _failure_fields_to_lift(request_data, provider_billed=True) + terminal_kwargs: Final = { + key: value + for key, value in {**request_data, **lifted}.items() + if key != "litellm_logging_obj" + } + row: Final = get_logging_payload( + kwargs=terminal_kwargs, + response_obj=failure, + start_time=request_start, + end_time=request_start + real_datetime.timedelta(seconds=1), + ) + + assert row["model_id"] == "deployment-a" + assert row["custom_llm_provider"] == "openai" + assert row["model"] == model + assert row["prompt_tokens"] == 4 + assert row["completion_tokens"] == 1 + lifted_payload: Final = TypeAdapter(StandardLoggingPayload).validate_python(lifted["standard_logging_object"]) + callback_payload: Final = TypeAdapter(StandardLoggingPayload).validate_python( + recorder.standard_logging_object + ) + + def baseline_reason(payload: StandardLoggingPayload) -> str: + encoded: Final = payload["autorouter_baseline_observation"] + assert isinstance(encoded, str) + return CapturedBaselineObservation.model_validate_json(encoded).observation.reason or "" + + assert baseline_reason(lifted_payload) == "failed_request" + assert baseline_reason(callback_payload) == "failed_request" + finally: + recorder.release.set() + await attempt + + +@pytest.mark.asyncio +async def test_queued_stale_attempt_dispatch_does_not_publish_over_current_attempt() -> None: + from litellm.litellm_core_utils.get_litellm_params import get_litellm_params + from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload + from litellm.types.utils import Usage + + request_start: Final = real_datetime.datetime(2026, 1, 1) + marker_a: Final = request_start + marker_b: Final = request_start + real_datetime.timedelta(seconds=1) + recorder: Final = _OrderedAttemptFailureRecorder() + logging_obj: Final = Logging( + model="deployment-model-a", + messages=ONE_USER_MESSAGE, + stream=True, + call_type="acompletion", + start_time=request_start, + litellm_call_id="call-lit145-stale-attempt", + function_id="fn-lit145-stale-attempt", + dynamic_async_failure_callbacks=[recorder], + ) + logging_obj.update_environment_variables( + model="deployment-model-a", + user="", + optional_params={}, + litellm_params=get_litellm_params(metadata={"model_info": {"id": "deployment-a"}}), + ) + logging_obj.model_call_details = { + **logging_obj.model_call_details, + "custom_llm_provider": "openai", + "first_api_call_start_time": request_start, + "current_attempt_api_call_start_time": marker_a, + "combined_usage_object": Usage(prompt_tokens=4, completion_tokens=1, total_tokens=5), + "response_cost": 0.001, + } + failure_a: Final = httpx.ReadError("attempt A dropped") + coro_a: Final = logging_obj.dispatch_attempt_failure_handlers(failure_a, "") + + logging_obj.carry_failed_attempt_usage() + logging_obj.update_environment_variables( + model="deployment-model-b", + user="", + optional_params={}, + litellm_params=get_litellm_params(metadata={"model_info": {"id": "deployment-b"}}), + ) + logging_obj.model_call_details = { + **logging_obj.model_call_details, + "model": "deployment-model-b", + "custom_llm_provider": "anthropic", + "current_attempt_api_call_start_time": marker_b, + "combined_usage_object": Usage(prompt_tokens=7, completion_tokens=2, total_tokens=9), + "response_cost": 0.002, + } + failure_b: Final = httpx.ReadError("attempt B dropped") + coro_b: Final = logging_obj.dispatch_attempt_failure_handlers(failure_b, "") + + task_a: Final = asyncio.create_task(coro_a) + task_b: Final = asyncio.create_task(coro_b) + try: + await asyncio.wait_for( + asyncio.gather(recorder.attempt_a_started.wait(), recorder.attempt_b_started.wait()), + timeout=10, + ) + request_data: Final = {"litellm_logging_obj": logging_obj} + lifted: Final = _failure_fields_to_lift(request_data, provider_billed=True) + terminal_kwargs: Final = { + key: value + for key, value in {**request_data, **lifted}.items() + if key != "litellm_logging_obj" + } + row: Final = get_logging_payload( + kwargs=terminal_kwargs, + response_obj=failure_b, + start_time=request_start, + end_time=request_start + real_datetime.timedelta(seconds=1), + ) + + assert row["model_id"] == "deployment-b" + assert row["custom_llm_provider"] == "anthropic" + assert row["model"] == "deployment-model-b" + finally: + recorder.release_attempt_a.set() + recorder.release_attempt_b.set() + await asyncio.gather(task_a, task_b) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("queued_start", [False, True], ids=["running-start", "queued-start"]) +async def test_dropped_stream_failure_does_not_overwrite_newer_terminal_payload(queued_start: bool) -> None: + from litellm.exceptions import MidStreamFallbackError + from litellm.litellm_core_utils.get_litellm_params import get_litellm_params + from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload + from litellm.types.utils import Delta, ModelResponseStream, StandardLoggingPayload, StreamingChoices, Usage + from pydantic import TypeAdapter + + request_start: Final = real_datetime.datetime(2026, 1, 1) + model: Final = "deployment-model" + recorder: Final = _ReplaceSharedFailurePayload() + logging_obj: Final = Logging( + model=model, + messages=ONE_USER_MESSAGE, + stream=True, + call_type="acompletion", + start_time=request_start, + litellm_call_id="call-lit145-identity", + function_id="fn-lit145-identity", + dynamic_async_failure_callbacks=[recorder], + ) + recorder.logging_obj = logging_obj + initial_payload: Final = { + "model_id": "older-deployment-id", + "custom_llm_provider": "older-provider", + "model": "older-model", + } + logging_obj.update_environment_variables( + model=model, + user="", + optional_params={}, + litellm_params=get_litellm_params(metadata={"model_info": {"id": "deployment-id"}}), + ) + logging_obj.model_call_details = { + **logging_obj.model_call_details, + "custom_llm_provider": "openai", + "first_api_call_start_time": request_start, + "current_attempt_api_call_start_time": request_start, + "standard_logging_object": initial_payload, + } + + def assert_newer_payload_attribution() -> None: + request_data: Final = {"litellm_logging_obj": logging_obj} + lifted: Final = _failure_fields_to_lift(request_data, provider_billed=True) + terminal_kwargs: Final = { + key: value + for key, value in {**request_data, **lifted}.items() + if key != "litellm_logging_obj" + } + row: Final = get_logging_payload( + kwargs=terminal_kwargs, + response_obj=exception, + start_time=request_start, + end_time=request_start + real_datetime.timedelta(seconds=1), + ) + + assert row["model_id"] == "newer-deployment-id" + assert row["custom_llm_provider"] == "newer-provider" + assert row["model"] == "newer-model" + + exception: Final = httpx.ReadError("connection dropped") + if queued_start: + dispatch: Final = logging_obj.dispatch_attempt_failure_handlers(exception, "") + logging_obj._failure_handler_helper_fn( + exception=httpx.ReadError("newer failure"), + traceback_exception="", + start_time=request_start, + end_time=request_start + real_datetime.timedelta(seconds=1), + ) + replacement_base: Final = logging_obj.model_call_details["standard_logging_object"] + replacement_payload: Final = TypeAdapter(StandardLoggingPayload).validate_python( + { + **replacement_base, + "model_id": "newer-deployment-id", + "custom_llm_provider": "newer-provider", + "model": "newer-model", + } + ) + recorder.pause_until_released = True + recorder.replacement_payload = replacement_payload + logging_obj.model_call_details = { + **logging_obj.model_call_details, + "standard_logging_object": replacement_payload, + } + assert logging_obj.model_call_details["current_attempt_api_call_start_time"] == request_start + + task: Final = asyncio.create_task(dispatch) + try: + await asyncio.wait_for(recorder.started.wait(), timeout=10) + assert_newer_payload_attribution() + finally: + recorder.release.set() + await task + + await asyncio.wait_for(recorder.replaced.wait(), timeout=10) + assert_newer_payload_attribution() + return + + wrapper: Final = CustomStreamWrapper( + completion_stream=None, + model=model, + logging_obj=logging_obj, + custom_llm_provider="openai", + ) + wrapper.chunks = [ + ModelResponseStream( + id="chatcmpl-identity", + created=1742056047, + model=model, + object="chat.completion.chunk", + choices=[StreamingChoices(finish_reason=None, index=0, delta=Delta(content="The"))], + usage=Usage(prompt_tokens=30, completion_tokens=1, total_tokens=31), + ) + ] + + with pytest.raises(MidStreamFallbackError): + wrapper._log_stream_failure_and_raise(exception) + await wrapper.wait_for_failure_logging(timeout=10) + await asyncio.wait_for(recorder.replaced.wait(), timeout=10) + assert_newer_payload_attribution() + + def test_a_timed_out_attempt_from_the_group_a_fallback_left_is_still_billed(): """After a fallback hop the router restarts ``attempted_retries`` at 0 for the new model group and re-logs the attempt that exhausted the previous group, so the final attempt (dep-b, retry 1) must be @@ -2655,6 +3336,47 @@ def test_a_request_that_timed_out_at_a_provider_bills_its_input_at_the_input_rat assert response_cost == pytest.approx(expected_input_cost) +@pytest.mark.parametrize( + "current_attempt_api_call_start_time", + [None, 1700000001.0], + ids=["retry-not-dispatched", "retry-dispatched"], +) +def test_a_retry_failure_bills_input_only_for_dispatched_attempts( + current_attempt_api_call_start_time: float | None, +) -> None: + from litellm.proxy.utils import _estimate_dispatched_failure_usage, _failure_fields_to_lift + from litellm.types.utils import Usage + + prior_usage: Final = Usage(prompt_tokens=4, completion_tokens=2, total_tokens=6) + estimated_retry_usage: Final = _estimate_dispatched_failure_usage(FAILURE_USAGE_MODEL, ONE_USER_MESSAGE, None) + assert estimated_retry_usage is not None + + lifted: Final = _failure_fields_to_lift( + { + "litellm_logging_obj": _LoggingObj( + { + "first_api_call_start_time": 1700000000.0, + "current_attempt_api_call_start_time": current_attempt_api_call_start_time, + "failed_attempts_usage": ((prior_usage, 0.25),), + "call_type": "acompletion", + "model": FAILURE_USAGE_MODEL, + "messages": ONE_USER_MESSAGE, + } + ) + } + ) + + usage: Final = lifted["combined_usage_object"] + assert isinstance(usage, Usage) + expected_retry_tokens: Final = ( + estimated_retry_usage.prompt_tokens if current_attempt_api_call_start_time is not None else 0 + ) + assert usage.prompt_tokens == prior_usage.prompt_tokens + expected_retry_tokens + assert usage.completion_tokens == prior_usage.completion_tokens + assert usage.total_tokens == prior_usage.total_tokens + expected_retry_tokens + assert lifted["response_cost"] == 0.25 + + @pytest.mark.parametrize( "model_call_details, dispatched", [ diff --git a/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py b/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py index f0f3fc37ddb..e88d951ff09 100644 --- a/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py +++ b/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py @@ -27,7 +27,7 @@ from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.xai.common_utils import xai_reported_cost_in_usd from litellm.proxy._experimental.mcp_server import mcp_server_manager from litellm.proxy._experimental.mcp_server.faults.list_outcomes import ServerListOk -from litellm.proxy.utils import _failure_fields_to_lift +from litellm.proxy.utils import _estimate_dispatched_failure_usage, _failure_fields_to_lift from litellm.responses.mcp.litellm_proxy_mcp_handler import LiteLLM_Proxy_MCP_Handler from litellm.types.llms.openai import ( ResponsesAPIResponse, @@ -1980,7 +1980,22 @@ async def test_non_stream_auto_execute_failing_every_retry_bills_each_paid_first failure_row_fields: Final = _failure_fields_to_lift(MappingProxyType({"litellm_logging_obj": logging_obj})) billed_usage: Final = failure_row_fields["combined_usage_object"] assert isinstance(billed_usage, Usage) - assert (billed_usage.prompt_tokens, billed_usage.completion_tokens) == billed_tokens + estimated_failure_usage: Final = ( + _estimate_dispatched_failure_usage( + model="gpt-5.6", + request_input=logging_obj.model_call_details.get("messages"), + system_input=None, + ) + if not retried_first_rounds + else None + ) + expected_prompt_tokens: Final = billed_tokens[0] + ( + estimated_failure_usage.prompt_tokens if estimated_failure_usage is not None else 0 + ) + assert (billed_usage.prompt_tokens, billed_usage.completion_tokens) == ( + expected_prompt_tokens, + billed_tokens[1], + ) assert failure_row_fields["response_cost"] == pytest.approx( _round_cost(_ROUND_ONE_BODY) + sum(_round_cost(body) for body in retried_first_rounds) ) diff --git a/tests/unit/router_utils/test_fallback_event_handlers.py b/tests/unit/router_utils/test_fallback_event_handlers.py index 09ce002353f..75b6ab04446 100644 --- a/tests/unit/router_utils/test_fallback_event_handlers.py +++ b/tests/unit/router_utils/test_fallback_event_handlers.py @@ -2,7 +2,7 @@ import asyncio import importlib import json import time -from collections.abc import AsyncIterator +from collections.abc import AsyncIterator, Mapping from datetime import datetime, timedelta from types import MappingProxyType from typing import Any, Final, Literal, NoReturn @@ -11,6 +11,7 @@ from unittest.mock import MagicMock, patch import httpx import pytest import respx +from pydantic import TypeAdapter import litellm from litellm import Router @@ -19,7 +20,8 @@ from litellm.llms.anthropic.pass_through.messages.fake_stream_iterator import Fa from litellm.litellm_core_utils import get_llm_provider_logic from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper -from litellm.router_utils.cooldown_handlers import mark_advisor_orchestration_failure +from litellm.router_utils.cooldown_handlers import async_get_cooldown_deployments, mark_advisor_orchestration_failure +from litellm.types.utils import Choices, Message, ModelResponse, StandardLoggingPayload, Usage from litellm.router_utils.fallback_event_handlers import ( MID_STREAM_FALLBACK_CONTROLS_KEY, PRE_ROUTING_SELECTED_MODEL_KEY, @@ -147,6 +149,220 @@ class RecordingRouter: return StreamingWrapper() +def _request_logging_obj() -> Logging: + return Logging( + model="primary-model", + messages=[{"role": "user", "content": "ping"}], + stream=True, + call_type="anthropic_messages", + start_time=datetime(2026, 1, 1), + litellm_call_id="fallback-hop", + function_id="fallback-hop", + ) + + +class _SuccessRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.logged: tuple[tuple[int, float], ...] = () + + async def async_log_success_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + payload: Final = TypeAdapter(StandardLoggingPayload).validate_python(kwargs["standard_logging_object"]) + self.logged = (*self.logged, (payload["completion_tokens"], payload["response_cost"])) + + +class HopLoggingItsOwnSuccessRouter(FakeRouter): + async def async_function_with_fallbacks(self, *args: object, **kwargs: object) -> StreamingWrapper: + hop_response: Final = ModelResponse( + id="chatcmpl-hop", + model="gpt-5.6", + choices=[Choices(index=0, message=Message(role="assistant", content="served by the hop"))], + usage=Usage(prompt_tokens=52, completion_tokens=7, total_tokens=59), + ) + now: Final = datetime(2026, 1, 1) + logging_obj: Final = kwargs["litellm_logging_obj"] + assert isinstance(logging_obj, Logging) + await logging_obj.async_success_handler(result=hop_response, start_time=now, end_time=now) + return StreamingWrapper() + + +@pytest.mark.asyncio +async def test_run_async_fallback_opens_the_hop_without_the_superseded_attempts_partial_usage() -> None: + """A stream the failed attempt dropped stashed the usage it consumed on the request's logging object. The + hop that replaces it logs its own usage, so the integrations see the hop's tokens and cost, not the stash.""" + recorder: Final = _SuccessRecorder() + logging_obj: Final = Logging( + model="gpt-5.6", + messages=[{"role": "user", "content": "ping"}], + stream=False, + call_type="completion", + start_time=datetime(2026, 1, 1), + litellm_call_id="fallback-hop-success", + function_id="fallback-hop-success", + dynamic_async_success_callbacks=[recorder], + ) + logging_obj.model_call_details["custom_llm_provider"] = "openai" + logging_obj.model_call_details["litellm_params"] = {"acompletion": True} + logging_obj.model_call_details["current_attempt_api_call_start_time"] = datetime(2026, 1, 1) + logging_obj.litellm_params = logging_obj.model_call_details["litellm_params"] + partial_usage: Final = Usage(prompt_tokens=52, completion_tokens=1, total_tokens=53) + logging_obj.record_partial_usage_for_failure(partial_usage, 0.01) + + await run_async_fallback( + litellm_router=HopLoggingItsOwnSuccessRouter(), + fallback_model_group=["fallback-model"], + original_model_group="primary-model", + original_exception=RuntimeError("stream dropped before content"), + max_fallbacks=3, + fallback_depth=0, + litellm_logging_obj=logging_obj, + ) + + assert tuple(tokens for tokens, _ in recorder.logged) == (7,) + assert recorder.logged[0][1] != 0.01 + assert logging_obj.model_call_details["failed_attempts_usage"] == ((partial_usage, 0.01),) + assert logging_obj.model_call_details["current_attempt_api_call_start_time"] is None + + +@pytest.mark.asyncio +async def test_run_async_fallback_preserves_superseded_usage_when_hop_fails_before_opening() -> None: + logging_obj: Final = _request_logging_obj() + partial_usage: Final = Usage(prompt_tokens=52, completion_tokens=1, total_tokens=53) + logging_obj.record_partial_usage_for_failure(partial_usage, 0.01) + + with pytest.raises(RuntimeError, match="fallback model also failed"): + await run_async_fallback( + litellm_router=AlwaysFailRouter(), + fallback_model_group=["fallback-model"], + original_model_group="primary-model", + original_exception=RuntimeError("stream dropped before content"), + max_fallbacks=3, + fallback_depth=0, + litellm_logging_obj=logging_obj, + ) + + assert logging_obj.model_call_details["failed_attempts_usage"] == ((partial_usage, 0.01),) + + +class _RateLimitedHop(Exception): + def __init__(self, deployment_id: str = "fallback-dep-1") -> None: + super().__init__("fallback deployment rate limited") + self.status_code = 429 + self.failed_deployment_id = deployment_id + + +_EARLIER_FAILURE_MARKER: Final = datetime(2026, 1, 1) +_FIRST_HOP_FAILURE_MARKER: Final = datetime(2026, 1, 2) +_FINAL_HOP_FAILURE_MARKER: Final = datetime(2026, 1, 3) + + +class _TwoFailureFallbackRouter(Router): + def __init__(self) -> None: + super().__init__( + model_list=[ + { + "model_name": "fallback-model", + "litellm_params": {"model": "openai/gpt-5.6", "api_key": "fake-key"}, + "model_info": {"id": f"fallback-dep-{index}"}, + } + for index in (1, 2) + ], + allowed_fails=0, + cooldown_time=60, + num_retries=0, + ) + + async def async_function_with_fallbacks(self, *args: object, **kwargs: object) -> NoReturn: + logging_obj: Final[object] = kwargs["litellm_logging_obj"] + assert isinstance(logging_obj, Logging) + assert logging_obj.model_call_details.get("has_logged_async_failure", False) is False + + logging_obj.model_call_details["current_attempt_api_call_start_time"] = _FIRST_HOP_FAILURE_MARKER + try: + raise _RateLimitedHop("fallback-dep-1") + except _RateLimitedHop: + logging_obj.mark_logging_complete("async_failure") + + logging_obj.model_call_details["current_attempt_api_call_start_time"] = _FINAL_HOP_FAILURE_MARKER + raise _RateLimitedHop("fallback-dep-2") + + +def _two_deployment_router() -> Router: + return Router( + model_list=[ + { + "model_name": "fallback-model", + "litellm_params": {"model": "openai/gpt-5.6", "api_key": "fake-key"}, + "model_info": {"id": f"fallback-dep-{index}"}, + } + for index in (1, 2) + ], + num_retries=0, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failure_already_logged", [False, True]) +async def test_run_async_fallback_cools_down_a_failed_hop_once(failure_already_logged: bool) -> None: + """The manual cooldown stands in for the failure callbacks the logging object blocks once a failure has + been logged. A hop whose own failure was the first one logged already ran them, so only a hop that failed + with the flag already set gets the manual cooldown.""" + router: Final = _two_deployment_router() + logging_obj: Final = _request_logging_obj() + if failure_already_logged: + logging_obj.model_call_details["current_attempt_api_call_start_time"] = _EARLIER_FAILURE_MARKER + logging_obj.mark_logging_complete("async_failure") + + async def rate_limited_hop(*args: object, **kwargs: object) -> NoReturn: + hop_logging_obj: Final = kwargs["litellm_logging_obj"] + assert isinstance(hop_logging_obj, Logging) + hop_logging_obj.model_call_details["current_attempt_api_call_start_time"] = _FINAL_HOP_FAILURE_MARKER + if hop_logging_obj.model_call_details.get("has_logged_async_failure", False) is not True: + hop_logging_obj.mark_logging_complete("async_failure") + raise _RateLimitedHop() + + with pytest.raises(_RateLimitedHop): + await run_async_fallback( + litellm_router=router, + fallback_model_group=["fallback-model"], + original_model_group="primary-model", + original_exception=RuntimeError("stream dropped before content"), + max_fallbacks=3, + fallback_depth=0, + litellm_logging_obj=logging_obj, + original_function=rate_limited_hop, + metadata={}, + ) + + cooled_down: Final = await async_get_cooldown_deployments(litellm_router_instance=router, parent_otel_span=None) + assert cooled_down == (["fallback-dep-1"] if failure_already_logged else []) + + +@pytest.mark.asyncio +async def test_run_async_fallback_cools_down_the_final_deployment_after_retry_failure_logging() -> None: + router: Final = _TwoFailureFallbackRouter() + logging_obj: Final = _request_logging_obj() + assert logging_obj.model_call_details.get("has_logged_async_failure", False) is False + + with pytest.raises(_RateLimitedHop): + await run_async_fallback( + litellm_router=router, + fallback_model_group=["fallback-model"], + original_model_group="primary-model", + original_exception=RuntimeError("stream dropped before content"), + max_fallbacks=3, + fallback_depth=0, + litellm_logging_obj=logging_obj, + metadata={}, + ) + + assert logging_obj.model_call_details.get("has_logged_async_failure", False) is True + cooled_down: Final = await async_get_cooldown_deployments(litellm_router_instance=router, parent_otel_span=None) + assert cooled_down == ["fallback-dep-2"] + + @pytest.mark.asyncio async def test_run_async_fallback_forwards_include_fallback_errors_to_nested_call(): """A nested fallback (multi-hop) must keep collecting errors, so the opt-in @@ -1163,11 +1379,18 @@ class TestRunAsyncFallbackTriggersCooldown: return None async def async_function_with_fallbacks(self, *args, **kwargs): + logging_obj: Final = kwargs.get("litellm_logging_obj") + if isinstance(logging_obj, Logging): + logging_obj.model_call_details["current_attempt_api_call_start_time"] = _FINAL_HOP_FAILURE_MARKER + if logging_obj.model_call_details.get("has_logged_async_failure", False) is not True: + logging_obj.mark_logging_complete("async_failure") raise RuntimeError("fallback model also failed") - def _logging_obj(self, has_logged_async_failure: bool) -> MagicMock: - logging_obj = MagicMock() - logging_obj.model_call_details = {"has_logged_async_failure": has_logged_async_failure} + def _logging_obj(self, has_logged_async_failure: bool) -> Logging: + logging_obj: Final = _request_logging_obj() + if has_logged_async_failure: + logging_obj.model_call_details["current_attempt_api_call_start_time"] = _EARLIER_FAILURE_MARKER + logging_obj.mark_logging_complete("async_failure") return logging_obj @pytest.mark.asyncio diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index b119d216542..e1f748fcb58 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -23,7 +23,7 @@ import pytest import respx from fastapi import HTTPException from opentelemetry import trace -from pydantic import BaseModel +from pydantic import BaseModel, TypeAdapter import litellm from litellm import APIConnectionError, Router @@ -93,7 +93,7 @@ from litellm.llms.base_llm.vector_store.transformation import( LiteLLMVectorStoreEmbeddingExecutor, RouterVectorStoreEmbeddingExecutor, ) -from litellm.types.utils import CallTypes, CredentialItem +from litellm.types.utils import CallTypes, CredentialItem, Usage from litellm.utils import _invalidate_model_cost_lowercase_map from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome from tests.fake_openai_endpoint import FAKE_OPENAI_API_BASE @@ -16726,6 +16726,7 @@ class _AnthropicMessagesFakeByteStream: self._hidden_params = {"additional_headers": {"x-amzn-requestid": "req-1"}} self.closed = False self.closed_as_superseded = False + self.close_contexts: tuple[bool, ...] = () def __aiter__(self): return self @@ -16736,6 +16737,7 @@ class _AnthropicMessagesFakeByteStream: return self._chunks.pop(0) async def aclose(self) -> None: + self.close_contexts = (*self.close_contexts, closing_superseded_attempt()) self.closed_as_superseded = self.closed_as_superseded or closing_superseded_attempt() self.closed = True @@ -17979,6 +17981,51 @@ async def test_anthropic_messages_stream_dropped_before_content_is_retried_withi assert [(attempted, budget) for _, attempted, budget in provider.calls] == [(0, 2), (1, 2)] +class _AnthropicMessagesRetryStashRecorder: + def __init__(self) -> None: + self.observed: tuple[Usage | None, ...] = () + + def record(self, usage: Usage | None) -> None: + self.observed = (*self.observed, usage) + + +@pytest.mark.asyncio +async def test_anthropic_messages_retry_opens_without_the_superseded_attempts_partial_usage() -> None: + """The dropped attempt's usage stays available for terminal billing without reaching the retry's success log.""" + router: Final = _anthropic_messages_retry_router(num_retries=1) + logging_obj: Final = LiteLLMLogging( + model="glm", + messages=[{"role": "user", "content": "ping"}], + stream=True, + call_type="anthropic_messages", + start_time=datetime(2026, 1, 1), + litellm_call_id="retry-own-usage", + function_id="retry-own-usage", + ) + partial_usage: Final = Usage(prompt_tokens=52, completion_tokens=1, total_tokens=53) + logging_obj.record_partial_usage_for_failure(partial_usage, 0.01) + stash_recorder: Final = _AnthropicMessagesRetryStashRecorder() + usage_adapter: Final = TypeAdapter(Usage | None) + + def retry_stream() -> _AnthropicMessagesFakeByteStream: + stash_recorder.record( + usage_adapter.validate_python(logging_obj.model_call_details.get("combined_usage_object")) + ) + return _anthropic_messages_retried_stream() + + provider: Final = _AnthropicMessagesScriptedProvider(_anthropic_messages_dropped_before_content, retry_stream) + + stream: Final = await _anthropic_messages_stream_through_router( + router, provider, litellm_logging_obj=logging_obj + ) + body: Final = tuple([chunk async for chunk in stream]) + + assert stash_recorder.observed == (None,) + assert body == (_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("pong")) + assert "response_cost" not in logging_obj.model_call_details + assert logging_obj.model_call_details["failed_attempts_usage"] == ((partial_usage, 0.01),) + + @pytest.mark.asyncio async def test_anthropic_messages_superseded_stream_is_closed_before_the_retry_opens(): """Both attempts log through the request's one logging object, so the superseded attempt's end-of-stream @@ -18001,6 +18048,7 @@ async def test_anthropic_messages_superseded_stream_is_closed_before_the_retry_o assert body == [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("pong")] assert superseded_closed_when_retry_opened == [True] + assert superseded.close_contexts[0] is True @pytest.mark.asyncio From 159b5b2f7243bba757c08e4ddc45a2514e4cd037 Mon Sep 17 00:00:00 2001 From: "berriai-litellm-provider-info-sync[bot]" <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> Date: Tue, 20 Oct 2026 21:52:01 -0700 Subject: [PATCH 04/14] chore(cost-map): add fireworks muse spark open priority prices from the pricing page (#47785) Price-Sync: litellm-providers Co-authored-by: berriai-litellm-provider-info-sync[bot] <328147090+berriai-litellm-provider-info-sync[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 6 +++++- model_prices_and_context_window.json | 6 +++++- 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 764cf397158..cb9bfa6ffe4 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -64195,13 +64195,17 @@ }, "fireworks_ai/accounts/fireworks/models/muse-spark-open": { "cache_read_input_token_cost": 1.5e-07, + "cache_read_input_token_cost_priority": 2.25e-07, "input_cost_per_token": 1.25e-06, + "input_cost_per_token_priority": 1.875e-06, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, "mode": "chat", "output_cost_per_token": 4.25e-06, - "source": "https://fireworks.ai/models/fireworks/muse-spark-open", + "output_cost_per_token_priority": 6.375e-06, + "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_reasoning": true, "supports_vision": true }, "fireworks_ai/accounts/fireworks/models/nemotron-lightning-3p5-30b-a3b": { diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 764cf397158..cb9bfa6ffe4 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -64195,13 +64195,17 @@ }, "fireworks_ai/accounts/fireworks/models/muse-spark-open": { "cache_read_input_token_cost": 1.5e-07, + "cache_read_input_token_cost_priority": 2.25e-07, "input_cost_per_token": 1.25e-06, + "input_cost_per_token_priority": 1.875e-06, "litellm_provider": "fireworks_ai", "max_input_tokens": 1048576, "mode": "chat", "output_cost_per_token": 4.25e-06, - "source": "https://fireworks.ai/models/fireworks/muse-spark-open", + "output_cost_per_token_priority": 6.375e-06, + "source": "https://docs.fireworks.ai/serverless/pricing", "supports_function_calling": true, + "supports_reasoning": true, "supports_vision": true }, "fireworks_ai/accounts/fireworks/models/nemotron-lightning-3p5-30b-a3b": { From b99d2cc06e9f5a05765228b85bc64d2a090ffd00 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Tue, 20 Oct 2026 22:31:08 -0700 Subject: [PATCH 05/14] test(perf): ramp the load cases until the pod stops keeping up and floor sustained_rps (#47758) * test(perf): add small_bodies_30rps and mixed_size_requests_4rps load cases Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(perf): run the load cases at 50 rps small and 5 rps mixed, the pod's knee Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(perf): ramp the load cases until the pod stops keeping up and report sustained_rps Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(perf): hand the ramp to the observation and widen the mixed ramp to 14 rps Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(perf): floor the small ramp at 20 rps and emit the per-step readings Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(perf): emit the ramp's hold bound and each step's unexpected statuses Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(perf): read each ramp step's cpu, redis latency and connections from the samplers' windows Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(perf): floor mixed_size_requests_ramp at 8 rps from the lane's 10 rps sustained step Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(perf): compare step readings with math.isclose so the e2e basedpyright gate stays clean Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: kerry Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/AGENTS.md | 2 +- .../perf/cases/mixed_size_requests_ramp.py | 24 +++ tests/e2e/perf/cases/small_bodies_ramp.py | 24 +++ tests/e2e/perf/setup/datastore.py | 18 +- tests/e2e/perf/setup/locustfile.py | 15 +- tests/e2e/perf/setup/observe.py | 129 ++++++++++++-- tests/e2e/perf/setup/performance_test_case.py | 78 +++++++- tests/e2e/perf/setup/plan.py | 4 + tests/e2e/perf/setup/podcpu.py | 18 +- tests/e2e/perf/setup/ramp.py | 116 ++++++++++++ tests/e2e/perf/setup/run.py | 1 + tests/e2e/perf/setup/series.py | 48 +++++ tests/e2e/perf/setup/telemetry.py | 51 +++++- tests/e2e/perf/setup/upgrade.py | 4 +- tests/e2e/perf/setup/wire.py | 2 + tests/e2e/perf/test_perf_cases.py | 15 +- tests/e2e_harness/perf/setup/test_ramp.py | 168 ++++++++++++++++++ tests/e2e_harness/perf/test_upgrade.py | 4 +- 18 files changed, 677 insertions(+), 44 deletions(-) create mode 100644 tests/e2e/perf/cases/mixed_size_requests_ramp.py create mode 100644 tests/e2e/perf/cases/small_bodies_ramp.py create mode 100644 tests/e2e/perf/setup/ramp.py create mode 100644 tests/e2e/perf/setup/series.py create mode 100644 tests/e2e_harness/perf/setup/test_ramp.py diff --git a/tests/e2e/AGENTS.md b/tests/e2e/AGENTS.md index d976a2b5ee3..ef57135a4d2 100644 --- a/tests/e2e/AGENTS.md +++ b/tests/e2e/AGENTS.md @@ -44,7 +44,7 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family - `logging/` - logging-integration delivery (datadog and friends) - `security/` - secret handling and log-leak protection - `router/` - routing and reliability behavior (fallbacks, cooldowns) plus the memory tests (`test_reliability_memory_e2e.py`: every worker's RSS as read at collection time, before any test traffic, must sit under a fixed idle budget, the release-gate check for a DB-backed boot that idles near the pod limit the way v1.100.x did, and a few hundred failing requests with retries and fallbacks must not grow proxy RSS past a loose ceiling over the measured batches nor store a request snapshot past a fixed size, the release-gate check for the v1.100.0 retry-breadcrumb leak) -- `perf/` - performance-category tests, kept OUT of the main suite: throughput/load SLO tests are a different testing category from functional e2e (variance-driven, historically flaky) and live outside this suite until re-implemented as their own pipeline (LIT-5163). Do not add a live load test that runs in the default collection. What lives here: the weekly session-anomaly test (`test_weekly_session_anomaly_e2e.py`, Claude Code-shaped multi-turn sessions against real providers with ceilings on error rate, cache read/write, turn time, and spend, marked `weekly` and deselected unless `E2E_WEEKLY_ANOMALY` is set, driven by `.github/workflows/weekly_load_anomaly.yml`), the Redis chaos test (`test_redis_chaos_e2e.py`, locust load against mock deployments split round robin over `/chat/completions` and `/v1/messages`, one endpoint per simulated user, with `CLIENT PAUSE ALL` on the proxy's Redis mid-run to simulate it being down outright, asserting zero failed requests on every endpoint, budgeting RSS and CPU-per-request as ratios against the same run's healthy phase, and holding p50/p90/p99 latency and log-bytes-per-request to flat ceilings (a ratio cannot bound those two: an open breaker skips Redis instead of waiting on it, so the chaos phase can measure cheaper than baseline while still being far slower than a user should see), needs a proxy booted from `gateway/redis_chaos_ci_config.yml` on the same host with `E2E_PROXY_PID` and `E2E_PROXY_LOG` set, marked `redis_chaos`, deselected unless `E2E_PERF_TESTS` is set and excluded from the per-PR selector like the rest of `perf/`, driven by `.github/workflows/test-e2e-redis-chaos.yml` and by the Buildkite `e2e-redis-chaos` step in project-releaser, which runs the proxy, Postgres and Valkey co-located with pytest in one pod and sets the opt-in). New perf tests go here too, marked `perf`, behind the same `E2E_PERF_TESTS` opt-in. The perf cases (`test_perf_cases.py`, one `PerformanceTestCase` per module in `perf/cases/`, no `_e2e` suffix) are the request-performance runs. A case names what the gateway was deployed with (model groups and guardrails, written out in the gateway's config: the `proxy_config` block of project-releaser's perf helm values in the lane, `gateway/perf_ci_config.yml` locally, with each deployment's mock knobs in an `x-mock-knobs` extra header so the model ids are real ones like `gpt-6-luna` and `anthropic.claude-sonnet-5`) and lists what the test creates and sends: teams, Keycloak JWT users, virtual keys, `RequestShape`s (a coding-agent session from `perf/setup/corpus/` whose input-token size is drawn per request from the shape's range, with or without tools, and the output cap), a `Traffic` plan (seed, rps, request count, endpoint, shape and stream shares) and `Limit`s (one metric at one percentile under one ceiling). The infra side lives in `perf/setup/`: `provision.py` checks the deployed config matches the case and builds the callers, `plan.py` draws the request list from the seed, `run.py` replays it from a separate Locust process (`locustfile.py`) that builds each body right before sending and records status, header, first-byte, first-token and total time per request, and `observe.py` joins all of that with what the mocks received and the workers' RSS to score every `Limit`. Limits are about what the caller sees and what the gateway costs: every non-blocked request is a 200 (`unexpected_status_count`), the seconds the gateway added per request at p50 and p99 (`gateway_added_seconds`, to the first token on streams), the client's time to the first token of a stream at p50 and p99 (`time_to_first_token_seconds`, mock delay included), worker memory growth (`memory_growth_megabytes`), and the Postgres statements and Redis commands a request waits on before it is forwarded (`postgres_calls_per_request`, `redis_calls_per_request`: auth and routing reads, the rate-limiter script), with the work the gateway does after the response or on a timer (spend writes, daily aggregates, config reloads) reported separately as `postgres_background_calls_per_request` and `redis_background_calls_per_request`; the same overhead as the gateway itself reports it, read from the `x-litellm-overhead-duration-ms` header per request (`litellm_overhead_seconds`, gated next to `gateway_added_seconds`, the outside reading of the same thing, so the two can be charted together) and the gateway's own end-to-end time from `x-litellm-response-duration-ms` (`litellm_response_seconds`), the run mean of the gateway's `litellm_redis_latency` histogram from a `/metrics` scrape before and after the run (`telemetry.py`, `prometheus_redis_latency_mean_seconds`), and the database's own view of the run from the same `pg_stat_statements` rows and a once-a-second `pg_stat_activity` sample: execution seconds per wall second (`postgres_active_sessions`, the average-active-sessions number RDS reports as DB load), the slowest statement's mean (`postgres_slowest_statement_seconds`, named in the detail) and the peak client connection count (`postgres_peak_connections`), each under a ceiling set from the first local runs and due a retune once the lane has its ten-run baseline; `datastore.py` reads `pg_stat_statements` before and after the run and holds a Redis `MONITOR` connection during it, sorts each statement by verb and table and each command by key, and the run output lists the top statements and commands per bucket so a red run names what grew. Whether a body, a guardrail call or a spend row is right is printed as a diagnostic, never a limit, that is the functional suites' job. The providers and the guardrail are out-of-process servers in `perf/setup/mocks/` (`server.py openai|anthropic|bedrock|guardrail`, each answering with a reply generated for the request in that provider's wire format, sized from the prompt and output cap, a tool call when tools were offered, waiting the case's `mock_latency` first, request logs at `/mock/requests` and `/mock/calls`), reached through `E2E_PERF_MOCK_*_URL` by both the gateway config and the test. Adding a perf test is adding a module with a `CASE` to `perf/cases/`; a new deployment or guardrail is a config change; a new metric is a name in `Metric` plus a branch in `observe.observe` +- `perf/` - performance-category tests, kept OUT of the main suite: throughput/load SLO tests are a different testing category from functional e2e (variance-driven, historically flaky) and live outside this suite until re-implemented as their own pipeline (LIT-5163). Do not add a live load test that runs in the default collection. What lives here: the weekly session-anomaly test (`test_weekly_session_anomaly_e2e.py`, Claude Code-shaped multi-turn sessions against real providers with ceilings on error rate, cache read/write, turn time, and spend, marked `weekly` and deselected unless `E2E_WEEKLY_ANOMALY` is set, driven by `.github/workflows/weekly_load_anomaly.yml`), the Redis chaos test (`test_redis_chaos_e2e.py`, locust load against mock deployments split round robin over `/chat/completions` and `/v1/messages`, one endpoint per simulated user, with `CLIENT PAUSE ALL` on the proxy's Redis mid-run to simulate it being down outright, asserting zero failed requests on every endpoint, budgeting RSS and CPU-per-request as ratios against the same run's healthy phase, and holding p50/p90/p99 latency and log-bytes-per-request to flat ceilings (a ratio cannot bound those two: an open breaker skips Redis instead of waiting on it, so the chaos phase can measure cheaper than baseline while still being far slower than a user should see), needs a proxy booted from `gateway/redis_chaos_ci_config.yml` on the same host with `E2E_PROXY_PID` and `E2E_PROXY_LOG` set, marked `redis_chaos`, deselected unless `E2E_PERF_TESTS` is set and excluded from the per-PR selector like the rest of `perf/`, driven by `.github/workflows/test-e2e-redis-chaos.yml` and by the Buildkite `e2e-redis-chaos` step in project-releaser, which runs the proxy, Postgres and Valkey co-located with pytest in one pod and sets the opt-in). New perf tests go here too, marked `perf`, behind the same `E2E_PERF_TESTS` opt-in. The perf cases (`test_perf_cases.py`, one `PerformanceTestCase` per module in `perf/cases/`, no `_e2e` suffix) are the request-performance runs. A case names what the gateway was deployed with (model groups and guardrails, written out in the gateway's config: the `proxy_config` block of project-releaser's perf helm values in the lane, `gateway/perf_ci_config.yml` locally, with each deployment's mock knobs in an `x-mock-knobs` extra header so the model ids are real ones like `gpt-6-luna` and `anthropic.claude-sonnet-5`) and lists what the test creates and sends: teams, Keycloak JWT users, virtual keys, `RequestShape`s (a coding-agent session from `perf/setup/corpus/` whose input-token size is drawn per request from the shape's range, with or without tools, and the output cap), a `Traffic` plan (seed, rps, request count, endpoint, shape and stream shares, or a `Ramp` that steps the rate up through the run, built with `ramped()`, whose plan carries each request's send second and whose run reports `sustained_rps`, the last step the pod held, scored in `perf/setup/ramp.py`, with every step's latency, completed rate, non-200s, CPU per request, Redis latency and Postgres connections cut from the samplers' time-stamped `Series` (`perf/setup/series.py`) to the step's window and emitted as `perf_ramp_step__` JUnit properties for the dashboard) and `Limit`s (one metric at one percentile under one ceiling, or over one floor for `sustained_rps`). The infra side lives in `perf/setup/`: `provision.py` checks the deployed config matches the case and builds the callers, `plan.py` draws the request list from the seed, `run.py` replays it from a separate Locust process (`locustfile.py`) that builds each body right before sending and records status, header, first-byte, first-token and total time per request, and `observe.py` joins all of that with what the mocks received and the workers' RSS to score every `Limit`. Limits are about what the caller sees and what the gateway costs: every non-blocked request is a 200 (`unexpected_status_count`), the seconds the gateway added per request at p50 and p99 (`gateway_added_seconds`, to the first token on streams), the client's time to the first token of a stream at p50 and p99 (`time_to_first_token_seconds`, mock delay included), worker memory growth (`memory_growth_megabytes`), and the Postgres statements and Redis commands a request waits on before it is forwarded (`postgres_calls_per_request`, `redis_calls_per_request`: auth and routing reads, the rate-limiter script), with the work the gateway does after the response or on a timer (spend writes, daily aggregates, config reloads) reported separately as `postgres_background_calls_per_request` and `redis_background_calls_per_request`; the same overhead as the gateway itself reports it, read from the `x-litellm-overhead-duration-ms` header per request (`litellm_overhead_seconds`, gated next to `gateway_added_seconds`, the outside reading of the same thing, so the two can be charted together) and the gateway's own end-to-end time from `x-litellm-response-duration-ms` (`litellm_response_seconds`), the run mean of the gateway's `litellm_redis_latency` histogram from a `/metrics` scrape before and after the run (`telemetry.py`, `prometheus_redis_latency_mean_seconds`), and the database's own view of the run from the same `pg_stat_statements` rows and a once-a-second `pg_stat_activity` sample: execution seconds per wall second (`postgres_active_sessions`, the average-active-sessions number RDS reports as DB load), the slowest statement's mean (`postgres_slowest_statement_seconds`, named in the detail) and the peak client connection count (`postgres_peak_connections`), each under a ceiling set from the first local runs and due a retune once the lane has its ten-run baseline; `datastore.py` reads `pg_stat_statements` before and after the run and holds a Redis `MONITOR` connection during it, sorts each statement by verb and table and each command by key, and the run output lists the top statements and commands per bucket so a red run names what grew. Whether a body, a guardrail call or a spend row is right is printed as a diagnostic, never a limit, that is the functional suites' job. The providers and the guardrail are out-of-process servers in `perf/setup/mocks/` (`server.py openai|anthropic|bedrock|guardrail`, each answering with a reply generated for the request in that provider's wire format, sized from the prompt and output cap, a tool call when tools were offered, waiting the case's `mock_latency` first, request logs at `/mock/requests` and `/mock/calls`), reached through `E2E_PERF_MOCK_*_URL` by both the gateway config and the test. Adding a perf test is adding a module with a `CASE` to `perf/cases/`; a new deployment or guardrail is a config change; a new metric is a name in `Metric` plus a branch in `observe.observe` - `other/` - the holding-pen suite for the `other.*` registry cluster with no home of its own yet: the master-key auth gate, JWT auth (access tokens issued by a real Keycloak realm, `idp.py` plus `idp_realm.json`, whose JWKS the proxy's `JWT_PUBLIC_KEY_URL` points at, see CONTRIBUTING.md for the start command and config block), and the process-lifecycle health probes (liveness, public readiness, authenticated readiness diagnostics). Promote a cluster out once it is large/stable enough for its own suite - `secret_manager/` - the gateway's `key_management_system` against a real secret manager: deployment keys resolved from it (`os.environ/` where the name exists only in the manager) and virtual keys written to and deleted from it. The tests are backend-agnostic and each backend is its own lane, because the setting is global to the proxy: `E2E_SECRET_MANAGER=` opts in and picks the backend from `secret_backends.BACKENDS`, the proxy is booted from `gateway/secret_manager__ci_config.yml` against the live manager, and the tests reach that manager through the backend's `SecretStore` (`secret_store_.py`). A test needing something not every backend does carries `requires_capability(...)` and is deselected on lanes that lack it. `secret_manager/backend.sh up ` runs a backend in Docker and writes the proxy's and the tests' env. Marked `secret_manager`, deselected unless `E2E_SECRET_MANAGER` is set, and kept out of the per-PR selector. Backends today: `hashicorp_vault` and `cyberark` (CyberArk Conjur, which cannot delete, so the delete test is Vault-only) - `gateway/` - proxy configuration only (`litellm-config.yml`). No tests diff --git a/tests/e2e/perf/cases/mixed_size_requests_ramp.py b/tests/e2e/perf/cases/mixed_size_requests_ramp.py new file mode 100644 index 00000000000..d7db545c031 --- /dev/null +++ b/tests/e2e/perf/cases/mixed_size_requests_ramp.py @@ -0,0 +1,24 @@ +"""The mixed-size traffic stepped up until the pod stops keeping up: the same callers, guardrails, mocks and body +mix (short chats, 300k to 1M-token contexts, tool calls, half streamed) as `mixed_size_requests_2rps`, 2 rps more +every 90 s from 2 to 14 on the same one 4-worker pod. Where the 2 rps case prices one big body on its own, this one +reports `sustained_rps`, the last step the pod held with big bodies overlapping in every worker, so a handler that +re-walks a 1M-token body on the event loop or a worker that holds several of them in memory at once moves the knee +down and fails the floor.""" + +from __future__ import annotations + +from dataclasses import replace +from typing import Final + +from setup.performance_test_case import Limit, Ramp, ramped + +from cases.mixed_size_requests_2rps import CASE as MIXED + +RAMP: Final = Ramp(first_rps=2.0, last_rps=14.0, step_rps=2.0, step_seconds=90.0, held_p99_seconds=5.0) + +CASE: Final = replace( + MIXED, + name="mixed_size_requests_ramp", + traffic=ramped(MIXED.traffic, RAMP, seed=9, concurrency=300), + limits=(Limit("sustained_rps", min=8),), +) diff --git a/tests/e2e/perf/cases/small_bodies_ramp.py b/tests/e2e/perf/cases/small_bodies_ramp.py new file mode 100644 index 00000000000..64308168cee --- /dev/null +++ b/tests/e2e/perf/cases/small_bodies_ramp.py @@ -0,0 +1,24 @@ +"""The small-bodies traffic stepped up until the pod stops keeping up: the same callers, guardrails, mocks and +2.5k to 4k-token bodies as `small_bodies_5rps`, 10 rps more every 90 s from 10 to 70 on the same one 4-worker pod. +The 5 rps case prices one request with the pod idle; this one reports `sustained_rps`, the last step the pod +held (gateway added p99 under 5 s, the queue still draining), so a change that moves the knee down, a lock held +across an await, a pool that fills, a flush that stalls the loop while dozens of requests are in flight, fails the +floor instead of hiding under a fixed rate the pod still manages.""" + +from __future__ import annotations + +from dataclasses import replace +from typing import Final + +from setup.performance_test_case import Limit, Ramp, ramped + +from cases.small_bodies_5rps import CASE as SMALL + +RAMP: Final = Ramp(first_rps=10.0, last_rps=70.0, step_rps=10.0, step_seconds=90.0, held_p99_seconds=5.0) + +CASE: Final = replace( + SMALL, + name="small_bodies_ramp", + traffic=ramped(SMALL.traffic, RAMP, seed=13, concurrency=400), + limits=(Limit("sustained_rps", min=20),), +) diff --git a/tests/e2e/perf/setup/datastore.py b/tests/e2e/perf/setup/datastore.py index 89739b6d13e..6ce37da3554 100644 --- a/tests/e2e/perf/setup/datastore.py +++ b/tests/e2e/perf/setup/datastore.py @@ -22,6 +22,8 @@ from typing import TYPE_CHECKING, Final, Literal from pydantic import TypeAdapter +from setup.series import Recorder, Series + if TYPE_CHECKING: import redis @@ -182,13 +184,22 @@ def slowest_statement(statements: tuple[StatementCalls, ...] | None) -> Statemen class ConnectionsSampler: """`pg_stat_activity` client backends on the gateway's database once a second on a background thread, minus - this sampler's own connection, so the peak is what the gateway (and the test's own reads) held open.""" + this sampler's own connection, so the peak is what the gateway (and the test's own reads) held open. + `series` keeps every count with its time for a ramped run's per-step peaks.""" def __init__(self, url: str) -> None: self._url: Final = url self._stop: Final = threading.Event() self._thread: Final = threading.Thread(target=self._run, name="postgres-connections", daemon=True) - self.peak: int | None = None + self._counts: Final[Recorder[int]] = Recorder() + + @property + def peak(self) -> int | None: + return max((sample.value for sample in self._counts.series.samples), default=None) + + @property + def series(self) -> Series[int]: + return self._counts.series def start(self) -> None: self._thread.start() @@ -203,12 +214,13 @@ class ConnectionsSampler: with psycopg.connect(self._url, autocommit=True) as conn: while not self._stop.is_set(): (count,) = _COMMITS_ROW.validate_python(conn.execute(_CONNECTIONS_QUERY).fetchone()) - self.peak = count if self.peak is None else max(self.peak, count) + self._counts.record(count) _ = self._stop.wait(_CONNECTIONS_SAMPLE_SECONDS) class NoConnectionsSampler: peak: None = None + series: None = None def start(self) -> None: return None diff --git a/tests/e2e/perf/setup/locustfile.py b/tests/e2e/perf/setup/locustfile.py index 1c38eae22a1..82a5749becf 100644 --- a/tests/e2e/perf/setup/locustfile.py +++ b/tests/e2e/perf/setup/locustfile.py @@ -4,7 +4,8 @@ time to headers, time to first byte, time to first token (streams), total time, and whether it matched. Runs in its own process (`setup/run.py` starts it), never in the gateway's. Reads `PERF_PLAN_FILE` (JSON list of wire requests), writes `PERF_RECORDS_FILE` (one JSON object per -line), paces to `PERF_RPS` over `PERF_USERS` users, and quits once the last request is recorded. +line), paces to `PERF_RPS` over `PERF_USERS` users (or, when the plan carries `send_at_s`, sends each +request at its scheduled second), and quits once the last request is recorded. """ from __future__ import annotations @@ -35,6 +36,8 @@ _JITTER: Final = 0.25 _RNG: Final = random.Random(_PLAN[0].marker if _PLAN else "empty-plan") _RECORDS: Final = _RECORDS_PATH.open("a", encoding="utf-8") _PROGRESS: Final = PlanProgress(len(_PLAN)) +_SCHEDULED: Final = bool(_PLAN) and _PLAN[0].send_at_s is not None +_RUN_STARTED: Final = time.monotonic() _STOP_FILE: Final = os.environ.get("PERF_STOP_FILE") BLOCK_REASON: Final = "perf-test-block" _TOKEN_MAX_AGE_S: Final = 240.0 @@ -211,7 +214,8 @@ def _authorization(user: FastHttpUser, item: WireRequest) -> str: class PlanUser(FastHttpUser): """One simulated client; all users share the plan and together hold the planned rate. Each user starts at a random point of its first interval and jitters every later interval by up to `_JITTER` either way, so the - users never fire in lockstep; the mean interval is unchanged and `rps` stays the total rate.""" + users never fire in lockstep; the mean interval is unchanged and `rps` stays the total rate. A scheduled + plan (a ramp) is paced by the plan instead: a user takes the next request and sleeps until its second.""" network_timeout = float(os.environ.get("PERF_REQUEST_TIMEOUT", "120")) connection_timeout = 30.0 @@ -219,9 +223,12 @@ class PlanUser(FastHttpUser): def on_start(self) -> None: self._last_run = 0.0 self._last_wait = 0.0 - gevent.sleep(_RNG.uniform(0.0, _INTERVAL_S)) + if not _SCHEDULED: + gevent.sleep(_RNG.uniform(0.0, _INTERVAL_S)) def wait_time(self) -> float: + if _SCHEDULED: + return 0.1 target: Final = _INTERVAL_S * _RNG.uniform(1.0 - _JITTER, 1.0 + _JITTER) now: Final = time.monotonic() run_s: Final = now - self._last_run - self._last_wait if self._last_run else 0.0 @@ -237,6 +244,8 @@ class PlanUser(FastHttpUser): if _PROGRESS.drained: _quit(self) return + if item.send_at_s is not None: + gevent.sleep(max(0.0, _RUN_STARTED + item.send_at_s - time.monotonic())) started_at: Final = time.time() started: Final = time.perf_counter() first_byte_s: float | None = None diff --git a/tests/e2e/perf/setup/observe.py b/tests/e2e/perf/setup/observe.py index 1532a78606e..dccc6e272b3 100644 --- a/tests/e2e/perf/setup/observe.py +++ b/tests/e2e/perf/setup/observe.py @@ -7,14 +7,16 @@ import hashlib import json import statistics from collections import Counter, defaultdict -from collections.abc import Iterable +from collections.abc import Iterable, Mapping from dataclasses import asdict, dataclass -from itertools import chain +from itertools import chain, groupby +from types import MappingProxyType from typing import Final, Literal from e2e_http import NoBody, get_external, post_external, unwrap from memory_readings import RssCapture, RssReading, WorkerKey from pydantic import BaseModel +from typing_extensions import ReadOnly, TypedDict from setup.datastore import ( DbCounters, @@ -27,7 +29,9 @@ from setup.datastore import ( statement_deltas, ) from setup.mocks.reply_common import Provider -from setup.performance_test_case import Limit, Metric, MockLatency, PerformanceTestCase +from setup.performance_test_case import Limit, Metric, MockLatency, PerformanceTestCase, Ramp +from setup.ramp import StepReading, render_steps, sustained_rps +from setup.series import Series from setup.telemetry import GatewayTelemetry from setup.wire import WireRecord @@ -122,6 +126,10 @@ class Observation: telemetry_before: GatewayTelemetry | None = None telemetry_after: GatewayTelemetry | None = None peak_connections: int | None = None + ramp: Ramp | None = None + cpu_series: Series[float] | None = None + connection_series: Series[int] | None = None + telemetry_series: Series[GatewayTelemetry] | None = None @property def telemetry(self) -> GatewayTelemetry | None: @@ -239,6 +247,67 @@ class Observation: post: Final = Counter(call.marker for call in self.guardrail_calls if call.input_type == "response") return sum(1 for record in self.served if pre[record.marker] != 1 or post[record.marker] != 1) + def steps(self) -> tuple[StepReading, ...]: + """One reading per step of a ramped run: the gateway's added latency over the step's requests, the rate + at which requests completed inside the step's window, the step's unexpected statuses, and what the pod + was doing in that window (CPU per completed request, Redis latency, Postgres connections).""" + ramp: Final = self.ramp + if ramp is None or not self.records: + return () + started: Final = min(record.started_at for record in self.records) + added: Final = _by_step(ramp, self.latency_pairs("gateway_added_seconds")) + first_token: Final = _by_step(ramp, self.latency_pairs("time_to_first_token_seconds")) + completed: Final = Counter( + int((record.started_at + record.total_s - started) // ramp.step_seconds) for record in self.records + ) + unexpected: Final = Counter( + ramp.step_of(record.seq) for record in self.records if not record.blocked and record.status != 200 + ) + return tuple( + StepReading( + rps=rate, + planned=count, + completed_rps=completed[index] / ramp.step_seconds, + gateway_added_p99=percentile(added.get(index, ()), 99), + unexpected=unexpected[index], + held_p99_seconds=ramp.held_p99_seconds, + gateway_added_p50=percentile(added.get(index, ()), 50), + time_to_first_token_p99=percentile(first_token.get(index, ()), 99), + **self._window_readings(started + index * ramp.step_seconds, ramp.step_seconds, completed[index]), + ) + for index, (rate, count) in enumerate(zip(ramp.rates, ramp.counts, strict=True)) + ) + + def _window_readings(self, start: float, seconds: float, completed: int) -> WindowReadings: + end: Final = start + seconds + cpu: Final = None if self.cpu_series is None else sum(self.cpu_series.between(start, end)) + conns: Final = None if self.connection_series is None else self.connection_series.between(start, end) + before: Final = None if self.telemetry_series is None else self.telemetry_series.last_before(start) + after: Final = None if self.telemetry_series is None else self.telemetry_series.last_before(end) + return WindowReadings( + cpu_seconds_per_request=None if cpu is None or completed == 0 else cpu / completed, + redis_latency_mean_seconds=None + if before is None or after is None + else after.minus(before).redis_latency.mean, + postgres_peak_connections=max(conns, default=None) if conns else None, + ) + + +class WindowReadings(TypedDict): + cpu_seconds_per_request: ReadOnly[float | None] + redis_latency_mean_seconds: ReadOnly[float | None] + postgres_peak_connections: ReadOnly[int | None] + + +def _by_step(ramp: Ramp, pairs: Iterable[tuple[WireRecord, float]]) -> Mapping[int, tuple[float, ...]]: + ordered: Final = sorted(pairs, key=lambda pair: pair[0].seq) + return MappingProxyType( + { + step: tuple(value for _, value in group) + for step, group in groupby(ordered, key=lambda pair: ramp.step_of(pair[0].seq)) + } + ) + def _text_sha1(text: str, *, intervened: bool) -> str: original: Final = text.removesuffix(" [post]") if intervened else text @@ -415,6 +484,8 @@ def observe( if metric == "cpu_seconds_per_request": cpu: Final = observation.gateway_cpu_seconds return None if cpu is None or not observation.records else cpu / len(observation.records) + if metric == "sustained_rps": + return None if observation.ramp is None else sustained_rps(observation.steps()) return worst_worker_growth_mb(observation.rss_before, observation.rss_after) @@ -424,6 +495,7 @@ class Row: percentile: int | None observed: float | None limit: float | None + floor: float | None = None detail: str = "" breakdown: tuple[Contributor, ...] = () @@ -437,11 +509,22 @@ class Row: @property def over(self) -> bool: - return self.limit is not None and self.observed is not None and self.observed > self.limit + """Past the limit: above a max, or below a floor.""" + if self.observed is None: + return False + return (self.limit is not None and self.observed > self.limit) or ( + self.floor is not None and self.observed < self.floor + ) @property def missing(self) -> bool: - return self.limit is not None and self.observed is None + return (self.limit is not None or self.floor is not None) and self.observed is None + + @property + def bound(self) -> str: + if self.limit is not None: + return f"{self.limit:.4f}" + return "" if self.floor is None else f">= {self.floor:.4f}" REPORTED: Final[tuple[tuple[Metric, int | None], ...]] = ( @@ -476,6 +559,8 @@ def evaluate(case: PerformanceTestCase, observation: Observation) -> tuple[list[ (limit.metric, limit.percentile): limit for limit in case.limits } keys: list[tuple[Metric, int | None]] = list(REPORTED) + if observation.ramp is not None: + keys.append(("sustained_rps", None)) keys.extend(key for key in limits if key not in keys) rows: list[Row] = [] for metric, rank in keys: @@ -486,20 +571,23 @@ def evaluate(case: PerformanceTestCase, observation: Observation) -> tuple[list[ percentile=rank, observed=observe(metric, rank, observation, case.background_cap_exempt_tables), limit=limit.max if limit else None, + floor=limit.min if limit else None, detail=_detail(metric, observation, case.background_cap_exempt_tables), breakdown=breakdown(metric, rank, observation, case.background_cap_exempt_tables), ) ) - exceeded: Final = [ - f"{row.label}: observed {row.observed:.4f} > limit {row.limit:.4f}{f' ({row.detail})' if row.detail else ''}" - if row.over and row.observed is not None and row.limit is not None - else f"{row.label}: limit set but nothing observed" - for row in rows - if row.over or row.missing - ] + exceeded: Final = [_breach(row) for row in rows if row.over or row.missing] return rows, exceeded +def _breach(row: Row) -> str: + if row.observed is None: + return f"{row.label}: limit set but nothing observed" + if row.floor is not None: + return f"{row.label}: observed {row.observed:.4f} < floor {row.floor:.4f}" + return f"{row.label}: observed {row.observed:.4f} > limit {row.bound}{f' ({row.detail})' if row.detail else ''}" + + def _detail(metric: Metric, observation: Observation, exempt_tables: frozenset[str]) -> str: if metric == "postgres_max_statement_per_request": busiest: Final = observation.busiest_awaited_statement @@ -516,12 +604,13 @@ def _detail(metric: Metric, observation: Observation, exempt_tables: frozenset[s def _row_properties(row: Row) -> tuple[tuple[str, float | str], ...]: observed: Final = ((row.property_name, row.observed),) if row.observed is not None else () limit: Final = ((f"{row.property_name}_limit", row.limit),) if row.limit is not None else () + floor: Final = ((f"{row.property_name}_floor", row.floor),) if row.floor is not None else () contributors: Final = ( ((f"{row.property_name}_breakdown", json.dumps([asdict(item) for item in row.breakdown])),) if row.breakdown else () ) - return observed + limit + contributors + return observed + limit + floor + contributors def junit_properties(rows: Iterable[Row]) -> tuple[tuple[str, float | str], ...]: @@ -546,9 +635,17 @@ def render(case: PerformanceTestCase, observation: Observation, rows: list[Row]) ] for row in rows: observed = "n/a" if row.observed is None else f"{row.observed:.4f}" - limit = "" if row.limit is None else f"{row.limit:.4f}" - flag = " OVER" if row.over else " MISSING" if row.missing else "" - lines.append(f" {row.label:32} {observed:>12} {limit:>12}{flag}") + flag = ( + " UNDER" + if row.over and row.floor is not None + else " OVER" + if row.over + else " MISSING" + if row.missing + else "" + ) + lines.append(f" {row.label:32} {observed:>12} {row.bound:>12}{flag}") + lines.extend(render_steps(observation.steps())) lines.extend(_diagnostics(observation)) lines.extend(_datastore_lines(observation)) for record in failures: diff --git a/tests/e2e/perf/setup/performance_test_case.py b/tests/e2e/perf/setup/performance_test_case.py index 3f39b4c38c8..512ef8b911f 100644 --- a/tests/e2e/perf/setup/performance_test_case.py +++ b/tests/e2e/perf/setup/performance_test_case.py @@ -9,8 +9,10 @@ so a case is a literal a reviewer can read top to bottom. from __future__ import annotations +import bisect from collections.abc import Mapping -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace +from itertools import accumulate from typing import Final, Literal from setup.event_loop import EventLoopStalls @@ -36,6 +38,7 @@ Metric = Literal[ "postgres_active_sessions", "postgres_slowest_statement_seconds", "postgres_peak_connections", + "sustained_rps", ] @@ -97,12 +100,50 @@ SHORT: Final = RequestShape("short", input_tokens=(5_000, 100_000), tools=True) LONG: Final = RequestShape("long", input_tokens=(300_000, 1_000_000), tools=True) +@dataclass(frozen=True, slots=True) +class Ramp: + """Rate stepped up through the run: `first_rps`, then `step_rps` more every `step_seconds`, up to `last_rps`. + Each step is scored on its own and the run reports `sustained_rps`, the last step the pod held before the + first one it did not: a step holds when its `gateway_added_seconds` p99 stayed under `held_p99_seconds`, it + completed at least `HELD_SHARE` of its rate and no more than `HELD_UNEXPECTED_SHARE` of its requests came back + with an unexpected status (`setup/ramp.py`). Requests are sent on a schedule drawn from the plan's seed, not + from per-user intervals, so the rate changes on the step boundary.""" + + first_rps: float + last_rps: float + step_rps: float + step_seconds: float + held_p99_seconds: float + + @property + def rates(self) -> tuple[float, ...]: + count: Final = round((self.last_rps - self.first_rps) / self.step_rps) + 1 + return tuple(self.first_rps + index * self.step_rps for index in range(count)) + + @property + def counts(self) -> tuple[int, ...]: + return tuple(round(rate * self.step_seconds) for rate in self.rates) + + @property + def total_requests(self) -> int: + return sum(self.counts) + + @property + def seconds(self) -> float: + return len(self.rates) * self.step_seconds + + def step_of(self, seq: int) -> int: + """Zero-based step of the request with 1-based plan number `seq`.""" + return bisect.bisect_left(tuple(accumulate(self.counts)), seq) + + @dataclass(frozen=True, slots=True) class Traffic: """The plan: `total_requests` drawn from `seed`, sent at `rps` by `concurrency` simulated - users. Shares are fractions of the requests: `endpoints` and `shapes` must each sum to 1, - `intervene_share` picks non-streaming requests the guardrail rewrites, `block_share` picks - requests it blocks.""" + users, or stepped up through a `ramp` (then `rps` is its first step and `total_requests` its + total, see `ramped`). Shares are fractions of the requests: `endpoints` and `shapes` must each + sum to 1, `intervene_share` picks non-streaming requests the guardrail rewrites, `block_share` + picks requests it blocks.""" seed: int rps: float @@ -113,20 +154,33 @@ class Traffic: intervene_share: float = 0.0 block_share: float = 0.0 concurrency: int = 20 + ramp: Ramp | None = None @property def expected_seconds(self) -> float: - return self.total_requests / self.rps + return self.total_requests / self.rps if self.ramp is None else self.ramp.seconds + + +def ramped(base: Traffic, ramp: Ramp, *, seed: int, concurrency: int) -> Traffic: + """`base`'s request mix sent along `ramp` instead of at one rate. `concurrency` caps the requests in + flight, so size it past the deepest step the ramp is expected to reach.""" + return replace( + base, seed=seed, rps=ramp.first_rps, total_requests=ramp.total_requests, concurrency=concurrency, ramp=ramp + ) @dataclass(frozen=True, slots=True) class Limit: - """One pass/fail rule: `metric` (at `percentile`, or the whole-run value) must not exceed `max`. - Metrics without a Limit are still reported.""" + """One pass/fail rule: `metric` (at `percentile`, or the whole-run value) must not exceed `max`, or must + reach `min` for a floor such as `sustained_rps`. Metrics without a Limit are still reported.""" metric: Metric - max: float + max: float | None = None percentile: int | None = None + min: float | None = None + + def __post_init__(self) -> None: + assert (self.max is None) != (self.min is None), f"{self.metric}: a Limit is either a max or a min" @property def label(self) -> str: @@ -201,6 +255,14 @@ class PerformanceTestCase: assert 0 <= self.traffic.stream_share <= 1, f"{self.name}: stream_share out of range" assert self.traffic.intervene_share + self.traffic.block_share <= 1, f"{self.name}: guardrail shares exceed 1" assert self.traffic.total_requests > 0 and self.traffic.rps > 0, f"{self.name}: empty traffic" + ramp: Final = self.traffic.ramp + assert ramp is None or (ramp.first_rps, ramp.total_requests) == ( + self.traffic.rps, + self.traffic.total_requests, + ), f"{self.name}: traffic rps and total_requests must be the ramp's, build it with ramped()" + assert ramp is None or (0 < ramp.first_rps <= ramp.last_rps and ramp.step_rps > 0 and ramp.step_seconds > 0), ( + f"{self.name}: ramp {ramp} does not step up" + ) assert self.model_groups, f"{self.name}: no model groups" assert 0 <= self.mock_latency.jitter_seconds <= self.mock_latency.seconds or self.mock_latency.seconds == 0, ( f"{self.name}: mock latency jitter {self.mock_latency.jitter_seconds} s exceeds {self.mock_latency.seconds} s" diff --git a/tests/e2e/perf/setup/plan.py b/tests/e2e/perf/setup/plan.py index 6b9a5a19662..4f28736a8e9 100644 --- a/tests/e2e/perf/setup/plan.py +++ b/tests/e2e/perf/setup/plan.py @@ -12,6 +12,7 @@ from dataclasses import dataclass from typing import Final, Literal, TypeVar from setup.performance_test_case import Endpoint, PerformanceTestCase, RequestShape, Team +from setup.ramp import send_offsets Flag = Literal["n", "i", "b"] CallerKind = Literal["jwt", "key"] @@ -39,6 +40,7 @@ class PlannedRequest: flag: Flag shape: RequestShape input_tokens: int + send_at_s: float | None = None @property def blocked(self) -> bool: @@ -98,6 +100,7 @@ def build_plan(case: PerformanceTestCase) -> tuple[PlannedRequest, ...]: case.validate() rng: Final = random.Random(case.traffic.seed) callers: Final = _callers_by_group(case) + offsets: Final = None if case.traffic.ramp is None else send_offsets(case.traffic.ramp, case.traffic.seed) planned: list[PlannedRequest] = [] for seq in range(1, case.traffic.total_requests + 1): endpoint = _pick(rng, case.traffic.endpoints) @@ -117,6 +120,7 @@ def build_plan(case: PerformanceTestCase) -> tuple[PlannedRequest, ...]: flag=flag, shape=shape, input_tokens=rng.randint(low, high), + send_at_s=None if offsets is None else offsets[seq - 1], ) ) return tuple(planned) diff --git a/tests/e2e/perf/setup/podcpu.py b/tests/e2e/perf/setup/podcpu.py index 400e5b7d38b..865fa14efe0 100644 --- a/tests/e2e/perf/setup/podcpu.py +++ b/tests/e2e/perf/setup/podcpu.py @@ -19,6 +19,8 @@ from urllib.parse import quote from e2e_http import AuthHeaders, Success, get_external from pydantic import BaseModel +from setup.series import Recorder, Series + SELECTOR_ENV: Final = "E2E_PERF_GATEWAY_POD_SELECTOR" _SERVICE_ACCOUNT: Final = Path("/var/run/secrets/kubernetes.io/serviceaccount") _INTERVAL_S: Final = 5.0 @@ -49,7 +51,8 @@ def cores(quantity: str) -> float: class GatewayCpuSampler: - """Integrates the matching pods' core use over the time it is running; `seconds` is CPU seconds.""" + """Integrates the matching pods' core use over the time it is running; `seconds` is CPU seconds and + `series` the CPU seconds of each sampling interval, stamped when the interval ended.""" def __init__(self, selector: str, namespace: str, token: str, ca_bundle: str) -> None: self._url: Final = ( @@ -60,12 +63,15 @@ class GatewayCpuSampler: self._ca_bundle: Final = ca_bundle self._stop: Final = threading.Event() self._thread: Final = threading.Thread(target=self._run, name="gateway-cpu-sampler", daemon=True) - self._cpu_seconds = 0.0 - self._samples = 0 + self._intervals: Final[Recorder[float]] = Recorder() @property def seconds(self) -> float | None: - return self._cpu_seconds if self._samples > 0 else None + return sum(sample.value for sample in self._intervals.series.samples) if len(self._intervals) else None + + @property + def series(self) -> Series[float]: + return self._intervals.series def _read_cores(self) -> float | None: result: Final = get_external( @@ -81,8 +87,7 @@ class GatewayCpuSampler: now: float = time.monotonic() current: float | None = self._read_cores() if current is not None: - self._cpu_seconds += current * (now - last) - self._samples += 1 + self._intervals.record(current * (now - last)) last = now # rebind-ok: the sampling clock advances every tick def start(self) -> None: @@ -95,6 +100,7 @@ class GatewayCpuSampler: class NoSampler: seconds: Final[None] = None + series: Final[None] = None def start(self) -> None: return None diff --git a/tests/e2e/perf/setup/ramp.py b/tests/e2e/perf/setup/ramp.py new file mode 100644 index 00000000000..217e140d470 --- /dev/null +++ b/tests/e2e/perf/setup/ramp.py @@ -0,0 +1,116 @@ +"""A ramped run step by step: the send schedule its plan carries (`send_offsets`), one `StepReading` per step +and `sustained_rps`, the last step the pod held before the first it did not. `Ramp` itself sits in +`performance_test_case.py` with the rest of the case schema.""" + +from __future__ import annotations + +import random +from collections.abc import Iterable, Iterator +from dataclasses import dataclass +from itertools import chain, takewhile +from typing import Final + +from setup.performance_test_case import Ramp + +HELD_SHARE: Final = 0.9 +"""A step that completed fewer than this share of its rate is a queue that stopped draining.""" +HELD_UNEXPECTED_SHARE: Final = 0.01 +"""A step with more than this share of unexpected statuses broke, whatever its latency.""" +_JITTER: Final = 0.25 + + +def _step_offsets(ramp: Ramp, index: int, count: int, rng: random.Random) -> Iterator[float]: + for slot in range(count): + yield (index + (slot + 0.5 + rng.uniform(-_JITTER, _JITTER)) / count) * ramp.step_seconds + + +def send_offsets(ramp: Ramp, seed: int) -> tuple[float, ...]: + """Seconds after the run starts to send each request of the plan, in plan order: a step's requests spread + evenly over its window, each moved by up to a quarter of a slot either way, drawn from `seed` so a rerun + sends at the same seconds.""" + rng: Final = random.Random(f"ramp-{seed}") + return tuple(chain.from_iterable(_step_offsets(ramp, index, count, rng) for index, count in enumerate(ramp.counts))) + + +@dataclass(frozen=True, slots=True) +class StepReading: + rps: float + planned: int + completed_rps: float + gateway_added_p99: float | None + unexpected: int + held_p99_seconds: float + gateway_added_p50: float | None = None + time_to_first_token_p99: float | None = None + cpu_seconds_per_request: float | None = None + redis_latency_mean_seconds: float | None = None + postgres_peak_connections: int | None = None + """The hold verdict reads the first six. The rest say what the pod was doing on that step: CPU per request + from the pod sampler, Redis latency from the gateway's `/metrics`, connections from `pg_stat_activity`, + all cut to the step's window, so a latency jump can be told apart as CPU, waiting or neither.""" + + @property + def held(self) -> bool: + return ( + self.gateway_added_p99 is not None + and self.gateway_added_p99 <= self.held_p99_seconds + and self.completed_rps >= HELD_SHARE * self.rps + and self.unexpected <= HELD_UNEXPECTED_SHARE * self.planned + ) + + +def sustained_rps(steps: Iterable[StepReading]) -> float: + """The rate of the last step that held before the first that did not, 0 when the first step did not hold.""" + held: Final = tuple(takewhile(lambda step: step.held, steps)) + return held[-1].rps if held else 0.0 + + +def step_properties(steps: Iterable[StepReading]) -> tuple[tuple[str, float], ...]: + """The per-step readings as JUnit properties (`perf_ramp_step__`, every reading the step has) + plus the hold bound, so the dashboard can draw the run's curves next to the sustained rate.""" + rows: Final = tuple(steps) + hold: Final = (("perf_ramp_hold_p99_seconds", rows[0].held_p99_seconds),) if rows else () + return (*hold, *chain.from_iterable(_step_properties(step) for step in rows)) + + +def _step_properties(step: StepReading) -> Iterator[tuple[str, float]]: + prefix: Final = f"perf_ramp_step_{step.rps:g}" + readings: Final[tuple[tuple[str, float | None], ...]] = ( + ("gateway_added_p99", step.gateway_added_p99), + ("completed_rps", step.completed_rps), + ("unexpected", float(step.unexpected)), + ("gateway_added_p50", step.gateway_added_p50), + ("time_to_first_token_p99", step.time_to_first_token_p99), + ("cpu_seconds_per_request", step.cpu_seconds_per_request), + ("redis_latency_mean_seconds", step.redis_latency_mean_seconds), + ( + "postgres_peak_connections", + None if step.postgres_peak_connections is None else float(step.postgres_peak_connections), + ), + ) + return ((f"{prefix}_{name}", value) for name, value in readings if value is not None) + + +def render_steps(steps: Iterable[StepReading]) -> tuple[str, ...]: + rows: Final = tuple(steps) + if not rows: + return () + header: Final = ( + f" steps: {'rps':>6} {'planned':>8} {'completed rps':>14} {'gateway added p50':>18} {'p99':>8}" + f" {'ttft p99':>9} {'unexpected':>11} {'cpu s/req':>10} {'redis mean':>11} {'pg conns':>9}" + ) + return (header, *(_step_line(step) for step in rows)) + + +def _seconds(value: float | None) -> str: + return "n/a" if value is None else f"{value:.4f}" + + +def _step_line(step: StepReading) -> str: + conns: Final = "n/a" if step.postgres_peak_connections is None else str(step.postgres_peak_connections) + return ( + f" {step.rps:6.1f} {step.planned:8d} {step.completed_rps:14.2f} {_seconds(step.gateway_added_p50):>18}" + f" {_seconds(step.gateway_added_p99):>8} {_seconds(step.time_to_first_token_p99):>9} {step.unexpected:11d}" + f" {_seconds(step.cpu_seconds_per_request):>10} {_seconds(step.redis_latency_mean_seconds):>11} {conns:>9}" + f" {'held' if step.held else 'broke'}" + ) diff --git a/tests/e2e/perf/setup/run.py b/tests/e2e/perf/setup/run.py index 272c64754c4..ec2927e90af 100644 --- a/tests/e2e/perf/setup/run.py +++ b/tests/e2e/perf/setup/run.py @@ -63,6 +63,7 @@ def wire_request(item: PlannedRequest, setup: CaseSetup) -> WireRequest: token_source=setup.token_source(item.caller), stream=item.stream, blocked=item.blocked, + send_at_s=item.send_at_s, ) diff --git a/tests/e2e/perf/setup/series.py b/tests/e2e/perf/setup/series.py new file mode 100644 index 00000000000..78e5cf272e1 --- /dev/null +++ b/tests/e2e/perf/setup/series.py @@ -0,0 +1,48 @@ +"""Readings a background sampler takes while the traffic runs, each stamped with the wall clock +(`time.time()`, the clock the Locust records carry in `started_at`), so a ramped run can cut them into its +steps' windows after the fact.""" + +from __future__ import annotations + +import time +from dataclasses import dataclass +from typing import Final, Generic, TypeVar + +T = TypeVar("T") + + +@dataclass(frozen=True, slots=True) +class Sample(Generic[T]): + at: float + value: T + + +@dataclass(frozen=True, slots=True) +class Series(Generic[T]): + samples: tuple[Sample[T], ...] + + def between(self, start: float, end: float) -> tuple[T, ...]: + """The values stamped inside `[start, end)`, in time order.""" + return tuple(sample.value for sample in self.samples if start <= sample.at < end) + + def last_before(self, at: float) -> T | None: + """The newest value stamped before `at`, None when the series starts later.""" + earlier: Final = tuple(sample.value for sample in self.samples if sample.at < at) + return earlier[-1] if earlier else None + + +class Recorder(Generic[T]): + """What a sampler thread appends to; `series` is a frozen copy for the observation.""" + + def __init__(self) -> None: + self._samples: list[Sample[T]] = [] + + def record(self, value: T) -> None: + self._samples.append(Sample(at=time.time(), value=value)) + + @property + def series(self) -> Series[T]: + return Series(samples=tuple(self._samples)) + + def __len__(self) -> int: + return len(self._samples) diff --git a/tests/e2e/perf/setup/telemetry.py b/tests/e2e/perf/setup/telemetry.py index e52bcabafda..185cc788601 100644 --- a/tests/e2e/perf/setup/telemetry.py +++ b/tests/e2e/perf/setup/telemetry.py @@ -1,5 +1,5 @@ -"""The gateway's own `/metrics`, read before and after a run for the numbers nothing outside the gateway -can see. Today that is `litellm_redis_latency`, the histogram customers' Grafana panels show when Redis +"""The gateway's own `/metrics`, read before and after a run (and every few seconds during a ramped run, see +`telemetry_sampler`) for the numbers nothing outside the gateway can see. Today that is `litellm_redis_latency`, the histogram customers' Grafana panels show when Redis slows the gateway down. The gateway's latency and request counts are read per request from its response headers instead (`setup/wire.py`), which gives percentiles where a histogram delta only gives a mean. @@ -10,7 +10,9 @@ on psutil (`proxy_usage.py`).""" from __future__ import annotations import math -from collections.abc import Iterable +import threading +from collections.abc import Generator, Iterable +from contextlib import contextmanager from dataclasses import dataclass from itertools import chain from typing import Final @@ -20,6 +22,10 @@ from prometheus_client.parser import text_string_to_metric_families from prometheus_client.samples import Sample from proxy_client import ProxyClient +from setup.series import Recorder, Series + +_SAMPLE_SECONDS: Final = 10.0 + @dataclass(frozen=True, slots=True) class HistogramTotals: @@ -63,3 +69,42 @@ def read_gateway_telemetry(proxy: ProxyClient) -> GatewayTelemetry | None: case every limit on a gateway-reported metric fails as missing rather than reading a silent zero.""" scrape: Final = proxy.probe("/metrics", params=NoBody()) return parse_gateway_telemetry(scrape.body) if scrape.status_code == 200 else None + + +class TelemetrySampler: + """`/metrics` every `_SAMPLE_SECONDS` on a background thread while traffic runs, so a ramped run can read + the Redis latency of each step's window as the difference between the scrapes around it.""" + + def __init__(self, proxy: ProxyClient) -> None: + self._proxy: Final = proxy + self._stop: Final = threading.Event() + self._thread: Final = threading.Thread(target=self._run, name="gateway-telemetry", daemon=True) + self._scrapes: Final[Recorder[GatewayTelemetry]] = Recorder() + + @property + def series(self) -> Series[GatewayTelemetry]: + return self._scrapes.series + + def _run(self) -> None: + while not self._stop.is_set(): + scrape: GatewayTelemetry | None = read_gateway_telemetry(self._proxy) + if scrape is not None: + self._scrapes.record(scrape) + _ = self._stop.wait(_SAMPLE_SECONDS) + + def start(self) -> None: + self._thread.start() + + def stop(self) -> None: + self._stop.set() + self._thread.join(timeout=30) + + +@contextmanager +def telemetry_sampler(proxy: ProxyClient) -> Generator[TelemetrySampler]: + sampler: Final = TelemetrySampler(proxy) + sampler.start() + try: + yield sampler + finally: + sampler.stop() diff --git a/tests/e2e/perf/setup/upgrade.py b/tests/e2e/perf/setup/upgrade.py index 36efc16e1fe..c9114ccbff2 100644 --- a/tests/e2e/perf/setup/upgrade.py +++ b/tests/e2e/perf/setup/upgrade.py @@ -145,7 +145,9 @@ class UpgradeMonitor: self.clock: Final = clock self.wall_clock: Final = wall_clock self.started: Final = clock() - self.status_limit: Final = next(limit.max for limit in CASE.limits if limit.metric == "unexpected_status_count") + self.status_limit: Final = next( + limit.max for limit in CASE.limits if limit.metric == "unexpected_status_count" and limit.max is not None + ) self.pending: tuple[WireRecord, ...] = () self.processed = 0 self.successes = 0 diff --git a/tests/e2e/perf/setup/wire.py b/tests/e2e/perf/setup/wire.py index 18351c18c6b..a1833fbd44b 100644 --- a/tests/e2e/perf/setup/wire.py +++ b/tests/e2e/perf/setup/wire.py @@ -51,6 +51,8 @@ class WireRequest(BaseModel): token_source: TokenSource | None = None stream: bool blocked: bool + send_at_s: float | None = None + """Seconds after the run starts to send this request, set when the case ramps; otherwise the users pace.""" class WireRecord(BaseModel): diff --git a/tests/e2e/perf/test_perf_cases.py b/tests/e2e/perf/test_perf_cases.py index 3e5c3c522c8..db657a04b00 100644 --- a/tests/e2e/perf/test_perf_cases.py +++ b/tests/e2e/perf/test_perf_cases.py @@ -44,8 +44,9 @@ from setup.performance_test_case import MockLatency, PerformanceTestCase from setup.plan import PlannedRequest, build_plan from setup.podcpu import gateway_cpu_sampler from setup.provision import MOCK_PROVIDERS, build_setup, guardrail_url, mock_url +from setup.ramp import step_properties from setup.run import run_plan -from setup.telemetry import read_gateway_telemetry +from setup.telemetry import read_gateway_telemetry, telemetry_sampler pytestmark = [pytest.mark.perf, pytest.mark.quiet_stack] _SPEND_SAMPLE: Final = 40 @@ -115,7 +116,12 @@ def test_perf_case( redis_before: Final = read_redis_counters() run_started: Final = time.monotonic() traffic_started_at: Final = datetime.now(tz=UTC) - with gateway_cpu_sampler() as cpu, redis_monitor() as monitor, connections_sampler() as connections: + with ( + gateway_cpu_sampler() as cpu, + redis_monitor() as monitor, + connections_sampler() as connections, + telemetry_sampler(proxy) as telemetry, + ): outcome: Final = run_plan(case, plan, setup) run_finished: Final = time.monotonic() db_after: Final = read_db_counters() @@ -151,6 +157,10 @@ def test_perf_case( telemetry_before=telemetry_before, telemetry_after=telemetry_after, peak_connections=connections.peak, + ramp=case.traffic.ramp, + cpu_series=cpu.series, + connection_series=connections.series, + telemetry_series=telemetry.series, ) rows, exceeded_limits = evaluate(case, observation) stall_rows: Final = _event_loop_rows(case, traffic_started_at) if case.event_loop is not None else () @@ -162,6 +172,7 @@ def test_perf_case( item.user_properties.extend( stall_properties(case.event_loop, stall_rows) if case.event_loop is not None else junit_properties(rows) ) + item.user_properties.extend(step_properties(observation.steps())) print(render(case, observation, rows)) if case.event_loop is not None: print(render_stalls(case.event_loop, stall_rows)) diff --git a/tests/e2e_harness/perf/setup/test_ramp.py b/tests/e2e_harness/perf/setup/test_ramp.py new file mode 100644 index 00000000000..6d6c1640b59 --- /dev/null +++ b/tests/e2e_harness/perf/setup/test_ramp.py @@ -0,0 +1,168 @@ +from __future__ import annotations + +import math +from dataclasses import replace +from typing import Final + +from cases.small_bodies_5rps import CASE +from memory_readings import RssCapture +from setup.observe import Observation, evaluate, junit_properties +from setup.performance_test_case import Limit, Ramp, ramped +from setup.plan import build_plan +from setup.ramp import StepReading, send_offsets, step_properties, sustained_rps +from setup.series import Sample, Series +from setup.telemetry import GatewayTelemetry, HistogramTotals +from setup.wire import WireRecord + +NO_RSS: Final = RssCapture(readings=(), failures=()) + + +def test_send_offsets_put_each_step_s_requests_inside_its_window_in_plan_order() -> None: + ramp: Final = Ramp(first_rps=10.0, last_rps=30.0, step_rps=10.0, step_seconds=10.0, held_p99_seconds=5.0) + offsets: Final = send_offsets(ramp, seed=13) + assert ramp.counts == (100, 200, 300) + assert len(offsets) == ramp.total_requests + assert offsets == tuple(sorted(offsets)) + for seq, offset in enumerate(offsets, start=1): + step = ramp.step_of(seq) + assert step * 10.0 <= offset < (step + 1) * 10.0, f"request {seq} at {offset:.2f}s is outside step {step}" + assert send_offsets(ramp, seed=13) == offsets + assert send_offsets(ramp, seed=14) != offsets + + +def test_build_plan_carries_the_schedule_only_for_a_ramped_case() -> None: + ramp: Final = Ramp(first_rps=2.0, last_rps=4.0, step_rps=2.0, step_seconds=5.0, held_p99_seconds=5.0) + ramped_case: Final = replace(CASE, traffic=ramped(CASE.traffic, ramp, seed=3, concurrency=20)) + assert ramped_case.traffic.expected_seconds == 10.0 + assert [item.send_at_s for item in build_plan(ramped_case)] == list(send_offsets(ramp, seed=3)) + assert all(item.send_at_s is None for item in build_plan(CASE)) + + +def _step(rps: float, p99: float, completed_rps: float | None = None, unexpected: int = 0) -> StepReading: + return StepReading( + rps=rps, + planned=int(rps * 90), + completed_rps=rps if completed_rps is None else completed_rps, + gateway_added_p99=p99, + unexpected=unexpected, + held_p99_seconds=5.0, + ) + + +def test_sustained_rps_is_the_last_step_held_before_the_first_that_broke() -> None: + assert sustained_rps((_step(10, 0.5), _step(20, 1.0), _step(30, 9.0), _step(40, 0.7))) == 20 + assert sustained_rps((_step(10, 0.5), _step(20, 1.0, completed_rps=17.0), _step(30, 1.0))) == 10 + assert sustained_rps((_step(10, 0.5), _step(20, 1.0, unexpected=19), _step(30, 1.0))) == 10 + assert sustained_rps((_step(10, 0.5), _step(20, 1.0, unexpected=18), _step(30, 1.0))) == 30 + assert sustained_rps((_step(10, 6.0), _step(20, 1.0))) == 0 + + +def _record(seq: int, started_at: float, total_s: float) -> WireRecord: + return WireRecord( + seq=seq, + marker=f"perfreq-{seq:06d}-n", + caller="key-a1", + endpoint="chat_completions", + stream=False, + blocked=False, + started_at=started_at, + status=200, + header_s=total_s, + first_byte_s=total_s, + first_token_s=None, + total_s=total_s, + text="", + ok=True, + mock_delay_s=1.0, + call_id=None, + error=None, + ) + + +def test_a_floor_fails_when_the_second_step_s_latency_breaks_it_and_the_properties_carry_the_floor() -> None: + ramp: Final = Ramp(first_rps=10.0, last_rps=20.0, step_rps=10.0, step_seconds=100.0, held_p99_seconds=2.0) + case: Final = replace( + CASE, + traffic=ramped(CASE.traffic, ramp, seed=1, concurrency=50), + limits=(Limit("sustained_rps", min=20),), + ) + offsets: Final = send_offsets(ramp, seed=1) + records: Final = tuple( + _record(seq, 1000.0 + offset, 1.5 if ramp.step_of(seq) == 0 else 4.0) + for seq, offset in enumerate(offsets, start=1) + ) + observation: Final = Observation( + records=records, + mock_requests=(), + guardrail_calls=(), + rss_before=NO_RSS, + rss_after=NO_RSS, + wall_seconds=200.0, + ramp=ramp, + ) + first, second = observation.steps() + assert first.held and first.gateway_added_p99 == 0.5 and first.completed_rps > 9.5 + assert not second.held and second.gateway_added_p99 == 3.0 + rows, exceeded = evaluate(case, observation) + assert exceeded == ["sustained_rps: observed 10.0000 < floor 20.0000"] + assert {("perf_sustained_rps", 10.0), ("perf_sustained_rps_floor", 20.0)} <= set(junit_properties(rows)) + assert dict(step_properties((first, second))) == { + "perf_ramp_hold_p99_seconds": 2.0, + "perf_ramp_step_10_gateway_added_p99": 0.5, + "perf_ramp_step_10_completed_rps": first.completed_rps, + "perf_ramp_step_10_unexpected": 0.0, + "perf_ramp_step_10_gateway_added_p50": 0.5, + "perf_ramp_step_20_gateway_added_p99": 3.0, + "perf_ramp_step_20_completed_rps": second.completed_rps, + "perf_ramp_step_20_unexpected": 0.0, + "perf_ramp_step_20_gateway_added_p50": 3.0, + } + assert not evaluate(replace(case, limits=(Limit("sustained_rps", min=10),)), observation)[1] + + +def test_each_step_reads_the_pod_s_cpu_redis_latency_and_connections_from_its_own_window() -> None: + ramp: Final = Ramp(first_rps=10.0, last_rps=20.0, step_rps=10.0, step_seconds=100.0, held_p99_seconds=2.0) + records: Final = tuple( + _record(seq, 1000.0 + offset, 1.5) for seq, offset in enumerate(send_offsets(ramp, seed=1), start=1) + ) + started: Final = min(record.started_at for record in records) + redis: Final = tuple( + Sample(at=started + at, value=GatewayTelemetry(redis_latency=HistogramTotals(total=total, count=count))) + for at, total, count in ((-1.0, 0.0, 0.0), (50.0, 1.0, 100.0), (99.0, 2.0, 200.0), (199.0, 42.0, 300.0)) + ) + observation: Final = Observation( + records=records, + mock_requests=(), + guardrail_calls=(), + rss_before=NO_RSS, + rss_after=NO_RSS, + wall_seconds=200.0, + ramp=ramp, + cpu_series=Series(samples=tuple(Sample(at=started + at, value=5.0) for at in (10.0, 60.0, 110.0, 160.0))), + connection_series=Series( + samples=tuple( + Sample(at=started + at, value=n) for at, n in ((1.0, 12), (50.0, 18), (150.0, 25), (199.0, 20)) + ) + ), + telemetry_series=Series(samples=redis), + ) + first, second = observation.steps() + first_completed: Final = round(first.completed_rps * ramp.step_seconds) + assert first.cpu_seconds_per_request is not None and math.isclose( + first.cpu_seconds_per_request, 10.0 / first_completed + ) + assert first.redis_latency_mean_seconds is not None and math.isclose(first.redis_latency_mean_seconds, 0.01) + assert first.postgres_peak_connections == 18 + assert second.redis_latency_mean_seconds is not None and math.isclose(second.redis_latency_mean_seconds, 0.4) + assert second.postgres_peak_connections == 25 + assert first.time_to_first_token_p99 is None + names: Final = {name for name, _ in step_properties((first, second)) if name.startswith("perf_ramp_step_10_")} + assert names == { + "perf_ramp_step_10_gateway_added_p99", + "perf_ramp_step_10_completed_rps", + "perf_ramp_step_10_unexpected", + "perf_ramp_step_10_gateway_added_p50", + "perf_ramp_step_10_cpu_seconds_per_request", + "perf_ramp_step_10_redis_latency_mean_seconds", + "perf_ramp_step_10_postgres_peak_connections", + } diff --git a/tests/e2e_harness/perf/test_upgrade.py b/tests/e2e_harness/perf/test_upgrade.py index 9a6c9d2af14..3414476ba4c 100644 --- a/tests/e2e_harness/perf/test_upgrade.py +++ b/tests/e2e_harness/perf/test_upgrade.py @@ -304,7 +304,9 @@ def test_invalid_200_is_not_an_allowed_non_200(failed: WireRecord) -> None: def _ceiling(metric: str, percentile: int | None = None) -> float: return next( - limit.max for limit in upgrade_case().limits if limit.metric == metric and limit.percentile == percentile + limit.max + for limit in upgrade_case().limits + if limit.metric == metric and limit.percentile == percentile and limit.max is not None ) From 7ae46f27687f8d927c4fc16d796af68eb8bda273 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Tue, 20 Oct 2026 22:37:21 -0700 Subject: [PATCH 06/14] fix(timing): carry litellm_received_at across to a planted litellm_metadata bucket so the overhead header stays anchored at receive time (#47790) A guardrail hook plants a litellm_metadata key on a chat completions request after the proxy stamped litellm_received_at into metadata. get_litellm_metadata_from_kwargs prefers that key when both exist and only copied the user_api_key_* fields across, so the response timing headers lost their receive anchor and x-litellm-overhead-duration-ms reported 2 ms on /v1/chat/completions against 58 ms on /v1/responses for the same proxy work. The receive stamp now rides along with the user_api_key_* fields. Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/litellm_core_utils/core_helpers.py | 7 +++- .../litellm_core_utils/test_core_helpers.py | 38 ++++++++++++++++++- 2 files changed, 41 insertions(+), 4 deletions(-) diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index b7c6f805233..1d61ef8519a 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -329,11 +329,14 @@ def add_missing_spend_metadata_to_litellm_metadata(litellm_metadata: dict, metad Helper to get litellm metadata for spend tracking PATCH for issue where both `litellm_metadata` and `metadata` are present in the kwargs - and user_api_key values are in 'metadata'. + and user_api_key values are in 'metadata'. The proxy receive stamp rides along for the same + reason: a guardrail hook can plant `litellm_metadata` on a chat request after the proxy stamped + `metadata`, and the response timing headers anchor on that stamp. """ potential_spend_tracking_metadata_substring: Final = "user_api_key" for key, value in dict(metadata).items(): - if potential_spend_tracking_metadata_substring in key: + carry_receive_stamp = key == "litellm_received_at" and key not in litellm_metadata + if potential_spend_tracking_metadata_substring in key or carry_receive_stamp: litellm_metadata[key] = value return litellm_metadata diff --git a/tests/unit/litellm_core_utils/test_core_helpers.py b/tests/unit/litellm_core_utils/test_core_helpers.py index 6459aa98e70..d66c37ba8c2 100644 --- a/tests/unit/litellm_core_utils/test_core_helpers.py +++ b/tests/unit/litellm_core_utils/test_core_helpers.py @@ -1,5 +1,6 @@ """Tests for litellm_core_utils.core_helpers module.""" +import datetime import json import logging import threading @@ -18,13 +19,13 @@ from litellm.litellm_core_utils.core_helpers import ( drop_params_env_flag, drop_params_flag, get_litellm_metadata_from_kwargs, - get_parent_otel_span_from_kwargs, get_or_create_metadata_bucket, + get_parent_otel_span_from_kwargs, get_provider_response_headers_from_hidden_params, map_finish_reason, normalize_drop_params, - process_response_headers, preserve_upstream_non_openai_attributes, + process_response_headers, reconstruct_model_name, redact_nested_match_and_regex_keys, set_provider_response_headers_in_hidden_params, @@ -685,3 +686,36 @@ def test_spend_metadata_merge_reads_request_metadata_that_is_a_read_only_mapping ) assert merged == {"user_api_key_hash": "hash", "model_group": "claude"} + + +class TestReceiveStampSurvivesASecondMetadataBucket: + """A guardrail hook can plant ``litellm_metadata`` on a chat request after the proxy stamped + ``litellm_received_at`` into ``metadata``. The resolved bucket must still carry the stamp, or + ``x-litellm-overhead-duration-ms`` starts at the router call and hides every pre-call cost.""" + + def test_receive_stamp_is_carried_into_litellm_metadata(self): + received_at = datetime.datetime(2026, 10, 21, 5, 0, tzinfo=datetime.timezone.utc) + kwargs = { + "litellm_params": { + "metadata": {"litellm_received_at": received_at, "user_api_key_user_id": "user-1"}, + "litellm_metadata": {"applied_guardrails": ["perf-pre"]}, + } + } + + resolved = get_litellm_metadata_from_kwargs(kwargs) + + assert resolved is kwargs["litellm_params"]["litellm_metadata"] + assert resolved["litellm_received_at"] is received_at + assert resolved["user_api_key_user_id"] == "user-1" + + def test_litellm_metadata_keeps_its_own_receive_stamp(self): + own = datetime.datetime(2026, 10, 21, 5, 0, tzinfo=datetime.timezone.utc) + other = datetime.datetime(2026, 10, 21, 5, 1, tzinfo=datetime.timezone.utc) + kwargs = { + "litellm_params": { + "metadata": {"litellm_received_at": other}, + "litellm_metadata": {"litellm_received_at": own}, + } + } + + assert get_litellm_metadata_from_kwargs(kwargs)["litellm_received_at"] is own From 230f653523fa55aa446e1a5b6469f0871cf7278d Mon Sep 17 00:00:00 2001 From: joshua-berri Date: Tue, 20 Oct 2026 22:50:28 -0700 Subject: [PATCH 07/14] fix(router): normalize native Responses latency per output token (#47727) * fix(router): normalize native Responses latency per output token * fix(router): isolate normalized latency samples during upgrades --------- Co-authored-by: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> --- litellm/router_strategy/lowest_latency.py | 61 +++-- .../router_strategy/test_lowest_latency.py | 222 ++++++++++++++++-- .../test_router_routing_groups.py | 4 +- 3 files changed, 249 insertions(+), 38 deletions(-) diff --git a/litellm/router_strategy/lowest_latency.py b/litellm/router_strategy/lowest_latency.py index 83bb22435e6..28ed062f1ec 100644 --- a/litellm/router_strategy/lowest_latency.py +++ b/litellm/router_strategy/lowest_latency.py @@ -14,7 +14,9 @@ from litellm._internal_context import with_service_target from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.core_helpers import get_parent_otel_span_from_kwargs, safe_divide_seconds +from litellm.responses.utils import ResponseAPILoggingUtils from litellm.router_utils.batch_utils import is_batch_retrieve_call_type +from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.utils import LiteLLMPydanticObjectBase if TYPE_CHECKING: @@ -86,7 +88,7 @@ class LowestLatencyLoggingHandler(CustomLogger): { {model_group}_map: { id: { - "latency": [..] + "latency_normalized": [..] f"{date:hour:minute}" : {"tpm": 34, "rpm": 3} } } @@ -116,11 +118,16 @@ class LowestLatencyLoggingHandler(CustomLogger): final_value: float = response_ms total_tokens = 0 - if isinstance(response_obj, ModelResponse): - _usage: Final = getattr(response_obj, "usage", None) + if isinstance(response_obj, (ModelResponse, ResponsesAPIResponse)): + _usage: Final = ( + ResponseAPILoggingUtils.transform_response_api_usage_to_chat_usage(response_obj.usage) + if isinstance(response_obj, ResponsesAPIResponse) + else getattr(response_obj, "usage", None) + ) if _usage is not None: completion_tokens: Final = _usage.completion_tokens - total_tokens = _usage.total_tokens + if isinstance(response_obj, ModelResponse): + total_tokens = _usage.total_tokens # response_ms is already normalized to float seconds above response_seconds: Final = response_ms @@ -143,10 +150,12 @@ class LowestLatencyLoggingHandler(CustomLogger): request_count_dict[id] = {} ## Latency - if len(request_count_dict[id].get("latency", [])) < self.routing_args.max_latency_list_size: - request_count_dict[id].setdefault("latency", []).append(final_value) + if len(request_count_dict[id].get("latency_normalized", [])) < self.routing_args.max_latency_list_size: + request_count_dict[id].setdefault("latency_normalized", []).append(final_value) else: - request_count_dict[id]["latency"] = request_count_dict[id]["latency"][1:] + [final_value] + request_count_dict[id]["latency_normalized"] = request_count_dict[id]["latency_normalized"][1:] + [ + final_value + ] ## Time to first token if time_to_first_token is not None: @@ -212,7 +221,7 @@ class LowestLatencyLoggingHandler(CustomLogger): { {model_group}_map: { id: { - "latency": [..] + "latency_normalized": [..] f"{date:hour:minute}" : {"tpm": 34, "rpm": 3} } } @@ -225,10 +234,15 @@ class LowestLatencyLoggingHandler(CustomLogger): request_count_dict[id] = {} ## Latency - give 1000s penalty for failing - if len(request_count_dict[id].get("latency", [])) < self.routing_args.max_latency_list_size: - request_count_dict[id].setdefault("latency", []).append(1000.0) + if ( + len(request_count_dict[id].get("latency_normalized", [])) + < self.routing_args.max_latency_list_size + ): + request_count_dict[id].setdefault("latency_normalized", []).append(1000.0) else: - request_count_dict[id]["latency"] = request_count_dict[id]["latency"][1:] + [1000.0] + request_count_dict[id]["latency_normalized"] = request_count_dict[id]["latency_normalized"][ + 1: + ] + [1000.0] await self.router_cache.async_set_cache( key=latency_key, @@ -270,7 +284,7 @@ class LowestLatencyLoggingHandler(CustomLogger): { {model_group}_map: { id: { - "latency": [..] + "latency_normalized": [..] "time_to_first_token_seconds": [..] f"{date:hour:minute}" : {"tpm": 34, "rpm": 3} } @@ -300,11 +314,16 @@ class LowestLatencyLoggingHandler(CustomLogger): final_value: float = response_ms total_tokens = 0 - if isinstance(response_obj, ModelResponse): - _usage: Final = getattr(response_obj, "usage", None) + if isinstance(response_obj, (ModelResponse, ResponsesAPIResponse)): + _usage: Final = ( + ResponseAPILoggingUtils.transform_response_api_usage_to_chat_usage(response_obj.usage) + if isinstance(response_obj, ResponsesAPIResponse) + else getattr(response_obj, "usage", None) + ) if _usage is not None: completion_tokens: Final = _usage.completion_tokens - total_tokens = _usage.total_tokens + if isinstance(response_obj, ModelResponse): + total_tokens = _usage.total_tokens # response_ms is already normalized to float seconds above response_seconds: Final = response_ms @@ -331,10 +350,12 @@ class LowestLatencyLoggingHandler(CustomLogger): request_count_dict[id] = {} ## Latency - if len(request_count_dict[id].get("latency", [])) < self.routing_args.max_latency_list_size: - request_count_dict[id].setdefault("latency", []).append(final_value) + if len(request_count_dict[id].get("latency_normalized", [])) < self.routing_args.max_latency_list_size: + request_count_dict[id].setdefault("latency_normalized", []).append(final_value) else: - request_count_dict[id]["latency"] = request_count_dict[id]["latency"][1:] + [final_value] + request_count_dict[id]["latency_normalized"] = request_count_dict[id]["latency_normalized"][1:] + [ + final_value + ] ## Time to first token if time_to_first_token is not None: @@ -403,7 +424,7 @@ class LowestLatencyLoggingHandler(CustomLogger): ## if healthy deployment not yet used if d["model_info"]["id"] not in all_deployments: all_deployments[d["model_info"]["id"]] = { - "latency": [0], + "latency_normalized": [0], precise_minute: {"tpm": 0, "rpm": 0}, } @@ -443,7 +464,7 @@ class LowestLatencyLoggingHandler(CustomLogger): or _deployment.get("model_info", {}).get("rpm", None) or float("inf") ) - item_latency = item_map.get("latency", []) + item_latency = item_map.get("latency_normalized", []) item_ttft_latency = item_map.get("time_to_first_token_seconds", []) item_rpm = item_map.get(precise_minute, {}).get("rpm", 0) item_tpm = item_map.get(precise_minute, {}).get("tpm", 0) diff --git a/tests/unit/router_strategy/test_lowest_latency.py b/tests/unit/router_strategy/test_lowest_latency.py index 0a3a92691ba..9af5def0d56 100644 --- a/tests/unit/router_strategy/test_lowest_latency.py +++ b/tests/unit/router_strategy/test_lowest_latency.py @@ -17,8 +17,10 @@ from pydantic import ValidationError import litellm from litellm.caching.caching import DualCache from litellm.caching.in_memory_cache import InMemoryCache +from litellm.responses.utils import ResponseAPILoggingUtils from litellm.router import Router from litellm.router_strategy.lowest_latency import LowestLatencyLoggingHandler, RoutingArgs +from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse DEPLOYMENT_ID = "9876" KWARGS = { @@ -43,7 +45,7 @@ def _embedding_response() -> litellm.EmbeddingResponse: def _recorded_latencies(cache: DualCache): cached = cache.get_cache(key="gemini-embedding-001_map") or {} - return cached.get(DEPLOYMENT_ID, {}).get("latency", []) + return cached.get(DEPLOYMENT_ID, {}).get("latency_normalized", []) def test_sync_embedding_latency_is_json_serializable(): @@ -69,7 +71,7 @@ def test_sync_embedding_latency_is_json_serializable(): ) assert latencies[-1] == pytest.approx(2.0) # the exact failure mode from production: redis cache sync json.dumps - json.dumps({"latency": latencies}) + json.dumps({"latency_normalized": latencies}) @pytest.mark.asyncio @@ -94,7 +96,7 @@ async def test_async_embedding_latency_is_json_serializable(): f"raw timedelta leaked into latency list: {latencies}" ) assert latencies[-1] == pytest.approx(3.0) - json.dumps({"latency": latencies}) + json.dumps({"latency_normalized": latencies}) def _chat_response(completion_tokens: int) -> litellm.ModelResponse: @@ -131,7 +133,7 @@ async def test_async_chat_latency_normalized_per_token(): latencies = _recorded_latencies(cache) assert latencies and latencies[-1] == pytest.approx(0.5) # 2s / 4 tokens - json.dumps({"latency": latencies}) + json.dumps({"latency_normalized": latencies}) @pytest.mark.asyncio @@ -151,7 +153,7 @@ async def test_async_chat_zero_completion_tokens_falls_back_to_seconds(): latencies = _recorded_latencies(cache) assert latencies and latencies[-1] == pytest.approx(3.0) assert not isinstance(latencies[-1], timedelta) - json.dumps({"latency": latencies}) + json.dumps({"latency_normalized": latencies}) def test_sync_chat_zero_completion_tokens_falls_back_to_seconds(): @@ -168,7 +170,7 @@ def test_sync_chat_zero_completion_tokens_falls_back_to_seconds(): latencies = _recorded_latencies(cache) assert latencies and latencies[-1] == pytest.approx(2.0) assert not isinstance(latencies[-1], timedelta) - json.dumps({"latency": latencies}) + json.dumps({"latency_normalized": latencies}) MODEL_GROUP = "gpt-4o-mini" @@ -362,8 +364,8 @@ async def test_ttft_percentile_does_not_change_non_streaming_routing(): cache.set_cache( key=f"{MODEL_GROUP}_map", value={ - FAST_TTFT_ID: {"latency": [1.0], "time_to_first_token_seconds": [0.1]}, - SLOW_TTFT_ID: {"latency": [0.2], "time_to_first_token_seconds": [1.5]}, + FAST_TTFT_ID: {"latency_normalized": [1.0], "time_to_first_token_seconds": [0.1]}, + SLOW_TTFT_ID: {"latency_normalized": [0.2], "time_to_first_token_seconds": [1.5]}, }, ) @@ -380,7 +382,7 @@ async def test_ttft_percentile_does_not_change_non_streaming_routing(): @pytest.mark.asyncio @pytest.mark.parametrize( "cached_entry", - [{"latency": []}, {"2026-09-05-15-39": {"tpm": 28, "rpm": 1}}], + [{"latency_normalized": []}, {"2026-09-05-15-39": {"tpm": 28, "rpm": 1}}], ids=["empty_latency_list", "minute_bucket_only_as_cost_based_routing_writes"], ) async def test_async_get_available_deployments_treats_missing_samples_as_zero_latency(cached_entry): @@ -388,7 +390,7 @@ async def test_async_get_available_deployments_treats_missing_samples_as_zero_la handler = LowestLatencyLoggingHandler(router_cache=cache) cache.set_cache( key="gemini-embedding-001_map", - value={DEPLOYMENT_ID: cached_entry, "slower": {"latency": [0.5]}}, + value={DEPLOYMENT_ID: cached_entry, "slower": {"latency_normalized": [0.5]}}, ) healthy_deployments = [ {"model_info": {"id": DEPLOYMENT_ID}, "litellm_params": {}}, @@ -605,7 +607,7 @@ def test_latency_updated(): end_time=end_time, ) latency_key = f"{model_group}_map" - assert end_time - start_time == test_cache.get_cache(key=latency_key)[deployment_id]["latency"][0] + assert end_time - start_time == test_cache.get_cache(key=latency_key)[deployment_id]["latency_normalized"][0] @pytest.mark.usefixtures("frozen_latency_clock") @@ -1002,7 +1004,7 @@ def test_latency_list_trimming_discards_oldest_entry(): latency_key = f"{model_group}_map" cached_data = test_cache.get_cache(key=latency_key) - latency_list = cached_data[deployment_id]["latency"] + latency_list = cached_data[deployment_id]["latency_normalized"] assert len(latency_list) == max_size, f"Expected {max_size} entries, got {len(latency_list)}" @@ -1062,7 +1064,7 @@ async def test_latency_list_trimming_discards_oldest_entry_async(): latency_key = f"{model_group}_map" cached_data = await test_cache.async_get_cache(key=latency_key) - latency_list = cached_data[deployment_id]["latency"] + latency_list = cached_data[deployment_id]["latency_normalized"] assert len(latency_list) == max_size @@ -1129,7 +1131,7 @@ async def test_timeout_penalty_discards_oldest_entry(): latency_key = f"{model_group}_map" cached_data = await test_cache.async_get_cache(key=latency_key) - latency_list = cached_data[deployment_id]["latency"] + latency_list = cached_data[deployment_id]["latency_normalized"] assert len(latency_list) == max_size @@ -1181,7 +1183,7 @@ def test_list_order_preserved_after_multiple_trims(): latency_key = f"{model_group}_map" cached_data = test_cache.get_cache(key=latency_key) - latency_list = cached_data[deployment_id]["latency"] + latency_list = cached_data[deployment_id]["latency_normalized"] assert len(latency_list) == max_size @@ -1423,7 +1425,7 @@ def test_latency_cache_honors_custom_ttl() -> None: latency_key: Final = "ttl-group_map" cached: Final = cache.get_cache(key=latency_key) assert isinstance(cached, dict) - assert cached["deployment"]["latency"] == [1.0] + assert cached["deployment"]["latency_normalized"] == [1.0] assert cache.in_memory_cache.ttl_dict[latency_key] == 103.0 clock.return_value = 104.0 @@ -1461,3 +1463,191 @@ async def test_router_model_group_usage_increases_by_logged_tokens() -> None: updated_usage: Final = await router.get_model_group_usage(model_group=model_group) assert updated_usage[0] == initial_usage[0] + 15 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("async_callback", [False, True]) +@pytest.mark.parametrize("normalized_usage", [False, True]) +@pytest.mark.parametrize("native_seconds", [2, 3]) +async def test_responses_and_chat_latency_use_comparable_token_units( + async_callback: bool, normalized_usage: bool, native_seconds: int +) -> None: + cache: Final = DualCache() + handler: Final = LowestLatencyLoggingHandler(router_cache=cache) + usage: Final = ResponseAPIUsage(input_tokens=10, output_tokens=100, total_tokens=110) + native: Final = ResponsesAPIResponse(id="resp-latency", created_at=0, output=[], usage=usage) + logged_native: Final = ( + native.model_copy(update={"usage": ResponseAPILoggingUtils.transform_response_api_usage_to_chat_usage(usage)}) + if normalized_usage + else native + ) + original_usage: Final = logged_native.usage + assert original_usage is not None + original_fields: Final = original_usage.model_dump() + start: Final = datetime(2026, 1, 1) + for deployment_id, response, seconds in ( + ("native", logged_native, native_seconds), + ("bridge", _chat_response(100), 3), + ): + kwargs: Final = { + "call_type": "aresponses", + "litellm_params": {"metadata": {"model_group": "mixed"}, "model_info": {"id": deployment_id}}, + } + if async_callback: + await handler.async_log_success_event(kwargs, response, start, start + timedelta(seconds=seconds)) + else: + handler.log_success_event(kwargs, response, start, start + timedelta(seconds=seconds)) + recorded: Final = cache.get_cache(key="mixed_map") + assert recorded["native"]["latency_normalized"] == pytest.approx([native_seconds / 100]) + assert recorded["bridge"]["latency_normalized"] == pytest.approx([3 / 100]) + assert logged_native.usage is original_usage + assert logged_native.usage.model_dump() == original_fields + assert native.usage is usage + deployments: Final = [ + {"model_name": "mixed", "model_info": {"id": name}, "litellm_params": {}} for name in ("native", "bridge") + ] + picked: Final = handler.get_available_deployments( + model_group="mixed", healthy_deployments=deployments, request_kwargs={"stream": False, "metadata": {}} + ) + assert picked["model_info"]["id"] in ({"native"} if native_seconds < 3 else {"native", "bridge"}) + native_counts: Final = tuple(value for key, value in recorded["native"].items() if key != "latency_normalized") + assert native_counts == ({"tpm": 0, "rpm": 1},) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("async_callback", [False, True]) +@pytest.mark.parametrize( + "response", + [ + litellm.ModelResponse(), + ResponsesAPIResponse(id="resp-missing", created_at=0, output=[], usage=None), + ResponsesAPIResponse( + id="resp-zero", created_at=0, output=[], + usage=ResponseAPIUsage(input_tokens=10, output_tokens=0, total_tokens=10), + ), + ], +) +async def test_responses_without_output_tokens_keep_elapsed_seconds( + async_callback: bool, response: litellm.ModelResponse | ResponsesAPIResponse +) -> None: + cache: Final = DualCache() + handler: Final = LowestLatencyLoggingHandler(router_cache=cache) + original_usage: Final = getattr(response, "usage", None) + if async_callback: + await handler.async_log_success_event(KWARGS, response, 1.0, 3.0) + else: + handler.log_success_event(KWARGS, response, 1.0, 3.0) + assert _recorded_latencies(cache) == [2.0] + assert getattr(response, "usage", None) is original_usage + + +@pytest.mark.usefixtures("frozen_latency_clock") +@pytest.mark.asyncio +@pytest.mark.parametrize("async_mode", [False, True]) +@pytest.mark.parametrize("preserve_new_samples", [False, True]) +async def test_legacy_latency_does_not_starve_deployments_after_upgrade( + async_mode: bool, preserve_new_samples: bool +) -> None: + cache: Final = DualCache() + handler: Final = LowestLatencyLoggingHandler(router_cache=cache) + minute: Final = FROZEN_NOW.strftime("%Y-%m-%d-%H-%M") + legacy: Final = { + "native": {"latency": [2.0], "time_to_first_token_seconds": [0.1], minute: {"tpm": 0, "rpm": 1}}, + "bridge": {"latency": [0.03], "time_to_first_token_seconds": [0.2], minute: {"tpm": 110, "rpm": 1}}, + } + cache.set_cache(key="mixed_map", value=copy.deepcopy(legacy)) + deployments: Final = [ + {"model_name": "mixed", "model_info": {"id": name}, "litellm_params": {}} for name in ("native", "bridge") + ] + native: Final = ResponsesAPIResponse( + id="resp-upgrade", created_at=0, output=[], + usage=ResponseAPIUsage(input_tokens=10, output_tokens=100, total_tokens=110), + ) + for iteration in range(2): + bridge_kwargs: Final = { + "litellm_params": {"metadata": {"model_group": "mixed"}, "model_info": {"id": "bridge"}}, + } + if async_mode: + await handler.async_log_success_event(bridge_kwargs, _chat_response(100), 0.0, 3.0) + else: + handler.log_success_event(bridge_kwargs, _chat_response(100), 0.0, 3.0) + selected: Final = ( + await handler.async_get_available_deployments(model_group="mixed", healthy_deployments=deployments) + if async_mode + else handler.get_available_deployments(model_group="mixed", healthy_deployments=deployments) + ) + assert selected["model_info"]["id"] == "native" + native_kwargs: Final = { + "litellm_params": {"metadata": {"model_group": "mixed"}, "model_info": {"id": "native"}}, + } + if async_mode: + await handler.async_log_success_event(native_kwargs, native, 0.0, 2.0) + else: + handler.log_success_event(native_kwargs, native, 0.0, 2.0) + recorded: Final = copy.deepcopy(cache.get_cache(key="mixed_map")) + assert recorded["native"]["latency_normalized"] == pytest.approx( + [0.02] * (iteration + 1 if preserve_new_samples else 1) + ) + assert recorded["native"]["latency"] == [2.0] + assert recorded["native"]["time_to_first_token_seconds"] == [0.1] + assert recorded["bridge"]["time_to_first_token_seconds"] == [0.2] + assert recorded["native"][minute] == {"tpm": 0, "rpm": iteration + 2} + assert recorded["bridge"][minute] == {"tpm": 110 + (iteration + 1) * 110, "rpm": iteration + 2} + streaming: Final = handler.get_available_deployments( + model_group="mixed", healthy_deployments=deployments, request_kwargs={"stream": True} + ) + assert streaming["model_info"]["id"] == "native" + cache.set_cache( + key="mixed_map", + value={ + name: {**(recorded[name] if preserve_new_samples else legacy[name]), minute: recorded[name][minute]} + for name in ("native", "bridge") + }, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("async_mode", [False, True]) +async def test_normalized_latency_window_evicts_samples_without_changing_legacy_data(async_mode: bool) -> None: + cache: Final = DualCache() + cache.set_cache(key="gemini-embedding-001_map", value={DEPLOYMENT_ID: {"latency": [2.0]}}) + handler: Final = LowestLatencyLoggingHandler(router_cache=cache, routing_args={"max_latency_list_size": 2}) + for elapsed, expected in ((2.0, [0.02]), (3.0, [0.02, 0.03]), (4.0, [0.03, 0.04])): + if async_mode: + await handler.async_log_success_event(KWARGS, _chat_response(100), 0.0, elapsed) + else: + handler.log_success_event(KWARGS, _chat_response(100), 0.0, elapsed) + recorded: Final = cache.get_cache(key="gemini-embedding-001_map")[DEPLOYMENT_ID] + assert recorded["latency_normalized"] == pytest.approx(expected) + assert recorded["latency"] == [2.0] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("existing_samples", [(), (0.02, 0.03)]) +async def test_timeout_penalty_updates_normalized_latency_and_avoids_failed_deployment( + existing_samples: tuple[float, ...], +) -> None: + cache: Final = DualCache() + handler: Final = LowestLatencyLoggingHandler(router_cache=cache, routing_args={"max_latency_list_size": 2}) + cache.set_cache( + key="mixed_map", + value={ + "native": {"latency": [0.001], "latency_normalized": list(existing_samples)}, + "bridge": {"latency": [2.0], "latency_normalized": [0.04]}, + }, + ) + await handler.async_log_failure_event( + kwargs={ + "exception": litellm.Timeout(message="request timed out", model="test", llm_provider="openai"), + "litellm_params": {"metadata": {"model_group": "mixed"}, "model_info": {"id": "native"}}, + }, + response_obj=None, start_time=0.0, end_time=3.0, + ) + recorded: Final = cache.get_cache(key="mixed_map")["native"] + assert recorded["latency_normalized"] == [*existing_samples[-1:], 1000.0] + assert recorded["latency"] == [0.001] + selected: Final = await handler.async_get_available_deployments( + model_group="mixed", + healthy_deployments=[{"model_info": {"id": name}, "litellm_params": {}} for name in ("native", "bridge")], + ) + assert selected["model_info"]["id"] == "bridge" diff --git a/tests/unit/router_strategy/test_router_routing_groups.py b/tests/unit/router_strategy/test_router_routing_groups.py index 7e4b470f1c1..7235333f59f 100644 --- a/tests/unit/router_strategy/test_router_routing_groups.py +++ b/tests/unit/router_strategy/test_router_routing_groups.py @@ -1249,7 +1249,7 @@ async def test_latency_based_override_learns_from_the_overriding_requests(): ) def samples(deployment_id: str) -> list[float]: - recorded = (router.cache.get_cache("grp_map") or {}).get(deployment_id, {}).get("latency", []) + recorded = (router.cache.get_cache("grp_map") or {}).get(deployment_id, {}).get("latency_normalized", []) return [latency for latency in recorded if latency > 0] async def overriding_call() -> str: @@ -1316,7 +1316,7 @@ async def test_override_matching_a_routing_group_strategy_records_each_request_o deployment_id = response._hidden_params["model_id"] await _async_wait_until(lambda: (router.cache.get_cache("grp_map") or {}).get(deployment_id) is not None) - assert len(router.cache.get_cache("grp_map")[deployment_id]["latency"]) == 1 + assert len(router.cache.get_cache("grp_map")[deployment_id]["latency_normalized"]) == 1 assert request["litellm_logging_obj"].dynamic_success_callbacks is None From 66c8f5b1a9cefaa7c55b758f86f4ab552f2e859e Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 21 Oct 2026 06:59:32 +0000 Subject: [PATCH 08/14] fix(llms): join api_base and route with one shared helper (#47535) * fix(llms): sanitize trailing slashes in provider get_complete_url * test(llms): move trailing-slash url tests into each provider's mapped unit file * fix(llms): strip trailing slashes in Perplexity search and Ollama embedding URLs and run the Topaz tests in CI * fix(llms): strip trailing slashes in search, DashScope, VolcEngine, GigaChat and Gradient AI URL builders * fix(fastcrw): strip a trailing slash from api_base and expect one slash in the GigaChat tests * fix(gigachat): strip a trailing slash from the pass-through api_base * fix(llms): join api_base and route with one shared helper A deployment api_base that ends in / made several get_complete_url builders emit a double slash, which Envoy and other gateways answer with 404. Add join_api_base in url_utils and use it in those builders, with a regression test that drives every get_complete_url under litellm/llms with a trailing-slash api_base * refactor(llms): move per-provider api_base strips onto join_api_base Use the shared join_api_base helper in the search, chat, embedding, image variation, GigaChat and VolcEngine URL builders instead of a per-provider rstrip and suffix check, so one helper owns the single-slash join * ci: run the trailing-slash URL builder test in the llms shard * test(integration): cover trailing-slash api_base across provider URL builders Wire cells for chat, streaming, messages, responses, embeddings, search and pass-through with and without a trailing slash, an owned-proxy module for the flag-gated providers, env bases and a worker kill, two chaos bursts, and a Topaz SDK cell * fix(llms): send every provider route through join_api_base Every place under litellm/llms and litellm/main.py that joined a caller-settable api_base with a route by concatenation or f-string now goes through join_api_base, so an api_base with a trailing slash sends the same upstream path as one without it Routes made only of a caller-supplied id, model, or passthrough path stay on a plain join over an already normalized base, because the helper's suffix dedup would turn a lookup by id into a list or create call The sweep test now also reaches the factory-made JSON provider builders, and a new loopback test drives the public entrypoints with an api_base ending in /gw/ and checks that every request path keeps /gw/ with no double slash * fix(llms): keep the route skip only where a builder already had it join_api_base is a pure join again. The new ensure_api_route skips the route only when the URL path already ends in it, and it is used only where the merge base or the trailing-slash test PR head already skipped. Evals cancel, Predibase, and Azure AI threads append their route again, EdenAI model listing and Gemini RAG ingestion join through join_api_base, and Azure AI Anthropic checks the path so a host named anthropic-* builds /anthropic/v1/messages * fix(azure_ai): read /anthropic from the URL path when building the count_tokens URL * test(llms): run the Gemini RAG trailing-slash case through aingest The sync ingest wrapper needs a current event loop, which an earlier async test in the same worker clears * test(llms): cover the Azure realtime, Bedrock rerank, and oobabooga joins with a trailing-slash api_base * test(integration): cover api_base joins across providers, query strings and named hosts * feat(url_utils): add join_api_base and ensure_api_route helpers * fix(url_utils): keep the scheme slashes when a stripped suffix was the whole host join_api_base and ensure_api_route now keep the // of a base left with only a scheme, so http://v1 with /v1 stripped joins to http://v1/files as before. The Azure v1 client base_url goes through ensure_api_route so a base ending in /openai/v1 no longer doubles it * fix(azure_ai): build the count_tokens URL from the messages URL * fix(azure_ai): match anthropic as a whole path segment in the messages URL * fix(llms): keep main's Azure AI anthropic match and non-string Bedrock endpoints Azure AI chat and messages go back to cutting the base at the first /anthropic, as on main, and count_tokens keeps its path check, so a gateway path like /anthropic-proxy builds the same URLs as main. A base that already ends in /v1/messages still gets count_tokens next to it. A non-string aws_bedrock_runtime_endpoint (an httpx.URL VPC endpoint) is used as its text again in AgentCore, Bedrock responses and the RAGFlow base, instead of falling back to the default host. get_volcengine_base_url and GigaChat's get_api_base strip a trailing slash again for outside callers. The cache-hit wire test is skipped as a known proxy bug, since one cache hit writes two cache_hit spend rows. * fix(llms): drop a slash before the api_base query and tighten the route test checks * test(llms): annotate new test locals as Final and flatten the builder generator * fix(url_utils): keep a scheme-only api_base hostless so it fails before sending --------- Co-authored-by: Raushankumar0720 --- .github/workflows/test-unit.yml | 2 + litellm/litellm_core_utils/url_utils.py | 24 +- .../aiml/image_generation/transformation.py | 3 +- .../aiohttp_openai/chat/transformation.py | 8 +- .../llms/anthropic/batches/transformation.py | 12 +- litellm/llms/anthropic/common_utils.py | 13 +- .../anthropic/count_tokens/transformation.py | 3 +- litellm/llms/anthropic/files/handler.py | 6 +- .../llms/anthropic/files/transformation.py | 14 +- .../pass_through/messages/transformation.py | 6 +- .../llms/anthropic/skills/transformation.py | 10 +- litellm/llms/anthropic/wif.py | 3 +- .../llms/apiserpent/search/transformation.py | 7 +- .../audio_transcription/transformation.py | 4 +- litellm/llms/azure/common_utils.py | 3 +- litellm/llms/azure/realtime/handler.py | 4 +- .../azure/realtime/http_transformation.py | 5 +- litellm/llms/azure/search/transformation.py | 6 +- .../azure/text_to_speech/transformation.py | 11 +- litellm/llms/azure_ai/agents/handler.py | 17 +- .../anthropic/count_tokens/transformation.py | 18 +- .../anthropic/messages_transformation.py | 21 +- .../llms/azure_ai/anthropic/transformation.py | 11 + litellm/llms/azure_ai/embed/handler.py | 10 +- .../image_generation/flux_transformation.py | 4 +- .../image_generation/mai_transformation.py | 9 +- .../azure_ai/vector_stores/transformation.py | 4 +- .../llms/base_llm/decisions/transformation.py | 3 +- litellm/llms/base_llm/evals/transformation.py | 3 +- .../base_llm/realtime/http_transformation.py | 8 +- .../llms/base_llm/skills/transformation.py | 3 +- .../bedrock/chat/agentcore/transformation.py | 5 +- .../chat/chat_completions/transformation.py | 10 +- litellm/llms/bedrock/chat/converse_handler.py | 9 +- .../chat/invoke_agent/transformation.py | 9 +- .../amazon_openai_transformation.py | 10 +- .../base_invoke_transformation.py | 14 +- .../bedrock/claude_platform/common_utils.py | 5 +- litellm/llms/bedrock/common_utils.py | 6 +- .../bedrock/count_tokens/mantle_handler.py | 3 +- .../bedrock/count_tokens/transformation.py | 5 +- litellm/llms/bedrock/embed/embedding.py | 7 +- litellm/llms/bedrock/image_edit/handler.py | 3 +- .../bedrock/image_generation/image_handler.py | 3 +- litellm/llms/bedrock/rerank/handler.py | 3 +- .../llms/bedrock/responses/transformation.py | 6 +- .../bedrock/vector_stores/transformation.py | 6 +- .../bedrock_mantle/messages/transformation.py | 3 +- .../responses/transformation.py | 3 +- .../image_edit/transformation.py | 7 +- .../image_generation/transformation.py | 6 +- .../llms/chatgpt/responses/transformation.py | 4 +- .../llms/cloudflare/chat/transformation.py | 3 +- litellm/llms/cloudflare/common_utils.py | 5 +- litellm/llms/cohere/rerank/transformation.py | 7 +- .../llms/cohere/rerank_v2/transformation.py | 7 +- litellm/llms/cometapi/chat/transformation.py | 11 +- litellm/llms/cometapi/embed/transformation.py | 5 +- .../image_generation/transformation.py | 8 +- litellm/llms/custom_httpx/llm_http_handler.py | 6 +- litellm/llms/dashscope/chat/transformation.py | 6 +- .../llms/dashscope/embed/transformation.py | 6 +- .../llms/dashscope/rerank/transformation.py | 5 +- .../llms/dashscope/videos/transformation.py | 6 +- .../llms/databricks/chat/transformation.py | 6 +- litellm/llms/databricks/common_utils.py | 12 +- .../databricks/decisions/transformation.py | 5 +- .../databricks/responses/transformation.py | 4 +- .../audio_transcription/transformation.py | 4 +- litellm/llms/deepgram/common_utils.py | 3 +- litellm/llms/deepseek/chat/transformation.py | 8 +- .../llms/deepseek/messages/transformation.py | 5 +- .../chat/transformation.py | 6 +- .../llms/duckduckgo/search/transformation.py | 3 +- litellm/llms/e2b/sandbox/transformation.py | 7 +- litellm/llms/edenai/chat/transformation.py | 4 +- litellm/llms/edenai/common_utils.py | 3 +- .../audio_transcription/transformation.py | 8 +- .../text_to_speech/transformation.py | 7 +- litellm/llms/exa_ai/search/transformation.py | 7 +- .../flux_lora_depth_transformation.py | 5 +- .../fal_ai/image_generation/transformation.py | 11 +- litellm/llms/fal_ai/videos/transformation.py | 6 +- litellm/llms/fastcrw/search/transformation.py | 7 +- .../llms/firecrawl/search/transformation.py | 7 +- .../llms/fireworks_ai/chat/transformation.py | 3 +- .../fireworks_ai/rerank/transformation.py | 5 +- .../fireworks_ai/responses/transformation.py | 5 +- litellm/llms/gdc/chat/transformation.py | 3 +- litellm/llms/gemini/agents/transformation.py | 13 +- .../audio_transcription/transformation.py | 4 +- litellm/llms/gemini/common_utils.py | 7 +- litellm/llms/gemini/count_tokens/handler.py | 3 +- litellm/llms/gemini/files/transformation.py | 14 +- .../llms/gemini/image_edit/transformation.py | 6 +- .../gemini/image_generation/transformation.py | 7 +- .../gemini/interactions/transformation.py | 14 +- .../llms/gemini/realtime/transformation.py | 6 +- .../gemini/vector_stores/transformation.py | 15 +- litellm/llms/gemini/videos/transformation.py | 7 +- litellm/llms/gigachat/chat/transformation.py | 4 +- .../llms/gigachat/embedding/transformation.py | 4 +- litellm/llms/gigachat/file_handler.py | 7 +- .../gigachat/passthrough/transformation.py | 5 +- litellm/llms/gigachat/utils.py | 2 +- .../embedding/transformation.py | 9 +- .../github_copilot/messages/transformation.py | 5 +- .../responses/transformation.py | 9 +- .../llms/gradient_ai/chat/transformation.py | 5 +- litellm/llms/heroku/chat/transformation.py | 10 +- .../hosted_vllm/embedding/transformation.py | 10 +- .../hosted_vllm/image_edit/transformation.py | 5 +- .../llms/hosted_vllm/rerank/transformation.py | 3 +- .../hosted_vllm/responses/transformation.py | 5 +- .../transcriptions/transformation.py | 7 +- .../llms/hosted_vllm/videos/transformation.py | 6 +- .../llms/huggingface/chat/transformation.py | 5 +- .../llms/huggingface/rerank/transformation.py | 11 +- .../llms/infinity/embedding/transformation.py | 7 +- .../llms/infinity/rerank/transformation.py | 7 +- .../llms/jina_ai/embedding/transformation.py | 8 +- litellm/llms/langflow/chat/transformation.py | 4 +- litellm/llms/langgraph/chat/transformation.py | 8 +- litellm/llms/lemonade/chat/transformation.py | 5 +- litellm/llms/linkup/search/transformation.py | 7 +- .../image_edit/transformation.py | 4 +- .../image_generation/transformation.py | 4 +- .../litellm_proxy/responses/transformation.py | 6 +- litellm/llms/manus/files/transformation.py | 8 +- .../llms/manus/responses/transformation.py | 6 +- .../milvus/vector_stores/transformation.py | 3 +- litellm/llms/minimax/chat/transformation.py | 12 +- .../llms/minimax/messages/transformation.py | 16 +- .../minimax/text_to_speech/transformation.py | 10 +- .../mistral/audio_speech/transformation.py | 5 +- .../audio_transcription/transformation.py | 4 +- .../llms/mistral/batches/transformation.py | 9 +- litellm/llms/mistral/chat/transformation.py | 4 +- litellm/llms/mistral/files/transformation.py | 8 +- .../llms/modelscope/chat/transformation.py | 8 +- .../image_generation/transformation.py | 6 +- .../mongodb/vector_stores/transformation.py | 3 +- litellm/llms/moonshot/chat/transformation.py | 8 +- litellm/llms/nimble/search/transformation.py | 6 +- .../rerank/ranking_transformation.py | 3 +- .../llms/nvidia_nim/rerank/transformation.py | 3 +- litellm/llms/oci/chat/transformation.py | 3 +- litellm/llms/oci/embed/transformation.py | 3 +- litellm/llms/ollama/chat/transformation.py | 6 +- litellm/llms/ollama/common_utils.py | 5 +- litellm/llms/ollama/completion/handler.py | 12 +- .../llms/ollama/completion/transformation.py | 6 +- litellm/llms/oobabooga/chat/oobabooga.py | 5 +- .../llms/openai/chat/gpt_transformation.py | 3 +- .../llms/openai/containers/transformation.py | 4 +- litellm/llms/openai/evals/transformation.py | 20 +- .../llms/openai/image_edit/transformation.py | 6 +- .../openai/realtime/http_transformation.py | 7 +- .../responses/count_tokens/transformation.py | 6 +- .../llms/openai/responses/transformation.py | 11 +- .../transcriptions/whisper_transformation.py | 5 +- .../vector_store_files/transformation.py | 9 +- .../openai/vector_stores/transformation.py | 9 +- litellm/llms/openai/videos/transformation.py | 16 +- litellm/llms/openai_like/common_utils.py | 6 +- litellm/llms/openai_like/dynamic_config.py | 19 +- .../openai_like/messages/transformation.py | 3 +- .../openai_like/responses/transformation.py | 4 +- .../openrouter/embedding/transformation.py | 11 +- .../openrouter/image_edit/transformation.py | 8 +- .../image_generation/transformation.py | 6 +- .../openrouter/responses/transformation.py | 5 +- .../opensandbox/sandbox/transformation.py | 9 +- .../audio_transcription/transformation.py | 6 +- litellm/llms/ovhcloud/chat/transformation.py | 6 +- .../llms/ovhcloud/embedding/transformation.py | 6 +- .../perplexity/embedding/transformation.py | 3 +- .../perplexity/responses/transformation.py | 3 +- .../llms/perplexity/search/transformation.py | 7 +- .../pg_vector/vector_stores/transformation.py | 9 +- litellm/llms/predibase/chat/transformation.py | 7 +- litellm/llms/ragflow/chat/transformation.py | 4 +- .../ragflow/vector_stores/transformation.py | 14 +- .../llms/recraft/image_edit/transformation.py | 8 +- .../image_generation/transformation.py | 8 +- .../image_generation/transformation.py | 15 +- .../runwayml/text_to_speech/transformation.py | 13 +- .../llms/runwayml/videos/transformation.py | 10 +- .../vector_stores/transformation.py | 3 +- .../sambanova/embedding/transformation.py | 7 +- litellm/llms/sap/chat/transformation.py | 7 +- litellm/llms/sap/credentials.py | 3 +- litellm/llms/sap/embed/transformation.py | 7 +- litellm/llms/scaledown/chat/transformation.py | 5 +- .../audio_transcription/transformation.py | 4 +- .../llms/scaleway/rerank/transformation.py | 5 +- litellm/llms/searxng/search/transformation.py | 8 +- litellm/llms/serper/search/transformation.py | 8 +- litellm/llms/snowflake/chat/transformation.py | 5 +- .../snowflake/embedding/transformation.py | 3 +- litellm/llms/snowflake/utils.py | 8 +- .../soniox/audio_transcription/handler.py | 27 +- .../image_generation/transformation.py | 6 +- litellm/llms/tavily/search/transformation.py | 7 +- litellm/llms/tencent/chat/transformation.py | 6 +- .../llms/tencent/messages/transformation.py | 3 +- .../topaz/image_variations/transformation.py | 4 +- .../vercel_ai_gateway/chat/transformation.py | 3 +- .../embedding/transformation.py | 9 +- .../gemini_transcribe_transformation.py | 5 +- .../audio_transcription/transformation.py | 5 +- litellm/llms/vertex_ai/batches/handler.py | 3 +- .../llms/vertex_ai/files/transformation.py | 5 +- .../vertex_imagen_transformation.py | 11 +- .../vertex_ai/interactions/transformation.py | 8 +- .../llms/vertex_ai/realtime/transformation.py | 6 +- .../text_to_speech/transformation.py | 11 +- .../vector_stores/rag_api/transformation.py | 3 +- .../count_tokens/handler.py | 13 +- litellm/llms/vertex_ai/vertex_llm_base.py | 43 +- .../llms/vertex_ai/videos/transformation.py | 14 +- .../volcengine/embedding/transformation.py | 9 +- .../volcengine/responses/transformation.py | 10 +- .../llms/voyage/embedding/transformation.py | 8 +- .../embedding/transformation_contextual.py | 8 +- .../embedding/transformation_multimodal.py | 8 +- litellm/llms/voyage/rerank/transformation.py | 5 +- .../audio_transcription/transformation.py | 4 +- litellm/llms/watsonx/chat/transformation.py | 3 +- .../llms/watsonx/completion/transformation.py | 3 +- litellm/llms/watsonx/embed/transformation.py | 3 +- litellm/llms/watsonx/rerank/transformation.py | 3 +- .../xai/audio_transcription/transformation.py | 3 +- litellm/llms/xai/batches/transformation.py | 9 +- litellm/llms/xai/common_utils.py | 3 +- litellm/llms/xai/files/transformation.py | 10 +- litellm/llms/xai/responses/transformation.py | 6 +- litellm/llms/you_com/search/transformation.py | 3 +- litellm/main.py | 24 +- litellm/rag/ingestion/gemini_ingestion.py | 7 +- .../providers/_api_base_url_join_support.py | 163 ++++++++ .../providers/test_api_base_url_join_chaos.py | 176 +++++++++ ...test_api_base_url_join_host_owned_proxy.py | 140 +++++++ .../providers/test_api_base_url_join_wire.py | 373 ++++++++++++++++++ .../unit/litellm_core_utils/test_url_utils.py | 81 ++++ .../realtime/test_azure_realtime_handler.py | 12 + .../llms/azure/test_azure_common_utils.py | 18 + .../unit/llms/azure_ai/agents/test_handler.py | 23 +- ...e_anthropic_count_tokens_transformation.py | 22 +- ...azure_anthropic_messages_transformation.py | 11 + .../test_agentcore_transformation.py | 14 + .../evals/test_openai_evals_transformation.py | 11 + .../test_api_base_trailing_slash_routes.py | 119 ++++++ .../test_get_complete_url_trailing_slash.py | 278 +++++++++++++ .../llms/test_predibase_transformation.py | 14 + .../llms/vertex_ai/test_vertex_llm_base.py | 16 + 256 files changed, 2260 insertions(+), 922 deletions(-) create mode 100644 tests/integration/providers/_api_base_url_join_support.py create mode 100644 tests/integration/providers/test_api_base_url_join_chaos.py create mode 100644 tests/integration/providers/test_api_base_url_join_host_owned_proxy.py create mode 100644 tests/integration/providers/test_api_base_url_join_wire.py create mode 100644 tests/unit/llms/test_api_base_trailing_slash_routes.py create mode 100644 tests/unit/llms/test_get_complete_url_trailing_slash.py diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index c593acbe8fb..2276efc4e34 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -216,10 +216,12 @@ jobs: - shard: llms-providers test-path: |- + tests/unit/llms/test_api_base_trailing_slash_routes.py tests/unit/llms/test_cache_control_and_reasoning.py tests/unit/llms/test_custom_llm.py tests/unit/llms/test_file_content_block.py tests/unit/llms/test_file_search_responses.py + tests/unit/llms/test_get_complete_url_trailing_slash.py tests/unit/llms/test_lifecycle_fix.py tests/unit/llms/test_oss_decision.py tests/unit/llms/test_polling_url_origin_match.py diff --git a/litellm/litellm_core_utils/url_utils.py b/litellm/litellm_core_utils/url_utils.py index 43e16599bf6..4a1eaa87d3a 100644 --- a/litellm/litellm_core_utils/url_utils.py +++ b/litellm/litellm_core_utils/url_utils.py @@ -23,7 +23,7 @@ import asyncio import socket from ipaddress import ip_address, ip_network from typing import Any, Final, Protocol -from urllib.parse import quote, urlparse, urlunparse +from urllib.parse import quote, urlparse, urlsplit, urlunparse import httpx from typing_extensions import ReadOnly, TypedDict @@ -129,6 +129,28 @@ def encode_url_path_segments(value: object, *, field_name: str = "path") -> str: return "/".join(encode_url_path_segment(segment, field_name=field_name) for segment in value_str.split("/")) +def _strip_trailing_slashes_keeping_the_scheme(api_base: str) -> str: + stripped: Final = api_base.rstrip("/") + parts: Final = urlsplit(stripped) + if parts.scheme and not (parts.netloc or parts.path or parts.query or parts.fragment): + return api_base + if not parts.query or not parts.path.endswith("/"): + return stripped + url_before_query, _, query = stripped.partition("?") + return f"{url_before_query.rstrip('/')}?{query}" + + +def join_api_base(api_base: str, route: str) -> str: + return f"{_strip_trailing_slashes_keeping_the_scheme(api_base)}/{route.lstrip('/')}" + + +def ensure_api_route(api_base: str, route: str) -> str: + base: Final = _strip_trailing_slashes_keeping_the_scheme(api_base) + if urlsplit(base).path.endswith(f"/{route.lstrip('/')}"): + return base + return join_api_base(base, route) + + def _is_blocked_ip(addr: str) -> bool: """Return True for any IP not safe to reach from a user-supplied URL. diff --git a/litellm/llms/aiml/image_generation/transformation.py b/litellm/llms/aiml/image_generation/transformation.py index 931038689cc..bb981983e98 100644 --- a/litellm/llms/aiml/image_generation/transformation.py +++ b/litellm/llms/aiml/image_generation/transformation.py @@ -4,6 +4,7 @@ from typing import TYPE_CHECKING, Any, Final import httpx from pydantic import ConfigDict, TypeAdapter +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.image_generation.transformation import ( BaseImageGenerationConfig, ) @@ -121,7 +122,7 @@ class AimlImageGenerationConfig(BaseImageGenerationConfig): complete_url = complete_url.rstrip("/") # Strip /v1 suffix if present since IMAGE_GENERATION_ENDPOINT already includes v1 complete_url = complete_url.removesuffix("/v1") - complete_url = f"{complete_url}/{self.IMAGE_GENERATION_ENDPOINT}" + complete_url = join_api_base(complete_url, self.IMAGE_GENERATION_ENDPOINT) return complete_url def validate_environment( diff --git a/litellm/llms/aiohttp_openai/chat/transformation.py b/litellm/llms/aiohttp_openai/chat/transformation.py index dba7a7405ab..74f782b56e7 100644 --- a/litellm/llms/aiohttp_openai/chat/transformation.py +++ b/litellm/llms/aiohttp_openai/chat/transformation.py @@ -13,6 +13,7 @@ from typing import TYPE_CHECKING, Any, Final from aiohttp import ClientResponse from pydantic import ConfigDict, TypeAdapter +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.openai_like.chat.transformation import OpenAILikeChatConfig from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import Choices, ModelResponse @@ -44,12 +45,7 @@ class AiohttpOpenAIChatConfig(OpenAILikeChatConfig): Ensure - /v1/chat/completions is at the end of the url """ - if api_base is None: - api_base = "https://api.openai.com" - - if not api_base.endswith("/chat/completions"): - api_base += "/chat/completions" - return api_base + return ensure_api_route("https://api.openai.com" if api_base is None else api_base, "chat/completions") def validate_environment( self, diff --git a/litellm/llms/anthropic/batches/transformation.py b/litellm/llms/anthropic/batches/transformation.py index b95b734c761..da6c2871289 100644 --- a/litellm/llms/anthropic/batches/transformation.py +++ b/litellm/llms/anthropic/batches/transformation.py @@ -7,7 +7,7 @@ import httpx from httpx import Headers, Response from typing_extensions import ReadOnly, TypedDict -from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.litellm_core_utils.url_utils import encode_url_path_segment, ensure_api_route, join_api_base from litellm.llms.base_llm.batches.transformation import BaseBatchesConfig from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.types.llms.openai import AllMessageValues, CreateBatchRequest @@ -106,10 +106,8 @@ class AnthropicBatchesConfig(BaseBatchesConfig): data: CreateBatchRequest, ) -> str: """Get the complete URL for batch creation request.""" - api_base = api_base or self.anthropic_model_info.get_api_base(api_base) - if not api_base.endswith("/v1/messages/batches"): - api_base = f"{api_base.rstrip('/')}/v1/messages/batches" - return api_base + api_base = self.anthropic_model_info.resolve_anthropic_api_base(api_base) + return ensure_api_route(api_base, "v1/messages/batches") def transform_create_batch_request( self, @@ -158,9 +156,9 @@ class AnthropicBatchesConfig(BaseBatchesConfig): Returns: Complete URL for Anthropic batch retrieval: {api_base}/v1/messages/batches/{batch_id} """ - api_base = api_base or self.anthropic_model_info.get_api_base(api_base) + api_base = self.anthropic_model_info.resolve_anthropic_api_base(api_base) encoded_batch_id: Final = encode_url_path_segment(batch_id, field_name="batch_id") - return f"{api_base.rstrip('/')}/v1/messages/batches/{encoded_batch_id}" + return join_api_base(api_base, f"v1/messages/batches/{encoded_batch_id}") def transform_retrieve_batch_request( self, diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index 6f08d6b0d61..1f5e4e7de39 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -35,6 +35,7 @@ from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import parts_of, ) from litellm.litellm_core_utils.reasoning_effort_utils import non_numeric_thinking_budget_error +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.anthropic.wif import ( aget_anthropic_wif_token, anthropic_base_without_chat_suffix, @@ -468,7 +469,7 @@ def _fetch_anthropic_models_page( # which a read-only one cannot support query: Final = f"?after_id={quote(after_id)}" if after_id else "" response: Final = litellm.module_level_client.get( - url=f"{api_base}/v1/models{query}", + url=f"{join_api_base(api_base, 'v1/models')}{query}", headers=headers, follow_redirects=False, ) @@ -1406,7 +1407,7 @@ class AnthropicModelInfo(BaseLLMModelInfo): return {**caller_headers, **anthropic_headers} @staticmethod - def get_api_base(api_base: str | None = None) -> str | None: + def resolve_anthropic_api_base(api_base: str | None = None) -> str: from litellm.secret_managers.main import get_secret_str return ( @@ -1416,6 +1417,10 @@ class AnthropicModelInfo(BaseLLMModelInfo): or "https://api.anthropic.com" ) + @staticmethod + def get_api_base(api_base: str | None = None) -> str | None: + return AnthropicModelInfo.resolve_anthropic_api_base(api_base) + @staticmethod def get_api_key(api_key: str | None = None) -> str | None: """An empty or whitespace-only key counts as unset: it can never authenticate anything, and @@ -1533,14 +1538,14 @@ class AnthropicModelInfo(BaseLLMModelInfo): api_base: str | None, litellm_params: Mapping[str, object] | None, ) -> list[str]: # mutable-ok: matches get_models' list[str] contract shared by every provider override - resolved_api_base: Final = AnthropicModelInfo.get_api_base(api_base) + resolved_api_base: Final = AnthropicModelInfo.resolve_anthropic_api_base(api_base) auth_header: Final = AnthropicModelInfo.get_auth_header( api_key, resolved_api_base, litellm_params=litellm_params, allow_workload_identity=config_allows_workload_identity(self), ) - if resolved_api_base is None or auth_header is None: + if auth_header is None: raise ValueError( "ANTHROPIC_API_BASE/ANTHROPIC_BASE_URL or ANTHROPIC_API_KEY/ANTHROPIC_AUTH_TOKEN (or workload " "identity federation via ANTHROPIC_FEDERATION_RULE_ID/ANTHROPIC_ORGANIZATION_ID/" diff --git a/litellm/llms/anthropic/count_tokens/transformation.py b/litellm/llms/anthropic/count_tokens/transformation.py index 592b5fb0f00..c2d87c641be 100644 --- a/litellm/llms/anthropic/count_tokens/transformation.py +++ b/litellm/llms/anthropic/count_tokens/transformation.py @@ -25,6 +25,7 @@ from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import anthropic_system_blocks, split_leading_system_run, ) +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.anthropic.common_utils import merge_anthropic_beta_headers from litellm.llms.anthropic.wif import resolve_anthropic_base from litellm.types.llms.openai import AllMessageValues, ChatCompletionImageObject, ValidUserMessageContentTypes @@ -317,7 +318,7 @@ class AnthropicCountTokensConfig: Returns: The endpoint URL for the CountTokens API """ - return resolve_anthropic_base(api_base) + "/v1/messages/count_tokens" + return join_api_base(resolve_anthropic_base(api_base), "v1/messages/count_tokens") def transform_request_to_count_tokens( self, diff --git a/litellm/llms/anthropic/files/handler.py b/litellm/llms/anthropic/files/handler.py index f1f9c204b58..1a793e32317 100644 --- a/litellm/llms/anthropic/files/handler.py +++ b/litellm/llms/anthropic/files/handler.py @@ -10,7 +10,7 @@ import litellm from litellm._logging import verbose_logger from litellm._uuid import uuid from litellm.litellm_core_utils.litellm_logging import Logging -from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.types.llms.openai import ( FileContentRequest, @@ -74,7 +74,7 @@ class AnthropicFilesHandler: batch_id = file_id # Get Anthropic API credentials - api_base = self.anthropic_model_info.get_api_base(api_base) + api_base = self.anthropic_model_info.resolve_anthropic_api_base(api_base) auth_header: Final = await self.anthropic_model_info.aget_auth_header( api_key, api_base, litellm_params=litellm_params, allow_workload_identity=True ) @@ -84,7 +84,7 @@ class AnthropicFilesHandler: # Construct the Anthropic batch results URL encoded_batch_id: Final = encode_url_path_segment(batch_id, field_name="batch_id") - results_url: Final = f"{api_base.rstrip('/')}/v1/messages/batches/{encoded_batch_id}/results" + results_url: Final = join_api_base(api_base, f"v1/messages/batches/{encoded_batch_id}/results") # Prepare headers headers: Final = { diff --git a/litellm/llms/anthropic/files/transformation.py b/litellm/llms/anthropic/files/transformation.py index 8433235dde4..db14d3366ff 100644 --- a/litellm/llms/anthropic/files/transformation.py +++ b/litellm/llms/anthropic/files/transformation.py @@ -22,7 +22,7 @@ from openai.types.file_deleted import FileDeleted from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.prompt_templates.common_utils import extract_file_data -from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.files.transformation import ( BaseFilesConfig, @@ -80,7 +80,7 @@ class AnthropicFilesConfig(BaseFilesConfig): stream: bool | None = None, ) -> str: api_base = AnthropicModelInfo.get_api_base(api_base) or ANTHROPIC_FILES_API_BASE - return f"{api_base.rstrip('/')}/v1/files" + return join_api_base(api_base, "v1/files") def get_error_class( self, @@ -226,7 +226,7 @@ class AnthropicFilesConfig(BaseFilesConfig): ) -> tuple[str, dict]: api_base: Final = AnthropicModelInfo.get_api_base(litellm_params.get("api_base")) or ANTHROPIC_FILES_API_BASE encoded_file_id: Final = encode_url_path_segment(file_id, field_name="file_id") - return f"{api_base.rstrip('/')}/v1/files/{encoded_file_id}", {} + return join_api_base(api_base, f"v1/files/{encoded_file_id}"), {} def transform_retrieve_file_response( self, @@ -245,7 +245,7 @@ class AnthropicFilesConfig(BaseFilesConfig): ) -> tuple[str, dict]: api_base: Final = AnthropicModelInfo.get_api_base(litellm_params.get("api_base")) or ANTHROPIC_FILES_API_BASE encoded_file_id: Final = encode_url_path_segment(file_id, field_name="file_id") - return f"{api_base.rstrip('/')}/v1/files/{encoded_file_id}", {} + return join_api_base(api_base, f"v1/files/{encoded_file_id}"), {} def transform_delete_file_response( self, @@ -264,7 +264,7 @@ class AnthropicFilesConfig(BaseFilesConfig): litellm_params: dict, ) -> tuple[str, dict]: api_base: Final = AnthropicModelInfo.get_api_base(litellm_params.get("api_base")) or ANTHROPIC_FILES_API_BASE - url: Final = f"{api_base.rstrip('/')}/v1/files" + url: Final = join_api_base(api_base, "v1/files") params: Final[dict[str, str]] = {} if purpose: params["purpose"] = purpose @@ -299,8 +299,8 @@ class AnthropicFilesConfig(BaseFilesConfig): api_base: Final = AnthropicModelInfo.get_api_base(litellm_params.get("api_base")) or ANTHROPIC_FILES_API_BASE encoded_file_id: Final = encode_url_path_segment(file_id, field_name="file_id") if file_id.startswith(ANTHROPIC_MESSAGE_BATCH_ID_PREFIX): - return f"{api_base.rstrip('/')}/v1/messages/batches/{encoded_file_id}/results", {} - return f"{api_base.rstrip('/')}/v1/files/{encoded_file_id}/content", {} + return join_api_base(api_base, f"v1/messages/batches/{encoded_file_id}/results"), {} + return join_api_base(api_base, f"v1/files/{encoded_file_id}/content"), {} def transform_file_content_response( self, diff --git a/litellm/llms/anthropic/pass_through/messages/transformation.py b/litellm/llms/anthropic/pass_through/messages/transformation.py index a71cd72b768..9f3aa0130d1 100644 --- a/litellm/llms/anthropic/pass_through/messages/transformation.py +++ b/litellm/llms/anthropic/pass_through/messages/transformation.py @@ -9,6 +9,7 @@ from litellm.anthropic_beta_headers_manager import get_unsupported_headers, is_b from litellm.exceptions import AuthenticationError from litellm.litellm_core_utils.core_helpers import drops_unsupported_params from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.anthropic.pass_through.convertible_blocks import drop_non_string_tool_result_texts from litellm.llms.anthropic.pass_through.safeguards import proxy_serves_safeguards from litellm.llms.base_llm.anthropic_messages.transformation import ( @@ -287,10 +288,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): litellm_params: dict, stream: bool | None = None, ) -> str: - api_base = AnthropicModelInfo.get_api_base(api_base) or "https://api.anthropic.com" - if not api_base.endswith("/v1/messages"): - api_base = f"{api_base}/v1/messages" - return api_base + return ensure_api_route(AnthropicModelInfo.get_api_base(api_base) or "https://api.anthropic.com", "v1/messages") def validate_anthropic_messages_environment( self, diff --git a/litellm/llms/anthropic/skills/transformation.py b/litellm/llms/anthropic/skills/transformation.py index 6d419a56889..8ea9d192c24 100644 --- a/litellm/llms/anthropic/skills/transformation.py +++ b/litellm/llms/anthropic/skills/transformation.py @@ -9,7 +9,7 @@ import httpx from pydantic import TypeAdapter from litellm._logging import verbose_logger -from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.llms.base_llm.skills.transformation import ( BaseSkillsAPIConfig, LiteLLMLoggingObj, @@ -75,13 +75,11 @@ class AnthropicSkillsConfig(BaseSkillsAPIConfig): """Get complete URL for Anthropic Skills API""" from litellm.llms.anthropic.common_utils import AnthropicModelInfo - if api_base is None: - api_base = AnthropicModelInfo.get_api_base() - + resolved_api_base: Final = AnthropicModelInfo.get_api_base(api_base) or "https://api.anthropic.com" if skill_id: encoded_skill_id: Final = encode_url_path_segment(skill_id, field_name="skill_id") - return f"{api_base}/v1/skills/{encoded_skill_id}" - return f"{api_base}/v1/{endpoint}" + return join_api_base(resolved_api_base, f"v1/skills/{encoded_skill_id}") + return join_api_base(resolved_api_base, f"v1/{endpoint}") def transform_create_skill_request( self, diff --git a/litellm/llms/anthropic/wif.py b/litellm/llms/anthropic/wif.py index b47c4f28eb2..ccb81c1244f 100644 --- a/litellm/llms/anthropic/wif.py +++ b/litellm/llms/anthropic/wif.py @@ -15,6 +15,7 @@ from typing_extensions import assert_never import litellm from litellm._logging import verbose_logger +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.auth.client_credentials import keycloak_assertion_source from litellm.llms.base_llm.auth.identity_source import ( AnthropicIdentitySourceKind, @@ -324,7 +325,7 @@ def anthropic_internal_issuer_jwks(credential_values: Mapping[str, object]) -> A def build_anthropic_wif_spec(params: AnthropicWifParams, api_base: str) -> TokenExchangeSpec: return TokenExchangeSpec( - token_url=api_base.rstrip("/") + ANTHROPIC_TOKEN_EXCHANGE_PATH, + token_url=join_api_base(api_base, ANTHROPIC_TOKEN_EXCHANGE_PATH), assertion_ref=params.assertion_ref, assertion_field="assertion", static_body=MappingProxyType( diff --git a/litellm/llms/apiserpent/search/transformation.py b/litellm/llms/apiserpent/search/transformation.py index bb8abff7272..faf1903ffe8 100644 --- a/litellm/llms/apiserpent/search/transformation.py +++ b/litellm/llms/apiserpent/search/transformation.py @@ -14,6 +14,7 @@ from urllib.parse import urlencode import httpx from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.apiserpent.search.defaults import ( DEEP_SEARCH_PATH, NUM_MAX, @@ -77,13 +78,11 @@ class APISerpentSearchConfig(BaseSearchConfig): Build the search URL. APISerpent uses GET, so the transformed request is serialized into the query string. The endpoint path (quick vs deep) is always applied; an ``api_base`` / ``APISERPENT_API_BASE`` override only - changes the host. The ``endswith`` guard keeps this idempotent, since the + changes the host. ``ensure_api_route`` keeps this idempotent, since the handler re-invokes this method with the already-resolved URL as api_base. """ - base = (api_base or get_secret_str("APISERPENT_API_BASE") or APISERPENT_BASE).rstrip("/") path: Final = DEEP_SEARCH_PATH if self._is_deep_search(optional_params) else QUICK_SEARCH_PATH - if not base.endswith(path): - base = f"{base}{path}" + base: Final = ensure_api_route(api_base or get_secret_str("APISERPENT_API_BASE") or APISERPENT_BASE, path) if data and isinstance(data, dict) and APISERPENT_PARAMS_KEY in data: query_string: Final = urlencode(data[APISERPENT_PARAMS_KEY], doseq=True) diff --git a/litellm/llms/azure/audio_transcription/transformation.py b/litellm/llms/azure/audio_transcription/transformation.py index 858f122f65a..a5fde7f7158 100644 --- a/litellm/llms/azure/audio_transcription/transformation.py +++ b/litellm/llms/azure/audio_transcription/transformation.py @@ -13,6 +13,7 @@ import httpx from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.audio_utils.utils import process_audio_file +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.audio_transcription.transformation import ( AudioTranscriptionRequestData, BaseAudioTranscriptionConfig, @@ -110,7 +111,7 @@ class AzureSpeechAudioTranscriptionConfig(BaseAudioTranscriptionConfig): "language": optional_params.get("language", self.DEFAULT_LANGUAGE), "format": self._get_azure_response_format(optional_params.get("response_format")), } - return f"{base_url}{self.STT_ENDPOINT_PATH}?{urlencode(query_params)}" + return f"{join_api_base(base_url, self.STT_ENDPOINT_PATH)}?{urlencode(query_params)}" def transform_audio_transcription_request( self, @@ -153,7 +154,6 @@ class AzureSpeechAudioTranscriptionConfig(BaseAudioTranscriptionConfig): ) def _resolve_stt_base_url(self, api_base: str) -> str: - api_base = api_base.rstrip("/") parsed_url: Final = urlparse(api_base) hostname: Final = parsed_url.hostname or "" diff --git a/litellm/llms/azure/common_utils.py b/litellm/llms/azure/common_utils.py index 52e3cdf380e..23c5aa35f4b 100644 --- a/litellm/llms/azure/common_utils.py +++ b/litellm/llms/azure/common_utils.py @@ -18,6 +18,7 @@ import litellm from litellm._logging import verbose_logger from litellm.caching.caching import DualCache from litellm.constants import DEFAULT_MAX_RETRIES +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.llms.openai.common_utils import BaseOpenAILLM, OpenAIAsyncHTTPClient, OpenAIHTTPClient @@ -616,7 +617,7 @@ class BaseAzureLLM(BaseOpenAILLM): v1_params: Final[_AzureV1ClientParams] = { "api_key": v1_api_key, - "base_url": f"{api_base}/openai/v1/", + "base_url": f"{ensure_api_route(api_base or '', 'openai/v1')}/", } if "timeout" in azure_client_params: v1_params["timeout"] = azure_client_params["timeout"] diff --git a/litellm/llms/azure/realtime/handler.py b/litellm/llms/azure/realtime/handler.py index 75f96bfdf63..3d4576fd7ac 100644 --- a/litellm/llms/azure/realtime/handler.py +++ b/litellm/llms/azure/realtime/handler.py @@ -10,6 +10,7 @@ from typing import Final, Protocol, cast from litellm._logging import verbose_proxy_logger from litellm.constants import REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.types.realtime import RealtimeQueryParams from ....litellm_core_utils.litellm_logging import Logging as LiteLLMLogging @@ -136,7 +137,8 @@ class AzureOpenAIRealtime(AzureChatCompletion): query_parts.append(urlencode({"intent": intent})) qs: Final = "&".join(query_parts) - return f"{api_base}{path}?{qs}" if qs else f"{api_base}{path}" + endpoint_url: Final = join_api_base(api_base, path) + return f"{endpoint_url}?{qs}" if qs else endpoint_url def construct_url( self, diff --git a/litellm/llms/azure/realtime/http_transformation.py b/litellm/llms/azure/realtime/http_transformation.py index 542a6a8c02f..545c3348182 100644 --- a/litellm/llms/azure/realtime/http_transformation.py +++ b/litellm/llms/azure/realtime/http_transformation.py @@ -3,6 +3,7 @@ from typing import Final import litellm +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.realtime.http_transformation import BaseRealtimeHTTPConfig from litellm.secret_managers.main import get_secret_str @@ -39,9 +40,9 @@ class AzureRealtimeHTTPConfig(BaseRealtimeHTTPConfig): return _azure_realtime_v1_url(self.get_api_base(api_base), "calls") def get_transcription_session_url(self, api_base: str | None, model: str, api_version: str | None = None) -> str: - base: Final = self.get_api_base(api_base).rstrip("/") + base: Final = self.get_api_base(api_base) version: Final = api_version or get_secret_str("AZURE_API_VERSION") or "2024-12-17" - return f"{base}/openai/realtime/transcription_sessions?api-version={version}" + return f"{join_api_base(base, 'openai/realtime/transcription_sessions')}?api-version={version}" def get_realtime_calls_headers(self, ephemeral_key: str) -> dict[str, str]: return { diff --git a/litellm/llms/azure/search/transformation.py b/litellm/llms/azure/search/transformation.py index f2bb2ecc68c..f7cdf88ade0 100644 --- a/litellm/llms/azure/search/transformation.py +++ b/litellm/llms/azure/search/transformation.py @@ -35,6 +35,7 @@ from typing import TYPE_CHECKING, Final, Literal import httpx from pydantic import ConfigDict, ValidationError +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.azure.common_utils import cached_azure_ad_token_refresh_provider from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.search.transformation import ( @@ -340,10 +341,7 @@ class BingGroundingSearchConfig(BaseSearchConfig): f"{PROJECT_ENDPOINT_ENV} is not set. Set it to your Microsoft Foundry project " f"endpoint, e.g. https://.services.ai.azure.com/api/projects/." ) - trimmed: Final = resolved_base.rstrip("/") - if trimmed.endswith(_RESPONSES_PATH): - return trimmed - return f"{trimmed}{_RESPONSES_PATH}" + return ensure_api_route(resolved_base, _RESPONSES_PATH) def transform_search_request( self, diff --git a/litellm/llms/azure/text_to_speech/transformation.py b/litellm/llms/azure/text_to_speech/transformation.py index eed7a3178ca..acae650ed75 100644 --- a/litellm/llms/azure/text_to_speech/transformation.py +++ b/litellm/llms/azure/text_to_speech/transformation.py @@ -11,6 +11,7 @@ from urllib.parse import urlparse import httpx import litellm +from litellm.litellm_core_utils.url_utils import ensure_api_route, join_api_base from litellm.llms.base_llm.text_to_speech.transformation import ( BaseTextToSpeechConfig, TextToSpeechRequestData, @@ -329,8 +330,6 @@ class AzureAVATextToSpeechConfig(BaseTextToSpeechConfig): f"https://{{region}}.{self.TTS_SPEECH_DOMAIN}" ) - # Remove trailing slash and parse URL - api_base = api_base.rstrip("/") parsed_url: Final = urlparse(api_base) hostname: Final = parsed_url.hostname or "" @@ -339,14 +338,10 @@ class AzureAVATextToSpeechConfig(BaseTextToSpeechConfig): region: Final = self._extract_region_from_hostname(hostname=hostname, domain=self.COGNITIVE_SERVICES_DOMAIN) return self._build_tts_url(region=region) - # Check if it's already a TTS endpoint if self._is_tts_endpoint(hostname=hostname): - if not api_base.endswith(self.TTS_ENDPOINT_PATH): - return f"{api_base}{self.TTS_ENDPOINT_PATH}" - return api_base + return ensure_api_route(api_base, self.TTS_ENDPOINT_PATH) - # Assume it's a custom endpoint, append the path - return f"{api_base}{self.TTS_ENDPOINT_PATH}" + return join_api_base(api_base, self.TTS_ENDPOINT_PATH) def _is_cognitive_services_endpoint(self, hostname: str) -> bool: """Check if hostname is a Cognitive Services endpoint""" diff --git a/litellm/llms/azure_ai/agents/handler.py b/litellm/llms/azure_ai/agents/handler.py index 0b0548dd3a2..94f1399f01b 100644 --- a/litellm/llms/azure_ai/agents/handler.py +++ b/litellm/llms/azure_ai/agents/handler.py @@ -30,7 +30,7 @@ from typing_extensions import ReadOnly from litellm._logging import verbose_logger from litellm.litellm_core_utils.hidden_params import set_hidden_param -from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.llms.azure_ai.agents.transformation import ( AzureAIAgentsConfig, AzureAIAgentsError, @@ -130,28 +130,29 @@ class AzureAIAgentsHandler: # See: https://learn.microsoft.com/en-us/azure/ai-foundry/agents/quickstart # ------------------------------------------------------------------------- def _build_thread_url(self, api_base: str, api_version: str) -> str: - return f"{api_base}/threads?api-version={api_version}" + return f"{join_api_base(api_base, 'threads')}?api-version={api_version}" def _build_messages_url(self, api_base: str, thread_id: str, api_version: str) -> str: encoded_thread_id: Final = encode_url_path_segment(thread_id, field_name="thread_id") - return f"{api_base}/threads/{encoded_thread_id}/messages?api-version={api_version}" + return f"{join_api_base(api_base, f'threads/{encoded_thread_id}/messages')}?api-version={api_version}" def _build_runs_url(self, api_base: str, thread_id: str, api_version: str) -> str: encoded_thread_id: Final = encode_url_path_segment(thread_id, field_name="thread_id") - return f"{api_base}/threads/{encoded_thread_id}/runs?api-version={api_version}" + return f"{join_api_base(api_base, f'threads/{encoded_thread_id}/runs')}?api-version={api_version}" def _build_run_status_url(self, api_base: str, thread_id: str, run_id: str, api_version: str) -> str: encoded_thread_id: Final = encode_url_path_segment(thread_id, field_name="thread_id") encoded_run_id: Final = encode_url_path_segment(run_id, field_name="run_id") - return f"{api_base}/threads/{encoded_thread_id}/runs/{encoded_run_id}?api-version={api_version}" + run_status_route: Final = f"threads/{encoded_thread_id}/runs/{encoded_run_id}" + return f"{join_api_base(api_base, run_status_route)}?api-version={api_version}" def _build_list_messages_url(self, api_base: str, thread_id: str, api_version: str) -> str: encoded_thread_id: Final = encode_url_path_segment(thread_id, field_name="thread_id") - return f"{api_base}/threads/{encoded_thread_id}/messages?api-version={api_version}" + return f"{join_api_base(api_base, f'threads/{encoded_thread_id}/messages')}?api-version={api_version}" def _build_create_thread_and_run_url(self, api_base: str, api_version: str) -> str: """URL for the create-thread-and-run endpoint (supports streaming).""" - return f"{api_base}/threads/runs?api-version={api_version}" + return f"{join_api_base(api_base, 'threads/runs')}?api-version={api_version}" # ------------------------------------------------------------------------- # Response Helpers @@ -281,8 +282,6 @@ class AzureAIAgentsHandler: api_version: Final = optional_params.get("api_version", self.config.DEFAULT_API_VERSION) agent_id: Final = self.config.get_agent_id(model, optional_params) thread_id: Final = optional_params.get("thread_id") - api_base = api_base.rstrip("/") - verbose_logger.debug("Azure AI Agents completion - api_base: %s, agent_id: %s", api_base, agent_id) return headers, api_version, agent_id, thread_id, api_base diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py b/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py index 5946a1f4480..6e3cfce5eb4 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py @@ -5,8 +5,10 @@ Extends the base Anthropic CountTokens transformation with Azure authentication. """ from typing import Final +from urllib.parse import urlsplit from litellm.constants import ANTHROPIC_TOKEN_COUNTING_BETA_VERSION +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.anthropic.count_tokens.transformation import ( AnthropicCountTokensConfig, ) @@ -76,14 +78,8 @@ class AzureAIAnthropicCountTokensConfig(AnthropicCountTokensConfig): Returns: The endpoint URL for the CountTokens API """ - # Azure AI Anthropic endpoint format: - # https://.services.ai.azure.com/anthropic/v1/messages/count_tokens - api_base = api_base.rstrip("/") - - # Ensure the URL has /anthropic path - if not api_base.endswith("/anthropic"): - if "/anthropic" not in api_base: - api_base = f"{api_base}/anthropic" - - # Add the count_tokens path - return f"{api_base}/v1/messages/count_tokens" + base: Final = api_base.rstrip("/") + if base.endswith("/v1/messages"): + return join_api_base(base, "count_tokens") + anthropic_base: Final = base if "/anthropic" in urlsplit(base).path else join_api_base(base, "anthropic") + return join_api_base(anthropic_base, "v1/messages/count_tokens") diff --git a/litellm/llms/azure_ai/anthropic/messages_transformation.py b/litellm/llms/azure_ai/anthropic/messages_transformation.py index 91c75661688..2aacc5771d0 100644 --- a/litellm/llms/azure_ai/anthropic/messages_transformation.py +++ b/litellm/llms/azure_ai/anthropic/messages_transformation.py @@ -9,6 +9,7 @@ from litellm.llms.anthropic.pass_through.messages.transformation import ( AnthropicMessagesConfig, ) from litellm.llms.azure.common_utils import BaseAzureLLM +from litellm.llms.azure_ai.anthropic.transformation import azure_anthropic_messages_url from litellm.types.router import GenericLiteLLMParams @@ -96,25 +97,7 @@ class AzureAnthropicMessagesConfig(AnthropicMessagesConfig): "Expected format: https://.services.ai.azure.com/anthropic" ) - # Ensure the URL ends with /v1/messages - api_base = api_base.rstrip("/") - if api_base.endswith("/v1/messages") or api_base.endswith("/anthropic/v1/messages"): - # Already correct - pass - else: - # Check if /anthropic is already in the path - if "/anthropic" in api_base: - # /anthropic exists, ensure we end with /anthropic/v1/messages - # Extract the base URL up to and including /anthropic - parts: Final = api_base.split("/anthropic", 1) - api_base = parts[0] + "/anthropic" - else: - # /anthropic not in path, add it - api_base = api_base + "/anthropic" - # Add /v1/messages - api_base = api_base + "/v1/messages" - - return api_base + return azure_anthropic_messages_url(api_base) def _remove_scope_from_cache_control(self, anthropic_messages_request: dict) -> None: """ diff --git a/litellm/llms/azure_ai/anthropic/transformation.py b/litellm/llms/azure_ai/anthropic/transformation.py index 18a4eeb7c0a..c75213c3e2c 100644 --- a/litellm/llms/azure_ai/anthropic/transformation.py +++ b/litellm/llms/azure_ai/anthropic/transformation.py @@ -4,12 +4,23 @@ Azure Anthropic transformation config - extends AnthropicConfig with Azure authe from typing import Final +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.anthropic.chat.transformation import AnthropicConfig from litellm.llms.azure.common_utils import BaseAzureLLM from litellm.types.llms.openai import AllMessageValues from litellm.types.router import GenericLiteLLMParams +def azure_anthropic_messages_url(api_base: str) -> str: + base: Final = api_base.rstrip("/") + if base.endswith("/v1/messages"): + return base + path_start: Final = base.find("/", base.find("://") + len("://")) if "://" in base else 0 + anthropic_start: Final = base.find("/anthropic", path_start) if path_start >= 0 else -1 + anthropic_root: Final = base[:anthropic_start] if anthropic_start >= 0 else base + return join_api_base(anthropic_root, "anthropic/v1/messages") + + def _promote_extra_body_to_optional_params(optional_params: dict) -> None: """Promote anthropic-native passthrough keys out of ``extra_body``. diff --git a/litellm/llms/azure_ai/embed/handler.py b/litellm/llms/azure_ai/embed/handler.py index 57e6ca51ade..dc3bc6bdc5e 100644 --- a/litellm/llms/azure_ai/embed/handler.py +++ b/litellm/llms/azure_ai/embed/handler.py @@ -5,6 +5,7 @@ from urllib.parse import urlsplit, urlunsplit from openai import OpenAI import litellm +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.azure_ai.common_utils import is_foundry_model_inference_base from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, @@ -74,13 +75,18 @@ class AzureAIEmbedding(OpenAIChatCompletion): api_base: str | None, client: HTTPHandler | AsyncHTTPHandler | None = None, ) -> EmbeddingResponse: + if api_base is None: + raise ValueError( + "api_base is None. Please set AZURE_AI_API_BASE or dynamically via `api_base` param, " + "to make the request." + ) if client is None or not isinstance(client, AsyncHTTPHandler): client = get_async_httpx_client( llm_provider=litellm.LlmProviders.AZURE_AI, params={"timeout": timeout}, ) - url: Final = f"{api_base}/images/embeddings" + url: Final = join_api_base(api_base, "images/embeddings") response: Final = await client.post( url=url, @@ -123,7 +129,7 @@ class AzureAIEmbedding(OpenAIChatCompletion): if client is None or not isinstance(client, HTTPHandler): client = HTTPHandler(timeout=timeout, concurrent_limit=1) - url: Final = f"{api_base}/images/embeddings" + url: Final = join_api_base(api_base, "images/embeddings") response: Final = client.post( url=url, diff --git a/litellm/llms/azure_ai/image_generation/flux_transformation.py b/litellm/llms/azure_ai/image_generation/flux_transformation.py index 8b4fa78cdf4..6a7bebd9f6b 100644 --- a/litellm/llms/azure_ai/image_generation/flux_transformation.py +++ b/litellm/llms/azure_ai/image_generation/flux_transformation.py @@ -3,6 +3,7 @@ from types import MappingProxyType from typing import Final from litellm.exceptions import BadRequestError, UnsupportedParamsError +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.openai.image_generation import GPTImageGenerationConfig from litellm.types.llms.openai import OpenAIImageGenerationOptionalParams @@ -52,7 +53,8 @@ class AzureFoundryFluxImageGenerationConfig(GPTImageGenerationConfig): return f"{api_base}?api-version={api_version}" provider_model_path: Final = AzureFoundryFluxImageGenerationConfig.get_flux2_provider_model_path(model) - return f"{api_base}/providers/blackforestlabs/v1/{provider_model_path}?api-version={api_version}" + flux_route: Final = f"providers/blackforestlabs/v1/{provider_model_path}" + return f"{join_api_base(api_base, flux_route)}?api-version={api_version}" @staticmethod def is_flux2_model(model: str) -> bool: diff --git a/litellm/llms/azure_ai/image_generation/mai_transformation.py b/litellm/llms/azure_ai/image_generation/mai_transformation.py index 18b4b6f456a..8bb5d3e6b8d 100644 --- a/litellm/llms/azure_ai/image_generation/mai_transformation.py +++ b/litellm/llms/azure_ai/image_generation/mai_transformation.py @@ -3,6 +3,7 @@ from typing import TYPE_CHECKING, Any, Final import httpx from litellm.exceptions import UnsupportedParamsError +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.image_generation.transformation import ( BaseImageGenerationConfig, ) @@ -41,9 +42,9 @@ class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig): if "/mai/" in path: prefix, _, _ = path.partition("/images/") - path = f"{prefix}/images/generations" + path = join_api_base(prefix, "images/generations") else: - path = f"{path}/mai/v1/images/generations" + path = join_api_base(path, "mai/v1/images/generations") if separator: return f"{path}?{query}" @@ -63,9 +64,9 @@ class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig): if "/mai/" in path: prefix, _, _ = path.partition("/images/") - path = f"{prefix}/images/edits" + path = join_api_base(prefix, "images/edits") else: - path = f"{path}/mai/v1/images/edits" + path = join_api_base(path, "mai/v1/images/edits") if separator: return f"{path}?{query}" diff --git a/litellm/llms/azure_ai/vector_stores/transformation.py b/litellm/llms/azure_ai/vector_stores/transformation.py index 7638229bc32..a4777be146f 100644 --- a/litellm/llms/azure_ai/vector_stores/transformation.py +++ b/litellm/llms/azure_ai/vector_stores/transformation.py @@ -5,6 +5,7 @@ from typing import TYPE_CHECKING, Any, Final import httpx +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.azure.common_utils import BaseAzureLLM from litellm.llms.base_llm.vector_store.transformation import ( BaseQueryEmbeddingVectorStoreConfig, @@ -172,7 +173,8 @@ class AzureAIVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig, BaseAzureLLM litellm_logging_obj.model_call_details["input"] = query_text litellm_logging_obj.model_call_details["embedding_model"] = litellm_params.get("litellm_embedding_model") litellm_logging_obj.model_call_details["top_k"] = top_k - return f"{api_base}/indexes/{vector_store_id}/docs/search?api-version=2024-07-01", { + search_route: Final = f"indexes/{vector_store_id}/docs/search" + return f"{join_api_base(api_base, search_route)}?api-version=2024-07-01", { "search": "*", "vectorQueries": [{"vector": query_vector, "fields": vector_field, "kind": "vector", "k": top_k}], "select": "id,content", diff --git a/litellm/llms/base_llm/decisions/transformation.py b/litellm/llms/base_llm/decisions/transformation.py index d74257a4a9d..a122aeb3033 100644 --- a/litellm/llms/base_llm/decisions/transformation.py +++ b/litellm/llms/base_llm/decisions/transformation.py @@ -12,6 +12,7 @@ from pydantic_core import ErrorDetails from typing_extensions import assert_never import litellm +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.secret_managers.main import get_secret_str from litellm.types.decisions import ( @@ -582,7 +583,7 @@ class BaseDecisionsConfig(ABC): ) def get_complete_url(self, api_base: str, model: str) -> str: - return f"{api_base.rstrip('/').removesuffix('/v1')}{self.path}" + return join_api_base(api_base.rstrip("/").removesuffix("/v1"), self.path) def signs_with_aws(self, api_base: str) -> bool: return False diff --git a/litellm/llms/base_llm/evals/transformation.py b/litellm/llms/base_llm/evals/transformation.py index 49124e271a6..40093953144 100644 --- a/litellm/llms/base_llm/evals/transformation.py +++ b/litellm/llms/base_llm/evals/transformation.py @@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Any import httpx +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.types.llms.openai_evals import ( CancelEvalResponse, @@ -79,7 +80,7 @@ class BaseEvalsAPIConfig(ABC): """ if api_base is None: raise ValueError("api_base is required") - return f"{api_base}/v1/{endpoint}" + return join_api_base(api_base, f"v1/{endpoint}") @abstractmethod def transform_create_eval_request( diff --git a/litellm/llms/base_llm/realtime/http_transformation.py b/litellm/llms/base_llm/realtime/http_transformation.py index 43a80edb493..a02054ec8cb 100644 --- a/litellm/llms/base_llm/realtime/http_transformation.py +++ b/litellm/llms/base_llm/realtime/http_transformation.py @@ -7,10 +7,11 @@ These are HTTP (not WebSocket) endpoints used by the WebRTC flow: """ from abc import ABC, abstractmethod -from typing import Final import httpx +from litellm.litellm_core_utils.url_utils import join_api_base + class BaseRealtimeHTTPConfig(ABC): """ @@ -61,7 +62,7 @@ class BaseRealtimeHTTPConfig(ABC): """Return the full URL for POST /realtime/transcription_sessions.""" base = (api_base or "").rstrip("/") base = base.removesuffix("/v1") - return f"{base}/v1/realtime/transcription_sessions" + return join_api_base(base, "v1/realtime/transcription_sessions") @abstractmethod def validate_environment( @@ -83,8 +84,7 @@ class BaseRealtimeHTTPConfig(ABC): def get_realtime_calls_url(self, api_base: str | None, model: str, api_version: str | None = None) -> str: """Return the full URL for POST /realtime/calls (SDP exchange).""" - base: Final = (api_base or "").rstrip("/") - return f"{base}/v1/realtime/calls" + return join_api_base(api_base or "", "v1/realtime/calls") def get_realtime_calls_headers(self, ephemeral_key: str) -> dict: """ diff --git a/litellm/llms/base_llm/skills/transformation.py b/litellm/llms/base_llm/skills/transformation.py index ba202e3db5d..64a8e055832 100644 --- a/litellm/llms/base_llm/skills/transformation.py +++ b/litellm/llms/base_llm/skills/transformation.py @@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Any import httpx +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.types.llms.anthropic_skills import ( CreateSkillRequest, @@ -71,7 +72,7 @@ class BaseSkillsAPIConfig(ABC): """ if api_base is None: raise ValueError("api_base is required") - return f"{api_base}/v1/{endpoint}" + return join_api_base(api_base, f"v1/{endpoint}") @abstractmethod def transform_create_skill_request( diff --git a/litellm/llms/bedrock/chat/agentcore/transformation.py b/litellm/llms/bedrock/chat/agentcore/transformation.py index 48b3795e274..bb5bb246a93 100644 --- a/litellm/llms/bedrock/chat/agentcore/transformation.py +++ b/litellm/llms/bedrock/chat/agentcore/transformation.py @@ -18,6 +18,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( convert_content_list_to_str, ) from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.a2a.common_utils import extract_text_from_a2a_response from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM @@ -98,14 +99,14 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): # Build the base endpoint URL for AgentCore # Note: We don't use get_runtime_endpoint as AgentCore has its own endpoint structure if aws_bedrock_runtime_endpoint: - base_url = aws_bedrock_runtime_endpoint + base_url = f"{aws_bedrock_runtime_endpoint}" else: base_url = f"https://bedrock-agentcore.{region}.{get_aws_dns_suffix(region)}" # Based on boto3 client.invoke_agent_runtime, the path is: # /runtimes/{URL-ENCODED-ARN}/invocations?qualifier= encoded_arn: Final = quote(agent_runtime_arn, safe="") - endpoint_url = f"{base_url}/runtimes/{encoded_arn}/invocations" + endpoint_url = join_api_base(base_url, f"runtimes/{encoded_arn}/invocations") # Add qualifier as query parameter if provided if "qualifier" in optional_params: diff --git a/litellm/llms/bedrock/chat/chat_completions/transformation.py b/litellm/llms/bedrock/chat/chat_completions/transformation.py index 907f8e50ac7..123e4a9e6b4 100644 --- a/litellm/llms/bedrock/chat/chat_completions/transformation.py +++ b/litellm/llms/bedrock/chat/chat_completions/transformation.py @@ -33,6 +33,7 @@ from litellm.litellm_core_utils.prompt_templates.image_handling import ( inline_remote_image_urls, inline_remote_media, ) +from litellm.litellm_core_utils.url_utils import ensure_api_route, join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, bedrock_bearer_token from litellm.llms.bedrock.common_utils import ( @@ -371,12 +372,9 @@ class AmazonBedrockRuntimeChatCompletionsConfig(OpenAILikeChatConfig): ), aws_region_name=aws_region_name, ) - base: Final = proxy_endpoint_url.rstrip("/") - if base.endswith("/openai/v1/chat/completions"): - return base - if base.endswith("/openai/v1"): - return f"{base}/chat/completions" - return f"{base}/openai/v1/chat/completions" + if proxy_endpoint_url.rstrip("/").endswith("/openai/v1"): + return join_api_base(proxy_endpoint_url, "chat/completions") + return ensure_api_route(proxy_endpoint_url, "openai/v1/chat/completions") def _params_with_region_from_path( self, optional_params: dict, model: str | None diff --git a/litellm/llms/bedrock/chat/converse_handler.py b/litellm/llms/bedrock/chat/converse_handler.py index 7d02b9ef230..a6f9192339f 100644 --- a/litellm/llms/bedrock/chat/converse_handler.py +++ b/litellm/llms/bedrock/chat/converse_handler.py @@ -10,6 +10,7 @@ from litellm.anthropic_beta_headers_manager import ( ) from litellm.litellm_core_utils.get_litellm_params import stored_control_options from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObject +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm import base_model_iterator from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, @@ -363,11 +364,11 @@ class BedrockConverseLLM(BaseAWSLLM): aws_region_name=aws_region_name, ) if (stream is not None and stream is True) and not fake_stream: - endpoint_url = f"{endpoint_url}/model/{modelId}/converse-stream" - proxy_endpoint_url = f"{proxy_endpoint_url}/model/{modelId}/converse-stream" + endpoint_url = join_api_base(endpoint_url, f"model/{modelId}/converse-stream") + proxy_endpoint_url = join_api_base(proxy_endpoint_url, f"model/{modelId}/converse-stream") else: - endpoint_url = f"{endpoint_url}/model/{modelId}/converse" - proxy_endpoint_url = f"{proxy_endpoint_url}/model/{modelId}/converse" + endpoint_url = join_api_base(endpoint_url, f"model/{modelId}/converse") + proxy_endpoint_url = join_api_base(proxy_endpoint_url, f"model/{modelId}/converse") ## COMPLETION CALL headers = {"Content-Type": "application/json"} diff --git a/litellm/llms/bedrock/chat/invoke_agent/transformation.py b/litellm/llms/bedrock/chat/invoke_agent/transformation.py index 6877af74494..baef89912fa 100644 --- a/litellm/llms/bedrock/chat/invoke_agent/transformation.py +++ b/litellm/llms/bedrock/chat/invoke_agent/transformation.py @@ -15,7 +15,7 @@ from litellm._uuid import uuid from litellm.litellm_core_utils.prompt_templates.common_utils import ( convert_content_list_to_str, ) -from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.llms.bedrock.common_utils import BedrockError @@ -101,9 +101,10 @@ class AmazonInvokeAgentConfig(BaseConfig, BaseAWSLLM): encoded_agent_alias_id: Final = encode_url_path_segment(agent_alias_id, field_name="agent_alias_id") encoded_session_id: Final = encode_url_path_segment(session_id, field_name="session_id") - endpoint_url = f"{endpoint_url}/agents/{encoded_agent_id}/agentAliases/{encoded_agent_alias_id}/sessions/{encoded_session_id}/text" - - return endpoint_url + return join_api_base( + endpoint_url, + f"agents/{encoded_agent_id}/agentAliases/{encoded_agent_alias_id}/sessions/{encoded_session_id}/text", + ) def sign_request( self, diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_openai_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_openai_transformation.py index 4bf1a1cba73..f1e1f3e8738 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_openai_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_openai_transformation.py @@ -11,6 +11,7 @@ from typing import TYPE_CHECKING, Any, Final import httpx +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.bedrock.request_metadata import ( @@ -97,13 +98,8 @@ class AmazonBedrockOpenAIConfig(OpenAIGPTConfig, BaseAWSLLM): # Encode model ID for ARNs (e.g., :imported-model/ -> :imported-model%2F) model_id = CommonUtils.encode_bedrock_runtime_modelid_arn(model_id) - # Build the invoke URL - if stream: - endpoint_url = f"{endpoint_url}/model/{model_id}/invoke-with-response-stream" - else: - endpoint_url = f"{endpoint_url}/model/{model_id}/invoke" - - return endpoint_url + invoke_route: Final = "invoke-with-response-stream" if stream else "invoke" + return join_api_base(endpoint_url, f"model/{model_id}/{invoke_route}") def sign_request( self, 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 87d4fc0b32e..51b5a99272e 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py @@ -18,6 +18,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import ( prompt_factory, render_custom_prompt_template, ) +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException from litellm.llms.bedrock.chat.invoke_handler import make_call, make_sync_call from litellm.llms.bedrock.common_utils import BedrockError @@ -136,20 +137,13 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM): aws_bedrock_runtime_endpoint: Final = optional_params.get( "aws_bedrock_runtime_endpoint", None ) # https://bedrock-runtime.{region_name}.amazonaws.com - endpoint_url, proxy_endpoint_url = self.get_runtime_endpoint( + endpoint_url, _ = self.get_runtime_endpoint( api_base=api_base, aws_bedrock_runtime_endpoint=aws_bedrock_runtime_endpoint, aws_region_name=self._get_aws_region_name(optional_params=optional_params, model=model), ) - - if (stream is not None and stream is True) and provider != "ai21": - endpoint_url = f"{endpoint_url}/model/{modelId}/invoke-with-response-stream" - proxy_endpoint_url = f"{proxy_endpoint_url}/model/{modelId}/invoke-with-response-stream" - else: - endpoint_url = f"{endpoint_url}/model/{modelId}/invoke" - proxy_endpoint_url = f"{proxy_endpoint_url}/model/{modelId}/invoke" - - return endpoint_url + invoke_route: Final = "invoke-with-response-stream" if stream is True and provider != "ai21" else "invoke" + return join_api_base(endpoint_url, f"model/{modelId}/{invoke_route}") def sign_request( self, diff --git a/litellm/llms/bedrock/claude_platform/common_utils.py b/litellm/llms/bedrock/claude_platform/common_utils.py index e25e58056f5..77cae177792 100644 --- a/litellm/llms/bedrock/claude_platform/common_utils.py +++ b/litellm/llms/bedrock/claude_platform/common_utils.py @@ -5,6 +5,7 @@ import httpx import litellm from litellm._logging import verbose_logger +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.llms.bedrock.common_utils import BedrockError from litellm.secret_managers.main import get_secret_str @@ -140,9 +141,7 @@ class BedrockClaudePlatformMixin(BaseAWSLLM): if api_base is None: aws_region_name: Final = self._get_required_aws_region_name(optional_params) api_base = f"https://{CLAUDE_PLATFORM_SERVICE_NAME}.{aws_region_name}.api.aws" - if not api_base.endswith("/v1/messages"): - api_base = f"{api_base.rstrip('/')}/v1/messages" - return api_base + return ensure_api_route(api_base, "v1/messages") def sign_request( self, diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 0cf3e895ad2..4954e646f58 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -31,6 +31,7 @@ from litellm import verbose_logger from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix from litellm.litellm_core_utils.optional_imports import ensure_optional_import from litellm.litellm_core_utils.prompt_templates.common_utils import schema_property_count +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm.anthropic_messages.transformation import ( BaseAnthropicMessagesConfig, ) @@ -1305,10 +1306,7 @@ def build_mantle_messages_url( """ override: Final = api_base or aws_bedrock_runtime_endpoint or _mantle_api_base_from_env() if override: - base: Final = override.rstrip("/") - if base.endswith(MANTLE_MESSAGES_PATH): - return base - return f"{base}{MANTLE_MESSAGES_PATH}" + return ensure_api_route(override, MANTLE_MESSAGES_PATH) return f"https://bedrock-mantle.{region}.api.aws{MANTLE_MESSAGES_PATH}" diff --git a/litellm/llms/bedrock/count_tokens/mantle_handler.py b/litellm/llms/bedrock/count_tokens/mantle_handler.py index 14edd39c386..34e622b5e8d 100644 --- a/litellm/llms/bedrock/count_tokens/mantle_handler.py +++ b/litellm/llms/bedrock/count_tokens/mantle_handler.py @@ -7,6 +7,7 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.anthropic.count_tokens.transformation import COUNT_TOKEN_OPTION_NAMES, AnthropicCountTokensConfig from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, run_aws_signing from litellm.llms.bedrock.common_utils import BedrockError, build_mantle_messages_url @@ -43,7 +44,7 @@ class BedrockMantleCountTokensHandler(BaseAWSLLM): messages_url: Final = build_mantle_messages_url( api_base=None, aws_bedrock_runtime_endpoint=None, region=aws_region_name ) - return f"{messages_url}{MANTLE_COUNT_TOKENS_SUFFIX}" + return join_api_base(messages_url, MANTLE_COUNT_TOKENS_SUFFIX) async def handle_count_tokens_request( self, diff --git a/litellm/llms/bedrock/count_tokens/transformation.py b/litellm/llms/bedrock/count_tokens/transformation.py index 7cca4f69acc..a9f95c1f9b9 100644 --- a/litellm/llms/bedrock/count_tokens/transformation.py +++ b/litellm/llms/bedrock/count_tokens/transformation.py @@ -11,6 +11,7 @@ from typing import Final, Literal from pydantic import JsonValue +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.llms.bedrock.common_utils import get_bedrock_base_model @@ -236,9 +237,7 @@ class BedrockCountTokensConfig(BaseAWSLLM): aws_bedrock_runtime_endpoint=aws_bedrock_runtime_endpoint, aws_region_name=aws_region_name, ) - endpoint: Final = f"{base_url}/model/{encoded_model_id}/count-tokens" - - return endpoint + return join_api_base(base_url, f"model/{encoded_model_id}/count-tokens") def transform_bedrock_response_to_anthropic( self, bedrock_response: Mapping[str, JsonValue] diff --git a/litellm/llms/bedrock/embed/embedding.py b/litellm/llms/bedrock/embed/embedding.py index 58cfb6cb95e..4d27d94664f 100644 --- a/litellm/llms/bedrock/embed/embedding.py +++ b/litellm/llms/bedrock/embed/embedding.py @@ -12,6 +12,7 @@ import httpx import litellm from litellm.constants import BEDROCK_EMBEDDING_PROVIDERS_LITERAL +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.cohere.embed.handler import embedding as cohere_embedding from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, @@ -512,9 +513,9 @@ class BedrockEmbedding(BaseAWSLLM): aws_region_name=aws_region_name, ) if has_async_invoke: - endpoint_url = f"{endpoint_url}/async-invoke" + endpoint_url = join_api_base(endpoint_url, "async-invoke") else: - endpoint_url = f"{endpoint_url}/model/{modelId}/invoke" + endpoint_url = join_api_base(endpoint_url, f"model/{modelId}/invoke") if batch_data is not None: if aembedding: @@ -611,7 +612,7 @@ class BedrockEmbedding(BaseAWSLLM): # Encode the ARN for use in URL path encoded_arn: Final = quote(invocation_arn, safe="") - status_url: Final = f"{endpoint_url.rstrip('/')}/async-invoke/{encoded_arn}" + status_url: Final = join_api_base(endpoint_url, f"async-invoke/{encoded_arn}") # Prepare headers for GET request headers: Final = {"Content-Type": "application/json"} diff --git a/litellm/llms/bedrock/image_edit/handler.py b/litellm/llms/bedrock/image_edit/handler.py index 53bb332f511..466993de9d1 100644 --- a/litellm/llms/bedrock/image_edit/handler.py +++ b/litellm/llms/bedrock/image_edit/handler.py @@ -14,6 +14,7 @@ import httpx import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.bedrock.image_edit.amazon_nova_canvas_image_edit_transformation import ( BedrockAmazonNovaCanvasImageEditConfig, ) @@ -226,7 +227,7 @@ class BedrockImageEdit(BaseAWSLLM): aws_bedrock_runtime_endpoint=boto3_credentials_info.aws_bedrock_runtime_endpoint, aws_region_name=boto3_credentials_info.aws_region_name, ) - proxy_endpoint_url = f"{proxy_endpoint_url}/model/{modelId}/invoke" + proxy_endpoint_url = join_api_base(proxy_endpoint_url, f"model/{modelId}/invoke") data: Final = self._get_request_body( model=model, image=image, diff --git a/litellm/llms/bedrock/image_generation/image_handler.py b/litellm/llms/bedrock/image_generation/image_handler.py index e81dfe76942..abd6d40f9e9 100644 --- a/litellm/llms/bedrock/image_generation/image_handler.py +++ b/litellm/llms/bedrock/image_generation/image_handler.py @@ -8,6 +8,7 @@ import httpx import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.bedrock.image_generation.amazon_nova_canvas_transformation import ( AmazonNovaCanvasConfig, ) @@ -248,7 +249,7 @@ class BedrockImageGeneration(BaseAWSLLM): aws_bedrock_runtime_endpoint=boto3_credentials_info.aws_bedrock_runtime_endpoint, aws_region_name=boto3_credentials_info.aws_region_name, ) - proxy_endpoint_url = f"{proxy_endpoint_url}/model/{modelId}/invoke" + proxy_endpoint_url = join_api_base(proxy_endpoint_url, f"model/{modelId}/invoke") data: Final = self._get_request_body( model=model, prompt=prompt, diff --git a/litellm/llms/bedrock/rerank/handler.py b/litellm/llms/bedrock/rerank/handler.py index 01ef5af2b67..0634275135b 100644 --- a/litellm/llms/bedrock/rerank/handler.py +++ b/litellm/llms/bedrock/rerank/handler.py @@ -6,6 +6,7 @@ from pydantic import ConfigDict, TypeAdapter import litellm from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, HTTPHandler, @@ -160,7 +161,7 @@ class BedrockRerankHandler(BaseAWSLLM): aws_region_name=boto3_credentials_info.aws_region_name, ) proxy_endpoint_url = proxy_endpoint_url.replace("bedrock-runtime", "bedrock-agent-runtime") - proxy_endpoint_url = f"{proxy_endpoint_url}/rerank" + proxy_endpoint_url = join_api_base(proxy_endpoint_url, "rerank") body: Final = json.dumps(data).encode("utf-8") headers = {"Content-Type": "application/json"} diff --git a/litellm/llms/bedrock/responses/transformation.py b/litellm/llms/bedrock/responses/transformation.py index 086d1211835..7cdd7b08f99 100644 --- a/litellm/llms/bedrock/responses/transformation.py +++ b/litellm/llms/bedrock/responses/transformation.py @@ -46,6 +46,7 @@ from litellm.litellm_core_utils.prompt_templates.image_handling import ( async_convert_url_to_base64, convert_url_to_base64, ) +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.responses.codex_compat import drop_unsupported_tools, normalize_codex_input_items from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM @@ -223,9 +224,10 @@ class BedrockOpenAIResponsesConfig(BaseAWSLLM, OpenAIResponsesAPIConfig): litellm_params: dict, # mutable-ok: signature fixed by the BaseResponsesAPIConfig override contract ) -> str: region: Final = self._get_aws_region_name(optional_params=litellm_params, model=None) + configured_endpoint: Final = litellm_params.get("aws_bedrock_runtime_endpoint") override: Final = ( api_base - or litellm_params.get("aws_bedrock_runtime_endpoint") + or (f"{configured_endpoint}" if configured_endpoint else None) or get_secret_str("AWS_BEDROCK_RUNTIME_ENDPOINT") ) # Partition-aware: bedrock-runtime is amazonaws.com.cn in China, and other @@ -237,7 +239,7 @@ class BedrockOpenAIResponsesConfig(BaseAWSLLM, OpenAIResponsesAPIConfig): (host[: -len(suffix)] for suffix in BEDROCK_RUNTIME_OPENAI_BASE_SUFFIXES if host.endswith(suffix)), host, ) - return f"{base}{BEDROCK_RUNTIME_OPENAI_RESPONSES_PATH}" + return join_api_base(base, BEDROCK_RUNTIME_OPENAI_RESPONSES_PATH) def supports_native_file_search(self) -> bool: return False diff --git a/litellm/llms/bedrock/vector_stores/transformation.py b/litellm/llms/bedrock/vector_stores/transformation.py index ba8ce7e5625..96d1da19bcc 100644 --- a/litellm/llms/bedrock/vector_stores/transformation.py +++ b/litellm/llms/bedrock/vector_stores/transformation.py @@ -6,7 +6,7 @@ from urllib.parse import urlparse import httpx from litellm._logging import verbose_logger -from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.llms.bedrock.common_utils import BedrockError @@ -196,7 +196,7 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): aws_region_name=self.get_aws_region_name_for_non_llm_api_calls(aws_region_name=aws_region_name), endpoint_type="agent", ) - return f"{endpoint_url}/knowledgebases" + return join_api_base(endpoint_url, "knowledgebases") def transform_search_vector_store_request( self, @@ -212,7 +212,7 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM): query = " ".join(query) encoded_vector_store_id: Final = encode_url_path_segment(vector_store_id, field_name="vector_store_id") - url: Final = f"{api_base}/{encoded_vector_store_id}/retrieve" + url: Final = join_api_base(api_base, f"{encoded_vector_store_id}/retrieve") request_body: Final[dict[str, object]] = { "retrievalQuery": BedrockKBRetrievalQuery(text=query), diff --git a/litellm/llms/bedrock_mantle/messages/transformation.py b/litellm/llms/bedrock_mantle/messages/transformation.py index 78881153e91..aa6510f0491 100644 --- a/litellm/llms/bedrock_mantle/messages/transformation.py +++ b/litellm/llms/bedrock_mantle/messages/transformation.py @@ -2,6 +2,7 @@ from collections.abc import Mapping from types import MappingProxyType from typing import Final +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.llms.bedrock.common_utils import MANTLE_MESSAGES_PATH from litellm.llms.bedrock.messages.mantle_transformation import AmazonMantleMessagesConfig @@ -33,7 +34,7 @@ def build_mantle_native_messages_url(api_base: str | None, litellm_params: Mappi configured, ) host: Final = f"https://bedrock-mantle.{region}.api.aws" if MANTLE_HOST_RE.match(stripped) else stripped - return f"{host}{MANTLE_MESSAGES_PATH}" + return join_api_base(host, MANTLE_MESSAGES_PATH) class BedrockMantleAnthropicMessagesConfig(BedrockMantleAuthMixin, AmazonMantleMessagesConfig): diff --git a/litellm/llms/bedrock_mantle/responses/transformation.py b/litellm/llms/bedrock_mantle/responses/transformation.py index 8a0ce58c8b2..3532cfbd963 100644 --- a/litellm/llms/bedrock_mantle/responses/transformation.py +++ b/litellm/llms/bedrock_mantle/responses/transformation.py @@ -22,6 +22,7 @@ import httpx import litellm from litellm._logging import verbose_logger +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.responses.codex_compat import drop_unsupported_tools, normalize_codex_input_items from litellm.llms.base_llm.responses.transformation import ReasoningSummaryLevel @@ -100,7 +101,7 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI if MANTLE_HOST_RE.match(base): base = f"https://bedrock-mantle.{region}.api.aws" path: Final = "/openai/v1/responses" if self.use_openai_path else "/v1/responses" - return f"{base}{path}" + return join_api_base(base, path) def validate_environment(self, headers: dict, model: str, litellm_params: GenericLiteLLMParams | None) -> dict: litellm_params = litellm_params or GenericLiteLLMParams() diff --git a/litellm/llms/black_forest_labs/image_edit/transformation.py b/litellm/llms/black_forest_labs/image_edit/transformation.py index 66eacc97fba..7cfb0daf334 100644 --- a/litellm/llms/black_forest_labs/image_edit/transformation.py +++ b/litellm/llms/black_forest_labs/image_edit/transformation.py @@ -19,7 +19,7 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH -from litellm.litellm_core_utils.url_utils import async_safe_get, safe_get +from litellm.litellm_core_utils.url_utils import async_safe_get, join_api_base, safe_get from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig from litellm.secret_managers.main import get_secret_str from litellm.types.images.main import ImageEditOptionalRequestParams @@ -189,11 +189,10 @@ class BlackForestLabsImageEditConfig(BaseImageEditConfig): """ Get the complete URL for the Black Forest Labs API request. """ - base_url: str = api_base or get_secret_str("BFL_API_BASE") or DEFAULT_API_BASE - base_url = base_url.rstrip("/") + base_url: Final = api_base or get_secret_str("BFL_API_BASE") or DEFAULT_API_BASE endpoint: Final = self._get_model_endpoint(model) - return f"{base_url}{endpoint}" + return join_api_base(base_url, endpoint) def _read_image_bytes( self, diff --git a/litellm/llms/black_forest_labs/image_generation/transformation.py b/litellm/llms/black_forest_labs/image_generation/transformation.py index 9d586fba26a..928af314506 100644 --- a/litellm/llms/black_forest_labs/image_generation/transformation.py +++ b/litellm/llms/black_forest_labs/image_generation/transformation.py @@ -14,6 +14,7 @@ from typing import TYPE_CHECKING, Any, Final import httpx from pydantic import ConfigDict, TypeAdapter +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.image_generation.transformation import ( BaseImageGenerationConfig, ) @@ -202,11 +203,10 @@ class BlackForestLabsImageGenerationConfig(BaseImageGenerationConfig): """ Get the complete URL for the Black Forest Labs API request. """ - base_url: str = api_base or get_secret_str("BFL_API_BASE") or DEFAULT_API_BASE - base_url = base_url.rstrip("/") + base_url: Final = api_base or get_secret_str("BFL_API_BASE") or DEFAULT_API_BASE endpoint: Final = self._get_model_endpoint(model) - return f"{base_url}{endpoint}" + return join_api_base(base_url, endpoint) def transform_image_generation_request( self, diff --git a/litellm/llms/chatgpt/responses/transformation.py b/litellm/llms/chatgpt/responses/transformation.py index 3e7db2981ef..34953a79c60 100644 --- a/litellm/llms/chatgpt/responses/transformation.py +++ b/litellm/llms/chatgpt/responses/transformation.py @@ -9,6 +9,7 @@ from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR, set_hid from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( safe_convert_created_field, ) +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.openai.common_utils import OpenAIError from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig from litellm.responses.sse_output_recovery import ( @@ -253,8 +254,7 @@ class ChatGPTResponsesAPIConfig(OpenAIResponsesAPIConfig): litellm_params: dict, ) -> str: api_base = api_base or self.authenticator.get_api_base() or CHATGPT_API_BASE - api_base = api_base.rstrip("/") - return f"{api_base}/responses" + return join_api_base(api_base, "responses") def supports_native_websocket(self) -> bool: """ChatGPT does not support native WebSocket for Responses API""" diff --git a/litellm/llms/cloudflare/chat/transformation.py b/litellm/llms/cloudflare/chat/transformation.py index ef11e51a5eb..0a76588f767 100644 --- a/litellm/llms/cloudflare/chat/transformation.py +++ b/litellm/llms/cloudflare/chat/transformation.py @@ -3,6 +3,7 @@ from typing import Final import httpx from litellm._logging import verbose_logger +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig from litellm.secret_managers.main import ( @@ -59,7 +60,7 @@ class CloudflareChatConfig(OpenAIGPTConfig): verbose_logger.warning( "Cloudflare api_base ending in '/ai/run' is the legacy Workers AI path and no longer serves OpenAI-compatible requests; rewriting to the '/ai/v1' endpoint" ) - return f"{trimmed[: -len('/ai/run')]}/ai/v1" + return join_api_base(trimmed[: -len("/ai/run")], "ai/v1") return api_base def validate_environment( diff --git a/litellm/llms/cloudflare/common_utils.py b/litellm/llms/cloudflare/common_utils.py index 9f1293cb9c4..02855a879a2 100644 --- a/litellm/llms/cloudflare/common_utils.py +++ b/litellm/llms/cloudflare/common_utils.py @@ -3,6 +3,7 @@ from typing import Final from pydantic import TypeAdapter, ValidationError +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.secret_managers.main import ( get_secret_str, normalize_nonempty_secret_str, @@ -34,10 +35,10 @@ def workers_ai_request_model(model: str) -> str: def workers_ai_run_url(api_base: str, model: str) -> str: normalized_api_base: Final = api_base.rstrip("/") if normalized_api_base.endswith("/ai/v1"): - return f"{normalized_api_base.removesuffix('/ai/v1')}/ai/run/{model}" + return join_api_base(normalized_api_base.removesuffix("/ai/v1"), f"ai/run/{model}") if normalized_api_base.endswith("/ai/run"): return f"{normalized_api_base}/{model}" - return f"{normalized_api_base}/ai/run/{model}" + return join_api_base(normalized_api_base, f"ai/run/{model}") def unwrap_workers_ai_result(payload: object) -> object: diff --git a/litellm/llms/cohere/rerank/transformation.py b/litellm/llms/cohere/rerank/transformation.py index b6d94d4853a..8b641c24e9b 100644 --- a/litellm/llms/cohere/rerank/transformation.py +++ b/litellm/llms/cohere/rerank/transformation.py @@ -6,6 +6,7 @@ from pydantic import ConfigDict, TypeAdapter import litellm from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig from litellm.secret_managers.main import get_secret_str @@ -31,11 +32,7 @@ class CohereRerankConfig(BaseRerankConfig): optional_params: dict | None = None, ) -> str: if api_base: - # Remove trailing slashes and ensure clean base URL - api_base = api_base.rstrip("/") - if not api_base.endswith("/v1/rerank"): - api_base = f"{api_base}/v1/rerank" - return api_base + return ensure_api_route(api_base, "v1/rerank") return "https://api.cohere.ai/v1/rerank" def get_supported_cohere_rerank_params(self, model: str) -> list: diff --git a/litellm/llms/cohere/rerank_v2/transformation.py b/litellm/llms/cohere/rerank_v2/transformation.py index e79357bf52f..009a220534c 100644 --- a/litellm/llms/cohere/rerank_v2/transformation.py +++ b/litellm/llms/cohere/rerank_v2/transformation.py @@ -1,5 +1,6 @@ from typing import Any, Final +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.cohere.rerank.transformation import CohereRerankConfig from litellm.types.rerank import OptionalRerankParams, RerankRequest @@ -19,11 +20,7 @@ class CohereRerankV2Config(CohereRerankConfig): optional_params: dict | None = None, ) -> str: if api_base: - # Remove trailing slashes and ensure clean base URL - api_base = api_base.rstrip("/") - if not api_base.endswith("/v2/rerank"): - api_base = f"{api_base}/v2/rerank" - return api_base + return ensure_api_route(api_base, "v2/rerank") return "https://api.cohere.ai/v2/rerank" def get_supported_cohere_rerank_params(self, model: str) -> list: diff --git a/litellm/llms/cometapi/chat/transformation.py b/litellm/llms/cometapi/chat/transformation.py index 26050700efc..2bf119a0a4a 100644 --- a/litellm/llms/cometapi/chat/transformation.py +++ b/litellm/llms/cometapi/chat/transformation.py @@ -10,6 +10,7 @@ from typing import Any, Final import httpx +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam @@ -112,16 +113,14 @@ class CometAPIConfig(OpenAIGPTConfig): # Ensure we include /v1 prefix when missing if api_base.endswith("/v1"): - return f"{api_base}/{endpoint}" - if api_base.endswith("/v1/"): - return f"{api_base}{endpoint}" + return join_api_base(api_base, endpoint) # If user provided https://api.cometapi.com, add /v1 if api_base == "https://api.cometapi.com": - return f"{api_base}/v1/{endpoint}" + return join_api_base(api_base, f"v1/{endpoint}") # Generic fallback: if '/v1' not in path, add it if "/v1" not in api_base.split("//", 1)[-1]: - return f"{api_base}/v1/{endpoint}" - return f"{api_base}/{endpoint}" + return join_api_base(api_base, f"v1/{endpoint}") + return join_api_base(api_base, endpoint) def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException: """ diff --git a/litellm/llms/cometapi/embed/transformation.py b/litellm/llms/cometapi/embed/transformation.py index 60f781be261..e3b5c722664 100644 --- a/litellm/llms/cometapi/embed/transformation.py +++ b/litellm/llms/cometapi/embed/transformation.py @@ -9,6 +9,7 @@ import httpx from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig from litellm.secret_managers.main import get_secret_str @@ -44,9 +45,7 @@ class CometAPIEmbeddingConfig(BaseEmbeddingConfig): """ Get the complete URL for the CometAPI embedding endpoint. """ - api_base = "https://api.cometapi.com/v1" if api_base is None else api_base.rstrip("/") - complete_url: Final = f"{api_base}/embeddings" - return complete_url + return join_api_base("https://api.cometapi.com/v1" if api_base is None else api_base, "embeddings") def validate_environment( self, diff --git a/litellm/llms/cometapi/image_generation/transformation.py b/litellm/llms/cometapi/image_generation/transformation.py index 567dbfd308e..ffaabbb28cd 100644 --- a/litellm/llms/cometapi/image_generation/transformation.py +++ b/litellm/llms/cometapi/image_generation/transformation.py @@ -4,6 +4,7 @@ from typing import TYPE_CHECKING, Any, Final import httpx from pydantic import ConfigDict, TypeAdapter +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.image_generation.transformation import ( BaseImageGenerationConfig, ) @@ -77,16 +78,13 @@ class CometAPIImageGenerationConfig(BaseImageGenerationConfig): """ Get the complete url for the request """ - complete_url: str = ( + base_url: Final = ( api_base or get_secret_str("COMETAPI_BASE_URL") or get_secret_str("COMETAPI_API_BASE") or self.DEFAULT_BASE_URL ) - - complete_url = complete_url.rstrip("/") - complete_url = f"{complete_url}/{self.IMAGE_GENERATION_ENDPOINT}" - return complete_url + return join_api_base(base_url, self.IMAGE_GENERATION_ENDPOINT) def validate_environment( self, diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index ac2ac1fe31b..a7ca0d5f49e 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -59,7 +59,7 @@ from litellm.litellm_core_utils.realtime_errors import ( websocket_close_reason, ) from litellm.litellm_core_utils.realtime_streaming import RealTimeStreaming -from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.llms.base_llm.anthropic_messages.transformation import ( BaseAnthropicMessagesConfig, ) @@ -6937,7 +6937,7 @@ class BaseLLMHTTPHandler: headers={}, model=model or "", api_key=api_key ) else: - url = f"{api_base.rstrip('/')}/v1/realtime/{endpoint}" + url = join_api_base(api_base, f"v1/realtime/{endpoint}") headers = { "Authorization": f"Bearer {api_key}", "Content-Type": "application/json", @@ -7006,7 +7006,7 @@ class BaseLLMHTTPHandler: url = provider_config.get_realtime_calls_url(api_base=api_base, model=model or "", api_version=api_version) headers: dict[str, object] = provider_config.get_realtime_calls_headers(ephemeral_key=openai_ephemeral_key) else: - url = f"{api_base.rstrip('/')}/v1/realtime/calls" + url = join_api_base(api_base, "v1/realtime/calls") headers = { "Authorization": f"Bearer {openai_ephemeral_key}", } diff --git a/litellm/llms/dashscope/chat/transformation.py b/litellm/llms/dashscope/chat/transformation.py index 9a1b381a2ba..966d1a01b8d 100644 --- a/litellm/llms/dashscope/chat/transformation.py +++ b/litellm/llms/dashscope/chat/transformation.py @@ -5,6 +5,7 @@ Translates from OpenAI's `/v1/chat/completions` to DashScope's `/v1/chat/complet from collections.abc import Coroutine from typing import Final, Literal, overload +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam @@ -82,7 +83,4 @@ class DashScopeChatConfig(OpenAIGPTConfig): """ If api_base is not provided, use the default DashScope /chat/completions endpoint. """ - resolved_api_base: Final = self._resolve_chat_api_base(api_base).rstrip("/") - if resolved_api_base.endswith("/chat/completions"): - return resolved_api_base - return f"{resolved_api_base}/chat/completions" + return ensure_api_route(self._resolve_chat_api_base(api_base), "chat/completions") diff --git a/litellm/llms/dashscope/embed/transformation.py b/litellm/llms/dashscope/embed/transformation.py index e5a987ca1ad..d0c2b4dff0c 100644 --- a/litellm/llms/dashscope/embed/transformation.py +++ b/litellm/llms/dashscope/embed/transformation.py @@ -16,6 +16,7 @@ from typing import Final import httpx from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig from litellm.secret_managers.main import get_secret_str @@ -98,10 +99,7 @@ class DashScopeEmbeddingConfig(BaseEmbeddingConfig): litellm_params: dict, stream: bool | None = None, ) -> str: - base: Final = self._resolve_embedding_api_base(api_base).rstrip("/") - if base.endswith("/embeddings"): - return base - return f"{base}/embeddings" + return ensure_api_route(self._resolve_embedding_api_base(api_base), "embeddings") def transform_embedding_request( self, diff --git a/litellm/llms/dashscope/rerank/transformation.py b/litellm/llms/dashscope/rerank/transformation.py index 4ec7d697f02..48ce4b5a07b 100644 --- a/litellm/llms/dashscope/rerank/transformation.py +++ b/litellm/llms/dashscope/rerank/transformation.py @@ -34,6 +34,7 @@ from pydantic import ConfigDict, TypeAdapter from litellm._uuid import uuid from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig from litellm.secret_managers.main import get_secret_str @@ -94,10 +95,10 @@ class DashScopeRerankConfig(BaseRerankConfig): return cleaned if cleaned.endswith("/v1"): - return f"{cleaned}/reranks" + return join_api_base(cleaned, "reranks") # Unknown base: append /reranks rather than silently ignoring the caller's api_base. - return f"{cleaned}/reranks" + return join_api_base(cleaned, "reranks") def validate_environment( self, diff --git a/litellm/llms/dashscope/videos/transformation.py b/litellm/llms/dashscope/videos/transformation.py index 1db91cbb127..5c738423a5a 100644 --- a/litellm/llms/dashscope/videos/transformation.py +++ b/litellm/llms/dashscope/videos/transformation.py @@ -16,7 +16,7 @@ from pydantic import BaseModel, Field, StrictBool, StrictInt, StrictStr, TypeAda import litellm from litellm.exceptions import UnsupportedParamsError -from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.videos.async_task_utils import ( input_reference_url, @@ -708,7 +708,7 @@ class DashScopeVideoConfig(BaseVideoConfig): if parameters: request_data["parameters"] = parameters - return request_data, (), f"{api_base}{DASHSCOPE_VIDEO_SYNTHESIS_PATH}" + return request_data, (), join_api_base(api_base, DASHSCOPE_VIDEO_SYNTHESIS_PATH) @staticmethod def _reference_input( @@ -773,7 +773,7 @@ class DashScopeVideoConfig(BaseVideoConfig): 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}{DASHSCOPE_TASKS_PATH}/{encoded_task_id}" + return join_api_base(api_base, f"{DASHSCOPE_TASKS_PATH}/{encoded_task_id}") def transform_video_status_retrieve_response( self, diff --git a/litellm/llms/databricks/chat/transformation.py b/litellm/llms/databricks/chat/transformation.py index d46d04578cc..61b41c2a173 100644 --- a/litellm/llms/databricks/chat/transformation.py +++ b/litellm/llms/databricks/chat/transformation.py @@ -22,6 +22,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( strip_name_from_message, strip_prefill_prefix_field, ) +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator from litellm.llms.base_llm.base_utils import ( type_to_response_format_param, # pyright: ignore[reportUnknownVariableType] # base_utils helper returns an untyped dict @@ -339,10 +340,7 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): stream: bool | None = None, ) -> str: use_ai_gateway: Final = self.uses_ai_gateway(model) - api_base = self._get_api_base(api_base, use_ai_gateway=use_ai_gateway) - url_base: Final = api_base.rstrip("/") if use_ai_gateway else api_base - complete_url: Final = f"{url_base}/chat/completions" - return complete_url + return join_api_base(self._get_api_base(api_base, use_ai_gateway=use_ai_gateway), "chat/completions") def get_supported_openai_params(self, model: str | None = None) -> list: return [ diff --git a/litellm/llms/databricks/common_utils.py b/litellm/llms/databricks/common_utils.py index 3112b2458b4..cd9432ac36f 100644 --- a/litellm/llms/databricks/common_utils.py +++ b/litellm/llms/databricks/common_utils.py @@ -20,6 +20,7 @@ from urllib.parse import urlsplit, urlunsplit from pydantic import BaseModel, ConfigDict import litellm +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.secret_managers.main import get_secret_str @@ -278,10 +279,10 @@ class DatabricksBase: if normalized_api_base.endswith("/ai-gateway/mlflow/v1"): return normalized_api_base if normalized_api_base.endswith("/serving-endpoints"): - return f"{normalized_api_base.removesuffix('/serving-endpoints')}/ai-gateway/mlflow/v1" + return join_api_base(normalized_api_base.removesuffix("/serving-endpoints"), "ai-gateway/mlflow/v1") api_base_parts: Final = urlsplit(normalized_api_base) if api_base_parts.path in ("", "/"): - return f"{normalized_api_base}/ai-gateway/mlflow/v1" + return join_api_base(normalized_api_base, "ai-gateway/mlflow/v1") return api_base def get_oauth_m2m_token( @@ -319,7 +320,7 @@ class DatabricksBase: api_base_parts: Final = urlsplit(api_base) workspace_url: Final = urlunsplit((api_base_parts.scheme, api_base_parts.netloc, "", "", "")) - token_url: Final = f"{workspace_url}/oidc/v1/token" + token_url: Final = join_api_base(workspace_url, "oidc/v1/token") try: response: Final = requests.post( @@ -477,5 +478,6 @@ class DatabricksBase: # Debug logging with redaction (never log actual tokens) verbose_logger.debug("Databricks request headers: %s", self.redact_headers_for_logging(headers)) - endpoint_path: Final = "" if custom_endpoint is True else DATABRICKS_ENDPOINT_PATHS[endpoint_type] - return f"{api_base}{endpoint_path}", headers + if custom_endpoint is True: + return api_base, headers + return join_api_base(api_base, DATABRICKS_ENDPOINT_PATHS[endpoint_type]), headers diff --git a/litellm/llms/databricks/decisions/transformation.py b/litellm/llms/databricks/decisions/transformation.py index c2ac1f220b8..d1698fdf5e2 100644 --- a/litellm/llms/databricks/decisions/transformation.py +++ b/litellm/llms/databricks/decisions/transformation.py @@ -6,6 +6,7 @@ from typing import Final from pydantic import TypeAdapter, ValidationError +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.decisions.transformation import BaseDecisionsConfig from litellm.types.decisions import DecisionsIRRequest, UnsupportedDecisionsRequest @@ -84,8 +85,8 @@ class DatabricksDecisionsConfig(BaseDecisionsConfig): def get_complete_url(self, api_base: str, model: str) -> str: if is_ai_decide_model(model): - return f"{databricks_workspace_host(api_base)}{_AI_DECIDE_PATH}" - return f"{api_base.rstrip('/')}/{validate_databricks_decisions_model(model)}/invocations" + return join_api_base(databricks_workspace_host(api_base), _AI_DECIDE_PATH) + return join_api_base(api_base, f"{validate_databricks_decisions_model(model)}/invocations") def transform_decisions_request( self, diff --git a/litellm/llms/databricks/responses/transformation.py b/litellm/llms/databricks/responses/transformation.py index 804c79ea79d..55c3edfefd0 100644 --- a/litellm/llms/databricks/responses/transformation.py +++ b/litellm/llms/databricks/responses/transformation.py @@ -10,6 +10,7 @@ Reference: https://docs.databricks.com/aws/en/machine-learning/foundation-model- import os from typing import TYPE_CHECKING, Any, Final +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.databricks.common_utils import DatabricksBase from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig from litellm.types.llms.openai import ResponseInputParam @@ -70,8 +71,7 @@ class DatabricksResponsesAPIConfig(DatabricksBase, OpenAIResponsesAPIConfig): ) -> str: api_base = api_base or os.getenv("DATABRICKS_API_BASE") api_base = self._get_api_base(api_base) - api_base = api_base.rstrip("/") - return f"{api_base}/responses" + return join_api_base(api_base, "responses") def transform_responses_api_request( self, diff --git a/litellm/llms/deepgram/audio_transcription/transformation.py b/litellm/llms/deepgram/audio_transcription/transformation.py index a3812dac8ef..a1c3021c9f5 100644 --- a/litellm/llms/deepgram/audio_transcription/transformation.py +++ b/litellm/llms/deepgram/audio_transcription/transformation.py @@ -10,6 +10,7 @@ from httpx import Headers, Response from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.audio_utils.utils import process_audio_file +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import ( @@ -179,7 +180,6 @@ class DeepgramAudioTranscriptionConfig(BaseAudioTranscriptionConfig): ) -> str: if api_base is None: api_base = get_secret_str("DEEPGRAM_API_BASE") or "https://api.deepgram.com/v1" - api_base = api_base.rstrip("/") # Remove trailing slash if present # Build query parameters including the model all_query_params: Final = {"model": model} @@ -189,7 +189,7 @@ class DeepgramAudioTranscriptionConfig(BaseAudioTranscriptionConfig): all_query_params.update(additional_params) # Construct URL with proper query string encoding - base_url: Final = f"{api_base}/listen" + base_url: Final = join_api_base(api_base, "listen") query_string: Final = urlencode(all_query_params) url: Final = f"{base_url}?{query_string}" diff --git a/litellm/llms/deepgram/common_utils.py b/litellm/llms/deepgram/common_utils.py index 9b071ac8321..29998895bbc 100644 --- a/litellm/llms/deepgram/common_utils.py +++ b/litellm/llms/deepgram/common_utils.py @@ -8,6 +8,7 @@ import httpx import litellm from litellm.constants import DEEPGRAM_DEFAULT_API_BASE, DEEPGRAM_LISTEN_DEFAULT_MODEL +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.types.utils import LlmProviders @@ -50,7 +51,7 @@ def _first_occurrences(query_string: str) -> httpx.QueryParams: def deepgram_listen_websocket_target(api_base: str | None, query_string: str) -> str: - listen_url: Final = httpx.URL(f"{(api_base or DEEPGRAM_DEFAULT_API_BASE).rstrip('/')}/listen") + listen_url: Final = httpx.URL(join_api_base(api_base or DEEPGRAM_DEFAULT_API_BASE, "listen")) websocket_url: Final = listen_url.copy_with(scheme=_WEBSOCKET_SCHEMES.get(listen_url.scheme, listen_url.scheme)) params: Final = _first_occurrences(query_string) query: Final = params if params.get("model") else params.remove("model").add("model", DEEPGRAM_LISTEN_DEFAULT_MODEL) diff --git a/litellm/llms/deepseek/chat/transformation.py b/litellm/llms/deepseek/chat/transformation.py index d3319856d29..a9a20101bb1 100644 --- a/litellm/llms/deepseek/chat/transformation.py +++ b/litellm/llms/deepseek/chat/transformation.py @@ -10,6 +10,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( convert_content_list_to_str, extract_search_results_text, ) +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues from litellm.utils import supports_reasoning, supports_vision @@ -363,9 +364,4 @@ class DeepSeekChatConfig(OpenAIGPTConfig): """ If api_base is not provided, use the default DeepSeek /chat/completions endpoint. """ - clean_api_base: Final = (api_base or "https://api.deepseek.com/beta").rstrip("/") - - if not clean_api_base.endswith("/chat/completions"): - return f"{clean_api_base}/chat/completions" - - return clean_api_base + return ensure_api_route(api_base or "https://api.deepseek.com/beta", "chat/completions") diff --git a/litellm/llms/deepseek/messages/transformation.py b/litellm/llms/deepseek/messages/transformation.py index 6a6b57131ad..6ddee99aea3 100644 --- a/litellm/llms/deepseek/messages/transformation.py +++ b/litellm/llms/deepseek/messages/transformation.py @@ -6,6 +6,7 @@ from collections.abc import Sequence from typing import Any, Final import litellm +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.anthropic.pass_through.messages.transformation import ( AnthropicMessagesConfig, ) @@ -90,9 +91,9 @@ class DeepSeekAnthropicMessagesConfig(AnthropicMessagesConfig): base_url = base_url.removesuffix("/beta") if not base_url.endswith("/anthropic") and "/anthropic/" not in base_url: - base_url = f"{base_url}/anthropic" + base_url = join_api_base(base_url, "anthropic") - return f"{base_url}/v1/messages" + return join_api_base(base_url, "v1/messages") @staticmethod def _sanitize_tools_for_deepseek(tools: Any) -> object: diff --git a/litellm/llms/docker_model_runner/chat/transformation.py b/litellm/llms/docker_model_runner/chat/transformation.py index adb40311efd..68b63812248 100644 --- a/litellm/llms/docker_model_runner/chat/transformation.py +++ b/litellm/llms/docker_model_runner/chat/transformation.py @@ -10,6 +10,7 @@ from typing import Final, Literal, overload from litellm.litellm_core_utils.prompt_templates.common_utils import ( handle_messages_with_content_list_to_str_conversion, ) +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues @@ -104,12 +105,9 @@ class DockerModelRunnerChatConfig(OpenAIGPTConfig): if not api_base: api_base = "http://localhost:22088/engines/llama.cpp" - # Remove trailing slashes from api_base - api_base = api_base.rstrip("/") - # Build the URL: {api_base}/v1/chat/completions # api_base is expected to already contain the engine path - complete_url: Final = f"{api_base}/v1/chat/completions" + complete_url: Final = join_api_base(api_base, "v1/chat/completions") return complete_url diff --git a/litellm/llms/duckduckgo/search/transformation.py b/litellm/llms/duckduckgo/search/transformation.py index 1816a0d24e4..36893b4bd27 100644 --- a/litellm/llms/duckduckgo/search/transformation.py +++ b/litellm/llms/duckduckgo/search/transformation.py @@ -12,6 +12,7 @@ import httpx from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.search.transformation import ( BaseSearchConfig, SearchResponse, @@ -91,7 +92,7 @@ class DuckDuckGoSearchConfig(BaseSearchConfig): if data and isinstance(data, dict) and "_duckduckgo_params" in data: params: Final = data["_duckduckgo_params"] query_string: Final = urlencode(params, doseq=True) - return f"{api_base}/?{query_string}" + return f"{join_api_base(api_base, '')}?{query_string}" return api_base diff --git a/litellm/llms/e2b/sandbox/transformation.py b/litellm/llms/e2b/sandbox/transformation.py index dbf3996248a..b5f42632eef 100644 --- a/litellm/llms/e2b/sandbox/transformation.py +++ b/litellm/llms/e2b/sandbox/transformation.py @@ -15,6 +15,7 @@ import httpx from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.sandbox.transformation import ( SANDBOX_MAX_OUTPUT_BYTES, BaseSandboxConfig, @@ -78,7 +79,7 @@ class E2BSandboxConfig(BaseSandboxConfig): body["metadata"] = metadata response: Final = await self._http(client).post( - url=f"{base}/sandboxes", + url=join_api_base(base, "sandboxes"), headers={"X-API-Key": key, "Content-Type": "application/json"}, json=body, ) @@ -151,10 +152,10 @@ class E2BSandboxConfig(BaseSandboxConfig): dict[str, object], getattr(handle, HIDDEN_PARAMS_ATTR) ) key: Final = api_key or hidden_params.get("api_key") or self.validate_environment() - base: Final = api_base or hidden_params.get("api_base") or E2B_API_BASE + base: Final = str(api_base or hidden_params.get("api_base") or E2B_API_BASE) try: response: Final = await self._http(client).delete( - url=f"{base}/sandboxes/{handle.id}", + url=join_api_base(base, f"sandboxes/{handle.id}"), headers={"X-API-Key": key}, ) except httpx.HTTPStatusError as e: diff --git a/litellm/llms/edenai/chat/transformation.py b/litellm/llms/edenai/chat/transformation.py index ae5da5d1cfd..480b382076d 100644 --- a/litellm/llms/edenai/chat/transformation.py +++ b/litellm/llms/edenai/chat/transformation.py @@ -23,7 +23,7 @@ from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import ModelResponse, ModelResponseStream, Usage -from ..common_utils import EdenAIException, reported_cost, resolve_api_base, resolve_api_key +from ..common_utils import EdenAIException, endpoint_url, reported_cost, resolve_api_base, resolve_api_key if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -138,7 +138,7 @@ class EdenAIChatConfig(OpenAIGPTConfig): def get_models( self, api_key: str | None = None, api_base: str | None = None ) -> list[str]: # mutable-ok: inherited contract - response: Final = litellm.module_level_client.get(url=f"{self.get_api_base(api_base)}/models") + response: Final = litellm.module_level_client.get(url=endpoint_url(api_base, "models")) if not response.is_success: raise EdenAIException(status_code=response.status_code, message=response.text, headers=response.headers) catalog: Final = _EdenAIModelCatalog.model_validate(response.json()) diff --git a/litellm/llms/edenai/common_utils.py b/litellm/llms/edenai/common_utils.py index 3216f798fbe..60c83f7ea75 100644 --- a/litellm/llms/edenai/common_utils.py +++ b/litellm/llms/edenai/common_utils.py @@ -11,6 +11,7 @@ from pydantic import AliasChoices, Field, ValidationError import litellm from litellm.exceptions import AuthenticationError +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.secret_managers.main import get_secret_str from litellm.types.llms.base import LiteLLMBaseModel @@ -74,7 +75,7 @@ def json_headers( def endpoint_url(api_base: str | None, path: str) -> str: - return f"{resolve_api_base(api_base).rstrip('/')}/{path}" + return join_api_base(resolve_api_base(api_base), path) def pick(params: Mapping[str, object], keys: Container[str]) -> Mapping[str, object]: diff --git a/litellm/llms/elevenlabs/audio_transcription/transformation.py b/litellm/llms/elevenlabs/audio_transcription/transformation.py index 04d2717d64f..5a673b5ac41 100644 --- a/litellm/llms/elevenlabs/audio_transcription/transformation.py +++ b/litellm/llms/elevenlabs/audio_transcription/transformation.py @@ -10,6 +10,7 @@ from pydantic import ConfigDict, TypeAdapter import litellm from litellm.litellm_core_utils.audio_utils.utils import process_audio_file +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import ( @@ -165,12 +166,7 @@ class ElevenLabsAudioTranscriptionConfig(BaseAudioTranscriptionConfig): ) -> str: if api_base is None: api_base = get_secret_str("ELEVENLABS_API_BASE") or "https://api.elevenlabs.io" - api_base = api_base.rstrip("/") # Remove trailing slash if present - - # ElevenLabs speech-to-text endpoint - url: Final = f"{api_base}/v1/speech-to-text" - - return url + return join_api_base(api_base, "v1/speech-to-text") def validate_environment( self, diff --git a/litellm/llms/elevenlabs/text_to_speech/transformation.py b/litellm/llms/elevenlabs/text_to_speech/transformation.py index eb93543df46..f910ef8c057 100644 --- a/litellm/llms/elevenlabs/text_to_speech/transformation.py +++ b/litellm/llms/elevenlabs/text_to_speech/transformation.py @@ -11,7 +11,7 @@ import httpx from httpx import Headers import litellm -from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.text_to_speech.transformation import ( BaseTextToSpeechConfig, @@ -297,15 +297,14 @@ class ElevenLabsTextToSpeechConfig(BaseTextToSpeechConfig): """ Construct the ElevenLabs endpoint URL, including path voice_id and query params. """ - base_url = api_base or get_secret_str("ELEVENLABS_API_BASE") or self.TTS_BASE_URL - base_url = base_url.rstrip("/") + base_url: Final = api_base or get_secret_str("ELEVENLABS_API_BASE") or self.TTS_BASE_URL voice_id: Final = litellm_params.get(self.ELEVENLABS_VOICE_ID_KEY) if not isinstance(voice_id, str) or not voice_id.strip(): raise ValueError("ElevenLabs voice_id is required. Pass `voice` when calling `litellm.speech()`.") encoded_voice_id: Final = encode_url_path_segment(voice_id, field_name="voice_id") - url = f"{base_url}{self.TTS_ENDPOINT_PATH}/{encoded_voice_id}" + url = join_api_base(base_url, f"{self.TTS_ENDPOINT_PATH}/{encoded_voice_id}") query_params: Final = litellm_params.get(self.ELEVENLABS_QUERY_PARAMS_KEY, {}) if query_params: diff --git a/litellm/llms/exa_ai/search/transformation.py b/litellm/llms/exa_ai/search/transformation.py index 89073409599..73e341d0033 100644 --- a/litellm/llms/exa_ai/search/transformation.py +++ b/litellm/llms/exa_ai/search/transformation.py @@ -11,6 +11,7 @@ import httpx from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm.search.transformation import ( BaseSearchConfig, SearchResponse, @@ -93,11 +94,7 @@ class ExaAISearchConfig(BaseSearchConfig): """ Get complete URL for Search endpoint. """ - resolved_api_base: Final = api_base or get_secret_str("EXA_API_BASE") or self.EXA_AI_API_BASE - clean_api_base: Final = resolved_api_base.rstrip("/") - if clean_api_base.endswith("/search"): - return clean_api_base - return f"{clean_api_base}/search" + return ensure_api_route(api_base or get_secret_str("EXA_API_BASE") or self.EXA_AI_API_BASE, "search") def transform_search_request( self, diff --git a/litellm/llms/fal_ai/image_edit/flux_lora_depth_transformation.py b/litellm/llms/fal_ai/image_edit/flux_lora_depth_transformation.py index 6d58caeb384..344d9d3c75b 100644 --- a/litellm/llms/fal_ai/image_edit/flux_lora_depth_transformation.py +++ b/litellm/llms/fal_ai/image_edit/flux_lora_depth_transformation.py @@ -4,6 +4,7 @@ from typing import Final from httpx._types import RequestFiles +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.secret_managers.main import get_secret_str from litellm.types.images.main import ImageEditOptionalRequestParams from litellm.types.router import GenericLiteLLMParams @@ -45,8 +46,8 @@ class FalAIFluxLoraDepthEditConfig(FalAIImageEditConfig): api_base: str | None, litellm_params: dict, # mutable-ok: base class contract ) -> str: - base_url: Final = (api_base or get_secret_str("FAL_AI_API_BASE") or DEFAULT_BASE_URL).rstrip("/") - return f"{base_url}/{FLUX_LORA_DEPTH_ENDPOINT}" + base_url: Final = api_base or get_secret_str("FAL_AI_API_BASE") or DEFAULT_BASE_URL + return join_api_base(base_url, FLUX_LORA_DEPTH_ENDPOINT) def transform_image_edit_request( self, diff --git a/litellm/llms/fal_ai/image_generation/transformation.py b/litellm/llms/fal_ai/image_generation/transformation.py index 0bb5dada071..79913d5c651 100644 --- a/litellm/llms/fal_ai/image_generation/transformation.py +++ b/litellm/llms/fal_ai/image_generation/transformation.py @@ -5,6 +5,7 @@ import httpx from pydantic import ConfigDict, TypeAdapter from typing_extensions import ReadOnly, TypedDict +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.image_generation.transformation import ( BaseImageGenerationConfig, ) @@ -85,12 +86,10 @@ class FalAIBaseConfig(BaseImageGenerationConfig): Some providers need `model` in `api_base` """ - complete_url: str = api_base or get_secret_str("FAL_AI_API_BASE") or self.DEFAULT_BASE_URL - - complete_url = complete_url.rstrip("/") - if self.IMAGE_GENERATION_ENDPOINT: - complete_url = f"{complete_url}/{self.IMAGE_GENERATION_ENDPOINT}" - return complete_url + base_url: Final = api_base or get_secret_str("FAL_AI_API_BASE") or self.DEFAULT_BASE_URL + if not self.IMAGE_GENERATION_ENDPOINT: + return base_url.rstrip("/") + return join_api_base(base_url, self.IMAGE_GENERATION_ENDPOINT) def validate_environment( self, diff --git a/litellm/llms/fal_ai/videos/transformation.py b/litellm/llms/fal_ai/videos/transformation.py index 3136763e1be..86aa6388ff0 100644 --- a/litellm/llms/fal_ai/videos/transformation.py +++ b/litellm/llms/fal_ai/videos/transformation.py @@ -10,7 +10,7 @@ import httpx from httpx._types import FileContent, RequestFiles from pydantic import TypeAdapter -from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.videos.transformation import BaseVideoConfig from litellm.llms.custom_httpx.http_handler import ( @@ -445,7 +445,7 @@ class FalAIVideoConfig(BaseVideoConfig): request_id, model_id = self._decode_video_id(video_id) encoded_request_id: Final[str] = encode_url_path_segment(request_id, field_name="video_id") return ( - f"{api_base.rstrip('/')}/{_queue_request_base_path(model_id)}/requests/{encoded_request_id}/status", + join_api_base(api_base, f"{_queue_request_base_path(model_id)}/requests/{encoded_request_id}/status"), {}, ) @@ -541,7 +541,7 @@ class FalAIVideoConfig(BaseVideoConfig): request_id, model_id = self._decode_video_id(video_id) encoded_request_id: Final[str] = encode_url_path_segment(request_id, field_name="video_id") return ( - f"{api_base.rstrip('/')}/{_queue_request_base_path(model_id)}/requests/{encoded_request_id}", + join_api_base(api_base, f"{_queue_request_base_path(model_id)}/requests/{encoded_request_id}"), {}, ) diff --git a/litellm/llms/fastcrw/search/transformation.py b/litellm/llms/fastcrw/search/transformation.py index 8f6f8d87300..27db6a95170 100644 --- a/litellm/llms/fastcrw/search/transformation.py +++ b/litellm/llms/fastcrw/search/transformation.py @@ -15,6 +15,7 @@ import httpx from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm.search.transformation import ( BaseSearchConfig, SearchResponse, @@ -83,11 +84,7 @@ class FastCRWSearchConfig(BaseSearchConfig): """ Get complete URL for Search endpoint. """ - resolved_api_base: Final = api_base or get_secret_str("CRW_API_BASE") or self.FASTCRW_API_BASE - clean_api_base: Final = resolved_api_base.rstrip("/") - if clean_api_base.endswith("/search"): - return clean_api_base - return f"{clean_api_base}/search" + return ensure_api_route(api_base or get_secret_str("CRW_API_BASE") or self.FASTCRW_API_BASE, "search") def transform_search_request( self, diff --git a/litellm/llms/firecrawl/search/transformation.py b/litellm/llms/firecrawl/search/transformation.py index ee99ab24413..6d1e2b22918 100644 --- a/litellm/llms/firecrawl/search/transformation.py +++ b/litellm/llms/firecrawl/search/transformation.py @@ -11,6 +11,7 @@ import httpx from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm.search.transformation import ( BaseSearchConfig, SearchResponse, @@ -85,11 +86,7 @@ class FirecrawlSearchConfig(BaseSearchConfig): """ Get complete URL for Search endpoint. """ - resolved_api_base: Final = api_base or get_secret_str("FIRECRAWL_API_BASE") or self.FIRECRAWL_API_BASE - clean_api_base: Final = resolved_api_base.rstrip("/") - if clean_api_base.endswith("/search"): - return clean_api_base - return f"{clean_api_base}/search" + return ensure_api_route(api_base or get_secret_str("FIRECRAWL_API_BASE") or self.FIRECRAWL_API_BASE, "search") def transform_search_request( self, diff --git a/litellm/llms/fireworks_ai/chat/transformation.py b/litellm/llms/fireworks_ai/chat/transformation.py index 14ae59d4f85..fa5794d71b3 100644 --- a/litellm/llms/fireworks_ai/chat/transformation.py +++ b/litellm/llms/fireworks_ai/chat/transformation.py @@ -17,6 +17,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( strip_litellm_response_only_fields, unpack_legacy_defs, ) +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import ( AllMessageValues, @@ -826,7 +827,7 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig): base = api_base.rstrip("/") base = base.removesuffix("/v1") response: Final = litellm.module_level_client.get( - url=f"{base}/v1/accounts/{account_id}/models", + url=join_api_base(base, f"v1/accounts/{account_id}/models"), headers={"Authorization": f"Bearer {api_key}"}, ) diff --git a/litellm/llms/fireworks_ai/rerank/transformation.py b/litellm/llms/fireworks_ai/rerank/transformation.py index 4eff64e351f..964f85c821e 100644 --- a/litellm/llms/fireworks_ai/rerank/transformation.py +++ b/litellm/llms/fireworks_ai/rerank/transformation.py @@ -12,6 +12,7 @@ from pydantic import ConfigDict, TypeAdapter from litellm._uuid import uuid from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig from litellm.llms.fireworks_ai.common_utils import FireworksAIMixin from litellm.types.llms.base import LiteLLMBaseModel @@ -61,9 +62,9 @@ class FireworksAIRerankConfig(FireworksAIMixin, BaseRerankConfig): api_base = api_base.rstrip("/") if not api_base.endswith("/rerank"): if api_base.endswith("/v1") or api_base.endswith("/inference/v1"): - api_base = f"{api_base}/rerank" + api_base = join_api_base(api_base, "rerank") else: - api_base = f"{api_base}/inference/v1/rerank" + api_base = join_api_base(api_base, "inference/v1/rerank") return api_base return "https://api.fireworks.ai/inference/v1/rerank" diff --git a/litellm/llms/fireworks_ai/responses/transformation.py b/litellm/llms/fireworks_ai/responses/transformation.py index 30eaa5a7ca2..db234c09930 100644 --- a/litellm/llms/fireworks_ai/responses/transformation.py +++ b/litellm/llms/fireworks_ai/responses/transformation.py @@ -6,6 +6,7 @@ from urllib.parse import unquote import httpx from openai.types.responses import EasyInputMessageParam, ResponseInputContentParam, ResponseInputItemParam +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.fireworks_ai.common_utils import ( FIREWORKS_FORWARD_USER_ID_PARAM, get_fireworks_forwarded_user_id, @@ -150,8 +151,8 @@ class FireworksAIResponsesAPIConfig(OpenAIResponsesAPIConfig): return dict(pinned) def get_complete_url(self, api_base: str | None, litellm_params: Mapping[str, object]) -> str: - base: Final = (api_base or get_secret_str("FIREWORKS_API_BASE") or FIREWORKS_AI_DEFAULT_API_BASE).rstrip("/") - return f"{base}/responses" + base: Final = api_base or get_secret_str("FIREWORKS_API_BASE") or FIREWORKS_AI_DEFAULT_API_BASE + return join_api_base(base, "responses") def transform_responses_api_request( self, diff --git a/litellm/llms/gdc/chat/transformation.py b/litellm/llms/gdc/chat/transformation.py index 6eac3ac79cd..cd9ffe9e03e 100644 --- a/litellm/llms/gdc/chat/transformation.py +++ b/litellm/llms/gdc/chat/transformation.py @@ -13,6 +13,7 @@ from urllib.parse import urlsplit from typing_extensions import ReadOnly, TypedDict, Unpack import litellm +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.openai_like.chat.transformation import OpenAILikeChatConfig from litellm.types.llms.openai import AllMessageValues @@ -140,7 +141,7 @@ class GDCGeminiConfig(OpenAILikeChatConfig): project = self._validate_path_id(project, "vertex_project", model) location = self._validate_path_id(location, "vertex_location", model) - return f"{api_base}/v1/projects/{project}/locations/{location}/chat/completions" + return join_api_base(api_base, f"v1/projects/{project}/locations/{location}/chat/completions") def _read_env_bool(self, val: bool | str | None, env_var: str, default: bool = True) -> bool | str: def _parse(s: str) -> bool | str: diff --git a/litellm/llms/gemini/agents/transformation.py b/litellm/llms/gemini/agents/transformation.py index cfdb55f2048..c19bd9ce87c 100644 --- a/litellm/llms/gemini/agents/transformation.py +++ b/litellm/llms/gemini/agents/transformation.py @@ -15,6 +15,7 @@ from typing import Any, Final import httpx from litellm._logging import verbose_logger +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.agents.transformation import BaseAgentsAPIConfig from litellm.llms.gemini.common_utils import GeminiError, GeminiModelInfo from litellm.types.agents import ( @@ -67,7 +68,7 @@ class GeminiAgentsConfig(BaseAgentsAPIConfig): return "v1beta" def _base_url(self, api_base: str | None) -> str: - return f"{GeminiModelInfo.get_api_base(api_base)}/{self.api_version}" + return join_api_base(GeminiModelInfo.get_api_base(api_base), self.api_version) # ------------------------------------------------------------------ # # Shared helpers # @@ -90,7 +91,7 @@ class GeminiAgentsConfig(BaseAgentsAPIConfig): api_base: str | None, litellm_params: Mapping[str, object], ) -> str: - return f"{self._base_url(api_base)}/agents" + return join_api_base(self._base_url(api_base), "agents") def validate_environment( self, @@ -177,7 +178,7 @@ class GeminiAgentsConfig(BaseAgentsAPIConfig): api_base: str | None, litellm_params: Mapping[str, object], ) -> tuple[str, dict[str, object]]: - url: Final = f"{self._base_url(api_base)}/agents" + url: Final = join_api_base(self._base_url(api_base), "agents") params: Final[dict[str, object]] = {} if litellm_params.get("page_size"): params["pageSize"] = litellm_params["page_size"] @@ -210,7 +211,7 @@ class GeminiAgentsConfig(BaseAgentsAPIConfig): api_base: str | None, litellm_params: Mapping[str, object], ) -> tuple[str, dict[str, object]]: - url: Final = f"{self._base_url(api_base)}/agents/{name}" + url: Final = join_api_base(self._base_url(api_base), f"agents/{name}") return url, {} def transform_get_response( @@ -239,7 +240,7 @@ class GeminiAgentsConfig(BaseAgentsAPIConfig): api_base: str | None, litellm_params: Mapping[str, object], ) -> str: - return f"{self._base_url(api_base)}/agents/{name}" + return join_api_base(self._base_url(api_base), f"agents/{name}") def transform_delete_response( self, @@ -265,7 +266,7 @@ class GeminiAgentsConfig(BaseAgentsAPIConfig): api_base: str | None, litellm_params: Mapping[str, object], ) -> tuple[str, dict[str, object]]: - url: Final = f"{self._base_url(api_base)}/agents/{name}/versions" + url: Final = join_api_base(self._base_url(api_base), f"agents/{name}/versions") params: Final[dict[str, object]] = {} if litellm_params.get("page_size"): params["pageSize"] = litellm_params["page_size"] diff --git a/litellm/llms/gemini/audio_transcription/transformation.py b/litellm/llms/gemini/audio_transcription/transformation.py index 48a345a4940..ea116a7e9de 100644 --- a/litellm/llms/gemini/audio_transcription/transformation.py +++ b/litellm/llms/gemini/audio_transcription/transformation.py @@ -9,6 +9,7 @@ from litellm.litellm_core_utils.audio_utils.utils import ( normalize_transcription_language_to_bcp47, process_audio_file, ) +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.audio_transcription.transformation import ( AudioTranscriptionRequestData, BaseAudioTranscriptionConfig, @@ -104,8 +105,7 @@ class GeminiAudioTranscriptionConfig(BaseAudioTranscriptionConfig): litellm_params: Mapping[str, object], stream: bool | None = None, ) -> str: - resolved_api_base: Final = GeminiModelInfo.get_api_base(api_base) - return f"{resolved_api_base}/v1beta/interactions" + return join_api_base(GeminiModelInfo.get_api_base(api_base), "v1beta/interactions") def transform_audio_transcription_request( self, diff --git a/litellm/llms/gemini/common_utils.py b/litellm/llms/gemini/common_utils.py index 8ee7a463828..13df3921d13 100644 --- a/litellm/llms/gemini/common_utils.py +++ b/litellm/llms/gemini/common_utils.py @@ -11,6 +11,7 @@ from pydantic import JsonValue, TypeAdapter import litellm from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.secret_managers.main import get_secret_str @@ -377,7 +378,7 @@ class GeminiModelInfo(BaseLLMModelInfo): return "v1beta" @staticmethod - def get_api_base(api_base: str | None = None) -> str | None: + def get_api_base(api_base: str | None = None) -> str: return api_base or get_secret_str("GEMINI_API_BASE") or "https://generativelanguage.googleapis.com" @staticmethod @@ -400,13 +401,13 @@ class GeminiModelInfo(BaseLLMModelInfo): api_base = GeminiModelInfo.get_api_base(api_base) api_key = GeminiModelInfo.get_api_key(api_key) endpoint: Final = f"/{self.api_version}/models" - if api_base is None or api_key is None: + if api_key is None: raise ValueError( "GEMINI_API_BASE or GEMINI_API_KEY/GOOGLE_API_KEY is not set. Please set the environment variable, to query Gemini's `/models` endpoint." ) response: Final = litellm.module_level_client.get( - url=f"{api_base}{endpoint}", + url=join_api_base(api_base, endpoint), headers={"x-goog-api-key": api_key}, ) diff --git a/litellm/llms/gemini/count_tokens/handler.py b/litellm/llms/gemini/count_tokens/handler.py index 1c9d09330a8..4ed47323d4f 100644 --- a/litellm/llms/gemini/count_tokens/handler.py +++ b/litellm/llms/gemini/count_tokens/handler.py @@ -6,6 +6,7 @@ from pydantic import ConfigDict, RootModel, with_config from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, get_async_httpx_client from litellm.types.utils import LlmProviders @@ -65,7 +66,7 @@ class GoogleAIStudioTokenCounter: Construct the URL for the Google Gen AI Studio countTokens endpoint. """ base_url: Final = api_base or "https://generativelanguage.googleapis.com" - return f"{base_url}/v1beta/models/{model}:countTokens" + return join_api_base(base_url, f"v1beta/models/{model}:countTokens") def __init__(self, client: AsyncHTTPHandler | None = None) -> None: super().__init__() diff --git a/litellm/llms/gemini/files/transformation.py b/litellm/llms/gemini/files/transformation.py index 2c62e04c5a3..f9a7e230ea9 100644 --- a/litellm/llms/gemini/files/transformation.py +++ b/litellm/llms/gemini/files/transformation.py @@ -15,7 +15,7 @@ from typing_extensions import ReadOnly, Required from litellm._logging import verbose_logger from litellm.litellm_core_utils.prompt_templates.common_utils import extract_file_data -from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.llms.base_llm.files.transformation import ( BaseFilesConfig, LiteLLMLoggingObj, @@ -106,8 +106,7 @@ class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig): if not final_api_key: raise ValueError("api_key is required") - url: Final = f"{api_base}/{endpoint}" - return url + return join_api_base(api_base, endpoint) def get_supported_openai_params(self, model: str) -> list[OpenAICreateFileRequestOptionalParams]: return [] @@ -235,10 +234,11 @@ class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig): file_part: Final = self._normalize_gemini_file_id(file_id) - api_base = self.get_api_base(litellm_params.get("api_base")) or "https://generativelanguage.googleapis.com" - api_base = api_base.rstrip("/") + api_base: Final = ( + self.get_api_base(litellm_params.get("api_base")) or "https://generativelanguage.googleapis.com" + ) - url: Final = f"{api_base}/v1beta/{file_part}" + url: Final = join_api_base(api_base, f"v1beta/{file_part}") # API key is passed via x-goog-api-key header (set in validate_environment) return url, {} @@ -345,7 +345,7 @@ class GoogleAIStudioFilesHandler(GeminiModelInfo, BaseFilesConfig): file_name: Final = self._normalize_gemini_file_id(file_id) # Construct the delete URL - url: Final = f"{api_base}/v1beta/{file_name}" + url: Final = join_api_base(api_base, f"v1beta/{file_name}") # Add API key as header (Google AI Studio uses x-goog-api-key header) params: Final[dict] = {} diff --git a/litellm/llms/gemini/image_edit/transformation.py b/litellm/llms/gemini/image_edit/transformation.py index 31d3963c70c..450c38a2206 100644 --- a/litellm/llms/gemini/image_edit/transformation.py +++ b/litellm/llms/gemini/image_edit/transformation.py @@ -7,6 +7,7 @@ import httpx from httpx._types import RequestFiles from litellm.images.utils import ImageEditRequestUtils +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig from litellm.llms.gemini.common_utils import ( get_gemini_image_generation_config, @@ -79,9 +80,8 @@ class GeminiImageEditConfig(BaseImageEditConfig): api_base: str | None, litellm_params: dict, ) -> str: - base_url = api_base or get_secret_str("GEMINI_API_BASE") or self.DEFAULT_BASE_URL - base_url = base_url.rstrip("/") - return f"{base_url}/models/{model}:generateContent" + base_url: Final = api_base or get_secret_str("GEMINI_API_BASE") or self.DEFAULT_BASE_URL + return join_api_base(base_url, f"models/{model}:generateContent") def transform_image_edit_request( self, diff --git a/litellm/llms/gemini/image_generation/transformation.py b/litellm/llms/gemini/image_generation/transformation.py index fe82a9cae81..1aa3d074b1a 100644 --- a/litellm/llms/gemini/image_generation/transformation.py +++ b/litellm/llms/gemini/image_generation/transformation.py @@ -4,6 +4,7 @@ from typing import TYPE_CHECKING, Any, Final import httpx from pydantic import ConfigDict, TypeAdapter +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.image_generation.transformation import ( BaseImageGenerationConfig, ) @@ -88,14 +89,12 @@ class GoogleImageGenConfig(BaseImageGenerationConfig): """ complete_url: str = api_base or get_secret_str("GEMINI_API_BASE") or self.DEFAULT_BASE_URL - complete_url = complete_url.rstrip("/") - # Gemini Flash Image Preview models use generateContent endpoint if is_gemini_image_model(model): - complete_url = f"{complete_url}/models/{model}:generateContent" + complete_url = join_api_base(complete_url, f"models/{model}:generateContent") else: # All other Imagen models use predict endpoint - complete_url = f"{complete_url}/models/{model}:predict" + complete_url = join_api_base(complete_url, f"models/{model}:predict") return complete_url diff --git a/litellm/llms/gemini/interactions/transformation.py b/litellm/llms/gemini/interactions/transformation.py index a58037f691f..6f3e9a56b92 100644 --- a/litellm/llms/gemini/interactions/transformation.py +++ b/litellm/llms/gemini/interactions/transformation.py @@ -16,7 +16,7 @@ from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_logger from litellm.litellm_core_utils.core_helpers import process_response_headers -from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.llms.base_llm.interactions.transformation import BaseInteractionsAPIConfig from litellm.llms.gemini.common_utils import GeminiError, GeminiModelInfo from litellm.types.interactions import ( @@ -153,10 +153,8 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): if not api_key: raise ValueError("Google API key is required. Set GOOGLE_API_KEY or GEMINI_API_KEY environment variable.") - if stream: - return f"{api_base}/{self.api_version}/interactions?alt=sse" - - return f"{api_base}/{self.api_version}/interactions" + collection_url: Final = join_api_base(api_base, f"{self.api_version}/interactions") + return f"{collection_url}?alt=sse" if stream else collection_url def transform_request( self, @@ -304,7 +302,7 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): raise ValueError("Google API key is required") encoded_interaction_id: Final = encode_url_path_segment(interaction_id, field_name="interaction_id") return ( - f"{resolved_api_base}/{self.api_version}/interactions/{encoded_interaction_id}", + join_api_base(resolved_api_base, f"{self.api_version}/interactions/{encoded_interaction_id}"), {}, ) @@ -344,7 +342,7 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): raise ValueError("Google API key is required") encoded_interaction_id: Final = encode_url_path_segment(interaction_id, field_name="interaction_id") return ( - f"{resolved_api_base}/{self.api_version}/interactions/{encoded_interaction_id}", + join_api_base(resolved_api_base, f"{self.api_version}/interactions/{encoded_interaction_id}"), {}, ) @@ -375,7 +373,7 @@ class GoogleAIStudioInteractionsConfig(BaseInteractionsAPIConfig): raise ValueError("Google API key is required") encoded_interaction_id: Final = encode_url_path_segment(interaction_id, field_name="interaction_id") return ( - f"{resolved_api_base}/{self.api_version}/interactions/{encoded_interaction_id}/cancel", + join_api_base(resolved_api_base, f"{self.api_version}/interactions/{encoded_interaction_id}/cancel"), {}, ) diff --git a/litellm/llms/gemini/realtime/transformation.py b/litellm/llms/gemini/realtime/transformation.py index f4e2941670a..9af7edeac86 100644 --- a/litellm/llms/gemini/realtime/transformation.py +++ b/litellm/llms/gemini/realtime/transformation.py @@ -13,6 +13,7 @@ import litellm from litellm import verbose_logger from litellm._uuid import uuid from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( VertexGeminiConfig, @@ -212,7 +213,10 @@ class GeminiRealtimeConfig(BaseRealtimeConfig): # so the API key must remain as a query parameter here. This is an accepted # limitation; httpx is not used for WebSocket so MaskedHTTPStatusError # already covers the main leak vector. - return f"{api_base}/ws/google.ai.generativelanguage.v1beta.GenerativeService.BidiGenerateContent?key={api_key}" + live_url: Final = join_api_base( + api_base, "ws/google.ai.generativelanguage.v1beta.GenerativeService.BidiGenerateContent" + ) + return f"{live_url}?key={api_key}" def map_model_turn_event(self, model_turn: HttpxContentType) -> OpenAIRealtimeEventTypes: """ diff --git a/litellm/llms/gemini/vector_stores/transformation.py b/litellm/llms/gemini/vector_stores/transformation.py index 34bd9660c43..11e619eb4c7 100644 --- a/litellm/llms/gemini/vector_stores/transformation.py +++ b/litellm/llms/gemini/vector_stores/transformation.py @@ -11,6 +11,7 @@ from typing import TYPE_CHECKING, Any, Final import httpx from typing_extensions import ReadOnly, TypedDict +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig from litellm.llms.gemini.common_utils import ( GeminiError, @@ -141,15 +142,7 @@ class GeminiVectorStoreConfig(BaseVectorStoreConfig): Note: This returns the base URL WITHOUT the API key. The API key will be appended to specific endpoint URLs in the transform methods. """ - if api_base is None: - api_base = GeminiModelInfo.get_api_base() - - if api_base is None: - raise ValueError("GEMINI_API_BASE is not set") - - # Ensure we're using the v1beta version for File Search - api_version: Final = "v1beta" - return f"{api_base}/{api_version}" + return join_api_base(GeminiModelInfo.get_api_base(api_base), "v1beta") def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> GeminiError: """Return Gemini-specific error class.""" @@ -187,7 +180,7 @@ class GeminiVectorStoreConfig(BaseVectorStoreConfig): api_key: Final = litellm_params.get("api_key") or GeminiModelInfo.get_api_key() if not api_key: raise ValueError("GEMINI_API_KEY or GOOGLE_API_KEY is required") - url: Final = f"{api_base}/models/{model}:generateContent" + url: Final = join_api_base(api_base, f"models/{model}:generateContent") # Build file_search tool configuration (using snake_case as per Gemini docs) file_search_config: Final[dict[str, object]] = {"file_search_store_names": [vector_store_id]} @@ -328,7 +321,7 @@ class GeminiVectorStoreConfig(BaseVectorStoreConfig): """ Transform create request to Gemini's fileSearchStores format. """ - url: Final = f"{api_base}/fileSearchStores" + url: Final = join_api_base(api_base, "fileSearchStores") # API key is passed via x-goog-api-key header (set in validate_environment) diff --git a/litellm/llms/gemini/videos/transformation.py b/litellm/llms/gemini/videos/transformation.py index 53ca5ed4adc..0656095927c 100644 --- a/litellm/llms/gemini/videos/transformation.py +++ b/litellm/llms/gemini/videos/transformation.py @@ -8,6 +8,7 @@ from httpx._types import RequestFiles import litellm from litellm.constants import DEFAULT_GOOGLE_VIDEO_DURATION_SECONDS from litellm.images.utils import ImageEditRequestUtils +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.videos.transformation import BaseVideoConfig from litellm.llms.vertex_ai.videos.transformation import veo_video_count_from_parameters from litellm.secret_managers.main import get_secret_str @@ -241,7 +242,7 @@ class GeminiVideoConfig(BaseVideoConfig): return api_base.rstrip("/") model_name: Final = model.replace("gemini/", "") - url: Final = f"{api_base.rstrip('/')}/v1beta/models/{model_name}:predictLongRunning" + url: Final = join_api_base(api_base, f"v1beta/models/{model_name}:predictLongRunning") return url @@ -377,7 +378,7 @@ class GeminiVideoConfig(BaseVideoConfig): GET https://generativelanguage.googleapis.com/v1beta/{operation_name} """ operation_name: Final = extract_original_video_id(video_id) - url: Final = f"{api_base.rstrip('/')}/v1beta/{operation_name}" + url: Final = join_api_base(api_base, f"v1beta/{operation_name}") params: Final[dict[str, object]] = {} return url, params @@ -451,7 +452,7 @@ class GeminiVideoConfig(BaseVideoConfig): """ operation_name: Final = extract_original_video_id(video_id) - status_url: Final = f"{api_base.rstrip('/')}/v1beta/{operation_name}" + status_url: Final = join_api_base(api_base, f"v1beta/{operation_name}") client: Final = litellm.module_level_client status_response: Final = client.get(url=status_url, headers=headers) status_response.raise_for_status() diff --git a/litellm/llms/gigachat/chat/transformation.py b/litellm/llms/gigachat/chat/transformation.py index 1fe7de0ba85..936e6339865 100644 --- a/litellm/llms/gigachat/chat/transformation.py +++ b/litellm/llms/gigachat/chat/transformation.py @@ -17,6 +17,7 @@ import httpx from pydantic import TypeAdapter from litellm._logging import verbose_logger +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException from litellm.llms.gigachat.utils import convert_usage, get_api_base from litellm.secret_managers.main import get_secret_str @@ -108,8 +109,7 @@ class GigaChatConfig(BaseConfig): stream: bool | None = None, ) -> str: """Get complete API URL for chat completions.""" - base: Final = get_api_base(api_base) - return f"{base}/chat/completions" + return join_api_base(get_api_base(api_base), "chat/completions") def validate_environment( self, diff --git a/litellm/llms/gigachat/embedding/transformation.py b/litellm/llms/gigachat/embedding/transformation.py index e0e746427ad..30ca7147865 100644 --- a/litellm/llms/gigachat/embedding/transformation.py +++ b/litellm/llms/gigachat/embedding/transformation.py @@ -16,6 +16,7 @@ from pydantic import ConfigDict, TypeAdapter from litellm import LlmProviders from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig from litellm.llms.gigachat.utils import get_api_base @@ -98,8 +99,7 @@ class GigaChatEmbeddingConfig(BaseEmbeddingConfig): stream: bool | None = None, ) -> str: """Get the complete URL for embeddings endpoint.""" - base: Final = get_api_base(api_base) - return f"{base}/embeddings" + return join_api_base(get_api_base(api_base), "embeddings") def transform_embedding_request( self, diff --git a/litellm/llms/gigachat/file_handler.py b/litellm/llms/gigachat/file_handler.py index c051d693720..aff0465d1ec 100644 --- a/litellm/llms/gigachat/file_handler.py +++ b/litellm/llms/gigachat/file_handler.py @@ -13,6 +13,7 @@ from collections.abc import Mapping from typing import Final from litellm._logging import verbose_logger +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, get_httpx_client, @@ -125,8 +126,7 @@ def upload_file_sync( access_token: Final = get_access_token(credentials=credentials, litellm_params=litellm_params) # Upload to GigaChat - base_url: Final = get_api_base(api_base) - upload_url: Final = f"{base_url}/files" + upload_url: Final = join_api_base(get_api_base(api_base), "files") client: Final = get_httpx_client(params={"ssl_verify": False}) response: Final = client.post( @@ -191,8 +191,7 @@ async def upload_file_async( access_token: Final = await get_access_token_async(credentials=credentials, litellm_params=litellm_params) # Upload to GigaChat - base_url: Final = get_api_base(api_base) - upload_url: Final = f"{base_url}/files" + upload_url: Final = join_api_base(get_api_base(api_base), "files") client: Final = get_async_httpx_client( llm_provider=LlmProviders.GIGACHAT, diff --git a/litellm/llms/gigachat/passthrough/transformation.py b/litellm/llms/gigachat/passthrough/transformation.py index f16d263fd4d..872eee6dbf9 100644 --- a/litellm/llms/gigachat/passthrough/transformation.py +++ b/litellm/llms/gigachat/passthrough/transformation.py @@ -40,7 +40,7 @@ class GigaChatPassthroughConfig(BasePassthroughConfig): if base_target_url is None: raise Exception("GigaChat api base not found") - complete_url: Final = f"{base_target_url}/{endpoint.lstrip('/')}" + complete_url: Final = f"{base_target_url.rstrip('/')}/{endpoint.lstrip('/')}" return ( httpx.URL(complete_url), @@ -193,8 +193,7 @@ class GigaChatPassthroughConfig(BasePassthroughConfig): @staticmethod def get_api_base(api_base: str | None = None) -> str | None: - resolved_api_base: Final = api_base or get_secret_str("GIGACHAT_API_BASE") or GIGACHAT_BASE_URL - return resolved_api_base.rstrip("/") if resolved_api_base else None + return api_base or get_secret_str("GIGACHAT_API_BASE") or GIGACHAT_BASE_URL @staticmethod def get_api_key( diff --git a/litellm/llms/gigachat/utils.py b/litellm/llms/gigachat/utils.py index ab1890ad73b..55ca4260221 100644 --- a/litellm/llms/gigachat/utils.py +++ b/litellm/llms/gigachat/utils.py @@ -21,5 +21,5 @@ def convert_usage(usage_data: Mapping[str, int]) -> Usage: ) -def get_api_base(api_base: str | None = None) -> str | None: +def get_api_base(api_base: str | None = None) -> str: return (api_base or get_secret_str("GIGACHAT_API_BASE") or GIGACHAT_BASE_URL).rstrip("/") diff --git a/litellm/llms/github_copilot/embedding/transformation.py b/litellm/llms/github_copilot/embedding/transformation.py index 7f9209afd65..f16564ba505 100644 --- a/litellm/llms/github_copilot/embedding/transformation.py +++ b/litellm/llms/github_copilot/embedding/transformation.py @@ -20,6 +20,7 @@ from pydantic import ConfigDict, TypeAdapter from litellm._logging import verbose_logger from litellm.exceptions import AuthenticationError +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig from litellm.types.llms.openai import AllEmbeddingInputValues @@ -121,7 +122,7 @@ class GithubCopilotEmbeddingConfig(BaseEmbeddingConfig): user_session: Final = require_github_copilot_user_session( cast("dict[str, object]", litellm_params) # cast-ok: litellm_params arrives as an untyped request dict ) - effective_api_base = ( + effective_api_base: Final = ( (user_session.api_base if user_session is not None else None) or api_base or self.authenticator.get_api_base() @@ -129,11 +130,7 @@ class GithubCopilotEmbeddingConfig(BaseEmbeddingConfig): or DEFAULT_GITHUB_COPILOT_API_BASE ) - # Remove trailing slashes - effective_api_base = effective_api_base.rstrip("/") - - # Return the embeddings endpoint - return f"{effective_api_base}/embeddings" + return join_api_base(effective_api_base, "embeddings") def transform_embedding_request( self, diff --git a/litellm/llms/github_copilot/messages/transformation.py b/litellm/llms/github_copilot/messages/transformation.py index be138755f85..56c3693a02c 100644 --- a/litellm/llms/github_copilot/messages/transformation.py +++ b/litellm/llms/github_copilot/messages/transformation.py @@ -1,6 +1,7 @@ from typing import Any, Final, cast # noqa: TID251 # narrows untyped request dicts at the session boundary from litellm.exceptions import AuthenticationError +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.anthropic.pass_through.messages.transformation import ( AnthropicMessagesConfig, ) @@ -140,5 +141,5 @@ class GithubCopilotAnthropicMessagesConfig(AnthropicMessagesConfig): or api_base or self.authenticator.get_api_base() or DEFAULT_GITHUB_COPILOT_API_BASE - ).rstrip("/") - return resolved if resolved.endswith("/v1/messages") else f"{resolved}/v1/messages" + ) + return ensure_api_route(resolved, "v1/messages") diff --git a/litellm/llms/github_copilot/responses/transformation.py b/litellm/llms/github_copilot/responses/transformation.py index d468889b626..bf0b8c79068 100644 --- a/litellm/llms/github_copilot/responses/transformation.py +++ b/litellm/llms/github_copilot/responses/transformation.py @@ -20,6 +20,7 @@ import litellm from litellm._logging import verbose_logger from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH from litellm.exceptions import AuthenticationError +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig from litellm.types.llms.openai import ( ResponseInputParam, @@ -266,7 +267,7 @@ class GithubCopilotResponsesAPIConfig(OpenAIResponsesAPIConfig): user_session: Final = require_github_copilot_user_session( cast("dict[str, object]", litellm_params) # cast-ok: litellm_params arrives as an untyped request dict ) - effective_api_base = ( + effective_api_base: Final = ( (user_session.api_base if user_session is not None else None) or api_base or self.authenticator.get_api_base() @@ -274,11 +275,7 @@ class GithubCopilotResponsesAPIConfig(OpenAIResponsesAPIConfig): or DEFAULT_GITHUB_COPILOT_API_BASE ) - # Remove trailing slashes - effective_api_base = effective_api_base.rstrip("/") - - # Return the responses endpoint - return f"{effective_api_base}/responses" + return join_api_base(effective_api_base, "responses") def _handle_reasoning_item(self, item: dict[str, object]) -> dict[str, object]: """ diff --git a/litellm/llms/gradient_ai/chat/transformation.py b/litellm/llms/gradient_ai/chat/transformation.py index fcee55a27e9..0ba5ca45c2a 100644 --- a/litellm/llms/gradient_ai/chat/transformation.py +++ b/litellm/llms/gradient_ai/chat/transformation.py @@ -1,5 +1,6 @@ from typing import Final, Literal +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import ( AllMessageValues, @@ -105,10 +106,10 @@ class GradientAIConfig(OpenAILikeChatConfig): ) -> str: clean_api_base: Final = (api_base or "").rstrip("/") if clean_api_base and clean_api_base != GRADIENT_AI_SERVERLESS_ENDPOINT: - return f"{clean_api_base}/api/v1/chat/completions" + return join_api_base(clean_api_base, "api/v1/chat/completions") agent_endpoint: Final = (get_secret_str("GRADIENT_AI_AGENT_ENDPOINT") or "").rstrip("/") if agent_endpoint and agent_endpoint != GRADIENT_AI_SERVERLESS_ENDPOINT: - return f"{agent_endpoint}/api/v1/chat/completions" + return join_api_base(agent_endpoint, "api/v1/chat/completions") return f"{GRADIENT_AI_SERVERLESS_ENDPOINT}/v1/chat/completions" def _get_openai_compatible_provider_info( diff --git a/litellm/llms/heroku/chat/transformation.py b/litellm/llms/heroku/chat/transformation.py index acd630c85d0..002078dd04b 100644 --- a/litellm/llms/heroku/chat/transformation.py +++ b/litellm/llms/heroku/chat/transformation.py @@ -6,11 +6,12 @@ this is OpenAI compatible - no translation needed / occurs import os from collections.abc import Coroutine -from typing import Final, Literal, overload +from typing import Literal, overload from litellm.litellm_core_utils.prompt_templates.common_utils import ( handle_messages_with_content_list_to_str_conversion, ) +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig from litellm.types.llms.openai import AllMessageValues @@ -78,9 +79,4 @@ class HerokuChatConfig(OpenAIGPTConfig): "No api base was set. Please provide an api_base, or set the HEROKU_API_BASE environment variable." ) - clean_api_base: Final = resolved_base.rstrip("/") - - if not clean_api_base.endswith("/v1/chat/completions"): - return f"{clean_api_base}/v1/chat/completions" - - return clean_api_base + return ensure_api_route(resolved_base, "v1/chat/completions") diff --git a/litellm/llms/hosted_vllm/embedding/transformation.py b/litellm/llms/hosted_vllm/embedding/transformation.py index 3b54145a67b..d72a8d06f21 100644 --- a/litellm/llms/hosted_vllm/embedding/transformation.py +++ b/litellm/llms/hosted_vllm/embedding/transformation.py @@ -11,6 +11,7 @@ from typing import TYPE_CHECKING, Any, Final import httpx +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig from litellm.secret_managers.main import get_secret_str @@ -81,14 +82,7 @@ class HostedVLLMEmbeddingConfig(BaseEmbeddingConfig): if api_base is None: raise ValueError("api_base is required for hosted_vllm embeddings") - # Remove trailing slashes - api_base = api_base.rstrip("/") - - # Ensure the URL ends with /embeddings - if not api_base.endswith("/embeddings"): - api_base = f"{api_base}/embeddings" - - return api_base + return ensure_api_route(api_base, "embeddings") def transform_embedding_request( self, diff --git a/litellm/llms/hosted_vllm/image_edit/transformation.py b/litellm/llms/hosted_vllm/image_edit/transformation.py index 6804d3a0fb8..8e05b6dfb46 100644 --- a/litellm/llms/hosted_vllm/image_edit/transformation.py +++ b/litellm/llms/hosted_vllm/image_edit/transformation.py @@ -1,5 +1,6 @@ from typing import Final +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.openai.image_edit.transformation import OpenAIImageEditConfig from litellm.secret_managers.main import get_secret_str @@ -39,5 +40,5 @@ class HostedVLLMImageEditConfig(OpenAIImageEditConfig): ) trimmed: Final = resolved_api_base.rstrip("/") if trimmed.endswith("/v1"): - return f"{trimmed}/images/edits" - return f"{trimmed}/v1/images/edits" + return join_api_base(trimmed, "images/edits") + return join_api_base(trimmed, "v1/images/edits") diff --git a/litellm/llms/hosted_vllm/rerank/transformation.py b/litellm/llms/hosted_vllm/rerank/transformation.py index a463fc1d3c9..4a5729945b4 100644 --- a/litellm/llms/hosted_vllm/rerank/transformation.py +++ b/litellm/llms/hosted_vllm/rerank/transformation.py @@ -12,6 +12,7 @@ from pydantic import ValidationError from litellm._uuid import uuid from litellm.exceptions import UnsupportedParamsError from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig from litellm.secret_managers.main import get_secret_str @@ -62,7 +63,7 @@ class HostedVLLMRerankConfig(BaseRerankConfig): if api_base.endswith("/v1/rerank"): api_base = api_base.replace("/v1/rerank", "/rerank") elif not api_base.endswith("/rerank"): - api_base = f"{api_base}/rerank" + api_base = join_api_base(api_base, "rerank") return api_base raise ValueError("api_base must be provided for Hosted VLLM rerank") diff --git a/litellm/llms/hosted_vllm/responses/transformation.py b/litellm/llms/hosted_vllm/responses/transformation.py index 20d755a2273..8489bc0da84 100644 --- a/litellm/llms/hosted_vllm/responses/transformation.py +++ b/litellm/llms/hosted_vllm/responses/transformation.py @@ -8,6 +8,7 @@ the chat completions → responses conversion pipeline. from typing import Final +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig from litellm.secret_managers.main import get_secret_str from litellm.types.router import GenericLiteLLMParams @@ -64,9 +65,9 @@ class HostedVLLMResponsesAPIConfig(OpenAIResponsesAPIConfig): # If api_base already ends with /v1, append /responses # Otherwise append /v1/responses if api_base.endswith("/v1"): - return f"{api_base}/responses" + return join_api_base(api_base, "responses") - return f"{api_base}/v1/responses" + return join_api_base(api_base, "v1/responses") def supports_native_websocket(self) -> bool: """Hosted vLLM does not support native WebSocket for Responses API""" diff --git a/litellm/llms/hosted_vllm/transcriptions/transformation.py b/litellm/llms/hosted_vllm/transcriptions/transformation.py index 9a8c67fadc0..785ccd52a69 100644 --- a/litellm/llms/hosted_vllm/transcriptions/transformation.py +++ b/litellm/llms/hosted_vllm/transcriptions/transformation.py @@ -6,6 +6,7 @@ from typing import Final import httpx +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm.audio_transcription.transformation import ( AudioTranscriptionRequestData, ) @@ -40,11 +41,7 @@ class HostedVLLMAudioTranscriptionConfig(OpenAIWhisperAudioTranscriptionConfig): stream: bool | None = None, ) -> str: if api_base: - # Remove trailing slashes and ensure clean base URL - api_base = api_base.rstrip("/") - if not api_base.endswith("/v1/audio/transcriptions"): - api_base = f"{api_base}/v1/audio/transcriptions" - return api_base + return ensure_api_route(api_base, "v1/audio/transcriptions") raise ValueError("api_base must be provided for Hosted VLLM rerank") def transform_audio_transcription_request( diff --git a/litellm/llms/hosted_vllm/videos/transformation.py b/litellm/llms/hosted_vllm/videos/transformation.py index 22e4f4876fb..30ae6835263 100644 --- a/litellm/llms/hosted_vllm/videos/transformation.py +++ b/litellm/llms/hosted_vllm/videos/transformation.py @@ -10,7 +10,7 @@ from urllib.parse import urlparse from httpx._types import FileTypes, RequestFiles from litellm.images.utils import ImageEditRequestUtils -from litellm.litellm_core_utils.url_utils import SSRFError, validate_url +from litellm.litellm_core_utils.url_utils import SSRFError, join_api_base, validate_url from litellm.llms.openai.videos.transformation import OpenAIVideoConfig from litellm.secret_managers.main import get_secret_str from litellm.types.router import GenericLiteLLMParams @@ -177,8 +177,8 @@ class HostedVLLMVideoConfig(OpenAIVideoConfig): ) trimmed: Final = resolved_api_base.rstrip("/") if trimmed.endswith("/v1"): - return f"{trimmed}/videos" - return f"{trimmed}/v1/videos" + return join_api_base(trimmed, "videos") + return join_api_base(trimmed, "v1/videos") def transform_video_create_request( self, diff --git a/litellm/llms/huggingface/chat/transformation.py b/litellm/llms/huggingface/chat/transformation.py index b45c706f8bf..3fcb738eacb 100644 --- a/litellm/llms/huggingface/chat/transformation.py +++ b/litellm/llms/huggingface/chat/transformation.py @@ -4,6 +4,7 @@ from typing import TYPE_CHECKING, Any, Final import httpx +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.types.llms.openai import AllMessageValues, ChatCompletionRequest if TYPE_CHECKING: @@ -29,11 +30,11 @@ def _build_chat_completion_url(model_url: str) -> str: # Append /chat/completions if not already present if model_url.endswith("/v1"): - model_url += "/chat/completions" + model_url = join_api_base(model_url, "chat/completions") # Append /v1/chat/completions if not already present if not model_url.endswith("/chat/completions"): - model_url += "/v1/chat/completions" + model_url = join_api_base(model_url, "v1/chat/completions") return model_url diff --git a/litellm/llms/huggingface/rerank/transformation.py b/litellm/llms/huggingface/rerank/transformation.py index 334c60ee848..1217a2540d4 100644 --- a/litellm/llms/huggingface/rerank/transformation.py +++ b/litellm/llms/huggingface/rerank/transformation.py @@ -7,6 +7,7 @@ from typing_extensions import TypedDict import litellm from litellm._uuid import uuid +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig from litellm.secret_managers.main import get_secret_str @@ -69,15 +70,7 @@ class HuggingFaceRerankConfig(BaseRerankConfig): """ Get the complete URL for the API call, including the /rerank suffix if necessary. """ - # Get base URL from api_base or default - base_url = self.get_api_base(model=model, api_base=api_base) - - # Remove trailing slashes and ensure we have the /rerank endpoint - base_url = base_url.rstrip("/") - if not base_url.endswith("/rerank"): - base_url = f"{base_url}/rerank" - - return base_url + return ensure_api_route(self.get_api_base(model=model, api_base=api_base), "rerank") def get_supported_cohere_rerank_params(self, model: str) -> list: return [ diff --git a/litellm/llms/infinity/embedding/transformation.py b/litellm/llms/infinity/embedding/transformation.py index f289918cb3c..a467c56fc74 100644 --- a/litellm/llms/infinity/embedding/transformation.py +++ b/litellm/llms/infinity/embedding/transformation.py @@ -5,6 +5,7 @@ import httpx from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig from litellm.secret_managers.main import get_secret_str @@ -36,11 +37,7 @@ class InfinityEmbeddingConfig(BaseEmbeddingConfig): ) -> str: if api_base is None: raise ValueError("api_base is required for Infinity embeddings") - # Remove trailing slashes and ensure clean base URL - api_base = api_base.rstrip("/") - if not api_base.endswith("/embeddings"): - api_base = f"{api_base}/embeddings" - return api_base + return ensure_api_route(api_base, "embeddings") def validate_environment( self, diff --git a/litellm/llms/infinity/rerank/transformation.py b/litellm/llms/infinity/rerank/transformation.py index e089b3fecbe..07f3bb2c26b 100644 --- a/litellm/llms/infinity/rerank/transformation.py +++ b/litellm/llms/infinity/rerank/transformation.py @@ -13,6 +13,7 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm from litellm._uuid import uuid from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.cohere.rerank.transformation import CohereRerankConfig from litellm.secret_managers.main import get_secret_str from litellm.types.rerank import ( @@ -61,11 +62,7 @@ class InfinityRerankConfig(CohereRerankConfig): ) -> str: if api_base is None: raise ValueError("api_base is required for Infinity rerank") - # Remove trailing slashes and ensure clean base URL - api_base = api_base.rstrip("/") - if not api_base.endswith("/rerank"): - api_base = f"{api_base}/rerank" - return api_base + return ensure_api_route(api_base, "rerank") def validate_environment( self, diff --git a/litellm/llms/jina_ai/embedding/transformation.py b/litellm/llms/jina_ai/embedding/transformation.py index 9b156e5ae21..04ff8f941a3 100644 --- a/litellm/llms/jina_ai/embedding/transformation.py +++ b/litellm/llms/jina_ai/embedding/transformation.py @@ -15,6 +15,7 @@ from pydantic import ConfigDict, TypeAdapter from litellm import LlmProviders from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm import BaseEmbeddingConfig from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.secret_managers.main import get_secret_str @@ -115,12 +116,7 @@ class JinaAIEmbeddingConfig(BaseEmbeddingConfig): litellm_params: dict, stream: bool | None = None, ) -> str: - if api_base: - clean_api_base: Final = api_base.rstrip("/") - if not clean_api_base.endswith("/embeddings"): - return f"{clean_api_base}/embeddings" - return clean_api_base - return "https://api.jina.ai/v1/embeddings" + return ensure_api_route(api_base or "https://api.jina.ai/v1", "embeddings") def transform_embedding_request( self, diff --git a/litellm/llms/langflow/chat/transformation.py b/litellm/llms/langflow/chat/transformation.py index 1781082d36f..0938cc0d07b 100644 --- a/litellm/llms/langflow/chat/transformation.py +++ b/litellm/llms/langflow/chat/transformation.py @@ -9,6 +9,7 @@ from litellm._logging import verbose_logger from litellm.litellm_core_utils.prompt_templates.common_utils import ( convert_content_list_to_str, ) +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import Choices, Message, ModelResponse, Usage @@ -107,9 +108,8 @@ class LangFlowConfig(BaseConfig): "api_base is required for LangFlow. Set it via LANGFLOW_API_BASE env var or api_base parameter." ) - api_base = api_base.rstrip("/") flow_id: Final = quote(self._get_flow_id(model, optional_params), safe="") - return f"{api_base}/api/v1/run/{flow_id}" + return join_api_base(api_base, f"api/v1/run/{flow_id}") def _get_last_user_message(self, messages: list[AllMessageValues]) -> str: """Extract the text of the last user message to use as input_value.""" diff --git a/litellm/llms/langgraph/chat/transformation.py b/litellm/llms/langgraph/chat/transformation.py index 4f93f435957..ce624eecc0e 100644 --- a/litellm/llms/langgraph/chat/transformation.py +++ b/litellm/llms/langgraph/chat/transformation.py @@ -18,6 +18,7 @@ from litellm._logging import verbose_logger from litellm.litellm_core_utils.prompt_templates.common_utils import ( convert_content_list_to_str, ) +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException from litellm.llms.langgraph.chat.sse_iterator import LangGraphSSEStreamIterator from litellm.types.llms.openai import AllMessageValues @@ -116,14 +117,11 @@ class LangGraphConfig(BaseConfig): "api_base is required for LangGraph. Set it via LANGGRAPH_API_BASE env var or api_base parameter." ) - # Remove trailing slash if present - api_base = api_base.rstrip("/") - # Choose endpoint based on streaming mode if stream: - return f"{api_base}/runs/stream" + return join_api_base(api_base, "runs/stream") else: - return f"{api_base}/runs/wait" + return join_api_base(api_base, "runs/wait") def _get_assistant_id(self, model: str, optional_params: dict) -> str: """ diff --git a/litellm/llms/lemonade/chat/transformation.py b/litellm/llms/lemonade/chat/transformation.py index a740050dd94..01fdc0b8bf8 100644 --- a/litellm/llms/lemonade/chat/transformation.py +++ b/litellm/llms/lemonade/chat/transformation.py @@ -12,6 +12,7 @@ from pydantic import ConfigDict, TypeAdapter import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import ( AllMessageValues, @@ -105,7 +106,7 @@ class LemonadeChatConfig(OpenAILikeChatConfig): # Getting the list of models from lemonade try: response: Final = litellm.module_level_client.get( - url=f"{api_base}/models", + url=join_api_base(api_base, "models"), headers=self._get_auth_headers(api_key), ) except Exception as e: @@ -193,7 +194,7 @@ class LemonadeChatConfig(OpenAILikeChatConfig): try: response: Final = litellm.module_level_client.get( - url=f"{api_base}/models/{encoded_model}", + url=join_api_base(str(api_base), f"models/{encoded_model}"), headers=self._get_auth_headers(api_key), ) response.raise_for_status() diff --git a/litellm/llms/linkup/search/transformation.py b/litellm/llms/linkup/search/transformation.py index e219a227b1f..fd1ca792302 100644 --- a/litellm/llms/linkup/search/transformation.py +++ b/litellm/llms/linkup/search/transformation.py @@ -11,6 +11,7 @@ import httpx from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm.search.transformation import ( BaseSearchConfig, SearchResponse, @@ -87,11 +88,7 @@ class LinkupSearchConfig(BaseSearchConfig): """ Get complete URL for Search endpoint. """ - resolved_api_base: Final = api_base or get_secret_str("LINKUP_API_BASE") or self.LINKUP_API_BASE - clean_api_base: Final = resolved_api_base.rstrip("/") - if clean_api_base.endswith("/search"): - return clean_api_base - return f"{clean_api_base}/search" + return ensure_api_route(api_base or get_secret_str("LINKUP_API_BASE") or self.LINKUP_API_BASE, "search") def transform_search_request( self, diff --git a/litellm/llms/litellm_proxy/image_edit/transformation.py b/litellm/llms/litellm_proxy/image_edit/transformation.py index 9e5a2a2c343..f319e31ff3f 100644 --- a/litellm/llms/litellm_proxy/image_edit/transformation.py +++ b/litellm/llms/litellm_proxy/image_edit/transformation.py @@ -1,3 +1,4 @@ +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.openai.image_edit.transformation import OpenAIImageEditConfig from litellm.secret_managers.main import get_secret_str @@ -21,5 +22,4 @@ class LiteLLMProxyImageEditConfig(OpenAIImageEditConfig): api_base = api_base or get_secret_str("LITELLM_PROXY_API_BASE") if api_base is None: raise ValueError("api_base not set for LiteLLM Proxy route. Set in env via `LITELLM_PROXY_API_BASE`") - api_base = api_base.rstrip("/") - return f"{api_base}/images/edits" + return join_api_base(api_base, "images/edits") diff --git a/litellm/llms/litellm_proxy/image_generation/transformation.py b/litellm/llms/litellm_proxy/image_generation/transformation.py index fcddeabfc85..05e536eece5 100644 --- a/litellm/llms/litellm_proxy/image_generation/transformation.py +++ b/litellm/llms/litellm_proxy/image_generation/transformation.py @@ -1,3 +1,4 @@ +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.openai.image_generation.gpt_transformation import ( GPTImageGenerationConfig, ) @@ -33,5 +34,4 @@ class LiteLLMProxyImageGenerationConfig(GPTImageGenerationConfig): api_base = api_base or get_secret_str("LITELLM_PROXY_API_BASE") if api_base is None: raise ValueError("api_base not set for LiteLLM Proxy route. Set in env via `LITELLM_PROXY_API_BASE`") - api_base = api_base.rstrip("/") - return f"{api_base}/images/generations" + return join_api_base(api_base, "images/generations") diff --git a/litellm/llms/litellm_proxy/responses/transformation.py b/litellm/llms/litellm_proxy/responses/transformation.py index 2928b07665a..3092b49b743 100644 --- a/litellm/llms/litellm_proxy/responses/transformation.py +++ b/litellm/llms/litellm_proxy/responses/transformation.py @@ -5,6 +5,7 @@ LiteLLM Proxy supports the OpenAI Responses API natively when the underlying mod This config enables pass-through behavior to the proxy's /v1/responses endpoint. """ +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig from litellm.secret_managers.main import get_secret_str from litellm.types.utils import LlmProviders @@ -40,10 +41,7 @@ class LiteLLMProxyResponsesAPIConfig(OpenAIResponsesAPIConfig): "Set via api_base parameter or LITELLM_PROXY_API_BASE environment variable" ) - # Remove trailing slashes - api_base = api_base.rstrip("/") - - return f"{api_base}/responses" + return join_api_base(api_base, "responses") def supports_native_websocket(self) -> bool: """LiteLLM Proxy does not support native WebSocket for Responses API""" diff --git a/litellm/llms/manus/files/transformation.py b/litellm/llms/manus/files/transformation.py index 94b0e62199d..991aa4cbc32 100644 --- a/litellm/llms/manus/files/transformation.py +++ b/litellm/llms/manus/files/transformation.py @@ -21,7 +21,7 @@ from pydantic import ConfigDict, TypeAdapter import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.prompt_templates.common_utils import extract_file_data -from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.files.transformation import ( BaseFilesConfig, @@ -140,8 +140,8 @@ class ManusFilesConfig(BaseFilesConfig): # Manus API uses /v1/files endpoint if api_base.endswith("/v1"): - return f"{api_base}/files" - return f"{api_base}/v1/files" + return join_api_base(api_base, "files") + return join_api_base(api_base, "v1/files") def get_error_class( self, @@ -423,7 +423,7 @@ class ManusFilesConfig(BaseFilesConfig): litellm_params=litellm_params, ) encoded_file_id: Final = encode_url_path_segment(file_id, field_name="file_id") - return f"{api_base}/{encoded_file_id}/content", {} + return join_api_base(api_base, f"{encoded_file_id}/content"), {} def transform_file_content_response( self, diff --git a/litellm/llms/manus/responses/transformation.py b/litellm/llms/manus/responses/transformation.py index 9785a50bbce..ed06391a260 100644 --- a/litellm/llms/manus/responses/transformation.py +++ b/litellm/llms/manus/responses/transformation.py @@ -10,7 +10,7 @@ from litellm.litellm_core_utils.core_helpers import process_response_headers from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( safe_convert_created_field, ) -from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.llms.openai.common_utils import OpenAIError from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig from litellm.secret_managers.main import get_secret_str @@ -119,8 +119,8 @@ class ManusResponsesAPIConfig(OpenAIResponsesAPIConfig): # Manus API uses /v1/responses endpoint (OpenAI-compatible) if api_base.endswith("/v1"): - return f"{api_base}/responses" - return f"{api_base}/v1/responses" + return join_api_base(api_base, "responses") + return join_api_base(api_base, "v1/responses") def transform_responses_api_request( self, diff --git a/litellm/llms/milvus/vector_stores/transformation.py b/litellm/llms/milvus/vector_stores/transformation.py index 70d3649debd..67de3a06af8 100644 --- a/litellm/llms/milvus/vector_stores/transformation.py +++ b/litellm/llms/milvus/vector_stores/transformation.py @@ -5,6 +5,7 @@ from typing import TYPE_CHECKING, Any, Final import httpx +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.vector_store.transformation import ( BaseQueryEmbeddingVectorStoreConfig, VectorStoreEmbeddingExecutor, @@ -185,7 +186,7 @@ class MilvusVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig): } litellm_logging_obj.model_call_details["input"] = query_text litellm_logging_obj.model_call_details["embedding_model"] = litellm_params.get("litellm_embedding_model") - return f"{api_base}/v2/vectordb/entities/search", { + return join_api_base(api_base, "v2/vectordb/entities/search"), { "collectionName": vector_store_id, "data": [query_vector], "annsField": "book_intro_vector", diff --git a/litellm/llms/minimax/chat/transformation.py b/litellm/llms/minimax/chat/transformation.py index c5aa8811f02..f397bf7fbb9 100644 --- a/litellm/llms/minimax/chat/transformation.py +++ b/litellm/llms/minimax/chat/transformation.py @@ -5,6 +5,7 @@ MiniMax OpenAI transformation config - extends OpenAI chat config for MiniMax's from typing import Final import litellm +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam @@ -55,17 +56,14 @@ class MinimaxChatConfig(OpenAIGPTConfig): Override to ensure we use MiniMax's endpoint. """ # Get the base URL (either provided or default MiniMax endpoint) - base_url: Final = self.get_api_base(api_base=api_base) + base_url: Final = self.get_api_base(api_base=api_base).rstrip("/") # Ensure it ends with /chat/completions if base_url.endswith("/chat/completions"): return base_url - elif base_url.endswith("/v1"): - return f"{base_url}/chat/completions" - elif base_url.endswith("/"): - return f"{base_url}v1/chat/completions" - else: - return f"{base_url}/v1/chat/completions" + if base_url.endswith("/v1"): + return join_api_base(base_url, "chat/completions") + return join_api_base(base_url, "v1/chat/completions") def remove_cache_control_flag_from_messages_and_tools( self, diff --git a/litellm/llms/minimax/messages/transformation.py b/litellm/llms/minimax/messages/transformation.py index d62b88a24c6..f9f3502237c 100644 --- a/litellm/llms/minimax/messages/transformation.py +++ b/litellm/llms/minimax/messages/transformation.py @@ -2,9 +2,10 @@ MiniMax Anthropic transformation config - extends AnthropicConfig for MiniMax's Anthropic-compatible API """ -from typing import Any, Final # noqa: TID251 # override below must mirror the legacy base signature +from typing import Any # noqa: TID251 # override below must mirror the legacy base signature import litellm +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.anthropic.pass_through.messages.transformation import ( AnthropicMessagesConfig, ) @@ -82,15 +83,4 @@ class MinimaxMessagesConfig(AnthropicMessagesConfig): Get the complete URL for MiniMax API. Override to ensure we use MiniMax's endpoint, not Anthropic's. """ - # Get the base URL (either provided or default MiniMax endpoint) - base_url: Final = self.get_api_base(api_base=api_base) - - # If the base URL already includes the full path, return it - if base_url.endswith("/v1/messages"): - return base_url - - # Otherwise append the messages endpoint - if base_url.endswith("/"): - return f"{base_url}v1/messages" - else: - return f"{base_url}/v1/messages" + return ensure_api_route(self.get_api_base(api_base=api_base), "v1/messages") diff --git a/litellm/llms/minimax/text_to_speech/transformation.py b/litellm/llms/minimax/text_to_speech/transformation.py index 1cabc43eb1c..6f2cb36ed03 100644 --- a/litellm/llms/minimax/text_to_speech/transformation.py +++ b/litellm/llms/minimax/text_to_speech/transformation.py @@ -13,6 +13,7 @@ from httpx import Headers from pydantic import ConfigDict, TypeAdapter import litellm +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.text_to_speech.transformation import ( BaseTextToSpeechConfig, @@ -404,10 +405,5 @@ class MinimaxTextToSpeechConfig(BaseTextToSpeechConfig): """ Construct the MiniMax endpoint URL. """ - base_url = api_base or get_secret_str("MINIMAX_API_BASE") or self.TTS_BASE_URL - base_url = base_url.rstrip("/") - - # MiniMax uses a simple endpoint path - url: Final = f"{base_url}{self.TTS_ENDPOINT_PATH}" - - return url + base_url: Final = api_base or get_secret_str("MINIMAX_API_BASE") or self.TTS_BASE_URL + return join_api_base(base_url, self.TTS_ENDPOINT_PATH) diff --git a/litellm/llms/mistral/audio_speech/transformation.py b/litellm/llms/mistral/audio_speech/transformation.py index 5a8960d0898..ac0e1a33620 100644 --- a/litellm/llms/mistral/audio_speech/transformation.py +++ b/litellm/llms/mistral/audio_speech/transformation.py @@ -13,6 +13,7 @@ from typing import TYPE_CHECKING, Final import httpx from pydantic import ConfigDict, TypeAdapter, ValidationError +from litellm.litellm_core_utils.url_utils import ensure_api_route, join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.text_to_speech.transformation import ( BaseTextToSpeechConfig, @@ -118,9 +119,7 @@ class MistralTextToSpeechConfig(BaseTextToSpeechConfig): api_base: str | None, litellm_params: Mapping[str, object], ) -> str: - configured_base: Final = (api_base or self.TTS_BASE_URL).rstrip("/") - versioned_base: Final = configured_base if configured_base.endswith("/v1") else f"{configured_base}/v1" - return f"{versioned_base}/audio/speech" + return join_api_base(ensure_api_route(api_base or self.TTS_BASE_URL, "v1"), "audio/speech") def transform_text_to_speech_request( self, diff --git a/litellm/llms/mistral/audio_transcription/transformation.py b/litellm/llms/mistral/audio_transcription/transformation.py index 4fce7a95350..3eb9e122265 100644 --- a/litellm/llms/mistral/audio_transcription/transformation.py +++ b/litellm/llms/mistral/audio_transcription/transformation.py @@ -9,6 +9,7 @@ from typing import Final import httpx from litellm.litellm_core_utils.audio_utils.utils import process_audio_file +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.audio_transcription.transformation import ( AudioTranscriptionRequestData, BaseAudioTranscriptionConfig, @@ -57,8 +58,7 @@ class MistralAudioTranscriptionConfig(BaseAudioTranscriptionConfig): litellm_params: dict, stream: bool | None = None, ) -> str: - api_base = "https://api.mistral.ai/v1" if api_base is None else api_base.rstrip("/") - return f"{api_base}/audio/transcriptions" + return join_api_base("https://api.mistral.ai/v1" if api_base is None else api_base, "audio/transcriptions") def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException: return MistralAudioTranscriptionException( diff --git a/litellm/llms/mistral/batches/transformation.py b/litellm/llms/mistral/batches/transformation.py index 86de2836048..18f5df9b6e9 100644 --- a/litellm/llms/mistral/batches/transformation.py +++ b/litellm/llms/mistral/batches/transformation.py @@ -18,7 +18,7 @@ from openai.types.batch_error import BatchError from pydantic import ConfigDict from typing_extensions import NotRequired, ReadOnly, TypedDict -from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.llms.base_llm.batches.transformation import BaseBatchesConfig from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.types.llms.base import LiteLLMBaseModel @@ -156,7 +156,7 @@ class MistralBatchesConfig(BaseBatchesConfig): litellm_params: Mapping[str, object], data: CreateBatchRequest, ) -> str: - return f"{get_mistral_api_base(api_base)}/v1/batch/jobs" + return join_api_base(get_mistral_api_base(api_base), "v1/batch/jobs") def transform_create_batch_request( self, @@ -199,7 +199,10 @@ class MistralBatchesConfig(BaseBatchesConfig): api_key: Final = litellm_params.get("api_key") request: Final = MistralPresignedRequest( method="GET", - url=f"{get_mistral_api_base(api_base if isinstance(api_base, str) else None)}/v1/batch/jobs/{encoded_batch_id}", + url=join_api_base( + get_mistral_api_base(api_base if isinstance(api_base, str) else None), + f"v1/batch/jobs/{encoded_batch_id}", + ), headers=get_mistral_auth_headers(_NO_HEADERS, api_key if isinstance(api_key, str) else None), ) return dict(request) diff --git a/litellm/llms/mistral/chat/transformation.py b/litellm/llms/mistral/chat/transformation.py index b8a6cb1beb9..2c0acacff23 100644 --- a/litellm/llms/mistral/chat/transformation.py +++ b/litellm/llms/mistral/chat/transformation.py @@ -19,6 +19,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( strip_litellm_internal_message_fields, strip_none_values_from_message, ) +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.openai.chat.gpt_transformation import ( OpenAIChatCompletionStreamingHandler, OpenAIGPTConfig, @@ -220,8 +221,7 @@ class MistralConfig(OpenAIGPTConfig): ) # if api_base does not end with /v1 we add it - if api_base is not None and not api_base.endswith("/v1"): # Mistral always needs a /v1 at the end - api_base = api_base + "/v1" + api_base = ensure_api_route(api_base, "v1") dynamic_api_key: Final = ( api_key or get_secret_str("MISTRAL_AZURE_API_KEY") # for Azure AI Mistral diff --git a/litellm/llms/mistral/files/transformation.py b/litellm/llms/mistral/files/transformation.py index 86bce763917..908567e88fb 100644 --- a/litellm/llms/mistral/files/transformation.py +++ b/litellm/llms/mistral/files/transformation.py @@ -18,7 +18,7 @@ from pydantic import ConfigDict from typing_extensions import ReadOnly, TypedDict from litellm.litellm_core_utils.prompt_templates.common_utils import extract_file_data -from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.files.transformation import BaseFilesConfig, LiteLLMLoggingObj from litellm.types.llms.base import LiteLLMBaseModel @@ -130,11 +130,11 @@ class MistralFilesConfig(BaseFilesConfig): litellm_params: Mapping[str, object], stream: bool | None = None, ) -> str: - return f"{get_mistral_api_base(api_base)}/v1/files" + return join_api_base(get_mistral_api_base(api_base), "v1/files") def _file_url(self, file_id: str, litellm_params: Mapping[str, object], suffix: str = "") -> str: encoded_file_id: Final = encode_url_path_segment(file_id, field_name="file_id") - return f"{_api_base_from(litellm_params)}/v1/files/{encoded_file_id}{suffix}" + return join_api_base(_api_base_from(litellm_params), f"v1/files/{encoded_file_id}{suffix}") def get_error_class( self, error_message: str, status_code: int, headers: Mapping[str, str] | httpx.Headers @@ -233,7 +233,7 @@ class MistralFilesConfig(BaseFilesConfig): optional_params: Mapping[str, object], litellm_params: Mapping[str, object], ) -> tuple[str, dict[str, str]]: # mutable-ok: BaseFilesConfig signature - url: Final = f"{_api_base_from(litellm_params)}/v1/files" + url: Final = join_api_base(_api_base_from(litellm_params), "v1/files") if not purpose: return url, _NO_QUERY_PARAMS return url, {"purpose": _to_mistral_purpose(purpose)} diff --git a/litellm/llms/modelscope/chat/transformation.py b/litellm/llms/modelscope/chat/transformation.py index c54437b9a21..ec48b344bd9 100644 --- a/litellm/llms/modelscope/chat/transformation.py +++ b/litellm/llms/modelscope/chat/transformation.py @@ -7,6 +7,7 @@ from typing import Final, Literal, cast, overload from typing_extensions import override +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues @@ -86,9 +87,4 @@ class ModelScopeChatConfig(OpenAIGPTConfig): """ If api_base is not provided, use the default ModelScope /chat/completions endpoint. """ - clean_api_base: Final = (api_base or self.DEFAULT_BASE_URL).rstrip("/") - - if not clean_api_base.endswith("/chat/completions"): - return f"{clean_api_base}/chat/completions" - - return clean_api_base + return ensure_api_route(api_base or self.DEFAULT_BASE_URL, "chat/completions") diff --git a/litellm/llms/modelscope/image_generation/transformation.py b/litellm/llms/modelscope/image_generation/transformation.py index 08cdbacdad0..36c8705f256 100644 --- a/litellm/llms/modelscope/image_generation/transformation.py +++ b/litellm/llms/modelscope/image_generation/transformation.py @@ -13,6 +13,7 @@ import httpx from pydantic import ConfigDict, TypeAdapter from typing_extensions import override +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.image_generation.transformation import ( BaseImageGenerationConfig, @@ -90,11 +91,10 @@ class ModelScopeImageGenerationConfig(BaseImageGenerationConfig): """ Get the complete URL for the ModelScope image generation API request. """ - base_url: str = api_base or get_secret_str("MODELSCOPE_API_BASE") or self.DEFAULT_BASE_URL - base_url = base_url.rstrip("/") + base_url: Final = api_base or get_secret_str("MODELSCOPE_API_BASE") or self.DEFAULT_BASE_URL # Return the images endpoint - return f"{base_url}/images/generations" + return join_api_base(base_url, "images/generations") @override def validate_environment( diff --git a/litellm/llms/mongodb/vector_stores/transformation.py b/litellm/llms/mongodb/vector_stores/transformation.py index 0d931ee78cb..7bf9d8b8441 100644 --- a/litellm/llms/mongodb/vector_stores/transformation.py +++ b/litellm/llms/mongodb/vector_stores/transformation.py @@ -9,6 +9,7 @@ import httpx from pydantic import ConfigDict, TypeAdapter, ValidationError from litellm.exceptions import AuthenticationError, BadRequestError, ServiceUnavailableError, Timeout +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.vector_store.transformation import ( BaseQueryEmbeddingVectorStoreConfig, @@ -283,7 +284,7 @@ class MongoDBVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig): raise config_error("The embedding model must return a non-empty, finite query vector.") limit: Final = cls._limit(optional_params) return ( - f"{api_base}/v1/vector_stores/{quote(vector_store_id, safe='')}/search", + join_api_base(api_base, f"v1/vector_stores/{quote(vector_store_id, safe='')}/search"), { "query": query_text, "query_vector": tuple(vector), diff --git a/litellm/llms/moonshot/chat/transformation.py b/litellm/llms/moonshot/chat/transformation.py index 93496e775db..698db8b7c25 100644 --- a/litellm/llms/moonshot/chat/transformation.py +++ b/litellm/llms/moonshot/chat/transformation.py @@ -9,6 +9,7 @@ import litellm from litellm.litellm_core_utils.prompt_templates.common_utils import ( handle_messages_with_content_list_to_str_conversion, ) +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues from litellm.utils import supports_reasoning @@ -93,12 +94,7 @@ class MoonshotChatConfig(OpenAIGPTConfig): """ If api_base is not provided, use the default Moonshot AI /chat/completions endpoint. """ - clean_api_base: Final = (api_base or "https://api.moonshot.ai/v1").rstrip("/") - - if not clean_api_base.endswith("/chat/completions"): - return f"{clean_api_base}/chat/completions" - - return clean_api_base + return ensure_api_route(api_base or "https://api.moonshot.ai/v1", "chat/completions") def get_supported_openai_params(self, model: str) -> list: """ diff --git a/litellm/llms/nimble/search/transformation.py b/litellm/llms/nimble/search/transformation.py index 30cfc38eb87..b4d9be909c3 100644 --- a/litellm/llms/nimble/search/transformation.py +++ b/litellm/llms/nimble/search/transformation.py @@ -13,6 +13,7 @@ from typing import TYPE_CHECKING, Final import httpx from pydantic import ConfigDict, TypeAdapter, ValidationError +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.search.transformation import ( BaseSearchConfig, @@ -124,10 +125,7 @@ class NimbleSearchConfig(BaseSearchConfig): data: dict[str, object] | list[dict[str, object]] | None = None, # mutable-ok: base signature **kwargs: object, # kwargs-ok: BaseSearchConfig.get_complete_url signature ) -> str: - resolved_base: Final = (api_base or get_secret_str("NIMBLE_API_BASE") or self.NIMBLE_API_BASE).rstrip("/") - if resolved_base.endswith("/search"): - return resolved_base - return f"{resolved_base}/search" + return ensure_api_route(api_base or get_secret_str("NIMBLE_API_BASE") or self.NIMBLE_API_BASE, "search") def transform_search_request( self, diff --git a/litellm/llms/nvidia_nim/rerank/ranking_transformation.py b/litellm/llms/nvidia_nim/rerank/ranking_transformation.py index 177d7883feb..39391be0fdb 100644 --- a/litellm/llms/nvidia_nim/rerank/ranking_transformation.py +++ b/litellm/llms/nvidia_nim/rerank/ranking_transformation.py @@ -11,6 +11,7 @@ from typing import Final import httpx from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.nvidia_nim.rerank.transformation import NvidiaNimRerankConfig from litellm.types.rerank import RerankResponse @@ -76,7 +77,7 @@ class NvidiaNimRankingConfig(NvidiaNimRerankConfig): api_base = api_base.removesuffix("/v1") - return f"{api_base}/v1/ranking" + return join_api_base(api_base, "v1/ranking") def map_cohere_rerank_params( self, diff --git a/litellm/llms/nvidia_nim/rerank/transformation.py b/litellm/llms/nvidia_nim/rerank/transformation.py index c8e1a331c8c..12377d8f854 100644 --- a/litellm/llms/nvidia_nim/rerank/transformation.py +++ b/litellm/llms/nvidia_nim/rerank/transformation.py @@ -8,6 +8,7 @@ from typing_extensions import Required, TypedDict import litellm from litellm._uuid import uuid from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig from litellm.secret_managers.main import get_secret_str @@ -106,7 +107,7 @@ class NvidiaNimRerankConfig(BaseRerankConfig): # Strip nvidia_nim/ prefix from model name if present clean_model: Final = self._get_clean_model_name(model) - return f"{api_base}/v1/retrieval/{clean_model}/reranking" + return join_api_base(api_base, f"v1/retrieval/{clean_model}/reranking") def get_supported_cohere_rerank_params(self, model: str) -> list: """ diff --git a/litellm/llms/oci/chat/transformation.py b/litellm/llms/oci/chat/transformation.py index e62e5e63642..8e1502b2610 100644 --- a/litellm/llms/oci/chat/transformation.py +++ b/litellm/llms/oci/chat/transformation.py @@ -21,6 +21,7 @@ from litellm._logging import verbose_logger from litellm.constants import DEFAULT_OCI_CHAT_MAX_TOKENS from litellm.litellm_core_utils.logging_utils import track_llm_api_timing from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, @@ -538,7 +539,7 @@ class OCIChatConfig(BaseConfig): stream: bool | None = None, ) -> str: base: Final = get_oci_base_url(optional_params, api_base or litellm.api_base) - return f"{base}/{OCI_API_VERSION}/actions/chat" + return join_api_base(base, f"{OCI_API_VERSION}/actions/chat") def _get_optional_params(self, vendor: OCIVendors, optional_params: dict, model: str = "") -> dict: param_map: Final = ( diff --git a/litellm/llms/oci/embed/transformation.py b/litellm/llms/oci/embed/transformation.py index edadc293771..ffcb156bb16 100644 --- a/litellm/llms/oci/embed/transformation.py +++ b/litellm/llms/oci/embed/transformation.py @@ -29,6 +29,7 @@ import httpx from pydantic import ConfigDict, TypeAdapter import litellm +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig from litellm.llms.oci.common_utils import ( @@ -156,7 +157,7 @@ class OCIEmbedConfig(BaseEmbeddingConfig): stream: bool | None = None, ) -> str: base: Final = get_oci_base_url(optional_params, api_base or litellm.api_base) - return f"{base}/{OCI_API_VERSION}/actions/embedText" + return join_api_base(base, f"{OCI_API_VERSION}/actions/embedText") def sign_request( self, diff --git a/litellm/llms/ollama/chat/transformation.py b/litellm/llms/ollama/chat/transformation.py index 644ab187f15..f9e7071db41 100644 --- a/litellm/llms/ollama/chat/transformation.py +++ b/litellm/llms/ollama/chat/transformation.py @@ -14,6 +14,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( extract_images_from_message, extract_reasoning_content, ) +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException from litellm.types.llms.ollama import ( @@ -218,10 +219,7 @@ class OllamaChatConfig(BaseConfig): Some providers need `model` in `api_base` """ - clean_api_base: Final = (api_base or "http://localhost:11434").rstrip("/") - if clean_api_base.endswith("/api/chat"): - return clean_api_base - return f"{clean_api_base}/api/chat" + return ensure_api_route(api_base or "http://localhost:11434", "api/chat") def transform_request( self, diff --git a/litellm/llms/ollama/common_utils.py b/litellm/llms/ollama/common_utils.py index 284a6ffea62..7894904c870 100644 --- a/litellm/llms/ollama/common_utils.py +++ b/litellm/llms/ollama/common_utils.py @@ -8,6 +8,7 @@ from pydantic import BaseModel, TypeAdapter from litellm import verbose_logger from litellm._uuid import uuid +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.types.utils import ModelResponse, ModelResponseStream @@ -149,7 +150,7 @@ class OllamaModelInfo(BaseLLMModelInfo): names: Final[set[str]] = set() try: - resp: Final = httpx.get(f"{base}/api/tags", headers=headers) + resp: Final = httpx.get(join_api_base(base, "api/tags"), headers=headers) resp.raise_for_status() data: Final = resp.json() # Expecting a dict with a 'models' list @@ -230,7 +231,7 @@ class OllamaModelInfo(BaseLLMModelInfo): try: response: Final = module_level_client.post( - url=f"{api_base}/api/show", + url=join_api_base(api_base, "api/show"), json={"name": model}, headers=headers, ) diff --git a/litellm/llms/ollama/completion/handler.py b/litellm/llms/ollama/completion/handler.py index 816028499e8..d4110e55bf0 100644 --- a/litellm/llms/ollama/completion/handler.py +++ b/litellm/llms/ollama/completion/handler.py @@ -10,6 +10,7 @@ from typing import Any, Final, Protocol, TypedDict from typing_extensions import NotRequired, ReadOnly import litellm +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.types.utils import EmbeddingResponse @@ -26,13 +27,6 @@ class OllamaEmbeddingResponse(TypedDict): prompt_eval_count: ReadOnly[NotRequired[int]] -def _ollama_embed_url(api_base: str) -> str: - clean_api_base: Final = api_base.rstrip("/") - if clean_api_base.endswith("/api/embed"): - return clean_api_base - return f"{clean_api_base}/api/embed" - - def _prepare_ollama_embedding_payload( model: str, prompts: list[str], optional_params: Mapping[str, object] ) -> dict[str, object]: @@ -99,7 +93,7 @@ async def ollama_aembeddings( ): data: Final = _prepare_ollama_embedding_payload(model, prompts, optional_params) - response: Final = await litellm.module_level_aclient.post(url=_ollama_embed_url(api_base), json=data) + response: Final = await litellm.module_level_aclient.post(url=ensure_api_route(api_base, "api/embed"), json=data) response_json: Final[OllamaEmbeddingResponse] = response.json() return _process_ollama_embedding_response( @@ -123,7 +117,7 @@ def ollama_embeddings( ): data: Final = _prepare_ollama_embedding_payload(model, prompts, optional_params) - response: Final = litellm.module_level_client.post(url=_ollama_embed_url(api_base), json=data) + response: Final = litellm.module_level_client.post(url=ensure_api_route(api_base, "api/embed"), json=data) response_json: Final[OllamaEmbeddingResponse] = response.json() return _process_ollama_embedding_response( diff --git a/litellm/llms/ollama/completion/transformation.py b/litellm/llms/ollama/completion/transformation.py index ac41310fb61..30de22b17ff 100644 --- a/litellm/llms/ollama/completion/transformation.py +++ b/litellm/llms/ollama/completion/transformation.py @@ -23,6 +23,7 @@ from litellm.litellm_core_utils.prompt_templates.image_handling import ( async_inline_remote_media, inline_remote_image_urls, ) +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException from litellm.types.llms.base import LiteLLMBaseModel @@ -500,10 +501,7 @@ class OllamaConfig(BaseConfig): Some providers need `model` in `api_base` """ - clean_api_base: Final = (api_base or "http://localhost:11434").rstrip("/") - if clean_api_base.endswith("/api/generate"): - return clean_api_base - return f"{clean_api_base}/api/generate" + return ensure_api_route(api_base or "http://localhost:11434", "api/generate") def get_model_response_iterator( self, diff --git a/litellm/llms/oobabooga/chat/oobabooga.py b/litellm/llms/oobabooga/chat/oobabooga.py index e21a165d757..8879c20c5eb 100644 --- a/litellm/llms/oobabooga/chat/oobabooga.py +++ b/litellm/llms/oobabooga/chat/oobabooga.py @@ -3,6 +3,7 @@ from collections.abc import Callable from typing import TYPE_CHECKING, Final import litellm +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.custom_httpx.http_handler import get_httpx_client from litellm.utils import EmbeddingResponse, ModelResponse, Usage @@ -49,7 +50,7 @@ def completion( ) model = model - completion_url = completion_url + "/v1/chat/completions" + completion_url = join_api_base(completion_url, "v1/chat/completions") data: Final = oobabooga_config.transform_request( model=model, messages=messages, @@ -103,7 +104,7 @@ def embedding( if model.startswith(("http://", "https://")): embeddings_url = model elif api_base: - embeddings_url = f"{api_base}/v1/embeddings" + embeddings_url = join_api_base(api_base, "v1/embeddings") else: raise OobaboogaError( status_code=404, diff --git a/litellm/llms/openai/chat/gpt_transformation.py b/litellm/llms/openai/chat/gpt_transformation.py index 4d7ef833181..b4c92848752 100644 --- a/litellm/llms/openai/chat/gpt_transformation.py +++ b/litellm/llms/openai/chat/gpt_transformation.py @@ -34,6 +34,7 @@ from litellm.litellm_core_utils.prompt_templates.image_handling import ( async_convert_url_to_base64, convert_url_to_base64, ) +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator from litellm.llms.base_llm.base_utils import BaseLLMModelInfo from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException @@ -807,7 +808,7 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): if endpoint in api_base: return api_base - return f"{api_base}/{endpoint}" + return join_api_base(api_base, endpoint) def validate_environment( self, diff --git a/litellm/llms/openai/containers/transformation.py b/litellm/llms/openai/containers/transformation.py index 6cb2396b85d..5d03585189d 100644 --- a/litellm/llms/openai/containers/transformation.py +++ b/litellm/llms/openai/containers/transformation.py @@ -9,7 +9,7 @@ from litellm.litellm_core_utils.hidden_params import get_or_create_hidden_params from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import ( StandardBuiltInToolCostTracking, ) -from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.secret_managers.main import get_secret_str from litellm.types.containers.main import ( ContainerCreateOptionalRequestParams, @@ -130,7 +130,7 @@ class OpenAIContainerConfig(BaseContainerConfig): or "https://api.openai.com/v1" ) - return f"{api_base.rstrip('/')}/containers" + return join_api_base(api_base, "containers") def transform_container_create_request( self, diff --git a/litellm/llms/openai/evals/transformation.py b/litellm/llms/openai/evals/transformation.py index 24de7f6d479..581e5e8ad87 100644 --- a/litellm/llms/openai/evals/transformation.py +++ b/litellm/llms/openai/evals/transformation.py @@ -8,7 +8,7 @@ from typing import Final import httpx from litellm._logging import verbose_logger -from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.llms.base_llm.evals.transformation import ( BaseEvalsAPIConfig, LiteLLMLoggingObj, @@ -67,8 +67,8 @@ class OpenAIEvalsConfig(BaseEvalsAPIConfig): if eval_id: encoded_eval_id: Final = encode_url_path_segment(eval_id, field_name="eval_id") - return f"{api_base}/v1/evals/{encoded_eval_id}" - return f"{api_base}/v1/{endpoint}" + return join_api_base(api_base, f"v1/evals/{encoded_eval_id}") + return join_api_base(api_base, f"v1/{endpoint}") def transform_create_eval_request( self, @@ -226,7 +226,9 @@ class OpenAIEvalsConfig(BaseEvalsAPIConfig): headers: dict, ) -> tuple[str, dict, dict]: """Transform cancel eval request for OpenAI""" - url: Final = f"{self.get_complete_url(api_base=api_base, endpoint='evals', eval_id=eval_id)}/cancel" + url: Final = join_api_base( + self.get_complete_url(api_base=api_base, endpoint="evals", eval_id=eval_id), "cancel" + ) # Empty body for cancel request request_body: Final[dict[str, object]] = {} @@ -260,7 +262,7 @@ class OpenAIEvalsConfig(BaseEvalsAPIConfig): api_base = litellm_params.api_base encoded_eval_id: Final = encode_url_path_segment(eval_id, field_name="eval_id") - url: Final = f"{api_base}/v1/evals/{encoded_eval_id}/runs" + url: Final = join_api_base(api_base, f"v1/evals/{encoded_eval_id}/runs") # Build request body request_body: Final = {k: v for k, v in create_request.items() if v is not None} @@ -293,7 +295,7 @@ class OpenAIEvalsConfig(BaseEvalsAPIConfig): api_base = litellm_params.api_base encoded_eval_id: Final = encode_url_path_segment(eval_id, field_name="eval_id") - url: Final = f"{api_base}/v1/evals/{encoded_eval_id}/runs" + url: Final = join_api_base(api_base, f"v1/evals/{encoded_eval_id}/runs") # Build query parameters query_params: Final[dict[str, object]] = {} @@ -335,7 +337,7 @@ class OpenAIEvalsConfig(BaseEvalsAPIConfig): """Transform get run request for OpenAI""" encoded_eval_id: Final = encode_url_path_segment(eval_id, field_name="eval_id") encoded_run_id: Final = encode_url_path_segment(run_id, field_name="run_id") - url: Final = f"{api_base}/v1/evals/{encoded_eval_id}/runs/{encoded_run_id}" + url: Final = join_api_base(api_base, f"v1/evals/{encoded_eval_id}/runs/{encoded_run_id}") verbose_logger.debug("Get run request - URL: %s", url) @@ -363,7 +365,7 @@ class OpenAIEvalsConfig(BaseEvalsAPIConfig): """Transform cancel run request for OpenAI""" encoded_eval_id: Final = encode_url_path_segment(eval_id, field_name="eval_id") encoded_run_id: Final = encode_url_path_segment(run_id, field_name="run_id") - url: Final = f"{api_base}/v1/evals/{encoded_eval_id}/runs/{encoded_run_id}/cancel" + url: Final = join_api_base(api_base, f"v1/evals/{encoded_eval_id}/runs/{encoded_run_id}/cancel") # Empty body for cancel request request_body: Final[dict[str, object]] = {} @@ -394,7 +396,7 @@ class OpenAIEvalsConfig(BaseEvalsAPIConfig): """Transform delete run request for OpenAI""" encoded_eval_id: Final = encode_url_path_segment(eval_id, field_name="eval_id") encoded_run_id: Final = encode_url_path_segment(run_id, field_name="run_id") - url: Final = f"{api_base}/v1/evals/{encoded_eval_id}/runs/{encoded_run_id}" + url: Final = join_api_base(api_base, f"v1/evals/{encoded_eval_id}/runs/{encoded_run_id}") # Empty body for delete request request_body: Final[dict[str, object]] = {} diff --git a/litellm/llms/openai/image_edit/transformation.py b/litellm/llms/openai/image_edit/transformation.py index a9e5b1b9f34..9def833052b 100644 --- a/litellm/llms/openai/image_edit/transformation.py +++ b/litellm/llms/openai/image_edit/transformation.py @@ -6,6 +6,7 @@ from httpx._types import RequestFiles import litellm from litellm.images.utils import ImageEditRequestUtils +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig from litellm.secret_managers.main import get_secret_str from litellm.types.images.main import ( @@ -184,7 +185,4 @@ class OpenAIImageEditConfig(BaseImageEditConfig): or "https://api.openai.com/v1" ) - # Remove trailing slashes - api_base = api_base.rstrip("/") - - return f"{api_base}/images/edits" + return join_api_base(api_base, "images/edits") diff --git a/litellm/llms/openai/realtime/http_transformation.py b/litellm/llms/openai/realtime/http_transformation.py index 61dbf20397f..49d33582c6f 100644 --- a/litellm/llms/openai/realtime/http_transformation.py +++ b/litellm/llms/openai/realtime/http_transformation.py @@ -1,6 +1,7 @@ """OpenAI realtime HTTP transformation config (client_secrets + realtime_calls).""" import litellm +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.realtime.http_transformation import BaseRealtimeHTTPConfig from litellm.secret_managers.main import get_secret_str @@ -15,17 +16,17 @@ class OpenAIRealtimeHTTPConfig(BaseRealtimeHTTPConfig): def get_complete_url(self, api_base: str | None, model: str, api_version: str | None = None) -> str: base = self.get_api_base(api_base).rstrip("/") base = base.removesuffix("/v1") - return f"{base}/v1/realtime/client_secrets" + return join_api_base(base, "v1/realtime/client_secrets") def get_realtime_calls_url(self, api_base: str | None, model: str, api_version: str | None = None) -> str: base = self.get_api_base(api_base).rstrip("/") base = base.removesuffix("/v1") - return f"{base}/v1/realtime/calls" + return join_api_base(base, "v1/realtime/calls") def get_transcription_session_url(self, api_base: str | None, model: str, api_version: str | None = None) -> str: base = self.get_api_base(api_base).rstrip("/") base = base.removesuffix("/v1") - return f"{base}/v1/realtime/transcription_sessions" + return join_api_base(base, "v1/realtime/transcription_sessions") def validate_environment( self, diff --git a/litellm/llms/openai/responses/count_tokens/transformation.py b/litellm/llms/openai/responses/count_tokens/transformation.py index 6d882e0eeb1..c1ee42444c5 100644 --- a/litellm/llms/openai/responses/count_tokens/transformation.py +++ b/litellm/llms/openai/responses/count_tokens/transformation.py @@ -13,6 +13,7 @@ from pydantic import JsonValue, TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict from litellm.litellm_core_utils.tool_fields import DropPolicy +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.anthropic.pass_through.convertible_blocks import malformed_tool_result_text_error from litellm.llms.anthropic.pass_through.messages.utils import normalize_native_messages from litellm.llms.anthropic.pass_through.responses_adapters.transformation import ( @@ -148,9 +149,8 @@ class OpenAICountTokensConfig: """ def get_openai_count_tokens_endpoint(self, api_base: str | None = None) -> str: - base = api_base or "https://api.openai.com/v1" - base = base.rstrip("/") - return f"{base}/responses/input_tokens" + base: Final = api_base or "https://api.openai.com/v1" + return join_api_base(base, "responses/input_tokens") def transform_request_to_count_tokens( self, diff --git a/litellm/llms/openai/responses/transformation.py b/litellm/llms/openai/responses/transformation.py index 8f29a71d63b..b90f6aa01da 100644 --- a/litellm/llms/openai/responses/transformation.py +++ b/litellm/llms/openai/responses/transformation.py @@ -23,7 +23,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( with_prompt_cache_breakpoint, ) from litellm.litellm_core_utils.safe_json_loads import safe_json_loads -from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig from litellm.llms.openai.chat.gpt_5_transformation import ( is_gpt_reasoning_series_name, @@ -881,10 +881,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): or "https://api.openai.com/v1" ) - # Remove trailing slashes - api_base = api_base.rstrip("/") - - return f"{api_base}/responses" + return join_api_base(api_base, "responses") def transform_streaming_response( self, @@ -1111,7 +1108,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): order: Literal["asc", "desc"] = "desc", ) -> tuple[str, dict]: encoded_response_id: Final = encode_url_path_segment(response_id, field_name="response_id") - url: Final = f"{api_base}/{encoded_response_id}/input_items" + url: Final = join_api_base(api_base, f"{encoded_response_id}/input_items") params: Final[dict[str, object]] = {} if after is not None: params["after"] = after @@ -1152,7 +1149,7 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): - POST /v1/responses/{response_id}/cancel """ encoded_response_id: Final = encode_url_path_segment(response_id, field_name="response_id") - url: Final = f"{api_base}/{encoded_response_id}/cancel" + url: Final = join_api_base(api_base, f"{encoded_response_id}/cancel") data: Final[dict] = {} return url, data diff --git a/litellm/llms/openai/transcriptions/whisper_transformation.py b/litellm/llms/openai/transcriptions/whisper_transformation.py index 8f96f688189..eb4f1b6c9ce 100644 --- a/litellm/llms/openai/transcriptions/whisper_transformation.py +++ b/litellm/llms/openai/transcriptions/whisper_transformation.py @@ -4,6 +4,7 @@ from typing import Final from httpx import Headers, Response +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.audio_transcription.transformation import ( AudioTranscriptionRequestData, BaseAudioTranscriptionConfig, @@ -46,9 +47,9 @@ class OpenAIWhisperAudioTranscriptionConfig(BaseAudioTranscriptionConfig): # if endswith "/v1" if api_base and api_base.endswith("/v1"): - api_base = f"{api_base}/audio/transcriptions" + api_base = join_api_base(api_base, "audio/transcriptions") else: - api_base = f"{api_base}/v1/audio/transcriptions" + api_base = join_api_base(api_base, "v1/audio/transcriptions") return api_base or "" diff --git a/litellm/llms/openai/vector_store_files/transformation.py b/litellm/llms/openai/vector_store_files/transformation.py index 7f65ade224e..bd8145bb460 100644 --- a/litellm/llms/openai/vector_store_files/transformation.py +++ b/litellm/llms/openai/vector_store_files/transformation.py @@ -4,7 +4,7 @@ from typing import Final, cast import httpx import litellm -from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.llms.base_llm.vector_store_files.transformation import ( BaseVectorStoreFilesConfig, ) @@ -92,16 +92,15 @@ class OpenAIVectorStoreFilesConfig(BaseVectorStoreFilesConfig): vector_store_id: str, litellm_params: Mapping[str, object], ) -> str: - base_url = ( + base_url: Final = ( api_base or litellm.api_base or get_secret_str("OPENAI_BASE_URL") or get_secret_str("OPENAI_API_BASE") or "https://api.openai.com/v1" ) - base_url = base_url.rstrip("/") encoded_vector_store_id: Final = encode_url_path_segment(vector_store_id, field_name="vector_store_id") - return f"{base_url}/vector_stores/{encoded_vector_store_id}/files" + return join_api_base(base_url, f"vector_stores/{encoded_vector_store_id}/files") def transform_create_vector_store_file_request( self, @@ -187,7 +186,7 @@ class OpenAIVectorStoreFilesConfig(BaseVectorStoreFilesConfig): api_base: str, ) -> tuple[str, dict[str, object]]: encoded_file_id: Final = encode_url_path_segment(file_id, field_name="file_id") - return f"{api_base}/{encoded_file_id}/content", {} + return join_api_base(api_base, f"{encoded_file_id}/content"), {} def transform_retrieve_vector_store_file_content_response( self, diff --git a/litellm/llms/openai/vector_stores/transformation.py b/litellm/llms/openai/vector_stores/transformation.py index 46105a66be3..d108802f7f9 100644 --- a/litellm/llms/openai/vector_stores/transformation.py +++ b/litellm/llms/openai/vector_stores/transformation.py @@ -3,7 +3,7 @@ from typing import TYPE_CHECKING, Any, Final, cast import httpx import litellm -from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig from litellm.llms.openai.workload_identity import resolve_openai_bearer_token from litellm.secret_managers.main import get_secret_str @@ -88,10 +88,7 @@ class OpenAIVectorStoreConfig(BaseVectorStoreConfig): or "https://api.openai.com/v1" ) - # Remove trailing slashes - api_base = api_base.rstrip("/") - - return f"{api_base}/vector_stores" + return join_api_base(api_base, "vector_stores") def transform_search_vector_store_request( self, @@ -105,7 +102,7 @@ class OpenAIVectorStoreConfig(BaseVectorStoreConfig): ) -> tuple[str, dict]: encoded_vector_store_id: Final = encode_url_path_segment(vector_store_id, field_name="vector_store_id") base_url, query_separator, query_string = api_base.partition("?") - url: Final = f"{base_url}/{encoded_vector_store_id}/search{query_separator}{query_string}" + url: Final = f"{join_api_base(base_url, f'{encoded_vector_store_id}/search')}{query_separator}{query_string}" typed_request_body: Final = VectorStoreSearchRequest( query=query, filters=vector_store_search_optional_params.get("filters", None), diff --git a/litellm/llms/openai/videos/transformation.py b/litellm/llms/openai/videos/transformation.py index d422e51f649..0560e1c2e3a 100644 --- a/litellm/llms/openai/videos/transformation.py +++ b/litellm/llms/openai/videos/transformation.py @@ -8,7 +8,7 @@ from urllib.parse import quote import httpx from httpx._types import FileContent, FileTypes, RequestFiles -from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.llms.base_llm.videos.transformation import BaseVideoConfig from litellm.llms.openai.image_edit.transformation import ImageEditRequestUtils from litellm.llms.openai.workload_identity import resolve_openai_bearer_token @@ -98,7 +98,7 @@ class OpenAIVideoConfig(BaseVideoConfig): if api_base is None: api_base = "https://api.openai.com/v1" - return f"{api_base.rstrip('/')}/videos" + return join_api_base(api_base, "videos") def use_multipart_form_data(self) -> bool: return True @@ -214,7 +214,7 @@ class OpenAIVideoConfig(BaseVideoConfig): encoded_video_id: Final = encode_url_path_segment(original_video_id, field_name="video_id") # Construct the URL for video content download - url = f"{api_base.rstrip('/')}/{encoded_video_id}/content" + url = join_api_base(api_base, f"{encoded_video_id}/content") if variant is not None: # Encode the user-controlled ``variant`` so a value like # ``thumbnail&extra=1`` cannot inject additional query params @@ -246,7 +246,7 @@ class OpenAIVideoConfig(BaseVideoConfig): encoded_video_id: Final = encode_url_path_segment(original_video_id, field_name="video_id") # Construct the URL for video remix - url: Final = f"{api_base.rstrip('/')}/{encoded_video_id}/remix" + url: Final = join_api_base(api_base, f"{encoded_video_id}/remix") # Prepare the request data data: Final[dict[str, object]] = {"prompt": prompt} @@ -465,7 +465,7 @@ class OpenAIVideoConfig(BaseVideoConfig): litellm_params: GenericLiteLLMParams, headers: dict, ) -> tuple[str, list]: - url: Final = f"{api_base.rstrip('/')}/characters" + url: Final = join_api_base(api_base, "characters") files_list: Final[list[tuple[str, FileTypes]]] = [("name", (None, name))] self._add_video_to_files(files_list, video, "video") return url, files_list @@ -486,7 +486,7 @@ class OpenAIVideoConfig(BaseVideoConfig): ) -> tuple[str, dict]: original_character_id: Final = extract_original_character_id(character_id) encoded_character_id: Final = encode_url_path_segment(original_character_id, field_name="character_id") - url: Final = f"{api_base.rstrip('/')}/characters/{encoded_character_id}" + url: Final = join_api_base(api_base, f"characters/{encoded_character_id}") return url, {} def transform_video_get_character_response( @@ -507,7 +507,7 @@ class OpenAIVideoConfig(BaseVideoConfig): extra_body: dict[str, object] | None = None, prefetched_source_data: dict[str, object] | None = None, ) -> tuple[str, Mapping[str, object], RequestFiles | None]: - url: Final = f"{api_base.rstrip('/')}/edits" + url: Final = join_api_base(api_base, "edits") if video_file is not None: files: Final[RequestFiles] = (self._video_file_tuple(video_file, "video"),) @@ -547,7 +547,7 @@ class OpenAIVideoConfig(BaseVideoConfig): extra_body: dict[str, object] | None = None, ) -> tuple[str, dict]: original_video_id: Final = extract_original_video_id(video_id) - url: Final = f"{api_base.rstrip('/')}/extensions" + url: Final = join_api_base(api_base, "extensions") data: Final[dict[str, object]] = { "prompt": prompt, "seconds": seconds, diff --git a/litellm/llms/openai_like/common_utils.py b/litellm/llms/openai_like/common_utils.py index 11b85e52af5..5d3085a3901 100644 --- a/litellm/llms/openai_like/common_utils.py +++ b/litellm/llms/openai_like/common_utils.py @@ -2,6 +2,8 @@ from typing import Literal import httpx +from litellm.litellm_core_utils.url_utils import join_api_base + class OpenAILikeError(Exception): def __init__(self, status_code, message): @@ -48,7 +50,7 @@ class OpenAILikeBase: if not custom_endpoint: if endpoint_type == "chat_completions": - api_base = f"{api_base}/chat/completions" + api_base = join_api_base(api_base, "chat/completions") elif endpoint_type == "embeddings": - api_base = f"{api_base}/embeddings" + api_base = join_api_base(api_base, "embeddings") return api_base, headers diff --git a/litellm/llms/openai_like/dynamic_config.py b/litellm/llms/openai_like/dynamic_config.py index f0bf710b951..685561f55a5 100644 --- a/litellm/llms/openai_like/dynamic_config.py +++ b/litellm/llms/openai_like/dynamic_config.py @@ -9,6 +9,7 @@ from litellm._logging import verbose_logger from litellm.litellm_core_utils.prompt_templates.common_utils import ( handle_messages_with_content_list_to_str_conversion, ) +from litellm.litellm_core_utils.url_utils import ensure_api_route, join_api_base from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig from litellm.llms.openai_like.chat.transformation import OpenAILikeChatConfig from litellm.secret_managers.main import get_secret_str @@ -91,13 +92,9 @@ def create_config_class(provider: SimpleProviderConfig): if not api_base: api_base = provider.base_url - if api_base is None: - raise ValueError(f"api_base is required for provider {provider.slug}") - - if not api_base.endswith("/chat/completions"): - api_base = f"{api_base}/chat/completions" - - return api_base + if isinstance(api_base, str): + return ensure_api_route(api_base, "chat/completions") + raise ValueError(f"api_base is required for provider {provider.slug}") def get_supported_openai_params(self, model: str) -> list: """Get supported OpenAI params, excluding tool-related params for models @@ -226,11 +223,9 @@ def create_responses_config_class(provider: SimpleProviderConfig) -> "type[OpenA if not api_base: api_base = provider.base_url - if api_base is None: - raise ValueError(f"api_base is required for provider {provider.slug}") - - api_base = api_base.rstrip("/") - return f"{api_base}/responses" + if isinstance(api_base, str): + return join_api_base(api_base, "responses") + raise ValueError(f"api_base is required for provider {provider.slug}") def transform_responses_api_request( self, diff --git a/litellm/llms/openai_like/messages/transformation.py b/litellm/llms/openai_like/messages/transformation.py index ef501720e2d..12c8b1da166 100644 --- a/litellm/llms/openai_like/messages/transformation.py +++ b/litellm/llms/openai_like/messages/transformation.py @@ -1,6 +1,7 @@ from typing import Any, Final import litellm +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.anthropic.common_utils import normalize_cache_control_in_anthropic_payload from litellm.llms.anthropic.pass_through.messages.transformation import ( AnthropicMessagesConfig, @@ -113,7 +114,7 @@ class OpenAILikeAnthropicMessagesConfig(AnthropicMessagesConfig): if base.endswith("/v1/messages"): return base base = base.removesuffix("/v1") - return f"{base}/v1/messages" + return join_api_base(base, "v1/messages") class JSONProviderAnthropicMessagesConfig(OpenAILikeAnthropicMessagesConfig): diff --git a/litellm/llms/openai_like/responses/transformation.py b/litellm/llms/openai_like/responses/transformation.py index 9655772e953..68de4065b09 100644 --- a/litellm/llms/openai_like/responses/transformation.py +++ b/litellm/llms/openai_like/responses/transformation.py @@ -8,6 +8,7 @@ provider-specific resolution (slug, API key env var, base URL). from typing import Final +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig from litellm.secret_managers.main import get_secret_str from litellm.types.router import GenericLiteLLMParams @@ -47,5 +48,4 @@ class OpenAILikeResponsesConfig(OpenAIResponsesAPIConfig): api_base = api_base or get_secret_str("OPENAI_LIKE_API_BASE") if not api_base: raise ValueError("api_base is required for openai_like provider") - api_base = api_base.rstrip("/") - return f"{api_base}/responses" + return join_api_base(api_base, "responses") diff --git a/litellm/llms/openrouter/embedding/transformation.py b/litellm/llms/openrouter/embedding/transformation.py index 14b0e462ea7..d1530a29515 100644 --- a/litellm/llms/openrouter/embedding/transformation.py +++ b/litellm/llms/openrouter/embedding/transformation.py @@ -11,6 +11,7 @@ from typing import TYPE_CHECKING, Any, Final import httpx +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig from litellm.types.llms.openai import AllEmbeddingInputValues from litellm.types.utils import EmbeddingResponse @@ -84,15 +85,7 @@ class OpenrouterEmbeddingConfig(BaseEmbeddingConfig): """ Get the complete URL for OpenRouter Embedding API endpoint. """ - # api_base is already set to https://openrouter.ai/api/v1 in main.py - # Remove trailing slashes - if api_base: - api_base = api_base.rstrip("/") - else: - api_base = "https://openrouter.ai/api/v1" - - # Return the embeddings endpoint - return f"{api_base}/embeddings" + return join_api_base(api_base or "https://openrouter.ai/api/v1", "embeddings") def transform_embedding_request( self, diff --git a/litellm/llms/openrouter/image_edit/transformation.py b/litellm/llms/openrouter/image_edit/transformation.py index a73a2fe4192..9644ee3f280 100644 --- a/litellm/llms/openrouter/image_edit/transformation.py +++ b/litellm/llms/openrouter/image_edit/transformation.py @@ -52,6 +52,7 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm from litellm.images.utils import ImageEditRequestUtils +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig from litellm.llms.openrouter.common_utils import OpenRouterException @@ -187,11 +188,8 @@ class OpenRouterImageEditConfig(BaseImageEditConfig): api_base: str | None, litellm_params: dict, ) -> str: - base_url = api_base or get_secret_str("OPENROUTER_API_BASE") or "https://openrouter.ai/api/v1" - base_url = base_url.rstrip("/") - if not base_url.endswith("/chat/completions"): - return f"{base_url}/chat/completions" - return base_url + base_url: Final = api_base or get_secret_str("OPENROUTER_API_BASE") or "https://openrouter.ai/api/v1" + return ensure_api_route(base_url, "chat/completions") def transform_image_edit_request( self, diff --git a/litellm/llms/openrouter/image_generation/transformation.py b/litellm/llms/openrouter/image_generation/transformation.py index 355df81b87a..ce6e62aaa88 100644 --- a/litellm/llms/openrouter/image_generation/transformation.py +++ b/litellm/llms/openrouter/image_generation/transformation.py @@ -35,6 +35,7 @@ from pydantic import ConfigDict, TypeAdapter import litellm from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR, set_hidden_params +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.image_generation.transformation import ( BaseImageGenerationConfig, @@ -269,10 +270,7 @@ class OpenRouterImageGenerationConfig(BaseImageGenerationConfig): Default: https://openrouter.ai/api/v1/chat/completions """ if api_base: - if not api_base.endswith("/chat/completions"): - api_base = api_base.rstrip("/") - return f"{api_base}/chat/completions" - return api_base + return ensure_api_route(api_base, "chat/completions") return "https://openrouter.ai/api/v1/chat/completions" diff --git a/litellm/llms/openrouter/responses/transformation.py b/litellm/llms/openrouter/responses/transformation.py index aeb5b09e14a..e166ccaad2d 100644 --- a/litellm/llms/openrouter/responses/transformation.py +++ b/litellm/llms/openrouter/responses/transformation.py @@ -18,6 +18,7 @@ import litellm from litellm.litellm_core_utils.core_helpers import RESPONSE_COST_HEADER from litellm.litellm_core_utils.hidden_params import set_hidden_params from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import ( @@ -138,9 +139,7 @@ class OpenRouterResponsesAPIConfig(OpenAIResponsesAPIConfig): api_base or litellm.api_base or get_secret_str("OPENROUTER_API_BASE") or "https://openrouter.ai/api/v1" ) - api_base = api_base.rstrip("/") - - return f"{api_base}/responses" + return join_api_base(api_base, "responses") def supports_native_websocket(self) -> bool: """OpenRouter does not support native WebSocket for Responses API""" diff --git a/litellm/llms/opensandbox/sandbox/transformation.py b/litellm/llms/opensandbox/sandbox/transformation.py index 13c9d805638..574fbfa9e5b 100644 --- a/litellm/llms/opensandbox/sandbox/transformation.py +++ b/litellm/llms/opensandbox/sandbox/transformation.py @@ -21,6 +21,7 @@ from litellm.constants import ( OPEN_SANDBOX_READY_TIMEOUT, ) from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR, set_hidden_params +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.sandbox.transformation import ( SANDBOX_MAX_OUTPUT_BYTES, BaseSandboxConfig, @@ -93,7 +94,7 @@ class OpenSandboxSandboxConfig(BaseSandboxConfig): ) response: Final = await self._http(client).post( - url=f"{base}/sandboxes", + url=join_api_base(base, "sandboxes"), headers=self._lifecycle_headers(key), json=body, ) @@ -192,7 +193,7 @@ class OpenSandboxSandboxConfig(BaseSandboxConfig): key: Final = self._api_key(api_key=api_key, handle=handle) try: response: Final = await self._http(client).delete( - url=f"{base}/sandboxes/{handle.id}", + url=join_api_base(base, f"sandboxes/{handle.id}"), headers=self._lifecycle_headers(key), ) except httpx.HTTPStatusError as e: @@ -258,7 +259,7 @@ class OpenSandboxSandboxConfig(BaseSandboxConfig): deadline: Final = time.monotonic() + ready_timeout while True: response = await self._http(client).get( - url=f"{api_base}/sandboxes/{sandbox_id}", + url=join_api_base(api_base, f"sandboxes/{sandbox_id}"), headers=headers, ) data = _JSON_VALUE.validate_python(response.json()) @@ -316,7 +317,7 @@ class OpenSandboxSandboxConfig(BaseSandboxConfig): client: AsyncHTTPHandler | None, ) -> tuple[str, dict[str, str]]: response: Final = await self._http(client).get( - url=f"{api_base}/sandboxes/{sandbox_id}/endpoints/{OPEN_SANDBOX_EXECD_PORT}", + url=join_api_base(api_base, f"sandboxes/{sandbox_id}/endpoints/{OPEN_SANDBOX_EXECD_PORT}"), headers=headers, params={"use_server_proxy": use_server_proxy}, ) diff --git a/litellm/llms/ovhcloud/audio_transcription/transformation.py b/litellm/llms/ovhcloud/audio_transcription/transformation.py index 3ae5f6fa69c..dd61bfb0c13 100644 --- a/litellm/llms/ovhcloud/audio_transcription/transformation.py +++ b/litellm/llms/ovhcloud/audio_transcription/transformation.py @@ -12,6 +12,7 @@ import httpx from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.audio_utils.utils import process_audio_file +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.audio_transcription.transformation import ( AudioTranscriptionRequestData, BaseAudioTranscriptionConfig, @@ -63,9 +64,8 @@ class OVHCloudAudioTranscriptionConfig(BaseAudioTranscriptionConfig): litellm_params: dict, stream: bool | None = None, ) -> str: - api_base = "https://oai.endpoints.kepler.ai.cloud.ovh.net/v1" if api_base is None else api_base.rstrip("/") - complete_url: Final = f"{api_base}/audio/transcriptions" - return complete_url + base_url: Final = "https://oai.endpoints.kepler.ai.cloud.ovh.net/v1" if api_base is None else api_base + return join_api_base(base_url, "audio/transcriptions") def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException: return OVHCloudException( diff --git a/litellm/llms/ovhcloud/chat/transformation.py b/litellm/llms/ovhcloud/chat/transformation.py index 3d3953eed34..91b3dbe6268 100644 --- a/litellm/llms/ovhcloud/chat/transformation.py +++ b/litellm/llms/ovhcloud/chat/transformation.py @@ -9,6 +9,7 @@ from typing import Final import httpx +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig @@ -31,9 +32,8 @@ class OVHCloudChatConfig(OpenAIGPTConfig): litellm_params: dict, stream: bool | None = None, ) -> str: - api_base = "https://oai.endpoints.kepler.ai.cloud.ovh.net/v1" if api_base is None else api_base.rstrip("/") - complete_url: Final = f"{api_base}/chat/completions" - return complete_url + base_url: Final = "https://oai.endpoints.kepler.ai.cloud.ovh.net/v1" if api_base is None else api_base + return join_api_base(base_url, "chat/completions") def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException: return OVHCloudException( diff --git a/litellm/llms/ovhcloud/embedding/transformation.py b/litellm/llms/ovhcloud/embedding/transformation.py index cdf095e6ad1..487c40afa52 100644 --- a/litellm/llms/ovhcloud/embedding/transformation.py +++ b/litellm/llms/ovhcloud/embedding/transformation.py @@ -10,6 +10,7 @@ import httpx from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig from litellm.secret_managers.main import get_secret_str @@ -35,9 +36,8 @@ class OVHCloudEmbeddingConfig(BaseEmbeddingConfig): litellm_params: dict, stream: bool | None = None, ) -> str: - api_base = "https://oai.endpoints.kepler.ai.cloud.ovh.net/v1" if api_base is None else api_base.rstrip("/") - complete_url: Final = f"{api_base}/embeddings" - return complete_url + base_url: Final = "https://oai.endpoints.kepler.ai.cloud.ovh.net/v1" if api_base is None else api_base + return join_api_base(base_url, "embeddings") def validate_environment( self, diff --git a/litellm/llms/perplexity/embedding/transformation.py b/litellm/llms/perplexity/embedding/transformation.py index 945a58d4d75..7b4da933976 100644 --- a/litellm/llms/perplexity/embedding/transformation.py +++ b/litellm/llms/perplexity/embedding/transformation.py @@ -20,6 +20,7 @@ import httpx from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig from litellm.secret_managers.main import get_secret_str @@ -68,7 +69,7 @@ class PerplexityEmbeddingConfig(BaseEmbeddingConfig): if api_base: clean_api_base: Final = api_base.rstrip("/") if not clean_api_base.endswith("/embeddings"): - return f"{clean_api_base}/v1/embeddings" + return join_api_base(clean_api_base, "v1/embeddings") return clean_api_base return "https://api.perplexity.ai/v1/embeddings" diff --git a/litellm/llms/perplexity/responses/transformation.py b/litellm/llms/perplexity/responses/transformation.py index 91d40d8247f..a2b7eceea40 100644 --- a/litellm/llms/perplexity/responses/transformation.py +++ b/litellm/llms/perplexity/responses/transformation.py @@ -14,6 +14,7 @@ from typing import Any, Final import httpx from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig from litellm.secret_managers.main import get_secret_str @@ -51,7 +52,7 @@ class PerplexityResponsesConfig(OpenAIResponsesAPIConfig): def get_complete_url(self, api_base: str | None, litellm_params: dict) -> str: api_base = api_base or get_secret_str("PERPLEXITY_API_BASE") or "https://api.perplexity.ai" - return f"{api_base.rstrip('/')}/v1/responses" + return join_api_base(api_base, "v1/responses") def _ensure_message_type(self, input: str | ResponseInputParam) -> str | ResponseInputParam: """Ensure list input items have type='message' (required by Perplexity).""" diff --git a/litellm/llms/perplexity/search/transformation.py b/litellm/llms/perplexity/search/transformation.py index b46e4225309..1f3b0fe729d 100644 --- a/litellm/llms/perplexity/search/transformation.py +++ b/litellm/llms/perplexity/search/transformation.py @@ -9,6 +9,7 @@ import httpx from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm.search.transformation import ( BaseSearchConfig, SearchResponse, @@ -78,11 +79,7 @@ class PerplexitySearchConfig(BaseSearchConfig): """ Get complete URL for Search endpoint. """ - resolved_api_base: Final = api_base or get_secret_str("PERPLEXITY_API_BASE") or self.PERPLEXITY_API_BASE - clean_api_base: Final = resolved_api_base.rstrip("/") - if clean_api_base.endswith("/search"): - return clean_api_base - return f"{clean_api_base}/search" + return ensure_api_route(api_base or get_secret_str("PERPLEXITY_API_BASE") or self.PERPLEXITY_API_BASE, "search") def transform_search_request( self, diff --git a/litellm/llms/pg_vector/vector_stores/transformation.py b/litellm/llms/pg_vector/vector_stores/transformation.py index e4b06c36bf4..eeecd02653d 100644 --- a/litellm/llms/pg_vector/vector_stores/transformation.py +++ b/litellm/llms/pg_vector/vector_stores/transformation.py @@ -1,6 +1,6 @@ from typing import TYPE_CHECKING, Any, Final -from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.llms.openai.vector_stores.transformation import OpenAIVectorStoreConfig from litellm.secret_managers.main import get_secret_str from litellm.types.router import GenericLiteLLMParams @@ -66,10 +66,7 @@ class PGVectorStoreConfig(OpenAIVectorStoreConfig): "PG Vector API base URL is required. Set PG_VECTOR_API_BASE environment variable or pass api_base in litellm_params." ) - # Remove trailing slashes - api_base = api_base.rstrip("/") - - return f"{api_base}/v1/vector_stores" + return join_api_base(api_base, "v1/vector_stores") def transform_search_vector_store_request( self, @@ -82,7 +79,7 @@ class PGVectorStoreConfig(OpenAIVectorStoreConfig): extra_body: dict[str, Any] | None = None, ) -> tuple[str, dict]: encoded_vector_store_id: Final = encode_url_path_segment(vector_store_id, field_name="vector_store_id") - url: Final = f"{api_base}/{encoded_vector_store_id}/search" + url: Final = join_api_base(api_base, f"{encoded_vector_store_id}/search") _, request_body = super().transform_search_vector_store_request( vector_store_id=vector_store_id, query=query, diff --git a/litellm/llms/predibase/chat/transformation.py b/litellm/llms/predibase/chat/transformation.py index 0dd7e27722e..b30713fa119 100644 --- a/litellm/llms/predibase/chat/transformation.py +++ b/litellm/llms/predibase/chat/transformation.py @@ -13,6 +13,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import ( prompt_factory, render_custom_prompt_template, ) +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import Choices, Message, ModelResponse, Usage @@ -337,12 +338,12 @@ class PredibaseConfig(BaseConfig): elif "PREDIBASE_API_BASE" in os.environ: base_url = os.getenv("PREDIBASE_API_BASE", "") - completion_url = f"{base_url}/{tenant_id}/deployments/v2/llms/{model}" + completion_url = join_api_base(base_url, f"{tenant_id}/deployments/v2/llms/{model}") should_stream: Final = stream if stream is not None else optional_params.get("stream", False) if should_stream is True: - completion_url += "/generate_stream" + completion_url = join_api_base(completion_url, "generate_stream") else: - completion_url += "/generate" + completion_url = join_api_base(completion_url, "generate") return completion_url def get_error_class(self, error_message: str, status_code: int, headers: dict | Headers) -> BaseLLMException: diff --git a/litellm/llms/ragflow/chat/transformation.py b/litellm/llms/ragflow/chat/transformation.py index 97b723f7dd5..f626cd8e55f 100644 --- a/litellm/llms/ragflow/chat/transformation.py +++ b/litellm/llms/ragflow/chat/transformation.py @@ -13,7 +13,7 @@ Model name format: from typing import Final import litellm -from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.llms.openai.openai import OpenAIConfig from litellm.secret_managers.main import get_secret, get_secret_str from litellm.types.llms.openai import AllMessageValues @@ -122,7 +122,7 @@ class RAGFlowConfig(OpenAIConfig): if not path.startswith("/"): path = "/" + path - return f"{api_base}{path}" + return join_api_base(api_base, path) def _get_openai_compatible_provider_info( self, diff --git a/litellm/llms/ragflow/vector_stores/transformation.py b/litellm/llms/ragflow/vector_stores/transformation.py index 007b047cc1b..ef44b6e4c0f 100644 --- a/litellm/llms/ragflow/vector_stores/transformation.py +++ b/litellm/llms/ragflow/vector_stores/transformation.py @@ -5,6 +5,7 @@ import httpx from pydantic import ConfigDict, RootModel, with_config from typing_extensions import NotRequired, ReadOnly, TypedDict +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig from litellm.secret_managers.main import get_secret_str from litellm.types.router import GenericLiteLLMParams @@ -95,15 +96,16 @@ class RAGFlowVectorStoreConfig(BaseVectorStoreConfig): - api_base in litellm_params - Default: http://localhost:9380 """ - api_base = ( - api_base or litellm_params.get("api_base") or get_secret_str("RAGFLOW_API_BASE") or "http://localhost:9380" + params_api_base: Final = litellm_params.get("api_base") + resolved_api_base: Final = ( + api_base + or (f"{params_api_base}" if params_api_base else None) + or get_secret_str("RAGFLOW_API_BASE") + or "http://localhost:9380" ) - # Remove trailing slashes - api_base = api_base.rstrip("/") - # RAGFlow datasets API endpoint - return f"{api_base}/api/v1/datasets" + return join_api_base(resolved_api_base, "api/v1/datasets") def transform_search_vector_store_request( self, diff --git a/litellm/llms/recraft/image_edit/transformation.py b/litellm/llms/recraft/image_edit/transformation.py index 46905580382..74bfa52c5e0 100644 --- a/litellm/llms/recraft/image_edit/transformation.py +++ b/litellm/llms/recraft/image_edit/transformation.py @@ -5,6 +5,7 @@ import httpx from httpx._types import RequestFiles from litellm.images.utils import ImageEditRequestUtils +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig from litellm.secret_managers.main import get_secret_str from litellm.types.images.main import ImageEditOptionalRequestParams @@ -68,11 +69,8 @@ class RecraftImageEditConfig(BaseImageEditConfig): Some providers need `model` in `api_base` """ - complete_url: str = api_base or get_secret_str("RECRAFT_API_BASE") or self.DEFAULT_BASE_URL - - complete_url = complete_url.rstrip("/") - complete_url = f"{complete_url}/{self.IMAGE_EDIT_ENDPOINT}" - return complete_url + base_url: Final = api_base or get_secret_str("RECRAFT_API_BASE") or self.DEFAULT_BASE_URL + return join_api_base(base_url, self.IMAGE_EDIT_ENDPOINT) def validate_environment( self, diff --git a/litellm/llms/recraft/image_generation/transformation.py b/litellm/llms/recraft/image_generation/transformation.py index 0fcfafb79fb..7b1a7714ec6 100644 --- a/litellm/llms/recraft/image_generation/transformation.py +++ b/litellm/llms/recraft/image_generation/transformation.py @@ -4,6 +4,7 @@ from typing import TYPE_CHECKING, Any, Final import httpx from pydantic import ConfigDict, TypeAdapter +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.image_generation.transformation import ( BaseImageGenerationConfig, ) @@ -72,11 +73,8 @@ class RecraftImageGenerationConfig(BaseImageGenerationConfig): Some providers need `model` in `api_base` """ - complete_url: str = api_base or get_secret_str("RECRAFT_API_BASE") or self.DEFAULT_BASE_URL - - complete_url = complete_url.rstrip("/") - complete_url = f"{complete_url}/{self.IMAGE_GENERATION_ENDPOINT}" - return complete_url + base_url: Final = api_base or get_secret_str("RECRAFT_API_BASE") or self.DEFAULT_BASE_URL + return join_api_base(base_url, self.IMAGE_GENERATION_ENDPOINT) def validate_environment( self, diff --git a/litellm/llms/runwayml/image_generation/transformation.py b/litellm/llms/runwayml/image_generation/transformation.py index ddc1b9474ac..e721ac678a1 100644 --- a/litellm/llms/runwayml/image_generation/transformation.py +++ b/litellm/llms/runwayml/image_generation/transformation.py @@ -11,6 +11,7 @@ from litellm.constants import ( RUNWAYML_DEFAULT_API_VERSION, RUNWAYML_POLLING_TIMEOUT, ) +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.image_generation.transformation import ( BaseImageGenerationConfig, ) @@ -62,12 +63,8 @@ class RunwayMLImageGenerationConfig(BaseImageGenerationConfig): Some providers need `model` in `api_base` """ - complete_url: str = api_base or get_secret_str("RUNWAYML_API_BASE") or self.DEFAULT_BASE_URL - - complete_url = complete_url.rstrip("/") - if self.IMAGE_GENERATION_ENDPOINT: - complete_url = f"{complete_url}/{self.IMAGE_GENERATION_ENDPOINT}" - return complete_url + base_url: Final = api_base or get_secret_str("RUNWAYML_API_BASE") or self.DEFAULT_BASE_URL + return join_api_base(base_url, self.IMAGE_GENERATION_ENDPOINT) def validate_environment( self, @@ -226,8 +223,7 @@ class RunwayMLImageGenerationConfig(BaseImageGenerationConfig): start_time: Final = time.time() # Build task status URL - api_base = api_base.rstrip("/") - task_url: Final = f"{api_base}/v1/tasks/{task_id}" + task_url: Final = join_api_base(api_base, f"v1/tasks/{task_id}") verbose_logger.debug("Polling RunwayML task: %s", task_url) @@ -275,8 +271,7 @@ class RunwayMLImageGenerationConfig(BaseImageGenerationConfig): start_time: Final = time.time() # Build task status URL - api_base = api_base.rstrip("/") - task_url: Final = f"{api_base}/v1/tasks/{task_id}" + task_url: Final = join_api_base(api_base, f"v1/tasks/{task_id}") verbose_logger.debug("Polling RunwayML task (async): %s", task_url) diff --git a/litellm/llms/runwayml/text_to_speech/transformation.py b/litellm/llms/runwayml/text_to_speech/transformation.py index d41ffea17dc..1f052648f13 100644 --- a/litellm/llms/runwayml/text_to_speech/transformation.py +++ b/litellm/llms/runwayml/text_to_speech/transformation.py @@ -18,6 +18,7 @@ from litellm.constants import ( RUNWAYML_DEFAULT_API_VERSION, RUNWAYML_POLLING_TIMEOUT, ) +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.text_to_speech.transformation import ( BaseTextToSpeechConfig, TextToSpeechRequestData, @@ -230,10 +231,8 @@ class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): """ Get the complete URL for RunwayML TTS request """ - complete_url = api_base or get_secret_str("RUNWAYML_API_BASE") or self.DEFAULT_BASE_URL - - complete_url = complete_url.rstrip("/") - return f"{complete_url}/{self.TTS_ENDPOINT_PATH}" + base_url: Final = api_base or get_secret_str("RUNWAYML_API_BASE") or self.DEFAULT_BASE_URL + return join_api_base(base_url, self.TTS_ENDPOINT_PATH) @staticmethod def _check_timeout(start_time: float, timeout_secs: float) -> None: @@ -311,8 +310,7 @@ class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): start_time: Final = time.time() # Build task status URL - api_base = api_base.rstrip("/") - task_url: Final = f"{api_base}/v1/tasks/{task_id}" + task_url: Final = join_api_base(api_base, f"v1/tasks/{task_id}") verbose_logger.debug("Polling RunwayML TTS task: %s", task_url) @@ -359,8 +357,7 @@ class RunwayMLTextToSpeechConfig(BaseTextToSpeechConfig): start_time: Final = time.time() # Build task status URL - api_base = api_base.rstrip("/") - task_url: Final = f"{api_base}/v1/tasks/{task_id}" + task_url: Final = join_api_base(api_base, f"v1/tasks/{task_id}") verbose_logger.debug("Polling RunwayML TTS task (async): %s", task_url) diff --git a/litellm/llms/runwayml/videos/transformation.py b/litellm/llms/runwayml/videos/transformation.py index 8b765b1390b..106567f6fba 100644 --- a/litellm/llms/runwayml/videos/transformation.py +++ b/litellm/llms/runwayml/videos/transformation.py @@ -9,7 +9,7 @@ from typing_extensions import ReadOnly, TypedDict import litellm from litellm.constants import RUNWAYML_DEFAULT_API_VERSION -from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.videos.transformation import BaseVideoConfig from litellm.llms.custom_httpx.http_handler import ( @@ -274,7 +274,7 @@ class RunwayMLVideoConfig(BaseVideoConfig): files_list: Final[RequestFiles] = [] - return request_data, files_list, f"{api_base}/{endpoint}" + return request_data, files_list, join_api_base(api_base, endpoint) def _select_generation_endpoint(self, request_data: Mapping[str, object]) -> str: if request_data.get("promptVideo") is not None or request_data.get("videoUri") is not None: @@ -413,7 +413,7 @@ class RunwayMLVideoConfig(BaseVideoConfig): encoded_video_id: Final = encode_url_path_segment(original_video_id, field_name="video_id") # Get task status to retrieve video URL - url: Final = f"{api_base}/tasks/{encoded_video_id}" + url: Final = join_api_base(api_base, f"tasks/{encoded_video_id}") return url, dict[str, str]() @@ -571,7 +571,7 @@ class RunwayMLVideoConfig(BaseVideoConfig): encoded_video_id: Final = encode_url_path_segment(original_video_id, field_name="video_id") # Construct the URL for task cancellation - url: Final = f"{api_base}/tasks/{encoded_video_id}/cancel" + url: Final = join_api_base(api_base, f"tasks/{encoded_video_id}/cancel") return url, dict[str, str]() @@ -608,7 +608,7 @@ class RunwayMLVideoConfig(BaseVideoConfig): encoded_video_id: Final = encode_url_path_segment(original_video_id, field_name="video_id") # Construct the full URL for task status retrieval - url: Final = f"{api_base}/tasks/{encoded_video_id}" + url: Final = join_api_base(api_base, f"tasks/{encoded_video_id}") # Empty dict for GET request (no body) return url, dict[str, str]() diff --git a/litellm/llms/s3_vectors/vector_stores/transformation.py b/litellm/llms/s3_vectors/vector_stores/transformation.py index 044315168d2..2b0de3dc9ea 100644 --- a/litellm/llms/s3_vectors/vector_stores/transformation.py +++ b/litellm/llms/s3_vectors/vector_stores/transformation.py @@ -3,6 +3,7 @@ from typing import TYPE_CHECKING, Any, Final import httpx +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.vector_store.transformation import ( BaseQueryEmbeddingVectorStoreConfig, VectorStoreEmbeddingExecutor, @@ -131,7 +132,7 @@ class S3VectorsVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig, BaseAWSLLM litellm_logging_obj: LiteLLMLoggingObj, ) -> tuple[str, dict[str, object]]: litellm_logging_obj.model_call_details["query"] = query_text - return f"{api_base}/QueryVectors", { + return join_api_base(api_base, "QueryVectors"), { "vectorBucketName": bucket_name, "indexName": index_name, "queryVector": {"float32": list(query_vector)}, diff --git a/litellm/llms/sambanova/embedding/transformation.py b/litellm/llms/sambanova/embedding/transformation.py index cf39c0a092f..93cbb13ac47 100644 --- a/litellm/llms/sambanova/embedding/transformation.py +++ b/litellm/llms/sambanova/embedding/transformation.py @@ -10,6 +10,7 @@ import httpx from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig from litellm.secret_managers.main import get_secret_str @@ -37,11 +38,7 @@ class SambaNovaEmbeddingConfig(BaseEmbeddingConfig): ) -> str: if api_base is None: raise ValueError("api_base is required for SambaNova embeddings") - # Remove trailing slashes and ensure clean base URL - api_base = api_base.rstrip("/") - if not api_base.endswith("/embeddings"): - api_base = f"{api_base}/embeddings" - return api_base + return ensure_api_route(api_base, "embeddings") def validate_environment( self, diff --git a/litellm/llms/sap/chat/transformation.py b/litellm/llms/sap/chat/transformation.py index a82dfe72322..8a7baf18e8f 100755 --- a/litellm/llms/sap/chat/transformation.py +++ b/litellm/llms/sap/chat/transformation.py @@ -11,6 +11,7 @@ import httpx from pydantic import BaseModel import litellm +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import ModelResponse @@ -185,12 +186,12 @@ class GenAIHubOrchestrationConfig(OpenAIGPTConfig): # Keep a short, tight client lifecycle here to avoid fd leaks client: Final = litellm.module_level_client # with httpx.Client(timeout=30) as client: - deployments: Final = client.get(f"{self.base_url}/lm/deployments", headers=self.headers).json() + deployments: Final = client.get(join_api_base(self.base_url, "lm/deployments"), headers=self.headers).json() valid: Final[list[tuple[str, str]]] = [] for dep in deployments.get("resources", []): if dep.get("scenarioId") == "orchestration": cfg = client.get( - f"{self.base_url}/lm/configurations/{dep['configurationId']}", + join_api_base(self.base_url, f"lm/configurations/{dep['configurationId']}"), headers=self.headers, ).json() if cfg.get("executableId") == "orchestration": @@ -263,7 +264,7 @@ class GenAIHubOrchestrationConfig(OpenAIGPTConfig): litellm_params: dict, stream: bool | None = None, ): - api_base_: Final = f"{self.deployment_url}/v2/completion" + api_base_: Final = join_api_base(self.deployment_url, "v2/completion") return api_base_ def _build_prompt_module( diff --git a/litellm/llms/sap/credentials.py b/litellm/llms/sap/credentials.py index b1618f38623..2622d0fa9b6 100644 --- a/litellm/llms/sap/credentials.py +++ b/litellm/llms/sap/credentials.py @@ -16,6 +16,7 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.custom_httpx.http_handler import HTTPHandler, get_httpx_client AUTH_ENDPOINT_SUFFIX: Final = "/oauth/token" @@ -112,7 +113,7 @@ CREDENTIAL_VALUES: Final[list[CredentialsValue]] = [ CredentialsValue( "base_url", ("serviceurls", "AI_API_URL"), - transform_fn=lambda url: url.rstrip("/") + ("" if url.endswith("/v2") else "/v2"), + transform_fn=lambda url: ensure_api_route(url, "v2"), ), CredentialsValue( "cert_url", diff --git a/litellm/llms/sap/embed/transformation.py b/litellm/llms/sap/embed/transformation.py index e91dfb91081..107fef0b355 100644 --- a/litellm/llms/sap/embed/transformation.py +++ b/litellm/llms/sap/embed/transformation.py @@ -8,6 +8,7 @@ from typing import Final, Literal import httpx from pydantic import Field +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.embedding.transformation import ( BaseEmbeddingConfig, LiteLLMLoggingObj, @@ -101,11 +102,11 @@ class GenAIHubEmbeddingConfig(BaseEmbeddingConfig): def deployment_url(self) -> str: with httpx.Client(timeout=30) as client: valid_deployments: Final = [] - deployments: Final = client.get(self.base_url + "/lm/deployments", headers=self.headers).json() + deployments: Final = client.get(join_api_base(self.base_url, "lm/deployments"), headers=self.headers).json() for deployment in deployments.get("resources", []): if deployment["scenarioId"] == "orchestration": config_details = client.get( - self.base_url + f"/lm/configurations/{deployment['configurationId']}", + join_api_base(self.base_url, f"lm/configurations/{deployment['configurationId']}"), headers=self.headers, ).json() if config_details["executableId"] == "orchestration": @@ -144,7 +145,7 @@ class GenAIHubEmbeddingConfig(BaseEmbeddingConfig): litellm_params: dict, stream: bool | None = None, ) -> str: - url: Final = self.deployment_url.rstrip("/") + "/v2/embeddings" + url: Final = join_api_base(self.deployment_url, "v2/embeddings") return url def transform_embedding_request( diff --git a/litellm/llms/scaledown/chat/transformation.py b/litellm/llms/scaledown/chat/transformation.py index 44eb4785615..3224bc97148 100644 --- a/litellm/llms/scaledown/chat/transformation.py +++ b/litellm/llms/scaledown/chat/transformation.py @@ -9,6 +9,7 @@ from urllib.parse import unquote import httpx from pydantic import JsonValue, TypeAdapter +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues @@ -103,8 +104,8 @@ class ScaleDownChatConfig(BaseConfig): base: Final = _root(api_base or _trusted_root()) operation: Final = _operation(model) if operation in DECISIONS_MODELS: - return f"{base}/v1/scaledown" - return f"{base}{NATIVE_PATHS[operation]}" + return join_api_base(base, "v1/scaledown") + return join_api_base(base, NATIVE_PATHS[operation]) def get_supported_openai_params(self, model: str) -> list[str]: # mutable-ok: BaseConfig signature operation: Final = _operation(model) diff --git a/litellm/llms/scaleway/audio_transcription/transformation.py b/litellm/llms/scaleway/audio_transcription/transformation.py index 7ebf4475355..8913cc362b3 100644 --- a/litellm/llms/scaleway/audio_transcription/transformation.py +++ b/litellm/llms/scaleway/audio_transcription/transformation.py @@ -9,6 +9,7 @@ from typing import Final import httpx from litellm.litellm_core_utils.audio_utils.utils import process_audio_file +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.audio_transcription.transformation import ( AudioTranscriptionRequestData, BaseAudioTranscriptionConfig, @@ -58,8 +59,7 @@ class ScalewayAudioTranscriptionConfig(BaseAudioTranscriptionConfig): litellm_params: dict, stream: bool | None = None, ) -> str: - api_base = "https://api.scaleway.ai/v1" if api_base is None else api_base.rstrip("/") - return f"{api_base}/audio/transcriptions" + return join_api_base("https://api.scaleway.ai/v1" if api_base is None else api_base, "audio/transcriptions") def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException: return ScalewayAudioTranscriptionException( diff --git a/litellm/llms/scaleway/rerank/transformation.py b/litellm/llms/scaleway/rerank/transformation.py index 921252d7091..1de5d71f17a 100644 --- a/litellm/llms/scaleway/rerank/transformation.py +++ b/litellm/llms/scaleway/rerank/transformation.py @@ -9,6 +9,7 @@ API reference: https://www.scaleway.com/en/developers/api/generative-apis/#path- from collections.abc import Mapping from typing import Final +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.jina_ai.rerank.transformation import JinaAIRerankConfig from litellm.secret_managers.main import get_secret_str @@ -25,8 +26,8 @@ class ScalewayRerankConfig(JinaAIRerankConfig): model: str, optional_params: Mapping[str, object] | None = None, ) -> str: - base: Final = SCALEWAY_API_BASE if api_base is None else api_base.rstrip("/") - return f"{base}/rerank" + base: Final = SCALEWAY_API_BASE if api_base is None else api_base + return join_api_base(base, "rerank") def validate_environment( self, diff --git a/litellm/llms/searxng/search/transformation.py b/litellm/llms/searxng/search/transformation.py index 265a3db72f5..b1753821cd2 100644 --- a/litellm/llms/searxng/search/transformation.py +++ b/litellm/llms/searxng/search/transformation.py @@ -11,6 +11,7 @@ import httpx from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm.search.transformation import ( BaseSearchConfig, SearchResponse, @@ -101,12 +102,7 @@ class SearXNGSearchConfig(BaseSearchConfig): "or pass `api_base` parameter. Example: os.environ['SEARXNG_API_BASE'] = 'https://your-searxng-instance.com'" ) - # Append "/search" to the api base if it's not already there - if not api_base.endswith("/search"): - if api_base.endswith("/"): - api_base = f"{api_base}search" - else: - api_base = f"{api_base}/search" + api_base = ensure_api_route(api_base, "search") # Build query parameters from the transformed request body if data and isinstance(data, dict) and "_searxng_params" in data: diff --git a/litellm/llms/serper/search/transformation.py b/litellm/llms/serper/search/transformation.py index 2539610ea59..e82427e4d0c 100644 --- a/litellm/llms/serper/search/transformation.py +++ b/litellm/llms/serper/search/transformation.py @@ -11,6 +11,7 @@ import httpx from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm.search.transformation import ( BaseSearchConfig, SearchResponse, @@ -84,12 +85,7 @@ class SerperSearchConfig(BaseSearchConfig): Get complete URL for Search endpoint. """ api_base = api_base or get_secret_str("SERPER_API_BASE") or self.SERPER_API_BASE - api_base = api_base.rstrip("/") - - if not api_base.endswith("/search"): - api_base = f"{api_base}/search" - - return api_base + return ensure_api_route(api_base, "search") def transform_search_request( self, diff --git a/litellm/llms/snowflake/chat/transformation.py b/litellm/llms/snowflake/chat/transformation.py index 6a480904a3e..030413030eb 100644 --- a/litellm/llms/snowflake/chat/transformation.py +++ b/litellm/llms/snowflake/chat/transformation.py @@ -25,6 +25,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import ( select_anthropic_content_block_type_for_file, ) from litellm.litellm_core_utils.prompt_templates.image_handling import async_inline_remote_media +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.anthropic.chat.handler import ModelResponseIterator as AnthropicStreamParser from litellm.llms.anthropic.chat.transformation import AnthropicConfig from litellm.llms.anthropic.common_utils import normalize_cache_control_in_anthropic_payload @@ -239,8 +240,8 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig): ) -> str: api_base = self._get_api_base(api_base, optional_params) if is_snowflake_claude_model(model): - return f"{api_base}/cortex/v1/messages" - return f"{api_base}/cortex/v1/chat/completions" + return join_api_base(api_base, "cortex/v1/messages") + return join_api_base(api_base, "cortex/v1/chat/completions") def validate_environment( self, diff --git a/litellm/llms/snowflake/embedding/transformation.py b/litellm/llms/snowflake/embedding/transformation.py index ec482337e52..45bec1283a3 100644 --- a/litellm/llms/snowflake/embedding/transformation.py +++ b/litellm/llms/snowflake/embedding/transformation.py @@ -5,6 +5,7 @@ import httpx from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig from litellm.types.llms.openai import AllEmbeddingInputValues @@ -31,7 +32,7 @@ class SnowflakeEmbeddingConfig(SnowflakeBaseConfig, BaseEmbeddingConfig): ) -> str: api_base = self._get_api_base(api_base, optional_params) - return f"{api_base}/cortex/inference:embed" + return join_api_base(api_base, "cortex/inference:embed") def transform_embedding_request( self, diff --git a/litellm/llms/snowflake/utils.py b/litellm/llms/snowflake/utils.py index a8c14d1d7f3..7cc88ab66e1 100644 --- a/litellm/llms/snowflake/utils.py +++ b/litellm/llms/snowflake/utils.py @@ -1,6 +1,7 @@ import re from typing import TYPE_CHECKING, Any, Final +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues @@ -53,7 +54,7 @@ class SnowflakeBaseConfig: optional_params[param] = value return optional_params - def _get_api_base(self, api_base, optional_params): + def _get_api_base(self, api_base: str | None, optional_params) -> str: if not api_base: if "account_id" in optional_params: account_id = optional_params.pop("account_id") @@ -65,10 +66,7 @@ class SnowflakeBaseConfig: raise ValueError("Invalid account_id format") api_base = f"https://{account_id}.snowflakecomputing.com/api/v2" - api_base = api_base.rstrip("/") - if not api_base.endswith("/api/v2"): - api_base += "/api/v2" - return api_base + return ensure_api_route(api_base, "api/v2") def validate_environment( self, diff --git a/litellm/llms/soniox/audio_transcription/handler.py b/litellm/llms/soniox/audio_transcription/handler.py index f86bc142898..a48e3e3ed4c 100644 --- a/litellm/llms/soniox/audio_transcription/handler.py +++ b/litellm/llms/soniox/audio_transcription/handler.py @@ -31,6 +31,7 @@ from litellm.litellm_core_utils.audio_utils.utils import ( get_audio_file_name, process_audio_file, ) +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, HTTPHandler, @@ -329,7 +330,7 @@ class SonioxAudioTranscriptionHandler: input=None, api_key=api_key, additional_args={ - "api_base": f"{api_base}/v1/transcriptions", + "api_base": join_api_base(api_base, "v1/transcriptions"), "atranscription": True, "complete_input_dict": SonioxAudioTranscriptionHandler._redact_body_for_logging(body), }, @@ -451,7 +452,7 @@ class SonioxAudioTranscriptionHandler: self._safe_log_pre_call(logging_obj, api_key, base_url, body) create_resp: Final = http_client.post( - url=f"{base_url}/v1/transcriptions", + url=join_api_base(base_url, "v1/transcriptions"), headers=auth_headers, json=body, timeout=timeout, @@ -472,7 +473,7 @@ class SonioxAudioTranscriptionHandler: ) transcript_resp: Final = http_client.get( - url=f"{base_url}/v1/transcriptions/{transcription_id}/transcript", + url=join_api_base(base_url, f"v1/transcriptions/{transcription_id}/transcript"), headers=auth_headers, timeout=timeout, ) @@ -528,7 +529,7 @@ class SonioxAudioTranscriptionHandler: # `Authorization` header is fine; httpx sets multipart Content-Type. upload_headers: Final = {"Authorization": auth_headers["Authorization"]} resp: Final = http_client.post( - url=f"{base_url}/v1/files", + url=join_api_base(base_url, "v1/files"), headers=upload_headers, files=files, timeout=timeout, @@ -550,7 +551,7 @@ class SonioxAudioTranscriptionHandler: ) -> _TranscriptionMeta: for _ in range(max_attempts): resp = http_client.get( - url=f"{base_url}/v1/transcriptions/{transcription_id}", + url=join_api_base(base_url, f"v1/transcriptions/{transcription_id}"), headers=auth_headers, timeout=timeout, ) @@ -593,7 +594,7 @@ class SonioxAudioTranscriptionHandler: if "transcription" in cleanup and transcription_id: try: http_client.delete( - url=f"{base_url}/v1/transcriptions/{transcription_id}", + url=join_api_base(base_url, f"v1/transcriptions/{transcription_id}"), headers=auth_headers, timeout=timeout, ) @@ -605,7 +606,7 @@ class SonioxAudioTranscriptionHandler: if "file" in cleanup and file_id_to_cleanup: try: http_client.delete( - url=f"{base_url}/v1/files/{file_id_to_cleanup}", + url=join_api_base(base_url, f"v1/files/{file_id_to_cleanup}"), headers=auth_headers, timeout=timeout, ) @@ -685,7 +686,7 @@ class SonioxAudioTranscriptionHandler: self._safe_log_pre_call(logging_obj, api_key, base_url, body) create_resp: Final = await http_client.post( - url=f"{base_url}/v1/transcriptions", + url=join_api_base(base_url, "v1/transcriptions"), headers=auth_headers, json=body, timeout=timeout, @@ -706,7 +707,7 @@ class SonioxAudioTranscriptionHandler: ) transcript_resp: Final = await http_client.get( - url=f"{base_url}/v1/transcriptions/{transcription_id}/transcript", + url=join_api_base(base_url, f"v1/transcriptions/{transcription_id}/transcript"), headers=auth_headers, timeout=timeout, ) @@ -761,7 +762,7 @@ class SonioxAudioTranscriptionHandler: } upload_headers: Final = {"Authorization": auth_headers["Authorization"]} resp: Final = await http_client.post( - url=f"{base_url}/v1/files", + url=join_api_base(base_url, "v1/files"), headers=upload_headers, files=files, timeout=timeout, @@ -783,7 +784,7 @@ class SonioxAudioTranscriptionHandler: ) -> _TranscriptionMeta: for _ in range(max_attempts): resp = await http_client.get( - url=f"{base_url}/v1/transcriptions/{transcription_id}", + url=join_api_base(base_url, f"v1/transcriptions/{transcription_id}"), headers=auth_headers, timeout=timeout, ) @@ -826,7 +827,7 @@ class SonioxAudioTranscriptionHandler: if "transcription" in cleanup and transcription_id: try: await http_client.delete( - url=f"{base_url}/v1/transcriptions/{transcription_id}", + url=join_api_base(base_url, f"v1/transcriptions/{transcription_id}"), headers=auth_headers, timeout=timeout, ) @@ -838,7 +839,7 @@ class SonioxAudioTranscriptionHandler: if "file" in cleanup and file_id_to_cleanup: try: await http_client.delete( - url=f"{base_url}/v1/files/{file_id_to_cleanup}", + url=join_api_base(base_url, f"v1/files/{file_id_to_cleanup}"), headers=auth_headers, timeout=timeout, ) diff --git a/litellm/llms/stability/image_generation/transformation.py b/litellm/llms/stability/image_generation/transformation.py index 1856523ce49..a53fd087098 100644 --- a/litellm/llms/stability/image_generation/transformation.py +++ b/litellm/llms/stability/image_generation/transformation.py @@ -12,6 +12,7 @@ from typing import TYPE_CHECKING, Any, Final import httpx from pydantic import ConfigDict, TypeAdapter +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.image_generation.transformation import ( BaseImageGenerationConfig, ) @@ -131,11 +132,10 @@ class StabilityImageGenerationConfig(BaseImageGenerationConfig): """ Get the complete URL for the Stability AI API request. """ - base_url: str = api_base or get_secret_str("STABILITY_API_BASE") or self.DEFAULT_BASE_URL - base_url = base_url.rstrip("/") + base_url: Final = api_base or get_secret_str("STABILITY_API_BASE") or self.DEFAULT_BASE_URL endpoint: Final = self._get_model_endpoint(model) - return f"{base_url}{endpoint}" + return join_api_base(base_url, endpoint) def validate_environment( self, diff --git a/litellm/llms/tavily/search/transformation.py b/litellm/llms/tavily/search/transformation.py index d0531e4c041..b46f974e0db 100644 --- a/litellm/llms/tavily/search/transformation.py +++ b/litellm/llms/tavily/search/transformation.py @@ -11,6 +11,7 @@ import httpx from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm.search.transformation import ( BaseSearchConfig, SearchResponse, @@ -90,11 +91,7 @@ class TavilySearchConfig(BaseSearchConfig): """ Get complete URL for Search endpoint. """ - resolved_api_base: Final = api_base or get_secret_str("TAVILY_API_BASE") or self.TAVILY_API_BASE - clean_api_base: Final = resolved_api_base.rstrip("/") - if clean_api_base.endswith("/search"): - return clean_api_base - return f"{clean_api_base}/search" + return ensure_api_route(api_base or get_secret_str("TAVILY_API_BASE") or self.TAVILY_API_BASE, "search") def transform_search_request( self, diff --git a/litellm/llms/tencent/chat/transformation.py b/litellm/llms/tencent/chat/transformation.py index 7e80b0012df..1f0bcd218d6 100644 --- a/litellm/llms/tencent/chat/transformation.py +++ b/litellm/llms/tencent/chat/transformation.py @@ -9,6 +9,7 @@ from typing import Final, TypedDict from typing_extensions import ReadOnly import litellm +from litellm.litellm_core_utils.url_utils import ensure_api_route, join_api_base from litellm.secret_managers.main import get_secret_str from litellm.utils import supports_reasoning @@ -137,7 +138,4 @@ class TencentChatConfig(OpenAIGPTConfig): if api_base.endswith("/chat/completions"): return api_base - if not api_base.endswith("/v1"): - api_base = f"{api_base}/v1" - - return f"{api_base}/chat/completions" + return join_api_base(ensure_api_route(api_base, "v1"), "chat/completions") diff --git a/litellm/llms/tencent/messages/transformation.py b/litellm/llms/tencent/messages/transformation.py index c56ecaeb51c..1e52bc37d47 100644 --- a/litellm/llms/tencent/messages/transformation.py +++ b/litellm/llms/tencent/messages/transformation.py @@ -8,6 +8,7 @@ alongside its standard OpenAI-compatible chat completions endpoint. from typing import Any import litellm +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.anthropic.pass_through.messages.transformation import ( AnthropicMessagesConfig, ) @@ -82,4 +83,4 @@ class TencentAnthropicMessagesConfig(AnthropicMessagesConfig): elif base_url.endswith("/v1"): base_url = base_url[: -len("/v1")] - return f"{base_url}/v1/messages" + return join_api_base(base_url, "v1/messages") diff --git a/litellm/llms/topaz/image_variations/transformation.py b/litellm/llms/topaz/image_variations/transformation.py index 8d03b73e94e..b3dded36f49 100644 --- a/litellm/llms/topaz/image_variations/transformation.py +++ b/litellm/llms/topaz/image_variations/transformation.py @@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Final from aiohttp import ClientResponse from httpx import Headers, Response +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import ( BaseLLMException, LiteLLMLoggingObj, @@ -39,8 +40,7 @@ class TopazImageVariationConfig(TopazModelInfo, BaseImageVariationConfig): litellm_params: dict, stream: bool | None = None, ) -> str: - clean_api_base: Final = (api_base or "https://api.topazlabs.com").rstrip("/") - return f"{clean_api_base}/image/v1/enhance" + return join_api_base(api_base or "https://api.topazlabs.com", "image/v1/enhance") def map_openai_params( self, diff --git a/litellm/llms/vercel_ai_gateway/chat/transformation.py b/litellm/llms/vercel_ai_gateway/chat/transformation.py index 19c4c2659ca..3fc74b6acc1 100644 --- a/litellm/llms/vercel_ai_gateway/chat/transformation.py +++ b/litellm/llms/vercel_ai_gateway/chat/transformation.py @@ -11,6 +11,7 @@ from typing import Final import httpx import litellm +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues @@ -92,7 +93,7 @@ class VercelAIGatewayConfig(OpenAIGPTConfig): if api_base is None: api_base = "https://ai-gateway.vercel.sh/v1" - models_url: Final = f"{api_base}/models" + models_url: Final = join_api_base(api_base, "models") response: Final = litellm.module_level_client.get(url=models_url) if response.status_code != 200: diff --git a/litellm/llms/vercel_ai_gateway/embedding/transformation.py b/litellm/llms/vercel_ai_gateway/embedding/transformation.py index 9cfaaab89f0..01d7b7766c0 100644 --- a/litellm/llms/vercel_ai_gateway/embedding/transformation.py +++ b/litellm/llms/vercel_ai_gateway/embedding/transformation.py @@ -12,6 +12,7 @@ from typing import TYPE_CHECKING, Any, Final import httpx +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig from litellm.secret_managers.main import get_secret_str @@ -77,12 +78,8 @@ class VercelAIGatewayEmbeddingConfig(BaseEmbeddingConfig): """ Get the complete URL for Vercel AI Gateway Embedding API endpoint. """ - if api_base: - api_base = api_base.rstrip("/") - else: - api_base = get_secret_str("VERCEL_AI_GATEWAY_API_BASE") or "https://ai-gateway.vercel.sh/v1" - - return f"{api_base}/embeddings" + base: Final = api_base or get_secret_str("VERCEL_AI_GATEWAY_API_BASE") or "https://ai-gateway.vercel.sh/v1" + return join_api_base(base, "embeddings") def transform_embedding_request( self, diff --git a/litellm/llms/vertex_ai/audio_transcription/gemini_transcribe_transformation.py b/litellm/llms/vertex_ai/audio_transcription/gemini_transcribe_transformation.py index f4db5eb110c..39d2e9b4c49 100644 --- a/litellm/llms/vertex_ai/audio_transcription/gemini_transcribe_transformation.py +++ b/litellm/llms/vertex_ai/audio_transcription/gemini_transcribe_transformation.py @@ -10,6 +10,7 @@ from litellm.litellm_core_utils.audio_utils.utils import ( normalize_transcription_language_to_bcp47, process_audio_file, ) +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.audio_transcription.transformation import ( AudioTranscriptionRequestData, BaseAudioTranscriptionConfig, @@ -129,10 +130,10 @@ class VertexGeminiAudioTranscriptionConfig(BaseAudioTranscriptionConfig, VertexB project_id: Final = validate_vertex_transcription_project_id( self.safe_get_vertex_ai_project(vertex_params) or self._resolve_project_id_from_credentials(vertex_params) ) - base_url: Final = (api_base or get_vertex_base_url(location)).rstrip("/") + base_url: Final = api_base or get_vertex_base_url(location) bare_model: Final = model.removeprefix("vertex_ai/") model_path: Final = f"projects/{project_id}/locations/{location}/publishers/google/models/{bare_model}" - return f"{base_url}/v1/{model_path}:generateContent" + return join_api_base(base_url, f"v1/{model_path}:generateContent") def _resolve_project_id_from_credentials(self, litellm_params: Mapping[str, object]) -> str: vertex_params: Final = dict(litellm_params) diff --git a/litellm/llms/vertex_ai/audio_transcription/transformation.py b/litellm/llms/vertex_ai/audio_transcription/transformation.py index 5603be81095..4fdf93a6307 100644 --- a/litellm/llms/vertex_ai/audio_transcription/transformation.py +++ b/litellm/llms/vertex_ai/audio_transcription/transformation.py @@ -9,6 +9,7 @@ from litellm.litellm_core_utils.audio_utils.utils import ( normalize_transcription_language_to_bcp47, process_audio_file, ) +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.audio_transcription.transformation import ( AudioTranscriptionRequestData, BaseAudioTranscriptionConfig, @@ -126,8 +127,8 @@ class VertexAIAudioTranscriptionConfig(BaseAudioTranscriptionConfig, VertexBase) project_id: Final = validate_vertex_transcription_project_id( self.safe_get_vertex_ai_project(litellm_params) or self._resolve_project_id_from_credentials(litellm_params) ) - base_url: Final = (api_base or f"https://{speech_to_text_host(location)}").rstrip("/") - return f"{base_url}/v2/projects/{project_id}/locations/{location}/recognizers/_:recognize" + base_url: Final = api_base or f"https://{speech_to_text_host(location)}" + return join_api_base(base_url, f"v2/projects/{project_id}/locations/{location}/recognizers/_:recognize") def _resolve_project_id_from_credentials(self, litellm_params: dict) -> str: _, project_id = self._ensure_access_token( diff --git a/litellm/llms/vertex_ai/batches/handler.py b/litellm/llms/vertex_ai/batches/handler.py index de5fe39cfc9..7c4238e3498 100644 --- a/litellm/llms/vertex_ai/batches/handler.py +++ b/litellm/llms/vertex_ai/batches/handler.py @@ -10,6 +10,7 @@ import litellm from litellm.litellm_core_utils.url_utils import ( async_safe_get, encode_url_path_segment, + join_api_base, safe_get, ) from litellm.llms.custom_httpx.http_handler import ( @@ -194,7 +195,7 @@ class VertexAIBatchPrediction(VertexLLM): api_base_path: Final = urlparse(api_base).path.rstrip("/") if api_base_path in ("/v1", "/v1beta1"): return graft_default_vertex_path(api_base=api_base, default_url=default_endpoint_url) - return api_base.rstrip("/") + urlparse(default_endpoint_url).path + return join_api_base(api_base, urlparse(default_endpoint_url).path) def _resolve_fine_tuned_endpoint_model( self, diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index d1e3f10fdd4..e36aba00293 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -37,6 +37,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( extract_file_data, extract_file_metadata, ) +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.base_utils import map_developer_role_to_system_role from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.files.batch_records import responses_batch_body_to_chat_body @@ -1083,12 +1084,12 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): if object_prefix: object_name = f"{object_prefix}/{object_name}" encoded_object_name: Final = encode_gcs_object_name_for_url(object_name) - endpoint: Final = f"upload/storage/v1/b/{bucket_name}/o?uploadType=media&name={encoded_object_name}" + upload_path: Final = f"upload/storage/v1/b/{bucket_name}/o" api_base = api_base or "https://storage.googleapis.com" if not api_base: raise ValueError("api_base is required") - return f"{api_base}/{endpoint}" + return f"{join_api_base(api_base, upload_path)}?uploadType=media&name={encoded_object_name}" def get_supported_openai_params(self, model: str) -> list[OpenAICreateFileRequestOptionalParams]: return [] diff --git a/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py b/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py index f916d494aad..085ff715c30 100644 --- a/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py +++ b/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py @@ -12,6 +12,7 @@ from pydantic import ConfigDict, TypeAdapter import litellm from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig from litellm.llms.vertex_ai.common_utils import get_vertex_base_url from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexLLM @@ -143,12 +144,12 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM): if model.startswith("vertex_ai/"): model_name = model.replace("vertex_ai/", "") - if api_base: - base_url = api_base.rstrip("/") - else: - base_url = get_vertex_base_url(vertex_location) + base_url: Final = api_base or get_vertex_base_url(vertex_location) - return f"{base_url}/v1/projects/{vertex_project}/locations/{vertex_location}/publishers/google/models/{model_name}:predict" + return join_api_base( + base_url, + f"v1/projects/{vertex_project}/locations/{vertex_location}/publishers/google/models/{model_name}:predict", + ) def transform_image_edit_request( self, diff --git a/litellm/llms/vertex_ai/interactions/transformation.py b/litellm/llms/vertex_ai/interactions/transformation.py index 36965fe666f..b8e13059917 100644 --- a/litellm/llms/vertex_ai/interactions/transformation.py +++ b/litellm/llms/vertex_ai/interactions/transformation.py @@ -2,7 +2,7 @@ from collections.abc import Callable, Mapping from dataclasses import dataclass from typing import Final -from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.llms.gemini.interactions.transformation import GoogleAIStudioInteractionsConfig from litellm.llms.vertex_ai.common_utils import validate_vertex_location from litellm.llms.vertex_ai.vertex_llm_base import VertexBase @@ -22,9 +22,9 @@ class VertexInteractionsTarget: @property def collection_url(self) -> str: - return ( - f"{self.base_url}/{VERTEX_INTERACTIONS_API_VERSION}" - f"/projects/{self.project_id}/locations/{self.location}/interactions" + return join_api_base( + self.base_url, + f"{VERTEX_INTERACTIONS_API_VERSION}/projects/{self.project_id}/locations/{self.location}/interactions", ) def interaction_url(self, interaction_id: str) -> str: diff --git a/litellm/llms/vertex_ai/realtime/transformation.py b/litellm/llms/vertex_ai/realtime/transformation.py index 55f1cf58cbe..13f547b3b92 100644 --- a/litellm/llms/vertex_ai/realtime/transformation.py +++ b/litellm/llms/vertex_ai/realtime/transformation.py @@ -21,6 +21,7 @@ from collections.abc import Awaitable, Callable from typing import Final from litellm import verbose_logger +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.gemini.realtime.transformation import GeminiRealtimeConfig from litellm.llms.vertex_ai.audio_transcription.realtime_transformation import ( VertexChirpRealtimeConfig, @@ -68,9 +69,8 @@ class VertexAIRealtimeConfig(GeminiRealtimeConfig): """ if api_base: # Allow callers to supply a fully-qualified wss:// base URL. - base = api_base.rstrip("/") - base = base.replace("https://", "wss://").replace("http://", "ws://") - return f"{base}/ws/google.cloud.aiplatform.v1.LlmBidiService/BidiGenerateContent" + base: Final = api_base.replace("https://", "wss://").replace("http://", "ws://") + return join_api_base(base, "ws/google.cloud.aiplatform.v1.LlmBidiService/BidiGenerateContent") host: Final = get_vertex_base_url(self._location).removeprefix("https://") return f"wss://{host}/ws/google.cloud.aiplatform.v1.LlmBidiService/BidiGenerateContent" diff --git a/litellm/llms/vertex_ai/text_to_speech/transformation.py b/litellm/llms/vertex_ai/text_to_speech/transformation.py index 1441d01e41e..c0082bc5ca6 100644 --- a/litellm/llms/vertex_ai/text_to_speech/transformation.py +++ b/litellm/llms/vertex_ai/text_to_speech/transformation.py @@ -23,7 +23,7 @@ from litellm.litellm_core_utils.audio_utils.utils import ( speech_media_type_from_audio_bytes, ) from litellm.litellm_core_utils.core_helpers import as_str_mapping, drops_unsupported_params -from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.text_to_speech.transformation import ( BaseTextToSpeechConfig, @@ -1075,13 +1075,14 @@ class VertexAILyriaTextToSpeechConfig(VertexAITextToSpeechConfig): }, ) location: Final = self.safe_get_vertex_ai_location(litellm_params) or self.get_default_vertex_location() - base_url: Final = self.get_api_base(api_base=api_base, vertex_location=location).rstrip("/") + base_url: Final = self.get_api_base(api_base=api_base, vertex_location=location) encoded_project: Final = encode_url_path_segment(project, field_name="project") encoded_location: Final = encode_url_path_segment(location, field_name="location") encoded_model: Final = encode_url_path_segment(base_model, field_name="model") - return ( - f"{base_url}/v1/projects/{encoded_project}/locations/{encoded_location}" - f"/publishers/google/models/{encoded_model}:predict" + return join_api_base( + base_url, + f"v1/projects/{encoded_project}/locations/{encoded_location}" + f"/publishers/google/models/{encoded_model}:predict", ) def transform_text_to_speech_request( diff --git a/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py b/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py index 5c250fc1a7e..440604688bb 100644 --- a/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py +++ b/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py @@ -4,6 +4,7 @@ from typing import TYPE_CHECKING, Any, Final, Protocol import httpx from typing_extensions import ReadOnly, TypedDict +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig from litellm.llms.vertex_ai.common_utils import get_vertex_base_url from litellm.llms.vertex_ai.vertex_llm_base import VertexBase @@ -288,7 +289,7 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase): """ Transform create request for Vertex AI RAG Corpus """ - url: Final = f"{api_base}/ragCorpora" # Base URL for creating RAG corpus + url: Final = join_api_base(api_base, "ragCorpora") # Base URL for creating RAG corpus # Add metadata if provided metadata: Final = vector_store_create_optional_params.get("metadata") diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py index 7100c9efde3..4be41e26326 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/count_tokens/handler.py @@ -8,6 +8,7 @@ their respective publisher-specific count-tokens endpoints. from typing import Any, Final +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.llms.vertex_ai.common_utils import get_vertex_base_url from litellm.llms.vertex_ai.vertex_llm_base import VertexBase @@ -59,16 +60,14 @@ class VertexAIPartnerModelsTokenCounter(VertexBase): publisher: Final = self._get_publisher_for_model(model) # Use custom api_base if provided, otherwise construct default - if api_base: - base_url = api_base - else: - base_url = get_vertex_base_url(vertex_location) + base_url: Final = api_base or get_vertex_base_url(vertex_location) # Construct the count-tokens endpoint # Format: /v1/projects/{project}/locations/{location}/publishers/{publisher}/models/count-tokens:rawPredict - endpoint: Final = ( - f"{base_url}/v1/projects/{project_id}/locations/{vertex_location}/" - f"publishers/{publisher}/models/count-tokens:rawPredict" + endpoint: Final = join_api_base( + base_url, + f"v1/projects/{project_id}/locations/{vertex_location}/" + f"publishers/{publisher}/models/count-tokens:rawPredict", ) return endpoint diff --git a/litellm/llms/vertex_ai/vertex_llm_base.py b/litellm/llms/vertex_ai/vertex_llm_base.py index 77508096e4a..6cff9729189 100644 --- a/litellm/llms/vertex_ai/vertex_llm_base.py +++ b/litellm/llms/vertex_ai/vertex_llm_base.py @@ -14,6 +14,7 @@ from urllib.parse import urlparse import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.asyncify import asyncify +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.secret_managers.main import get_secret_str from litellm.types.llms.vertex_ai import VERTEX_CREDENTIALS_TYPES, VertexPartnerProvider @@ -295,23 +296,22 @@ class VertexBase: if api_base is None: api_base = get_vertex_base_url(vertex_location) + location_path: Final = f"projects/{vertex_project}/locations/{vertex_location}" + raw_predict_method: Final = "streamRawPredict" if stream else "rawPredict" if partner == VertexPartnerProvider.llama: - return f"{api_base}/v1/projects/{vertex_project}/locations/{vertex_location}/endpoints/openapi/chat/completions" + return join_api_base(api_base, f"v1/{location_path}/endpoints/openapi/chat/completions") elif partner == VertexPartnerProvider.mistralai: - if stream: - return f"{api_base}/v1/projects/{vertex_project}/locations/{vertex_location}/publishers/mistralai/models/{model}:streamRawPredict" - else: - return f"{api_base}/v1/projects/{vertex_project}/locations/{vertex_location}/publishers/mistralai/models/{model}:rawPredict" + return join_api_base( + api_base, f"v1/{location_path}/publishers/mistralai/models/{model}:{raw_predict_method}" + ) elif partner == VertexPartnerProvider.ai21: - if stream: - return f"{api_base}/v1beta1/projects/{vertex_project}/locations/{vertex_location}/publishers/ai21/models/{model}:streamRawPredict" - else: - return f"{api_base}/v1beta1/projects/{vertex_project}/locations/{vertex_location}/publishers/ai21/models/{model}:rawPredict" + return join_api_base( + api_base, f"v1beta1/{location_path}/publishers/ai21/models/{model}:{raw_predict_method}" + ) elif partner == VertexPartnerProvider.claude: - if stream: - return f"{api_base}/v1/projects/{vertex_project}/locations/{vertex_location}/publishers/anthropic/models/{model}:streamRawPredict" - else: - return f"{api_base}/v1/projects/{vertex_project}/locations/{vertex_location}/publishers/anthropic/models/{model}:rawPredict" + return join_api_base( + api_base, f"v1/{location_path}/publishers/anthropic/models/{model}:{raw_predict_method}" + ) def get_complete_vertex_url( self, @@ -658,7 +658,7 @@ class VertexBase: # For Gemini (Google AI Studio), construct the full path like other providers if model is None: raise ValueError("Model parameter is required for Gemini custom API base URLs") - url = f"{api_base}/models/{model}:{endpoint}" + url = join_api_base(api_base, f"models/{model}:{endpoint}") if gemini_api_key is None: raise ValueError( "Missing Gemini API key. Set the GEMINI_API_KEY or GOOGLE_API_KEY environment variable." @@ -678,20 +678,17 @@ class VertexBase: model_for_url: Final = get_vertex_base_model_name(model=model) # Format: {api_base}/v1/projects/{project}/locations/{location}/endpoints/{model}:{endpoint} version: Final = vertex_api_version or "v1" - url = "{}/{}/projects/{}/locations/{}/endpoints/{}:{}".format( - api_base.rstrip("/"), - version, - vertex_project, - vertex_location, - model_for_url, - endpoint, + url = join_api_base( + api_base, + f"{version}/projects/{vertex_project}/locations/{vertex_location}" + f"/endpoints/{model_for_url}:{endpoint}", ) elif urlparse(api_base).path in ("", "/"): - url = api_base.rstrip("/") + urlparse(url).path + url = join_api_base(api_base, urlparse(url).path) elif urlparse(api_base).path.rstrip("/") in ("/v1", "/v1beta1") and "/projects/" in urlparse(url).path: url = graft_default_vertex_path(api_base=api_base, default_url=url) else: - url = f"{api_base}:{endpoint}" + url = f"{api_base.rstrip('/')}:{endpoint}" if stream is True: parsed_stream_url: Final = urlparse(url) stream_query: Final = f"{parsed_stream_url.query}&alt=sse" if parsed_stream_url.query else "alt=sse" diff --git a/litellm/llms/vertex_ai/videos/transformation.py b/litellm/llms/vertex_ai/videos/transformation.py index de368cea5ac..080899d08b0 100644 --- a/litellm/llms/vertex_ai/videos/transformation.py +++ b/litellm/llms/vertex_ai/videos/transformation.py @@ -18,6 +18,7 @@ from typing_extensions import ReadOnly import litellm from litellm.constants import DEFAULT_GOOGLE_VIDEO_DURATION_SECONDS from litellm.images.utils import ImageEditRequestUtils +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.videos.transformation import BaseVideoConfig from litellm.llms.vertex_ai.common_utils import ( convert_vertex_datetime_to_openai_datetime, @@ -325,13 +326,10 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase): model_name: Final = model.replace("vertex_ai/", "") # Construct the URL - if api_base: - base_url = api_base.rstrip("/") - else: - base_url = get_vertex_base_url(vertex_location) + base_url: Final = api_base or get_vertex_base_url(vertex_location) - url: Final = ( - f"{base_url}/v1/projects/{vertex_project}/locations/{vertex_location}/publishers/google/models/{model_name}" + url: Final = join_api_base( + base_url, f"v1/projects/{vertex_project}/locations/{vertex_location}/publishers/google/models/{model_name}" ) return url @@ -480,7 +478,7 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase): # Construct the full URL including model ID # URL format: https://LOCATION-aiplatform.googleapis.com/v1/projects/PROJECT/locations/LOCATION/publishers/google/models/MODEL:fetchPredictOperation # Strip trailing slashes from api_base and append model - url: Final = f"{api_base.rstrip('/')}/{model}:fetchPredictOperation" + url: Final = join_api_base(api_base, f"{model}:fetchPredictOperation") # Request body contains the operation name params: Final = {"operationName": operation_name} @@ -778,7 +776,7 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase): if vertex_params: request_data["parameters"] = vertex_params - edit_url: Final = f"{api_base.rstrip('/')}/{model}:predictLongRunning" + edit_url: Final = join_api_base(api_base, f"{model}:predictLongRunning") return edit_url, request_data, None def transform_video_edit_response( diff --git a/litellm/llms/volcengine/embedding/transformation.py b/litellm/llms/volcengine/embedding/transformation.py index b298c9996f7..cc1e7d59caa 100644 --- a/litellm/llms/volcengine/embedding/transformation.py +++ b/litellm/llms/volcengine/embedding/transformation.py @@ -10,6 +10,7 @@ from typing import Final import httpx from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import ensure_api_route, join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig from litellm.types.llms.openai import AllEmbeddingInputValues, AllMessageValues @@ -76,12 +77,8 @@ class VolcEngineEmbeddingConfig(BaseEmbeddingConfig): Returns: Complete URL for the embedding API endpoint """ - base_url: Final = get_volcengine_base_url(api_base) - # Construct the complete URL with /embeddings endpoint - if base_url.endswith("/api/v3"): - return f"{base_url}/embeddings" - else: - return f"{base_url}/api/v3/embeddings" + versioned_base_url: Final = ensure_api_route(get_volcengine_base_url(api_base), "api/v3") + return join_api_base(versioned_base_url, "embeddings") def map_openai_params( self, diff --git a/litellm/llms/volcengine/responses/transformation.py b/litellm/llms/volcengine/responses/transformation.py index 6e0aa88b6b6..74f6182bd89 100644 --- a/litellm/llms/volcengine/responses/transformation.py +++ b/litellm/llms/volcengine/responses/transformation.py @@ -11,7 +11,7 @@ from litellm.litellm_core_utils.core_helpers import process_response_headers from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( safe_convert_created_field, ) -from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import ( @@ -137,8 +137,8 @@ class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig): if base_url.endswith("/responses"): return base_url if base_url.endswith("/api/v3"): - return f"{base_url}/responses" - return f"{base_url}/api/v3/responses" + return join_api_base(base_url, "responses") + return join_api_base(base_url, "api/v3/responses") def map_openai_params( self, @@ -345,7 +345,7 @@ class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig): order: Literal["asc", "desc"] = "desc", ) -> tuple[str, dict]: encoded_response_id: Final = encode_url_path_segment(response_id, field_name="response_id") - url: Final = f"{api_base}/{encoded_response_id}/input_items" + url: Final = join_api_base(api_base, f"{encoded_response_id}/input_items") params: Final[dict[str, str | int]] = {} if after is not None: params["after"] = after @@ -380,7 +380,7 @@ class VolcEngineResponsesAPIConfig(OpenAIResponsesAPIConfig): headers: dict, ) -> tuple[str, dict]: encoded_response_id: Final = encode_url_path_segment(response_id, field_name="response_id") - url: Final = f"{api_base}/{encoded_response_id}/cancel" + url: Final = join_api_base(api_base, f"{encoded_response_id}/cancel") data: Final[dict] = {} return url, data diff --git a/litellm/llms/voyage/embedding/transformation.py b/litellm/llms/voyage/embedding/transformation.py index 7e8425bb3d0..1ba66e5ce65 100644 --- a/litellm/llms/voyage/embedding/transformation.py +++ b/litellm/llms/voyage/embedding/transformation.py @@ -5,6 +5,7 @@ import httpx from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig from litellm.llms.voyage.common_utils import get_default_base_url, get_voyage_api_key @@ -50,12 +51,7 @@ class VoyageEmbeddingConfig(BaseEmbeddingConfig): litellm_params: dict, stream: bool | None = None, ) -> str: - if api_base: - clean_api_base: Final = api_base.rstrip("/") - if not clean_api_base.endswith("/embeddings"): - return f"{clean_api_base}/embeddings" - return clean_api_base - return f"{get_default_base_url(api_key)}/embeddings" + return ensure_api_route(api_base or get_default_base_url(api_key), "embeddings") def get_supported_openai_params(self, model: str) -> list: return [ diff --git a/litellm/llms/voyage/embedding/transformation_contextual.py b/litellm/llms/voyage/embedding/transformation_contextual.py index 8e3802ccbc9..62c03139fbe 100644 --- a/litellm/llms/voyage/embedding/transformation_contextual.py +++ b/litellm/llms/voyage/embedding/transformation_contextual.py @@ -10,6 +10,7 @@ import httpx from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig from litellm.llms.voyage.common_utils import get_default_base_url, get_voyage_api_key @@ -58,12 +59,7 @@ class VoyageContextualEmbeddingConfig(BaseEmbeddingConfig): litellm_params: dict, stream: bool | None = None, ) -> str: - if api_base: - clean_api_base: Final = api_base.rstrip("/") - if not clean_api_base.endswith("/contextualizedembeddings"): - return f"{clean_api_base}/contextualizedembeddings" - return clean_api_base - return f"{get_default_base_url(api_key)}/contextualizedembeddings" + return ensure_api_route(api_base or get_default_base_url(api_key), "contextualizedembeddings") def get_supported_openai_params(self, model: str) -> list: # mutable-ok: base class signature return ["encoding_format", "dimensions"] diff --git a/litellm/llms/voyage/embedding/transformation_multimodal.py b/litellm/llms/voyage/embedding/transformation_multimodal.py index 29fdb9ecfe1..c04a2e1f26b 100644 --- a/litellm/llms/voyage/embedding/transformation_multimodal.py +++ b/litellm/llms/voyage/embedding/transformation_multimodal.py @@ -13,6 +13,7 @@ import httpx from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig from litellm.llms.voyage.common_utils import get_default_base_url, get_voyage_api_key @@ -59,12 +60,7 @@ class VoyageMultimodalEmbeddingConfig(BaseEmbeddingConfig): litellm_params: dict, stream: bool | None = None, ) -> str: - if api_base: - clean_api_base: Final = api_base.rstrip("/") - if not clean_api_base.endswith("/multimodalembeddings"): - return f"{clean_api_base}/multimodalembeddings" - return clean_api_base - return f"{get_default_base_url(api_key)}/multimodalembeddings" + return ensure_api_route(api_base or get_default_base_url(api_key), "multimodalembeddings") def get_supported_openai_params(self, model: str) -> list: return ["dimensions"] diff --git a/litellm/llms/voyage/rerank/transformation.py b/litellm/llms/voyage/rerank/transformation.py index 121c18e82ae..7fe643870d5 100644 --- a/litellm/llms/voyage/rerank/transformation.py +++ b/litellm/llms/voyage/rerank/transformation.py @@ -11,6 +11,7 @@ import httpx from pydantic import ConfigDict, TypeAdapter from litellm._uuid import uuid +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import LiteLLMLoggingObj from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig from litellm.llms.voyage.common_utils import get_default_base_url, get_voyage_api_key @@ -80,9 +81,9 @@ class VoyageRerankConfig(BaseRerankConfig): api_base = api_base.rstrip("/") if not api_base.endswith("/v1/rerank"): if api_base.endswith("/v1"): - api_base = f"{api_base}/rerank" + api_base = join_api_base(api_base, "rerank") else: - api_base = f"{api_base}/v1/rerank" + api_base = join_api_base(api_base, "v1/rerank") return api_base def transform_rerank_request( diff --git a/litellm/llms/watsonx/audio_transcription/transformation.py b/litellm/llms/watsonx/audio_transcription/transformation.py index 77ef03b8e10..e5da7dd412f 100644 --- a/litellm/llms/watsonx/audio_transcription/transformation.py +++ b/litellm/llms/watsonx/audio_transcription/transformation.py @@ -12,6 +12,7 @@ from pydantic import ConfigDict, TypeAdapter import litellm from litellm.litellm_core_utils.audio_utils.utils import process_audio_file +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.types.llms.openai import ( AllMessageValues, OpenAIAudioTranscriptionOptionalParams, @@ -150,10 +151,9 @@ class IBMWatsonXAudioTranscriptionConfig(IBMWatsonXMixin, OpenAIWhisperAudioTran """ # Get base URL url = self._get_base_url(api_base=api_base) - url = url.rstrip("/") # Add the audio transcription endpoint - url = f"{url}/ml/v1/audio/transcriptions" + url = join_api_base(url, "ml/v1/audio/transcriptions") # Add version parameter (only version in query string, not project_id) api_version: Final = optional_params.get("api_version", None) or litellm.WATSONX_DEFAULT_API_VERSION diff --git a/litellm/llms/watsonx/chat/transformation.py b/litellm/llms/watsonx/chat/transformation.py index 4ff2212711d..071ff7108b9 100644 --- a/litellm/llms/watsonx/chat/transformation.py +++ b/litellm/llms/watsonx/chat/transformation.py @@ -7,6 +7,7 @@ Docs: https://cloud.ibm.com/apidocs/watsonx-ai#text-chat from typing import Final from litellm import verbose_logger +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.secret_managers.main import get_secret_str from litellm.types.llms.watsonx import ( WatsonXAIEndpoint, @@ -96,7 +97,7 @@ class IBMWatsonXChatConfig(IBMWatsonXMixin, OpenAIGPTConfig): endpoint = endpoint.format(deployment_id=deployment_id) else: endpoint = WatsonXAIEndpoint.CHAT_STREAM.value if stream else WatsonXAIEndpoint.CHAT.value - url = url.rstrip("/") + endpoint + url = join_api_base(url, endpoint) ## add api version url = self._add_api_version_to_url(url=url, api_version=optional_params.pop("api_version", None)) diff --git a/litellm/llms/watsonx/completion/transformation.py b/litellm/llms/watsonx/completion/transformation.py index 5c0416bd6ca..9ea062ea2f8 100644 --- a/litellm/llms/watsonx/completion/transformation.py +++ b/litellm/llms/watsonx/completion/transformation.py @@ -5,6 +5,7 @@ from typing import TYPE_CHECKING, Any, Final import httpx +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator from litellm.types.llms.openai import AllMessageValues, ChatCompletionUsageBlock from litellm.types.llms.watsonx import WatsonXAIEndpoint @@ -342,7 +343,7 @@ class IBMWatsonXAIConfig(IBMWatsonXMixin, BaseConfig): endpoint = endpoint.format(deployment_id=deployment_id) else: endpoint = WatsonXAIEndpoint.TEXT_GENERATION_STREAM if stream else WatsonXAIEndpoint.TEXT_GENERATION - url = url.rstrip("/") + endpoint + url = join_api_base(url, endpoint) ## add api version url = self._add_api_version_to_url(url=url, api_version=optional_params.pop("api_version", None)) diff --git a/litellm/llms/watsonx/embed/transformation.py b/litellm/llms/watsonx/embed/transformation.py index 1b870c7f639..6ec3b50bbe9 100644 --- a/litellm/llms/watsonx/embed/transformation.py +++ b/litellm/llms/watsonx/embed/transformation.py @@ -9,6 +9,7 @@ import httpx from pydantic import ConfigDict, TypeAdapter from typing_extensions import ReadOnly, TypedDict +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.embedding.transformation import ( BaseEmbeddingConfig, LiteLLMLoggingObj, @@ -92,7 +93,7 @@ class IBMWatsonXEmbeddingConfig(IBMWatsonXMixin, BaseEmbeddingConfig): if model.startswith("deployment/"): deployment_id: Final = "/".join(model.split("/")[1:]) endpoint = endpoint.format(deployment_id=deployment_id) - url = url.rstrip("/") + endpoint + url = join_api_base(url, endpoint) ## add api version url = self._add_api_version_to_url(url=url, api_version=optional_params.pop("api_version", None)) diff --git a/litellm/llms/watsonx/rerank/transformation.py b/litellm/llms/watsonx/rerank/transformation.py index 5f337603763..591ddda3cbd 100644 --- a/litellm/llms/watsonx/rerank/transformation.py +++ b/litellm/llms/watsonx/rerank/transformation.py @@ -11,6 +11,7 @@ from typing import Final, cast import httpx from pydantic import ConfigDict, TypeAdapter +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import LiteLLMLoggingObj from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig from litellm.secret_managers.main import get_secret_str @@ -45,7 +46,7 @@ class IBMWatsonXRerankConfig(IBMWatsonXMixin, BaseRerankConfig): base_url: Final = self._get_base_url(api_base=api_base) endpoint: Final = WatsonXAIEndpoint.RERANK.value - url: Final = base_url.rstrip("/") + endpoint + url: Final = join_api_base(base_url, endpoint) params: Final = optional_params or {} diff --git a/litellm/llms/xai/audio_transcription/transformation.py b/litellm/llms/xai/audio_transcription/transformation.py index e34034b5ae9..e29c3d9334a 100644 --- a/litellm/llms/xai/audio_transcription/transformation.py +++ b/litellm/llms/xai/audio_transcription/transformation.py @@ -10,6 +10,7 @@ from pydantic import ConfigDict, TypeAdapter, ValidationError import litellm from litellm.litellm_core_utils.audio_utils.utils import process_audio_file +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.llms.openai import ( @@ -185,7 +186,7 @@ class XAIAudioTranscriptionConfig(BaseAudioTranscriptionConfig): ) -> str: base: Final = (XAIModelInfo.get_api_base(api_base) or "").rstrip("/") normalized: Final = base.removesuffix("/v1") - return f"{normalized}/v1/stt" + return join_api_base(normalized, "v1/stt") def validate_environment( self, diff --git a/litellm/llms/xai/batches/transformation.py b/litellm/llms/xai/batches/transformation.py index fdd9939526e..1e803057e3d 100644 --- a/litellm/llms/xai/batches/transformation.py +++ b/litellm/llms/xai/batches/transformation.py @@ -20,7 +20,7 @@ from pydantic import ConfigDict from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm.constants import XAI_API_BASE -from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.xai.common_utils import XAIModelInfo from litellm.secret_managers.main import get_secret_str @@ -76,10 +76,11 @@ def get_xai_auth_headers( def xai_batches_url(api_base: str | None, batch_id: str | None = None, suffix: str = "") -> str: - base: Final = f"{get_xai_api_base(api_base)}/v1/batches" + root: Final = get_xai_api_base(api_base) if batch_id is None: - return base - return f"{base}/{encode_url_path_segment(batch_id, field_name='batch_id')}{suffix}" + return join_api_base(root, "v1/batches") + encoded_batch_id: Final = encode_url_path_segment(batch_id, field_name="batch_id") + return join_api_base(root, f"v1/batches/{encoded_batch_id}{suffix}") def is_xai_batch_results_id(file_id: str) -> bool: diff --git a/litellm/llms/xai/common_utils.py b/litellm/llms/xai/common_utils.py index 4843caeac59..06e3f32a767 100644 --- a/litellm/llms/xai/common_utils.py +++ b/litellm/llms/xai/common_utils.py @@ -5,6 +5,7 @@ import httpx from pydantic import ConfigDict, TypeAdapter import litellm +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.base_utils import BaseLLMModelInfo from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues @@ -91,7 +92,7 @@ class XAIModelInfo(BaseLLMModelInfo): "XAI API base or key is not set. Set XAI_API_BASE and provide an xAI API key via api_key, litellm.xai_key, or XAI_API_KEY." ) response: Final = litellm.module_level_client.get( - url=f"{api_base}/v1/models", + url=join_api_base(api_base, "v1/models"), headers={"Authorization": f"Bearer {api_key}"}, ) diff --git a/litellm/llms/xai/files/transformation.py b/litellm/llms/xai/files/transformation.py index 33baca79005..88f2fc51983 100644 --- a/litellm/llms/xai/files/transformation.py +++ b/litellm/llms/xai/files/transformation.py @@ -14,7 +14,7 @@ from pydantic import ConfigDict from typing_extensions import ReadOnly, TypedDict from litellm.litellm_core_utils.prompt_templates.common_utils import extract_file_data -from litellm.litellm_core_utils.url_utils import encode_url_path_segment +from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.files.transformation import BaseFilesConfig, LiteLLMLoggingObj from litellm.types.llms.base import LiteLLMBaseModel @@ -101,11 +101,11 @@ class XAIFilesConfig(BaseFilesConfig): litellm_params: Mapping[str, object], stream: bool | None = None, ) -> str: - return f"{get_xai_api_base(api_base)}/v1/files" + return join_api_base(get_xai_api_base(api_base), "v1/files") def _file_url(self, file_id: str, litellm_params: Mapping[str, object], suffix: str = "") -> str: encoded_file_id: Final = encode_url_path_segment(file_id, field_name="file_id") - return f"{_api_base_from(litellm_params)}/v1/files/{encoded_file_id}{suffix}" + return join_api_base(_api_base_from(litellm_params), f"v1/files/{encoded_file_id}{suffix}") def get_error_class( self, error_message: str, status_code: int, headers: Mapping[str, str] | httpx.Headers @@ -204,7 +204,7 @@ class XAIFilesConfig(BaseFilesConfig): optional_params: Mapping[str, object], litellm_params: Mapping[str, object], ) -> tuple[str, dict[str, str]]: # mutable-ok: BaseFilesConfig signature - return f"{_api_base_from(litellm_params)}/v1/files", _NO_QUERY_PARAMS + return join_api_base(_api_base_from(litellm_params), "v1/files"), _NO_QUERY_PARAMS def transform_list_files_next_request( self, @@ -215,7 +215,7 @@ class XAIFilesConfig(BaseFilesConfig): page: Final = XAIFileList.model_validate(raw_response.json()) if not page.pagination_token or not page.data: return None - return f"{_api_base_from(litellm_params)}/v1/files", {"pagination_token": page.pagination_token} + return join_api_base(_api_base_from(litellm_params), "v1/files"), {"pagination_token": page.pagination_token} def transform_list_files_response( self, diff --git a/litellm/llms/xai/responses/transformation.py b/litellm/llms/xai/responses/transformation.py index 2f638da49c8..b6ea516fb79 100644 --- a/litellm/llms/xai/responses/transformation.py +++ b/litellm/llms/xai/responses/transformation.py @@ -9,6 +9,7 @@ import litellm from litellm._logging import verbose_logger from litellm.constants import XAI_API_BASE from litellm.exceptions import AuthenticationError +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig from litellm.llms.xai.common_utils import XAIModelInfo, xai_reported_cost_in_usd from litellm.secret_managers.main import get_secret_str @@ -249,10 +250,7 @@ class XAIResponsesAPIConfig(OpenAIResponsesAPIConfig): else: api_base = api_base or litellm.api_base or get_secret_str("XAI_API_BASE") or XAI_API_BASE - # Remove trailing slashes - api_base = api_base.rstrip("/") - - return f"{api_base}/responses" + return join_api_base(api_base, "responses") def transform_response_api_response( self, diff --git a/litellm/llms/you_com/search/transformation.py b/litellm/llms/you_com/search/transformation.py index 8b2df949a93..9228deda26b 100644 --- a/litellm/llms/you_com/search/transformation.py +++ b/litellm/llms/you_com/search/transformation.py @@ -10,6 +10,7 @@ from typing import Final, TypedDict import httpx from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.base_llm.search.transformation import ( BaseSearchConfig, SearchResponse, @@ -109,7 +110,7 @@ class YouComSearchConfig(BaseSearchConfig): api_base = api_base.rstrip("/") if not api_base.endswith("/v1/search") and not api_base.endswith("/v1/agents/search"): - api_base = f"{api_base}/v1/search" + api_base = join_api_base(api_base, "v1/search") return api_base diff --git a/litellm/main.py b/litellm/main.py index 079b0bcbd6c..27700cf9def 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -112,6 +112,7 @@ from litellm.litellm_core_utils.request_timeout_resolver import ( get_configured_request_timeout, ) from litellm.litellm_core_utils.tokenizer import Encoding as Tokenizer +from litellm.litellm_core_utils.url_utils import ensure_api_route from litellm.llms.base_llm import BaseConfig, BaseImageGenerationConfig from litellm.llms.base_llm.base_model_iterator import ( convert_model_response_to_streaming, @@ -230,6 +231,7 @@ from .llms.azure.azure import AzureChatCompletion, check_dynamic_azure_params from .llms.azure.chat.o_series_handler import AzureOpenAIO1ChatCompletion from .llms.azure.completion.handler import AzureTextCompletion from .llms.azure_ai.anthropic.handler import AzureAnthropicChatCompletion +from .llms.azure_ai.anthropic.transformation import azure_anthropic_messages_url from .llms.azure_ai.embed import AzureAIEmbedding from .llms.bedrock.chat import BedrockConverseLLM from .llms.bedrock.chat.converse_transformation import AmazonConverseConfig @@ -1666,16 +1668,8 @@ def _complete_azure_ai(ctx: CompletionDispatchContext) -> _CompletionDispatchRes ) api_key = AzureFoundryModelInfo.get_api_key(api_key) - # Ensure the URL ends with /v1/messages for Anthropic if api_base: - api_base = api_base.rstrip("/") - if not api_base.endswith("/v1/messages"): - if "/anthropic" in api_base: - parts: Final = api_base.split("/anthropic", 1) - api_base = parts[0] + "/anthropic" - else: - api_base = api_base + "/anthropic" - api_base = api_base + "/v1/messages" + api_base = azure_anthropic_messages_url(api_base) response = azure_anthropic_chat_completions.completion( model=model, @@ -2895,8 +2889,8 @@ def _complete_anthropic_text( # Check if we should disable automatic URL suffix appending disable_url_suffix: Final = get_secret_bool("LITELLM_ANTHROPIC_DISABLE_URL_SUFFIX") - if api_base is not None and not disable_url_suffix and not api_base.endswith("/v1/complete"): - api_base += "/v1/complete" + if api_base is not None and not disable_url_suffix: + api_base = ensure_api_route(api_base, "v1/complete") elif disable_url_suffix: verbose_logger.debug("LITELLM_ANTHROPIC_DISABLE_URL_SUFFIX is set, skipping /v1/complete suffix") @@ -2951,8 +2945,8 @@ def _complete_anthropic(ctx: CompletionDispatchContext) -> _CompletionDispatchRe # Check if we should disable automatic URL suffix appending disable_url_suffix: Final = get_secret_bool("LITELLM_ANTHROPIC_DISABLE_URL_SUFFIX") - if api_base is not None and not disable_url_suffix and not api_base.endswith("/v1/messages"): - api_base += "/v1/messages" + if api_base is not None and not disable_url_suffix: + api_base = ensure_api_route(api_base, "v1/messages") elif disable_url_suffix: verbose_logger.debug("LITELLM_ANTHROPIC_DISABLE_URL_SUFFIX is set, skipping /v1/messages suffix") @@ -4090,9 +4084,7 @@ def _complete_text_completion_inception( api_base = passed_api_base or get_secret_str("INCEPTION_API_BASE") or "https://api.inceptionlabs.ai/v1" # FIM is served at `/v1/fim/completions`; the OpenAI client appends # `/completions`, so point it at the `/v1/fim` base. - api_base = api_base.rstrip("/") - if not api_base.endswith("/fim"): - api_base += "/fim" + api_base = ensure_api_route(api_base, "fim") # Don't forward the server-managed Inception key to a caller-supplied # api_base; only resolve it for the default/server base, or when the diff --git a/litellm/rag/ingestion/gemini_ingestion.py b/litellm/rag/ingestion/gemini_ingestion.py index 4ee16685e03..6320e477bf2 100644 --- a/litellm/rag/ingestion/gemini_ingestion.py +++ b/litellm/rag/ingestion/gemini_ingestion.py @@ -12,6 +12,7 @@ from typing import TYPE_CHECKING, Final, cast from pydantic import ConfigDict from litellm._logging import verbose_logger +from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, header_value, @@ -107,7 +108,7 @@ class GeminiRAGIngestion(BaseRAGIngestion): raise ValueError("GEMINI_API_BASE is required") api_version: Final = "v1beta" - base_url: Final = f"{api_base}/{api_version}" + base_url: Final = join_api_base(api_base, api_version) # Create File Search store if not provided if not vector_store_id: @@ -148,7 +149,7 @@ class GeminiRAGIngestion(BaseRAGIngestion): Returns: Store name (format: fileSearchStores/xxxxxxx) """ - url: Final = f"{base_url}/fileSearchStores" + url: Final = join_api_base(base_url, "fileSearchStores") request_body: Final = {"displayName": display_name} @@ -236,7 +237,7 @@ class GeminiRAGIngestion(BaseRAGIngestion): # base_url is like: https://generativelanguage.googleapis.com/v1beta # We need: https://generativelanguage.googleapis.com/upload/v1beta/{store_id}:uploadToFileSearchStore api_base: Final = base_url.replace("/v1beta", "") # Get base without version - url: Final = f"{api_base}/upload/v1beta/{vector_store_id}:uploadToFileSearchStore" + url: Final = join_api_base(api_base, f"upload/v1beta/{vector_store_id}:uploadToFileSearchStore") # Build request body with chunking config and metadata if provided request_body: Final[dict[str, object]] = {"displayName": filename} diff --git a/tests/integration/providers/_api_base_url_join_support.py b/tests/integration/providers/_api_base_url_join_support.py new file mode 100644 index 00000000000..44bd6ec4dca --- /dev/null +++ b/tests/integration/providers/_api_base_url_join_support.py @@ -0,0 +1,163 @@ +import json +from collections.abc import Callable, Mapping +from types import MappingProxyType +from typing import Final +from urllib.parse import urlsplit + +from integration._support.anthropic_sse import message_stream +from integration._support.openai_wire import chat_reply +from integration._support.upstream import aws_event_stream_frame +from integration._support.wire import Reply, Request +from integration.providers._api_base_trailing_slash_support import VECTOR, marker_of, reply_text, request_json +from pydantic import JsonValue + +COUNTED_TOKENS: Final = 7 +RELEVANCE: Final = 0.75 +_USAGE: Final = {"input_tokens": 3, "output_tokens": 2} + + +def _json(payload: Mapping[str, JsonValue]) -> Reply: + return Reply(body=json.dumps(payload).encode()) + + +def _path(request: Request) -> str: + return urlsplit(request.target).path + + +def _wants_stream(request: Request) -> bool: + return request_json(request).get("stream") is True + + +def anthropic_message(marker: str) -> dict[str, JsonValue]: + return { + "id": f"msg_{marker}", + "type": "message", + "role": "assistant", + "model": "scripted", + "content": [{"type": "text", "text": reply_text(marker)}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": dict(_USAGE), + } + + +def _anthropic(request: Request) -> Reply: + if _path(request).endswith("/count_tokens"): + return _json({"input_tokens": COUNTED_TOKENS}) + marker: Final = marker_of(request) + if _wants_stream(request): + return Reply( + content_type="text/event-stream", chunks=message_stream(f"msg_{marker}", "scripted", reply_text(marker)) + ) + return _json(anthropic_message(marker)) + + +def _gemini_candidate(text: str, finished: bool) -> dict[str, JsonValue]: + finish: Final[dict[str, JsonValue]] = {"finishReason": "STOP"} if finished else {} + return {"content": {"role": "model", "parts": [{"text": text}]}, "index": 0, **finish} + + +def _gemini_frame(text: str, finished: bool) -> bytes: + usage: Final = {"promptTokenCount": 3, "candidatesTokenCount": 2, "totalTokenCount": 5} + frame: Final = {"candidates": [_gemini_candidate(text, finished)], "usageMetadata": usage} + return b"data: " + json.dumps(frame).encode() + b"\r\n\r\n" + + +def _gemini(request: Request) -> Reply: + path: Final = _path(request) + marker: Final = marker_of(request) + if path.endswith(":batchEmbedContents"): + return _json({"embeddings": [{"values": list(VECTOR)}]}) + if path.endswith(":streamGenerateContent"): + return Reply( + content_type="text/event-stream", + chunks=(_gemini_frame("reply ", False), _gemini_frame(marker, True)), + ) + return _json( + { + "candidates": [_gemini_candidate(reply_text(marker), True)], + "usageMetadata": {"promptTokenCount": 3, "candidatesTokenCount": 2, "totalTokenCount": 5}, + } + ) + + +def _embedding_list(model: str) -> Reply: + return _json( + { + "object": "list", + "data": [{"object": "embedding", "index": 0, "embedding": list(VECTOR)}], + "model": model, + "usage": {"prompt_tokens": 2, "total_tokens": 2}, + } + ) + + +def _openai_like(model: str) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + path: Final = _path(request) + marker: Final = marker_of(request) + if path.endswith("/embeddings"): + return _embedding_list(model) + if path.endswith("/audio/transcriptions"): + return _json({"text": reply_text(marker), "model": model}) + return chat_reply(f"chatcmpl-{marker}", model, reply_text(marker), stream=_wants_stream(request)) + + return respond + + +def _converse_event(kind: str, payload: Mapping[str, JsonValue]) -> bytes: + headers: Final = {":event-type": kind, ":content-type": "application/json", ":message-type": "event"} + return aws_event_stream_frame(headers, json.dumps(payload, separators=(",", ":")).encode()) + + +def _bedrock_converse(request: Request) -> Reply: + marker: Final = marker_of(request) + usage: Final[dict[str, JsonValue]] = {"inputTokens": 3, "outputTokens": 2, "totalTokens": 5} + if _path(request).endswith("/converse-stream"): + events: Final[tuple[tuple[str, Mapping[str, JsonValue]], ...]] = ( + ("messageStart", {"role": "assistant"}), + ("contentBlockDelta", {"delta": {"text": "reply "}, "contentBlockIndex": 0}), + ("contentBlockDelta", {"delta": {"text": marker}, "contentBlockIndex": 0}), + ("contentBlockStop", {"contentBlockIndex": 0}), + ("messageStop", {"stopReason": "end_turn"}), + ("metadata", {"usage": usage, "metrics": {"latencyMs": 1}}), + ) + return Reply( + content_type="application/vnd.amazon.eventstream", + chunks=tuple(_converse_event(kind, payload) for kind, payload in events), + ) + return _json( + { + "output": {"message": {"role": "assistant", "content": [{"text": reply_text(marker)}]}}, + "stopReason": "end_turn", + "usage": usage, + "metrics": {"latencyMs": 1}, + } + ) + + +def _cohere_rerank(request: Request) -> Reply: + return _json( + { + "id": f"rerank-{marker_of(request)}", + "results": [{"index": 0, "relevance_score": RELEVANCE}], + "meta": {"billed_units": {"search_units": 1}}, + } + ) + + +_BY_PREFIX: Final[Mapping[str, Callable[[Request], Reply]]] = MappingProxyType( + { + "an": _anthropic, + "bi": lambda request: _json(anthropic_message(marker_of(request))), + "bc": _bedrock_converse, + "gm": _gemini, + "mi": _openai_like("mistral-small-latest"), + "db": _openai_like("databricks-meta-llama-3-3-70b-instruct"), + "co": _cohere_rerank, + } +) + + +def join_peer(request: Request) -> Reply: + return _BY_PREFIX[request.target.split("/")[1]](request) diff --git a/tests/integration/providers/test_api_base_url_join_chaos.py b/tests/integration/providers/test_api_base_url_join_chaos.py new file mode 100644 index 00000000000..8ebcd165f4d --- /dev/null +++ b/tests/integration/providers/test_api_base_url_join_chaos.py @@ -0,0 +1,176 @@ +import asyncio +from collections.abc import Awaitable, Callable +from itertools import chain, repeat +from typing import Final, TypeAlias + +import anthropic +import httpx +import openai +from integration._support.client import Gateway, Scenario, eventually +from integration._support.database import read_rows +from integration._support.wire import Reply, Request, wire_server +from integration.providers._api_base_trailing_slash_support import ( + KEY, + VECTOR, + assert_control_answers, + assert_one_request, + body_of, + marker_of, + new_marker, + of_marker, + reply_text, +) +from integration.providers._api_base_url_join_support import join_peer +from pydantic import JsonValue + +_PER_KIND: Final = 3 +_HAIKU: Final = "anthropic.claude-haiku-5-5" +_GENERATE: Final = "/gm/v1beta/models/gemini-3.8-flash:generateContent" +_STREAM_GENERATE: Final = "/gm/v1beta/models/gemini-3.8-flash:streamGenerateContent?alt=sse" + +Answer: TypeAlias = tuple[str | None, str | None] +Call: TypeAlias = Callable[[str], Awaitable[Answer]] + + +def _deploy(scenario: Scenario, model: str, api_base: str) -> str: + if model.startswith("bedrock/"): + return scenario.model(model=model, api_base=api_base, api_key=KEY, aws_region_name="us-east-1") + return scenario.model(model=model, api_base=api_base, api_key=KEY) + + +def _spend_api_bases(identities: tuple[str, ...]) -> dict[str, JsonValue]: + rows: Final = eventually( + lambda: read_rows( + "SELECT request_id, api_base FROM \"LiteLLM_SpendLogs\" WHERE request_id = ANY(string_to_array(%s, ','))", + (",".join(identities),), + ), + lambda found: len(found) == len(identities), + seconds=70, + ) + return {str(row["request_id"]): row["api_base"] for row in rows} + + +def test_mixed_burst_on_slash_bases_lands_each_call_once_on_its_single_slash_url(gateway: Gateway) -> None: + with wire_server(join_peer) as wire, gateway.scenario() as scenario: + claude: Final = _deploy(scenario, "anthropic/claude-haiku-5-5", f"{wire.url}/an/") + gemini: Final = _deploy(scenario, "gemini/gemini-3.8-flash", f"{wire.url}/gm/v1beta/") + mistral: Final = _deploy(scenario, "mistral/mistral-small-latest", f"{wire.url}/mi/v1/") + converse: Final = _deploy(scenario, f"bedrock/converse/{_HAIKU}", f"{wire.url}/bc/") + databricks: Final = _deploy( + scenario, "databricks/databricks-meta-llama-3-3-70b-instruct", f"{wire.url}/db/serving-endpoints/" + ) + embed_model: Final = _deploy(scenario, "mistral/mistral-embed", f"{wire.url}/mi/v1/") + + def openai_client() -> openai.AsyncOpenAI: + return openai.AsyncOpenAI(base_url=str(gateway.client.base_url) + "/v1", api_key=gateway.key, max_retries=0) + + def chat(model: str) -> Call: + async def call(marker: str) -> Answer: + async with openai_client() as client: + completion: Final = await client.chat.completions.create( + model=model, messages=[{"role": "user", "content": marker}] + ) + return completion.choices[0].message.content, completion.id + + return call + + def streamed(model: str) -> Call: + async def call(marker: str) -> Answer: + async with openai_client() as client: + stream: Final = await client.chat.completions.create( + model=model, messages=[{"role": "user", "content": marker}], stream=True + ) + parts: Final = [chunk.choices[0].delta.content or "" async for chunk in stream if chunk.choices] + return "".join(parts), None + + return call + + async def messages(marker: str) -> Answer: + async with anthropic.AsyncAnthropic( + base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0 + ) as client: + message: Final = await client.messages.create( + model=claude, max_tokens=16, messages=[{"role": "user", "content": marker}] + ) + return "".join(block.text for block in message.content if block.type == "text"), None + + async def embed(marker: str) -> Answer: + async with httpx.AsyncClient(base_url=str(gateway.client.base_url), timeout=60, trust_env=False) as client: + response: Final = await client.post( + "/v1/embeddings", + json={"model": embed_model, "input": marker}, + headers={"Authorization": f"Bearer {gateway.key}"}, + ) + data: Final = body_of(response)["data"] + vector_matches: Final = data == [{"object": "embedding", "index": 0, "embedding": list(VECTOR)}] + return (reply_text(marker) if vector_matches else str(data)), None + + kinds: Final[tuple[tuple[str, Call], ...]] = ( + ("/an/v1/messages", chat(claude)), + ("/an/v1/messages", streamed(claude)), + ("/an/v1/messages", messages), + (_GENERATE, chat(gemini)), + (_STREAM_GENERATE, streamed(gemini)), + ("/mi/v1/chat/completions", chat(mistral)), + ("/mi/v1/chat/completions", streamed(mistral)), + (f"/bc/model/{_HAIKU}/converse", chat(converse)), + (f"/bc/model/{_HAIKU}/converse-stream", streamed(converse)), + ("/db/serving-endpoints/chat/completions", chat(databricks)), + ("/mi/v1/embeddings", embed), + ) + repeated_kinds: Final = tuple(chain.from_iterable(repeat(kind, _PER_KIND) for kind in kinds)) + calls: Final = tuple((target, call, new_marker()) for target, call in repeated_kinds) + + async def burst() -> list[Answer]: + return await asyncio.gather(*(call(marker) for _, call, marker in calls)) + + answers: Final = asyncio.run(burst()) + assert [text for text, _ in answers] == [reply_text(marker) for *_, marker in calls] + received: Final = wire.drain() + assert len(received) == len(calls), sorted(request.target for request in received) + for target, _, marker in calls: + assert_one_request(of_marker(received, marker), target, marker) + logged: Final = { + identity: f"{wire.url}{target}" + for (target, _, _), (_, identity) in zip(calls, answers, strict=True) + if identity is not None + } + assert _spend_api_bases(tuple(logged)) == logged + + +def test_upstream_dropping_doomed_connections_mid_burst_fails_only_those_calls(gateway: Gateway) -> None: + markers: Final = tuple(new_marker() for _ in range(16)) + doomed: Final = frozenset(markers[::2]) + + def respond(request: Request) -> Reply: + return Reply(drop_connection=True) if marker_of(request) in doomed else join_peer(request) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + claude: Final = _deploy(scenario, "anthropic/claude-haiku-5-5", f"{wire.url}/an/") + gemini: Final = _deploy(scenario, "gemini/gemini-3.8-flash", f"{wire.url}/gm/v1beta/") + routed: Final = { + marker: (claude, "/an/v1/messages") if index % 4 < 2 else (gemini, _GENERATE) + for index, marker in enumerate(markers) + } + + async def call(marker: str) -> tuple[int, str]: + async with httpx.AsyncClient(base_url=str(gateway.client.base_url), timeout=60, trust_env=False) as client: + response: Final = await client.post( + "/v1/chat/completions", + json={"model": routed[marker][0], "messages": [{"role": "user", "content": marker}]}, + headers={"Authorization": f"Bearer {gateway.key}"}, + ) + return response.status_code, response.text + + async def burst() -> list[tuple[int, str]]: + return await asyncio.gather(*(call(marker) for marker in markers)) + + answers: Final = dict(zip(markers, asyncio.run(burst()), strict=True)) + received: Final = wire.drain() + assert_control_answers(gateway, scenario) + for marker in doomed: + assert answers[marker][0] >= 500 and "error" in answers[marker][1], answers[marker] + assert {request.target for request in of_marker(received, marker)} == {routed[marker][1]}, marker + for marker in frozenset(markers) - doomed: + assert answers[marker][0] == 200 and reply_text(marker) in answers[marker][1], answers[marker] + assert_one_request(of_marker(received, marker), routed[marker][1], marker) diff --git a/tests/integration/providers/test_api_base_url_join_host_owned_proxy.py b/tests/integration/providers/test_api_base_url_join_host_owned_proxy.py new file mode 100644 index 00000000000..4a075138305 --- /dev/null +++ b/tests/integration/providers/test_api_base_url_join_host_owned_proxy.py @@ -0,0 +1,140 @@ +import os +from collections.abc import Callable, Iterator +from pathlib import Path +from typing import Final + +import pytest +import yaml +from integration._support.client import JSON_OBJECT, Gateway, gateway_from_environment +from integration._support.forward_proxy import ForwardProxy, refusing_forward_proxy +from integration._support.process import graceful_stop_seconds, owned_proxy +from integration.providers._api_base_trailing_slash_support import KEY, new_marker +from pydantic import JsonValue + +_AZURE_HOST: Final = "anthropic-prod.services.ai.azure.com" +_EMBEDDINGS: Final = "host-embeddings" +_RERANK: Final = "host-rerank" +_AZURE_CLAUDE: Final = "host-azure-claude" +_AZURE_COUNT: Final = "host-azure-count" +_EMPTY_BASE: Final = "host-empty-base" +_SCRUBBED_PREFIXES: Final = ("ANTHROPIC_", "AZURE_AI_", "INFINITY_") +_PROXY_VARIABLES: Final = frozenset({"HTTPS_PROXY", "HTTP_PROXY", "ALL_PROXY", "NO_PROXY"}) +_load_yaml: Final[Callable[[str], object]] = yaml.safe_load + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) + + +def _deployment(name: str, model: str, api_base: str) -> dict[str, JsonValue]: + return {"model_name": name, "litellm_params": {"model": model, "api_base": api_base, "api_key": KEY}} + + +def _config(directory: Path) -> Path: + base: Final = JSON_OBJECT.validate_python(_load_yaml(Path("tests/integration/proxy_config.yaml").read_text())) + configuration: Final = { + **base, + "model_list": [ + _deployment(_EMBEDDINGS, "infinity/BAAI/bge-small-en-v1.5", "http://embeddings"), + _deployment(_RERANK, "infinity/BAAI/bge-reranker-base", "http://rerank"), + _deployment(_AZURE_CLAUDE, "azure_ai/claude-haiku-5-5", f"http://{_AZURE_HOST}/anthropic"), + _deployment(_AZURE_COUNT, "azure_ai/claude-haiku-5-5", f"http://{_AZURE_HOST}"), + _deployment(_EMPTY_BASE, "anthropic/claude-haiku-5-5", ""), + ], + "router_settings": {"disable_cooldowns": True, "num_retries": 0}, + } + path: Final = directory / "host-named-like-the-route.yaml" + path.write_text(yaml.safe_dump(configuration)) + return path + + +def _scrubbed() -> tuple[str, ...]: + return tuple( + name + for name in os.environ + if name.upper() in _PROXY_VARIABLES or name.upper().startswith(_SCRUBBED_PREFIXES) + ) + + +@pytest.fixture(scope="module") +def forward() -> Iterator[ForwardProxy]: + with refusing_forward_proxy() as proxy: + yield proxy + + +@pytest.fixture(scope="module") +def isolated(forward: ForwardProxy, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Gateway]: + directory: Final = tmp_path_factory.mktemp("url-join-host") + with gateway_from_environment() as shared, owned_proxy( + shared, + directory, + { + "HTTP_PROXY": forward.url, + "HTTPS_PROXY": forward.url, + "NO_PROXY": "127.0.0.1,localhost", + "PROXY_BATCH_POLLING_ENABLED": "false", + }, + config=_config(directory), + remove_environment=_scrubbed(), + ) as candidate: + yield candidate + + +def _targets_after(forward: ForwardProxy, call: Callable[[], object]) -> frozenset[str]: + forward.drain() + call() + targets: Final = frozenset(forward.targets()) + assert targets, "the forward proxy saw no outbound request" + return targets + + +def _chat(isolated: Gateway, model: str) -> object: + return isolated.request( + "POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": new_marker()}]} + ) + + +def _messages(isolated: Gateway, model: str, path: str) -> object: + return isolated.request( + "POST", + path, + {"model": model, "max_tokens": 16, "messages": [{"role": "user", "content": new_marker()}]}, + headers={"anthropic-version": "2023-06-01"}, + ) + + +def test_embedding_host_named_embeddings_gets_the_route_appended(isolated: Gateway, forward: ForwardProxy) -> None: + targets: Final = _targets_after( + forward, lambda: isolated.request("POST", "/v1/embeddings", {"model": _EMBEDDINGS, "input": new_marker()}) + ) + assert targets == {"http://embeddings/embeddings"} + + +def test_rerank_host_named_rerank_gets_the_route_appended(isolated: Gateway, forward: ForwardProxy) -> None: + body: Final[dict[str, JsonValue]] = {"model": _RERANK, "query": new_marker(), "documents": ["only document"]} + targets: Final = _targets_after(forward, lambda: isolated.request("POST", "/v1/rerank", body)) + assert targets == {"http://rerank/rerank"} + + +def test_azure_claude_chat_on_a_host_starting_with_anthropic_keeps_the_host( + isolated: Gateway, forward: ForwardProxy +) -> None: + targets: Final = _targets_after(forward, lambda: _chat(isolated, _AZURE_CLAUDE)) + assert targets == {f"http://{_AZURE_HOST}/anthropic/v1/messages"} + + +def test_azure_claude_messages_on_a_host_starting_with_anthropic_keeps_the_host( + isolated: Gateway, forward: ForwardProxy +) -> None: + targets: Final = _targets_after(forward, lambda: _messages(isolated, _AZURE_CLAUDE, "/v1/messages")) + assert targets == {f"http://{_AZURE_HOST}/anthropic/v1/messages"} + + +def test_azure_claude_count_tokens_on_a_bare_host_starting_with_anthropic_adds_the_anthropic_path( + isolated: Gateway, forward: ForwardProxy +) -> None: + targets: Final = _targets_after(forward, lambda: _messages(isolated, _AZURE_COUNT, "/v1/messages/count_tokens")) + assert targets == {f"http://{_AZURE_HOST}/anthropic/v1/messages/count_tokens"} + + +def test_empty_api_base_falls_back_to_the_provider_default_host(isolated: Gateway, forward: ForwardProxy) -> None: + targets: Final = _targets_after(forward, lambda: _chat(isolated, _EMPTY_BASE)) + assert targets == {"api.anthropic.com:443"} diff --git a/tests/integration/providers/test_api_base_url_join_wire.py b/tests/integration/providers/test_api_base_url_join_wire.py new file mode 100644 index 00000000000..42e53ff9f95 --- /dev/null +++ b/tests/integration/providers/test_api_base_url_join_wire.py @@ -0,0 +1,373 @@ +import json +from typing import Final +from urllib.parse import urlsplit + +import anthropic +import openai +import pytest +from integration._support.client import Gateway, Scenario, eventually, list_value, object_value, string_value +from integration._support.database import read_rows +from integration._support.wire import Reply, Request, wire_server +from integration.providers._api_base_trailing_slash_support import ( + KEY, + TRAILING, + TRAILING_IDS, + VECTOR, + Row, + assert_control_answers, + assert_one_request, + body_of, + first_item, + new_marker, + reply_text, + request_json, +) +from integration.providers._api_base_url_join_support import COUNTED_TOKENS, RELEVANCE, join_peer +from pydantic import JsonValue + +_HAIKU: Final = "anthropic.claude-haiku-5-5" +_ANTHROPIC: Final = Row("anthropic/claude-haiku-5-5", "/an", "/v1/messages") +_GEMINI: Final = Row("gemini/gemini-3.8-flash", "/gm/v1beta", "/models/gemini-3.8-flash:generateContent") +_MISTRAL: Final = Row("mistral/mistral-small-latest", "/mi/v1", "/chat/completions") +_BEDROCK_INVOKE: Final = Row(f"bedrock/invoke/{_HAIKU}", "/bi", f"/model/{_HAIKU}/invoke") +_CHAT_ROWS: Final = ( + _ANTHROPIC, + _GEMINI, + _MISTRAL, + Row("databricks/databricks-meta-llama-3-3-70b-instruct", "/db/serving-endpoints", "/chat/completions"), + _BEDROCK_INVOKE, + Row(f"bedrock/converse/{_HAIKU}", "/bc", f"/model/{_HAIKU}/converse"), +) +_STREAM_ROWS: Final = ( + _ANTHROPIC, + Row("gemini/gemini-3.8-flash", "/gm/v1beta", "/models/gemini-3.8-flash:streamGenerateContent?alt=sse"), + _MISTRAL, + Row(f"bedrock/converse/{_HAIKU}", "/bc", f"/model/{_HAIKU}/converse-stream"), +) +_MESSAGES_ROWS: Final = (_ANTHROPIC, _BEDROCK_INVOKE, _GEMINI, _MISTRAL) +_RESPONSES_ROWS: Final = (_ANTHROPIC, _GEMINI) +_EMBEDDING_ROWS: Final = ( + Row("mistral/mistral-embed", "/mi/v1", "/embeddings"), + Row("gemini/gemini-embedding-2", "/gm/v1beta", "/models/gemini-embedding-2:batchEmbedContents"), +) +_TRANSCRIPTION: Final = Row("mistral/voxtral-mini-latest", "/mi/v1", "/audio/transcriptions") +_FULL_ROUTE_SUFFIXES: Final = ("", "/", "?beta=true") +_FULL_ROUTE_IDS: Final = ("exact", "slash", "query") + + +def _ids(rows: tuple[Row, ...]) -> tuple[str, ...]: + return tuple(row.model for row in rows) + + +def _deploy(scenario: Scenario, row: Row, api_base: str) -> str: + if row.model.startswith("bedrock/"): + return scenario.model(model=row.model, api_base=api_base, api_key=KEY, aws_region_name="us-east-1") + return scenario.model(model=row.model, api_base=api_base, api_key=KEY) + + +def _openai(gateway: Gateway) -> openai.OpenAI: + return openai.OpenAI(base_url=str(gateway.client.base_url) + "/v1", api_key=gateway.key, max_retries=0) + + +def _anthropic(gateway: Gateway) -> anthropic.Anthropic: + return anthropic.Anthropic(base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0) + + +def _chat(gateway: Gateway, model: str, marker: str) -> str | None: + completion: Final = _openai(gateway).chat.completions.create( + model=model, messages=[{"role": "user", "content": marker}] + ) + return completion.choices[0].message.content + + +def _message_text(gateway: Gateway, model: str, marker: str) -> list[str]: + message: Final = _anthropic(gateway).messages.create( + model=model, max_tokens=16, messages=[{"role": "user", "content": marker}] + ) + return [block.text for block in message.content if block.type == "text"] + + +def _spend_rows(pattern: str) -> list[dict[str, JsonValue]]: + return read_rows( + 'SELECT request_id, api_base, cache_hit FROM "LiteLLM_SpendLogs" WHERE request_id LIKE %s', (pattern,) + ) + + +def _one_spend_row(pattern: str) -> dict[str, JsonValue]: + (row,) = eventually(lambda: _spend_rows(pattern), lambda rows: len(rows) == 1, seconds=70) + return row + + +def _gemini_contents(marker: str) -> dict[str, JsonValue]: + return {"contents": [{"role": "user", "parts": [{"text": marker}]}]} + + +def _gemini_text(payload: dict[str, JsonValue]) -> JsonValue: + content: Final = object_value(first_item(payload["candidates"])["content"]) + return first_item(content["parts"])["text"] + + +@pytest.mark.parametrize("trailing", TRAILING, ids=TRAILING_IDS) +@pytest.mark.parametrize("row", _CHAT_ROWS, ids=_ids(_CHAT_ROWS)) +def test_chat_completion_reaches_one_single_slash_url(gateway: Gateway, row: Row, trailing: str) -> None: + marker: Final = new_marker() + with wire_server(join_peer) as wire, gateway.scenario() as scenario: + model: Final = _deploy(scenario, row, f"{wire.url}{row.base}{trailing}") + assert _chat(gateway, model, marker) == reply_text(marker) + assert_one_request(wire.drain(), row.target, marker) + + +@pytest.mark.parametrize("trailing", TRAILING, ids=TRAILING_IDS) +@pytest.mark.parametrize("row", _STREAM_ROWS, ids=_ids(_STREAM_ROWS)) +def test_streamed_chat_reaches_one_single_slash_url(gateway: Gateway, row: Row, trailing: str) -> None: + marker: Final = new_marker() + with wire_server(join_peer) as wire, gateway.scenario() as scenario: + model: Final = _deploy(scenario, row, f"{wire.url}{row.base}{trailing}") + stream: Final = _openai(gateway).chat.completions.create( + model=model, messages=[{"role": "user", "content": marker}], stream=True + ) + text: Final = "".join(chunk.choices[0].delta.content or "" for chunk in stream if chunk.choices) + assert text == reply_text(marker) + assert_one_request(wire.drain(), row.target, marker) + + +@pytest.mark.parametrize("trailing", TRAILING, ids=TRAILING_IDS) +@pytest.mark.parametrize("row", _MESSAGES_ROWS, ids=_ids(_MESSAGES_ROWS)) +def test_messages_reaches_one_single_slash_url(gateway: Gateway, row: Row, trailing: str) -> None: + marker: Final = new_marker() + with wire_server(join_peer) as wire, gateway.scenario() as scenario: + model: Final = _deploy(scenario, row, f"{wire.url}{row.base}{trailing}") + assert _message_text(gateway, model, marker) == [reply_text(marker)] + assert_one_request(wire.drain(), row.target, marker) + + +@pytest.mark.parametrize("trailing", TRAILING, ids=TRAILING_IDS) +def test_streamed_messages_reach_one_single_slash_url(gateway: Gateway, trailing: str) -> None: + marker: Final = new_marker() + with wire_server(join_peer) as wire, gateway.scenario() as scenario: + model: Final = _deploy(scenario, _ANTHROPIC, f"{wire.url}/an{trailing}") + with _anthropic(gateway).messages.stream( + model=model, max_tokens=16, messages=[{"role": "user", "content": marker}] + ) as stream: + text: Final = "".join(stream.text_stream) + assert text == reply_text(marker) + assert request_json(assert_one_request(wire.drain(), "/an/v1/messages", marker))["stream"] is True + + +@pytest.mark.parametrize("trailing", TRAILING, ids=TRAILING_IDS) +@pytest.mark.parametrize("row", _RESPONSES_ROWS, ids=_ids(_RESPONSES_ROWS)) +def test_responses_bridge_reaches_one_single_slash_url(gateway: Gateway, row: Row, trailing: str) -> None: + marker: Final = new_marker() + with wire_server(join_peer) as wire, gateway.scenario() as scenario: + model: Final = _deploy(scenario, row, f"{wire.url}{row.base}{trailing}") + assert _openai(gateway).responses.create(model=model, input=marker).output_text == reply_text(marker) + assert_one_request(wire.drain(), row.target, marker) + + +@pytest.mark.parametrize("trailing", TRAILING, ids=TRAILING_IDS) +@pytest.mark.parametrize("row", _EMBEDDING_ROWS, ids=_ids(_EMBEDDING_ROWS)) +def test_embedding_reaches_one_single_slash_url(gateway: Gateway, row: Row, trailing: str) -> None: + marker: Final = new_marker() + with wire_server(join_peer) as wire, gateway.scenario() as scenario: + model: Final = _deploy(scenario, row, f"{wire.url}{row.base}{trailing}") + response: Final = gateway.request("POST", "/v1/embeddings", {"model": model, "input": marker}) + assert first_item(body_of(response)["data"])["embedding"] == list(VECTOR) + assert_one_request(wire.drain(), row.target, marker) + + +@pytest.mark.parametrize("trailing", TRAILING, ids=TRAILING_IDS) +def test_mistral_transcription_reaches_one_single_slash_url(gateway: Gateway, trailing: str) -> None: + marker: Final = new_marker() + with wire_server(join_peer) as wire, gateway.scenario() as scenario: + model: Final = _deploy(scenario, _TRANSCRIPTION, f"{wire.url}/mi/v1{trailing}") + response: Final = gateway.request_multipart( + "/v1/audio/transcriptions", + {"model": model}, + {"file": ("speech.wav", f"RIFF {marker}".encode(), "audio/wav")}, + ) + assert body_of(response)["text"] == reply_text(marker) + assert_one_request(wire.drain(), _TRANSCRIPTION.target, marker) + + +@pytest.mark.parametrize("trailing", TRAILING, ids=TRAILING_IDS) +def test_google_genai_generate_content_reaches_one_single_slash_url(gateway: Gateway, trailing: str) -> None: + marker: Final = new_marker() + with wire_server(join_peer) as wire, gateway.scenario() as scenario: + model: Final = _deploy(scenario, _GEMINI, f"{wire.url}/gm/v1beta{trailing}") + response: Final = gateway.request( + "POST", f"/v1beta/models/{model}:generateContent", _gemini_contents(marker) + ) + assert _gemini_text(body_of(response)) == reply_text(marker) + assert_one_request(wire.drain(), _GEMINI.target, marker) + + +@pytest.mark.parametrize("trailing", TRAILING, ids=TRAILING_IDS) +def test_google_genai_stream_generate_content_reaches_one_single_slash_url(gateway: Gateway, trailing: str) -> None: + marker: Final = new_marker() + with wire_server(join_peer) as wire, gateway.scenario() as scenario: + model: Final = _deploy(scenario, _GEMINI, f"{wire.url}/gm/v1beta{trailing}") + response: Final = gateway.request( + "POST", + f"/v1beta/models/{model}:streamGenerateContent", + _gemini_contents(marker), + params={"alt": "sse"}, + ) + assert response.status_code == 200, response.text + assert marker in response.text, response.text + assert_one_request( + wire.drain(), "/gm/v1beta/models/gemini-3.8-flash:streamGenerateContent?alt=sse", marker + ) + + +@pytest.mark.parametrize("trailing", TRAILING, ids=TRAILING_IDS) +def test_count_tokens_reaches_one_single_slash_url(gateway: Gateway, trailing: str) -> None: + marker: Final = new_marker() + with wire_server(join_peer) as wire, gateway.scenario() as scenario: + model: Final = _deploy(scenario, _ANTHROPIC, f"{wire.url}/an{trailing}") + response: Final = gateway.request( + "POST", "/v1/messages/count_tokens", {"model": model, "messages": [{"role": "user", "content": marker}]} + ) + assert body_of(response)["input_tokens"] == COUNTED_TOKENS + assert_one_request(wire.drain(), "/an/v1/messages/count_tokens", marker) + + +@pytest.mark.parametrize("trailing", TRAILING, ids=TRAILING_IDS) +def test_health_check_reaches_one_single_slash_url(gateway: Gateway, trailing: str) -> None: + with wire_server(join_peer) as wire, gateway.scenario() as scenario: + model: Final = _deploy(scenario, _ANTHROPIC, f"{wire.url}/an{trailing}") + report: Final = body_of(gateway.request("GET", "/health", params={"model": model})) + assert (report["healthy_count"], report["unhealthy_count"]) == (1, 0), report + assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/an/v1/messages")] + + +@pytest.mark.parametrize("trailing", TRAILING, ids=TRAILING_IDS) +def test_spend_log_records_the_single_slash_url_it_called(gateway: Gateway, trailing: str) -> None: + marker: Final = new_marker() + with wire_server(join_peer) as wire, gateway.scenario() as scenario: + model: Final = _deploy(scenario, _ANTHROPIC, f"{wire.url}/an{trailing}") + response: Final = gateway.request( + "POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": marker}]} + ) + identity: Final = string_value(body_of(response)["id"]) + assert_one_request(wire.drain(), "/an/v1/messages", marker) + assert _one_spend_row(identity)["api_base"] == f"{wire.url}/an/v1/messages" + + +def test_cache_hit_on_a_slash_base_skips_the_upstream_and_logs_its_own_row(gateway: Gateway) -> None: + pytest.skip("BUG: a response-cache hit writes two cache_hit spend rows for one chat completion") + marker: Final = new_marker() + with wire_server(join_peer) as wire, gateway.scenario() as scenario: + model: Final = _deploy(scenario, _ANTHROPIC, f"{wire.url}/an/") + body: Final[dict[str, JsonValue]] = {"model": model, "messages": [{"role": "user", "content": marker}]} + first: Final = body_of(gateway.request("POST", "/v1/chat/completions", body)) + second: Final = body_of(gateway.request("POST", "/v1/chat/completions", body)) + assert first_item(second["choices"])["message"] == first_item(first["choices"])["message"] + assert_one_request(wire.drain(), "/an/v1/messages", marker) + identity: Final = string_value(first["id"]) + assert _one_spend_row(identity)["api_base"] == f"{wire.url}/an/v1/messages" + hit: Final = _one_spend_row(f"{identity}_cache_hit%") + assert hit["cache_hit"] == "True", hit + assert "//" not in urlsplit(string_value(hit["api_base"])).path, hit + + +@pytest.mark.parametrize("suffix", _FULL_ROUTE_SUFFIXES, ids=_FULL_ROUTE_IDS) +def test_chat_base_that_already_names_the_route_keeps_its_path_and_query(gateway: Gateway, suffix: str) -> None: + marker: Final = new_marker() + with wire_server(join_peer) as wire, gateway.scenario() as scenario: + model: Final = _deploy(scenario, _ANTHROPIC, f"{wire.url}/an/v1/messages{suffix}") + assert _chat(gateway, model, marker) == reply_text(marker) + assert_one_request(wire.drain(), f"/an/v1/messages{suffix.rstrip('/')}", marker) + + +@pytest.mark.parametrize("suffix", _FULL_ROUTE_SUFFIXES, ids=_FULL_ROUTE_IDS) +def test_messages_base_that_already_names_the_route_keeps_its_path_and_query(gateway: Gateway, suffix: str) -> None: + marker: Final = new_marker() + with wire_server(join_peer) as wire, gateway.scenario() as scenario: + model: Final = _deploy(scenario, _ANTHROPIC, f"{wire.url}/an/v1/messages{suffix}") + assert _message_text(gateway, model, marker) == [reply_text(marker)] + assert_one_request(wire.drain(), f"/an/v1/messages{suffix.rstrip('/')}", marker) + + +@pytest.mark.parametrize("suffix", ("", "/", "?tier=a"), ids=("exact", "slash", "query")) +def test_rerank_base_that_already_names_the_route_keeps_its_path_and_query(gateway: Gateway, suffix: str) -> None: + marker: Final = new_marker() + with wire_server(join_peer) as wire, gateway.scenario() as scenario: + model: Final = _deploy(scenario, Row("cohere/rerank-v3.5", "/co", ""), f"{wire.url}/co/v2/rerank{suffix}") + response: Final = gateway.request( + "POST", "/v1/rerank", {"model": model, "query": marker, "documents": ["only document"]} + ) + assert [object_value(result)["relevance_score"] for result in list_value(body_of(response)["results"])] == [ + RELEVANCE + ] + assert_one_request(wire.drain(), f"/co/v2/rerank{suffix.rstrip('/')}", marker) + + +@pytest.mark.parametrize( + ("row", "trailing"), + ( + (_ANTHROPIC, "//"), + (Row("anthropic/claude-haiku-5-5", f"/an/{'p' * 5000}", "/v1/messages"), "/"), + (_GEMINI, "//"), + (Row("gemini/gemini-3.8-flash", f"/gm/{'p' * 5000}", "/models/gemini-3.8-flash:generateContent"), "/"), + ), + ids=("anthropic-two-slashes", "anthropic-5kb-prefix", "gemini-two-slashes", "gemini-5kb-prefix"), +) +def test_base_with_repeated_slashes_or_a_long_prefix_reaches_one_url(gateway: Gateway, row: Row, trailing: str) -> None: + marker: Final = new_marker() + with wire_server(join_peer) as wire, gateway.scenario() as scenario: + model: Final = _deploy(scenario, row, f"{wire.url}{row.base}{trailing}") + assert _chat(gateway, model, marker) == reply_text(marker) + assert_one_request(wire.drain(), row.target, marker) + + +def _anthropic_error(status: int, marker: str) -> bytes: + return json.dumps( + {"type": "error", "error": {"type": "authentication_error", "message": f"scripted {status} {marker}"}} + ).encode() + + +def _gemini_error(status: int, marker: str) -> bytes: + error: Final = {"code": status, "message": f"scripted {status} {marker}", "status": "NOT_FOUND"} + return json.dumps({"error": error}).encode() + + +@pytest.mark.parametrize( + ("row", "status"), + ((_ANTHROPIC, 401), (_GEMINI, 404)), + ids=("anthropic-401", "gemini-404"), +) +def test_upstream_error_on_a_slash_base_reaches_the_caller_from_one_url( + gateway: Gateway, row: Row, status: int +) -> None: + marker: Final = new_marker() + error: Final = _anthropic_error(status, marker) if row is _ANTHROPIC else _gemini_error(status, marker) + + def respond(_: Request) -> Reply: + return Reply(status=status, body=error) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = _deploy(scenario, row, f"{wire.url}{row.base}/") + response: Final = gateway.request( + "POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": marker}]} + ) + assert response.status_code == status, response.text + assert marker in response.text, response.text + assert_one_request(wire.drain(), row.target, marker) + assert_control_answers(gateway, scenario) + + +def test_streamed_chat_upstream_error_on_a_slash_base_reaches_the_caller(gateway: Gateway) -> None: + marker: Final = new_marker() + + def respond(_: Request) -> Reply: + return Reply(status=401, body=_anthropic_error(401, marker)) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = _deploy(scenario, _ANTHROPIC, f"{wire.url}/an/") + with pytest.raises(openai.AuthenticationError) as raised: + _openai(gateway).chat.completions.create( + model=model, messages=[{"role": "user", "content": marker}], stream=True + ) + assert marker in str(raised.value) + assert_one_request(wire.drain(), "/an/v1/messages", marker) + assert_control_answers(gateway, scenario) diff --git a/tests/unit/litellm_core_utils/test_url_utils.py b/tests/unit/litellm_core_utils/test_url_utils.py index 46bb91959c0..9b29eac208e 100644 --- a/tests/unit/litellm_core_utils/test_url_utils.py +++ b/tests/unit/litellm_core_utils/test_url_utils.py @@ -1,6 +1,8 @@ import asyncio import socket import threading +from collections.abc import Callable +from urllib.parse import urlsplit import httpx import pytest @@ -13,6 +15,8 @@ from litellm.litellm_core_utils.url_utils import ( assert_same_origin, encode_url_path_segment, encode_url_path_segments, + ensure_api_route, + join_api_base, validate_url, ) @@ -571,3 +575,80 @@ async def test_async_safe_get_resolves_dns_off_the_event_loop(monkeypatch): assert response.status_code == 200 assert resolver_threads and all(thread is not loop_thread for thread in resolver_threads) assert released_by_the_loop and all(released_by_the_loop) + + +@pytest.mark.parametrize( + ("api_base", "route", "expected"), + [ + ("https://gw.test/p/", "v1/messages", "https://gw.test/p/v1/messages"), + ("https://gw.test/p//", "v1/messages", "https://gw.test/p/v1/messages"), + ("https://gw.test/p", "v1/messages", "https://gw.test/p/v1/messages"), + ("https://gw.test/p", "/v1/messages", "https://gw.test/p/v1/messages"), + ("https://gw.test/p/", "/v1/messages", "https://gw.test/p/v1/messages"), + ("https://gw.test/p/v1/messages", "v1/messages", "https://gw.test/p/v1/messages/v1/messages"), + ("https://api.test/v1/evals/cancel", "cancel", "https://api.test/v1/evals/cancel/cancel"), + ("http://embeddings", "embeddings", "http://embeddings/embeddings"), + ("http:/", "v1/files", "http://v1/files"), + ("https://gw.test/p/?x=1", "v1/messages", "https://gw.test/p?x=1/v1/messages"), + ], + ids=( + "trailing-slash-base", + "repeated-trailing-slashes", + "base-without-slash", + "route-with-leading-slash", + "slashes-on-both-sides", + "base-already-ends-in-route", + "id-equal-to-the-route", + "host-named-like-the-route", + "host-stripped-as-a-suffix-leaves-only-the-scheme", + "slash-before-the-query-matches-the-base-without-it", + ), +) +def test_join_api_base_always_appends_the_route_after_exactly_one_slash( + api_base: str, route: str, expected: str +) -> None: + assert join_api_base(api_base, route) == expected + + +@pytest.mark.parametrize("api_base", ["https://", "http://", "https:///"]) +@pytest.mark.parametrize("route", ["v1/messages", "models/gemini-flash:generateContent"]) +@pytest.mark.parametrize("join", [join_api_base, ensure_api_route], ids=["join_api_base", "ensure_api_route"]) +def test_a_base_that_is_only_a_scheme_never_turns_the_route_into_the_host( + api_base: str, route: str, join: Callable[[str, str], str] +) -> None: + assert urlsplit(join(api_base, route)).netloc == "" + + +@pytest.mark.parametrize( + ("api_base", "route", "expected"), + [ + ("https://gw.test/p/v1/messages", "v1/messages", "https://gw.test/p/v1/messages"), + ("https://gw.test/p/v1/messages/", "/v1/messages", "https://gw.test/p/v1/messages"), + ("https://gw.test/x/listen", "listen", "https://gw.test/x/listen"), + ("https://gw.test/p/", "v1/messages", "https://gw.test/p/v1/messages"), + ("https://gw.test/p", "/v1/messages", "https://gw.test/p/v1/messages"), + ("https://gw.test/research", "search", "https://gw.test/research/search"), + ("http://embeddings", "embeddings", "http://embeddings/embeddings"), + ("http://embeddings:8080/", "embeddings", "http://embeddings:8080/embeddings"), + ("http:/", "anthropic/v1/messages", "http://anthropic/v1/messages"), + ("https://gw.test/p/v1/messages/?beta=true", "v1/messages", "https://gw.test/p/v1/messages?beta=true"), + ("https://gw.test/p/v1/messages?beta=true", "v1/messages", "https://gw.test/p/v1/messages?beta=true"), + ], + ids=( + "path-already-ends-in-route", + "path-ends-in-route-before-trailing-slash", + "last-segment-is-the-route", + "trailing-slash-base", + "route-with-leading-slash", + "path-ends-in-lookalike-suffix", + "host-named-like-the-route", + "host-with-port-named-like-the-route", + "host-stripped-as-a-suffix-leaves-only-the-scheme", + "route-path-with-a-slash-before-the-query", + "route-path-before-the-query", + ), +) +def test_ensure_api_route_appends_the_route_only_when_the_url_path_lacks_it( + api_base: str, route: str, expected: str +) -> None: + assert ensure_api_route(api_base, route) == expected diff --git a/tests/unit/llms/azure/realtime/test_azure_realtime_handler.py b/tests/unit/llms/azure/realtime/test_azure_realtime_handler.py index 5c1ff1b46a3..f1db76c1e21 100644 --- a/tests/unit/llms/azure/realtime/test_azure_realtime_handler.py +++ b/tests/unit/llms/azure/realtime/test_azure_realtime_handler.py @@ -740,3 +740,15 @@ async def test_arealtime_forwards_deployment_azure_ad_token(monkeypatch, no_ambi ) assert mock_async_realtime.call_args.kwargs["azure_ad_token"] == "deployment-entra-token" + + +@pytest.mark.parametrize("protocol", [None, "GA"]) +def test_construct_url_with_a_trailing_slash_api_base_matches_the_url_without_it(protocol): + from litellm.llms.azure.realtime.handler import AzureOpenAIRealtime + + def url_for(api_base: str) -> str: + return AzureOpenAIRealtime().construct_url( + api_base=api_base, model="gpt-realtime", api_version="2024-10-01-preview", realtime_protocol=protocol + ) + + assert url_for("https://gw.example.com/azure/") == url_for("https://gw.example.com/azure") diff --git a/tests/unit/llms/azure/test_azure_common_utils.py b/tests/unit/llms/azure/test_azure_common_utils.py index b3bd1b8743a..1dc7ff74323 100644 --- a/tests/unit/llms/azure/test_azure_common_utils.py +++ b/tests/unit/llms/azure/test_azure_common_utils.py @@ -1702,6 +1702,24 @@ def test_azure_v1_api_uses_openai_client(api_version): ), f"base_url should contain /openai/v1/, got {async_client.base_url}" +@pytest.mark.parametrize( + "api_base", + [ + "https://test.openai.azure.com", + "https://test.openai.azure.com/", + "https://test.openai.azure.com/openai/v1", + "https://test.openai.azure.com/openai/v1/", + ], +) +@pytest.mark.parametrize("is_async", [False, True]) +def test_azure_v1_api_client_targets_openai_v1_once_whatever_the_api_base_ends_in(api_base: str, is_async: bool): + client: Final = BaseAzureLLM().get_azure_openai_client( + api_key="test-key", api_base=api_base, api_version="v1", _is_async=is_async + ) + + assert str(client.base_url) == "https://test.openai.azure.com/openai/v1/" + + @pytest.mark.parametrize("api_version", ["v1", "latest", "preview"]) def test_azure_v1_api_with_azure_ad_token_provider(api_version): """ diff --git a/tests/unit/llms/azure_ai/agents/test_handler.py b/tests/unit/llms/azure_ai/agents/test_handler.py index f92f5359c49..71466bb172c 100644 --- a/tests/unit/llms/azure_ai/agents/test_handler.py +++ b/tests/unit/llms/azure_ai/agents/test_handler.py @@ -22,10 +22,10 @@ def _httpx_only_transport(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") -def _assistants_api(thread_id: str, answer: str) -> respx.MockRouter: - thread_path: Final = f"/api/projects/litellm-test/threads/{thread_id}" +def _assistants_api(thread_id: str, answer: str, project_path: str = "/api/projects/litellm-test") -> respx.MockRouter: + thread_path: Final = f"{project_path}/threads/{thread_id}" router: Final = respx.mock(assert_all_called=False) - router.post(path="/api/projects/litellm-test/threads").respond( + router.post(path=f"{project_path}/threads").respond( 200, json={"id": thread_id, "object": "thread", "created_at": 1767225600, "metadata": {}} ) router.post(path=f"{thread_path}/messages").respond( @@ -116,6 +116,23 @@ async def test_acompletion_creates_thread_posts_message_runs_and_returns_the_rep assert response._hidden_params["thread_id"] == "thread_new" +@pytest.mark.asyncio +async def test_acompletion_creates_the_thread_under_a_project_named_threads() -> None: + router: Final = _assistants_api("thread_new", "hello", project_path="/api/projects/threads") + + with router: + await litellm.acompletion( + model=f"azure_ai/agents/{AGENT_ID}", + messages=[{"role": "user", "content": "hi"}], + api_base="https://agents-test.services.ai.azure.com/api/projects/threads", + api_key="azure-ad-token", + stream=False, + ) + recorded: Final = _recorded(router) + + assert recorded[0][:2] == ("POST", f"/api/projects/threads/threads?api-version={API_VERSION}") + + @pytest.mark.asyncio async def test_handler_acompletion_with_thread_id_reuses_that_thread() -> None: router: Final = _assistants_api("thread_existing", "Your name is Alice.") diff --git a/tests/unit/llms/azure_ai/claude/test_azure_anthropic_count_tokens_transformation.py b/tests/unit/llms/azure_ai/claude/test_azure_anthropic_count_tokens_transformation.py index 4b317cff975..8971bd31a60 100644 --- a/tests/unit/llms/azure_ai/claude/test_azure_anthropic_count_tokens_transformation.py +++ b/tests/unit/llms/azure_ai/claude/test_azure_anthropic_count_tokens_transformation.py @@ -4,8 +4,9 @@ Tests for Azure AI Anthropic CountTokens transformation. Verifies that the CountTokens API uses the correct authentication headers. """ +from typing import Final - +import pytest from litellm.llms.azure_ai.anthropic.count_tokens.transformation import ( AzureAIAnthropicCountTokensConfig, @@ -105,3 +106,22 @@ class TestAzureAIAnthropicCountTokensConfig: endpoint == "https://my-resource.services.ai.azure.com/anthropic/v1/messages/count_tokens" ) + + def test_get_count_tokens_endpoint_for_a_resource_named_anthropic(self): + config: Final = AzureAIAnthropicCountTokensConfig() + + endpoint: Final = config.get_count_tokens_endpoint("https://anthropic-res.services.ai.azure.com") + + assert endpoint == "https://anthropic-res.services.ai.azure.com/anthropic/v1/messages/count_tokens" + + @pytest.mark.parametrize( + "api_base", + [ + "https://my-resource.services.ai.azure.com/anthropic/v1/messages", + "https://my-resource.services.ai.azure.com/anthropic/v1/messages/", + ], + ) + def test_get_count_tokens_endpoint_sits_next_to_the_messages_url(self, api_base: str): + endpoint: Final = AzureAIAnthropicCountTokensConfig().get_count_tokens_endpoint(api_base) + + assert endpoint == "https://my-resource.services.ai.azure.com/anthropic/v1/messages/count_tokens" diff --git a/tests/unit/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py b/tests/unit/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py index 38483b8a1fa..5e0395f84dc 100644 --- a/tests/unit/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py +++ b/tests/unit/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py @@ -198,6 +198,17 @@ class TestAzureAnthropicMessagesConfig: assert url == "https://test.services.ai.azure.com/anthropic/v1/messages" + def test_get_complete_url_adds_the_anthropic_path_for_a_host_named_anthropic(self): + url: Final = AzureAnthropicMessagesConfig().get_complete_url( + api_base="https://anthropic-res.services.ai.azure.com", + api_key="test-api-key", + model="claude-sonnet-4-5", + optional_params={}, + litellm_params={}, + ) + + assert url == "https://anthropic-res.services.ai.azure.com/anthropic/v1/messages" + def test_get_complete_url_raises_error_when_api_base_missing(self): """Test get_complete_url raises error when api_base is None""" config = AzureAnthropicMessagesConfig() diff --git a/tests/unit/llms/bedrock/chat/agentcore/test_agentcore_transformation.py b/tests/unit/llms/bedrock/chat/agentcore/test_agentcore_transformation.py index 053ba477212..413061a732d 100644 --- a/tests/unit/llms/bedrock/chat/agentcore/test_agentcore_transformation.py +++ b/tests/unit/llms/bedrock/chat/agentcore/test_agentcore_transformation.py @@ -1269,3 +1269,17 @@ def test_agentcore_runtime_response_becomes_assistant_message( assert json.loads(route.calls[0].request.content)["prompt"] == "Explain machine learning in simple terms" assert response.choices[0].message.content == "Machine learning lets computers learn from data." assert response.choices[0].finish_reason == "stop" + + +def test_get_complete_url_sends_a_non_string_runtime_endpoint_to_that_endpoint(): + vpc_endpoint: Final = "https://vpce-1.bedrock-agentcore.us-west-2.vpce.amazonaws.com" + + url: Final = AmazonAgentCoreConfig().get_complete_url( + api_base=None, + api_key=None, + model="agentcore/arn:aws:bedrock-agentcore:us-west-2:111122223333:runtime/my-agent", + optional_params={"aws_bedrock_runtime_endpoint": httpx.URL(f"{vpc_endpoint}/")}, + litellm_params={}, + ) + + assert url.startswith(f"{vpc_endpoint}/runtimes/") diff --git a/tests/unit/llms/openai/evals/test_openai_evals_transformation.py b/tests/unit/llms/openai/evals/test_openai_evals_transformation.py index f359f214f97..ed8f86d4727 100644 --- a/tests/unit/llms/openai/evals/test_openai_evals_transformation.py +++ b/tests/unit/llms/openai/evals/test_openai_evals_transformation.py @@ -264,6 +264,17 @@ def test_transform_cancel_eval_request(config: OpenAIEvalsConfig): assert request_body == {} +def test_transform_cancel_eval_request_appends_cancel_to_an_eval_id_named_cancel(config: OpenAIEvalsConfig) -> None: + eval_url, _ = config.transform_get_eval_request( + eval_id="cancel", api_base="https://api.openai.com", litellm_params=GenericLiteLLMParams(), headers={} + ) + cancel_url, _, _ = config.transform_cancel_eval_request( + eval_id="cancel", api_base="https://api.openai.com", litellm_params=GenericLiteLLMParams(), headers={} + ) + + assert cancel_url == f"{eval_url}/cancel" + + def test_transform_cancel_eval_response(config: OpenAIEvalsConfig): """Test transformation of cancel eval response""" response = httpx.Response( diff --git a/tests/unit/llms/test_api_base_trailing_slash_routes.py b/tests/unit/llms/test_api_base_trailing_slash_routes.py new file mode 100644 index 00000000000..ce8cecf11a5 --- /dev/null +++ b/tests/unit/llms/test_api_base_trailing_slash_routes.py @@ -0,0 +1,119 @@ +import asyncio +import contextlib +from collections.abc import Callable, Iterator, Mapping +from types import MappingProxyType +from typing import Final + +import pytest +import respx + +import litellm +from litellm.evals.main import cancel_run, create_run, get_run + +FAKE_KEY: Final = "sk-fake" +FAKE_AWS: Final[Mapping[str, str]] = MappingProxyType( + {"aws_access_key_id": "AKIAFAKE", "aws_secret_access_key": "fake-secret", "aws_region_name": "us-east-1"} +) +MESSAGES: Final = ({"role": "user", "content": "hi"},) +GATEWAY_PATH: Final = "/gw/" + + +GATEWAY_HOST: Final = "gateway.test" + + +@pytest.fixture +def received_paths() -> Iterator[Callable[[], tuple[str, ...]]]: + with respx.mock(assert_all_called=False) as router: + route: Final = router.route(host=GATEWAY_HOST).respond( + 404, json={"error": {"message": "not found", "code": 404}} + ) + yield lambda: tuple(call.request.url.raw_path.decode().split("?", 1)[0] for call in route.calls) + + +def _drain(result: object) -> None: + if isinstance(result, litellm.CustomStreamWrapper): + for _ in result: + pass + + +def _completion(model: str, stream: bool = False, **kwargs: object) -> Callable[[str], object]: + return lambda api_base: _drain( + litellm.completion(model=model, messages=list(MESSAGES), api_base=api_base, stream=stream, timeout=5, **kwargs) + ) + + +def _embedding(model: str, **kwargs: object) -> Callable[[str], object]: + return lambda api_base: litellm.embedding(model=model, input=["hi"], api_base=api_base, timeout=5, **kwargs) + + +ROUTES: Final[Mapping[str, Callable[[str], object]]] = MappingProxyType( + { + "anthropic-chat": _completion("anthropic/claude-opus-5-5", api_key=FAKE_KEY), + "bedrock-converse": _completion("bedrock/converse/anthropic.claude-sonnet-5-5", **FAKE_AWS), + "bedrock-converse-stream": _completion("bedrock/converse/anthropic.claude-sonnet-5-5", stream=True, **FAKE_AWS), + "gemini-chat": _completion("gemini/gemini-pro-latest", api_key=FAKE_KEY), + "gemini-chat-stream": _completion("gemini/gemini-pro-latest", stream=True, api_key=FAKE_KEY), + "mistral-chat": _completion("mistral/mistral-large-latest", api_key=FAKE_KEY), + "gigachat-chat": _completion("gigachat/GigaChat-2-Max", api_key=FAKE_KEY), + "bedrock-titan-embedding": _embedding("bedrock/amazon.titan-embed-text-v2:0", **FAKE_AWS), + "gemini-embedding": _embedding("gemini/gemini-embedding-2", api_key=FAKE_KEY), + "gigachat-embedding": _embedding("gigachat/Embeddings-2", api_key=FAKE_KEY), + "volcengine-embedding": _embedding("volcengine/doubao-embedding-text-240715", api_key=FAKE_KEY), + "ollama-embedding": _embedding("ollama/nomic-embed-text"), + "oobabooga-chat": _completion("oobabooga/local-model"), + "oobabooga-embedding": _embedding("oobabooga/local-model"), + "bedrock-rerank": lambda api_base: litellm.rerank( + model="bedrock/arn:aws:bedrock:us-west-2::foundation-model/amazon.rerank-v1:0", + query="hi", + documents=["a"], + api_base=api_base, + **FAKE_AWS, + ), + "bedrock-image-generation": lambda api_base: litellm.image_generation( + model="bedrock/amazon.nova-canvas-v1:0", prompt="a cat", api_base=api_base, timeout=5, **FAKE_AWS + ), + "gemini-interaction-get": lambda api_base: litellm.interactions.get( + interaction_id="interaction-1", api_base=api_base, api_key=FAKE_KEY + ), + "gemini-interaction-delete": lambda api_base: litellm.interactions.delete( + interaction_id="interaction-1", api_base=api_base, api_key=FAKE_KEY + ), + "gemini-interaction-cancel": lambda api_base: litellm.interactions.cancel( + interaction_id="interaction-1", api_base=api_base, api_key=FAKE_KEY + ), + "openai-eval-run-create": lambda api_base: create_run( + eval_id="eval-1", data_source={"type": "jsonl"}, api_base=api_base, api_key=FAKE_KEY + ), + "openai-eval-run-get": lambda api_base: get_run( + eval_id="eval-1", run_id="run-1", api_base=api_base, api_key=FAKE_KEY + ), + "openai-eval-run-cancel": lambda api_base: cancel_run( + eval_id="eval-1", run_id="run-1", api_base=api_base, api_key=FAKE_KEY + ), + "edenai-model-list": lambda api_base: litellm.get_valid_models( + check_provider_endpoint=True, custom_llm_provider="edenai", api_key=FAKE_KEY, api_base=api_base + ), + "gemini-rag-ingest": lambda api_base: asyncio.run( + litellm.aingest( + ingest_options={ + "vector_store": {"custom_llm_provider": "gemini", "api_base": api_base, "api_key": FAKE_KEY} + }, + file_data=("notes.txt", b"hello", "text/plain"), + ) + ), + } +) + + +@pytest.mark.parametrize("route", ROUTES.values(), ids=ROUTES.keys()) +def test_api_base_with_a_trailing_slash_reaches_the_gateway_path_with_single_slashes( + received_paths: Callable[[], tuple[str, ...]], route: Callable[[str], object], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("GIGACHAT_ACCESS_TOKEN", FAKE_KEY) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + with contextlib.suppress(Exception): + route(f"http://{GATEWAY_HOST}{GATEWAY_PATH}") + + paths: Final = received_paths() + assert paths, "the route never reached the api_base" + assert all(path.startswith(GATEWAY_PATH) and "//" not in path for path in paths), paths diff --git a/tests/unit/llms/test_get_complete_url_trailing_slash.py b/tests/unit/llms/test_get_complete_url_trailing_slash.py new file mode 100644 index 00000000000..9118498c2c2 --- /dev/null +++ b/tests/unit/llms/test_get_complete_url_trailing_slash.py @@ -0,0 +1,278 @@ +import copy +import importlib +import inspect +import itertools +from collections.abc import Mapping +from dataclasses import dataclass, field +from pathlib import Path +from types import MappingProxyType, ModuleType +from typing import Final +from urllib.parse import urlsplit, urlunsplit + +import pytest +from _pytest.mark.structures import ParameterSet + +import litellm +from litellm.llms.elevenlabs.text_to_speech.transformation import ElevenLabsTextToSpeechConfig +from litellm.llms.meta.realtime.transformation import MUSE_MODEL +from litellm.llms.openai_like.dynamic_config import create_config_class, create_responses_config_class +from litellm.llms.openai_like.json_loader import SimpleProviderConfig + +API_BASE_WITH_TRAILING_SLASH: Final = "https://gw.example.com/p/" +API_BASE_WITHOUT_TRAILING_SLASH: Final = API_BASE_WITH_TRAILING_SLASH.rstrip("/") +LLMS_ROOT: Final = Path(litellm.__file__).parent / "llms" +PACKAGE_ROOT: Final = LLMS_ROOT.parent.parent + +ARGUMENTS_BY_NAME: Final[Mapping[str, object]] = MappingProxyType( + { + "api_key": "sk-test", + "model": "test-model", + "optional_params": {}, + "litellm_params": {"api_key": "sk-test", "vertex_project": "project", "vertex_location": "us-central1"}, + "stream": False, + "endpoint": "v1/messages", + "request_query_params": None, + "data": None, + "api_version": "2024-01-01", + "vector_store_id": "vs_test", + } +) + + +SAMPLE_JSON_PROVIDER: Final = SimpleProviderConfig( + "json-provider", {"base_url": "https://json-provider.example.com", "api_key_env": "UNSET_KEY"} +) +FACTORY_BUILDERS: Final[tuple[type, ...]] = ( + create_config_class(SAMPLE_JSON_PROVIDER), + create_responses_config_class(SAMPLE_JSON_PROVIDER), +) + + +async def _resolve_access_token() -> str: + return "token" + + +@dataclass(frozen=True, slots=True) +class BuilderInputs: + constructor: Mapping[str, object] = field(default_factory=lambda: MappingProxyType({})) + call: Mapping[str, object] = field(default_factory=lambda: MappingProxyType({})) + api_base_option: str | None = None + + +INPUTS_BY_BUILDER: Final[Mapping[str, BuilderInputs]] = MappingProxyType( + { + "anthropic.skills.transformation.AnthropicSkillsConfig": BuilderInputs(call={"endpoint": "skills"}), + "azure_ai.image_edit.transformation.AzureFoundryFluxImageEditConfig": BuilderInputs( + call={"litellm_params": {"api_version": "2025-04-01-preview"}} + ), + "bedrock.chat.agentcore.transformation.AmazonAgentCoreConfig": BuilderInputs( + call={"model": "agentcore/arn:aws:bedrock-agentcore:us-west-2:123456789012:runtime/test-runtime"}, + api_base_option="aws_bedrock_runtime_endpoint", + ), + "bedrock.chat.invoke_agent.transformation.AmazonInvokeAgentConfig": BuilderInputs( + call={"model": "agent/AGENT_ID/ALIAS_ID", "optional_params": {"sessionID": "session-id"}} + ), + "black_forest_labs.image_edit.transformation.BlackForestLabsImageEditConfig": BuilderInputs( + call={"model": "flux-kontext-pro"} + ), + "black_forest_labs.image_generation.transformation.BlackForestLabsImageGenerationConfig": BuilderInputs( + call={"model": "flux-pro-1.1"} + ), + "elevenlabs.text_to_speech.transformation.ElevenLabsTextToSpeechConfig": BuilderInputs( + call={"litellm_params": {ElevenLabsTextToSpeechConfig.ELEVENLABS_VOICE_ID_KEY: "voice-id"}} + ), + "meta.realtime.transformation.MetaRealtimeConfig": BuilderInputs(call={"model": MUSE_MODEL}), + "openai.evals.transformation.OpenAIEvalsConfig": BuilderInputs(call={"endpoint": "evals"}), + "openai_like.messages.transformation.JSONProviderAnthropicMessagesConfig": BuilderInputs( + constructor={"provider": SAMPLE_JSON_PROVIDER} + ), + "predibase.chat.transformation.PredibaseConfig": BuilderInputs( + call={"litellm_params": {"predibase_tenant_id": "tenant"}} + ), + "ragflow.chat.transformation.RAGFlowConfig": BuilderInputs(call={"model": "ragflow/chat/chat-id/model"}), + "scaledown.chat.transformation.ScaleDownChatConfig": BuilderInputs(call={"model": "classify"}), + "vertex_ai.audio_transcription.realtime_transformation.VertexChirpRealtimeConfig": BuilderInputs( + constructor={"resolve_access_token": _resolve_access_token, "project": "project", "location": None}, + call={"model": "chirp_3"}, + ), + "vertex_ai.interactions.transformation.VertexAIInteractionsConfig": BuilderInputs( + constructor={"mint_access_token": lambda _credentials, _project: ("token", "project")} + ), + "vertex_ai.realtime.transformation.VertexAIRealtimeConfig": BuilderInputs( + constructor={"access_token": "token", "project": "project", "location": "us-central1"} + ), + "vertex_ai.text_to_speech.transformation.VertexAILyriaTextToSpeechConfig": BuilderInputs( + call={"model": "lyria-002"} + ), + } +) + +UNCALLABLE_BUILDERS: Final[Mapping[str, str]] = MappingProxyType( + { + "base_llm.search.transformation.BaseSearchConfig": "raises-NotImplementedError-for-providers-to-override", + "bedrock.image_edit.amazon_nova_canvas_image_edit_transformation.BedrockAmazonNovaCanvasImageEditConfig": ( + "raises-because-BedrockImageEdit-builds-the-url" + ), + "sap.chat.transformation.GenAIHubOrchestrationConfig": "url-comes-from-a-live-SAP-deployment-lookup", + "sap.embed.transformation.GenAIHubEmbeddingConfig": "url-comes-from-a-live-SAP-deployment-lookup", + "triton.completion.transformation.TritonConfig": "api_base-must-be-the-full-generate-or-infer-url", + "vertex_ai.rerank.transformation.VertexAIRerankConfig": "ignores-api_base-and-needs-Vertex-credentials", + } +) + + +API_BASE_IS_THE_FULL_URL: Final[Mapping[str, str]] = MappingProxyType( + { + builder_id: "sends-the-configured-api_base-as-the-request-url" + for builder_id in ( + "base_llm.audio_transcription.transformation.BaseAudioTranscriptionConfig", + "base_llm.chat.transformation.BaseConfig", + "base_llm.completion.transformation.BaseTextCompletionConfig", + "base_llm.embedding.transformation.BaseEmbeddingConfig", + "base_llm.image_generation.transformation.BaseImageGenerationConfig", + "base_llm.image_variations.transformation.BaseImageVariationConfig", + "base_llm.vector_store.transformation.BaseDirectVectorStoreConfig", + "bedrock.realtime.transformation.BedrockRealtimeConfig", + "bedrock.search.transformation.AgentCoreSearchConfig", + "brave.search.transformation.BraveSearchConfig", + "cohere.chat.v2_transformation.CohereV2ChatConfig", + "cohere.embed.transformation.CohereEmbeddingConfig", + "dashscope.image_generation.transformation.DashScopeImageGenerationConfig", + "dataforseo.search.transformation.DataForSEOSearchConfig", + "datarobot.chat.transformation.DataRobotConfig", + "duckduckgo.search.transformation.DuckDuckGoSearchConfig", + "google_pse.search.transformation.GooglePSESearchConfig", + "searchapi.search.transformation.SearchAPIConfig", + "tinyfish.search.transformation.TinyfishSearchConfig", + "vertex_ai.text_to_speech.transformation.VertexAITextToSpeechConfig", + "vertex_ai.vertex_ai_partner_models.anthropic.experimental_pass_through.transformation." + "VertexAIPartnerModelsAnthropicMessagesConfig", + ) + } +) + + +HOST_NAMED_LIKE_A_STRIPPED_SUFFIX_IS_NOT_KEPT: Final[Mapping[str, str]] = MappingProxyType( + { + "deepinfra.rerank.transformation.DeepinfraRerankConfig": "removes-every-openai-substring-on-main-too", + "minimax.videos.transformation.MinimaxVideoConfig": "strips-surface-suffixes-to-the-scheme-on-main-too", + "volcengine.videos.transformation.VolcEngineVideoConfig": "strips-surface-suffixes-to-the-scheme-on-main-too", + "vertex_ai.audio_transcription.realtime_transformation.VertexChirpRealtimeConfig": "returns-a-grpc-host", + "azure.realtime.http_transformation.AzureRealtimeHTTPConfig": "replaces-the-host-with-a-route-segment-on-main-too", + "bedrock.responses.transformation.BedrockOpenAIResponsesConfig": "replaces-the-host-with-a-route-segment-on-main-too", + "bedrock_mantle.chat.claude_transformation.BedrockMantleClaudeChatConfig": "replaces-the-host-with-a-route-segment-on-main-too", + "bedrock_mantle.messages.transformation.BedrockMantleAnthropicMessagesConfig": "replaces-the-host-with-a-route-segment-on-main-too", + "bedrock_mantle.responses.transformation.BedrockMantleResponsesAPIConfig": "replaces-the-host-with-a-route-segment-on-main-too", + "deepseek.messages.transformation.DeepSeekAnthropicMessagesConfig": "replaces-the-host-with-a-route-segment-on-main-too", + "ragflow.chat.transformation.RAGFlowConfig": "replaces-the-host-with-a-route-segment-on-main-too", + "scaledown.chat.transformation.ScaleDownChatConfig": "replaces-the-host-with-a-route-segment-on-main-too", + } +) + + +def _builder_id(cls: type) -> str: + return f"{cls.__module__.removeprefix('litellm.llms.')}.{cls.__qualname__}" + + +def _module_name(path: Path) -> str: + return ".".join(path.relative_to(PACKAGE_ROOT).with_suffix("").parts) + + +def _builder_modules() -> tuple[ModuleType, ...]: + builder_paths: Final = tuple(path for path in LLMS_ROOT.rglob("*.py") if "def get_complete_url" in path.read_text()) + return tuple(importlib.import_module(_module_name(path)) for path in sorted(builder_paths)) + + +def _defined_builders(module: ModuleType) -> tuple[type, ...]: + classes: Final = (*(cls for _, cls in inspect.getmembers(module, inspect.isclass)), *FACTORY_BUILDERS) + return tuple(cls for cls in classes if cls.__module__ == module.__name__ and "get_complete_url" in vars(cls)) + + +def _all_builders() -> tuple[type, ...]: + builders: Final = itertools.chain.from_iterable(_defined_builders(module) for module in _builder_modules()) + return tuple(sorted(builders, key=_builder_id)) + + +def _skip_reason(cls: type) -> str | None: + if getattr(vars(cls)["get_complete_url"], "__isabstractmethod__", False): + return "abstract-get_complete_url" + return UNCALLABLE_BUILDERS.get(_builder_id(cls)) + + +def _builder_param(cls: type) -> ParameterSet: + reason: Final = _skip_reason(cls) + if reason is None: + return pytest.param(cls, id=_builder_id(cls)) + return pytest.param(cls, id=f"{_builder_id(cls)}-skipped-{reason}", marks=pytest.mark.skip(reason=reason)) + + +def _unimplemented(*_args: object, **_kwargs: object) -> object: + raise NotImplementedError + + +def _instantiable(cls: type) -> type: + abstract_methods: Final[frozenset[str]] = getattr(cls, "__abstractmethods__", frozenset()) + if not abstract_methods: + return cls + return type(cls.__name__, (cls,), dict.fromkeys(abstract_methods, _unimplemented)) + + +def _call_arguments(builder: object, inputs: BuilderInputs, api_base: str) -> dict[str, object]: + parameters: Final = inspect.signature(getattr(builder, "get_complete_url")).parameters.values() + named: Final = tuple(p for p in parameters if p.kind not in (p.VAR_POSITIONAL, p.VAR_KEYWORD)) + api_base_options: Final = ( + {} if inputs.api_base_option is None else {"optional_params": {inputs.api_base_option: api_base}} + ) + values: Final = {**ARGUMENTS_BY_NAME, "api_base": api_base, **inputs.call, **api_base_options} + missing: Final = tuple(p.name for p in named if p.default is p.empty and p.name not in values) + assert not missing, f"add a value for {missing} to ARGUMENTS_BY_NAME or INPUTS_BY_BUILDER" + return {p.name: copy.deepcopy(values[p.name]) for p in named if p.name in values} + + +def _built_url(cls: type, api_base: str) -> str: + inputs: Final = INPUTS_BY_BUILDER.get(_builder_id(cls), BuilderInputs()) + builder: Final = _instantiable(cls)(**inputs.constructor) + result: Final = getattr(builder, "get_complete_url")(**_call_arguments(builder, inputs, api_base)) + return str(result[0] if isinstance(result, tuple) else result) + + +def _without_one_trailing_path_slash(url: str) -> str: + parts: Final = urlsplit(url) + return urlunsplit(parts._replace(path=parts.path.removesuffix("/"))) + + +@pytest.mark.parametrize("builder_class", [_builder_param(cls) for cls in _all_builders()]) +def test_get_complete_url_keeps_one_slash_after_a_trailing_slash_api_base(builder_class: type) -> None: + url: Final = _built_url(builder_class, API_BASE_WITH_TRAILING_SLASH) + + assert "//" not in urlsplit(url).path, url + + +@pytest.mark.parametrize("builder_class", [_builder_param(cls) for cls in _all_builders()]) +def test_get_complete_url_sends_the_same_url_with_or_without_a_trailing_slash(builder_class: type) -> None: + with_slash: Final = _built_url(builder_class, API_BASE_WITH_TRAILING_SLASH) + without_slash: Final = _built_url(builder_class, API_BASE_WITHOUT_TRAILING_SLASH) + + if _builder_id(builder_class) in API_BASE_IS_THE_FULL_URL: + assert _without_one_trailing_path_slash(with_slash) == _without_one_trailing_path_slash(without_slash) + else: + assert with_slash == without_slash + + +def _host_preserving_builder_param(cls: type) -> ParameterSet: + reason: Final = HOST_NAMED_LIKE_A_STRIPPED_SUFFIX_IS_NOT_KEPT.get(_builder_id(cls)) + if reason is None: + return _builder_param(cls) + return pytest.param(cls, id=f"{_builder_id(cls)}-skipped-{reason}", marks=pytest.mark.skip(reason=reason)) + + +@pytest.mark.parametrize("api_base", ["https://v1", "https://api/v1", "https://openai/v1/", "https://anthropic/v1"]) +@pytest.mark.parametrize("builder_class", [_host_preserving_builder_param(cls) for cls in _all_builders()]) +def test_get_complete_url_keeps_a_host_named_like_the_suffix_it_strips(builder_class: type, api_base: str) -> None: + url: Final = _built_url(builder_class, api_base) + follows_the_api_base_host: Final = ( + urlsplit(_built_url(builder_class, "https://gateway.test/v1")).netloc == "gateway.test" + ) + + assert urlsplit(url).netloc, url + assert not follows_the_api_base_host or urlsplit(url).netloc == urlsplit(api_base).netloc, url diff --git a/tests/unit/llms/test_predibase_transformation.py b/tests/unit/llms/test_predibase_transformation.py index 20d547c3b63..a02840ed4ef 100644 --- a/tests/unit/llms/test_predibase_transformation.py +++ b/tests/unit/llms/test_predibase_transformation.py @@ -1,3 +1,4 @@ +from typing import Final from unittest.mock import AsyncMock, Mock import httpx @@ -84,6 +85,19 @@ def test_predibase_get_complete_url_stream_and_non_stream(): assert stream_url.endswith("/generate_stream") +@pytest.mark.parametrize(("model", "stream"), [("generate", False), ("generate_stream", True)]) +def test_predibase_get_complete_url_appends_the_route_to_a_model_named_like_it(model: str, stream: bool) -> None: + url: Final = PredibaseConfig().get_complete_url( + api_base="https://serving.example.com", + api_key="test-key", + model=model, + optional_params={"stream": stream}, + litellm_params={"predibase_tenant_id": "tenant-123"}, + ) + + assert url.endswith(f"/llms/{model}/{model}") + + def test_predibase_get_complete_url_missing_tenant_id(): config = PredibaseConfig() diff --git a/tests/unit/llms/vertex_ai/test_vertex_llm_base.py b/tests/unit/llms/vertex_ai/test_vertex_llm_base.py index 17a92b31284..dd2b86b6f81 100644 --- a/tests/unit/llms/vertex_ai/test_vertex_llm_base.py +++ b/tests/unit/llms/vertex_ai/test_vertex_llm_base.py @@ -977,6 +977,22 @@ class TestVertexBase: assert result_url == f"{gateway_api_base}:embedContent" + def test_check_custom_proxy_vertex_api_base_with_path_and_trailing_slash_appends_endpoint_once(self): + gateway_api_base: Final = "https://gateway.example.com/vertex/v1/projects/p/locations/us-central1/endpoints/123" + + _, result_url = VertexBase()._check_custom_proxy( + api_base=f"{gateway_api_base}/", + custom_llm_provider="vertex_ai", + gemini_api_key=None, + endpoint="predict", + stream=None, + auth_header="Bearer token123", + url="https://us-central1-aiplatform.googleapis.com/v1/projects/p/locations/us-central1/endpoints/123:predict", + model="123", + ) + + assert result_url == f"{gateway_api_base}:predict" + def test_check_custom_proxy_vertex_api_base_with_version_path_grafts_default_path(self): vertex_base = VertexBase() From a8113b0dd721804c752a6f1952367fb642205224 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 21 Oct 2026 08:50:19 +0000 Subject: [PATCH 09/14] fix(responses)!: round-trip Anthropic compaction blocks as compaction items (#46754) * fix(responses): round-trip Anthropic compaction blocks as compaction items /v1/responses on Anthropic routes now returns Claude's compaction block as a compaction output item (and streams it as output_item.added/done), and a compaction input item is replayed to Claude as its native compaction block, so the next turn reuses the summary instead of compacting again. Other routes receive the summary as assistant text. The Anthropic stream handler now keeps the summary text that arrives in compaction_delta. Fixes #41456 * fix(responses): replay compaction blocks natively only when the request compacts * fix(responses): carry compaction through dict-form edits, SDK replay, affinity, and guardrails Read Anthropic's dict-form context_management, accept SDK ResponseCompactionItem objects on replay, strip the encrypted_content affinity wrap before decoding, scan and rewrite compaction summaries in Responses guardrails, drop the history a summary replay replaces, send the summary as a user message so thinking ordering holds, and send the streamed compaction .done with the newest block. * fix(responses): count, route, guard, and redact compaction items consistently A compaction item now reads as its summary everywhere the proxy counts or routes a Responses request: pre-call context checks, the TPM limiter, the input token endpoint, the stream usage estimate, and the previous_response_id session rebuild all drop the history the item replaces. A Claude compaction item sent to a non-Anthropic model (including a fallback) becomes a user summary message instead of an encrypted blob the provider rejects, and Claude on Bedrock Claude Platform takes the summary path. Output guardrails see the summary of a returned compaction item, and logging redaction hides it. * fix(responses): start previous_response_id replay at a stored compaction and count reservations from it * fix(responses): keep a stored compaction summary whole in spend logs * fix(responses): keep stored compaction items summary-only, redact compaction_delta, and finish compaction before reasoning A spend-log row now stores a top-level compaction input item as its summary alone, so extra keys in a client-built item no longer reach the database whole. Message redaction also nulls the streamed compaction_delta, and the streaming bridge sends the compaction item's done event before the reasoning item that follows it. Adds the audit's integration tests for compaction items on the wire. * test(responses): cover stored compaction items, compaction_delta redaction, and compaction done order Adds integration cells for the summary-only spend row of a caller-sent compaction item and its previous_response_id replay, a redacted chat stream row that holds no streamed compaction_delta, and a compaction then thinking stream whose done events follow output order. * fix(spend_logs): honor the redaction header on failure rows for litellm_metadata routes A failed /v1/responses or /v1/messages call keeps its proxy request headers in litellm_metadata, but the failure row read redaction headers from metadata only, so x-litellm-enable-message-redaction never redacted those rows. * fix(spend_logs): lift only the redaction opt-in header into failure rows The failure hook on litellm_metadata routes copied every real request header into the caller's own metadata dict, so the stored request body gained the headers, a caller's metadata.headers was replaced, and a disable header loosened global redaction on failed calls. It now copies only the enable headers into a new litellm_params metadata dict, and only when the caller sent no metadata.headers, so no failure row is less redacted than on main * test(responses): cover failure-row redaction headers and caller metadata on litellm_metadata routes * fix(responses): keep earlier-turn instructions after compaction and count compaction items by their summary A previous_response_id rebuild that starts at a compaction now keeps the system and developer messages of every earlier turn. Token estimates for raw Responses input count a compaction item as its summary text instead of the serialized item. * fix(responses): drop the stored history a caller-sent compaction item replaces on previous_response_id * fix(responses): count only message items with a string role as compaction instructions * test(responses): flatten the last two-for comprehension in the compaction wire test * fix(responses): replay from a compaction in stored Responses output and SDK compaction items A previous_response_id whose stored row is a Responses output holding a compaction item now starts its replay at that compaction, and the handler dumps SDK compaction objects before checking for a summary. Adds tests for the uncovered compaction paths. * fix(responses): keep the stored tool call a function_call_output answers when a compaction item replaces history * fix(responses): keep a stored tool call on every turn after a compaction item and bound the stored summary A previous_response_id rebuild that starts at a turn holding a compaction item now sends again each stored call a replayed tool output answers, so a retry or a later turn no longer loses the turn-1 call. The start is picked only from rows up to the referenced response. A spend row keeps a compaction summary up to the larger of 64K chars and MAX_STRING_LENGTH_PROMPT_IN_DB instead of with no bound. * fix(responses): replay a compacted session only through the referenced response A later sibling turn that sent a compaction item put its summary between a stored tool call and its tool output, so Claude answered 400 on a turn that continued the first of two retried turns, and on a retry sent without the item * fix(spend-logs): keep main's stored value for a prompt length limit under 2 At a MAX_STRING_LENGTH_PROMPT_IN_DB of 0 or 1 the shared truncation helper stored only the truncation marker, where main stores the marker followed by the whole value * fix(spend-logs): keep main's stored tail for a negative prompt length limit At a negative MAX_STRING_LENGTH_PROMPT_IN_DB the shared truncation helper stored one more character of each value than main, since it dropped main's clamp of the end slice to the limit minus the start * fix(responses): keep a compaction item when a guardrail rewrites its summary through structured_messages * test(responses): run the compaction wire tests on the worker database under isolated workers * fix(router): keep the encrypted content affinity pin out of provider requests * fix(proxy): inspect compaction summaries in guardrails, strip a client-sent affinity pin, and price failures from the newest compaction * fix(proxy): count a failed chat or /v1/messages call on its whole history, not from a compaction-typed message * fix(proxy): trim a failed call's input at the newest compaction only when the client called /v1/responses * fix(proxy): trim a failed call's input at the newest compaction when the proxy dispatched it as a Responses call * test(proxy): run the dispatched-call-type failure test through the real request setup * fix(proxy): count a failed Cursor stream's input from the newest compaction item The Cursor route maps to no call type, so the stream generator's failure hook fell back to the route and counted the whole history. Pass the processor's route_type. * fix(proxy): count a broken Cursor stream's prompt from the newest compaction item * fix(responses): keep a tool call a later output answers when trimming before a compaction item * fix(responses): move a compaction-split tool call's turn to right before its output on both compaction routes * perf(responses): look up moved turn items in a set so a long compaction-split turn trims in linear time * test(responses): measure the compaction-split turn move as a CodSpeed benchmark instead of a clock read in unit tests * fix(responses): move only calls still open at the compaction, and bill a disconnected stream from the newest compaction item * fix(responses): restore only the newest stored call when a reused call_id is answered after a compaction item * fix(spend logs): bound a Responses compaction item's summary like a chat block instead of cutting its encrypted_content * test(proxy): price a failed dispatched call through a real ProxyLogging instead of mocked hooks * fix(responses): finish a stream whose Claude compaction failed and group answered call ids by turn once * fix(responses): count input_tokens after a compaction item the way inference sends it --- litellm/constants.py | 1 + litellm/litellm_core_utils/redact_messages.py | 57 +- litellm/llms/anthropic/chat/handler.py | 14 +- .../guardrail_translation/handler.py | 82 +- litellm/proxy/_lazy_openapi_snapshot.json | 3 + litellm/proxy/common_request_processing.py | 6 + litellm/proxy/guardrails/_content_utils.py | 17 +- .../hooks/parallel_request_limiter_v3.py | 7 +- .../proxy/hooks/proxy_track_cost_callback.py | 49 +- litellm/proxy/litellm_pre_call_utils.py | 1 + litellm/proxy/proxy_server.py | 2 + .../proxy/response_api_endpoints/endpoints.py | 21 +- litellm/proxy/spend_tracking/input_tokens.py | 11 +- .../spend_tracking/spend_tracking_utils.py | 126 +- litellm/proxy/utils.py | 31 +- .../compaction_items.py | 329 +++++ .../count_request.py | 21 +- .../handler.py | 3 + .../session_handler.py | 164 ++- .../streaming_iterator.py | 78 +- .../tool_search.py | 18 +- .../transformation.py | 265 ++++- litellm/responses/main.py | 14 +- litellm/responses/streaming_iterator.py | 1 + litellm/router.py | 1 + litellm/types/litellm_params.py | 3 + litellm/types/llms/anthropic.py | 6 + litellm/types/llms/openai.py | 2 + .../test_responses_compaction_benchmarks.py | 35 + .../test_responses_compaction_items_wire.py | 1057 +++++++++++++++++ .../test_redact_messages.py | 22 + ...test_openai_responses_guardrail_handler.py | 180 +++ .../proxy/guardrails/test_content_utils.py | 41 + .../hooks/test_proxy_track_cost_callback.py | 68 ++ tests/unit/proxy/hooks/test_tpm_concurrent.py | 50 +- .../response_api_endpoints/test_endpoints.py | 122 ++ .../proxy/spend_tracking/test_input_tokens.py | 57 + .../test_spend_tracking_utils.py | 233 +++- .../proxy/test_common_request_processing.py | 67 ++ .../unit/proxy/test_litellm_pre_call_utils.py | 2 + ..._utils_model_creation_and_error_logging.py | 135 ++- .../test_anthropic_responses_bridge.py | 639 +++++++++- .../test_compaction_items.py | 87 ++ .../test_count_request.py | 75 ++ .../test_litellm_completion_responses.py | 402 +++++++ .../test_session_handler.py | 367 +++++- .../test_streaming_iterator_transformation.py | 34 + .../test_responses_api_request_body.py | 56 + .../unit/responses/test_streaming_iterator.py | 21 + tests/unit/test_router/test_router.py | 23 + tests/unit/types/test_litellm_params.py | 1 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 +- 52 files changed, 4879 insertions(+), 230 deletions(-) create mode 100644 litellm/responses/litellm_completion_transformation/compaction_items.py create mode 100644 tests/benchmarks/test_responses_compaction_benchmarks.py create mode 100644 tests/integration/providers/test_responses_compaction_items_wire.py create mode 100644 tests/unit/responses/litellm_completion_transformation/test_compaction_items.py diff --git a/litellm/constants.py b/litellm/constants.py index 6a503463a8c..c0fe0491775 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -611,6 +611,7 @@ NON_LLM_CONNECTION_TIMEOUT: Final = get_env_int( ) # timeout for adjacent services (e.g. jwt auth) MAX_EXCEPTION_MESSAGE_LENGTH: Final = get_env_int("MAX_EXCEPTION_MESSAGE_LENGTH", 2000) MAX_STRING_LENGTH_PROMPT_IN_DB: Final = get_env_int("MAX_STRING_LENGTH_PROMPT_IN_DB", 2048) +MAX_STRING_LENGTH_COMPACTION_SUMMARY_IN_DB: Final = 64 * 1024 BEDROCK_MAX_POLICY_SIZE: Final = get_env_int("BEDROCK_MAX_POLICY_SIZE", 75) # One entry per distinct AWS credential-argument set. Per-user cost attribution passes the attributed # identity as aws_session_name, so this bounds how many attributed identities keep a cached STS session. diff --git a/litellm/litellm_core_utils/redact_messages.py b/litellm/litellm_core_utils/redact_messages.py index 0168cdad7c0..7f8a38cc8ec 100644 --- a/litellm/litellm_core_utils/redact_messages.py +++ b/litellm/litellm_core_utils/redact_messages.py @@ -24,6 +24,7 @@ from litellm.litellm_core_utils.core_helpers import ( get_metadata_variable_name_from_kwargs, ) from litellm.litellm_core_utils.served_output_texts import SERVED_OUTPUT_TEXTS_KEY +from litellm.llms.base_llm.guardrail_translation.utils import stream_item_field, stream_item_items from litellm.llms.vertex_ai.common_utils import ( redact_vertex_ai_metadata_from_litellm_params, redact_vertex_ai_metadata_from_logged_object, @@ -41,6 +42,8 @@ if TYPE_CHECKING: else: LiteLLMLoggingObject = Any +MESSAGE_REDACTION_ENABLE_HEADERS: Final = ("litellm-enable-message-redaction", "x-litellm-enable-message-redaction") + def _opted_out_of_message_logging(custom_logger: CustomLogger) -> bool: return hasattr(custom_logger, "message_logging") and custom_logger.message_logging is not True @@ -195,33 +198,45 @@ def _redacted_provider_specific_fields(fields: Mapping[str, object]) -> Mapping[ return { **fields, **({"reasoning_content": REDACTED_BY_LITELLM} if fields.get("reasoning_content") is not None else {}), - **{key: None for key in ("thinking_blocks", "reasoningContentBlocks", "citations") if key in fields}, + **{ + key: None + for key in ( + "thinking_blocks", + "reasoningContentBlocks", + "citations", + "compaction_blocks", + "compaction_delta", + ) + if key in fields + }, } +def _redact_field_when_present(item: object, field: str) -> None: + if getattr(item, field, None) is not None: + setattr(item, field, REDACTED_BY_LITELLM) + + def _redact_responses_api_output(output_items): """Helper to redact ResponsesAPIResponse output items.""" for output_item in output_items: + output_type = stream_item_field(output_item, "type") if getattr(output_item, "text", None) is not None: output_item.text = REDACTED_BY_LITELLM - if hasattr(output_item, "content") and isinstance(output_item.content, list): - for content_part in output_item.content: - if getattr(content_part, "text", None) is not None: - content_part.text = REDACTED_BY_LITELLM - if getattr(content_part, "refusal", None) is not None: - content_part.refusal = REDACTED_BY_LITELLM + for content_part in stream_item_items(output_item, "content"): + _redact_field_when_present(content_part, "text") + _redact_field_when_present(content_part, "refusal") - # Redact reasoning items in output array - if hasattr(output_item, "type") and output_item.type == "reasoning": - if hasattr(output_item, "summary") and isinstance(output_item.summary, list): - for summary_item in output_item.summary: - if getattr(summary_item, "text", None) is not None: - summary_item.text = REDACTED_BY_LITELLM + if output_type == "reasoning": + for summary_item in stream_item_items(output_item, "summary"): + _redact_field_when_present(summary_item, "text") - if hasattr(output_item, "type") and output_item.type == "function_call" and hasattr(output_item, "arguments"): + if output_type == "compaction" and hasattr(output_item, "encrypted_content"): + output_item.encrypted_content = REDACTED_BY_LITELLM + if output_type == "function_call" and hasattr(output_item, "arguments"): output_item.arguments = REDACTED_BY_LITELLM - if hasattr(output_item, "type") and output_item.type == "custom_tool_call" and hasattr(output_item, "input"): + if output_type == "custom_tool_call" and hasattr(output_item, "input"): output_item.input = REDACTED_BY_LITELLM @@ -262,6 +277,11 @@ def _redacted_responses_api_output_item(item: object, redacted_str: str) -> obje if item_mapping.get("type") == "reasoning" and isinstance(summary, list) else {} ), + **( + {"encrypted_content": redacted_str} + if item_mapping.get("type") == "compaction" and "encrypted_content" in item_mapping + else {} + ), **( {"arguments": redacted_str} if item_mapping.get("type") == "function_call" and "arguments" in item_mapping @@ -556,13 +576,8 @@ def should_redact_message_logging(model_call_details: dict) -> bool: # User explicitly disabled redaction via header return False - possible_enable_headers: Final = [ - "litellm-enable-message-redaction", # old header. maintain backwards compatibility - "x-litellm-enable-message-redaction", # new header - ] - is_redaction_enabled_via_header = False - for header in possible_enable_headers: + for header in MESSAGE_REDACTION_ENABLE_HEADERS: if bool(request_headers.get(header, False)): is_redaction_enabled_via_header = True break diff --git a/litellm/llms/anthropic/chat/handler.py b/litellm/llms/anthropic/chat/handler.py index f0143f64be8..13ceea3e31b 100644 --- a/litellm/llms/anthropic/chat/handler.py +++ b/litellm/llms/anthropic/chat/handler.py @@ -856,11 +856,21 @@ class ModelResponseIterator: if reasoning_content is None: reasoning_content = "" elif "content" in content_block["delta"] and content_block["delta"].get("type") == "compaction_delta": - # Handle compaction delta + compaction_delta: Final = content_block["delta"]["content"] provider_specific_fields["compaction_delta"] = { "type": "compaction_delta", - "content": content_block["delta"]["content"], + "content": compaction_delta, } + if self.compaction_blocks and isinstance(compaction_delta, str): + streamed_summary: Final = self.compaction_blocks[-1].get("content") + self.compaction_blocks = [ + *self.compaction_blocks[:-1], + { + **self.compaction_blocks[-1], + "content": (streamed_summary if isinstance(streamed_summary, str) else "") + compaction_delta, + }, + ] + provider_specific_fields["compaction_blocks"] = self.compaction_blocks return ( text, diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 84ff790ae48..3ea126c5a2a 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -65,6 +65,11 @@ from litellm.llms.base_llm.guardrail_translation.utils import ( unappliable_request_rewrite, ) from litellm.llms.openai.responses.guardrail_translation.tool_merge import merge_guardrailed_tools +from litellm.responses.litellm_completion_transformation.compaction_items import ( + compaction_summary, + decode_compaction_item, + with_compaction_summary, +) from litellm.responses.litellm_completion_transformation.custom_tools import without_malformed_tool_flags from litellm.responses.litellm_completion_transformation.tool_search import is_server_tool_search_tool from litellm.responses.litellm_completion_transformation.transformation import ( @@ -110,6 +115,16 @@ if TYPE_CHECKING: from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth +_COMPACTION_SUMMARY_POSITION: Final = -1 + + +def _output_compaction_summary(item: object) -> str | None: + if stream_item_field(item, "type") != "compaction": + return None + block: Final = decode_compaction_item(stream_item_field(item, "encrypted_content")) + return None if block is None else block.summary + + class _ToolCallShape(NamedTuple): name: str | None arguments: str @@ -266,7 +281,17 @@ def _item_rewrite_field(item: Mapping[str, object]) -> str | None: return _PATCHABLE_ITEM_FIELDS.get(item_type) +def _rewritten_compaction_item(item: Mapping[str, object], rewritten: object) -> Mapping[str, object] | None: + encrypted_content: Final = item.get("encrypted_content") + summary: Final = rewritten.get("content") if isinstance(rewritten, Mapping) else None + if not isinstance(encrypted_content, str) or not isinstance(summary, str) or not summary: + return None + return {**item, "encrypted_content": with_compaction_summary(encrypted_content, summary)} + + def _rewritten_input_item(item: Mapping[str, object], rewritten: object) -> Mapping[str, object] | None: + if compaction_summary(item) is not None: + return _rewritten_compaction_item(item, rewritten) field: Final = _item_rewrite_field(item) if field is None or not isinstance(rewritten, Mapping): return None @@ -804,6 +829,11 @@ class OpenAIResponsesHandler(BaseTranslation): Override this method to customize text/image extraction logic. """ + summary: Final = compaction_summary(message) + if summary is not None: + texts_to_check.append(summary) + task_mappings.append((msg_idx, None)) + return content: Final = message.get("content", None) if content is None: return @@ -849,6 +879,12 @@ class OpenAIResponsesHandler(BaseTranslation): msg_idx = cast(int, mapping[0]) content_idx_optional = cast(int | None, mapping[1]) + if compaction_summary(messages[msg_idx]) is not None: + messages[msg_idx]["encrypted_content"] = with_compaction_summary( # rebind-ok: in-place input rewrite + messages[msg_idx]["encrypted_content"], guardrail_response + ) + continue + content = messages[msg_idx].get("content", None) if content is None: continue @@ -1116,13 +1152,17 @@ class OpenAIResponsesHandler(BaseTranslation): # delivered, so a truncated buffer still reports text undeliverable. # # ------------------------------------------------------------------ # if final_chunk.get("type") == "response.output_item.done": + finished_summary: Final = _output_compaction_summary(stream_item_field(final_chunk, "item")) model_response_stream: Final = ( OpenAiResponsesToChatCompletionStreamIterator.translate_responses_chunk_to_openai_stream(final_chunk) ) tool_calls: Final = model_response_stream.choices[0].delta.tool_calls - if tool_calls: + if tool_calls or finished_summary is not None: inputs = GenericGuardrailAPIInputs() - inputs["tool_calls"] = cast(list[ChatCompletionToolCallChunk], tool_calls) + if finished_summary is not None: + inputs["texts"] = [finished_summary] + if tool_calls: + inputs["tool_calls"] = cast(list[ChatCompletionToolCallChunk], tool_calls) if hasattr(model_response_stream, "model") and model_response_stream.model: inputs["model"] = model_response_stream.model await guardrail_to_apply.apply_guardrail( @@ -1214,6 +1254,14 @@ class OpenAIResponsesHandler(BaseTranslation): else: setattr(event, field, value) + @staticmethod + def _write_compaction_summary(item: object, summary: str) -> None: + encrypted_content: Final = stream_item_field(item, "encrypted_content") + if isinstance(encrypted_content, str): + OpenAIResponsesHandler._write_event_field( + item, "encrypted_content", with_compaction_summary(encrypted_content, summary) + ) + def _sync_stream_events_with_rewrites( self, stream_events: Sequence[Any], @@ -1256,11 +1304,13 @@ class OpenAIResponsesHandler(BaseTranslation): output_index: int, rewrites_by_position: Mapping[tuple[int, int], str], ) -> None: - content: Final = item.get("content") if isinstance(item, dict) else getattr(item, "content", None) - if not isinstance(content, list): + summary_rewrite: Final = rewrites_by_position.get((output_index, _COMPACTION_SUMMARY_POSITION)) + if summary_rewrite is not None: + OpenAIResponsesHandler._write_compaction_summary(item, summary_rewrite) return + content: Final = stream_item_items(item, "content") for (item_idx, content_idx), rewritten in rewrites_by_position.items(): - if item_idx != output_index or content_idx >= len(content): + if item_idx != output_index or not 0 <= content_idx < len(content): continue OpenAIResponsesHandler._write_event_field(content[content_idx], "text", rewritten) @@ -1481,13 +1531,13 @@ class OpenAIResponsesHandler(BaseTranslation): message_items: Final = tuple( item for item in output_items if stream_item_field(item, "type") not in _OUTPUT_TOOL_CALL_ITEM_TYPES ) + message_parts: Final = chain.from_iterable(stream_item_items(item, "content") for item in message_items) + part_texts: Final = tuple( + text for part in message_parts if isinstance(text := stream_item_field(part, "text"), str) and text + ) + summaries: Final = tuple(summary for item in output_items if (summary := _output_compaction_summary(item))) return StreamingScanKey( - texts=tuple( - text - for item in message_items - for part in stream_item_items(item, "content") - if isinstance(text := stream_item_field(part, "text"), str) and text - ), + texts=(*part_texts, *summaries), tool_calls=tuple( stream_item_fingerprint(item) for item in output_items @@ -1610,6 +1660,12 @@ class OpenAIResponsesHandler(BaseTranslation): ) return + summary: Final = _output_compaction_summary(output_item) + if summary is not None: + texts_to_check.append(summary) + task_mappings.append((output_idx, _COMPACTION_SUMMARY_POSITION)) + return + # Handle both GenericResponseOutputItem and dict content: list[OutputText] | list[dict] | None = None if isinstance(output_item, BaseModel): @@ -1677,6 +1733,10 @@ class OpenAIResponsesHandler(BaseTranslation): output_item = response_output[output_idx] + if content_idx == _COMPACTION_SUMMARY_POSITION: + self._write_compaction_summary(output_item, guardrail_response) + continue + # Handle both GenericResponseOutputItem, BaseModel, and dict if isinstance(output_item, GenericResponseOutputItem): if output_item.content and content_idx < len(output_item.content): diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 545b317055a..4572e59944b 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -28240,6 +28240,9 @@ }, { "$ref": "#/components/schemas/OutputToolSearchCall" + }, + { + "$ref": "#/components/schemas/ResponseCompactionItem" } ] }, diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 6fda5780f63..23faa964d5f 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -1985,6 +1985,7 @@ class ProxyBaseLLMRequestProcessing: self.data = data self.served_model = served_model self._tags_before_guardrails: frozenset[str] | None = None + self.route_type: str | None = None @property def litellm_call_id(self) -> str | None: @@ -2434,6 +2435,7 @@ class ProxyBaseLLMRequestProcessing: ## LOGGING OBJECT ## - initialize logging object for logging success/failure events for call ## IMPORTANT Note: - initialize this before running pre-call checks. Ensures we log rejected requests to langfuse. + self.route_type = route_type logging_obj, self.data = litellm.utils.function_setup( original_function=route_type, rules_obj=litellm.utils.Rules(), @@ -2800,6 +2802,7 @@ class ProxyBaseLLMRequestProcessing: user_api_key_dict=user_api_key_dict, original_exception=exc, request_data=self.data, + dispatched_call_type=self.route_type, ) return await open_sse_before_first_byte( @@ -2864,6 +2867,7 @@ class ProxyBaseLLMRequestProcessing: "skip_pre_call_logic=True requires litellm_logging_obj to be set in data. " "Ensure common_processing_pre_call_logic was called before using this parameter." ) + self.route_type = route_type else: self.data, logging_obj = await self._pre_call_with_fallbacks( request=request, @@ -3473,6 +3477,7 @@ class ProxyBaseLLMRequestProcessing: user_api_key_dict=user_api_key_dict, original_exception=exc, request_data=request_data, + dispatched_call_type=self.route_type, ) except HTTPException: return @@ -4142,6 +4147,7 @@ class ProxyBaseLLMRequestProcessing: user_api_key_dict=user_api_key_dict, original_exception=e, request_data=self.data, + dispatched_call_type=self.route_type, ) # Use transformed exception if callback returned one, otherwise use original if transformed_exception is not None: diff --git a/litellm/proxy/guardrails/_content_utils.py b/litellm/proxy/guardrails/_content_utils.py index 64c935ce0d0..17bea031794 100644 --- a/litellm/proxy/guardrails/_content_utils.py +++ b/litellm/proxy/guardrails/_content_utils.py @@ -11,6 +11,12 @@ every text fragment. from collections.abc import Callable, Iterator, Mapping, Sequence from typing import Any, Final +from litellm.responses.litellm_completion_transformation.compaction_items import ( + compaction_summary, + input_items, + with_visited_compaction_summary, +) + # Call types whose body carries free-form chat / prompt text that # text-content guardrails (banned keywords, content moderation, secret # detection, …) should inspect. The proxy ingress passes ``route_type`` @@ -95,12 +101,12 @@ def _coerce_input_to_messages(input_value: object) -> list[dict[str, object]]: """Coerce a Responses-API ``data["input"]`` value into chat-style messages.""" if isinstance(input_value, str): return [{"role": "user", "content": input_value}] - if not isinstance(input_value, list): - return [] messages: Final[list[dict[str, object]]] = [] - for item in input_value: + for item in input_items(input_value): if isinstance(item, str): messages.append({"role": "user", "content": item}) + elif (summary := compaction_summary(item)) is not None: + messages.append({"role": "user", "content": summary}) elif isinstance(item, dict): if _part_text(item) is not None: messages.append({"role": item.get("role") or "user", "content": [item]}) @@ -189,11 +195,14 @@ def walk_user_text(data: dict[str, object], visit: Callable[[str], str]) -> int: data["input"] = visit(input_value) return visited if isinstance(input_value, list): - for idx, item in enumerate(input_value): + for idx, item in enumerate(input_items(input_value)): if isinstance(item, str): if item: visited += 1 input_value[idx] = visit(item) + elif (compaction := with_visited_compaction_summary(item, visit)) is not None: + visited += 1 + input_value[idx] = compaction elif isinstance(item, dict): if _part_text(item) is not None: visited += 1 diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index a8ba8d65785..a2bb691b401 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -73,6 +73,7 @@ from litellm.proxy.hooks.batch_enqueued_tokens import ( canonical_provider_batch_id, ) from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit +from litellm.responses.litellm_completion_transformation.compaction_items import billable_input_since_newest_compaction from litellm.router_utils.add_retry_fallback_headers import ( ensure_response_additional_headers, response_has_hidden_params, @@ -1247,7 +1248,10 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): estimable_data.get("input"), ) ) - messages, prompt, input_text = selected_fields + messages, prompt, raw_input = selected_fields + input_text: Final = ( + billable_input_since_newest_compaction(raw_input) if call_type in RESPONSES_API_CALL_TYPES else raw_input + ) total_chars: Final = ( len(get_str_from_messages(messages)) @@ -3835,6 +3839,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): return LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( input=data.get("input") or "", responses_api_request=data, + only_since_compaction=True, ) @staticmethod diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 994310c6619..6801c8521e5 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -22,6 +22,7 @@ from litellm.litellm_core_utils.core_helpers import ( ) from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import guardrail_information_cost +from litellm.litellm_core_utils.redact_messages import MESSAGE_REDACTION_ENABLE_HEADERS from litellm.llms.azure.ptu_capacity import is_str_keyed_mapping from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.auth_checks import ( @@ -127,6 +128,9 @@ _CAPTURED_IDENTITY_CALL_TYPES: Final[frozenset[str]] = frozenset( ) +_OBJECT_MAPPING: Final = TypeAdapter(ObjectMapping) + + def _proxy_stamped_used_client_oauth_token( request_data: Mapping[str, object], request_route: str | None ) -> bool | None: @@ -183,6 +187,26 @@ def _failure_snapshot_with_guardrails_recorded_after_it( return {**failure_snapshot, "guardrail_information": recorded_guardrail_information} +def _guardrail_info_lifted_from_litellm_metadata( + litellm_metadata: ObjectMapping, *, guardrail_info_already_set: bool +) -> ObjectMapping: + guardrail_info: Final = litellm_metadata.get("standard_logging_guardrail_information") + if guardrail_info is None or guardrail_info_already_set: + return {} + return {"standard_logging_guardrail_information": guardrail_info} + + +def _redaction_opt_in_lifted_from_litellm_metadata( + litellm_metadata: ObjectMapping, *, caller_set_headers: bool +) -> ObjectMapping: + request_headers: Final = litellm_metadata.get("headers") + if caller_set_headers or not isinstance(request_headers, dict): + return {} + headers: Final = _OBJECT_MAPPING.validate_python(request_headers) + opt_in_headers: Final = {name: headers[name] for name in MESSAGE_REDACTION_ENABLE_HEADERS if name in headers} + return {"headers": opt_in_headers} if opt_in_headers else {} + + def _proxy_spend_writer() -> DBSpendUpdateWriter: from litellm.proxy.proxy_server import proxy_logging_obj @@ -417,13 +441,17 @@ class _ProxyDBLogger(CustomLogger): existing_metadata.update({**_routed_metadata_fallback(request_fields), **_metadata}) litellm_metadata_bucket: Final = request_data.get("litellm_metadata") - if ( - isinstance(litellm_metadata_bucket, dict) - and "standard_logging_guardrail_information" not in existing_metadata - ): - guardrail_info: Final = litellm_metadata_bucket.get("standard_logging_guardrail_information") - if guardrail_info is not None: - existing_metadata["standard_logging_guardrail_information"] = guardrail_info + litellm_metadata: Final[ObjectMapping] = ( + _OBJECT_MAPPING.validate_python(litellm_metadata_bucket) + if isinstance(litellm_metadata_bucket, dict) + else {} + ) + existing_metadata.update( + _guardrail_info_lifted_from_litellm_metadata( + litellm_metadata, + guardrail_info_already_set="standard_logging_guardrail_information" in existing_metadata, + ) + ) if "litellm_params" not in request_data: request_data["litellm_params"] = {} @@ -438,7 +466,12 @@ class _ProxyDBLogger(CustomLogger): request_data["litellm_params"]["proxy_server_request"] = ( request_data.get("proxy_server_request") or existing_litellm_params.get("proxy_server_request") or {} ) - request_data["litellm_params"]["metadata"] = existing_metadata + request_data["litellm_params"]["metadata"] = { + **existing_metadata, + **_redaction_opt_in_lifted_from_litellm_metadata( + litellm_metadata, caller_set_headers="headers" in existing_metadata + ), + } # Preserve model name and custom_llm_provider if "model" not in request_data: diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index d34787cd5e9..8a45f658f36 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -353,6 +353,7 @@ LITELLM_TRACE_CONTROL_METADATA_FIELDS: Final = frozenset( _UNTRUSTED_ROOT_CONTROL_FIELDS: Final = ( "weights", "_router_weights", + "_encrypted_content_affinity_pinned", "fallback_depth", "_target_order", "attempted_targets", diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index b48b5ab1c1e..f632317d9a8 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -10384,6 +10384,7 @@ async def async_data_generator( request: Request | None = None, *, responses_stream_errors: bool = False, + dispatched_call_type: str | None = None, ): verbose_proxy_logger.debug("inside generator") stream_completed = False @@ -10599,6 +10600,7 @@ async def async_data_generator( user_api_key_dict=user_api_key_dict, original_exception=e, request_data=request_data, + dispatched_call_type=dispatched_call_type, ) client_error: Final[Exception] = e if replacement is None else replacement verbose_proxy_logger.debug( diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index c918ea3f1d1..312c5124ca5 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -37,6 +37,7 @@ from litellm.constants import EMPTY_MAPPING, MANAGED_BACKGROUND_RESPONSE_METADAT from litellm.integrations.custom_guardrail import ModifyResponseException from litellm.litellm_core_utils.core_helpers import as_str_mapping, get_metadata_variable_name_from_kwargs from litellm.litellm_core_utils.internal_call_metadata import is_background_response, requests_background +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup from litellm.litellm_core_utils.tool_fields import TOOL_FIELDS, first_validation_error from litellm.llms.openai.responses.guardrail_translation.handler import ( @@ -79,7 +80,8 @@ from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attributio request_tags_from_request_data, ) from litellm.proxy.route_llm_request import raise_if_required_body_param_missing -from litellm.proxy.utils import handle_exception_on_proxy +from litellm.proxy.utils import count_request_input_tokens, handle_exception_on_proxy +from litellm.responses.litellm_completion_transformation.compaction_items import messages_since_newest_compaction from litellm.responses.litellm_completion_transformation.count_request import MODEL_FILE_ID_MAPPING_FIELD from litellm.responses.streaming_iterator import GuardedResponsesWsFrame, ResponsesWsFrameGuard, spawn_logging_task from litellm.router import Router @@ -94,7 +96,6 @@ from litellm.types.token_count import TokenCountSourceFields from litellm.types.utils import ADDRESSED_RESPONSE_ID_FIELD, TokenCountResponse if TYPE_CHECKING: - from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy.utils import ProxyLogging router: Final = APIRouter() @@ -532,6 +533,14 @@ async def cursor_model_list( return await model_list(user_api_key_dict=user_api_key_dict) +def _dispatched_prompt_token_counter(logging_obj: LiteLLMLoggingObj) -> Callable[[], int]: + return lambda: count_request_input_tokens( + model=str(logging_obj.model_call_details.get("model") or ""), + request_input=messages_since_newest_compaction(logging_obj.model_call_details.get("messages")), + system_input=None, + ) + + @router.post( "/cursor/chat/completions", dependencies=[ @@ -656,6 +665,11 @@ async def cursor_chat_completions( model=request_data.get("model", ""), custom_llm_provider=None, logging_obj=logging_obj, + count_prompt_tokens=( + _dispatched_prompt_token_counter(logging_obj) + if isinstance(logging_obj, LiteLLMLoggingObj) + else None + ), ) # Use async_data_generator to format as SSE return async_data_generator( @@ -663,6 +677,7 @@ async def cursor_chat_completions( user_api_key_dict=user_api_key_dict, request_data=request_data, request=request, + dispatched_call_type=processor.route_type, ) # Otherwise, use the default generator return async_data_generator( @@ -670,6 +685,7 @@ async def cursor_chat_completions( user_api_key_dict=user_api_key_dict, request_data=request_data, request=request, + dispatched_call_type=processor.route_type, ) try: @@ -1585,6 +1601,7 @@ def _responses_input_as_token_count_messages( transformed: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( input=input_value, responses_api_request=request_params, + only_since_compaction=True, ) return tuple( message if isinstance(message, dict) else message.model_dump(exclude_none=True) for message in transformed diff --git a/litellm/proxy/spend_tracking/input_tokens.py b/litellm/proxy/spend_tracking/input_tokens.py index 6c7083fb6db..927152296e9 100644 --- a/litellm/proxy/spend_tracking/input_tokens.py +++ b/litellm/proxy/spend_tracking/input_tokens.py @@ -19,6 +19,10 @@ from typing import Final import litellm from litellm._logging import verbose_proxy_logger +from litellm.responses.litellm_completion_transformation.compaction_items import ( + billable_input_since_newest_compaction, + holds_compaction_summary, +) from litellm.rust_bridge import runtime from litellm.rust_bridge.catalog import Route, RouteContext from litellm.rust_bridge.token_counter import ( @@ -54,10 +58,11 @@ async def count_input_tokens( (model, rust_tokenizer(model)) for model in models ) groups: Final[tuple[RustTokenizer | None, ...]] = tuple(dict.fromkeys(tokenizer for _, tokenizer in tokenizers)) + counted_raw_body: Final = None if holds_compaction_summary(request_body.get("input")) else raw_body group_counts: Final = [ await _count_group( request_body=request_body, - raw_body=raw_body, + raw_body=counted_raw_body, tokenizer=tokenizer, models=tuple(model for model, selected in tokenizers if selected == tokenizer), ) @@ -140,7 +145,9 @@ def count_input_tokens_for_model(request_body: dict, model: str) -> int | None: if "prompt" in request_body: return _count_text_tokens(model=model, text=request_body.get("prompt")) if "input" in request_body: - return _count_text_tokens(model=model, text=request_body.get("input")) + return _count_text_tokens( + model=model, text=billable_input_since_newest_compaction(request_body.get("input")) + ) if "query" in request_body or "documents" in request_body: query_tokens: Final = _count_text_tokens(model=model, text=request_body.get("query")) document_tokens: Final = _count_text_tokens( diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 62aee788ed0..94b079163e9 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -27,6 +27,7 @@ from litellm.constants import ( LITTELM_CLI_SERVICE_ACCOUNT_NAME, LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME, MAX_SPEND_LOG_MODEL_NAME_LENGTH, + MAX_STRING_LENGTH_COMPACTION_SUMMARY_IN_DB, MCP_SPEND_LOG_MODEL_PREFIX, NON_INFERENCE_CALL_TYPES, REDACTED_BY_LITELM_STRING, @@ -70,6 +71,11 @@ from litellm.proxy._types import ( from litellm.proxy.route_llm_request import ProxyModelNotFoundError from litellm.proxy.spend_tracking.spend_log_error_logger import spend_log_error from litellm.proxy.utils import PrismaClient, hash_token +from litellm.responses.litellm_completion_transformation.compaction_items import ( + compaction_summary, + summary_only_compaction_content, + with_compaction_summary, +) from litellm.types.passthrough_endpoints.pass_through_endpoints import ( SETTLES_BACKGROUND_RESPONSE_KWARG, BackgroundResponseSettlement, @@ -1227,11 +1233,42 @@ def _is_spend_log_content(parent: Mapping[str, object], key: str) -> bool: ) +def _truncated_for_db(value: str, max_length: int) -> str: + if len(value) <= max_length: + return value + start_chars: Final = int(max_length * 0.35) + end_chars: Final = int(max_length * 0.65) + return ( + f"{value[:start_chars]}" + f"... ({LITELLM_TRUNCATED_PAYLOAD_FIELD} skipped {len(value) - start_chars - end_chars} chars. " + f"{LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE}) ..." + f"{value[-min(end_chars, max_length - start_chars) :]}" + ) + + +def _compaction_summary_for_db(summary: str) -> str: + return _truncated_for_db( + summary, max(_get_max_string_length_prompt_in_db(), MAX_STRING_LENGTH_COMPACTION_SUMMARY_IN_DB) + ) + + +def _bounded_compaction_summary( + parent: Mapping[str, object], key: str, value: object, keep_compaction_blocks: bool +) -> str | None: + if not keep_compaction_blocks or not isinstance(value, str) or parent.get("type") != "compaction": + return None + if key == "content": + return _compaction_summary_for_db(value) + summary: Final = compaction_summary(parent) if key == "encrypted_content" else None + return None if summary is None else with_compaction_summary(value, _compaction_summary_for_db(summary)) + + def _sanitize_request_body_for_spend_logs_payload( request_body: Mapping[str, object], visited: set | None = None, max_string_length_prompt_in_db: int | None = None, mask_credentials: bool = True, + keep_compaction_blocks: bool = False, ) -> dict: """ Recursively sanitize request body to prevent logging large base64 strings or other large values. @@ -1242,11 +1279,6 @@ def _sanitize_request_body_for_spend_logs_payload( SensitiveDataMasker classifies as credentials with REDACTED_BY_LITELM_STRING, except inside tool payloads and logprobs. """ - from litellm.constants import ( - LITELLM_TRUNCATED_PAYLOAD_FIELD, - LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE, - ) - if visited is None: visited = set() if max_string_length_prompt_in_db is None: @@ -1261,51 +1293,23 @@ def _sanitize_request_body_for_spend_logs_payload( def _sanitize_value(value: object, mask_credentials: bool) -> object: if isinstance(value, Mapping): return _sanitize_request_body_for_spend_logs_payload( - value, visited, max_string_length_prompt_in_db, mask_credentials + value, visited, max_string_length_prompt_in_db, mask_credentials, keep_compaction_blocks ) elif isinstance(value, list): return [_sanitize_value(item, mask_credentials) for item in value] elif isinstance(value, str): - if len(value) > max_string_length_prompt_in_db: - # Keep 35% from beginning and 65% from end (end is usually more important) - # This split ensures we keep more context from the end of conversations - start_ratio: Final = 0.35 - end_ratio: Final = 0.65 - - # Calculate character distribution - start_chars: Final = int(max_string_length_prompt_in_db * start_ratio) - end_chars = int(max_string_length_prompt_in_db * end_ratio) - - # Ensure we don't exceed the total limit - total_keep: Final = start_chars + end_chars - if total_keep > max_string_length_prompt_in_db: - end_chars = max_string_length_prompt_in_db - start_chars - - # If the string length is less than what we want to keep, just truncate normally - if len(value) <= max_string_length_prompt_in_db: - return value - - # Calculate how many characters are being skipped - skipped_chars: Final = len(value) - total_keep - - # Build the truncated string: beginning + truncation marker + end - truncated_value: Final = ( - f"{value[:start_chars]}" - f"... ({LITELLM_TRUNCATED_PAYLOAD_FIELD} skipped {skipped_chars} chars. " - f"{LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE}) ..." - f"{value[-end_chars:]}" - ) - return truncated_value - return value + return _truncated_for_db(value, max_string_length_prompt_in_db) return value - return { - k: REDACTED_BY_LITELM_STRING - if mask_credentials and _is_request_body_credential(k, v) - else _sanitize_value(v, mask_credentials and not _is_spend_log_content(request_body, k)) - for k, v in request_body.items() - if k not in _SENSITIVE_REQUEST_BODY_KEYS - } + def _stored_value(key: str, value: object) -> object: + if mask_credentials and _is_request_body_credential(key, value): + return REDACTED_BY_LITELM_STRING + bounded_summary: Final = _bounded_compaction_summary(request_body, key, value, keep_compaction_blocks) + if bounded_summary is not None: + return bounded_summary + return _sanitize_value(value, mask_credentials and not _is_spend_log_content(request_body, key)) + + return {k: _stored_value(k, v) for k, v in request_body.items() if k not in _SENSITIVE_REQUEST_BODY_KEYS} # Quoted-key form: ``"input"`` / ``'messages'`` / ``"prompt"`` followed by @@ -1820,6 +1824,33 @@ def _stored_request_without_requester_ip_address(request_body: Mapping[object, o } +def _with_bounded_compaction_summary(item: object, stored_item: object) -> object: + summary: Final = compaction_summary(item) + if summary is None or not isinstance(stored_item, Mapping): + return stored_item + return { + **_OBJECT_MAPPING.validate_python(stored_item), + "encrypted_content": summary_only_compaction_content(_compaction_summary_for_db(summary)), + } + + +def _sanitize_proxy_server_request_body(request_body: Mapping[str, object]) -> Mapping[object, object]: + sanitized: Final = _OBJECT_MAPPING.validate_python(_sanitize_request_body_for_spend_logs_payload(request_body)) + items: Final = request_body.get("input") + stored_items: Final = sanitized.get("input") + if not isinstance(items, list) or not isinstance(stored_items, list): + return sanitized + return { + **sanitized, + "input": [ + _with_bounded_compaction_summary(item, stored_item) + for item, stored_item in zip( + _OBJECT_LIST.validate_python(items), _OBJECT_LIST.validate_python(stored_items) + ) + ], + } + + def _get_proxy_server_request_for_spend_logs_payload( metadata: dict, litellm_params: Mapping[str, object], @@ -1864,8 +1895,7 @@ def _get_proxy_server_request_for_spend_logs_payload( _request_body = _convert_mapping_to_json_serializable(without_classifier_audit(_request_body)) perform_redaction(model_call_details=_request_body, result=None) - _request_body = _sanitize_request_body_for_spend_logs_payload(_request_body) - stored_request_body: Final = _OBJECT_MAPPING.validate_python(_request_body) + stored_request_body: Final = _sanitize_proxy_server_request_body(_request_body) _request_body_json_str: Final = safe_dumps( _stored_request_without_requester_ip_address(stored_request_body) if drop_requester_ip_address @@ -1939,7 +1969,9 @@ def _get_response_for_spend_logs_payload( response_obj = _convert_to_json_serializable_dict(response_obj) response_obj = perform_redaction(model_call_details={}, result=response_obj) - sanitized_wrapper: Final = _sanitize_request_body_for_spend_logs_payload({"response": response_obj}) + sanitized_wrapper: Final = _sanitize_request_body_for_spend_logs_payload( + {"response": response_obj}, keep_compaction_blocks=True + ) sanitized_response: Final = sanitized_wrapper.get("response", response_obj) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 2a328971817..52b3866d04d 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -89,6 +89,7 @@ from litellm.proxy.common_utils.openai_error_payload import ( with_litellm_call_id, ) from litellm.proxy.spend_tracking.spend_log_error_logger import spend_log_error +from litellm.responses.litellm_completion_transformation.compaction_items import messages_since_newest_compaction from litellm.types.guardrails import GuardrailEventHooks from litellm.types.proxy.model_listing import ModelInfoResponse from litellm.types.utils import MCP_GUARDRAIL_CALL_TYPES, CallTypes, CallTypesLiteral, ModelInfo, Usage @@ -937,7 +938,7 @@ def _system_prompt_text(system_input: object) -> str: return "".join(_prompt_block_text(block) for block in system_input) -def _count_request_input_tokens(model: str, request_input: object, system_input: object) -> int: +def count_request_input_tokens(model: str, request_input: object, system_input: object) -> int: system_text: Final = _system_prompt_text(system_input) system_tokens: Final = litellm.token_counter(model=model, text=system_text) if system_text else 0 if isinstance(request_input, str): @@ -958,7 +959,7 @@ def _estimate_dispatched_failure_usage(model: str, request_input: object, system the same tokenizer fallback interrupted streams use, so the spend log's failure row records what was sent instead of zero.""" try: - input_tokens: Final = _count_request_input_tokens( + input_tokens: Final = count_request_input_tokens( model=model, request_input=request_input, system_input=system_input ) except Exception: @@ -997,6 +998,8 @@ _INPUT_ESTIMABLE_CALL_TYPES: Final = frozenset( ) +_RESPONSES_CALL_TYPES: Final = frozenset((CallTypes.responses.value, CallTypes.aresponses.value)) + _FAILED_ATTEMPTS_USAGE: Final = TypeAdapter(tuple[tuple[Usage, float], ...]) @@ -1022,6 +1025,7 @@ def _current_failure_usage_to_lift( provider_billed: bool = False, deployment_id: str | None = None, base_model: str | None = None, + dispatched_call_type: str | None = None, ) -> tuple[object, object] | None: """A stream that broke mid-flight still billed the provider for the chunks already delivered; the streaming handler stashes that recovered usage and @@ -1045,7 +1049,7 @@ def _current_failure_usage_to_lift( # The proxy-only error handler stamps this flag on a client disconnect too, even though the # request had already reached the provider, so a billed failure keeps its estimate. return None - estimated_usage: Final = _estimated_prompt_usage(model_call_details, request_body) + estimated_usage: Final = _estimated_prompt_usage(model_call_details, request_body, dispatched_call_type) if estimated_usage is None: return None if not provider_billed: @@ -1056,13 +1060,14 @@ def _current_failure_usage_to_lift( def _estimated_prompt_usage( - model_call_details: Mapping[str, object], request_body: Mapping[str, object] + model_call_details: Mapping[str, object], request_body: Mapping[str, object], dispatched_call_type: str | None ) -> Usage | None: """The input-only Usage estimate for the prompt this request sent, or None when the call type has no estimable input. The raw request body backfills the system prompt when the SDK bridges an endpoint (e.g. /v1/messages on a chat-completions provider) without filling optional_params.""" if str(model_call_details.get("call_type")) not in _INPUT_ESTIMABLE_CALL_TYPES: return None + dispatched_input: Final = model_call_details.get("messages") optional_params: Final = model_call_details.get("optional_params") dispatched_system: Final = ( (optional_params.get("system") or optional_params.get("instructions")) @@ -1079,7 +1084,11 @@ def _estimated_prompt_usage( ) return _estimate_dispatched_failure_usage( model=str(model_call_details.get("model") or ""), - request_input=model_call_details.get("messages"), + request_input=( + messages_since_newest_compaction(dispatched_input) + if dispatched_call_type in _RESPONSES_CALL_TYPES + else dispatched_input + ), system_input=system_input, ) @@ -1444,7 +1453,9 @@ def _provider_billed_the_failure(exception: Exception) -> bool: return _timed_out_reading_from_provider(exception) or _client_disconnected(exception) -def _failure_fields_to_lift(request_data: Mapping[str, object], provider_billed: bool = False) -> Mapping[str, object]: +def _failure_fields_to_lift( + request_data: Mapping[str, object], provider_billed: bool = False, dispatched_call_type: str | None = None +) -> Mapping[str, object]: """Failure-path callbacks run after ``litellm_logging_obj`` is popped from request_data (it is not serialisable), so the caller merges these fields onto request_data first: the request start and first-handoff instants for @@ -1468,9 +1479,10 @@ def _failure_fields_to_lift(request_data: Mapping[str, object], provider_billed: provider_billed=_billed, deployment_id=_deployment_id, base_model=_base_model, + dispatched_call_type=dispatched_call_type, ) _prompt_estimate: Final = ( - _estimated_prompt_usage(_model_call_details, request_data) + _estimated_prompt_usage(_model_call_details, request_data, dispatched_call_type) if _billed and _first_handoff is not None and not _model_call_details.get("failed_attempts_usage") else None ) @@ -3634,6 +3646,7 @@ class ProxyLogging: error_type: ProxyErrorTypes | None = None, route: str | None = None, traceback_str: str | None = None, + dispatched_call_type: str | None = None, ) -> HTTPException | None: """ Allows users to raise custom exceptions/log when a call fails, without having to deal with parsing Request body. @@ -3704,7 +3717,9 @@ class ProxyLogging: request_data.update( await offload_token_count(_failure_fields_to_lift)( - request_data, provider_billed=_provider_billed_the_failure(original_exception) + request_data, + provider_billed=_provider_billed_the_failure(original_exception), + dispatched_call_type=dispatched_call_type or _call_type_for_route(user_api_key_dict.request_route), ) ) diff --git a/litellm/responses/litellm_completion_transformation/compaction_items.py b/litellm/responses/litellm_completion_transformation/compaction_items.py new file mode 100644 index 00000000000..8bd82c49090 --- /dev/null +++ b/litellm/responses/litellm_completion_transformation/compaction_items.py @@ -0,0 +1,329 @@ +import json +import uuid +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass +from itertools import chain, groupby +from types import MappingProxyType +from typing import Final, TypeVar + +from openai.types.responses.easy_input_message_param import EasyInputMessageParam +from openai.types.responses.response_compaction_item import ResponseCompactionItem +from openai.types.responses.response_create_params import ResponseInputParam +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError + +from litellm.responses.utils import ResponsesAPIRequestUtils + +from .tool_search import chat_config_builds_anthropic_request, provider_chat_config + +_JSON_OBJECT: Final = TypeAdapter(dict[str, object]) +_OBJECT_LIST: Final = TypeAdapter(list[object]) +_Item: Final = TypeVar("_Item") +_INSTRUCTION_ROLES: Final = frozenset(("system", "developer")) + + +def _json_object(value: object) -> Mapping[str, object] | None: + try: + return _JSON_OBJECT.validate_python(value) + except ValidationError: + return None + + +@dataclass(frozen=True, slots=True) +class CompactionBlock: + summary: str + wire_block: Mapping[str, object] + + +def _compaction_block(value: object) -> CompactionBlock | None: + block: Final = _json_object(value) + if block is None: + return None + summary: Final = block.get("content") + if block.get("type") != "compaction" or not isinstance(summary, str) or not summary: + return None + return CompactionBlock(summary=summary, wire_block=block) + + +def _newest_compaction_block(compaction_blocks: object) -> CompactionBlock | None: + if not isinstance(compaction_blocks, list): + return None + blocks: Final = _OBJECT_LIST.validate_python(compaction_blocks) + return next((block for value in reversed(blocks) if (block := _compaction_block(value)) is not None), None) + + +def _provider_compaction_block(provider_specific_fields: object) -> CompactionBlock | None: + fields: Final = _json_object(provider_specific_fields) + return None if fields is None else _newest_compaction_block(fields.get("compaction_blocks")) + + +def provider_compaction_summary(provider_specific_fields: object) -> str | None: + block: Final = _provider_compaction_block(provider_specific_fields) + return None if block is None else block.summary + + +class _AssistantTurnFields(BaseModel): + model_config = ConfigDict(frozen=True, from_attributes=True) + + content: object = None + tool_calls: object = None + provider_specific_fields: object = None + + +def carries_only_compaction(message: object) -> bool: + fields: Final = _AssistantTurnFields.model_validate(message) + return ( + not fields.content + and not fields.tool_calls + and provider_compaction_summary(fields.provider_specific_fields) is not None + ) + + +def _choice_provider_specific_fields(choice: object) -> object: + choice_fields: Final = _json_object(choice) + message: Final = None if choice_fields is None else _json_object(choice_fields.get("message")) + return None if message is None else message.get("provider_specific_fields") + + +def response_holds_compaction_block(response: object) -> bool: + fields: Final = _json_object(response) + if fields is None: + return False + if fields.get("object") == "response": + return holds_compaction_summary(fields.get("output")) + try: + choices: Final = _OBJECT_LIST.validate_python(fields.get("choices")) + except ValidationError: + return False + return any(provider_compaction_summary(_choice_provider_specific_fields(choice)) is not None for choice in choices) + + +def compaction_output_item(provider_specific_fields: object) -> ResponseCompactionItem | None: + block: Final = _provider_compaction_block(provider_specific_fields) + if block is None: + return None + return ResponseCompactionItem( + id=f"cmp_{uuid.uuid4()}", + encrypted_content=json.dumps(dict(block.wire_block), separators=(",", ":")), + type="compaction", + ) + + +def _entry_type(entry: object) -> object: + fields: Final = _json_object(entry) + return None if fields is None else fields.get("type") + + +def _has_entry_of_type(entries: object, entry_type: str) -> bool: + try: + rows: Final = _OBJECT_LIST.validate_python(entries) + except ValidationError: + return False + return any(_entry_type(row) == entry_type for row in rows) + + +def requests_compaction(context_management: object) -> bool: + if _has_entry_of_type(context_management, "compaction"): + return True + anthropic_form: Final = _json_object(context_management) + return anthropic_form is not None and _has_entry_of_type(anthropic_form.get("edits"), "compact_20260112") + + +def provider_takes_compaction_blocks(model: str, custom_llm_provider: str | None) -> bool: + from litellm.llms.bedrock.claude_platform.transformation import BedrockClaudePlatformConfig + + chat_config: Final = provider_chat_config(model, custom_llm_provider) + return chat_config_builds_anthropic_request(chat_config) and not isinstance( + chat_config, BedrockClaudePlatformConfig + ) + + +def _without_affinity_wrap(encrypted_content: str) -> str: + _, original = ResponsesAPIRequestUtils.unwrap_encrypted_content_with_model_id(encrypted_content) + return original + + +def decode_compaction_item(encrypted_content: object) -> CompactionBlock | None: + if not isinstance(encrypted_content, str): + return None + try: + decoded: Final = _JSON_OBJECT.validate_json(_without_affinity_wrap(encrypted_content)) + except ValidationError: + return None + return _compaction_block(decoded) + + +def compaction_summary(item: object) -> str | None: + fields: Final = _json_object(item) + if fields is None or fields.get("type") != "compaction": + return None + block: Final = decode_compaction_item(fields.get("encrypted_content")) + return None if block is None else block.summary + + +def summary_only_compaction_content(summary: str) -> str: + return json.dumps({"type": "compaction", "content": summary}, separators=(",", ":")) + + +def with_compaction_summary(encrypted_content: str, summary: str) -> str: + original: Final = _without_affinity_wrap(encrypted_content) + block: Final = decode_compaction_item(original) + if block is None: + return encrypted_content + rewritten: Final = json.dumps({**block.wire_block, "content": summary}, separators=(",", ":")) + return encrypted_content.removesuffix(original) + rewritten + + +def with_visited_compaction_summary(item: object, visit: Callable[[str], str]) -> Mapping[str, object] | None: + fields: Final = _json_object(item) + if fields is None: + return None + summary: Final = compaction_summary(fields) + encrypted_content: Final = fields.get("encrypted_content") + if summary is None or not isinstance(encrypted_content, str): + return None + return {**fields, "encrypted_content": with_compaction_summary(encrypted_content, visit(summary))} + + +def holds_compaction_summary(input: object) -> bool: + try: + items: Final = _OBJECT_LIST.validate_python(input) + except ValidationError: + return False + return any(compaction_summary(item) is not None for item in items) + + +def _is_instruction_item(item: object) -> bool: + fields: Final = _json_object(item) + if fields is None or fields.get("type", "message") != "message": + return False + role: Final = fields.get("role") + return isinstance(role, str) and role in _INSTRUCTION_ROLES + + +def input_items(input: object) -> tuple[object, ...]: + try: + return tuple(_OBJECT_LIST.validate_python(input)) + except ValidationError: + return () + + +def instruction_items(input: object) -> tuple[object, ...]: + return tuple(item for item in input_items(input) if _is_instruction_item(item)) + + +def _call_id(item: object) -> str | None: + fields: Final = _json_object(item) + call_id: Final = None if fields is None else fields.get("call_id") + return call_id if isinstance(call_id, str) else None + + +def _is_turn_item(item: object) -> bool: + fields: Final = _json_object(item) + item_type: Final = None if fields is None else fields.get("type", "message") + if fields is None or not isinstance(item_type, str): + return False + if item_type == "message": + return fields.get("role") == "assistant" + return item_type == "reasoning" or (_call_id(fields) is not None and not item_type.endswith("_output")) + + +def _turn_indices_answered_later( + turn: tuple[tuple[int, object], ...], calls_answered_later: frozenset[int] +) -> tuple[int, ...]: + if not any(index in calls_answered_later for index, _ in turn): + return () + return tuple(index for index, item in turn if _call_id(item) is None or index in calls_answered_later) + + +def _is_turn_entry(entry: tuple[int, object]) -> bool: + return _is_turn_item(entry[1]) + + +def _calls_answered_only_later(earlier: Sequence[object], answered: frozenset[str]) -> frozenset[int]: + last_index_by_call_id: Final = MappingProxyType( + {call_id: index for index, item in enumerate(earlier) if (call_id := _call_id(item)) is not None} + ) + return frozenset(index for call_id, index in last_index_by_call_id.items() if call_id in answered) + + +def _earlier_turn_indices_answered_later(earlier: Sequence[object], answered: frozenset[str]) -> tuple[int, ...]: + calls_answered_later: Final = _calls_answered_only_later(earlier, answered) + turns: Final = (tuple(turn) for is_turn, turn in groupby(enumerate(earlier), key=_is_turn_entry) if is_turn) + return tuple(chain.from_iterable(_turn_indices_answered_later(turn, calls_answered_later) for turn in turns)) + + +def _call_ids(items: Sequence[object]) -> frozenset[str]: + return frozenset(call_id for item in items if (call_id := _call_id(item)) is not None) + + +def _keeps_every_item(item: object) -> bool: + return True + + +def _with_answered_turns_after_newest_compaction( + input: Sequence[_Item], keeps_earlier_item: Callable[[object], bool] +) -> tuple[_Item, ...]: + newest: Final = next( + (index for index in range(len(input) - 1, 0, -1) if compaction_summary(input[index]) is not None), 0 + ) + later: Final = input[newest:] + moved: Final = _earlier_turn_indices_answered_later(input[:newest], _call_ids(later)) + moved_indices: Final = frozenset(moved) + moved_call_ids: Final = _call_ids([input[index] for index in moved]) + first_answer: Final = next( + (index for index, item in enumerate(later) if _call_id(item) in moved_call_ids), len(later) + ) + return ( + *(item for index, item in enumerate(input[:newest]) if index not in moved_indices and keeps_earlier_item(item)), + *later[:first_answer], + *(input[index] for index in moved), + *later[first_answer:], + ) + + +def _items_from_newest_compaction(input: Sequence[_Item]) -> tuple[_Item, ...]: + return _with_answered_turns_after_newest_compaction(input, _is_instruction_item) + + +def with_answered_turns_after_newest_compaction(input: str | ResponseInputParam) -> str | ResponseInputParam: + return ( + list(_with_answered_turns_after_newest_compaction(input, _keeps_every_item)) + if isinstance(input, list) + else input + ) + + +def without_items_before_compaction(input: str | ResponseInputParam) -> str | ResponseInputParam: + return list(_items_from_newest_compaction(input)) if isinstance(input, list) else input + + +def _summary_or_item(item: object) -> object: + summary: Final = compaction_summary(item) + return item if summary is None else summary + + +def billable_input_since_newest_compaction(input: object) -> object: + try: + items: Final = _OBJECT_LIST.validate_python(input) + except ValidationError: + return input + return [_summary_or_item(item) for item in _items_from_newest_compaction(items)] + + +def _as_summary_message(item: _Item) -> _Item | EasyInputMessageParam: + summary: Final = compaction_summary(item) + return item if summary is None else EasyInputMessageParam(role="user", content=summary, type="message") + + +def with_compaction_items_as_summaries(input: str | ResponseInputParam) -> str | ResponseInputParam: + if not isinstance(input, list) or not holds_compaction_summary(input): + return input + return [_as_summary_message(item) for item in _items_from_newest_compaction(input)] + + +def messages_since_newest_compaction(input: object) -> object: + try: + items: Final = _OBJECT_LIST.validate_python(input) + except ValidationError: + return input + return [_as_summary_message(item) for item in _items_from_newest_compaction(items)] diff --git a/litellm/responses/litellm_completion_transformation/count_request.py b/litellm/responses/litellm_completion_transformation/count_request.py index 56a0b9024c5..7679b4ede71 100644 --- a/litellm/responses/litellm_completion_transformation/count_request.py +++ b/litellm/responses/litellm_completion_transformation/count_request.py @@ -1,4 +1,4 @@ -from collections.abc import Mapping +from collections.abc import Awaitable, Callable, Mapping from dataclasses import dataclass from typing import Final, cast # noqa: TID251 # input items are forwarded as sent, as inference does @@ -14,7 +14,11 @@ from litellm.litellm_core_utils.tool_fields import DropPolicy from litellm.llms.base_llm.files.batch_records import responses_batch_body_to_chat_body from litellm.responses.additional_tools import hoist_additional_tools from litellm.responses.file_search.emulated_handler import replace_file_search_tools -from litellm.responses.litellm_completion_transformation.transformation import LiteLLMCompletionResponsesConfig +from litellm.responses.litellm_completion_transformation.compaction_items import holds_compaction_summary +from litellm.responses.litellm_completion_transformation.transformation import ( + ChatCompletionSession, + LiteLLMCompletionResponsesConfig, +) from litellm.responses.main import previous_response_client_ids from litellm.types.llms.openai import ResponseInputParam from litellm.utils import ProviderConfigManager @@ -132,12 +136,17 @@ def _with_emulated_file_search(responses_request: Mapping[str, object]) -> Mappi return {**responses_request, "tools": replace_file_search_tools(_ITEMS.validate_python(tools))[0]} +def _with_compaction_as_summary_text(responses_request: Mapping[str, object]) -> Mapping[str, object]: + return {name: value for name, value in responses_request.items() if name != "context_management"} + + async def _with_previous_turns( chat_body: Mapping[str, object], model: str, custom_llm_provider: str, responses_request: Mapping[str, object], policy: DropPolicy, + session_lookup: Callable[[str], Awaitable[ChatCompletionSession]] | None, ) -> Mapping[str, object]: previous_response_id: Final = responses_request.get("previous_response_id") if not isinstance(previous_response_id, str) or not previous_response_id: @@ -147,8 +156,10 @@ async def _with_previous_turns( previous_response_id=previous_response_id, litellm_completion_request={**chat_body, "model": model, "custom_llm_provider": custom_llm_provider}, instructions=instructions if isinstance(instructions, str) else None, + session_lookup=session_lookup, request_kwargs={"drop_params": policy.drop_params}, previous_response_client_ids=previous_response_client_ids(previous_response_id, responses_request), + input_replaces_history=holds_compaction_summary(responses_request.get("input")), ) @@ -171,6 +182,7 @@ async def chat_count_request( api_base: str | None, responses_request: Mapping[str, object], policy: DropPolicy, + session_lookup: Callable[[str], Awaitable[ChatCompletionSession]] | None = None, ) -> ChatCountRequest | ChatCountBridgeError | None: """The chat request Responses inference sends a provider with no native Responses API, or None when inference sends this provider the Responses body itself.""" @@ -184,7 +196,9 @@ async def chat_count_request( try: chat_body: Final = await _with_previous_turns( responses_batch_body_to_chat_body( - _with_emulated_file_search(_with_hoisted_additional_tools(responses_request)), + _with_compaction_as_summary_text( + _with_emulated_file_search(_with_hoisted_additional_tools(responses_request)) + ), custom_llm_provider, policy.drop_params, ), @@ -192,6 +206,7 @@ async def chat_count_request( custom_llm_provider, responses_request, policy, + session_lookup, ) tools: Final = ( () if "tools" in policy.additional_drop_params else tuple(map(_unmarked, _mappings(chat_body.get("tools")))) diff --git a/litellm/responses/litellm_completion_transformation/handler.py b/litellm/responses/litellm_completion_transformation/handler.py index 5489d2223ac..c9d800421f7 100644 --- a/litellm/responses/litellm_completion_transformation/handler.py +++ b/litellm/responses/litellm_completion_transformation/handler.py @@ -9,11 +9,13 @@ from pydantic import TypeAdapter import litellm from litellm.responses.additional_tools import hoist_additional_tools +from litellm.responses.litellm_completion_transformation.compaction_items import holds_compaction_summary from litellm.responses.litellm_completion_transformation.streaming_iterator import ( LiteLLMCompletionStreamingIterator, ) from litellm.responses.litellm_completion_transformation.transformation import ( LiteLLMCompletionResponsesConfig, + with_compaction_items_as_params, ) from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator from litellm.types.llms.openai import ( @@ -120,6 +122,7 @@ class LiteLLMCompletionTransformationHandler: instructions=responses_api_request.get("instructions"), request_kwargs=_REQUEST_KWARGS_ADAPTER.validate_python(kwargs), previous_response_client_ids=previous_response_client_ids, + input_replaces_history=holds_compaction_summary(with_compaction_items_as_params(request_input)), ) acompletion_args: Final = {} diff --git a/litellm/responses/litellm_completion_transformation/session_handler.py b/litellm/responses/litellm_completion_transformation/session_handler.py index 35160830fdc..d50e9a61edc 100644 --- a/litellm/responses/litellm_completion_transformation/session_handler.py +++ b/litellm/responses/litellm_completion_transformation/session_handler.py @@ -1,12 +1,20 @@ import asyncio import json from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from itertools import chain from typing import TYPE_CHECKING, Any, Final, cast import litellm from litellm._logging import verbose_proxy_logger from litellm.constants import REDACTED_BY_LITELLM, REDACTED_TOOL_CALL_ARGUMENTS_PLACEHOLDER from litellm.litellm_core_utils.cold_storage_handler import ColdStorageHandler +from litellm.responses.litellm_completion_transformation.compaction_items import ( + holds_compaction_summary, + input_items, + instruction_items, + response_holds_compaction_block, +) from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.llms.openai import ( AllMessageValues, @@ -20,6 +28,7 @@ if TYPE_CHECKING: from litellm.proxy._types import SpendLogsMetadata, SpendLogsPayload from litellm.responses.litellm_completion_transformation.transformation import ( ChatCompletionSession, + SessionMessage, ) else: ChatCompletionSession = Any @@ -66,6 +75,95 @@ def _stored_responses_api_output(stored_response: object) -> ResponseInputParam return None +@dataclass(frozen=True, slots=True) +class _ReplayStart: + index: int + at_output: bool + + +def _compaction_point( + index: int, spend_log: "SpendLogsPayload", proxy_server_request: Mapping[str, object] | None +) -> _ReplayStart | None: + if response_holds_compaction_block(spend_log.get("response")): + return _ReplayStart(index=index, at_output=True) + if proxy_server_request is not None and holds_compaction_summary(proxy_server_request.get("input")): + return _ReplayStart(index=index, at_output=False) + return None + + +def _replay_start( + spend_logs: Sequence["SpendLogsPayload"], proxy_server_requests: Sequence[Mapping[str, object] | None] +) -> _ReplayStart: + return next( + ( + start + for index in range(len(spend_logs) - 1, -1, -1) + if (start := _compaction_point(index, spend_logs[index], proxy_server_requests[index])) is not None + ), + _ReplayStart(index=0, at_output=False), + ) + + +def _holds_compaction_point( + spend_logs: Sequence["SpendLogsPayload"], proxy_server_requests: Sequence[Mapping[str, object] | None] +) -> bool: + return any( + _compaction_point(index, spend_log, proxy_server_request) is not None + for index, (spend_log, proxy_server_request) in enumerate(zip(spend_logs, proxy_server_requests, strict=True)) + ) + + +def _referenced_request_ids(previous_response_id: str, client_response_ids: tuple[str, ...]) -> tuple[str, ...]: + decoded_response_id: Final = ResponsesAPIRequestUtils.decode_responses_api_response_id(previous_response_id) + response_id: Final = decoded_response_id.get("response_id", previous_response_id) + return tuple(dict.fromkeys((response_id, previous_response_id, *client_response_ids))) + + +def _turns_through_referenced_response( + spend_logs: Sequence["SpendLogsPayload"], referenced_request_ids: tuple[str, ...] +) -> int: + return next( + ( + index + 1 + for index in range(len(spend_logs) - 1, -1, -1) + if spend_logs[index].get("request_id") in referenced_request_ids + ), + len(spend_logs), + ) + + +def _replayed_request( + index: int, + proxy_server_request: Mapping[str, object] | None, + start: _ReplayStart, + earlier_instructions: tuple[object, ...], +) -> Mapping[str, object] | None: + if proxy_server_request is None or index != start.index: + return proxy_server_request + stored_input: Final = proxy_server_request.get("input") + if start.at_output: + return { + **proxy_server_request, + "input": [*earlier_instructions, *instruction_items(stored_input)], + "messages": None, + } + if not earlier_instructions: + return proxy_server_request + return {**proxy_server_request, "input": [*earlier_instructions, *input_items(stored_input)]} + + +def _instruction_items_before( + proxy_server_requests: Sequence[Mapping[str, object] | None], start: _ReplayStart +) -> tuple[object, ...]: + return tuple( + chain.from_iterable( + instruction_items(request.get("input")) + for request in proxy_server_requests[: start.index] + if request is not None + ) + ) + + def _stored_instructions(proxy_server_request: Mapping[str, object] | None) -> str | None: instructions: Final = None if proxy_server_request is None else proxy_server_request.get("instructions") return instructions if isinstance(instructions, str) and instructions else None @@ -91,6 +189,7 @@ class ResponsesSessionHandler: """ from litellm.responses.litellm_completion_transformation.transformation import ( ChatCompletionSession, + with_stored_calls_the_input_answers, ) verbose_proxy_logger.debug("inside get_chat_completion_message_history_for_previous_response_id") @@ -107,26 +206,34 @@ class ResponsesSessionHandler: if len(all_spend_logs) > 0: litellm_session_id = all_spend_logs[0].get("session_id") - chat_completion_message_history: list[ - AllMessageValues - | GenericChatCompletionMessage - | ChatCompletionMessageToolCall - | ChatCompletionResponseMessage - | Message - ] = [] proxy_server_requests: Final = [ await ResponsesSessionHandler.get_proxy_server_request_from_spend_log(spend_log=spend_log) for spend_log in all_spend_logs ] - for spend_log, proxy_server_request_dict in zip(all_spend_logs, proxy_server_requests): - chat_completion_message_history = ( - ResponsesSessionHandler.extend_chat_completion_message_with_spend_log_payload( - spend_log=spend_log, - proxy_server_request_dict=proxy_server_request_dict, - chat_completion_message_history=chat_completion_message_history, - item_cache_control=item_cache_control, - ) - ) + replayable_turns: Final = _turns_through_referenced_response( + all_spend_logs, _referenced_request_ids(previous_response_id, client_response_ids) + ) + start: Final = _replay_start(all_spend_logs[:replayable_turns], proxy_server_requests[:replayable_turns]) + replayed_turns: Final = ( + replayable_turns if _holds_compaction_point(all_spend_logs, proxy_server_requests) else len(all_spend_logs) + ) + earlier_instructions: Final = _instruction_items_before(proxy_server_requests, start) + replayed_history: Final = ResponsesSessionHandler._history_of_turns( + all_spend_logs[start.index : replayed_turns], + [ + _replayed_request(index, proxy_server_requests[index], start, earlier_instructions) + for index in range(start.index, replayed_turns) + ], + item_cache_control, + ) + history_before_replay: Final = ResponsesSessionHandler._history_of_turns( + all_spend_logs[: start.index], proxy_server_requests[: start.index], item_cache_control + ) + chat_completion_message_history: Final = ( + list(with_stored_calls_the_input_answers(replayed_history, history_before_replay)) + if history_before_replay + else replayed_history + ) verbose_proxy_logger.debug( "chat_completion_message_history %s", @@ -138,6 +245,24 @@ class ResponsesSessionHandler: instructions=_latest_replayable_instructions(proxy_server_requests), ) + @staticmethod + def _history_of_turns( + spend_logs: Sequence["SpendLogsPayload"], + proxy_server_requests: Sequence[Mapping[str, object] | None], + item_cache_control: bool, + ) -> list["SessionMessage"]: + return list( + chain.from_iterable( + ResponsesSessionHandler.extend_chat_completion_message_with_spend_log_payload( + spend_log=spend_log, + proxy_server_request_dict=proxy_server_request, + chat_completion_message_history=[], + item_cache_control=item_cache_control, + ) + for spend_log, proxy_server_request in zip(spend_logs, proxy_server_requests, strict=True) + ) + ) + @staticmethod def extend_chat_completion_message_with_spend_log_payload( spend_log: "SpendLogsPayload", @@ -182,6 +307,7 @@ class ResponsesSessionHandler: responses_api_request={}, replay_reasoning=True, item_cache_control=item_cache_control, + only_since_compaction=True, ) ) @@ -196,6 +322,7 @@ class ResponsesSessionHandler: responses_api_request={}, replay_reasoning=True, item_cache_control=item_cache_control, + only_since_compaction=True, ) ) elif isinstance(_response_output, dict) and _response_output and _response_output != {}: @@ -342,12 +469,11 @@ class ResponsesSessionHandler: verbose_proxy_logger.debug("decoding response id=%s", previous_response_id) - decoded_response_id: Final = ResponsesAPIRequestUtils.decode_responses_api_response_id(previous_response_id) - response_id: Final = decoded_response_id.get("response_id", previous_response_id) + request_ids: Final = list(_referenced_request_ids(previous_response_id, client_response_ids)) + response_id: Final = request_ids[0] if prisma_client is None: return [] - request_ids: Final = list(dict.fromkeys((response_id, previous_response_id, *client_response_ids))) query: Final = """ WITH matching_session AS ( SELECT session_id, api_key, "user", team_id diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index 4ef68aa072a..c25dd4a405e 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -8,6 +8,7 @@ import litellm from litellm.litellm_core_utils.hidden_params import get_or_create_hidden_params, served_from_cache from litellm.litellm_core_utils.litellm_logging import BilledRoundsSpend from litellm.main import stream_chunk_builder +from litellm.responses.litellm_completion_transformation.compaction_items import compaction_output_item from litellm.responses.litellm_completion_transformation.custom_tools import ( build_tool_call_item_kwargs, extract_custom_tool_names, @@ -83,6 +84,11 @@ def _output_items_with_id(items: tuple[Any, ...], item_type: str, item_id: str | ) +def _streams_compaction(chunk: ModelResponseStream) -> bool: + delta_fields: Final = chunk.choices[0].delta.provider_specific_fields if chunk.choices else None + return delta_fields is not None and "compaction_blocks" in delta_fields + + def _delta_has_signed_thinking_block(delta: object) -> bool: blocks: Final = getattr(delta, "thinking_blocks", None) or () return any(isinstance(b, dict) and (b.get("signature") or b.get("data")) for b in blocks) @@ -146,6 +152,9 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._cached_item_id: str | None = None self._message_output_index: int | None = None self._reasoning_output_index: int | None = None + self._compaction_output_index: int | None = None + self._compaction_item_id: str | None = None + self._sent_compaction_item_done: bool = False self._cached_response_id: str | None = None self._buffered_chunk: ModelResponseStream | None = None self._upstream_exhausted: bool = False @@ -186,7 +195,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._tool_search_call_ids = self._tool_search_call_ids | {call_id} return True - def _queue_output_item_added(self, output_index: int, item: BaseLiteLLMOpenAIResponseObject) -> None: + def _next_output_item_added(self, output_index: int, item: BaseLiteLLMOpenAIResponseObject) -> OutputItemAddedEvent: self._sequence_number += 1 added: Final = OutputItemAddedEvent( type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, @@ -194,7 +203,10 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): item=item, ) added.__dict__["sequence_number"] = self._sequence_number - self._pending_tool_events.append(added) + return added + + def _queue_output_item_added(self, output_index: int, item: BaseLiteLLMOpenAIResponseObject) -> None: + self._pending_tool_events.append(self._next_output_item_added(output_index, item)) def _queue_tool_search_call_added(self, call_id: str, output_index: int, args_delta: str) -> None: if call_id not in self._tool_args_by_call_id: @@ -248,6 +260,8 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): return self._message_output_index case "reasoning": return self._reasoning_output_index + case "compaction": + return self._compaction_output_index case _: call_id: Final = getattr(item, "call_id", None) or str(getattr(item, "id", "")).removeprefix("ws_") return self._tool_output_index_by_call_id.get(str(call_id)) @@ -256,6 +270,45 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): index: Final = self._streamed_output_index(item) return (index is None, index or 0) + def _streamed_compaction_item(self) -> BaseLiteLLMOpenAIResponseObject | None: + item: Final = compaction_output_item(self._accumulated_provider_specific_fields) + if item is None: + return None + self._compaction_item_id = self._compaction_item_id or item.id + return BaseLiteLLMOpenAIResponseObject(**{**item.model_dump(exclude_none=True), "id": self._compaction_item_id}) + + def _queue_compaction_item_added(self) -> None: + if ( + self._compaction_output_index is not None + or "compaction_blocks" not in self._accumulated_provider_specific_fields + ): + return + streamed_item: Final = self._streamed_compaction_item() + if streamed_item is None: + return + self._compaction_output_index = self._allocate_output_index() + unfinished_item: Final = streamed_item.model_copy(update={"encrypted_content": ""}) + self._pending_response_events.append( + self._next_output_item_added(self._compaction_output_index, unfinished_item) + ) + + def _queue_compaction_item_done(self) -> None: + if self._compaction_output_index is None or self._sent_compaction_item_done: + return + streamed_item: Final = self._streamed_compaction_item() + if streamed_item is None: + return + self._sent_compaction_item_done = True + self._sequence_number += 1 + self._pending_response_events.append( + OutputItemDoneEvent( + type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, + output_index=self._compaction_output_index, + sequence_number=self._sequence_number, + item=streamed_item, + ) + ) + def _normalize_tool_call_index(self, tool_call: object) -> int | None: idx_raw: Final = tool_call.get("index") if isinstance(tool_call, dict) else getattr(tool_call, "index", None) if idx_raw is None: @@ -592,6 +645,12 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): } ) + def _ordered_chunk(self, chunk: ModelResponseStream) -> ModelResponseStream: + ordered: Final = self._reasoning_part_first(chunk) + if not _streams_compaction(ordered): + self._queue_compaction_item_added() + return ordered + def _with_encoded_response_id(self, response: ResponsesAPIResponse) -> ResponsesAPIResponse: return ResponsesAPIRequestUtils.update_responses_api_response_id_with_model_id( responses_api_response=response, @@ -1004,6 +1063,10 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): if not self.litellm_model_response or isinstance(self.litellm_model_response, TextCompletionResponse): self.litellm_model_response = self.create_litellm_model_response() if self.litellm_model_response: + self._queue_compaction_item_added() + self._queue_compaction_item_done() + if self._pending_response_events: + return self._pending_response_events.pop(0) # If tool calls exist, emit tool events before finishing/response.completed. if isinstance(self.litellm_model_response, ModelResponse): self._queue_final_tool_call_done_events(self.litellm_model_response) @@ -1091,6 +1154,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._reasoning_item_id or self._cached_reasoning_item_id or mint_reasoning_item_id() ) reasoning_item_id: Final = self._cached_reasoning_item_id + self._queue_compaction_item_done() self._sequence_number += 1 text_done_event: Final = self.create_reasoning_summary_text_done_event( reasoning_item_id=reasoning_item_id, @@ -1139,7 +1203,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): raise StopAsyncIteration chunk = await self.litellm_custom_stream_wrapper.__anext__() if chunk is not None: - chunk = self._reasoning_part_first(cast(ModelResponseStream, chunk)) + chunk = self._ordered_chunk(cast(ModelResponseStream, chunk)) for src in ( getattr(chunk, "provider_specific_fields", None), getattr( @@ -1195,11 +1259,11 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): try: buffered_chunk = self._take_buffered_chunk() if buffered_chunk is not None: - chunk = self._reasoning_part_first(buffered_chunk) + chunk = self._ordered_chunk(buffered_chunk) elif self._upstream_exhausted: raise StopIteration else: - chunk = self._reasoning_part_first(self.litellm_custom_stream_wrapper.__next__()) + chunk = self._ordered_chunk(self.litellm_custom_stream_wrapper.__next__()) for src in ( getattr(chunk, "provider_specific_fields", None), getattr( @@ -1356,7 +1420,8 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): "reasoning", self._cached_reasoning_item_id, ) - streamed_items: Final = filterfalse(self._is_unstreamed_empty_message, reasoning_aligned) + compaction_aligned: Final = _output_items_with_id(reasoning_aligned, "compaction", self._compaction_item_id) + streamed_items: Final = filterfalse(self._is_unstreamed_empty_message, compaction_aligned) return tuple(sorted(streamed_items, key=self._streamed_output_position)) def _is_unstreamed_empty_message(self, item: object) -> bool: @@ -1388,6 +1453,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): messages=LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( # pyright: ignore[reportUnknownMemberType] # the transformer's signature is partially untyped input=self.request_input, responses_api_request=dict(self.responses_api_request), + only_since_compaction=True, ), logging_obj=self.litellm_logging_obj, ) diff --git a/litellm/responses/litellm_completion_transformation/tool_search.py b/litellm/responses/litellm_completion_transformation/tool_search.py index 0a94acdd461..780a23e02b2 100644 --- a/litellm/responses/litellm_completion_transformation/tool_search.py +++ b/litellm/responses/litellm_completion_transformation/tool_search.py @@ -1,5 +1,5 @@ from collections.abc import Iterable, Iterator, Mapping, Sequence, Set -from typing import Final, Literal, TypeAlias +from typing import TYPE_CHECKING, Final, Literal, TypeAlias from pydantic import ConfigDict, TypeAdapter, ValidationError @@ -12,6 +12,9 @@ from litellm.types.utils import LlmProviders, Message from .custom_tools import openai_shaped_tool_call_item_id +if TYPE_CHECKING: + from litellm.llms.base_llm.chat.transformation import BaseConfig + __all__ = [ "ANTHROPIC_TOOL_SEARCH_TOOL_NAME", "ANTHROPIC_TOOL_SEARCH_TOOL_TYPE", @@ -31,6 +34,7 @@ __all__ = [ "is_openai_tool_search_tool", "is_server_tool_search_tool", "provider_builds_anthropic_request", + "provider_chat_config", "runs_tool_search_on_the_server", "tool_search_call_ids", "tools_with_anthropic_tool_search_mapped", @@ -159,7 +163,7 @@ def input_without_unanswered_anthropic_tool_search_calls( return input if len(kept) == len(input) else kept -def provider_builds_anthropic_request(model: str, custom_llm_provider: str | None) -> bool: +def provider_chat_config(model: str, custom_llm_provider: str | None) -> "BaseConfig | None": from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.utils import ProviderConfigManager @@ -167,10 +171,12 @@ def provider_builds_anthropic_request(model: str, custom_llm_provider: str | Non try: llm_provider: Final = LlmProviders(provider) except ValueError: - return False - return chat_config_builds_anthropic_request( - ProviderConfigManager.get_provider_chat_config(model=model, provider=llm_provider) - ) + return None + return ProviderConfigManager.get_provider_chat_config(model=model, provider=llm_provider) + + +def provider_builds_anthropic_request(model: str, custom_llm_provider: str | None) -> bool: + return chat_config_builds_anthropic_request(provider_chat_config(model, custom_llm_provider)) def chat_config_builds_anthropic_request(chat_config: object) -> bool: diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index d9ebf50cc1a..02d9566d396 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -8,7 +8,7 @@ import uuid from collections.abc import Awaitable, Callable, Iterable, Iterator, Mapping, Sequence from dataclasses import dataclass from functools import partial -from itertools import accumulate, chain, islice +from itertools import accumulate, chain, groupby, islice from types import MappingProxyType from typing import ( TYPE_CHECKING, @@ -29,6 +29,7 @@ from openai.types.chat.chat_completion_named_tool_choice_param import ( Function as NamedToolChoiceFunction, ) from openai.types.responses import ResponseFunctionToolCall, ResponseFunctionWebSearch +from openai.types.responses.response_compaction_item import ResponseCompactionItem from openai.types.responses.response_create_params import ResponseInputParam from openai.types.responses.tool_choice_custom_param import ToolChoiceCustomParam from openai.types.responses.tool_choice_function_param import ToolChoiceFunctionParam @@ -120,6 +121,17 @@ from litellm.types.utils import ( Usage, ) +from .compaction_items import ( + carries_only_compaction, + compaction_output_item, + decode_compaction_item, + holds_compaction_summary, + provider_compaction_summary, + provider_takes_compaction_blocks, + requests_compaction, + with_answered_turns_after_newest_compaction, + without_items_before_compaction, +) from .custom_tools import ( convert_custom_tool_to_function_tool, extract_custom_tool_names, @@ -224,6 +236,7 @@ TOOL_CALLS_CACHE: Final = InMemoryCache() _ANY_KEY_DICT_ADAPTER: Final = TypeAdapter(dict[object, object]) _STR_KEY_DICT_ADAPTER: Final = TypeAdapter(dict[str, object]) _OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object]) +_MessageT = TypeVar("_MessageT") _CACHED_CONTENT_ADAPTER: Final = TypeAdapter(ChatCompletionCachedContent) _DICT_ITEMS_LIST_ADAPTER: Final = TypeAdapter(list[dict[object, object]]) _TEXT_ADAPTER: Final = TypeAdapter(str) @@ -306,6 +319,126 @@ def _without_stored_tool_search_uses(message: SessionMessage, search_call_ids: f return message +def _instruction_messages(messages: Sequence[SessionMessage]) -> tuple[SessionMessage, ...]: + return tuple( + message + for message in messages + if isinstance(message, Mapping) and message.get("role") in ("system", "developer") + ) + + +def _message_field(message: object, key: str) -> object: + if isinstance(message, Mapping): + return _STR_KEY_DICT_ADAPTER.validate_python(message).get(key) + return getattr(message, key, None) + + +def _str_field(message: object, key: str) -> str | None: + value: Final = _message_field(message, key) + return value if isinstance(value, str) else None + + +def _tool_output_call_id(message: object) -> str | None: + return _str_field(message, "tool_call_id") if _message_field(message, "role") == "tool" else None + + +def _assistant_tool_calls(message: object) -> tuple[object, ...]: + if _message_field(message, "role") != "assistant": + return () + return tuple(_OBJECT_LIST_ADAPTER.validate_python(_message_field(message, "tool_calls") or [])) + + +def _answered_tool_call_ids(messages: Sequence[object]) -> frozenset[str]: + called: Final = frozenset( + _str_field(call, "id") for call in chain.from_iterable(map(_assistant_tool_calls, messages)) + ) + return frozenset( + call_id for call_id in map(_tool_output_call_id, messages) if call_id is not None and call_id not in called + ) + + +def _stored_turn_with_calls(message: SessionMessage, call_ids: frozenset[str]) -> SessionMessage | None: + calls: Final = [call for call in _assistant_tool_calls(message) if _str_field(call, "id") in call_ids] + if not calls: + return None + if isinstance(message, Message): + return message.model_copy( + update={"tool_calls": calls, "provider_specific_fields": _provider_fields_without_compaction(message)} + ) + return cast( # cast-ok: the stored assistant turn keeps every field but the calls the new input leaves unanswered + SessionMessage, {**_STR_KEY_DICT_ADAPTER.validate_python(message), "tool_calls": calls} + ) + + +def _call_ids_at_turn(indexed_message: tuple[int, object]) -> tuple[tuple[str | None, int], ...]: + index, message = indexed_message + return tuple((_str_field(call, "id"), index) for call in _assistant_tool_calls(message)) + + +def _newest_stored_calls_the_input_answers( + session_history: Sequence[SessionMessage], answered: frozenset[str] +) -> Iterator[SessionMessage | None]: + newest_turn_by_call_id: Final = MappingProxyType( + dict(chain.from_iterable(map(_call_ids_at_turn, enumerate(session_history)))) + ) + answered_turns: Final = tuple( + sorted((newest_turn_by_call_id[call_id], call_id) for call_id in answered if call_id in newest_turn_by_call_id) + ) + answered_by_turn: Final = MappingProxyType( + { + turn: frozenset(call_id for _, call_id in calls) + for turn, calls in groupby(answered_turns, key=lambda turn_and_call: turn_and_call[0]) + } + ) + return ( + _stored_turn_with_calls(message, answered_by_turn.get(index, frozenset())) + for index, message in enumerate(session_history) + ) + + +def with_stored_calls_the_input_answers( + new_messages: Sequence[_MessageT], session_history: Sequence[SessionMessage] +) -> tuple[_MessageT | SessionMessage, ...]: + answered: Final = _answered_tool_call_ids(new_messages) + stored_calls: Final = tuple( + turn for turn in _newest_stored_calls_the_input_answers(session_history, answered) if turn is not None + ) + if not stored_calls: + return tuple(new_messages) + first_output: Final = next( + index for index, message in enumerate(new_messages) if _tool_output_call_id(message) in answered + ) + return (*new_messages[:first_output], *stored_calls, *new_messages[first_output:]) + + +def _stored_compaction_summary(message: SessionMessage) -> str | None: + return provider_compaction_summary(message.provider_specific_fields) if isinstance(message, Message) else None + + +def _provider_fields_without_compaction(message: Message) -> Mapping[str, object] | None: + if not message.provider_specific_fields: + return None + other_fields: Final = { + key: value + for key, value in _STR_KEY_DICT_ADAPTER.validate_python(message.provider_specific_fields).items() + if key != "compaction_blocks" + } + return other_fields or None + + +def _with_stored_compaction_as_summary(message: SessionMessage) -> tuple[SessionMessage, ...]: + if not isinstance(message, Message): + return (message,) + summary: Final = provider_compaction_summary(message.provider_specific_fields) + if summary is None: + return (message,) + answer: Final = message.model_copy( + update={"provider_specific_fields": _provider_fields_without_compaction(message)} + ) + summary_message: Final = ChatCompletionUserMessage(role="user", content=summary) + return (summary_message, answer) if answer.content or answer.tool_calls else (summary_message,) + + def _input_item_cache_control(input_item: Mapping[str, object]) -> ChatCompletionCachedContent | None: cache_control: Final = input_item.get("cache_control") if not cache_control: @@ -340,6 +473,32 @@ def without_item_cache_control(input: str | ResponseInputParam) -> str | Respons ) +def _as_input_param(item: object) -> object: + return item.model_dump(exclude_none=True) if isinstance(item, ResponseCompactionItem) else item + + +def with_compaction_items_as_params(input: str | ResponseInputParam) -> str | ResponseInputParam: + if not isinstance(input, list): + return input + return cast( # cast-ok: the same input items, a compaction item replayed from SDK output dumped to its param dict + ResponseInputParam, [_as_input_param(item) for item in input] + ) + + +def _input_to_read( + input: str | ResponseInputParam, only_since_compaction: bool, native_compaction: bool +) -> str | ResponseInputParam: + if only_since_compaction: + return without_items_before_compaction(input) + if native_compaction: + return with_answered_turns_after_newest_compaction(input) + return input + + +def _replays_compaction_block(input: str | ResponseInputParam, context_management: object) -> bool: + return requests_compaction(context_management) and holds_compaction_summary(input) + + _PROVIDERS_KEEPING_MESSAGE_CACHE_CONTROL: Final = frozenset(("anthropic", "dashscope", "minimax", "sap", "zai")) @@ -426,9 +585,6 @@ def _block_with_cache_control(block: object, rewrite: _CacheControlRewrite) -> o return _fields_with_cache_control(_STR_KEY_DICT_ADAPTER.validate_python(block), rewrite) -_MessageT = TypeVar("_MessageT") - - def _with_cache_control(message: _MessageT, rewrite: _CacheControlRewrite) -> _MessageT: if not isinstance(message, dict): return message @@ -870,14 +1026,21 @@ class LiteLLMCompletionResponsesConfig: ) request_kwargs: Final = _STR_KEY_DICT_ADAPTER.validate_python({**kwargs, "drop_params": drop_params}) + context_management: Final = responses_api_request.get("context_management") keeps_item_marks: Final = _target_keeps_item_marks( model, custom_llm_provider, drops_unsupported_params(request_kwargs.get("drop_params")) ) + request_input: Final = with_compaction_items_as_params(input) + native_compaction: Final = _replays_compaction_block( + request_input, context_management + ) and provider_takes_compaction_blocks(model, custom_llm_provider) bridged_messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( - input=input, + input=request_input, responses_api_request=responses_api_request, replay_reasoning=True, item_cache_control=keeps_item_marks, + native_compaction=native_compaction, + only_since_compaction=not native_compaction, ) litellm_completion_request: dict = { @@ -909,7 +1072,7 @@ class LiteLLMCompletionResponsesConfig: "response_format": response_format, "reasoning_effort": reasoning.effort, "reasoning_summary": reasoning.summary, - "context_management": responses_api_request.get("context_management"), + "context_management": context_management, # litellm specific params "custom_llm_provider": custom_llm_provider, "extra_headers": extra_headers, @@ -939,6 +1102,8 @@ class LiteLLMCompletionResponsesConfig: responses_api_request: ResponsesAPIOptionalRequestParams | dict, replay_reasoning: bool = False, item_cache_control: bool = False, + native_compaction: bool = False, + only_since_compaction: bool = False, ) -> list[ AllMessageValues | GenericChatCompletionMessage @@ -960,6 +1125,19 @@ class LiteLLMCompletionResponsesConfig: guardrail scanning) leave both off, because they need every piece of text in the request to stay readable as message ``content``, with no item mark a strict message schema could refuse. + + ``native_compaction`` belongs to callers whose target builds an + Anthropic request that also asks for compaction in + ``context_management``: a ``compaction`` item then goes back as the + provider's own compaction block, which Anthropic rejects without that + compaction strategy. Everyone else reads its summary as a user message. + + ``only_since_compaction`` belongs to callers that send, count, or + rebuild what the model reads: the items before the newest + ``compaction`` item are dropped, keeping system and developer + messages, the way Anthropic ignores content before a compaction block. + Guardrail scans leave it off because they write rewrites back by input + item position. """ messages: list[ AllMessageValues @@ -975,10 +1153,13 @@ class LiteLLMCompletionResponsesConfig: ) ) + dumped_input: Final = with_compaction_items_as_params(input) + read_input: Final = _input_to_read(dumped_input, only_since_compaction, native_compaction) messages.extend( LiteLLMCompletionResponsesConfig._transform_response_input_param_to_chat_completion_message( - input=input if item_cache_control else without_item_cache_control(input), + input=read_input if item_cache_control else without_item_cache_control(read_input), replay_reasoning=replay_reasoning, + native_compaction=native_compaction, ) ) @@ -997,6 +1178,18 @@ class LiteLLMCompletionResponsesConfig: if (kept := _without_stored_tool_search_uses(message, search_call_ids)) is not None ) + @staticmethod + def _with_compaction_blocks_the_route_takes( + messages: Sequence[SessionMessage], route: CompletionRoute, context_management: object + ) -> tuple[SessionMessage, ...]: + if not any(_stored_compaction_summary(message) is not None for message in messages): + return tuple(messages) + if requests_compaction(context_management) and provider_takes_compaction_blocks( + route.model, route.custom_llm_provider + ): + return tuple(messages) + return tuple(chain.from_iterable(map(_with_stored_compaction_as_summary, messages))) + @staticmethod async def async_responses_api_session_handler( previous_response_id: str, @@ -1005,6 +1198,7 @@ class LiteLLMCompletionResponsesConfig: session_lookup: Callable[[str], Awaitable[ChatCompletionSession]] | None = None, request_kwargs: Mapping[str, object] = MappingProxyType({}), previous_response_client_ids: tuple[str, ...] = (), + input_replaces_history: bool = False, ) -> dict[str, object]: """ Async hook to get the chain of previous input and output pairs and return a list of Chat Completion messages @@ -1027,11 +1221,15 @@ class LiteLLMCompletionResponsesConfig: else ChatCompletionSession(messages=[], litellm_session_id=None) ) _messages: Final = litellm_completion_request.get("messages") or [] - stored_messages: Final = chat_completion_session.get("messages") or [] + route: Final = CompletionRoute(model=model, custom_llm_provider=provider) + session_history: Final = chat_completion_session.get("messages") or [] + stored_messages: Final = LiteLLMCompletionResponsesConfig._with_compaction_blocks_the_route_takes( + _instruction_messages(session_history) if input_replaces_history else session_history, + route, + request_fields.get("context_management"), + ) replayed_messages: Final = ( - LiteLLMCompletionResponsesConfig._without_tool_search_uses_off_anthropic( - stored_messages, CompletionRoute.model_validate(litellm_completion_request) - ) + LiteLLMCompletionResponsesConfig._without_tool_search_uses_off_anthropic(stored_messages, route) if any(holds_tool_search_use(message) for message in stored_messages) else stored_messages ) @@ -1062,11 +1260,16 @@ class LiteLLMCompletionResponsesConfig: # Store original _messages before combining for safety check original_new_messages: Final = _messages.copy() if _messages else [] - combined_messages = [*leading_system_messages, *session_messages, *_messages[instructions_end:]] + new_messages: Final = ( + with_stored_calls_the_input_answers(_messages[instructions_end:], session_history) + if input_replaces_history + else _messages[instructions_end:] + ) + combined_messages = [*leading_system_messages, *session_messages, *new_messages] # Fix: Ensure tool_results have corresponding tool_calls in previous assistant message # Pass tools parameter to help reconstruct tool_calls if not in cache - tools: Final = litellm_completion_request.get("tools") or [] + tools: Final = _OBJECT_LIST_ADAPTER.validate_python(litellm_completion_request.get("tools") or []) combined_messages = LiteLLMCompletionResponsesConfig._ensure_tool_results_have_corresponding_tool_calls( messages=combined_messages, tools=tools ) @@ -1115,6 +1318,7 @@ class LiteLLMCompletionResponsesConfig: def _transform_response_input_param_to_chat_completion_message( input: str | ResponseInputParam, replay_reasoning: bool = False, + native_compaction: bool = False, ) -> list[ AllMessageValues | GenericChatCompletionMessage | ChatCompletionMessageToolCall | ChatCompletionResponseMessage ]: @@ -1140,6 +1344,7 @@ class LiteLLMCompletionResponsesConfig: LiteLLMCompletionResponsesConfig._transform_responses_api_input_item_to_chat_completion_message( input_item=_input, replay_reasoning=replay_reasoning, + native_compaction=native_compaction, ) ) @@ -1523,7 +1728,7 @@ class LiteLLMCompletionResponsesConfig: """Find the index of the previous assistant message.""" for j in range(current_idx - 1, -1, -1): if messages[j].get("role") == "assistant": - return j + return None if carries_only_compaction(messages[j]) else j return None @staticmethod @@ -1839,6 +2044,7 @@ class LiteLLMCompletionResponsesConfig: def _transform_responses_api_input_item_to_chat_completion_message( input_item: Mapping[str, object], replay_reasoning: bool = False, + native_compaction: bool = False, ) -> list[AllMessageValues | GenericChatCompletionMessage | ChatCompletionResponseMessage]: """ Transform a Responses API input item into a Chat Completion message @@ -1877,6 +2083,12 @@ class LiteLLMCompletionResponsesConfig: Mapping[str, str], input_item ) ) + elif input_item.get("type") == "compaction": + return list( + LiteLLMCompletionResponsesConfig._transform_compaction_item_to_chat_completion_messages( + input_item, native_compaction + ) + ) elif input_item.get("type") == "reasoning": # A ResponseReasoningItemParam carries the prior-turn chain-of-thought. # Chat-completions providers (DeepSeek V4, Kimi K2.6, ...) expect this @@ -1928,6 +2140,23 @@ class LiteLLMCompletionResponsesConfig: return [] return [LiteLLMCompletionResponsesConfig._transform_responses_api_message_item(input_item, content)] + @staticmethod + def _transform_compaction_item_to_chat_completion_messages( + input_item: Mapping[str, object], native_compaction: bool + ) -> tuple[ChatCompletionUserMessage | ChatCompletionResponseMessage, ...]: + compaction: Final = decode_compaction_item(input_item.get("encrypted_content")) + if compaction is None: + return () + if native_compaction: + return ( + ChatCompletionResponseMessage( + role="assistant", + content=None, + provider_specific_fields={"compaction_blocks": [dict(compaction.wire_block)]}, + ), + ) + return (ChatCompletionUserMessage(role="user", content=compaction.summary),) + @staticmethod def _transform_responses_api_message_item( input_item: Mapping[str, object], content: object @@ -3038,6 +3267,7 @@ class LiteLLMCompletionResponsesConfig: | ResponseFunctionWebSearch | CustomToolCallOutputItem | OutputToolSearchCall + | ResponseCompactionItem ]: responses_output: list[ GenericResponseOutputItem @@ -3048,8 +3278,10 @@ class LiteLLMCompletionResponsesConfig: | ResponseFunctionWebSearch | CustomToolCallOutputItem | OutputToolSearchCall + | ResponseCompactionItem ] = [] + responses_output.extend(LiteLLMCompletionResponsesConfig._extract_compaction_output_items(choices)) responses_output.extend( LiteLLMCompletionResponsesConfig._extract_reasoning_output_items(chat_completion_response, choices) ) @@ -3115,6 +3347,11 @@ class LiteLLMCompletionResponsesConfig: output_items.append(item) return output_items + @staticmethod + def _extract_compaction_output_items(choices: Sequence[Choices]) -> tuple[ResponseCompactionItem, ...]: + item: Final = compaction_output_item(choices[0].message.provider_specific_fields) if choices else None + return () if item is None else (item,) + @staticmethod def _encode_thinking_blocks(message: Message) -> str | None: thinking_blocks: Final[Sequence[Mapping[str, object]]] = getattr(message, "thinking_blocks", None) or () diff --git a/litellm/responses/main.py b/litellm/responses/main.py index 3bdb1c9c375..09eea7cb323 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -40,9 +40,15 @@ from litellm.litellm_core_utils.provider_affinity import add_provider_affinity_h from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.llms.openai_like.responses.transformation import OpenAILikeResponsesConfig +from litellm.responses.litellm_completion_transformation.compaction_items import ( + with_compaction_items_as_summaries, +) from litellm.responses.litellm_completion_transformation.handler import ( LiteLLMCompletionTransformationHandler, ) +from litellm.responses.litellm_completion_transformation.transformation import ( + with_compaction_items_as_params, +) from litellm.responses.mcp.request_context import MCPRequestContext from litellm.responses.mcp.round_billing import ( billed_for_every_round, @@ -1669,7 +1675,9 @@ def responses( ) # Decode any litellm-encoded encrypted-content item IDs back to their original IDs - input = ResponsesAPIRequestUtils.restore_encrypted_content_item_ids_in_input(input) + input = with_compaction_items_as_summaries( + with_compaction_items_as_params(ResponsesAPIRequestUtils.restore_encrypted_content_item_ids_in_input(input)) + ) # Call the handler with _is_async flag instead of directly calling the async handler if custom_llm_provider is None: @@ -2516,7 +2524,9 @@ def compact_responses( # Decode any litellm-encoded encrypted-content item IDs back to their original IDs # before forwarding to the upstream provider. - input = ResponsesAPIRequestUtils.restore_encrypted_content_item_ids_in_input(input) + input = with_compaction_items_as_summaries( + with_compaction_items_as_params(ResponsesAPIRequestUtils.restore_encrypted_content_item_ids_in_input(input)) + ) # Call the handler with _is_async flag instead of directly calling the async handler response = base_llm_http_handler.compact_response_api_handler( diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index bc30da72df9..0d59c93786c 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -1707,6 +1707,7 @@ def _estimate_usage_from_text( messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( # pyright: ignore[reportUnknownMemberType] # the transformer's signature is partially untyped input=request_input, # pyright: ignore[reportArgumentType] # the raw Responses API input is a str or ResponseInputParam list, matching the helper's declared union responses_api_request=dict(responses_api_request), + only_since_compaction=True, ) input_tokens: Final = litellm.token_counter( # pyright: ignore[reportUnknownMemberType] # token_counter's public signature is untyped model=model, messages=messages diff --git a/litellm/router.py b/litellm/router.py index fe418b5830b..da04e59f18a 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -14069,6 +14069,7 @@ class Router: input_messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( input=typed_input, responses_api_request={"instructions": instructions} if instructions is not None else {}, + only_since_compaction=True, ) return litellm.token_counter( messages=cast(list, input_messages), # cast-ok: transformed chat messages diff --git a/litellm/types/litellm_params.py b/litellm/types/litellm_params.py index 0ed156524be..089f21a33c1 100644 --- a/litellm/types/litellm_params.py +++ b/litellm/types/litellm_params.py @@ -386,6 +386,9 @@ class AgenticLoopState: @dataclass(frozen=True, slots=True, kw_only=True) class RouterState: weights: "RouterWeights | None" = field(default=None, metadata=wire("_router_weights")) + encrypted_content_affinity_pinned: bool | None = field( + default=None, metadata=wire("_encrypted_content_affinity_pinned") + ) fallback_depth: int | None = None max_fallbacks: int | None = None attempted_targets: "AttemptedFallbackTargets | None" = None diff --git a/litellm/types/llms/anthropic.py b/litellm/types/llms/anthropic.py index 6aeead7679f..ba422a0a901 100644 --- a/litellm/types/llms/anthropic.py +++ b/litellm/types/llms/anthropic.py @@ -481,6 +481,11 @@ class ContentThinkingSignatureBlockDelta(TypedDict): signature: str +class ContentCompactionBlockDelta(TypedDict): + type: ReadOnly[Literal["compaction_delta"]] + content: ReadOnly[str | None] + + StreamingContentBlockDeltaType = Literal["text_delta", "input_json_delta", "thinking_delta", "signature_delta"] @@ -493,6 +498,7 @@ class ContentBlockDelta(TypedDict): | ContentCitationsBlockDelta | ContentThinkingBlockDelta | ContentThinkingSignatureBlockDelta + | ContentCompactionBlockDelta ) diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index 8dbc0d97d54..62dc152ca84 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -55,6 +55,7 @@ except (ImportError, AttributeError): from typing import Annotated +from openai.types.responses.response_compaction_item import ResponseCompactionItem from openai.types.responses.response_create_params import ( Reasoning, ResponseIncludable, @@ -1488,6 +1489,7 @@ class ResponsesAPIResponse(BaseLiteLLMOpenAIResponseObject): | ResponseFunctionWebSearch | CustomToolCallOutputItem | OutputToolSearchCall + | ResponseCompactionItem ] ) parallel_tool_calls: bool | None = None diff --git a/tests/benchmarks/test_responses_compaction_benchmarks.py b/tests/benchmarks/test_responses_compaction_benchmarks.py new file mode 100644 index 00000000000..62a5e95255e --- /dev/null +++ b/tests/benchmarks/test_responses_compaction_benchmarks.py @@ -0,0 +1,35 @@ +import json +from collections.abc import Mapping +from typing import Final + +import pytest + +from litellm.responses.litellm_completion_transformation.compaction_items import ( + with_answered_turns_after_newest_compaction, +) + +COMPACTION_ITEM: Final[Mapping[str, object]] = { + "type": "compaction", + "id": "cmp_1", + "encrypted_content": json.dumps({"type": "compaction", "content": "Turn one planned Tokyo."}), +} +LONG_TURN_OF_ASSISTANT_TEXT: Final[tuple[Mapping[str, object], ...]] = tuple( + {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": f"Step {index}."}]} + for index in range(5_000) +) +CALL: Final[Mapping[str, object]] = { + "type": "function_call", + "call_id": "call_osaka", + "name": "get_weather", + "arguments": "{}", +} +OUTPUT: Final[Mapping[str, object]] = {"type": "function_call_output", "call_id": "call_osaka", "output": "Sunny."} + + +@pytest.mark.benchmark +def test_with_answered_turns_after_newest_compaction_long_turn_split_from_its_output() -> None: + replayed: Final = with_answered_turns_after_newest_compaction( + [*LONG_TURN_OF_ASSISTANT_TEXT, CALL, COMPACTION_ITEM, OUTPUT] + ) + + assert replayed == [COMPACTION_ITEM, *LONG_TURN_OF_ASSISTANT_TEXT, CALL, OUTPUT] diff --git a/tests/integration/providers/test_responses_compaction_items_wire.py b/tests/integration/providers/test_responses_compaction_items_wire.py new file mode 100644 index 00000000000..5b8e2b534d7 --- /dev/null +++ b/tests/integration/providers/test_responses_compaction_items_wire.py @@ -0,0 +1,1057 @@ +import json +import os +import re +import uuid +from collections.abc import Iterator, Mapping, Sequence +from dataclasses import dataclass +from itertools import chain +from pathlib import Path +from typing import Final, Literal, TypeAlias +from urllib.parse import urlsplit + +import httpx +import openai +import pytest +import yaml +from integration._support import claude_code as cc +from integration._support.client import ( + Gateway, + Scenario, + eventually, + gateway_from_environment, + object_value, + string_value, +) +from integration._support.database import read_rows +from integration._support.process import graceful_stop_seconds, owned_proxy +from integration._support.wire import Reply, Request, Wire, wire_server +from integration._support.worker_readiness import wait_for_worker_readiness +from pydantic import JsonValue, TypeAdapter + +from litellm.constants import LITELLM_TRUNCATED_PAYLOAD_FIELD +from litellm.responses.utils import ResponsesAPIRequestUtils + +pytestmark: Final = [pytest.mark.timeout(2 * graceful_stop_seconds() + 120), pytest.mark.shares_worker_state] + +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_CLAUDE: Final = cc.OPUS +_GEMINI: Final = "gemini-3.8-flash" +_OPENAI: Final = "gpt-5.6" +_PROXY_WORKERS: Final = max(1, int(os.environ.get("INTEGRATION_PROXY_WORKERS", "1"))) +_CACHE_BUST: Final[Mapping[str, JsonValue]] = {"cache": {"no-cache": True}} +_COMPACTING: Final[list[JsonValue]] = [{"type": "compaction", "compact_threshold": 150000}] +_COMPACT_BETA: Final = "compact-2026-01-12" +_CITY: Final = "Tokyo" +_MASKED_CITY: Final = "[CITY]" +_MASK_INPUT: Final = "compaction-mask-input" +_MASK_OUTPUT: Final = "compaction-mask-output" +_BLOCK_OUTPUT: Final = "compaction-block-output" +_REDACT_HEADER: Final = {"x-litellm-enable-message-redaction": "true"} +_UNREDACT_HEADER: Final = {"litellm-disable-message-redaction": "true"} +_SCRIPTED_REJECTION: Final = json.dumps( + {"type": "error", "error": {"type": "invalid_request_error", "message": "scripted rejection"}} +).encode() +_ASK: Final = re.compile(r"ask-([0-9a-f]{32})") + +Route: TypeAlias = Literal["claude-compacting", "claude-plain", "gemini", "openai", "openai-compact", "openai-opaque"] + + +def _new_marker() -> str: + return uuid.uuid4().hex + + +def _ask(marker: str, *, compact: bool = False) -> str: + return f"ask-{marker} compact-now-{marker}" if compact else f"ask-{marker}" + + +def _fail(marker: str) -> str: + return f"fail-now-{marker}" + + +def _think(marker: str) -> str: + return f"think-now-{marker}" + + +def _thinking(marker: str) -> dict[str, JsonValue]: + return {"type": "thinking", "thinking": f"thought-{marker}", "signature": f"signature-{marker}"} + + +def _answer(marker: str) -> str: + return f"answer-{marker}" + + +def _summary(marker: str) -> str: + return f"kept-{marker} {'a' * 1500} middle-{marker} {'b' * 1500} the user lives in {_CITY}" + + +def _block(summary: str) -> dict[str, JsonValue]: + return {"type": "compaction", "content": summary} + + +def _compaction_item(summary: str) -> dict[str, JsonValue]: + return { + "type": "compaction", + "id": f"cmp_{uuid.uuid4().hex}", + "encrypted_content": json.dumps(_block(summary), separators=(",", ":")), + } + + +def _newest_marker(raw: bytes) -> str: + found: Final = _ASK.findall(raw.decode()) + assert found, raw[:400] + return str(found[-1]) + + +def _anthropic_usage() -> dict[str, JsonValue]: + return {"input_tokens": 20, "output_tokens": 5} + + +def _anthropic_block_frames(index: int, block: Mapping[str, JsonValue]) -> tuple[bytes, ...]: + def start(content_block: Mapping[str, JsonValue]) -> bytes: + return cc.sse_frame( + "content_block_start", {"type": "content_block_start", "index": index, "content_block": dict(content_block)} + ) + + def delta(fields: Mapping[str, JsonValue]) -> bytes: + return cc.sse_frame( + "content_block_delta", {"type": "content_block_delta", "index": index, "delta": dict(fields)} + ) + + stop: Final = cc.sse_frame("content_block_stop", {"type": "content_block_stop", "index": index}) + match block["type"]: + case "compaction": + summary: Final = string_value(block["content"]) + half: Final = len(summary) // 2 + return ( + start({"type": "compaction", "content": ""}), + delta({"type": "compaction_delta", "content": summary[:half]}), + delta({"type": "compaction_delta", "content": summary[half:]}), + stop, + ) + case "thinking": + return ( + start({"type": "thinking", "thinking": ""}), + delta({"type": "thinking_delta", "thinking": block["thinking"]}), + delta({"type": "signature_delta", "signature": block["signature"]}), + stop, + ) + case _: + return (start({"type": "text", "text": ""}), delta({"type": "text_delta", "text": block["text"]}), stop) + + +def _anthropic_stream(identity: str, content: Sequence[Mapping[str, JsonValue]]) -> tuple[bytes, ...]: + start: Final = cc.sse_frame( + "message_start", + { + "type": "message_start", + "message": { + "id": identity, + "type": "message", + "role": "assistant", + "model": _CLAUDE, + "content": [], + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": 20, "output_tokens": 1}, + }, + }, + ) + blocks: Final = tuple( + chain.from_iterable(_anthropic_block_frames(index, block) for index, block in enumerate(content)) + ) + end: Final = cc.sse_frame( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 5}, + }, + ) + return (start, *blocks, end, cc.sse_frame("message_stop", {"type": "message_stop"})) + + +def _claude_reply(body: Mapping[str, JsonValue], raw: bytes) -> Reply: + marker: Final = _newest_marker(raw) + if _fail(marker).encode() in raw: + return Reply(status=400, body=_SCRIPTED_REJECTION) + compaction: Final = (_block(_summary(marker)),) if f"compact-now-{marker}".encode() in raw else () + thinking: Final = (_thinking(marker),) if _think(marker).encode() in raw else () + content: Final = (*compaction, *thinking, {"type": "text", "text": _answer(marker)}) + identity: Final = f"msg_{uuid.uuid4().hex}" + if body.get("stream") is True: + return Reply(content_type="text/event-stream", chunks=_anthropic_stream(identity, content)) + return Reply(body=cc.message_reply(identity, _CLAUDE, content, _anthropic_usage())) + + +def _gemini_reply(raw: bytes) -> Reply: + frame: Final = { + "responseId": f"gemini-{uuid.uuid4().hex}", + "candidates": [ + { + "content": {"role": "model", "parts": [{"text": _answer(_newest_marker(raw))}]}, + "index": 0, + "finishReason": "STOP", + } + ], + "usageMetadata": {"promptTokenCount": 20, "candidatesTokenCount": 5, "totalTokenCount": 25}, + "modelVersion": _GEMINI, + } + return Reply(body=json.dumps(frame).encode()) + + +_OPENAI_USAGE: Final = { + "input_tokens": 20, + "input_tokens_details": {"cached_tokens": 0}, + "output_tokens": 5, + "output_tokens_details": {"reasoning_tokens": 0}, + "total_tokens": 25, +} + + +def _openai_reply(raw: bytes) -> Reply: + tag: Final = uuid.uuid4().hex + message: Final = { + "id": f"msg_{tag}", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": _answer(_newest_marker(raw)), "annotations": []}], + } + response: Final = { + "id": f"resp_{tag}", + "object": "response", + "created_at": 1, + "status": "completed", + "model": _OPENAI, + "output": [message], + "usage": _OPENAI_USAGE, + } + return Reply(body=json.dumps(response).encode()) + + +def _openai_compact_reply() -> Reply: + compacted: Final = { + "id": f"resp_compact_{uuid.uuid4().hex}", + "created_at": 1, + "object": "response.compaction", + "output": [], + "usage": _OPENAI_USAGE, + } + return Reply(body=json.dumps(compacted).encode()) + + +def _vendor(request: Request) -> Reply: + if request.method == "GET": + return Reply(body=json.dumps({"object": "list", "data": [{"id": _OPENAI, "object": "model"}]}).encode()) + path: Final = urlsplit(request.target).path + if path.endswith("/v1/messages/count_tokens"): + return Reply(body=json.dumps({"input_tokens": _word_count(_JSON_OBJECT.validate_json(request.body))}).encode()) + if path.endswith("/v1/messages"): + return _claude_reply(_JSON_OBJECT.validate_json(request.body), request.body) + if path.endswith(":generateContent"): + return _gemini_reply(request.body) + if path.endswith("/responses/compact"): + return _openai_compact_reply() + assert path.endswith("/responses"), request.target + return _openai_reply(request.body) + + +def _provider_calls(wire: Wire, marker: str) -> tuple[Request, ...]: + needle: Final = f"ask-{marker}".encode() + return tuple(request for request in wire.drain() if request.method == "POST" and needle in request.body) + + +def _only_call(wire: Wire, marker: str) -> tuple[Request, dict[str, JsonValue]]: + calls: Final = _provider_calls(wire, marker) + assert len(calls) == 1, [call.target for call in calls] + return calls[0], _JSON_OBJECT.validate_json(calls[0].body) + + +def _objects(value: JsonValue) -> tuple[dict[str, JsonValue], ...]: + assert isinstance(value, list), value + return tuple(object_value(item) for item in value) + + +def _message_texts(message: Mapping[str, JsonValue], key: str = "content") -> tuple[str, ...]: + content: Final = message.get(key) + if isinstance(content, str): + return (content,) + if not isinstance(content, list): + return () + return tuple(string_value(part["text"]) for part in _objects(content) if isinstance(part.get("text"), str)) + + +def _anthropic_texts(body: Mapping[str, JsonValue], role: str) -> tuple[str, ...]: + messages: Final = tuple(message for message in _objects(body["messages"]) if message["role"] == role) + return tuple(chain.from_iterable(map(_message_texts, messages))) + + +def _word_count(body: Mapping[str, JsonValue]) -> int: + texts: Final = (*_anthropic_texts(body, "user"), *_anthropic_texts(body, "assistant")) + return sum(len(text.split()) for text in texts) + + +def _anthropic_block_types(body: Mapping[str, JsonValue]) -> tuple[JsonValue, ...]: + contents: Final = tuple(message["content"] for message in _objects(body["messages"])) + blocks: Final = chain.from_iterable(_objects(content) for content in contents if isinstance(content, list)) + return tuple(block.get("type") for block in blocks) + + +def _gemini_texts(body: Mapping[str, JsonValue], role: str) -> tuple[str, ...]: + turns: Final = tuple(turn for turn in _objects(body["contents"]) if turn["role"] == role) + return tuple(chain.from_iterable(_message_texts(turn, "parts") for turn in turns)) + + +def _openai_user_texts(body: Mapping[str, JsonValue]) -> tuple[str, ...]: + items: Final = tuple(item for item in _objects(body["input"]) if item.get("role") == "user") + return tuple(chain.from_iterable(map(_message_texts, items))) + + +def _openai_item_types(body: Mapping[str, JsonValue]) -> tuple[JsonValue, ...]: + return tuple(item.get("type") for item in _objects(body["input"])) + + +def _register(scenario: Scenario, wire: Wire, provider: Literal["claude", "gemini", "openai"]) -> str: + match provider: + case "claude": + model = scenario.model(model=f"anthropic/{_CLAUDE}", api_base=wire.url, api_key=cc.ANTHROPIC_API_KEY) + case "gemini": + model = scenario.model(model=f"gemini/{_GEMINI}", api_base=wire.url, api_key="synthetic-gemini-key") + case "openai": + model = scenario.model(model=f"openai/{_OPENAI}", api_base=f"{wire.url}/v1", api_key="synthetic-openai-key") + wait_for_worker_readiness(scenario.gateway, model, 1, _PROXY_WORKERS) + return model + + +def _openai_client(gateway: Gateway) -> openai.OpenAI: + return openai.OpenAI( + base_url=f"{gateway.client.base_url}/v1", + api_key=gateway.key, + max_retries=0, + http_client=httpx.Client(trust_env=False, timeout=60), + ) + + +def _decoded(encrypted_content: str) -> dict[str, JsonValue]: + return _JSON_OBJECT.validate_json(encrypted_content) + + +def _responses( + gateway: Gateway, body: Mapping[str, JsonValue], headers: Mapping[str, str] | None = None +) -> dict[str, JsonValue]: + response: Final = gateway.request("POST", "/v1/responses", {**body, **_CACHE_BUST}, headers=headers) + assert response.status_code == 200, response.text + return _JSON_OBJECT.validate_json(response.content) + + +def _sse_events(text: str) -> tuple[dict[str, JsonValue], ...]: + payloads: Final = tuple( + line.removeprefix("data:").strip() for line in text.splitlines() if line.startswith("data:") + ) + return tuple(_JSON_OBJECT.validate_json(payload) for payload in payloads if payload != "[DONE]") + + +def _stream( + gateway: Gateway, path: str, body: Mapping[str, JsonValue], headers: Mapping[str, str] | None = None +) -> tuple[dict[str, JsonValue], ...]: + sent_headers: Final = {"Authorization": f"Bearer {gateway.key}", **(headers or {})} + with gateway.client.stream( + "POST", path, json={**body, **_CACHE_BUST, "stream": True}, headers=sent_headers + ) as response: + text: Final = response.read().decode() + assert response.status_code == 200, text + return _sse_events(text) + + +def _delta_provider_fields(chunk: Mapping[str, JsonValue]) -> tuple[dict[str, JsonValue], ...]: + return tuple( + fields + for choice in _objects(chunk.get("choices", [])) + if isinstance(fields := object_value(choice.get("delta", {})).get("provider_specific_fields"), dict) + ) + + +def _streamed_compaction_blocks(chunks: Sequence[Mapping[str, JsonValue]]) -> tuple[JsonValue, ...]: + fields: Final = chain.from_iterable(_delta_provider_fields(chunk) for chunk in chunks) + return tuple(field["compaction_blocks"] for field in fields if "compaction_blocks" in field) + + +def _output_items(response: Mapping[str, JsonValue], item_type: str) -> tuple[dict[str, JsonValue], ...]: + return tuple(item for item in _objects(response["output"]) if item["type"] == item_type) + + +def _listed(value: JsonValue) -> list[JsonValue]: + assert isinstance(value, list), value + return value + + +def _only(items: Sequence[dict[str, JsonValue]]) -> dict[str, JsonValue]: + assert len(items) == 1, items + return items[0] + + +@pytest.mark.parametrize("stream", (False, True), ids=("json", "stream")) +def test_claude_compaction_block_is_returned_as_a_compaction_output_item(gateway: Gateway, stream: bool) -> None: + marker: Final = _new_marker() + with wire_server(_vendor) as wire, gateway.scenario() as scenario: + model: Final = _register(scenario, wire, "claude") + client: Final = _openai_client(gateway) + request: Final = { + "input": _ask(marker, compact=True), + "extra_body": {"context_management": _COMPACTING, **_CACHE_BUST}, + } + if not stream: + reply: Final = client.responses.create(model=model, **request) + assert [item.type for item in reply.output] == ["compaction", "message"] + compaction: Final = reply.output[0] + assert compaction.type == "compaction" + assert _decoded(compaction.encrypted_content) == _block(_summary(marker)) + assert reply.output_text == _answer(marker) + else: + events: Final = tuple(client.responses.create(model=model, stream=True, **request)) + added: Final = [e for e in events if e.type == "response.output_item.added" and e.item.type == "compaction"] + done: Final = [e for e in events if e.type == "response.output_item.done" and e.item.type == "compaction"] + completed: Final = [e for e in events if e.type == "response.completed"] + assert len(added) == 1 and len(done) == 1 and len(completed) == 1, [e.type for e in events] + assert added[0].item.encrypted_content == "" + assert done[0].output_index == added[0].output_index + assert done[0].item.id == added[0].item.id + assert _decoded(done[0].item.encrypted_content) == _block(_summary(marker)) + finished: Final = [item for item in completed[0].response.output if item.type == "compaction"] + assert [(item.id, item.encrypted_content) for item in finished] == [ + (done[0].item.id, done[0].item.encrypted_content) + ] + message_indexes: Final = { + e.output_index for e in events if e.type == "response.output_item.added" and e.item.type == "message" + } + assert message_indexes and added[0].output_index not in message_indexes + assert completed[0].response.output_text == _answer(marker) + call, body = _only_call(wire, marker) + assert _COMPACT_BETA in call.headers.get("anthropic-beta", ""), dict(call.headers) + + +def _history(marker: str, compaction: Mapping[str, JsonValue]) -> list[JsonValue]: + return [ + {"role": "developer", "content": f"rules-{marker}"}, + {"role": "user", "content": f"old-{marker}"}, + { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": f"old-answer-{marker}", "annotations": []}], + }, + dict(compaction), + { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": f"recap-{marker}", "annotations": []}], + }, + {"role": "user", "content": _ask(marker)}, + ] + + +_ROUTES: Final[tuple[Route, ...]] = ( + "claude-compacting", + "claude-plain", + "gemini", + "openai", + "openai-compact", + "openai-opaque", +) + + +def _route_provider(route: Route) -> Literal["claude", "gemini", "openai"]: + if route.startswith("claude"): + return "claude" + return "gemini" if route == "gemini" else "openai" + + +def _assert_reached( + route: Route, marker: str, sent: Mapping[str, JsonValue], call: Request, body: Mapping[str, JsonValue] +) -> None: + summary: Final = _summary(marker) + dumped: Final = json.dumps(body) + if route == "openai-opaque": + assert urlsplit(call.target).path.endswith("/responses"), call.target + assert sent in _objects(body["input"]), body["input"] + assert f"old-{marker}" in dumped + return + if route != "claude-compacting": + assert f"old-{marker}" not in dumped and f"old-answer-{marker}" not in dumped, dumped[:2000] + assert f"rules-{marker}" in dumped + assert f"recap-{marker}" in dumped + match route: + case "claude-compacting": + assistant: Final = tuple( + message for message in _objects(body["messages"]) if message["role"] == "assistant" + ) + blocks: Final = tuple( + chain.from_iterable( + _objects(message["content"]) for message in assistant if isinstance(message["content"], list) + ) + ) + assert _block(summary) in blocks, body["messages"] + assert summary not in _anthropic_texts(body, "user") + assert f"rules-{marker}" in json.dumps(body["system"]) + assert _COMPACT_BETA in call.headers.get("anthropic-beta", ""), dict(call.headers) + case "claude-plain": + assert "compaction" not in _anthropic_block_types(body), body["messages"] + assert summary in _anthropic_texts(body, "user"), body["messages"] + case "gemini": + assert summary in _gemini_texts(body, "user"), body["contents"] + case "openai" | "openai-compact": + suffix: Final = "/responses/compact" if route == "openai-compact" else "/responses" + assert urlsplit(call.target).path.endswith(suffix), call.target + assert "compaction" not in _openai_item_types(body), body["input"] + assert summary in _openai_user_texts(body), body["input"] + + +@pytest.mark.parametrize("route", _ROUTES) +def test_compaction_item_reaches_each_route_in_the_shape_it_reads(gateway: Gateway, route: Route) -> None: + marker: Final = _new_marker() + compaction: Final = ( + {"type": "compaction", "id": f"cmp_{marker}", "encrypted_content": f"gAAAAA-opaque-{marker}"} + if route == "openai-opaque" + else _compaction_item(_summary(marker)) + ) + compacting: Final[dict[str, JsonValue]] = ( + {"context_management": _COMPACTING} if route == "claude-compacting" else {} + ) + dropping: Final[dict[str, JsonValue]] = ( + {"context_management": _COMPACTING, "drop_params": True} if route == "gemini" else {} + ) + path: Final = "/v1/responses/compact" if route == "openai-compact" else "/v1/responses" + with wire_server(_vendor) as wire, gateway.scenario() as scenario: + model: Final = _register(scenario, wire, _route_provider(route)) + response: Final = gateway.request( + "POST", + path, + {"model": model, "input": _history(marker, compaction), **compacting, **dropping, **_CACHE_BUST}, + ) + assert response.status_code == 200, response.text + if route != "openai-compact": + assert _answer(marker) in response.text, response.text + call, body = _only_call(wire, marker) + _assert_reached(route, marker, compaction, call, body) + + +def _input_tokens(gateway: Gateway, model: str, items: list[JsonValue]) -> int: + counted: Final = gateway.post("/v1/responses/input_tokens", {"model": model, "input": items}) + tokens: Final = counted["input_tokens"] + assert isinstance(tokens, int), counted + return tokens + + +def test_input_tokens_count_from_the_newest_compaction(gateway: Gateway) -> None: + marker: Final = _new_marker() + summary: Final = _summary(marker) + old_history: Final = " ".join(f"history-{index}" for index in range(600)) + compacted: Final[list[JsonValue]] = [ + {"role": "user", "content": old_history}, + _compaction_item(summary), + {"role": "user", "content": _ask(marker)}, + ] + equivalent: Final[list[JsonValue]] = [ + {"role": "user", "content": summary}, + {"role": "user", "content": _ask(marker)}, + ] + with wire_server(_vendor) as wire, gateway.scenario() as scenario: + model: Final = _register(scenario, wire, "claude") + assert _input_tokens(gateway, model, compacted) == _input_tokens(gateway, model, equivalent) + counted: Final = tuple( + request.body for request in wire.drain() if urlsplit(request.target).path.endswith("/count_tokens") + ) + assert len(counted) == 2, counted + assert all(b"history-" not in body for body in counted), counted[0][:2000] + + +def test_streamed_compaction_item_is_done_before_the_reasoning_after_it(gateway: Gateway) -> None: + marker: Final = _new_marker() + with wire_server(_vendor) as wire, gateway.scenario() as scenario: + model: Final = _register(scenario, wire, "claude") + events: Final = _stream( + gateway, + "/v1/responses", + { + "model": model, + "input": f"{_ask(marker, compact=True)} {_think(marker)}", + "context_management": _COMPACTING, + }, + ) + done: Final = tuple( + (position, event["output_index"], object_value(event["item"])["type"]) + for position, event in enumerate(events) + if event["type"] == "response.output_item.done" + ) + completed: Final = _only([event for event in events if event["type"] == "response.completed"]) + assert [(index, kind) for _, index, kind in done] == [(0, "compaction"), (1, "reasoning"), (2, "message")], [ + event["type"] for event in events + ] + assert [item["type"] for item in _objects(object_value(completed["response"])["output"])] == [ + "compaction", + "reasoning", + "message", + ] + reasoning_done: Final = done[1][0] + sequence_numbers: Final = [ + number for event in events[: reasoning_done + 1] if isinstance(number := event.get("sequence_number"), int) + ] + assert sequence_numbers == sorted(set(sequence_numbers)), sequence_numbers + _only_call(wire, marker) + + +def test_chat_stream_reports_the_whole_compaction_block(gateway: Gateway) -> None: + marker: Final = _new_marker() + with wire_server(_vendor) as wire, gateway.scenario() as scenario: + model: Final = _register(scenario, wire, "claude") + chunks: Final = _stream( + gateway, + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": _ask(marker, compact=True)}], + "context_management": {"edits": [{"type": "compact_20260112"}]}, + }, + ) + reported: Final = _streamed_compaction_blocks(chunks) + assert reported, chunks + assert _objects(reported[-1])[-1] == _block(_summary(marker)), reported[-1] + assert len(_provider_calls(wire, marker)) == 1 + + +def _guardrail_peer(request: Request) -> Reply: + payload: Final = _JSON_OBJECT.validate_json(request.body) + texts: Final = payload.get("texts") or [] + assert isinstance(texts, list), payload + if urlsplit(request.target).path.startswith("/block/"): + found: Final = any(isinstance(text, str) and _CITY in text for text in texts) + verdict: Final = {"action": "BLOCKED", "blocked_reason": "city in summary"} if found else {"action": "NONE"} + return Reply(body=json.dumps(verdict).encode()) + masked: Final = [text.replace(_CITY, _MASKED_CITY) if isinstance(text, str) else text for text in texts] + return Reply(body=json.dumps({"action": "GUARDRAIL_INTERVENED", "texts": masked}).encode()) + + +def _guardrail(name: str, mode: str, api_base: str, **fields: JsonValue) -> dict[str, JsonValue]: + return { + "guardrail_name": name, + "litellm_params": { + "guardrail": "generic_guardrail_api", + "mode": mode, + "default_on": False, + "api_base": api_base, + "api_key": "synthetic-guardrail-key", + **fields, + }, + } + + +def _deployment(name: str, model: str, api_base: str, **fields: JsonValue) -> dict[str, JsonValue]: + return {"model_name": name, "litellm_params": {"model": model, "api_base": api_base, **fields}} + + +_STORED_CLAUDE: Final = "compaction-stored-claude" +_SMALL_CLAUDE: Final = "compaction-small-claude" +_STORED_GEMINI: Final = "compaction-stored-gemini" +_STORED_OPENAI: Final = "compaction-stored-openai" +_SMALL_CONTEXT: Final = 2000 + + +def _stored_config(directory: Path, provider_url: str, guardrail_url: str, **litellm_settings: JsonValue) -> Path: + stock: Final = _JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + claude: Final = {"api_key": cc.ANTHROPIC_API_KEY} + config: Final[dict[str, JsonValue]] = { + **stock, + "model_list": [ + _deployment(_STORED_CLAUDE, f"anthropic/{_CLAUDE}", provider_url, **claude), + { + **_deployment(_SMALL_CLAUDE, f"anthropic/{_CLAUDE}", provider_url, **claude), + "model_info": {"max_input_tokens": _SMALL_CONTEXT}, + }, + _deployment(_STORED_GEMINI, f"gemini/{_GEMINI}", provider_url, api_key="synthetic-gemini-key"), + _deployment(_STORED_OPENAI, f"openai/{_OPENAI}", f"{provider_url}/v1", api_key="synthetic-openai-key"), + ], + "general_settings": { + **object_value(stock["general_settings"]), + "store_prompts_in_spend_logs": True, + "disable_responses_id_security": True, + }, + "litellm_settings": {**object_value(stock["litellm_settings"]), **litellm_settings}, + "router_settings": {**object_value(stock["router_settings"]), "num_retries": 0, "enable_pre_call_checks": True}, + "guardrails": [ + _guardrail(_MASK_INPUT, "pre_call", f"{guardrail_url}/input"), + _guardrail(_MASK_OUTPUT, "post_call", f"{guardrail_url}/output"), + _guardrail(_BLOCK_OUTPUT, "post_call", f"{guardrail_url}/block", streaming_buffer_until_moderated=True), + ], + } + path: Final = directory / "compaction-stored.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +@dataclass(frozen=True, slots=True) +class _StoredRig: + gateway: Gateway + provider: Wire + guardrail: Wire + + +def _stored_rig(directory: Path, **litellm_settings: JsonValue) -> Iterator[_StoredRig]: + with ( + wire_server(_vendor) as provider, + wire_server(_guardrail_peer) as guardrail, + gateway_from_environment() as gateway, + ): + config: Final = _stored_config(directory, provider.url, guardrail.url, **litellm_settings) + with owned_proxy(gateway, directory, {}, config=config, workers=2) as owned: + yield _StoredRig(owned, provider, guardrail) + + +@pytest.fixture(scope="module") +def stored(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_StoredRig]: + yield from _stored_rig(tmp_path_factory.mktemp("compaction-stored")) + + +@pytest.fixture(scope="module") +def redacting(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_StoredRig]: + yield from _stored_rig(tmp_path_factory.mktemp("compaction-redacting"), turn_off_message_logging=True) + + +def _logged_row(request_id: str, status: Literal["success", "failure"]) -> dict[str, JsonValue]: + rows: Final = eventually( + lambda: read_rows( + 'SELECT request_id, proxy_server_request, response, metadata, status FROM "LiteLLM_SpendLogs"' + " WHERE request_id=%s", + (request_id,), + ), + lambda values: len(values) == 1, + seconds=70, + ) + assert rows[0]["status"] == status, rows[0] + return rows[0] + + +def _spend_row(response_id: str) -> dict[str, JsonValue]: + decoded: Final = ResponsesAPIRequestUtils.decode_responses_api_response_id(response_id).get("response_id") + return _logged_row(decoded if isinstance(decoded, str) else response_id, "success") + + +_REPLAY_ROUTES: Final = ( + pytest.param(_STORED_CLAUDE, True, id="claude-compacting"), + pytest.param(_STORED_CLAUDE, False, id="claude-plain"), + pytest.param(_STORED_GEMINI, True, id="gemini"), +) + + +@pytest.mark.parametrize(("follow_up_model", "compacting"), _REPLAY_ROUTES) +def test_previous_response_id_after_a_compaction_turn_replays_the_whole_summary( + stored: _StoredRig, follow_up_model: str, compacting: bool +) -> None: + first: Final = _new_marker() + second: Final = _new_marker() + summary: Final = _summary(first) + long_text: Final = f"long-{first} {'z' * 3000} tail-{first}" + compaction_request: Final[dict[str, JsonValue]] = {"context_management": _COMPACTING} + opened: Final = _responses( + stored.gateway, + { + "model": _STORED_CLAUDE, + "input": [{"role": "user", "content": long_text}, {"role": "user", "content": _ask(first, compact=True)}], + **compaction_request, + }, + ) + assert _decoded(string_value(_only(_output_items(opened, "compaction"))["encrypted_content"])) == _block(summary) + row: Final = _spend_row(string_value(opened["id"])) + assert summary in json.dumps(row["response"]), json.dumps(row["response"])[:600] + stored_request: Final = json.dumps(row["proxy_server_request"]) + assert LITELLM_TRUNCATED_PAYLOAD_FIELD in stored_request and "z" * 3000 not in stored_request + _only_call(stored.provider, first) + + follow_up: Final = _responses( + stored.gateway, + { + "model": follow_up_model, + "input": _ask(second), + "previous_response_id": opened["id"], + **(compaction_request if compacting else {}), + **({"drop_params": True} if follow_up_model == _STORED_GEMINI else {}), + }, + ) + assert _answer(second) in json.dumps(follow_up) + call, body = _only_call(stored.provider, second) + dumped: Final = json.dumps(body) + assert f"long-{first}" not in dumped and f"ask-{first}" not in dumped, dumped[:2000] + assert _answer(first) in dumped + if follow_up_model == _STORED_GEMINI: + assert summary in _gemini_texts(body, "user"), body["contents"] + elif compacting: + replayed: Final = _objects(body["messages"])[0] + assert replayed["role"] == "assistant", body["messages"] + assert _objects(replayed["content"])[0] == _block(summary), replayed + assert _COMPACT_BETA in call.headers.get("anthropic-beta", ""), dict(call.headers) + else: + assert "compaction" not in _anthropic_block_types(body), body["messages"] + assert summary in _anthropic_texts(body, "user"), body["messages"] + + +def test_router_pre_call_check_counts_from_the_newest_compaction(stored: _StoredRig) -> None: + compacted_marker: Final = _new_marker() + control_marker: Final = _new_marker() + old_history: Final = " ".join(f"history-{index}" for index in range(3 * _SMALL_CONTEXT)) + + def send(marker: str, items: list[JsonValue]) -> httpx.Response: + return stored.gateway.request( + "POST", + "/v1/responses", + {"model": _SMALL_CLAUDE, "input": [*items, {"role": "user", "content": _ask(marker)}], **_CACHE_BUST}, + ) + + compacted: Final = send( + compacted_marker, + [{"role": "user", "content": old_history}, _compaction_item(f"short summary {compacted_marker}")], + ) + assert compacted.status_code == 200, compacted.text + _only_call(stored.provider, compacted_marker) + control: Final = send(control_marker, [{"role": "user", "content": old_history}]) + assert control.status_code == 400, control.text + assert _provider_calls(stored.provider, control_marker) == () + + +_GUARDED_ROUTES: Final = ( + pytest.param(_STORED_OPENAI, id="openai"), + pytest.param(_STORED_CLAUDE, id="claude-compacting"), +) + + +@pytest.mark.parametrize("model", _GUARDED_ROUTES) +def test_pre_call_guardrail_rewrites_the_compaction_summary(stored: _StoredRig, model: str) -> None: + marker: Final = _new_marker() + summary: Final = _summary(marker) + masked: Final = summary.replace(_CITY, _MASKED_CITY) + compacting: Final[dict[str, JsonValue]] = {"context_management": _COMPACTING} if model == _STORED_CLAUDE else {} + _responses( + stored.gateway, + { + "model": model, + "input": [_compaction_item(summary), {"role": "user", "content": _ask(marker)}], + "guardrails": [_MASK_INPUT], + **compacting, + }, + ) + scanned: Final = tuple( + _JSON_OBJECT.validate_json(request.body) + for request in stored.guardrail.drain() + if f"kept-{marker}".encode() in request.body + ) + assert any(summary in _listed(payload["texts"]) for payload in scanned), scanned + call, body = _only_call(stored.provider, marker) + assert _CITY not in json.dumps(body), json.dumps(body)[:2000] + if model == _STORED_OPENAI: + assert masked in _openai_user_texts(body), body["input"] + else: + assert _objects(_objects(body["messages"])[0]["content"])[0] == _block(masked), body["messages"] + + +def test_post_call_guardrail_rewrites_the_returned_compaction_summary(stored: _StoredRig) -> None: + marker: Final = _new_marker() + returned: Final = _responses( + stored.gateway, + { + "model": _STORED_CLAUDE, + "input": _ask(marker, compact=True), + "context_management": _COMPACTING, + "guardrails": [_MASK_OUTPUT], + }, + ) + masked: Final = _block(_summary(marker).replace(_CITY, _MASKED_CITY)) + assert _decoded(string_value(_only(_output_items(returned, "compaction"))["encrypted_content"])) == masked + _only_call(stored.provider, marker) + + +def test_buffered_stream_guardrail_blocks_on_the_compaction_summary(stored: _StoredRig) -> None: + marker: Final = _new_marker() + body: Final[dict[str, JsonValue]] = { + "model": _STORED_CLAUDE, + "input": _ask(marker, compact=True), + "context_management": _COMPACTING, + "guardrails": [_BLOCK_OUTPUT], + "stream": True, + **_CACHE_BUST, + } + headers: Final = {"Authorization": f"Bearer {stored.gateway.key}"} + with stored.gateway.client.stream("POST", "/v1/responses", json=body, headers=headers) as response: + text: Final = response.read().decode() + assert response.status_code == 200, text + peer_calls: Final = tuple( + (urlsplit(request.target).path, _JSON_OBJECT.validate_json(request.body).get("texts")) + for request in stored.guardrail.drain() + if marker.encode() in request.body + ) + assert any(path.startswith("/block/") and _summary(marker) in _listed(texts) for path, texts in peer_calls), [ + (path, [str(value)[:80] for value in _listed(texts or [])]) for path, texts in peer_calls + ] + event_types: Final = tuple(event["type"] for event in _sse_events(text)) + assert "response.failed" in event_types, event_types + assert "response.completed" not in event_types, event_types + assert f"kept-{marker}" not in text and _CITY not in text, text[:2000] + _only_call(stored.provider, marker) + + +def test_redacted_spend_row_holds_no_compaction_summary(stored: _StoredRig) -> None: + marker: Final = _new_marker() + returned: Final = _responses( + stored.gateway, + {"model": _STORED_CLAUDE, "input": _ask(marker, compact=True), "context_management": _COMPACTING}, + headers=_REDACT_HEADER, + ) + assert _decoded(string_value(_only(_output_items(returned, "compaction"))["encrypted_content"])) == _block( + _summary(marker) + ) + row: Final = _spend_row(string_value(returned["id"])) + assert f"kept-{marker}" not in json.dumps(row["response"]), json.dumps(row["response"])[:2000] + _only_call(stored.provider, marker) + + +def test_stored_compaction_input_item_keeps_its_whole_summary_and_nothing_else(stored: _StoredRig) -> None: + first: Final = _new_marker() + second: Final = _new_marker() + summary: Final = _summary(first) + padded_item: Final = { + **_compaction_item(summary), + "encrypted_content": json.dumps({**_block(summary), "padding": f"pad-{first} {'p' * 3000}"}), + } + opened: Final = _responses( + stored.gateway, + { + "model": _STORED_CLAUDE, + "input": [padded_item, {"role": "user", "content": _ask(first)}], + "context_management": _COMPACTING, + }, + ) + stored_input: Final = _objects( + object_value(_spend_row(string_value(opened["id"]))["proxy_server_request"])["input"] + ) + assert _decoded(string_value(stored_input[0]["encrypted_content"])) == _block(summary), stored_input[0] + _only_call(stored.provider, first) + + _responses( + stored.gateway, + { + "model": _STORED_CLAUDE, + "input": _ask(second), + "previous_response_id": opened["id"], + "context_management": _COMPACTING, + }, + ) + _, body = _only_call(stored.provider, second) + assert summary in _anthropic_texts(body, "user"), body["messages"] + assert f"pad-{first}" not in json.dumps(body), body["messages"] + + +def test_redacted_chat_stream_spend_row_holds_no_compaction_delta(stored: _StoredRig) -> None: + marker: Final = _new_marker() + summary: Final = _summary(marker) + chunks: Final = _stream( + stored.gateway, + "/v1/chat/completions", + { + "model": _STORED_CLAUDE, + "messages": [{"role": "user", "content": _ask(marker, compact=True)}], + "context_management": {"edits": [{"type": "compact_20260112"}]}, + }, + headers=_REDACT_HEADER, + ) + streamed_blocks: Final = _streamed_compaction_blocks(chunks) + assert streamed_blocks and _objects(streamed_blocks[-1])[-1] == _block(summary), chunks + secrets: Final = (f"kept-{marker}", summary[len(summary) // 2 :]) + logged: Final = {column: json.dumps(value) for column, value in _spend_row(string_value(chunks[0]["id"])).items()} + leaked: Final = {column: text for column, text in logged.items() if any(secret in text for secret in secrets)} + assert not leaked, {column: text[max(0, text.find(secrets[0]) - 600) :][:1500] for column, text in leaked.items()} + _only_call(stored.provider, marker) + + +def _failing_request(path: str, marker: str) -> dict[str, JsonValue]: + ask: Final = f"{_ask(marker)} {_fail(marker)}" + if path == "/v1/responses": + return { + "model": _STORED_CLAUDE, + "input": [_compaction_item(_summary(marker)), {"role": "user", "content": ask}], + } + return {"model": _STORED_CLAUDE, "max_tokens": 64, "messages": [{"role": "user", "content": ask}]} + + +def _failed_call_row( + rig: _StoredRig, + path: str, + marker: str, + headers: Mapping[str, str], + key: str | None = None, + **fields: JsonValue, +) -> dict[str, JsonValue]: + body: Final = {**_failing_request(path, marker), **fields} + response: Final = rig.gateway.request("POST", path, body, headers=headers, key=key) + assert response.status_code == 400, response.text + _only_call(rig.provider, marker) + return _logged_row(response.headers["x-litellm-call-id"], "failure") + + +_LEGACY_REDACT_HEADER: Final = {"litellm-enable-message-redaction": "true"} +_REDACTED_FAILURES: Final = ( + pytest.param("/v1/responses", _REDACT_HEADER, id="responses"), + pytest.param("/v1/responses", _LEGACY_REDACT_HEADER, id="responses-legacy-header"), + pytest.param("/v1/messages", _REDACT_HEADER, id="messages"), + pytest.param("/v1/chat/completions", _REDACT_HEADER, id="chat"), +) + + +@pytest.mark.parametrize(("path", "headers"), _REDACTED_FAILURES) +def test_failed_call_spend_row_holds_no_request_text_with_the_redaction_header( + stored: _StoredRig, path: str, headers: Mapping[str, str] +) -> None: + marker: Final = _new_marker() + logged: Final = json.dumps(_failed_call_row(stored, path, marker, headers)) + assert marker not in logged, logged[max(0, logged.find(marker) - 600) :][:1500] + + +def test_failed_responses_call_spend_row_keeps_the_whole_compaction_summary_without_the_header( + stored: _StoredRig, +) -> None: + marker: Final = _new_marker() + row: Final = _failed_call_row(stored, "/v1/responses", marker, {}) + stored_input: Final = _objects(object_value(row["proxy_server_request"])["input"]) + assert _decoded(string_value(stored_input[0]["encrypted_content"])) == _block(_summary(marker)), stored_input[0] + assert f"{_ask(marker)} {_fail(marker)}" in _message_texts(stored_input[1]), stored_input[1] + + +def test_failed_responses_call_stores_the_caller_metadata_without_the_request_headers(stored: _StoredRig) -> None: + marker: Final = _new_marker() + row: Final = _failed_call_row(stored, "/v1/responses", marker, _REDACT_HEADER, metadata={"caller": marker}) + stored_metadata: Final = object_value(object_value(row["proxy_server_request"])["metadata"]) + assert stored_metadata.get("caller") == marker, stored_metadata + assert "headers" not in stored_metadata, stored_metadata + + +def test_failed_responses_call_keeps_the_caller_metadata_headers_over_the_redaction_header(stored: _StoredRig) -> None: + marker: Final = _new_marker() + caller_headers: Final = {"caller": marker} + row: Final = _failed_call_row(stored, "/v1/responses", marker, _REDACT_HEADER, metadata={"headers": caller_headers}) + stored_request: Final = object_value(row["proxy_server_request"]) + assert object_value(stored_request["metadata"])["headers"] == caller_headers, stored_request["metadata"] + assert _ask(marker) in json.dumps(stored_request), json.dumps(stored_request)[:1500] + + +_GLOBALLY_REDACTED_FAILURES: Final = ( + pytest.param("/v1/responses", _UNREDACT_HEADER, False, id="responses-disable-header"), + pytest.param("/v1/chat/completions", _UNREDACT_HEADER, True, id="chat-disable-header"), + pytest.param("/v1/responses", {}, False, id="responses-no-header"), +) + + +@pytest.mark.parametrize(("path", "headers", "kept"), _GLOBALLY_REDACTED_FAILURES) +def test_opted_out_key_failed_call_spend_row_takes_the_disable_header_over_global_redaction_only_on_chat( + redacting: _StoredRig, path: str, headers: Mapping[str, str], kept: bool +) -> None: + marker: Final = _new_marker() + with redacting.gateway.scenario() as scenario: + key: Final = scenario.key(metadata={"allow_client_message_redaction_opt_out": True}) + row: Final = _failed_call_row(redacting, path, marker, headers, key) + stored_request: Final = json.dumps(row["proxy_server_request"]) + assert (_ask(marker) in stored_request) is kept, stored_request[:1500] diff --git a/tests/unit/litellm_core_utils/test_redact_messages.py b/tests/unit/litellm_core_utils/test_redact_messages.py index a91888fea79..6a734ca3b70 100644 --- a/tests/unit/litellm_core_utils/test_redact_messages.py +++ b/tests/unit/litellm_core_utils/test_redact_messages.py @@ -57,6 +57,8 @@ def _provider_fields() -> dict[str, object]: "reasoningContentBlocks": [{"reasoningText": {"text": SECRET}}], "citations": [{"cited_text": SECRET}], "thought_signatures": [SIGNATURE], + "compaction_blocks": [{"type": "compaction", "content": SECRET}], + "compaction_delta": {"type": "compaction_delta", "content": SECRET}, } @@ -165,6 +167,26 @@ def test_streaming_redaction_redacts_delta_provider_fields_without_touching_the_ assert fields == _provider_fields() +@pytest.mark.parametrize("shape", ["object", "dict"]) +def test_redaction_hides_a_compaction_summary_in_responses_api_output(shape: str) -> None: + response: Final = ResponsesAPIResponse( + id="resp_1", + created_at=1, + model="claude-sonnet-5-5", + object="response", + status="completed", + output=[ + {"type": "compaction", "id": "cmp_1", "encrypted_content": json.dumps({"type": "compaction", "content": SECRET})} + ], + ) + result: Final = response if shape == "object" else response.model_dump() + + redacted: Final = perform_redaction({}, result) + + assert SECRET not in str(redacted), redacted + assert SECRET in str(result) + + @pytest.mark.parametrize("surface", ("typed", "dict", "standard", "callback", "helper")) def test_responses_redaction_removes_instructions_without_changing_the_response(surface: str) -> None: response: Final = litellm.ResponsesAPIResponse.model_validate( diff --git a/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py index bf871210074..2f10fd6fd0d 100644 --- a/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -34,6 +34,7 @@ from litellm.llms.openai.responses.guardrail_translation.handler import ( from litellm.llms.openai.responses.guardrail_translation.tool_merge import merge_guardrailed_tools from litellm.proxy.guardrails.guardrail_hooks.generic_guardrail_api import GenericGuardrailAPI from litellm.types.llms.openai import ChatCompletionToolCallChunk +from litellm.responses.litellm_completion_transformation.compaction_items import compaction_summary from litellm.responses.litellm_completion_transformation.transformation import ( LiteLLMCompletionResponsesConfig, ) @@ -298,6 +299,32 @@ class TestOpenAIResponsesHandlerInputProcessing: assert result["input"][1]["content"] == "World [GUARDRAILED]" assert result["model"] == "gpt-4" + @pytest.mark.asyncio + async def test_compaction_summary_is_scanned_and_rewritten_inside_its_item(self) -> None: + affinity_prefix: Final = "litellm_enc:bW9kZWxfaWQ6ZGVwLWE=;" + data: Final = { + "input": [ + { + "type": "compaction", + "id": "cmp_1", + "encrypted_content": affinity_prefix + + json.dumps({"type": "compaction", "content": "Cat is Biscuit."}), + }, + {"role": "user", "content": "Cat name?", "type": "message"}, + ], + "model": "claude-sonnet-5-5", + } + + result: Final = await OpenAIResponsesHandler().process_input_messages(data, MockGuardrail(guardrail_name="t")) + + rewritten: Final = result["input"][0]["encrypted_content"] + assert rewritten.startswith(affinity_prefix) + assert json.loads(rewritten.removeprefix(affinity_prefix)) == { + "type": "compaction", + "content": "Cat is Biscuit. [GUARDRAILED]", + } + assert result["input"][1]["content"] == "Cat name? [GUARDRAILED]" + @pytest.mark.asyncio async def test_process_input_list_with_multimodal_content(self): """Test processing list input with multimodal content""" @@ -3072,6 +3099,21 @@ class DictStructuredMessagesGuardrail(CustomGuardrail): return {**inputs, "structured_messages": {"evaluation": "allowed", "messages": []}} +class NameMaskingStructuredGuardrail(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + rewritten: Final = [ + {**m, "content": m["content"].replace("Biscuit", "[NAME]")} if isinstance(m.get("content"), str) else m + for m in inputs.get("structured_messages") or [] + ] + return {**inputs, "structured_messages": rewritten} + + def _parallel_tool_call_input() -> list: return [ {"id": "fc_1", "type": "function_call", "call_id": "call_1", "name": "read_a", "arguments": "{}"}, @@ -3298,6 +3340,30 @@ class TestProvenancePatching: assistant_items = [item for item in result["input"] if item.get("role") == "assistant"] assert [_texts(item) for item in assistant_items] == [[COMPRESSED_MARKER]] + @pytest.mark.asyncio + async def test_masked_compaction_summary_stays_a_compaction_item_that_drops_older_history(self) -> None: + data: Final = { + "model": "claude-sonnet-5-5", + "input": [ + {"role": "user", "content": "My cat is Biscuit. " * 200}, + {"role": "assistant", "content": "Noted."}, + { + "type": "compaction", + "id": "cmp_1", + "encrypted_content": json.dumps({"type": "compaction", "content": "Cat is Biscuit."}), + }, + {"role": "user", "content": "Cat name?"}, + ], + } + + result: Final = await OpenAIResponsesHandler().process_input_messages(data, NameMaskingStructuredGuardrail()) + + assert compaction_summary(result["input"][2]) == "Cat is [NAME]." + provider_messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( + input=result["input"], responses_api_request={}, only_since_compaction=True + ) + assert [message["content"] for message in provider_messages] == ["Cat is [NAME].", "Cat name?"] + def test_provenance_bails_on_non_mapping_item(self): from litellm.llms.openai.responses.guardrail_translation.handler import _input_item_provenance @@ -4174,3 +4240,117 @@ class TestOpenAIResponsesHandlerStreamingScanKey: assert [inputs.get("texts") for inputs in guardrail.seen_inputs] == [["hi"]], guardrail.seen_inputs tool_calls = guardrail.seen_inputs[0].get("tool_calls") or [] assert [call["function"]["arguments"] for call in tool_calls] == ['{"city": "Paris"}'], tool_calls + + @pytest.mark.asyncio + async def test_a_tool_call_finished_mid_stream_reaches_the_guardrail_without_texts(self): + guardrail: Final = MockRecordingGuardrail(guardrail_name="test") + await OpenAIResponsesHandler().process_output_streaming_response( + responses_so_far=[self._finished_function_call(1, "fc_1", "Paris")[-1]], + guardrail_to_apply=guardrail, + request_data={}, + ) + assert [ + (sorted(inputs), [call["function"]["arguments"] for call in inputs.get("tool_calls") or []]) + for inputs in guardrail.seen_inputs + ] == [(["tool_calls"], ['{"city": "Paris"}'])] + + +class TestCompactionOutputSummaryReachesOutputGuardrails: + SUMMARY: Final = "The user lives at 1 Secret Lane." + + @classmethod + def _compaction_item(cls, summary: str | None = None) -> dict[str, str]: + return { + "type": "compaction", + "id": "cmp_1", + "encrypted_content": json.dumps({"type": "compaction", "content": summary or cls.SUMMARY}), + } + + @staticmethod + def _summary(item: object) -> str: + encrypted_content: Final = item["encrypted_content"] if isinstance(item, dict) else item.encrypted_content + return json.loads(encrypted_content)["content"] + + @staticmethod + def _lane_masking_guardrail() -> CustomGuardrail: + class MaskLane(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: LiteLLMLoggingObj | None = None, + ) -> GenericGuardrailAPIInputs: + return { + **inputs, + "texts": [text.replace("Secret Lane", "[MASKED]") for text in inputs.get("texts", [])], + } + + return MaskLane(guardrail_name="mask-lane") + + @pytest.mark.asyncio + async def test_non_streaming_summary_is_masked_inside_its_item(self) -> None: + response: Final = ResponsesAPIResponse( + id="resp_1", + created_at=1, + model="claude-sonnet-5-5", + object="response", + status="completed", + output=[self._compaction_item()], + ) + + result: Final = await OpenAIResponsesHandler().process_output_response(response, self._lane_masking_guardrail()) + + assert self._summary(result.output[0]) == "The user lives at 1 [MASKED]." + + @pytest.mark.asyncio + async def test_ended_stream_summary_is_masked_in_the_done_item_and_the_envelope(self) -> None: + events: Final = [ + { + "type": "response.output_item.added", + "output_index": 0, + "item": {**self._compaction_item(), "encrypted_content": ""}, + }, + {"type": "response.output_item.done", "output_index": 0, "item": self._compaction_item()}, + { + "type": "response.completed", + "response": { + "id": "resp_1", + "created_at": 1, + "model": "claude-sonnet-5-5", + "output": [self._compaction_item()], + "status": "completed", + }, + }, + ] + + await OpenAIResponsesHandler().process_output_streaming_response( + responses_so_far=events, + guardrail_to_apply=self._lane_masking_guardrail(), + deliver_ended_stream_rewrites=True, + ) + + assert self._summary(events[1]["item"]) == "The user lives at 1 [MASKED]." + assert self._summary(events[2]["response"]["output"][0]) == "The user lives at 1 [MASKED]." + + @pytest.mark.asyncio + async def test_a_blocking_guardrail_rejects_a_finished_compaction_item_mid_stream(self) -> None: + with pytest.raises(HTTPException) as exc_info: + await OpenAIResponsesHandler().process_output_streaming_response( + responses_so_far=[ + {"type": "response.output_item.done", "output_index": 0, "item": self._compaction_item()} + ], + guardrail_to_apply=MockGuardrail(guardrail_name="block"), + ) + + assert exc_info.value.detail["texts"] == [self.SUMMARY] + + def test_a_changed_summary_changes_the_ended_stream_scan_key(self) -> None: + handler: Final = OpenAIResponsesHandler() + + def ended_key(summary: str) -> object: + return handler.get_streaming_scan_key( + [{"type": "response.completed", "response": {"output": [self._compaction_item(summary)]}}] + ) + + assert ended_key(self.SUMMARY) != ended_key("The user lives elsewhere.") diff --git a/tests/unit/proxy/guardrails/test_content_utils.py b/tests/unit/proxy/guardrails/test_content_utils.py index 920ffc77095..c14375135dc 100644 --- a/tests/unit/proxy/guardrails/test_content_utils.py +++ b/tests/unit/proxy/guardrails/test_content_utils.py @@ -1,5 +1,7 @@ """Tests for the shared guardrail content extraction helpers.""" +import json + from litellm.proxy.guardrails._content_utils import ( apply_redacted_messages_back, build_inspection_messages, @@ -167,6 +169,21 @@ def test_iter_message_text_inspects_reasoning_content_and_summary(): assert list(iter_message_text(data)) == ["content secret", "summary secret"] + +def test_iter_message_text_inspects_compaction_summary(): + data = { + "input": [ + { + "type": "compaction", + "id": "cmp_1", + "encrypted_content": json.dumps({"type": "compaction", "content": "banned word"}), + }, + {"role": "user", "content": "next turn"}, + ] + } + assert list(iter_message_text(data)) == ["banned word", "next turn"] + + # ── walk_user_text ──────────────────────────────────────────────────────────── @@ -347,6 +364,30 @@ def test_walk_user_text_redacts_reasoning_content_and_summary(): assert item["id"] == "rs_1" + +def test_walk_user_text_redacts_compaction_summary_and_keeps_the_block(): + data = { + "input": [ + { + "type": "compaction", + "id": "cmp_1", + "encrypted_content": json.dumps( + {"type": "compaction", "content": "AKIAEXAMPLE summary", "signature": "sig"} + ), + } + ] + } + visited = walk_user_text(data, lambda s: s.replace("AKIAEXAMPLE", "[REDACTED]")) + item = data["input"][0] + assert visited == 1 + assert (item["type"], item["id"]) == ("compaction", "cmp_1") + assert json.loads(item["encrypted_content"]) == { + "type": "compaction", + "content": "[REDACTED] summary", + "signature": "sig", + } + + # ── build_inspection_messages ───────────────────────────────────────────────── diff --git a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py index bd997f10322..76ad1e08e28 100644 --- a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py @@ -415,6 +415,74 @@ async def test_async_post_call_failure_hook_does_not_clobber_guardrail_info_in_m assert metadata["standard_logging_guardrail_information"] == metadata_bucket_info +@pytest.mark.asyncio +@pytest.mark.parametrize( + "request_headers, caller_metadata, redaction_is_global, code_word_is_stored", + [ + ({"x-litellm-enable-message-redaction": "true"}, {"user_id": "claude-code-session"}, False, False), + ({"litellm-enable-message-redaction": "true"}, {}, False, False), + ({"user-agent": "claude-cli"}, {"user_id": "claude-code-session"}, False, True), + ({"user-agent": "claude-cli"}, {"headers": {"x-litellm-enable-message-redaction": "true"}}, False, False), + ({"litellm-disable-message-redaction": "true"}, {"user_id": "claude-code-session"}, True, False), + ], + ids=[ + "enable_header_redacts", + "legacy_enable_header_redacts", + "no_header_stores_prompt", + "caller_metadata_headers_still_apply", + "disable_header_never_loosens_global_redaction", + ], +) +async def test_async_post_call_failure_hook_honors_the_redaction_header_on_a_litellm_metadata_route( + monkeypatch: pytest.MonkeyPatch, + request_headers: dict[str, str], + caller_metadata: dict[str, object], + redaction_is_global: bool, + code_word_is_stored: bool, +): + monkeypatch.setenv("STORE_PROMPTS_IN_SPEND_LOGS", "true") + monkeypatch.setattr(litellm, "turn_off_message_logging", redaction_is_global) + writer: Final = MagicMock(spec=DBSpendUpdateWriter) + writer.update_database = AsyncMock() + logger: Final = ProxyDBLogger(spend_writer=lambda: writer) + request_input: Final = "The code word is VERMILION-42. Reply OK." + body_metadata: Final = dict(caller_metadata) + litellm_metadata_headers: Final = dict(request_headers) + + await logger.async_post_call_failure_hook( + request_data={ + "model": "gpt-5.6", + "input": request_input, + "metadata": body_metadata, + "litellm_metadata": {"headers": litellm_metadata_headers}, + "proxy_server_request": { + "url": "http://localhost/v1/responses", + "body": { + "model": "gpt-5.6", + "input": request_input, + "temperature": 7, + "metadata": body_metadata, + "litellm_metadata": {"headers": litellm_metadata_headers}, + }, + }, + }, + original_exception=Exception("temperature must be at most 2"), + user_api_key_dict=UserAPIKeyAuth(api_key="test_api_key"), + ) + + row: Final = get_logging_payload( + kwargs=writer.update_database.call_args.kwargs["kwargs"], + response_obj={}, + start_time=datetime.now(), + end_time=datetime.now(), + ) + stored_request: Final = json.loads(row["proxy_server_request"]) + assert stored_request["temperature"] == 7 + assert ("VERMILION-42" in row["proxy_server_request"]) is code_word_is_stored + assert stored_request["metadata"].get("headers") == caller_metadata.get("headers") + assert stored_request["litellm_metadata"]["headers"] == request_headers + + @pytest.mark.asyncio @pytest.mark.parametrize( "used_client_oauth_token, custom_llm_provider, expected", diff --git a/tests/unit/proxy/hooks/test_tpm_concurrent.py b/tests/unit/proxy/hooks/test_tpm_concurrent.py index bb2a89b1d98..3d1e494e36c 100644 --- a/tests/unit/proxy/hooks/test_tpm_concurrent.py +++ b/tests/unit/proxy/hooks/test_tpm_concurrent.py @@ -15,8 +15,9 @@ Redis. """ import asyncio +import json from datetime import datetime, timedelta -from typing import Any, Dict +from typing import Any, Dict, Final import pytest @@ -3727,3 +3728,50 @@ async def test_the_project_itpm_reservation_counts_the_request_off_the_event_loo if __name__ == "__main__": pytest.main([__file__, "-v", "-s"]) + + +_COMPACTION_SUMMARY: Final = "The user planned a Tokyo trip." +_NEXT_TURN: Final = {"type": "message", "role": "user", "content": "And Osaka?"} +_COMPACTED_INPUT: Final = [ + {"type": "message", "role": "user", "content": "Plan a Tokyo trip. " * 200}, + {"type": "message", "role": "assistant", "content": "Day one, Asakusa. " * 200}, + { + "type": "compaction", + "id": "cmp_1", + "encrypted_content": json.dumps({"type": "compaction", "content": _COMPACTION_SUMMARY}), + }, + _NEXT_TURN, +] +_SUMMARY_INPUT: Final = [{"type": "message", "role": "user", "content": _COMPACTION_SUMMARY}, _NEXT_TURN] + + +def test_reservation_skips_the_history_a_compaction_item_replaces(rate_limiter): + handler, _cache = rate_limiter + compacted_request: Final = {"model": "gpt-4o", "input": _COMPACTED_INPUT, "max_output_tokens": 1} + + assert handler._estimate_input_and_output_tokens( + data=compacted_request, call_type="aresponses" + ) == handler._estimate_input_and_output_tokens( + data={**compacted_request, "input": _COMPACTED_INPUT[2:]}, call_type="aresponses" + ) + assert handler._estimate_precise_input_tokens( + compacted_request, model="gpt-4o", call_type="aresponses" + ) == handler._estimate_precise_input_tokens( + {**compacted_request, "input": _SUMMARY_INPUT}, model="gpt-4o", call_type="aresponses" + ) + + +def test_reservation_counts_a_compaction_item_by_its_summary(rate_limiter): + handler, _cache = rate_limiter + padded_item: Final = { + **_COMPACTED_INPUT[2], + "encrypted_content": json.dumps( + {"type": "compaction", "content": _COMPACTION_SUMMARY, "opaque_state": "x" * 4000} + ), + } + + assert handler._estimate_input_and_output_tokens( + data={"model": "gpt-4o", "input": [padded_item, _NEXT_TURN], "max_output_tokens": 1}, call_type="aresponses" + ) == handler._estimate_input_and_output_tokens( + data={"model": "gpt-4o", "input": _COMPACTED_INPUT[2:], "max_output_tokens": 1}, call_type="aresponses" + ) diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index abc1fe88c7a..985aee061bc 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -176,6 +176,105 @@ async def test_streaming_upstream_errors_keep_the_client_protocol( assert "error" in events[-1] +_COMPACTION_ITEM: Final = { + "type": "compaction", + "encrypted_content": json.dumps({"type": "compaction", "content": "the cat is named Biscuit"}), +} +_HISTORY_BEFORE_COMPACTION: Final = ( + {"role": "user", "content": "a long turn " * 200}, + {"role": "assistant", "content": "a long reply " * 200}, +) +_TURN_AFTER_COMPACTION: Final = {"role": "user", "content": "what is the cat's name?"} + + +async def _cursor_stream_failure_prompt_tokens( + monkeypatch: pytest.MonkeyPatch, request_input: list[Mapping[str, str]], fails_after_output: bool +) -> int: + import litellm.proxy.proxy_server as ps + from litellm.integrations.custom_logger import CustomLogger + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + class _FailureRowRecorder(CustomLogger): + prompt_tokens: int | None = None + + async def async_post_call_failure_hook( + self, + request_data: dict, + original_exception: Exception, + user_api_key_dict: object, + traceback_str: str | None = None, + ) -> None: + usage: Final = request_data.get("combined_usage_object") + self.prompt_tokens = usage.prompt_tokens if isinstance(usage, litellm.Usage) else None + + model: Final = "gpt-6-astra" + upstream_url: Final = "https://streaming.example/v1" + created: Final = { + "type": "response.created", + "sequence_number": 0, + "response": { + "id": "resp_upstream", + "object": "response", + "created_at": 1, + "status": "in_progress", + "model": model, + "output": [], + }, + } + output_delta: Final = { + "type": "response.output_text.delta", + "sequence_number": 1, + "item_id": "msg_partial", + "output_index": 0, + "content_index": 0, + "delta": "Biscuit", + } + failed: Final = {"type": "error", "error": {"message": "upstream broke", "code": "server_error", "type": None}} + events: Final = (created, output_delta, failed) if fails_after_output else (created, failed) + wire: Final = "".join("data: " + json.dumps(event) + "\n\n" for event in events) + recorder: Final = _FailureRowRecorder() + router: Final = litellm.Router( + model_list=[ + { + "model_name": model, + "litellm_params": {"model": "openai/" + model, "api_base": upstream_url, "api_key": "fixture-key"}, + } + ], + num_retries=0, + ) + monkeypatch.setattr(ps, "llm_router", router) + monkeypatch.setattr(litellm, "callbacks", [recorder]) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setitem(app.dependency_overrides, user_api_key_auth, _auth_override) + with respx.mock as transport: + transport.post(upstream_url + "/responses").respond( + 200, content=wire, headers={"Content-Type": "text/event-stream"} + ) + async with httpx.AsyncClient(transport=httpx.ASGITransport(app), base_url="http://testserver") as client: + result: Final = await client.post( + "/cursor/chat/completions", json={"model": model, "stream": True, "input": request_input} + ) + assert recorder.prompt_tokens is not None, result.text + return recorder.prompt_tokens + + +@pytest.mark.asyncio +@pytest.mark.parametrize("fails_after_output", [False, True]) +async def test_a_failed_cursor_stream_counts_its_input_from_the_newest_compaction_item( + monkeypatch: pytest.MonkeyPatch, fails_after_output: bool +) -> None: + """Cursor's route maps to no call type, yet it dispatches a Responses call, which sends + the provider only the newest compaction item onward, so the failure row counts that much.""" + resent: Final = await _cursor_stream_failure_prompt_tokens( + monkeypatch, [*_HISTORY_BEFORE_COMPACTION, _COMPACTION_ITEM, _TURN_AFTER_COMPACTION], fails_after_output + ) + dispatched_after_trim: Final = await _cursor_stream_failure_prompt_tokens( + monkeypatch, [_COMPACTION_ITEM, _TURN_AFTER_COMPACTION], fails_after_output + ) + + assert resent == dispatched_after_trim + + @pytest.mark.asyncio async def test_responses_api_background_polling_rejects_missing_input(): from fastapi import Response as FastAPIResponse @@ -4832,3 +4931,26 @@ def test_responses_routes_document_response_models_in_openapi_schema(): assert "output" in ok_200_properties("/v1/responses/{response_id}", "get") assert "deleted" in ok_200_properties("/v1/responses/{response_id}", "delete") assert "data" in ok_200_properties("/v1/responses/{response_id}/input_items", "get") + + +_COMPACTION_SUMMARY: Final = "The user planned a Tokyo trip." +_NEXT_TURN: Final = {"type": "message", "role": "user", "content": "And Osaka?"} +_COMPACTED_INPUT: Final = [ + {"type": "message", "role": "user", "content": "Plan a Tokyo trip. " * 200}, + {"type": "message", "role": "assistant", "content": "Day one, Asakusa. " * 200}, + { + "type": "compaction", + "id": "cmp_1", + "encrypted_content": json.dumps({"type": "compaction", "content": _COMPACTION_SUMMARY}), + }, + _NEXT_TURN, +] +_SUMMARY_INPUT: Final = [{"type": "message", "role": "user", "content": _COMPACTION_SUMMARY}, _NEXT_TURN] + + +def test_input_token_count_reads_a_compaction_item_as_its_summary(): + from litellm.proxy.response_api_endpoints.endpoints import _responses_input_as_token_count_messages + + assert _responses_input_as_token_count_messages(_COMPACTED_INPUT, None) == ( + _responses_input_as_token_count_messages(_SUMMARY_INPUT, None) + ) diff --git a/tests/unit/proxy/spend_tracking/test_input_tokens.py b/tests/unit/proxy/spend_tracking/test_input_tokens.py index a1da6c79cd7..71f4dec2345 100644 --- a/tests/unit/proxy/spend_tracking/test_input_tokens.py +++ b/tests/unit/proxy/spend_tracking/test_input_tokens.py @@ -2,6 +2,7 @@ from __future__ import annotations +import json from types import MappingProxyType from typing import Final @@ -27,3 +28,59 @@ async def test_large_input_is_still_counted() -> None: assert counts[CL100K_MODEL] == count_input_tokens_for_model(request_body=request_body, model=CL100K_MODEL) assert isinstance(counts, MappingProxyType) + + +_COMPACTION_ITEM: Final = { + "type": "compaction", + "id": "cmp_1", + "encrypted_content": json.dumps({"type": "compaction", "content": "Turn one planned Tokyo."}), +} +_TURN_AFTER_COMPACTION: Final = {"role": "user", "content": "And Osaka?"} + + +@pytest.mark.parametrize( + "earlier_item", + [ + pytest.param({"role": "user", "content": "Plan a Tokyo trip. " * 500}, id="user_turn"), + pytest.param({"role": ["system"], "content": "Plan a Tokyo trip."}, id="list_role"), + pytest.param({"role": {"name": "system"}, "content": "Plan a Tokyo trip."}, id="dict_role"), + ], +) +@pytest.mark.parametrize("model", [CL100K_MODEL, "claude-sonnet-5-5"]) +@pytest.mark.asyncio +async def test_history_before_a_compaction_item_is_not_counted(model: str, earlier_item: dict[str, object]) -> None: + compacted_body: Final = {"model": model, "input": [_COMPACTION_ITEM, _TURN_AFTER_COMPACTION]} + full_body: Final = {"model": model, "input": [earlier_item, _COMPACTION_ITEM, _TURN_AFTER_COMPACTION]} + + full_counts: Final = await count_input_tokens( + request_body=full_body, raw_body=json.dumps(full_body).encode(), models=(model,) + ) + compacted_counts: Final = await count_input_tokens( + request_body=compacted_body, raw_body=json.dumps(compacted_body).encode(), models=(model,) + ) + + assert full_counts[model] == compacted_counts[model] + assert count_input_tokens_for_model(request_body=full_body, model=model) == compacted_counts[model] + + +def _compaction_item_carrying(extra_block_fields: dict[str, str]) -> dict[str, str]: + return { + "type": "compaction", + "id": "cmp_1", + "encrypted_content": json.dumps( + {"type": "compaction", "content": "Turn one planned Tokyo.", **extra_block_fields} + ), + } + + +@pytest.mark.parametrize("model", [CL100K_MODEL, "claude-sonnet-5-5"]) +def test_a_compaction_item_counts_only_its_summary(model: str) -> None: + plain_body: Final = {"model": model, "input": [_compaction_item_carrying({}), _TURN_AFTER_COMPACTION]} + padded_body: Final = { + "model": model, + "input": [_compaction_item_carrying({"opaque_state": "x" * 4000}), _TURN_AFTER_COMPACTION], + } + + assert count_input_tokens_for_model(request_body=padded_body, model=model) == count_input_tokens_for_model( + request_body=plain_body, model=model + ) 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 55358a3b0de..78f5e1ace0c 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,23 @@ 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 KUBERNETES_POD_ROUTING_KEY, 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.constants import ( + KUBERNETES_POD_ROUTING_KEY, + 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 @@ -366,7 +381,10 @@ def test_caller_session_keeps_the_legacy_trace_id_with_correlation_logging_off(m 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"}} + 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, @@ -380,7 +398,11 @@ def test_session_displaced_by_a_trace_header_wins_only_with_correlation_logging( ("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"), + 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( @@ -394,7 +416,10 @@ def test_caller_session_postgres_cannot_store_keeps_the_trace_id( [ pytest.param({"metadata": {}}, id="no-session"), pytest.param( - {"litellm_session_id": "minted-1", "metadata": {"session_id": "minted-1", SESSION_ID_GENERATED_METADATA_KEY: True}}, + { + "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"), @@ -732,6 +757,19 @@ def test_large_request_no_truncation_threshold(): assert "skipped" in truncated_content +@pytest.mark.parametrize( + ("max_string_length_prompt_in_db", "stored_tail"), + [(0, "claude-sonnet-5-5"), (1, "claude-sonnet-5-5"), (-1, "laude-sonnet-5-5"), (-2, "aude-sonnet-5-5")], +) +def test_a_limit_under_two_keeps_mains_tail_after_the_marker(max_string_length_prompt_in_db: int, stored_tail: str): + sanitized: Final = _sanitize_request_body_for_spend_logs_payload( + {"model": "claude-sonnet-5-5"}, max_string_length_prompt_in_db=max_string_length_prompt_in_db + ) + + assert sanitized["model"].startswith("...") + assert sanitized["model"].endswith(f") ...{stored_tail}") + + def test_small_request_no_truncation(): """ Test that small strings are not truncated by MAX_STRING_LENGTH_PROMPT_IN_DB @@ -750,9 +788,7 @@ def test_small_request_no_truncation(): # Verify the content was NOT truncated assert sanitized["messages"][0]["content"] == small_content - assert ( - len(sanitized["messages"][0]["content"]) == MAX_STRING_LENGTH_PROMPT_IN_DB - 100 - ) + assert len(sanitized["messages"][0]["content"]) == MAX_STRING_LENGTH_PROMPT_IN_DB - 100 def test_configurable_string_length_env_var(monkeypatch): @@ -1144,6 +1180,187 @@ def test_get_response_for_spend_logs_payload_truncates_large_base64(mock_should_ assert parsed["data"][0]["other_field"] == "value" +def test_get_response_for_spend_logs_payload_keeps_a_long_compaction_summary_whole(monkeypatch: pytest.MonkeyPatch): + from litellm.constants import MAX_STRING_LENGTH_PROMPT_IN_DB + + monkeypatch.setenv("STORE_PROMPTS_IN_SPEND_LOGS", "true") + filler = "detail " * (MAX_STRING_LENGTH_PROMPT_IN_DB // 7) + summary = f"{filler}The cat is named Biscuit. {filler}" + answer = f"{filler}OK{filler}" + payload = cast( + StandardLoggingPayload, + { + "response": { + "choices": [ + { + "message": { + "role": "assistant", + "content": answer, + "provider_specific_fields": { + "compaction_blocks": [{"type": "compaction", "content": summary}] + }, + } + } + ] + } + }, + ) + + message = json.loads(_get_response_for_spend_logs_payload(payload))["choices"][0]["message"] + assert message["provider_specific_fields"]["compaction_blocks"] == [{"type": "compaction", "content": summary}] + assert LITELLM_TRUNCATED_PAYLOAD_FIELD in message["content"] + + +def test_stored_request_keeps_only_the_summary_of_a_long_claude_compaction_input_item(monkeypatch: pytest.MonkeyPatch): + from litellm.constants import MAX_STRING_LENGTH_PROMPT_IN_DB + + monkeypatch.setenv("STORE_PROMPTS_IN_SPEND_LOGS", "true") + filler = "detail " * (MAX_STRING_LENGTH_PROMPT_IN_DB // 7) + summary = f"{filler}The cat is named Biscuit. {filler}" + claude_item = json.dumps({"type": "compaction", "content": summary, "padding": filler * 10}) + opaque_item = "gAAAA" + "x" * (MAX_STRING_LENGTH_PROMPT_IN_DB * 2) + request_body = { + "input": [ + {"type": "compaction", "id": "cmp_1", "encrypted_content": claude_item}, + {"type": "compaction", "id": "cmp_2", "encrypted_content": opaque_item}, + {"role": "assistant", "content": [{"type": "compaction", "content": f"{filler}Biscuit{filler}"}]}, + {"role": "user", "content": f"{filler}What is my cat called?{filler}"}, + ], + "metadata": {"note": {"type": "compaction", "encrypted_content": claude_item}}, + } + + stored = json.loads( + _get_proxy_server_request_for_spend_logs_payload( + metadata={}, litellm_params={"proxy_server_request": {"body": request_body}} + ) + ) + + claude_stored, opaque_stored, messages_block_stored, message_stored = stored["input"] + assert {**claude_stored, "encrypted_content": json.loads(claude_stored["encrypted_content"])} == { + "type": "compaction", + "id": "cmp_1", + "encrypted_content": {"type": "compaction", "content": summary}, + } + assert LITELLM_TRUNCATED_PAYLOAD_FIELD in opaque_stored["encrypted_content"] + assert LITELLM_TRUNCATED_PAYLOAD_FIELD in messages_block_stored["content"][0]["content"] + assert LITELLM_TRUNCATED_PAYLOAD_FIELD in message_stored["content"] + assert LITELLM_TRUNCATED_PAYLOAD_FIELD in stored["metadata"]["note"]["encrypted_content"] + + +def _compaction_summary_past_the_summary_cap() -> str: + from litellm.constants import MAX_STRING_LENGTH_COMPACTION_SUMMARY_IN_DB + + filler: Final = "detail " * (MAX_STRING_LENGTH_COMPACTION_SUMMARY_IN_DB // 7) + return f"The cat is named Biscuit. {filler}The dog is named Rex." + + +def _assert_bounded_like_a_db_string(stored: str, summary: str) -> None: + from litellm.constants import MAX_STRING_LENGTH_COMPACTION_SUMMARY_IN_DB + + assert LITELLM_TRUNCATED_PAYLOAD_FIELD in stored + assert len(stored) < MAX_STRING_LENGTH_COMPACTION_SUMMARY_IN_DB + 500 < len(summary) + 500 + assert stored.startswith("The cat is named Biscuit.") + assert stored.endswith("The dog is named Rex.") + + +@pytest.mark.parametrize("operator_cap_above_the_summary", [False, True]) +def test_spend_row_bounds_a_compaction_summary_at_the_larger_of_the_summary_cap_and_the_operator_cap( + monkeypatch: pytest.MonkeyPatch, operator_cap_above_the_summary: bool +): + monkeypatch.setenv("STORE_PROMPTS_IN_SPEND_LOGS", "true") + summary: Final = _compaction_summary_past_the_summary_cap() + if operator_cap_above_the_summary: + monkeypatch.setenv("MAX_STRING_LENGTH_PROMPT_IN_DB", str(len(summary) * 2)) + payload: Final = cast( + StandardLoggingPayload, + { + "response": { + "choices": [ + { + "message": { + "role": "assistant", + "content": "OK", + "provider_specific_fields": { + "compaction_blocks": [{"type": "compaction", "content": summary}] + }, + } + } + ] + } + }, + ) + request_body: Final = { + "input": [ + { + "type": "compaction", + "id": "cmp_1", + "encrypted_content": json.dumps({"type": "compaction", "content": summary}), + } + ] + } + + response_message: Final = json.loads(_get_response_for_spend_logs_payload(payload))["choices"][0]["message"] + stored_item: Final = json.loads( + _get_proxy_server_request_for_spend_logs_payload( + metadata={}, litellm_params={"proxy_server_request": {"body": request_body}} + ) + )["input"][0] + + stored_response_block: Final = response_message["provider_specific_fields"]["compaction_blocks"][0] + stored_request_block: Final = json.loads(stored_item["encrypted_content"]) + assert stored_request_block["type"] == stored_response_block["type"] == "compaction" + if operator_cap_above_the_summary: + assert stored_response_block["content"] == stored_request_block["content"] == summary + return + _assert_bounded_like_a_db_string(stored_response_block["content"], summary) + _assert_bounded_like_a_db_string(stored_request_block["content"], summary) + + +@pytest.mark.parametrize("summary_past_the_summary_cap", [False, True]) +def test_responses_shaped_spend_row_keeps_a_compaction_output_item_summary_bounded_like_a_chat_block( + monkeypatch: pytest.MonkeyPatch, summary_past_the_summary_cap: bool +): + from litellm.constants import MAX_STRING_LENGTH_PROMPT_IN_DB + from litellm.responses.litellm_completion_transformation.compaction_items import compaction_summary + from litellm.responses.utils import ResponsesAPIRequestUtils + + monkeypatch.setenv("STORE_PROMPTS_IN_SPEND_LOGS", "true") + filler: Final = "detail " * (MAX_STRING_LENGTH_PROMPT_IN_DB // 7) + summary: Final = ( + _compaction_summary_past_the_summary_cap() + if summary_past_the_summary_cap + else f"{filler}The cat is named Biscuit. {filler}" + ) + payload: Final = cast( + StandardLoggingPayload, + { + "response": { + "object": "response", + "output": [ + { + "type": "compaction", + "id": "cmp_1", + "encrypted_content": ResponsesAPIRequestUtils.wrap_encrypted_content_with_model_id( + json.dumps({"type": "compaction", "content": summary}), "deployment-1" + ), + }, + {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "OK"}]}, + ], + } + }, + ) + + stored_item: Final = json.loads(_get_response_for_spend_logs_payload(payload))["output"][0] + model_id, _ = ResponsesAPIRequestUtils.unwrap_encrypted_content_with_model_id(stored_item["encrypted_content"]) + stored_summary: Final = compaction_summary(stored_item) + assert model_id == "deployment-1" + assert stored_summary is not None + if summary_past_the_summary_cap: + _assert_bounded_like_a_db_string(stored_summary, summary) + return + assert stored_summary == summary + + @patch("litellm.proxy.spend_tracking.spend_tracking_utils.should_store_prompts_and_responses_in_spend_logs") def test_get_response_for_spend_logs_payload_strips_null_bytes(mock_should_store): """Regression for PostgreSQL 22P05: NUL bytes must be stripped from response.""" diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index a226ce8a4a9..92bf2b45fab 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -4066,6 +4066,72 @@ class TestHasAttributeErrorInChain: assert _has_attribute_error_in_chain(exc_a) is False +_COMPACTION_ITEM: Final = { + "type": "compaction", + "encrypted_content": json.dumps({"type": "compaction", "content": "the cat is named Biscuit"}), +} +_HISTORY_BEFORE_COMPACTION: Final = ( + {"role": "user", "content": "a long turn " * 200}, + {"role": "assistant", "content": "a long reply " * 200}, +) +_TURN_AFTER_COMPACTION: Final = {"role": "user", "content": "what is the cat's name?"} + + +class TestHandleLLMApiExceptionPricesWhatTheServerDispatched: + async def _failed_prompt_tokens( + self, + route_type: Literal["acompletion", "aresponses"], + logged_call_type: str, + request_input: list[dict[str, str]], + ) -> int: + processing_obj: Final = ProxyBaseLLMRequestProcessing(data={"model": "gpt-4o", "input": request_input}) + cursor_caller: Final = ProxyUserAPIKeyAuth(request_route="/cursor/chat/completions") + mock_request: Final = MagicMock(spec=Request) + mock_request.headers = {} + proxy_logging: Final = ProxyLogging(user_api_key_cache=DualCache()) + proxy_logging.alert_types = [] + _, logging_obj = await processing_obj.common_processing_pre_call_logic( + request=mock_request, + general_settings={}, + user_api_key_dict=cursor_caller, + proxy_logging_obj=proxy_logging, + proxy_config=MagicMock(spec=ProxyConfig), + route_type=route_type, + ) + logging_obj.model_call_details.update( + first_api_call_start_time=datetime.datetime.now(), + messages=request_input, + call_type=logged_call_type, + ) + with pytest.raises(ProxyException): + await processing_obj.handle_llm_api_exception( + e=Exception("boom"), + user_api_key_dict=cursor_caller, + proxy_logging_obj=proxy_logging, + ) + return processing_obj.data["combined_usage_object"].prompt_tokens + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "route_type, logged_call_type, trims_at_the_compaction", + [("aresponses", "aresponses", True), ("acompletion", "responses", False)], + ) + async def test_a_failed_call_counts_from_the_newest_compaction_only_when_dispatched_as_responses( + self, + route_type: Literal["acompletion", "aresponses"], + logged_call_type: str, + trims_at_the_compaction: bool, + ) -> None: + async def prompt_tokens(request_input: list[dict[str, str]]) -> int: + return await self._failed_prompt_tokens(route_type, logged_call_type, request_input) + + resent: Final = await prompt_tokens([*_HISTORY_BEFORE_COMPACTION, _COMPACTION_ITEM, _TURN_AFTER_COMPACTION]) + dispatched_after_trim: Final = await prompt_tokens([_COMPACTION_ITEM, _TURN_AFTER_COMPACTION]) + + assert (resent == dispatched_after_trim) is trims_at_the_compaction + assert resent >= dispatched_after_trim + + @pytest.mark.asyncio class TestHandleLLMApiExceptionDictDetail: """ @@ -10817,6 +10883,7 @@ class TestDetachedStreamFailureHook: "user_api_key_dict": user_api_key_dict, "original_exception": failure, "request_data": processing_obj.data, + "dispatched_call_type": "anthropic_messages", } ] diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index 3f8e3bc93d1..d9def23bde8 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -1229,6 +1229,7 @@ async def test_add_litellm_data_to_request_strips_user_control_fields(): "litellm_metadata": copy.deepcopy(malicious_metadata), "weights": {"gpt-3.5-turbo": {"forged-deployment-id": 100}}, "_router_weights": {"gpt-3.5-turbo": {"forged-deployment-id": 100}}, + "_encrypted_content_affinity_pinned": True, } updated = await add_litellm_data_to_request( @@ -1251,6 +1252,7 @@ async def test_add_litellm_data_to_request_strips_user_control_fields(): assert "_router_weights" not in updated assert "weights" not in updated["proxy_server_request"]["body"] assert "_router_weights" not in updated["proxy_server_request"]["body"] + assert "_encrypted_content_affinity_pinned" not in updated stripped_keys = { "disable_global_guardrails", diff --git a/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py b/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py index ab6a874a010..b5940e4c212 100644 --- a/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py +++ b/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py @@ -1,5 +1,6 @@ import asyncio import datetime as real_datetime +import json import smtplib import traceback from collections.abc import Mapping @@ -745,7 +746,9 @@ class TestPostCallFailureHookEstimatesDispatchedInputTokens: post-dispatch failures keep zero spend. """ - async def _run(self, request_data, original_exception: Exception | None = None): + async def _run( + self, request_data, original_exception: Exception | None = None, request_route=None, dispatched_call_type=None + ): from unittest.mock import AsyncMock, patch from litellm.proxy._types import UserAPIKeyAuth @@ -756,7 +759,8 @@ class TestPostCallFailureHookEstimatesDispatchedInputTokens: await proxy_logging_obj.post_call_failure_hook( request_data=request_data, original_exception=original_exception or Exception("boom"), - user_api_key_dict=UserAPIKeyAuth(), + user_api_key_dict=UserAPIKeyAuth(request_route=request_route), + dispatched_call_type=dispatched_call_type, ) def _timeout(self, model: str, cause: Exception) -> Exception: @@ -773,6 +777,38 @@ class TestPostCallFailureHookEstimatesDispatchedInputTokens: logging_obj.model_call_details = model_call_details return logging_obj + @pytest.mark.asyncio + @pytest.mark.parametrize( + "request_route, dispatched_call_type", + [("/v1/responses", None), ("/cursor/chat/completions", "aresponses")], + ) + async def test_a_failed_responses_request_counts_its_input_from_the_newest_compaction_item( + self, request_route, dispatched_call_type + ): + from datetime import datetime + + async def prompt_tokens(request_input): + request_data = { + "litellm_logging_obj": self._logging_obj( + { + "first_api_call_start_time": datetime.now(), + "model": FAILURE_USAGE_MODEL, + "messages": request_input, + "call_type": "aresponses", + } + ), + "metadata": {}, + } + await self._run(request_data, request_route=request_route, dispatched_call_type=dispatched_call_type) + return request_data["combined_usage_object"].prompt_tokens + + resent = await prompt_tokens( + [COMPACTION_INSTRUCTIONS, *COMPACTED_HISTORY, COMPACTION_ITEM, TURN_AFTER_COMPACTION] + ) + + assert resent == await prompt_tokens([COMPACTION_INSTRUCTIONS, COMPACTION_ITEM, TURN_AFTER_COMPACTION]) + assert resent < await prompt_tokens([COMPACTION_INSTRUCTIONS, *COMPACTED_HISTORY, TURN_AFTER_COMPACTION]) + @pytest.mark.asyncio async def test_timed_out_request_estimates_input_tokens_priced_at_input_rate(self): from datetime import datetime @@ -2164,29 +2200,29 @@ def test_a_system_prompt_reads_the_same_whatever_shape_it_arrived_in(system_inpu def test_a_system_prompt_is_counted_on_top_of_the_request(): - from litellm.proxy.utils import _count_request_input_tokens + from litellm.proxy.utils import count_request_input_tokens - without = _count_request_input_tokens(FAILURE_USAGE_MODEL, "hello world", None) - with_system = _count_request_input_tokens(FAILURE_USAGE_MODEL, "hello world", "be brief") + without = count_request_input_tokens(FAILURE_USAGE_MODEL, "hello world", None) + with_system = count_request_input_tokens(FAILURE_USAGE_MODEL, "hello world", "be brief") assert without > 0 assert with_system > without def test_a_request_with_nothing_in_it_counts_zero(): - from litellm.proxy.utils import _count_request_input_tokens + from litellm.proxy.utils import count_request_input_tokens - assert _count_request_input_tokens(FAILURE_USAGE_MODEL, [], None) == 0 - assert _count_request_input_tokens(FAILURE_USAGE_MODEL, None, None) == 0 + assert count_request_input_tokens(FAILURE_USAGE_MODEL, [], None) == 0 + assert count_request_input_tokens(FAILURE_USAGE_MODEL, None, None) == 0 def test_a_failed_dispatch_is_estimated_as_input_only(): - from litellm.proxy.utils import _count_request_input_tokens, _estimate_dispatched_failure_usage + from litellm.proxy.utils import count_request_input_tokens, _estimate_dispatched_failure_usage usage = _estimate_dispatched_failure_usage(FAILURE_USAGE_MODEL, ONE_USER_MESSAGE, None) assert usage is not None - assert usage.prompt_tokens == _count_request_input_tokens(FAILURE_USAGE_MODEL, ONE_USER_MESSAGE, None) + assert usage.prompt_tokens == count_request_input_tokens(FAILURE_USAGE_MODEL, ONE_USER_MESSAGE, None) assert usage.completion_tokens == 0 assert usage.total_tokens == usage.prompt_tokens @@ -2450,7 +2486,7 @@ def test_earlier_attempts_are_priced_from_the_prompt_estimate_not_the_final_stre lifted = _failure_fields_to_lift({"litellm_logging_obj": logging_obj}, provider_billed=True) - estimate = _estimated_prompt_usage(logging_obj.model_call_details, {}) + estimate = _estimated_prompt_usage(logging_obj.model_call_details, {}, None) assert lifted["combined_usage_object"] is recovered assert [(a.usage.prompt_tokens, a.usage.completion_tokens) for a in lifted["billed_failed_attempts"]] == [ (estimate.prompt_tokens, 0) @@ -3336,6 +3372,83 @@ def test_a_request_that_timed_out_at_a_provider_bills_its_input_at_the_input_rat assert response_cost == pytest.approx(expected_input_cost) + +COMPACTION_INSTRUCTIONS = {"role": "system", "content": "answer briefly"} +COMPACTED_HISTORY = [ + {"role": "user", "content": "a long turn " * 200}, + {"role": "assistant", "content": "a long reply " * 200}, +] +COMPACTION_ITEM = { + "type": "compaction", + "encrypted_content": json.dumps({"type": "compaction", "content": "the cat is named Biscuit"}), +} +TURN_AFTER_COMPACTION = {"role": "user", "content": "what is the cat's name?"} + + +def _failed_call_prompt_tokens(call_type, dispatched_call_type, request_input): + from litellm.proxy.utils import _failure_usage_to_lift + + lifted = _failure_usage_to_lift( + model_call_details={"call_type": call_type, "model": FAILURE_USAGE_MODEL, "messages": request_input}, + request_body={}, + dispatched=True, + dispatched_call_type=dispatched_call_type, + ) + assert lifted is not None + return lifted[0].prompt_tokens + + +@pytest.mark.parametrize( + "call_type, dispatched_call_type", + [ + ("acompletion", "acompletion"), + ("aanthropic_messages", "anthropic_messages"), + ("responses", "acompletion"), + ], +) +def test_a_failed_call_off_the_responses_route_counts_the_history_before_a_compaction_typed_message( + call_type, dispatched_call_type +): + def prompt_tokens(request_input): + return _failed_call_prompt_tokens(call_type, dispatched_call_type, request_input) + + sent = prompt_tokens([COMPACTION_INSTRUCTIONS, *COMPACTED_HISTORY, COMPACTION_ITEM, TURN_AFTER_COMPACTION]) + + assert sent > prompt_tokens([COMPACTION_INSTRUCTIONS, COMPACTION_ITEM, TURN_AFTER_COMPACTION]) + + +def test_an_earlier_timed_out_responses_attempt_counts_its_input_from_the_newest_compaction_item(): + from litellm.litellm_core_utils.get_litellm_params import get_litellm_params + from litellm.proxy.utils import _failure_fields_to_lift + + def earlier_attempt_prompt_tokens(request_input): + logging_obj = _dispatched_failure_logging_obj( + FAILURE_USAGE_MODEL, + get_litellm_params( + metadata={ + "attempted_retries": 1, + "model_info": {"id": "dep-compaction-final"}, + "previous_models": ( + {"deployment_id": "dep-compaction-final", "exception_type": "Timeout", "attempted_retries": 0}, + ), + } + ), + ) + logging_obj.model_call_details["call_type"] = "aresponses" + logging_obj.model_call_details["messages"] = request_input + lifted = _failure_fields_to_lift( + {"litellm_logging_obj": logging_obj}, provider_billed=True, dispatched_call_type="aresponses" + ) + (earlier_attempt,) = lifted["billed_failed_attempts"] + return earlier_attempt.usage.prompt_tokens + + sent = earlier_attempt_prompt_tokens( + [COMPACTION_INSTRUCTIONS, *COMPACTED_HISTORY, COMPACTION_ITEM, TURN_AFTER_COMPACTION] + ) + + assert sent == earlier_attempt_prompt_tokens([COMPACTION_INSTRUCTIONS, COMPACTION_ITEM, TURN_AFTER_COMPACTION]) + + @pytest.mark.parametrize( "current_attempt_api_call_start_time", [None, 1700000001.0], diff --git a/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py b/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py index 94efef050a6..36df738c0bb 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py +++ b/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py @@ -1,15 +1,17 @@ import json from collections.abc import Mapping, Sequence from itertools import chain -from typing import Final +from typing import Final, cast from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest import respx +from openai.types.responses.response_compaction_item import ResponseCompactionItem from pydantic import BaseModel import litellm +from litellm import CustomLLM from litellm.constants import REDACTED_BY_LITELLM from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.responses.litellm_completion_transformation.handler import ( @@ -52,6 +54,7 @@ async def test_async_response_api_handler_merges_trace_id_without_error(): instructions: str | None = None, request_kwargs: Mapping[str, object] | None = None, previous_response_client_ids: tuple[str, ...] = (), + input_replaces_history: bool = False, ) -> dict[str, object]: litellm_completion_request["litellm_trace_id"] = "session-trace" return litellm_completion_request @@ -561,3 +564,637 @@ def test_streamed_responses_over_anthropic_emit_the_responses_event_sequence(res assert deltas == "Argentina won." assert completed.response.status == "completed" assert completed.response.output[0].content[0].text == "Argentina won." + + +_ANTHROPIC_MESSAGES_URL: Final = "https://api.anthropic.com/v1/messages" +_COMPACTION_BLOCK: Final = {"type": "compaction", "content": "The user pasted a long document and asked about it."} +_FIRST_INPUT: Final = ({"role": "user", "content": "What is the first word of the document?"},) +_NEXT_INPUT: Final = ({"role": "user", "content": "And the last word?"},) +_COMPACT_AT_50K: Final = ({"type": "compaction", "compact_threshold": 50000},) + + +def _anthropic_message(*content: Mapping[str, object]) -> httpx.Response: + return httpx.Response( + 200, + json={ + "id": "msg_compacted", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5-5", + "content": list(content), + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 1, "output_tokens": 1}, + }, + ) + + +def _replay(output: Sequence[BaseModel]) -> list[object]: + return [*_FIRST_INPUT, *(item.model_dump(exclude_none=True) for item in output), *_NEXT_INPUT] + + +def _sent_assistant_content(route: respx.Route) -> list[object]: + sent: Final = json.loads(route.calls.last.request.content) + return next(message["content"] for message in sent["messages"] if message["role"] == "assistant") + + +def test_responses_over_anthropic_return_the_compaction_block_and_send_it_back( + respx_mock: respx.MockRouter, +) -> None: + route: Final = respx_mock.post(_ANTHROPIC_MESSAGES_URL).mock( + side_effect=[ + _anthropic_message(_COMPACTION_BLOCK, {"type": "text", "text": "The."}), + _anthropic_message({"type": "text", "text": "dog."}), + ] + ) + + first: Final = litellm.responses( + model=_MODEL, input=list(_FIRST_INPUT), context_management=list(_COMPACT_AT_50K), api_key="mock_api_key" + ) + litellm.responses( + model=_MODEL, input=_replay(first.output), context_management=list(_COMPACT_AT_50K), api_key="mock_api_key" + ) + + assert [item.type for item in first.output] == ["compaction", "message"] + assert _sent_assistant_content(route) == [_COMPACTION_BLOCK, {"type": "text", "text": "The."}] + + +def test_streamed_responses_over_anthropic_emit_the_compaction_item_and_send_it_back( + respx_mock: respx.MockRouter, +) -> None: + route: Final = respx_mock.post(_ANTHROPIC_MESSAGES_URL).mock( + side_effect=[ + httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + content=_anthropic_sse( + { + "type": "message_start", + "message": { + "id": "msg_compacted_stream", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5-5", + "content": [], + "usage": {"input_tokens": 9, "output_tokens": 1}, + }, + }, + { + "type": "content_block_start", + "index": 0, + "content_block": {"type": "compaction", "content": None}, + }, + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "compaction_delta", "content": _COMPACTION_BLOCK["content"]}, + }, + {"type": "content_block_stop", "index": 0}, + {"type": "content_block_start", "index": 1, "content_block": {"type": "text", "text": ""}}, + {"type": "content_block_delta", "index": 1, "delta": {"type": "text_delta", "text": "The."}}, + {"type": "content_block_stop", "index": 1}, + {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 4}}, + {"type": "message_stop"}, + ), + ), + _anthropic_message({"type": "text", "text": "dog."}), + ] + ) + + events: Final = tuple( + litellm.responses( + model=_MODEL, + input=list(_FIRST_INPUT), + context_management=list(_COMPACT_AT_50K), + stream=True, + api_key="mock_api_key", + ) + ) + item_events: Final = tuple( + event + for event in events + if event.type in (ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE) + ) + completed: Final = events[-1].response + litellm.responses( + model=_MODEL, input=_replay(completed.output), context_management=list(_COMPACT_AT_50K), api_key="mock_api_key" + ) + + assert [(event.type, event.output_index, event.item.type) for event in item_events] == [ + (ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, 0, "compaction"), + (ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, 1, "message"), + (ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, 0, "compaction"), + (ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, 1, "message"), + ] + assert item_events[0].item.encrypted_content == "" + assert item_events[0].item.id == item_events[2].item.id + assert [(item.type, getattr(item, "encrypted_content", None)) for item in completed.output[:1]] == [ + ("compaction", item_events[2].item.encrypted_content) + ] + assert completed.output[0].id == item_events[2].item.id + assert _sent_assistant_content(route) == [_COMPACTION_BLOCK, {"type": "text", "text": "The."}] + + +def test_streamed_responses_over_anthropic_answer_after_a_failed_compaction(respx_mock: respx.MockRouter) -> None: + respx_mock.post(_ANTHROPIC_MESSAGES_URL).mock( + return_value=httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + content=_anthropic_sse( + { + "type": "message_start", + "message": { + "id": "msg_failed_compaction", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5-5", + "content": [], + "usage": {"input_tokens": 9, "output_tokens": 1}, + }, + }, + {"type": "content_block_start", "index": 0, "content_block": {"type": "compaction", "content": None}}, + {"type": "content_block_delta", "index": 0, "delta": {"type": "compaction_delta", "content": None}}, + {"type": "content_block_stop", "index": 0}, + {"type": "content_block_start", "index": 1, "content_block": {"type": "text", "text": ""}}, + {"type": "content_block_delta", "index": 1, "delta": {"type": "text_delta", "text": "The."}}, + {"type": "content_block_stop", "index": 1}, + {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 4}}, + {"type": "message_stop"}, + ), + ) + ) + + events: Final = tuple( + litellm.responses( + model=_MODEL, + input=list(_FIRST_INPUT), + context_management=list(_COMPACT_AT_50K), + stream=True, + api_key="mock_api_key", + ) + ) + + assert events[-1].type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED + assert [item.type for item in events[-1].response.output] == ["message"] + assert events[-1].response.output[0].content[0].text == "The." + + +@pytest.mark.asyncio +async def test_async_streamed_compaction_item_is_done_before_the_reasoning_that_follows_it() -> None: + anthropic: Final = httpx.MockTransport( + lambda _request: httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + content=_anthropic_sse( + { + "type": "message_start", + "message": { + "id": "msg_compacted_thinking", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5-5", + "content": [], + "usage": {"input_tokens": 9, "output_tokens": 1}, + }, + }, + {"type": "content_block_start", "index": 0, "content_block": {"type": "compaction", "content": None}}, + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "compaction_delta", "content": _COMPACTION_BLOCK["content"]}, + }, + {"type": "content_block_stop", "index": 0}, + {"type": "content_block_start", "index": 1, "content_block": {"type": "thinking", "thinking": ""}}, + {"type": "content_block_delta", "index": 1, "delta": {"type": "thinking_delta", "thinking": "Recall."}}, + {"type": "content_block_delta", "index": 1, "delta": {"type": "signature_delta", "signature": "sig"}}, + {"type": "content_block_stop", "index": 1}, + {"type": "content_block_start", "index": 2, "content_block": {"type": "text", "text": ""}}, + {"type": "content_block_delta", "index": 2, "delta": {"type": "text_delta", "text": "The."}}, + {"type": "content_block_stop", "index": 2}, + {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 4}}, + {"type": "message_stop"}, + ), + ) + ) + + stream: Final = await litellm.aresponses( + model=_MODEL, + input=list(_FIRST_INPUT), + context_management=list(_COMPACT_AT_50K), + stream=True, + api_key="mock_api_key", + client=AsyncHTTPHandler(transport=anthropic), + ) + events: Final = [event async for event in stream] + done_items: Final = [ + (event.output_index, event.item.type) + for event in events + if event.type == ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE + ] + reasoning_done: Final = next( + index + for index, event in enumerate(events) + if event.type == ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE and event.item.type == "reasoning" + ) + sequence_numbers: Final = [ + event.sequence_number for event in events[: reasoning_done + 1] if getattr(event, "sequence_number", None) + ] + + assert done_items == [(0, "compaction"), (1, "reasoning"), (2, "message")] + assert [item.type for item in events[-1].response.output] == ["compaction", "reasoning", "message"] + assert sequence_numbers == sorted(set(sequence_numbers)) + + +def _compaction_stream_events(*summaries: str) -> tuple[Mapping[str, object], ...]: + return tuple( + chain.from_iterable( + ( + { + "type": "content_block_start", + "index": 2 * index, + "content_block": {"type": "compaction", "content": None}, + }, + { + "type": "content_block_delta", + "index": 2 * index, + "delta": {"type": "compaction_delta", "content": summary}, + }, + {"type": "content_block_stop", "index": 2 * index}, + {"type": "content_block_start", "index": 2 * index + 1, "content_block": {"type": "text", "text": ""}}, + {"type": "content_block_delta", "index": 2 * index + 1, "delta": {"type": "text_delta", "text": "Ok."}}, + {"type": "content_block_stop", "index": 2 * index + 1}, + ) + for index, summary in enumerate(summaries) + ) + ) + + +def test_streamed_compaction_item_done_carries_the_same_newest_block_as_the_completed_response( + respx_mock: respx.MockRouter, +) -> None: + respx_mock.post(_ANTHROPIC_MESSAGES_URL).mock( + return_value=httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + content=_anthropic_sse( + { + "type": "message_start", + "message": { + "id": "msg_compacted_twice", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5-5", + "content": [], + "usage": {"input_tokens": 9, "output_tokens": 1}, + }, + }, + *_compaction_stream_events("First summary.", "Second summary."), + {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 4}}, + {"type": "message_stop"}, + ), + ) + ) + + events: Final = tuple( + litellm.responses( + model=_MODEL, + input=list(_FIRST_INPUT), + context_management=list(_COMPACT_AT_50K), + stream=True, + api_key="mock_api_key", + ) + ) + streamed: Final = tuple( + (event.item.id, json.loads(event.item.encrypted_content)["content"]) + for event in events + if event.type == ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE and event.item.type == "compaction" + ) + completed: Final = tuple( + (item.id, json.loads(item.encrypted_content)["content"]) + for item in events[-1].response.output + if item.type == "compaction" + ) + + assert streamed == completed + assert [summary for _, summary in completed] == ["Second summary."] + + +def test_responses_over_anthropic_send_back_a_compaction_item_replayed_as_sdk_output( + respx_mock: respx.MockRouter, +) -> None: + route: Final = respx_mock.post(_ANTHROPIC_MESSAGES_URL).mock( + side_effect=[ + _anthropic_message(_COMPACTION_BLOCK, {"type": "text", "text": "The."}), + _anthropic_message({"type": "text", "text": "dog."}), + ] + ) + + first: Final = litellm.responses( + model=_MODEL, input=list(_FIRST_INPUT), context_management=list(_COMPACT_AT_50K), api_key="mock_api_key" + ) + litellm.responses( + model=_MODEL, + input=[*_FIRST_INPUT, *first.output, *_NEXT_INPUT], + context_management=list(_COMPACT_AT_50K), + api_key="mock_api_key", + ) + + assert _sent_assistant_content(route)[0] == _COMPACTION_BLOCK + + +@pytest.mark.asyncio +async def test_router_with_encrypted_content_affinity_sends_the_compaction_block_back( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + route: Final = respx_mock.post(_ANTHROPIC_MESSAGES_URL).mock( + side_effect=[ + _anthropic_message(_COMPACTION_BLOCK, {"type": "text", "text": "The."}), + _anthropic_message({"type": "text", "text": "dog."}), + ] + ) + router: Final = litellm.Router( + model_list=[ + { + "model_name": "claude", + "litellm_params": {"model": _MODEL, "api_key": "mock_api_key"}, + "model_info": {"id": "deployment-a"}, + } + ], + optional_pre_call_checks=["encrypted_content_affinity"], + ) + + first: Final = await router.aresponses( + model="claude", input=list(_FIRST_INPUT), context_management=list(_COMPACT_AT_50K) + ) + await router.aresponses(model="claude", input=_replay(first.output), context_management=list(_COMPACT_AT_50K)) + + assert first.output[0].encrypted_content.startswith("litellm_enc:") + assert _sent_assistant_content(route) == [_COMPACTION_BLOCK, {"type": "text", "text": "The."}] + assert [key for key in json.loads(route.calls.last.request.content) if key.startswith("_")] == [] + + +_DOCUMENT_BEFORE_COMPACTION: Final = "The design document pasted before the compaction." +_TURN_BEFORE_COMPACTION: Final = { + "request_id": "chatcmpl-before-compaction", + "call_type": "aresponses", + "session_id": "session-compacted", + "proxy_server_request": { + "model": _MODEL, + "input": [ + {"role": "developer", "content": "Use metric units."}, + {"role": "user", "content": _DOCUMENT_BEFORE_COMPACTION}, + ], + }, + "response": { + "id": "chatcmpl-before-compaction", + "object": "chat.completion", + "created": 0, + "model": _MODEL, + "choices": [{"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": "Read it."}}], + }, +} + + +class _RecordingRawAnthropicMessages: + def __init__(self) -> None: + self.body: str = "" + + def __call__(self, request: httpx.Request) -> httpx.Response: + self.body = request.content.decode() + return _anthropic_message({"type": "text", "text": "dog."}) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("context_management", [None, list(_COMPACT_AT_50K)], ids=["summary_text", "native_block"]) +@pytest.mark.parametrize( + "compaction_item", + [ + pytest.param( + {"type": "compaction", "id": "cmp_caller", "encrypted_content": json.dumps(_COMPACTION_BLOCK)}, + id="param_dict", + ), + pytest.param( + ResponseCompactionItem(type="compaction", id="cmp_caller", encrypted_content=json.dumps(_COMPACTION_BLOCK)), + id="sdk_output_item", + ), + ], +) +async def test_previous_response_id_with_a_compaction_item_in_input_replays_only_the_stored_instructions( + context_management: list[Mapping[str, object]] | None, + compaction_item: Mapping[str, object] | ResponseCompactionItem, +) -> None: + anthropic: Final = _RecordingRawAnthropicMessages() + + with patch("litellm.proxy.proxy_server.prisma_client", _FakePrismaClient([_TURN_BEFORE_COMPACTION])): + await litellm.aresponses( + model=_MODEL, + previous_response_id="chatcmpl-before-compaction", + input=[compaction_item, *_NEXT_INPUT], + context_management=context_management, + api_key="sk-ant-fake", + client=AsyncHTTPHandler(transport=httpx.MockTransport(anthropic)), + ) + + assert _DOCUMENT_BEFORE_COMPACTION not in anthropic.body + assert "Read it." not in anthropic.body + assert "Use metric units." in anthropic.body + assert _COMPACTION_BLOCK["content"] in anthropic.body + assert _NEXT_INPUT[0]["content"] in anthropic.body + + +class _RecordingCustomLLM(CustomLLM): + def __init__(self) -> None: + super().__init__() + self.sent_messages: list[object] = [] + + def completion(self, *args: object, **kwargs: object) -> ModelResponse: + self.sent_messages.extend(cast(list[object], kwargs["messages"])) + return litellm.completion( + model="gpt-5.6", messages=[{"role": "user", "content": "Hi."}], mock_response="dog." + ) + + +def test_responses_over_a_custom_provider_send_a_compaction_item_as_summary_text( + monkeypatch: pytest.MonkeyPatch, +) -> None: + custom_llm: Final = _RecordingCustomLLM() + monkeypatch.setattr(litellm, "provider_list", list(litellm.provider_list)) + monkeypatch.setattr(litellm, "_custom_providers", list(litellm._custom_providers)) + monkeypatch.setattr( + litellm, "custom_provider_map", [{"provider": "compaction-test-llm", "custom_handler": custom_llm}] + ) + compaction_item: Final = { + "type": "compaction", + "id": "cmp_earlier", + "encrypted_content": json.dumps(_COMPACTION_BLOCK), + } + + litellm.responses( + model="compaction-test-llm/my-model", + input=[*_FIRST_INPUT, {"role": "assistant", "content": "The."}, compaction_item, *_NEXT_INPUT], + context_management=list(_COMPACT_AT_50K), + drop_params=True, + ) + + assert custom_llm.sent_messages == [ + {"role": "user", "content": _COMPACTION_BLOCK["content"]}, + *_NEXT_INPUT, + ] + + +_WEATHER_CALL: Final = { + "id": "toolu_weather", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city": "Paris"}'}, +} +_TOOL_TURN_BEFORE_COMPACTION: Final = { + **_TURN_BEFORE_COMPACTION, + "request_id": "chatcmpl-tool-turn", + "response": { + "id": "chatcmpl-tool-turn", + "object": "chat.completion", + "created": 0, + "model": _MODEL, + "choices": [ + { + "index": 0, + "finish_reason": "tool_calls", + "message": {"role": "assistant", "content": None, "tool_calls": [_WEATHER_CALL]}, + } + ], + }, +} + + +def _content_blocks(message: Mapping[str, object]) -> list[Mapping[str, object]]: + content: Final = message["content"] + return content if isinstance(content, list) else [{"type": "text", "text": content}] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("context_management", [None, list(_COMPACT_AT_50K)], ids=["summary_text", "native_block"]) +async def test_previous_response_id_with_a_compaction_item_keeps_the_stored_call_its_tool_output_answers( + context_management: list[Mapping[str, object]] | None, +) -> None: + anthropic: Final = _RecordingRawAnthropicMessages() + compaction_item: Final = { + "type": "compaction", + "id": "cmp_caller", + "encrypted_content": json.dumps(_COMPACTION_BLOCK), + } + tool_output: Final = {"type": "function_call_output", "call_id": "toolu_weather", "output": "18C and sunny"} + + with patch("litellm.proxy.proxy_server.prisma_client", _FakePrismaClient([_TOOL_TURN_BEFORE_COMPACTION])): + await litellm.aresponses( + model=_MODEL, + previous_response_id="chatcmpl-tool-turn", + input=[compaction_item, tool_output], + context_management=context_management, + api_key="sk-ant-fake", + client=AsyncHTTPHandler(transport=httpx.MockTransport(anthropic)), + ) + + messages: Final = json.loads(anthropic.body)["messages"] + call_turn: Final = next( + index + for index, message in enumerate(messages) + if any(block["type"] == "tool_use" for block in _content_blocks(message)) + ) + tool_use: Final = next(block for block in _content_blocks(messages[call_turn]) if block["type"] == "tool_use") + assert (tool_use["id"], tool_use["input"]) == ("toolu_weather", {"city": "Paris"}) + assert _content_blocks(messages[call_turn + 1])[0]["type"] == "tool_result" + assert _content_blocks(messages[call_turn + 1])[0]["tool_use_id"] == "toolu_weather" + assert _COMPACTION_BLOCK["content"] in json.dumps(messages[: call_turn + 1]) + assert _DOCUMENT_BEFORE_COMPACTION not in anthropic.body + + +_OSAKA_CALL: Final = {**_WEATHER_CALL, "function": {"name": "get_weather", "arguments": '{"city": "Osaka"}'}} +_TOOL_TURN_REUSING_THE_CALL_ID: Final = { + **_TOOL_TURN_BEFORE_COMPACTION, + "request_id": "chatcmpl-osaka-turn", + "proxy_server_request": { + "model": _MODEL, + "input": [{"type": "function_call_output", "call_id": "toolu_weather", "output": "18C and sunny"}], + }, + "response": { + "id": "chatcmpl-osaka-turn", + "object": "chat.completion", + "created": 0, + "model": _MODEL, + "choices": [ + { + "index": 0, + "finish_reason": "tool_calls", + "message": {"role": "assistant", "content": None, "tool_calls": [_OSAKA_CALL]}, + } + ], + }, +} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("context_management", [None, list(_COMPACT_AT_50K)], ids=["summary_text", "native_block"]) +async def test_previous_response_id_with_a_compaction_item_restores_only_the_newest_stored_call_for_a_reused_call_id( + context_management: list[Mapping[str, object]] | None, +) -> None: + anthropic: Final = _RecordingRawAnthropicMessages() + compaction_item: Final = { + "type": "compaction", + "id": "cmp_caller", + "encrypted_content": json.dumps(_COMPACTION_BLOCK), + } + tool_output: Final = {"type": "function_call_output", "call_id": "toolu_weather", "output": "12C and rain"} + spend_logs: Final = [_TOOL_TURN_BEFORE_COMPACTION, _TOOL_TURN_REUSING_THE_CALL_ID] + + with patch("litellm.proxy.proxy_server.prisma_client", _FakePrismaClient(spend_logs)): + await litellm.aresponses( + model=_MODEL, + previous_response_id="chatcmpl-osaka-turn", + input=[compaction_item, tool_output], + context_management=context_management, + api_key="sk-ant-fake", + client=AsyncHTTPHandler(transport=httpx.MockTransport(anthropic)), + ) + + messages: Final = json.loads(anthropic.body)["messages"] + blocks: Final = list( + chain.from_iterable( + ((index, block) for block in _content_blocks(message)) for index, message in enumerate(messages) + ) + ) + tool_uses: Final = [(index, block["id"], block["input"]) for index, block in blocks if block["type"] == "tool_use"] + tool_results: Final = [(index, block["tool_use_id"]) for index, block in blocks if block["type"] == "tool_result"] + assert [(tool_id, tool_input) for _, tool_id, tool_input in tool_uses] == [("toolu_weather", {"city": "Osaka"})] + assert tool_results == [(tool_uses[0][0] + 1, "toolu_weather")] + assert "12C and rain" in json.dumps(messages[tool_uses[0][0] + 1]) + + +def test_responses_over_anthropic_send_a_compaction_item_as_summary_text_when_the_request_does_not_compact( + respx_mock: respx.MockRouter, +) -> None: + route: Final = respx_mock.post(_ANTHROPIC_MESSAGES_URL).mock( + return_value=_anthropic_message({"type": "text", "text": "dog."}) + ) + compaction_item: Final = { + "type": "compaction", + "id": "cmp_earlier", + "encrypted_content": json.dumps(_COMPACTION_BLOCK), + } + + litellm.responses( + model=_MODEL, + input=[*_FIRST_INPUT, {"role": "assistant", "content": "The."}, compaction_item, *_NEXT_INPUT], + api_key="mock_api_key", + ) + + sent: Final = json.loads(route.calls.last.request.content) + assert sent["messages"] == [ + { + "role": "user", + "content": [ + {"type": "text", "text": _COMPACTION_BLOCK["content"]}, + {"type": "text", "text": _NEXT_INPUT[0]["content"]}, + ], + } + ] diff --git a/tests/unit/responses/litellm_completion_transformation/test_compaction_items.py b/tests/unit/responses/litellm_completion_transformation/test_compaction_items.py new file mode 100644 index 00000000000..e0ef3fbb552 --- /dev/null +++ b/tests/unit/responses/litellm_completion_transformation/test_compaction_items.py @@ -0,0 +1,87 @@ +import json +from typing import Final + +import pytest + +from litellm.responses.litellm_completion_transformation.compaction_items import ( + compaction_output_item, + holds_compaction_summary, + response_holds_compaction_block, + with_compaction_summary, +) +from litellm.responses.utils import ResponsesAPIRequestUtils + +_BLOCK: Final = {"type": "compaction", "content": "Turn one planned Tokyo."} +_ITEM: Final = {"type": "compaction", "id": "cmp_1", "encrypted_content": json.dumps(_BLOCK)} + + +@pytest.mark.parametrize( + "newer_entry", + [ + pytest.param("not a block", id="not_a_mapping"), + pytest.param({"type": "compaction", "content": None}, id="summary_not_streamed_yet"), + pytest.param({"type": "compaction", "content": ""}, id="empty_summary"), + pytest.param({"type": "text", "content": "Hi."}, id="not_a_compaction"), + ], +) +def test_compaction_output_item_carries_the_newest_block_that_holds_a_summary(newer_entry: object) -> None: + item: Final = compaction_output_item({"compaction_blocks": [_BLOCK, newer_entry]}) + + assert item is not None + assert json.loads(item.encrypted_content) == _BLOCK + + +@pytest.mark.parametrize( + ("response", "holds"), + [ + pytest.param( + {"choices": [{"message": {"role": "assistant", "provider_specific_fields": {"compaction_blocks": [_BLOCK]}}}]}, + True, + id="chat_answer_compacted", + ), + pytest.param({"choices": [{"message": {"role": "assistant", "content": "Hi."}}]}, False, id="chat_answer"), + pytest.param({"object": "response", "output": [_ITEM]}, True, id="responses_output_compacted"), + pytest.param( + {"object": "response", "output": [{"type": "compaction", "id": "cmp_openai", "encrypted_content": "gAAAA"}]}, + False, + id="responses_output_with_an_opaque_compaction", + ), + pytest.param({"choices": "not a list"}, False, id="malformed_choices"), + pytest.param(None, False, id="no_stored_response"), + ], +) +def test_response_holds_compaction_block_reads_chat_and_responses_shaped_rows(response: object, holds: bool) -> None: + assert response_holds_compaction_block(response) is holds + + +@pytest.mark.parametrize( + ("encrypted_content", "holds"), + [ + pytest.param(json.dumps(_BLOCK), True, id="litellm_block"), + pytest.param( + ResponsesAPIRequestUtils.wrap_encrypted_content_with_model_id(json.dumps(_BLOCK), "deployment-1"), + True, + id="affinity_wrapped_block", + ), + pytest.param("gAAAAABopaque", False, id="openai_token"), + pytest.param(None, False, id="missing"), + ], +) +def test_only_a_compaction_item_litellm_can_decode_holds_a_summary(encrypted_content: object, holds: bool) -> None: + assert holds_compaction_summary([{"type": "compaction", "id": "cmp_1", "encrypted_content": encrypted_content}]) is ( + holds + ) + + +def test_with_compaction_summary_rewrites_the_summary_and_keeps_the_affinity_wrap() -> None: + wrapped: Final = ResponsesAPIRequestUtils.wrap_encrypted_content_with_model_id(json.dumps(_BLOCK), "deployment-1") + + rewritten: Final = with_compaction_summary(wrapped, "[REDACTED]") + model_id, original = ResponsesAPIRequestUtils.unwrap_encrypted_content_with_model_id(rewritten) + + assert model_id == "deployment-1" + assert json.loads(original) == {**_BLOCK, "content": "[REDACTED]"} + + +def test_with_compaction_summary_leaves_a_token_it_cannot_read_unchanged() -> None: + assert with_compaction_summary("gAAAAABopaque", "[REDACTED]") == "gAAAAABopaque" diff --git a/tests/unit/responses/litellm_completion_transformation/test_count_request.py b/tests/unit/responses/litellm_completion_transformation/test_count_request.py index 464bdb1a367..0477ae15b8e 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_count_request.py +++ b/tests/unit/responses/litellm_completion_transformation/test_count_request.py @@ -10,6 +10,7 @@ from litellm.responses.litellm_completion_transformation.count_request import ( ChatCountRequest, chat_count_request, ) +from litellm.responses.litellm_completion_transformation.transformation import ChatCompletionSession _MARK: Final[dict[str, JsonValue]] = {"type": "ephemeral"} _MARKED_RESPONSES_REQUEST: Final[dict[str, JsonValue]] = { @@ -125,3 +126,77 @@ async def test_claude_count_request_leaves_out_tools_the_deployment_drops() -> N assert isinstance(dropped, ChatCountRequest) assert dropped.tools is None assert dropped == await _claude_count({"input": [_QUESTION]}, policy) + + +_SUMMARY: Final = "Earlier turns covered the Tokyo trip." +_COMPACTED_INPUT: Final[list[JsonValue]] = [ + {"type": "compaction", "id": "cmp_1", "encrypted_content": json.dumps({"type": "compaction", "content": _SUMMARY})}, + {"type": "message", "role": "user", "content": "And Osaka?"}, +] +_COMPACTION_REQUESTS: Final = ( + pytest.param({}, id="summary_route"), + pytest.param({"context_management": [{"type": "compaction", "compact_threshold": 50000}]}, id="native_route"), +) + + +async def _stored_turns_the_compaction_replaces(previous_response_id: str) -> ChatCompletionSession: + return ChatCompletionSession( + messages=[ + {"role": "user", "content": "Weather in SF?"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": "call_0", "type": "function", "function": {"name": "get_weather", "arguments": "{}"}} + ], + }, + ], + litellm_session_id=None, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("compaction_request", _COMPACTION_REQUESTS) +async def test_claude_count_request_counts_a_compaction_item_as_its_summary_text( + compaction_request: dict[str, JsonValue], +) -> None: + request: Final = await chat_count_request( + model="claude-opus-5-5", + custom_llm_provider="anthropic", + api_base=None, + responses_request={"model": "claude-opus-5-5", "input": _COMPACTED_INPUT, **compaction_request}, + policy=DropPolicy(drop_params=False, additional_drop_params=frozenset()), + ) + + assert isinstance(request, ChatCountRequest) + assert [(message.get("role"), message.get("content")) for message in request.messages] == [ + ("user", _SUMMARY), + ("user", "And Osaka?"), + ] + assert "compaction_blocks" not in json.dumps(request.messages, default=str) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("compaction_request", _COMPACTION_REQUESTS) +async def test_claude_count_request_leaves_out_the_stored_turns_a_compaction_item_replaces( + compaction_request: dict[str, JsonValue], +) -> None: + request: Final = await chat_count_request( + model="claude-opus-5-5", + custom_llm_provider="anthropic", + api_base=None, + responses_request={ + "model": "claude-opus-5-5", + "previous_response_id": "resp_before_the_compaction", + "input": _COMPACTED_INPUT, + **compaction_request, + }, + policy=DropPolicy(drop_params=False, additional_drop_params=frozenset()), + session_lookup=_stored_turns_the_compaction_replaces, + ) + + assert isinstance(request, ChatCountRequest) + counted: Final = json.dumps(request.messages, default=str) + assert _SUMMARY in counted and "And Osaka?" in counted + assert "Weather in SF?" not in counted + assert "call_0" not in counted 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 f34ca7d0ff8..30e13442c92 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 @@ -24,6 +24,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import ( _bedrock_converse_messages_pt, anthropic_messages_pt, ) +from litellm.responses.litellm_completion_transformation.reasoning_items import encode_thinking_blocks from litellm.responses.litellm_completion_transformation.transformation import ( TOOL_CALLS_CACHE, ChatCompletionSession, @@ -6502,6 +6503,363 @@ def test_transform_request_reads_service_tier_from_responses_api_request() -> No assert result["service_tier"] == "priority" +_COMPACTION_BLOCK: Final = {"type": "compaction", "content": "Earlier turns covered the Tokyo trip."} + + +_COMPACTION_REQUESTED: Final = ResponsesAPIOptionalRequestParams( + context_management=[{"type": "compaction", "compact_threshold": 50000}] +) +_SUMMARY_AS_TEXT: Final = {"role": "user", "content": _COMPACTION_BLOCK["content"]} + + +def _anthropic_form_context_management(edit: Mapping[str, object]) -> ResponsesAPIOptionalRequestParams: + return cast( # cast-ok: Anthropic's native {"edits": [...]} form, which the param type does not declare + ResponsesAPIOptionalRequestParams, {"context_management": {"edits": [edit]}} + ) + + +@pytest.mark.parametrize( + ("model", "responses_api_request", "replayed_compaction"), + [ + ( + "anthropic/claude-sonnet-5-5", + _COMPACTION_REQUESTED, + { + "role": "assistant", + "content": None, + "provider_specific_fields": {"compaction_blocks": [_COMPACTION_BLOCK]}, + }, + ), + ( + "anthropic/claude-sonnet-5-5", + _anthropic_form_context_management({"type": "compact_20260112"}), + { + "role": "assistant", + "content": None, + "provider_specific_fields": {"compaction_blocks": [_COMPACTION_BLOCK]}, + }, + ), + ("anthropic/claude-sonnet-5-5", ResponsesAPIOptionalRequestParams(), _SUMMARY_AS_TEXT), + ( + "anthropic/claude-sonnet-5-5", + _anthropic_form_context_management({"type": "clear_tool_uses_20250919"}), + _SUMMARY_AS_TEXT, + ), + ("gemini/gemini-3.8-flash", _COMPACTION_REQUESTED, _SUMMARY_AS_TEXT), + ("bedrock/claude_platform/claude-sonnet-5-5", _COMPACTION_REQUESTED, _SUMMARY_AS_TEXT), + ], +) +def test_compaction_input_item_replays_as_a_compaction_block_only_when_anthropic_compacts( + model: str, responses_api_request: ResponsesAPIOptionalRequestParams, replayed_compaction: Mapping[str, object] +) -> None: + result: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( + model=model, + input=cast( + ResponseInputParam, + [ + {"type": "compaction", "id": "cmp_openai", "encrypted_content": "gAAAAABopaque"}, + {"type": "compaction", "id": "cmp_litellm", "encrypted_content": json.dumps(_COMPACTION_BLOCK)}, + {"role": "user", "content": "And Osaka?"}, + ], + ), + responses_api_request=responses_api_request, + ) + + assert [dict(message) for message in result["messages"]] == [ + replayed_compaction, + {"role": "user", "content": "And Osaka?"}, + ] + + +def _replayed_compacted_answer(model: str, context_management: object) -> list[tuple[object, object, object]]: + async def session_lookup(previous_response_id: str) -> ChatCompletionSession: + return ChatCompletionSession( + messages=[ + Message( + role="assistant", + content="Day one, Asakusa.", + provider_specific_fields={"compaction_blocks": [_COMPACTION_BLOCK]}, + ), + {"role": "user", "content": "Day two?"}, + Message(role="assistant", content="Ueno."), + ], + litellm_session_id=None, + ) + + request: Final = asyncio.run( + LiteLLMCompletionResponsesConfig.async_responses_api_session_handler( + previous_response_id="resp_compacted", + litellm_completion_request={ + "model": model, + "messages": [{"role": "user", "content": "And Osaka?"}], + **({"context_management": context_management} if context_management else {}), + }, + instructions=None, + session_lookup=session_lookup, + ) + ) + return [ + ( + message.get("role"), + message.get("content"), + (message.get("provider_specific_fields") or {}).get("compaction_blocks"), + ) + for message in request["messages"] + ] + + +_TURNS_AFTER_THE_COMPACTED_ANSWER: Final = [ + ("user", "Day two?", None), + ("assistant", "Ueno.", None), + ("user", "And Osaka?", None), +] +_ANSWER_AFTER_SUMMARY: Final = [ + ("user", _COMPACTION_BLOCK["content"], None), + ("assistant", "Day one, Asakusa.", None), + *_TURNS_AFTER_THE_COMPACTED_ANSWER, +] + + +@pytest.mark.parametrize( + ("model", "context_management", "replayed"), + [ + ( + "anthropic/claude-sonnet-5-5", + _COMPACTION_REQUESTED["context_management"], + [("assistant", "Day one, Asakusa.", [_COMPACTION_BLOCK]), *_TURNS_AFTER_THE_COMPACTED_ANSWER], + ), + ("anthropic/claude-sonnet-5-5", None, _ANSWER_AFTER_SUMMARY), + ( + "bedrock/claude_platform/claude-sonnet-5-5", + _COMPACTION_REQUESTED["context_management"], + _ANSWER_AFTER_SUMMARY, + ), + ("gpt-5.6", _COMPACTION_REQUESTED["context_management"], _ANSWER_AFTER_SUMMARY), + ], +) +def test_stored_compacted_answer_replays_its_block_only_when_anthropic_compacts( + model: str, context_management: object, replayed: list[tuple[object, object, object]] +) -> None: + assert _replayed_compacted_answer(model, context_management) == replayed + + +def test_summary_replay_drops_the_history_its_compaction_item_replaces() -> None: + result: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( + model="gemini/gemini-3.8-flash", + input=cast( + ResponseInputParam, + [ + {"role": "system", "content": "Answer briefly."}, + {"role": "user", "content": "Plan a Tokyo trip."}, + {"role": "assistant", "content": "Day one, Asakusa."}, + {"type": "compaction", "id": "cmp_litellm", "encrypted_content": json.dumps(_COMPACTION_BLOCK)}, + {"role": "user", "content": "And Osaka?"}, + ], + ), + responses_api_request=_COMPACTION_REQUESTED, + ) + + assert [dict(message) for message in result["messages"]] == [ + {"role": "system", "content": "Answer briefly."}, + _SUMMARY_AS_TEXT, + {"role": "user", "content": "And Osaka?"}, + ] + + + +@pytest.mark.parametrize( + ("call", "output"), + [ + pytest.param( + {"type": "function_call", "call_id": "call_weather", "name": "get_weather", "arguments": "{}"}, + {"type": "function_call_output", "call_id": "call_weather", "output": "Rain."}, + id="function_call", + ), + pytest.param( + {"type": "custom_tool_call", "call_id": "call_patch", "name": "apply_patch", "input": "*** Begin Patch"}, + {"type": "custom_tool_call_output", "call_id": "call_patch", "output": "Done."}, + id="custom_tool_call", + ), + ], +) +def test_summary_replay_keeps_a_call_its_compaction_item_replaces_when_a_later_output_answers_it( + call: Mapping[str, str], output: Mapping[str, str] +) -> None: + messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( + model="gemini/gemini-3.8-flash", + input=cast( + ResponseInputParam, + [ + {"role": "system", "content": "Answer briefly."}, + {"role": "user", "content": "Plan a Tokyo trip."}, + call, + {"type": "compaction", "id": "cmp_litellm", "encrypted_content": json.dumps(_COMPACTION_BLOCK)}, + output, + ], + ), + responses_api_request=_COMPACTION_REQUESTED, + )["messages"] + + assert [message["role"] for message in messages] == ["system", "user", "assistant", "tool"] + assert dict(messages[1]) == _SUMMARY_AS_TEXT + assert [tool_call["id"] for tool_call in messages[2]["tool_calls"]] == [output["call_id"]] + assert messages[3]["tool_call_id"] == output["call_id"] + + +_ANSWERED_THINKING: Final = {"type": "thinking", "thinking": "Need Osaka weather.", "signature": "sig_osaka"} +_CHECKING_OSAKA: Final = { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "Checking Osaka."}], +} + + +def _reasoning_item(item_id: str, thinking: Mapping[str, str]) -> Mapping[str, object]: + return {"type": "reasoning", "id": item_id, "summary": [], "encrypted_content": encode_thinking_blocks([thinking])} + + +def test_summary_replay_keeps_only_the_reasoning_of_the_turn_whose_call_a_later_output_answers() -> None: + tokyo_thinking: Final = {"type": "thinking", "thinking": "Need Tokyo weather.", "signature": "sig_tokyo"} + messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( + model="anthropic/claude-sonnet-5-5", + input=cast( + ResponseInputParam, + [ + {"role": "user", "content": "Plan a Tokyo trip."}, + _reasoning_item("rs_tokyo", tokyo_thinking), + {"type": "function_call", "call_id": "call_tokyo", "name": "get_weather", "arguments": "{}"}, + {"type": "function_call_output", "call_id": "call_tokyo", "output": "Sun."}, + _reasoning_item("rs_osaka", _ANSWERED_THINKING), + _CHECKING_OSAKA, + {"type": "function_call", "call_id": "call_osaka", "name": "get_weather", "arguments": "{}"}, + {"type": "function_call", "call_id": "call_kyoto", "name": "get_weather", "arguments": "{}"}, + {"type": "function_call_output", "call_id": "call_kyoto", "output": "Clouds."}, + {"type": "compaction", "id": "cmp_litellm", "encrypted_content": json.dumps(_COMPACTION_BLOCK)}, + {"type": "function_call_output", "call_id": "call_osaka", "output": "Rain."}, + ], + ), + responses_api_request=ResponsesAPIOptionalRequestParams(), + )["messages"] + + assert [message["role"] for message in messages] == ["user", "assistant", "tool"] + assert dict(messages[0]) == _SUMMARY_AS_TEXT + assert messages[1]["thinking_blocks"] == [_ANSWERED_THINKING] + assert [tool_call["id"] for tool_call in messages[1]["tool_calls"]] == ["call_osaka"] + + +def test_native_compaction_replay_moves_a_call_its_compaction_item_splits_from_its_output_after_the_block() -> None: + messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( + model="anthropic/claude-sonnet-5-5", + input=cast( + ResponseInputParam, + [ + {"role": "user", "content": "Plan a Tokyo trip."}, + _reasoning_item("rs_osaka", _ANSWERED_THINKING), + _CHECKING_OSAKA, + {"type": "function_call", "call_id": "call_osaka", "name": "get_weather", "arguments": "{}"}, + {"type": "compaction", "id": "cmp_litellm", "encrypted_content": json.dumps(_COMPACTION_BLOCK)}, + {"type": "function_call_output", "call_id": "call_osaka", "output": "Rain."}, + ], + ), + responses_api_request=_COMPACTION_REQUESTED, + )["messages"] + + assert [message["role"] for message in messages] == ["user", "assistant", "assistant", "tool"] + assert messages[0]["content"] == "Plan a Tokyo trip." + assert messages[1]["provider_specific_fields"] == {"compaction_blocks": [_COMPACTION_BLOCK]} + assert messages[2]["thinking_blocks"] == [_ANSWERED_THINKING] + assert [tool_call["id"] for tool_call in messages[2]["tool_calls"]] == ["call_osaka"] + assert messages[3]["tool_call_id"] == "call_osaka" + + +@pytest.mark.parametrize( + "responses_api_request", + [ + pytest.param(_COMPACTION_REQUESTED, id="native_compaction_block"), + pytest.param(ResponsesAPIOptionalRequestParams(), id="summary_as_text"), + ], +) +def test_compaction_replay_invents_no_call_for_an_output_whose_call_was_never_stored( + responses_api_request: ResponsesAPIOptionalRequestParams, +) -> None: + async def no_stored_turns(previous_response_id: str) -> ChatCompletionSession: + return ChatCompletionSession(messages=[], litellm_session_id=None) + + chat_request: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( + model="anthropic/claude-sonnet-5-5", + input=cast( + ResponseInputParam, + [ + {"type": "compaction", "id": "cmp_litellm", "encrypted_content": json.dumps(_COMPACTION_BLOCK)}, + {"type": "function_call_output", "call_id": "call_never_stored", "output": "Rain."}, + ], + ), + responses_api_request=ResponsesAPIOptionalRequestParams( + **responses_api_request, + tools=[{"type": "function", "name": "get_weather", "parameters": {"type": "object"}, "strict": False}], + ), + ) + + request: Final = asyncio.run( + LiteLLMCompletionResponsesConfig.async_responses_api_session_handler( + previous_response_id="resp_unknown", + litellm_completion_request=chat_request, + instructions=None, + session_lookup=no_stored_turns, + ) + ) + + assert not any(message.get("tool_calls") for message in request["messages"]) + + +@pytest.mark.parametrize( + ("responses_api_request", "expected_turns"), + [ + pytest.param( + _COMPACTION_REQUESTED, + [("user", ()), ("assistant", ("call_0",)), ("tool", "Sunny."), ("assistant", ()), ("assistant", ("call_0",)), ("tool", "Rain.")], + id="native_compaction_block", + ), + pytest.param( + ResponsesAPIOptionalRequestParams(), + [("user", ()), ("assistant", ("call_0",)), ("tool", "Rain.")], + id="summary_as_text", + ), + ], +) +def test_compaction_replay_moves_only_the_call_a_later_output_answers_when_an_earlier_call_reused_its_id( + responses_api_request: ResponsesAPIOptionalRequestParams, expected_turns: list[tuple[str, object]] +) -> None: + messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( + model="anthropic/claude-sonnet-5-5", + input=cast( + ResponseInputParam, + [ + {"role": "user", "content": "Plan a Tokyo trip."}, + {"type": "function_call", "call_id": "call_0", "name": "get_weather", "arguments": "{}"}, + {"type": "function_call_output", "call_id": "call_0", "output": "Sunny."}, + _reasoning_item("rs_osaka", _ANSWERED_THINKING), + _CHECKING_OSAKA, + {"type": "function_call", "call_id": "call_0", "name": "get_weather", "arguments": "{}"}, + {"type": "compaction", "id": "cmp_litellm", "encrypted_content": json.dumps(_COMPACTION_BLOCK)}, + {"type": "function_call_output", "call_id": "call_0", "output": "Rain."}, + ], + ), + responses_api_request=responses_api_request, + )["messages"] + + assert [ + ( + message["role"], + message["content"] + if message["role"] == "tool" + else tuple(tool_call["id"] for tool_call in message.get("tool_calls") or ()), + ) + for message in messages + ] == expected_turns + assert messages[-2]["thinking_blocks"] == [_ANSWERED_THINKING] + + 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" @@ -6603,3 +6961,47 @@ async def test_bridge_partial_response_on_disconnect_bills_the_relayed_chat_chun assert incomplete_event.response.usage is not None assert incomplete_event.response.usage.output_tokens == litellm.token_counter(model="gpt-4o-mini", text=relayed_text) assert incomplete_event.response.usage.input_tokens > 0 + + +async def _input_tokens_billed_on_disconnect(request_input: ResponseInputParam) -> int: + from litellm.responses.litellm_completion_transformation.streaming_iterator import ( + LiteLLMCompletionStreamingIterator, + ) + + chat_stream: Final = await litellm.acompletion( + model="anthropic/claude-sonnet-5-5", + messages=[{"role": "user", "content": "Which day is driest?"}], + mock_response="Thursday looks driest, so plan the walking tour then.", + stream=True, + api_key="test-key", + ) + iterator: Final = LiteLLMCompletionStreamingIterator( + model="anthropic/claude-sonnet-5-5", + litellm_custom_stream_wrapper=chat_stream, + request_input=request_input, + responses_api_request={}, + custom_llm_provider="anthropic", + ) + _ = [await iterator.__anext__() for _ in range(6)] + incomplete_event: Final = iterator.partial_response_on_disconnect() + assert incomplete_event is not None and incomplete_event.response.usage is not None + return incomplete_event.response.usage.input_tokens + + +@pytest.mark.asyncio +async def test_bridge_partial_response_on_disconnect_bills_only_the_input_since_the_newest_compaction_item(): + since_compaction: Final = cast( + ResponseInputParam, + [ + {"type": "compaction", "id": "cmp_litellm", "encrypted_content": json.dumps(_COMPACTION_BLOCK)}, + {"role": "user", "content": "Which day is driest?"}, + ], + ) + history_it_replaces: Final = cast( + ResponseInputParam, + [{"role": "user", "content": "Plan a Tokyo trip, day by day, with every museum and food stall. " * 200}], + ) + + assert await _input_tokens_billed_on_disconnect( + cast(ResponseInputParam, [*history_it_replaces, *since_compaction]) + ) == await _input_tokens_billed_on_disconnect(since_compaction) 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 7a146c1b665..0eb1aa26a65 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_session_handler.py +++ b/tests/unit/responses/litellm_completion_transformation/test_session_handler.py @@ -1,6 +1,7 @@ import json import sqlite3 from collections.abc import Sequence +from itertools import chain from types import SimpleNamespace from typing import Final from unittest.mock import AsyncMock, patch @@ -18,6 +19,7 @@ from litellm.responses.litellm_completion_transformation.session_handler import ResponsesSessionHandler, _normalize_redacted_tool_call_arguments, ) +from litellm.responses.litellm_completion_transformation.transformation import ChatCompletionSession from litellm.responses.mcp.litellm_proxy_mcp_handler import LiteLLM_Proxy_MCP_Handler from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.llms.openai import ResponsesAPIResponse @@ -233,11 +235,10 @@ async def test_e2e_cold_storage_successful_retrieval(): patch.object(session_handler, "COLD_STORAGE_HANDLER") as mock_cold_storage, patch("litellm.cold_storage_custom_logger", return_value="s3"), ): - # Setup mocks mock_get_spend_logs.return_value = mock_spend_logs - mock_cold_storage.get_proxy_server_request_from_cold_storage_with_object_key = ( - AsyncMock(return_value=full_proxy_request) + mock_cold_storage.get_proxy_server_request_from_cold_storage_with_object_key = AsyncMock( + return_value=full_proxy_request ) # Call the main function @@ -291,7 +292,6 @@ async def test_e2e_cold_storage_fallback_to_truncated_payload(): ) as mock_get_spend_logs, patch.object(session_handler, "COLD_STORAGE_HANDLER") as mock_cold_storage, ): - # Setup mocks mock_get_spend_logs.return_value = mock_spend_logs @@ -358,41 +358,25 @@ async def test_should_check_cold_storage_for_full_payload(): with patch("litellm.cold_storage_custom_logger", return_value="s3"): # Test case 1: Should return True for truncated content - result1 = ResponsesSessionHandler._should_check_cold_storage_for_full_payload( - proxy_request_with_truncated_pdf - ) - assert ( - result1 == True - ), "Should return True for proxy request with truncated PDF content" + result1 = ResponsesSessionHandler._should_check_cold_storage_for_full_payload(proxy_request_with_truncated_pdf) + assert result1 == True, "Should return True for proxy request with truncated PDF content" # Test case 2: Should return False for regular content - result2 = ResponsesSessionHandler._should_check_cold_storage_for_full_payload( - proxy_request_regular - ) - assert ( - result2 == False - ), "Should return False for regular proxy request without truncation" + result2 = ResponsesSessionHandler._should_check_cold_storage_for_full_payload(proxy_request_regular) + assert result2 == False, "Should return False for regular proxy request without truncation" # Test case 3: Should return True for empty request - result3 = ResponsesSessionHandler._should_check_cold_storage_for_full_payload( - proxy_request_empty - ) + result3 = ResponsesSessionHandler._should_check_cold_storage_for_full_payload(proxy_request_empty) assert result3 == True, "Should return True for empty proxy request" # Test case 4: Should return True for None request - result4 = ResponsesSessionHandler._should_check_cold_storage_for_full_payload( - proxy_request_none - ) + result4 = ResponsesSessionHandler._should_check_cold_storage_for_full_payload(proxy_request_none) assert result4 == True, "Should return True for None proxy request" # Test case 5: Should return False when cold storage is not configured with patch.object(litellm, "cold_storage_custom_logger", None): - result5 = ResponsesSessionHandler._should_check_cold_storage_for_full_payload( - proxy_request_with_truncated_pdf - ) - assert ( - result5 == False - ), "Should return False when cold storage is not configured, even with truncated content" + result5 = ResponsesSessionHandler._should_check_cold_storage_for_full_payload(proxy_request_with_truncated_pdf) + assert result5 == False, "Should return False when cold storage is not configured, even with truncated content" @pytest.mark.asyncio @@ -422,9 +406,7 @@ async def test_get_chat_completion_message_history_empty_response_dict(): } ] - with patch.object( - ResponsesSessionHandler, "get_all_spend_logs_for_previous_response_id" - ) as mock_get_spend_logs: + with patch.object(ResponsesSessionHandler, "get_all_spend_logs_for_previous_response_id") as mock_get_spend_logs: mock_get_spend_logs.return_value = mock_spend_logs # Call the function @@ -518,9 +500,7 @@ async def test_message_history_reconstructs_list_shaped_input(): ) as mock_get_spend_logs: mock_get_spend_logs.return_value = mock_spend_logs - result = await ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id( - request_id - ) + result = await ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id(request_id) messages = result["messages"] assert [(message.get("role"), message.get("content")) for message in messages] == [ @@ -549,9 +529,7 @@ async def test_message_history_retries_a_spend_log_the_batch_writer_has_not_flus fake_prisma_client = _FakePrismaClient(results=[[], [spend_log]]) with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma_client): - result = await ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id( - request_id - ) + result = await ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id(request_id) messages = result["messages"] assert [(message.get("role"), message.get("content")) for message in messages] == [ @@ -618,14 +596,10 @@ async def test_message_history_looks_up_the_decoded_chat_completion_id(): model_id="e0f302a1412e78470ebb28cbed01fff5f88c0d331c667e9f2ba4b413c6fbd282", response_id=request_id, ) - fake_prisma_client = _FakePrismaClient( - results=[[_spend_log(request_id, "session-a", "Hello.", "Hi.")]] - ) + fake_prisma_client = _FakePrismaClient(results=[[_spend_log(request_id, "session-a", "Hello.", "Hi.")]]) with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma_client): - await ResponsesSessionHandler.get_all_spend_logs_for_previous_response_id( - encoded_response_id - ) + await ResponsesSessionHandler.get_all_spend_logs_for_previous_response_id(encoded_response_id) assert len(fake_prisma_client.db.calls) == 1 assert request_id in fake_prisma_client.db.calls[0] @@ -724,8 +698,9 @@ async def test_session_lookup_does_not_retry_when_spend_logs_are_disabled( """ fake_prisma_client = _FakePrismaClient(results=[]) - with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma_client), patch( - "litellm.proxy.proxy_server.disable_spend_logs", True + with ( + patch("litellm.proxy.proxy_server.prisma_client", fake_prisma_client), + patch("litellm.proxy.proxy_server.disable_spend_logs", True), ): spend_logs = await ResponsesSessionHandler.get_all_spend_logs_for_previous_response_id( "chatcmpl-does-not-exist" @@ -741,7 +716,11 @@ def test_normalize_redacted_arguments_skips_custom_tool_calls(): content=None, tool_calls=[ {"id": "call_c", "type": "custom", "custom": {"name": "run_code", "input": "print(1)"}}, - {"id": "call_f", "type": "function", "function": {"name": "get_weather", "arguments": "redacted-by-litellm"}}, + { + "id": "call_f", + "type": "function", + "function": {"name": "get_weather", "arguments": "redacted-by-litellm"}, + }, ], ) @@ -796,11 +775,13 @@ async def test_message_history_normalizes_redacted_tool_call_arguments(): } ] - with patch.object( # test-quality-ok: the handler has no DI seam for the spend-log fetch; every test in this file stubs this same boundary - ResponsesSessionHandler, - "get_all_spend_logs_for_previous_response_id", - new_callable=AsyncMock, - ) as mock_get_spend_logs: + with ( + patch.object( # test-quality-ok: the handler has no DI seam for the spend-log fetch; every test in this file stubs this same boundary + ResponsesSessionHandler, + "get_all_spend_logs_for_previous_response_id", + new_callable=AsyncMock, + ) as mock_get_spend_logs + ): mock_get_spend_logs.return_value = mock_spend_logs result = await ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id( @@ -878,6 +859,286 @@ async def test_message_history_replays_real_key_named_tool_payloads() -> None: assert json.loads(tool_message["content"]) == function_output +def _turn_spend_log( + turn_input: str | list[dict[str, str]], answer: str, compaction_blocks: list[dict[str, str]] | None = None +) -> dict[str, object]: + compaction_fields: Final = ( + {"provider_specific_fields": {"compaction_blocks": compaction_blocks}} if compaction_blocks else {} + ) + return { + "request_id": f"resp_{len(turn_input)}", + "session_id": "session-compacted", + "proxy_server_request": json.dumps({"input": turn_input}), + "response": { + "id": f"chatcmpl-{len(turn_input)}", + "object": "chat.completion", + "choices": [{"index": 0, "message": {"role": "assistant", "content": answer, **compaction_fields}}], + }, + } + + +def _role_and_content(message: object) -> tuple[object, object]: + if isinstance(message, dict): + return message.get("role"), message.get("content") + return getattr(message, "role", None), getattr(message, "content", None) + + +async def _history_from_stored_turns( + spend_logs: tuple[dict[str, object], ...], previous_response_id: str +) -> ChatCompletionSession: + with patch("litellm.proxy.proxy_server.prisma_client", SimpleNamespace(db=_StoredSpendLogs(spend_logs))): + return await ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id( + previous_response_id + ) + + +@pytest.mark.asyncio +async def test_previous_response_history_starts_at_the_newest_turn_that_sent_a_compaction_item(): + first_turn: Final = [{"role": "user", "content": "Plan a Tokyo trip."}] + compaction_item: Final = { + "type": "compaction", + "id": "cmp_1", + "encrypted_content": json.dumps({"type": "compaction", "content": "Turn one planned Tokyo."}), + } + second_turn: Final = [ + *first_turn, + {"role": "assistant", "content": "Day one, Asakusa."}, + compaction_item, + {"role": "user", "content": "And Osaka?"}, + ] + spend_logs: Final = ( + _turn_spend_log(first_turn, "Day one, Asakusa."), + _turn_spend_log(second_turn, "Dotonbori."), + ) + + session: Final = await _history_from_stored_turns(spend_logs, "resp_4") + + assert [_role_and_content(message) for message in session["messages"]] == [ + ("user", "Turn one planned Tokyo."), + ("user", "And Osaka?"), + ("assistant", "Dotonbori."), + ] + + +@pytest.mark.parametrize( + "first_turn_input", + [pytest.param([{"role": "user", "content": "Hi."}], id="item_list"), pytest.param("Hi.", id="plain_string")], +) +@pytest.mark.asyncio +async def test_previous_response_history_starts_at_the_newest_answer_that_compacted( + first_turn_input: str | list[dict[str, str]], +): + compaction_blocks: Final = [{"type": "compaction", "content": "Turn one planned Tokyo."}] + spend_logs: Final = ( + _turn_spend_log(first_turn_input, "Hello."), + _turn_spend_log( + [{"role": "developer", "content": "Answer briefly."}, {"role": "user", "content": "Plan a Tokyo trip."}], + "Day one, Asakusa.", + compaction_blocks, + ), + _turn_spend_log([{"role": "user", "content": "And Osaka?"}], "Dotonbori."), + ) + + session: Final = await _history_from_stored_turns(spend_logs, "resp_1") + + assert [_role_and_content(message) for message in session["messages"]] == [ + ("developer", "Answer briefly."), + ("assistant", "Day one, Asakusa."), + ("user", "And Osaka?"), + ("assistant", "Dotonbori."), + ] + assert getattr(session["messages"][1], "provider_specific_fields", None) == {"compaction_blocks": compaction_blocks} + + +_COMPACTION_ITEM_SENT: Final = { + "type": "compaction", + "id": "cmp_1", + "encrypted_content": json.dumps({"type": "compaction", "content": "Turn one planned Tokyo."}), +} + + +@pytest.mark.parametrize( + "compacting_turn", + [ + pytest.param( + _turn_spend_log( + [{"role": "user", "content": "Plan a Tokyo trip."}], + "Day one, Asakusa.", + [{"type": "compaction", "content": "Turn one planned Tokyo."}], + ), + id="answer_compacted", + ), + pytest.param( + _turn_spend_log( + [_COMPACTION_ITEM_SENT, {"role": "user", "content": "Plan a Tokyo trip."}], "Day one, Asakusa." + ), + id="input_sent_a_compaction_item", + ), + ], +) +@pytest.mark.asyncio +async def test_previous_response_history_keeps_instructions_from_turns_the_compaction_covers( + compacting_turn: dict[str, object], +): + spend_logs: Final = ( + _turn_spend_log( + [{"role": "developer", "content": "Use metric units."}, {"role": "user", "content": "Hi."}], "Hello." + ), + compacting_turn, + _turn_spend_log([{"role": "user", "content": "And Osaka?"}], "Dotonbori."), + ) + + session: Final = await _history_from_stored_turns(spend_logs, "resp_1") + + roles_and_contents: Final = [_role_and_content(message) for message in session["messages"]] + assert roles_and_contents[0] == ("developer", "Use metric units.") + assert ("user", "Hi.") not in roles_and_contents + assert ("assistant", "Hello.") not in roles_and_contents + assert roles_and_contents[-2:] == [("user", "And Osaka?"), ("assistant", "Dotonbori.")] + + +@pytest.mark.asyncio +async def test_previous_response_history_drops_a_tool_output_carrying_a_system_role_the_compaction_covers(): + spend_logs: Final = ( + _turn_spend_log( + [ + {"role": "user", "content": "Look up the code."}, + {"type": "function_call", "call_id": "call_1", "name": "lookup", "arguments": "{}"}, + {"type": "function_call_output", "call_id": "call_1", "output": "4242", "role": "system"}, + ], + "Found it.", + ), + _turn_spend_log([_COMPACTION_ITEM_SENT, {"role": "user", "content": "Plan a Tokyo trip."}], "Day one."), + _turn_spend_log([{"role": "user", "content": "And Osaka?"}], "Dotonbori."), + ) + + session: Final = await _history_from_stored_turns(spend_logs, "resp_1") + + contents: Final = [str(_role_and_content(message)[1]) for message in session["messages"]] + assert not any("4242" in content for content in contents) + assert contents[-2:] == ["And Osaka?", "Dotonbori."] + + +def _output_message(text: str) -> dict[str, object]: + return {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": text}]} + + +@pytest.mark.asyncio +async def test_previous_response_history_starts_at_a_compaction_item_in_a_stored_responses_output(): + responses_shaped_turn: Final = { + "request_id": "resp_mcp", + "session_id": "session-compacted", + "proxy_server_request": json.dumps({"input": [{"role": "user", "content": "Plan a Tokyo trip."}]}), + "response": { + "id": "resp_mcp", + "object": "response", + "output": [_output_message("Checking the map."), _COMPACTION_ITEM_SENT, _output_message("Day one.")], + }, + } + spend_logs: Final = ( + _turn_spend_log([{"role": "user", "content": "Hi."}], "Hello."), + responses_shaped_turn, + _turn_spend_log([{"role": "user", "content": "And Osaka?"}], "Dotonbori."), + ) + + session: Final = await _history_from_stored_turns(spend_logs, "resp_1") + + replayed: Final = json.dumps([_role_and_content(message) for message in session["messages"]], default=str) + assert "Hello." not in replayed + assert "Plan a Tokyo trip." not in replayed + assert "Checking the map." not in replayed + assert "Turn one planned Tokyo." in replayed + assert [str(_role_and_content(message)[1]) for message in session["messages"]][-2:] == [ + "And Osaka?", + "Dotonbori.", + ] + + +def _tool_call_ids(message: object) -> tuple[str, ...]: + tool_calls: Final = message.get("tool_calls") if isinstance(message, dict) else getattr(message, "tool_calls", None) + return tuple(call["id"] if isinstance(call, dict) else call.id for call in tool_calls or ()) + + +_WEATHER_CALL_TURN: Final = { + "request_id": "resp_weather", + "session_id": "session-compacted", + "proxy_server_request": json.dumps({"input": [{"role": "user", "content": "Weather in Paris?"}]}), + "response": { + "id": "chatcmpl-weather", + "object": "chat.completion", + "choices": [ + { + "index": 0, + "finish_reason": "tool_calls", + "message": { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city": "Paris"}'}, + } + ], + }, + } + ], + }, +} + + +def _compacted_weather_answer(request_id: str) -> dict[str, object]: + return { + **_turn_spend_log( + [_COMPACTION_ITEM_SENT, {"type": "function_call_output", "call_id": "call_1", "output": "sunny"}], + "It is sunny.", + ), + "request_id": request_id, + } + + +@pytest.mark.asyncio +async def test_previous_response_history_ignores_a_compaction_item_sent_after_the_referenced_response(): + session: Final = await _history_from_stored_turns( + (_WEATHER_CALL_TURN, _compacted_weather_answer("resp_first_send")), "resp_weather" + ) + + assert _role_and_content(session["messages"][0]) == ("user", "Weather in Paris?") + assert tuple(chain.from_iterable(map(_tool_call_ids, session["messages"]))) == ("call_1",) + + +@pytest.mark.asyncio +async def test_a_turn_after_a_compaction_item_keeps_the_stored_call_its_tool_output_answers(): + session: Final = await _history_from_stored_turns( + (_WEATHER_CALL_TURN, _compacted_weather_answer("resp_compacted")), "resp_compacted" + ) + + roles: Final = [_role_and_content(message)[0] for message in session["messages"]] + assert roles == ["user", "assistant", "tool", "assistant"] + assert _tool_call_ids(session["messages"][1]) == ("call_1",) + assert ("user", "Weather in Paris?") not in map(_role_and_content, session["messages"]) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("spend_logs", "previous_response_id", "replayed_roles"), + [ + ((_WEATHER_CALL_TURN, _compacted_weather_answer("resp_first_send")), "resp_weather", ["user", "assistant"]), + ( + (_WEATHER_CALL_TURN, _compacted_weather_answer("resp_first_send"), _compacted_weather_answer("resp_retry")), + "resp_first_send", + ["user", "assistant", "tool", "assistant"], + ), + ], +) +async def test_a_compacted_session_replays_no_turn_sent_after_the_referenced_response( + spend_logs: tuple[dict[str, object], ...], previous_response_id: str, replayed_roles: list[str] +): + session: Final = await _history_from_stored_turns(spend_logs, previous_response_id) + + assert [_role_and_content(message)[0] for message in session["messages"]] == replayed_roles + + class _SqliteBackedPrismaDB: def __init__(self, rows: list[tuple[str, str, str, str, str | None, str | None]]) -> None: self._connection = sqlite3.connect(":memory:") @@ -889,7 +1150,9 @@ class _SqliteBackedPrismaDB: self._connection.executemany('INSERT INTO "LiteLLM_SpendLogs" VALUES (?, ?, ?, ?, ?, ?)', rows) async def query_raw(self, query: str, *args: object) -> list[dict[str, object]]: - sqlite_query: Final = query.replace("= ANY($1::text[])", "IN (SELECT value FROM json_each(?))").replace("$1", "?") + sqlite_query: Final = query.replace("= ANY($1::text[])", "IN (SELECT value FROM json_each(?))").replace( + "$1", "?" + ) sqlite_args: Final = tuple(json.dumps(arg) if isinstance(arg, list) else arg for arg in args) return [dict(row) for row in self._connection.execute(sqlite_query, sqlite_args)] diff --git a/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py b/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py index 1a3430224d1..ed90acbd2fa 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py +++ b/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py @@ -1524,3 +1524,37 @@ async def test_caller_function_named_like_the_search_tool_streams_as_function_ca assert [(item["type"], item["name"]) for item in output] == [("function_call", "tool_search_tool_regex")] assert "response.function_call_arguments.done" in [event["type"] for event in events] + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +async def test_a_compaction_block_without_its_summary_streams_no_compaction_item(sync_mode: bool): + deltas: Final = ( + Delta( + role="assistant", + content="", + provider_specific_fields={"compaction_blocks": [{"type": "compaction", "content": None}]}, + ), + Delta(content="Ok."), + ) + chunks: Final = tuple( + ModelResponseStream( + id=CHAT_COMPLETION_ID, + created=1748575031, + model="claude-sonnet-5-5", + object="chat.completion.chunk", + choices=[ + StreamingChoices(index=0, delta=delta, finish_reason="stop" if index == len(deltas) - 1 else None) + ], + ) + for index, delta in enumerate(deltas) + ) + iterator: Final = _build_iterator(chunks) + events: Final = ( + [event.model_dump(exclude_none=True) for event in iterator] + if sync_mode + else [event.model_dump(exclude_none=True) async for event in iterator] + ) + + assert [event["item"]["type"] for event in events if "item" in event] == ["message", "message"] + assert [item["type"] for item in events[-1]["response"]["output"]] == ["message"] diff --git a/tests/unit/responses/test_responses_api_request_body.py b/tests/unit/responses/test_responses_api_request_body.py index ea3eab7e002..436367bd82d 100644 --- a/tests/unit/responses/test_responses_api_request_body.py +++ b/tests/unit/responses/test_responses_api_request_body.py @@ -7,6 +7,7 @@ in expected_responses_api_request/. import copy import json import logging +from collections.abc import Awaitable, Callable from pathlib import Path from importlib import import_module from typing import Final @@ -1141,6 +1142,61 @@ async def test_a_native_responses_provider_places_every_point_itself(): assert with_unmatchable_system == user_only +_CLAUDE_COMPACTION_ITEM: Final = { + "type": "compaction", + "id": "cmp_claude", + "encrypted_content": json.dumps({"type": "compaction", "content": "The user's cat is named Biscuit."}), +} +_OPENAI_COMPACTION_ITEM: Final = {"type": "compaction", "id": "cmp_openai", "encrypted_content": "gAAAAABopaque"} +_HISTORY: Final = ( + {"role": "system", "content": "Answer briefly."}, + {"role": "user", "content": "My cat is named Biscuit."}, + {"role": "assistant", "content": "Noted."}, +) +_NEXT_TURN: Final = {"role": "user", "content": "What is my cat's name?"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("call", "upstream_url"), + [ + pytest.param(litellm.aresponses, "https://api.openai.com/v1/responses", id="create"), + pytest.param(litellm.acompact_responses, "https://api.openai.com/v1/responses/compact", id="compact"), + ], +) +@pytest.mark.parametrize( + ("compaction_item", "sent_input"), + [ + ( + _CLAUDE_COMPACTION_ITEM, + [ + _HISTORY[0], + {"role": "user", "content": "The user's cat is named Biscuit.", "type": "message"}, + _NEXT_TURN, + ], + ), + (_OPENAI_COMPACTION_ITEM, [*_HISTORY, _OPENAI_COMPACTION_ITEM, _NEXT_TURN]), + ], +) +async def test_aresponses_sends_a_claude_compaction_item_to_openai_as_its_summary( + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, + call: Callable[..., Awaitable[object]], + upstream_url: str, + compaction_item: dict[str, str], + sent_input: list[dict[str, str]], +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + upstream: Final = respx_mock.post(upstream_url).mock( + return_value=httpx.Response(200, json=_minimal_responses_api_payload("resp_compacted", "gpt-5.6")) + ) + + await call(model="openai/gpt-5.6", api_key="fake-api-key", input=[*_HISTORY, compaction_item, _NEXT_TURN]) + + assert json.loads(upstream.calls[0].request.read())["input"] == sent_input + + def test_non_streamed_response_ignores_a_stream_timeout_it_never_uses() -> None: answer: Final = json.dumps(_minimal_responses_api_payload("resp_unstreamed", "gpt-5.6")).encode() client: Final = HTTPHandler( diff --git a/tests/unit/responses/test_streaming_iterator.py b/tests/unit/responses/test_streaming_iterator.py index 62d89a55e85..6c3e6177503 100644 --- a/tests/unit/responses/test_streaming_iterator.py +++ b/tests/unit/responses/test_streaming_iterator.py @@ -2605,6 +2605,27 @@ def test_faked_stream_builds_each_refusal_summary_and_custom_tool_input_once_for } +_COMPACTION_SUMMARY: Final = "The user planned a Tokyo trip." +_NEXT_TURN: Final = {"type": "message", "role": "user", "content": "And Osaka?"} +_COMPACTED_INPUT: Final = [ + {"type": "message", "role": "user", "content": "Plan a Tokyo trip. " * 200}, + {"type": "message", "role": "assistant", "content": "Day one, Asakusa. " * 200}, + { + "type": "compaction", + "id": "cmp_1", + "encrypted_content": json.dumps({"type": "compaction", "content": _COMPACTION_SUMMARY}), + }, + _NEXT_TURN, +] +_SUMMARY_INPUT: Final = [{"type": "message", "role": "user", "content": _COMPACTION_SUMMARY}, _NEXT_TURN] + + +def test_usage_estimate_counts_a_compaction_item_as_its_summary(): + assert _estimate_usage_from_text("gpt-4o-mini", _COMPACTED_INPUT, {}, "Dotonbori.") == _estimate_usage_from_text( + "gpt-4o-mini", _SUMMARY_INPUT, {}, "Dotonbori." + ) + + @pytest.mark.asyncio async def test_shutdown_drain_waits_for_a_pending_logging_task(): release: Final = asyncio.Event() diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index e1f748fcb58..9cc1bf6ce79 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -7889,6 +7889,29 @@ def test_count_pre_call_check_tokens_across_api_surfaces(): router._count_pre_call_check_tokens(messages=None, input=None) +def test_count_pre_call_check_tokens_skips_the_history_a_compaction_item_replaces(): + router = litellm.Router(model_list=[{"model_name": "m", "litellm_params": {"model": "gpt-3.5-turbo"}}]) + summary: Final = "The user planned a Tokyo trip." + compaction_item: Final = { + "type": "compaction", + "id": "cmp_1", + "encrypted_content": json.dumps({"type": "compaction", "content": summary}), + } + history: Final = [ + {"role": "user", "content": "Plan a Tokyo trip. " * 200}, + {"role": "assistant", "content": "Day one, Asakusa. " * 200}, + ] + next_turn: Final = {"role": "user", "content": "And Osaka?"} + + compacted_tokens: Final = router._count_pre_call_check_tokens( + messages=None, input=[*history, compaction_item, next_turn] + ) + + assert compacted_tokens == router._count_pre_call_check_tokens( + messages=None, input=[{"role": "user", "content": summary}, next_turn] + ) + + def test_pre_call_checks_no_messages_or_input_does_not_crash(monkeypatch): """ When neither messages nor input is provided (e.g. endpoints without prompt text), diff --git a/tests/unit/types/test_litellm_params.py b/tests/unit/types/test_litellm_params.py index 5a2cefd08b1..b0086662f1c 100644 --- a/tests/unit/types/test_litellm_params.py +++ b/tests/unit/types/test_litellm_params.py @@ -254,6 +254,7 @@ INTERNAL_STATE_NAMES: Final = ( "_context_compaction_state", *AGENTIC_LOOP_STATE_NAMES, "_router_weights", + "_encrypted_content_affinity_pinned", "fallback_depth", "max_fallbacks", "attempted_targets", diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 2b174a1141a..bd08b743c7d 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -44279,7 +44279,7 @@ export interface components { /** Output */ output: (components["schemas"]["ResponseOutputMessage"] | components["schemas"]["ResponseFileSearchToolCall"] | components["schemas"]["ResponseFunctionToolCall"] | components["schemas"]["ResponseFunctionToolCallOutputItem"] | components["schemas"]["ResponseFunctionWebSearch"] | components["schemas"]["ResponseComputerToolCall"] | components["schemas"]["ResponseComputerToolCallOutputItem"] | components["schemas"]["ResponseReasoningItem"] | components["schemas"]["ResponseToolSearchCall"] | components["schemas"]["ResponseToolSearchOutputItem"] | components["schemas"]["ResponseCompactionItem"] | components["schemas"]["ImageGenerationCall"] | components["schemas"]["ResponseCodeInterpreterToolCall"] | components["schemas"]["LocalShellCall"] | components["schemas"]["LocalShellCallOutput"] | components["schemas"]["ResponseFunctionShellToolCall"] | components["schemas"]["ResponseFunctionShellToolCallOutput"] | components["schemas"]["ResponseApplyPatchToolCall"] | components["schemas"]["ResponseApplyPatchToolCallOutput"] | components["schemas"]["McpCall"] | components["schemas"]["McpListTools"] | components["schemas"]["McpApprovalRequest"] | components["schemas"]["McpApprovalResponse"] | components["schemas"]["ResponseCustomToolCall"] | components["schemas"]["ResponseCustomToolCallOutputItem"] | { [key: string]: unknown; - })[] | (components["schemas"]["GenericResponseOutputItem"] | components["schemas"]["OutputCodeInterpreterCall"] | components["schemas"]["OutputFunctionToolCall"] | components["schemas"]["OutputImageGenerationCall"] | components["schemas"]["ResponseFunctionToolCall"] | components["schemas"]["ResponseFunctionWebSearch"] | components["schemas"]["CustomToolCallOutputItem"] | components["schemas"]["OutputToolSearchCall"])[]; + })[] | (components["schemas"]["GenericResponseOutputItem"] | components["schemas"]["OutputCodeInterpreterCall"] | components["schemas"]["OutputFunctionToolCall"] | components["schemas"]["OutputImageGenerationCall"] | components["schemas"]["ResponseFunctionToolCall"] | components["schemas"]["ResponseFunctionWebSearch"] | components["schemas"]["CustomToolCallOutputItem"] | components["schemas"]["OutputToolSearchCall"] | components["schemas"]["ResponseCompactionItem"])[]; /** Parallel Tool Calls */ parallel_tool_calls?: boolean | null; /** Previous Response Id */ From a31c5230de303d342f62e0e2ecd62af83e5fcd4e Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 21 Oct 2026 09:42:05 +0000 Subject: [PATCH 10/14] fix(proxy): forward dispatched_call_type through _failure_usage_to_lift (#47815) --- litellm/proxy/utils.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 52b3866d04d..28be11509ae 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1294,6 +1294,7 @@ def _failure_usage_to_lift( provider_billed: bool = False, deployment_id: str | None = None, base_model: str | None = None, + dispatched_call_type: str | None = None, ) -> tuple[object, object] | None: current_attempt: Final = _current_failure_usage_to_lift( model_call_details, @@ -1302,6 +1303,7 @@ def _failure_usage_to_lift( provider_billed, deployment_id, base_model, + dispatched_call_type=dispatched_call_type, ) failed_attempts: Final = _FAILED_ATTEMPTS_USAGE.validate_python(model_call_details.get("failed_attempts_usage", ())) if not failed_attempts: From 9d8fe1ddb51c0cb4c3190aa51274fbdde98eb05f Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 21 Oct 2026 11:58:31 +0000 Subject: [PATCH 11/14] refactor(types): replace Any with proven types in 50 files (#47818) * refactor(types): replace Any with proven types in 100 files Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(types): keep untyped parsing where strict validation changed accepted shapes Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(types): keep the OpenAI-compatible response parse untyped Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(types): keep only annotation changes, drop new runtime validation 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/integrations/deepeval/types.py | 4 ++-- .../streaming_iterator.py | 6 +++--- litellm/llms/base_llm/ocr/transformation.py | 6 +++--- litellm/llms/base_llm/sandbox/transformation.py | 6 +++--- .../llms/github_copilot/embedding/transformation.py | 4 +++- litellm/llms/ollama/completion/handler.py | 4 ++-- litellm/models/managed_files.py | 8 ++++---- litellm/proxy/a2a/endpoints.py | 6 +++--- litellm/proxy/auth/auth_exception_handler.py | 2 +- litellm/proxy/management_endpoints/types.py | 4 ++-- litellm/proxy/spend_tracking/vantage_endpoints.py | 4 +++- litellm/rust_bridge/detectors.py | 2 +- litellm/types/caching.py | 4 ++-- litellm/types/completion.py | 2 +- litellm/types/extraction.py | 3 ++- litellm/types/google_genai/main.py | 2 +- litellm/types/images/main.py | 4 ++-- litellm/types/integrations/argilla.py | 4 ++-- litellm/types/integrations/cloudzero.py | 7 ++----- .../types/integrations/compression_interception.py | 4 ++-- litellm/types/integrations/gcs_bucket.py | 4 ++-- litellm/types/integrations/langsmith.py | 8 ++++---- litellm/types/integrations/posthog.py | 4 ++-- .../llms/anthropic_messages/anthropic_response.py | 4 ++-- litellm/types/llms/base.py | 4 ++-- litellm/types/llms/databricks.py | 6 +++--- .../management_endpoints/cache_settings_endpoints.py | 6 +++--- .../router_settings_endpoints.py | 6 +++--- litellm/types/mcp.py | 10 +++++----- litellm/types/mcp_server/mcp_server_manager.py | 4 ++-- litellm/types/prompts/init_prompts.py | 6 +++--- litellm/types/proxy/callback_logs_endpoints.py | 4 ++-- litellm/types/proxy/cloudzero_endpoints.py | 2 +- .../guardrail_hooks/azure/azure_prompt_shield.py | 4 +--- .../guardrail_hooks/azure/azure_text_moderation.py | 4 ++-- .../proxy/guardrails/guardrail_hooks/compresr.py | 4 ++-- .../guardrail_hooks/litellm_content_filter.py | 6 +++--- .../proxy/guardrails/guardrail_hooks/presidio.py | 8 ++++---- .../management_endpoints/common_daily_activity.py | 5 ++--- .../proxy/management_endpoints/config_overrides.py | 6 ++---- .../management_endpoints/internal_user_endpoints.py | 4 ++-- .../model_management_endpoints.py | 3 +-- .../proxy/management_endpoints/team_endpoints.py | 8 ++++---- litellm/types/proxy/prompt_endpoints.py | 4 +--- .../types/proxy/public_endpoints/public_endpoints.py | 4 ++-- litellm/types/proxy/vantage_endpoints.py | 6 +++--- litellm/types/rag.py | 6 +++--- litellm/types/responses/main.py | 4 ++-- litellm/types/vector_store_files.py | 12 +++++++----- litellm/types/vector_stores.py | 2 +- 50 files changed, 120 insertions(+), 124 deletions(-) diff --git a/litellm/integrations/deepeval/types.py b/litellm/integrations/deepeval/types.py index eb4bb1f9b2a..a547863e26a 100644 --- a/litellm/integrations/deepeval/types.py +++ b/litellm/integrations/deepeval/types.py @@ -1,6 +1,6 @@ # Duplicate -> https://github.com/confident-ai/deepeval/blob/main/deepeval/tracing/api.py from enum import Enum -from typing import Any, ClassVar, Literal +from typing import ClassVar, Literal from pydantic import ConfigDict, Field @@ -55,7 +55,7 @@ class TraceApi(LiteLLMBaseModel): tool_spans: list[BaseApiSpan] = Field(alias="toolSpans") start_time: str = Field(alias="startTime") end_time: str = Field(alias="endTime") - metadata: dict[str, Any] | None = Field(None) + metadata: dict[str, object] | None = Field(None) tags: list[str] | None = Field(None) environment: str | None = Field(None) diff --git a/litellm/interactions/litellm_responses_transformation/streaming_iterator.py b/litellm/interactions/litellm_responses_transformation/streaming_iterator.py index dab12d447df..46e633d2b96 100644 --- a/litellm/interactions/litellm_responses_transformation/streaming_iterator.py +++ b/litellm/interactions/litellm_responses_transformation/streaming_iterator.py @@ -3,8 +3,8 @@ Streaming iterator for transforming Responses API stream to Interactions API str """ from collections import deque -from collections.abc import AsyncIterator, Iterator -from typing import Any, Final, cast +from collections.abc import AsyncIterator, Iterator, Mapping +from typing import Final, cast from litellm.responses.streaming_iterator import ( BaseResponsesAPIStreamingIterator, @@ -44,7 +44,7 @@ class LiteLLMResponsesInteractionsStreamingIterator: request_input: InteractionInput | None, optional_params: InteractionsAPIOptionalRequestParams, custom_llm_provider: str | None = None, - litellm_metadata: dict[str, Any] | None = None, + litellm_metadata: Mapping[str, object] | None = None, ): self.model = model self.responses_stream_iterator = litellm_custom_stream_wrapper diff --git a/litellm/llms/base_llm/ocr/transformation.py b/litellm/llms/base_llm/ocr/transformation.py index 97d402e7a1b..171b193f60c 100644 --- a/litellm/llms/base_llm/ocr/transformation.py +++ b/litellm/llms/base_llm/ocr/transformation.py @@ -4,7 +4,7 @@ Base OCR types shared by the Rust OCR route and Python consumers. import builtins from collections.abc import Mapping -from typing import Any, Final, Literal +from typing import Final, Literal from pydantic import PrivateAttr @@ -45,7 +45,7 @@ class OCRPageImage(LiteLLMPydanticObjectBase): """Image extracted from OCR page.""" image_base64: str | None = None - bbox: dict[str, Any] | None = None + bbox: dict[str, builtins.object] | None = None model_config = {"extra": "allow"} @@ -80,7 +80,7 @@ class OCRResponse(LiteLLMPydanticObjectBase): pages: list[OCRPage] model: str - document_annotation: Any | None = None + document_annotation: builtins.object | None = None usage_info: OCRUsageInfo | None = None content: str | None = None tables: list[dict[str, builtins.object]] | None = None diff --git a/litellm/llms/base_llm/sandbox/transformation.py b/litellm/llms/base_llm/sandbox/transformation.py index 877edc15ec7..fe0e3abb87e 100644 --- a/litellm/llms/base_llm/sandbox/transformation.py +++ b/litellm/llms/base_llm/sandbox/transformation.py @@ -7,7 +7,7 @@ run code -> delete container; `code_interpreter_tool` combines all three. """ import builtins -from typing import Any, Final +from typing import Final import httpx from pydantic import Field, PrivateAttr @@ -42,8 +42,8 @@ class CodeExecutionResult(LiteLLMPydanticObjectBase): stdout: str = "" stderr: str = "" - results: list[dict[str, Any]] = Field(default_factory=list) - error: dict[str, Any] | None = None + results: list[dict[str, object]] = Field(default_factory=list) + error: dict[str, object] | None = None execution_count: int | None = None object: str = "code_execution" diff --git a/litellm/llms/github_copilot/embedding/transformation.py b/litellm/llms/github_copilot/embedding/transformation.py index f16564ba505..4f308b01159 100644 --- a/litellm/llms/github_copilot/embedding/transformation.py +++ b/litellm/llms/github_copilot/embedding/transformation.py @@ -202,7 +202,9 @@ class GithubCopilotEmbeddingConfig(BaseEmbeddingConfig): optional_params[param] = value return optional_params - def get_error_class(self, error_message: str, status_code: int, headers: Any) -> BaseLLMException: + def get_error_class( + self, error_message: str, status_code: int, headers: dict[str, str] | httpx.Headers + ) -> BaseLLMException: from litellm.llms.openai.openai import OpenAIConfig return OpenAIConfig().get_error_class(error_message=error_message, status_code=status_code, headers=headers) diff --git a/litellm/llms/ollama/completion/handler.py b/litellm/llms/ollama/completion/handler.py index d4110e55bf0..0c13b1cd1d3 100644 --- a/litellm/llms/ollama/completion/handler.py +++ b/litellm/llms/ollama/completion/handler.py @@ -87,7 +87,7 @@ async def ollama_aembeddings( model: str, prompts: list[str], model_response: EmbeddingResponse, - optional_params: dict, + optional_params: Mapping[str, object], logging_obj: object, encoding: TokenEncoder | None, ): @@ -110,7 +110,7 @@ def ollama_embeddings( api_base: str, model: str, prompts: list[str], - optional_params: dict, + optional_params: Mapping[str, object], model_response: EmbeddingResponse, logging_obj: object, encoding: TokenEncoder | None = None, diff --git a/litellm/models/managed_files.py b/litellm/models/managed_files.py index 435db632026..b2bfb572b7d 100644 --- a/litellm/models/managed_files.py +++ b/litellm/models/managed_files.py @@ -6,7 +6,7 @@ Canonical definitions for the ``litellm_managed*`` tables. Re-exported from """ from datetime import datetime -from typing import Any, Literal +from typing import Literal from litellm.types.llms.base import LiteLLMPydanticObjectBase from litellm.types.llms.openai import OpenAIFileObject, ResponsesAPIResponse @@ -39,7 +39,7 @@ class LiteLLM_ManagedVectorStoreTable(LiteLLMPydanticObjectBase): """Table for managing vector stores with target_model_names support.""" unified_resource_id: str - resource_object: Any | None = None + resource_object: object | None = None model_mappings: dict[str, str] flat_model_resource_ids: list[str] created_by: str | None = None @@ -54,11 +54,11 @@ class LiteLLM_ManagedVectorStoresTable(LiteLLMPydanticObjectBase): custom_llm_provider: str vector_store_name: str | None = None vector_store_description: str | None = None - vector_store_metadata: dict[str, Any] | None = None + vector_store_metadata: dict[str, object] | None = None created_at: datetime | None = None updated_at: datetime | None = None litellm_credential_name: str | None = None - litellm_params: dict[str, Any] | None = None + litellm_params: dict[str, object] | None = None team_id: str | None = None user_id: str | None = None is_config: bool = False diff --git a/litellm/proxy/a2a/endpoints.py b/litellm/proxy/a2a/endpoints.py index 56b512df934..730d40b46d0 100644 --- a/litellm/proxy/a2a/endpoints.py +++ b/litellm/proxy/a2a/endpoints.py @@ -9,7 +9,7 @@ admin pick which ones to expose through the proxy. The actual merge into a LiteLLM-fronted card happens when the agent is saved via ``POST /v1/agents``. """ -from typing import Any, Final +from typing import Final from fastapi import APIRouter, Depends, HTTPException from fastapi.responses import JSONResponse @@ -50,7 +50,7 @@ class DiscoverAgentRequest(LiteLLMBaseModel): "query parameter." ), ) - params: dict[str, Any] | None = Field( + params: dict[str, object] | None = Field( default=None, description=( "Mode-specific parameters. ``langgraph_platform`` requires " @@ -61,7 +61,7 @@ class DiscoverAgentRequest(LiteLLMBaseModel): class DiscoverAgentResponse(LiteLLMBaseModel): url: str - agent_card: dict[str, Any] + agent_card: dict[str, object] @router.post( diff --git a/litellm/proxy/auth/auth_exception_handler.py b/litellm/proxy/auth/auth_exception_handler.py index ae473a62271..f6637a1fdf3 100644 --- a/litellm/proxy/auth/auth_exception_handler.py +++ b/litellm/proxy/auth/auth_exception_handler.py @@ -41,7 +41,7 @@ DB_UNAVAILABLE_FALLBACK_USER_ID: Final = "__db_unavailable_fallback__" if TYPE_CHECKING: from opentelemetry.trace import Span as _Span - Span = _Span | Any + Span = _Span else: Span = Any diff --git a/litellm/proxy/management_endpoints/types.py b/litellm/proxy/management_endpoints/types.py index 4414eed97b2..221a6221fab 100644 --- a/litellm/proxy/management_endpoints/types.py +++ b/litellm/proxy/management_endpoints/types.py @@ -5,7 +5,7 @@ Might include fastapi/proxy requirements.txt related imports """ from collections.abc import Iterable, Sequence -from typing import Any, Final, cast +from typing import Final, cast from fastapi_sso.sso.base import OpenID @@ -94,4 +94,4 @@ def get_litellm_user_role(role_str: object) -> LitellmUserRoles | None: class CustomOpenID(OpenID): team_ids: list[str] user_role: LitellmUserRoles | None = None - extra_fields: dict[str, Any] | None = None + extra_fields: dict[str, object] | None = None diff --git a/litellm/proxy/spend_tracking/vantage_endpoints.py b/litellm/proxy/spend_tracking/vantage_endpoints.py index d3858f95243..6938588444c 100644 --- a/litellm/proxy/spend_tracking/vantage_endpoints.py +++ b/litellm/proxy/spend_tracking/vantage_endpoints.py @@ -122,7 +122,9 @@ async def _get_vantage_settings(): encrypted_api_key: Final = settings.get("api_key") if encrypted_api_key: - decrypted_api_key = decrypt_value_helper(encrypted_api_key, key="vantage_api_key", exception_type="error") + decrypted_api_key: Final = decrypt_value_helper( + encrypted_api_key, key="vantage_api_key", exception_type="error" + ) if decrypted_api_key is None: raise HTTPException( status_code=500, diff --git a/litellm/rust_bridge/detectors.py b/litellm/rust_bridge/detectors.py index 0ad1b787e38..0af3a2e6fa6 100644 --- a/litellm/rust_bridge/detectors.py +++ b/litellm/rust_bridge/detectors.py @@ -128,7 +128,7 @@ def _custom_loggers(callbacks: Iterable[object]) -> tuple[CustomLogger, ...]: def overridden(method: str, callbacks: Iterable[object]) -> bool: - base: Final = getattr(CustomLogger, method) + base: Final[object] = getattr(CustomLogger, method) return any(getattr(type(callback), method, base) is not base for callback in _custom_loggers(callbacks)) diff --git a/litellm/types/caching.py b/litellm/types/caching.py index b9bcbbb93dd..ced2065a674 100644 --- a/litellm/types/caching.py +++ b/litellm/types/caching.py @@ -75,7 +75,7 @@ class RedisPipelineSetOperation(TypedDict): """ key: str - value: Any + value: object ttl: int | None @@ -133,7 +133,7 @@ class HealthCheckCacheParams(LiteLLMBaseModel): host: str | None = None port: str | int | None = None - redis_kwargs: dict[str, Any] | None = None + redis_kwargs: dict[str, object] | None = None namespace: str | None = None redis_version: str | int | float | None = None diff --git a/litellm/types/completion.py b/litellm/types/completion.py index a6693030e89..0f4bdbd75ce 100644 --- a/litellm/types/completion.py +++ b/litellm/types/completion.py @@ -246,7 +246,7 @@ _CompletionDispatchContext = CompletionDispatchContext _CompletionDispatchResult = Union[ - Coroutine[Any, Any, Union["ModelResponse", "CustomStreamWrapper"]], + Coroutine[object, object, Union["ModelResponse", "CustomStreamWrapper"]], "ModelResponse", "CustomStreamWrapper", ] diff --git a/litellm/types/extraction.py b/litellm/types/extraction.py index 4311e549b9a..562bc41f9de 100644 --- a/litellm/types/extraction.py +++ b/litellm/types/extraction.py @@ -13,6 +13,7 @@ from pydantic import ( model_validator, with_config, ) +from pydantic.fields import FieldInfo from typing_extensions import ReadOnly, Required, Self, TypedDict from litellm.types.llms.base import LiteLLMPydanticObjectBase @@ -30,7 +31,7 @@ def _is_none(value: object) -> bool: return value is None -_OMITTED_WHEN_NONE: Final = Field(exclude_if=_is_none) +_OMITTED_WHEN_NONE: Final = FieldInfo.from_field(exclude_if=_is_none) OptionalJson: TypeAlias = Annotated[JsonValue, _OMITTED_WHEN_NONE] OptionalFloat: TypeAlias = Annotated[float | None, _OMITTED_WHEN_NONE] OptionalInt: TypeAlias = Annotated[int | None, _OMITTED_WHEN_NONE] diff --git a/litellm/types/google_genai/main.py b/litellm/types/google_genai/main.py index 0241c9fa809..dcde0a031e3 100644 --- a/litellm/types/google_genai/main.py +++ b/litellm/types/google_genai/main.py @@ -17,7 +17,7 @@ if TYPE_CHECKING: ToolConfigDict = _genai_types.ToolConfigDict class GenerateContentRequestDict(GenerateContentRequestParametersDict): - generationConfig: Any | None + generationConfig: object | None tools: ToolConfigDict | None class GenerateContentResponse(GoogleGenAIGenerateContentResponse, BaseLiteLLMOpenAIResponseObject): diff --git a/litellm/types/images/main.py b/litellm/types/images/main.py index 5d80135a8a1..963c980a4b8 100644 --- a/litellm/types/images/main.py +++ b/litellm/types/images/main.py @@ -1,4 +1,4 @@ -from typing import Any, Literal +from typing import Literal from typing_extensions import TypedDict @@ -20,7 +20,7 @@ class ImageEditOptionalRequestParams(TypedDict, total=False): response_format: Literal["url", "b64_json"] | None size: str | None user: str | None - imageConfig: dict[str, Any] | None + imageConfig: dict[str, object] | None class ImageEditRequestParams(ImageEditOptionalRequestParams, total=False): diff --git a/litellm/types/integrations/argilla.py b/litellm/types/integrations/argilla.py index 2c98486c7a7..754916c2f3a 100644 --- a/litellm/types/integrations/argilla.py +++ b/litellm/types/integrations/argilla.py @@ -1,10 +1,10 @@ -from typing import Any, Final +from typing import Final from typing_extensions import TypedDict class ArgillaItem(TypedDict): - fields: dict[str, Any] + fields: dict[str, object] class ArgillaPayload(TypedDict): diff --git a/litellm/types/integrations/cloudzero.py b/litellm/types/integrations/cloudzero.py index 988d501600f..dcc5d03ccca 100644 --- a/litellm/types/integrations/cloudzero.py +++ b/litellm/types/integrations/cloudzero.py @@ -1,7 +1,4 @@ -from typing import Any - - -class CBFRecord(dict[str, Any]): +class CBFRecord(dict[str, object]): """CloudZero Billing Format (CBF) record structure. This class represents a CBF record that is created from LiteLLM usage data @@ -31,4 +28,4 @@ class CBFRecord(dict[str, Any]): # Type alias for better readability in function signatures -CBFRecordDict = dict[str, Any] +CBFRecordDict = dict[str, object] diff --git a/litellm/types/integrations/compression_interception.py b/litellm/types/integrations/compression_interception.py index a1e66e5057e..e9c40fbf838 100644 --- a/litellm/types/integrations/compression_interception.py +++ b/litellm/types/integrations/compression_interception.py @@ -2,7 +2,7 @@ Type definitions for Compression Interception integration. """ -from typing import Any, Literal, TypedDict +from typing import Literal, TypedDict class CompressionInterceptionConfig(TypedDict, total=False): @@ -24,7 +24,7 @@ class CompressionInterceptionConfig(TypedDict, total=False): compression_trigger: int compression_target: int | None embedding_model: str | None - embedding_model_params: dict[str, Any] | None + embedding_model_params: dict[str, object] | None class CompressionSavingsMetadata(TypedDict): diff --git a/litellm/types/integrations/gcs_bucket.py b/litellm/types/integrations/gcs_bucket.py index 3840d7681a5..79563a0a6e4 100644 --- a/litellm/types/integrations/gcs_bucket.py +++ b/litellm/types/integrations/gcs_bucket.py @@ -31,5 +31,5 @@ class GCSLogQueueItem(TypedDict): """ payload: StandardLoggingPayload - kwargs: dict[str, Any] - response_obj: Any | None + kwargs: dict[str, object] + response_obj: object | None diff --git a/litellm/types/integrations/langsmith.py b/litellm/types/integrations/langsmith.py index 594190ddb5c..f053a5864c1 100644 --- a/litellm/types/integrations/langsmith.py +++ b/litellm/types/integrations/langsmith.py @@ -1,6 +1,6 @@ from dataclasses import dataclass from datetime import datetime -from typing import Any, NamedTuple +from typing import NamedTuple from typing_extensions import TypedDict @@ -9,7 +9,7 @@ from litellm.types.llms.base import LiteLLMBaseModel class LangsmithInputs(LiteLLMBaseModel): model: str | None = None - messages: list[Any] | None = None + messages: list[object] | None = None stream: bool | None = None call_type: str | None = None litellm_call_id: str | None = None @@ -17,9 +17,9 @@ class LangsmithInputs(LiteLLMBaseModel): temperature: float | None = None max_tokens: int | None = None custom_llm_provider: str | None = None - input: list[Any] | None = None + input: list[object] | None = None log_event_type: str | None = None - original_response: Any | None = None + original_response: object | None = None response_cost: float | None = None # LiteLLM Virtual Key specific fields diff --git a/litellm/types/integrations/posthog.py b/litellm/types/integrations/posthog.py index ac04d08107d..d9739df7df7 100644 --- a/litellm/types/integrations/posthog.py +++ b/litellm/types/integrations/posthog.py @@ -1,4 +1,4 @@ -from typing import Any, Final, TypedDict +from typing import Final, TypedDict POSTHOG_MAX_BATCH_SIZE: Final = 100 @@ -7,7 +7,7 @@ class PostHogEventPayload(TypedDict): """PostHog event payload structure""" event: str # "$ai_generation" or "$ai_embedding" - properties: dict[str, Any] + properties: dict[str, object] distinct_id: str diff --git a/litellm/types/llms/anthropic_messages/anthropic_response.py b/litellm/types/llms/anthropic_messages/anthropic_response.py index 03956ae7a54..9ad13243451 100644 --- a/litellm/types/llms/anthropic_messages/anthropic_response.py +++ b/litellm/types/llms/anthropic_messages/anthropic_response.py @@ -1,5 +1,5 @@ from collections.abc import Sequence -from typing import Any, Literal, TypeAlias +from typing import Literal, TypeAlias from typing_extensions import NotRequired, ReadOnly, TypedDict @@ -20,7 +20,7 @@ class AnthropicResponseTextBlock(TypedDict, total=False): Anthropic Response Text Block: https://docs.anthropic.com/en/api/messages """ - citations: list[dict[str, Any]] | None + citations: list[dict[str, object]] | None text: str type: Literal["text"] diff --git a/litellm/types/llms/base.py b/litellm/types/llms/base.py index 47fec914f9c..5a2da071cec 100644 --- a/litellm/types/llms/base.py +++ b/litellm/types/llms/base.py @@ -1,4 +1,4 @@ -from typing import Any, Final +from typing import Final from openai._models import BaseModel as OpenAIObject from pydantic import ConfigDict @@ -45,7 +45,7 @@ class BaseLiteLLMOpenAIResponseObject(LiteLLMBaseModel): class HiddenParams(OpenAIObject): - original_response: str | Any | None = None + original_response: str | object | None = None model_id: str | None = None # used in Router for individual deployments api_base: str | None = None # returns api base used for making completion call _response_ms: float | None = None diff --git a/litellm/types/llms/databricks.py b/litellm/types/llms/databricks.py index c7023e35777..bc229421d91 100644 --- a/litellm/types/llms/databricks.py +++ b/litellm/types/llms/databricks.py @@ -1,4 +1,4 @@ -from typing import Any, Literal +from typing import Literal from pydantic import BaseModel from typing_extensions import ( @@ -23,7 +23,7 @@ class GenericStreamingChunk(TypedDict, total=False): class DatabricksTextContent(TypedDict, total=False): type: Literal["text"] text: Required[str] - citations: list[dict[str, Any]] | None + citations: list[dict[str, object]] | None class DatabricksReasoningSummary(TypedDict): @@ -35,7 +35,7 @@ class DatabricksReasoningSummary(TypedDict): class DatabricksReasoningContent(TypedDict, total=False): type: Literal["reasoning"] summary: Required[list[DatabricksReasoningSummary]] - citations: list[dict[str, Any]] | None + citations: list[dict[str, object]] | None AllDatabricksContentListValues = DatabricksTextContent | DatabricksReasoningContent diff --git a/litellm/types/management_endpoints/cache_settings_endpoints.py b/litellm/types/management_endpoints/cache_settings_endpoints.py index f6394f52c99..3584f8ce1da 100644 --- a/litellm/types/management_endpoints/cache_settings_endpoints.py +++ b/litellm/types/management_endpoints/cache_settings_endpoints.py @@ -2,7 +2,7 @@ Types and field definitions for cache settings management endpoints """ -from typing import Any, Final +from typing import Final from litellm.types.llms.base import LiteLLMBaseModel @@ -10,9 +10,9 @@ from litellm.types.llms.base import LiteLLMBaseModel class CacheSettingsField(LiteLLMBaseModel): field_name: str field_type: str - field_value: Any + field_value: object field_description: str - field_default: Any = None + field_default: object = None options: list[str] | None = None # For fields with predefined options/enum values ui_field_name: str # User-friendly display name link: str | None = None # Documentation link for the field diff --git a/litellm/types/management_endpoints/router_settings_endpoints.py b/litellm/types/management_endpoints/router_settings_endpoints.py index 9c8fe446d72..64068f30a16 100644 --- a/litellm/types/management_endpoints/router_settings_endpoints.py +++ b/litellm/types/management_endpoints/router_settings_endpoints.py @@ -2,7 +2,7 @@ Types and field definitions for router settings management endpoints """ -from typing import Any, Final, Literal +from typing import Final, Literal from pydantic import Field, field_validator @@ -72,9 +72,9 @@ class FallbackDeleteResponse(LiteLLMBaseModel): class RouterSettingsField(LiteLLMBaseModel): field_name: str field_type: str - field_value: Any + field_value: object field_description: str - field_default: Any = None + field_default: object = None options: list[str] | None = None # For fields with predefined options/enum values ui_field_name: str # User-friendly display name link: str | None = None # Documentation link for the field diff --git a/litellm/types/mcp.py b/litellm/types/mcp.py index f0be66f055d..3fa9adef913 100644 --- a/litellm/types/mcp.py +++ b/litellm/types/mcp.py @@ -113,7 +113,7 @@ class MCPPublicServer(LiteLLMBaseModel): transport: MCPTransportType spec_path: str | None = None auth_type: MCPAuthType | None = None - mcp_info: dict[str, Any] | None = None + mcp_info: dict[str, object] | None = None class MCPAllowedClient(LiteLLMBaseModel): @@ -448,11 +448,11 @@ class MCPPreCallRequestObject(LiteLLMBaseModel): """ tool_name: str - arguments: dict[str, Any] + arguments: dict[str, object] server_name: str | None = None tool_description: str | None = None tool_input_schema: Mapping[str, object] | None = None - user_api_key_auth: dict[str, Any] | None = None + user_api_key_auth: dict[str, object] | None = None hidden_params: HiddenParams = HiddenParams() @@ -462,7 +462,7 @@ class MCPPreCallResponseObject(LiteLLMBaseModel): """ should_proceed: bool = True - modified_arguments: dict[str, Any] | None = None + modified_arguments: dict[str, object] | None = None error_message: str | None = None hidden_params: HiddenParams = HiddenParams() @@ -473,7 +473,7 @@ class MCPDuringCallRequestObject(LiteLLMBaseModel): """ tool_name: str - arguments: dict[str, Any] + arguments: dict[str, object] server_name: str | None = None tool_description: str | None = None tool_input_schema: Mapping[str, object] | None = None diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index 219f23200d2..d29e4db55a4 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -117,7 +117,7 @@ class MCPServer(LiteLLMBaseModel): # Admin-configured env vars. Each entry is {name, value, scope, description}. # scope=="global" values are interpolated into static_headers using ${NAME}. # scope=="user" values must be supplied per-user. - env_vars: list[dict[str, Any]] | None = None + env_vars: list[dict[str, object]] | None = None # OAuth-specific fields authorization_server_metadata: Mapping[str, object] | None = Field(default=None, exclude=True) client_id: str | None = None @@ -240,7 +240,7 @@ class MCPServer(LiteLLMBaseModel): # token_validation: key-value pairs that must match fields in the OAuth token # response (supports dot-notation for nested fields, e.g. "team.enterprise_id"). # Tokens that fail validation are rejected before storage. - token_validation: dict[str, Any] | None = None + token_validation: dict[str, object] | None = None oauth_identity_binding: MCPOAuthIdentityBinding | None = None # Optional TTL override (seconds) for the Redis per-user token cache, capped # at the token's expires_in minus the expiry buffer so a cached entry never diff --git a/litellm/types/prompts/init_prompts.py b/litellm/types/prompts/init_prompts.py index fc4c36e760a..c1a3ab56311 100644 --- a/litellm/types/prompts/init_prompts.py +++ b/litellm/types/prompts/init_prompts.py @@ -1,6 +1,6 @@ from datetime import datetime from enum import Enum -from typing import Any, Literal +from typing import Literal from pydantic import ConfigDict @@ -31,7 +31,7 @@ class PromptLiteLLMParams(LiteLLMBaseModel): api_base: str | None = None api_key: str | None = None - provider_specific_query_params: dict[str, Any] | None = None + provider_specific_query_params: dict[str, object] | None = None ignore_prompt_manager_model: bool | None = False ignore_prompt_manager_optional_params: bool | None = False @@ -66,7 +66,7 @@ class PromptSpec(LiteLLMBaseModel): class PromptTemplateBase(LiteLLMBaseModel): litellm_prompt_id: str content: str - metadata: dict[str, Any] | None = None + metadata: dict[str, object] | None = None class PromptInfoResponse(LiteLLMBaseModel): diff --git a/litellm/types/proxy/callback_logs_endpoints.py b/litellm/types/proxy/callback_logs_endpoints.py index b2519d97122..0291cc902bc 100644 --- a/litellm/types/proxy/callback_logs_endpoints.py +++ b/litellm/types/proxy/callback_logs_endpoints.py @@ -5,7 +5,7 @@ External producers (e.g. the litellm-rust gateway) POST finished logging payloads here; the proxy replays them through the standard callback fan-out. """ -from typing import Any, Literal +from typing import Literal from pydantic import Field @@ -17,7 +17,7 @@ class CallbackLogRecord(LiteLLMBaseModel): """A single finished logging event to replay through the callbacks.""" status: Literal["success", "failure"] - standard_logging_payload: dict[str, Any] + standard_logging_payload: dict[str, object] error: str | None = None diff --git a/litellm/types/proxy/cloudzero_endpoints.py b/litellm/types/proxy/cloudzero_endpoints.py index ece294674b9..f87537494bd 100644 --- a/litellm/types/proxy/cloudzero_endpoints.py +++ b/litellm/types/proxy/cloudzero_endpoints.py @@ -46,7 +46,7 @@ class CloudZeroExportResponse(LiteLLMBaseModel): dry_run_data: dict[str, Any] | None = Field( None, description="Dry run data including usage data and CBF transformed data" ) - summary: dict[str, Any] | None = Field(None, description="Summary statistics for dry run") + summary: dict[str, object] | None = Field(None, description="Summary statistics for dry run") class CloudZeroSettingsView(LiteLLMBaseModel): diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/azure/azure_prompt_shield.py b/litellm/types/proxy/guardrails/guardrail_hooks/azure/azure_prompt_shield.py index 60846b2a1bd..4f4bebf9704 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/azure/azure_prompt_shield.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/azure/azure_prompt_shield.py @@ -1,5 +1,3 @@ -from typing import Any - from pydantic import Field from typing_extensions import TypedDict @@ -23,7 +21,7 @@ class AzurePromptShieldGuardrailResponse(TypedDict): """Configuration parameters for the Azure Prompt Shield guardrail""" userPromptAnalysis: UserPromptAnalysis - documentsAnalysis: list[dict[str, Any]] + documentsAnalysis: list[dict[str, object]] class AzurePromptShieldGuardrailConfigModel( diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/azure/azure_text_moderation.py b/litellm/types/proxy/guardrails/guardrail_hooks/azure/azure_text_moderation.py index 90662e84ed5..6ee5741baf5 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/azure/azure_text_moderation.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/azure/azure_text_moderation.py @@ -1,4 +1,4 @@ -from typing import Any, Final, Literal +from typing import Final, Literal from pydantic import Field from typing_extensions import Required, TypedDict @@ -37,7 +37,7 @@ class AzureTextModerationGuardrailResponseCategoriesAnalysis(TypedDict): class AzureTextModerationGuardrailResponse(TypedDict): """Response from the Azure Text Moderation guardrail""" - blocklistsMatch: list[dict[str, Any]] + blocklistsMatch: list[dict[str, object]] categoriesAnalysis: list[AzureTextModerationGuardrailResponseCategoriesAnalysis] diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/compresr.py b/litellm/types/proxy/guardrails/guardrail_hooks/compresr.py index 6598f1de62b..c0f5ea5346e 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/compresr.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/compresr.py @@ -1,4 +1,4 @@ -from typing import Any, Literal +from typing import Literal from pydantic import Field @@ -93,7 +93,7 @@ class CompresrGuardrailOptionalParams(LiteLLMBaseModel): "Unset lets the server default apply (~10.0)." ), ) - compression_params: dict[str, Any] | None = Field( + compression_params: dict[str, object] | None = Field( default=None, description=( "Passthrough of extra parameters forwarded verbatim in the Compresr " diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/litellm_content_filter.py b/litellm/types/proxy/guardrails/guardrail_hooks/litellm_content_filter.py index 78800a5e9a2..d49a3afae04 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/litellm_content_filter.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/litellm_content_filter.py @@ -1,5 +1,5 @@ from enum import Enum -from typing import Any, Literal, TypedDict +from typing import Literal, TypedDict from pydantic import Field @@ -77,7 +77,7 @@ class CompetitorIntentDetection(TypedDict): action_hint: str entities: dict[str, list[str]] signals: list[str] - evidence: list[dict[str, Any]] + evidence: list[dict[str, object]] ContentFilterDetection = PatternDetection | BlockedWordDetection | CategoryKeywordDetection | CompetitorIntentDetection @@ -157,7 +157,7 @@ class LitellmContentFilterGuardrailConfigModel(GuardrailConfigModel): ) # Competitor intent blocker (generic; industry presets add domain_words, etc.) - competitor_intent_config: dict[str, Any] | None = Field( + competitor_intent_config: dict[str, object] | None = Field( default=None, description="Optional config for intent-based competitor comparison detection. " "Keys: brand_self (list), competitors (list), competitor_aliases (dict), " diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/types/proxy/guardrails/guardrail_hooks/presidio.py index 9d1e375af72..6fe324ce6b0 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/presidio.py @@ -1,6 +1,6 @@ -from typing import Any +from collections.abc import Mapping -from typing_extensions import TypedDict +from typing_extensions import ReadOnly, TypedDict from litellm.types.guardrails import PiiEntityType @@ -17,5 +17,5 @@ class PresidioAnalyzeResponseItem(TypedDict, total=False): start: int | None end: int | None score: float | None - analysis_explanation: dict[str, Any] | None - recognition_metadata: dict[str, Any] | None + analysis_explanation: ReadOnly[Mapping[str, object] | None] + recognition_metadata: ReadOnly[Mapping[str, object] | None] diff --git a/litellm/types/proxy/management_endpoints/common_daily_activity.py b/litellm/types/proxy/management_endpoints/common_daily_activity.py index 66702abfb34..213c911cfdc 100644 --- a/litellm/types/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/types/proxy/management_endpoints/common_daily_activity.py @@ -1,7 +1,6 @@ -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from datetime import date from enum import Enum -from typing import Any from pydantic import ConfigDict, Field from typing_extensions import TypedDict @@ -61,7 +60,7 @@ class KeyMetricWithMetadata(MetricBase): class MetricWithMetadata(MetricBase): - metadata: dict[str, Any] = Field(default_factory=dict) + metadata: Mapping[str, object] = Field(default_factory=dict) # API key breakdown for this metric (e.g., which API keys are using this MCP server) api_key_breakdown: dict[str, KeyMetricWithMetadata] = Field(default_factory=dict) # api_key -> {metrics, metadata} diff --git a/litellm/types/proxy/management_endpoints/config_overrides.py b/litellm/types/proxy/management_endpoints/config_overrides.py index c82e616cf9a..b807b5537d5 100644 --- a/litellm/types/proxy/management_endpoints/config_overrides.py +++ b/litellm/types/proxy/management_endpoints/config_overrides.py @@ -1,5 +1,3 @@ -from typing import Any - from pydantic import Field from litellm.types.llms.base import LiteLLMBaseModel @@ -103,5 +101,5 @@ class ConfigOverrideSettingsResponse(LiteLLMBaseModel): """Response model for config override settings GET endpoints.""" config_type: str = Field(description="The type of config override") - values: dict[str, Any] = Field(description="Current configuration values (sensitive fields decrypted)") - field_schema: dict[str, Any] = Field(description="Schema information for UI rendering") + values: dict[str, object] = Field(description="Current configuration values (sensitive fields decrypted)") + field_schema: dict[str, object] = Field(description="Schema information for UI rendering") diff --git a/litellm/types/proxy/management_endpoints/internal_user_endpoints.py b/litellm/types/proxy/management_endpoints/internal_user_endpoints.py index febd4485d10..fc0b73f77b6 100644 --- a/litellm/types/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/types/proxy/management_endpoints/internal_user_endpoints.py @@ -1,5 +1,5 @@ from collections.abc import Mapping, Sequence -from typing import Any, Final, Literal +from typing import Final, Literal from pydantic import ConfigDict, Field, field_validator from typing_extensions import ReadOnly, TypedDict @@ -80,7 +80,7 @@ class UserUpdateResult(LiteLLMBaseModel): user_email: str | None = None success: bool error: str | None = None - updated_user: dict[str, Any] | None = None + updated_user: dict[str, object] | None = None class BulkUpdateUserResponse(LiteLLMBaseModel): diff --git a/litellm/types/proxy/management_endpoints/model_management_endpoints.py b/litellm/types/proxy/management_endpoints/model_management_endpoints.py index 6a4160c3c50..c012f872280 100644 --- a/litellm/types/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/types/proxy/management_endpoints/model_management_endpoints.py @@ -1,6 +1,5 @@ from collections.abc import Mapping from datetime import datetime -from typing import Any from pydantic import ConfigDict, Field @@ -22,7 +21,7 @@ class UpdateUsefulLinksRequest(LiteLLMBaseModel): # Supports both old format (Dict[str, str]) and new format (Dict[str, Dict[str, Any]]) # New format: { "displayName": { "url": "...", "index": 0 } } # Old format: { "displayName": "url" } (for backward compatibility) - useful_links: dict[str, str | dict[str, Any]] + useful_links: dict[str, str | dict[str, object]] class AutoRouterClassifierDefaultPromptResponse(LiteLLMBaseModel): diff --git a/litellm/types/proxy/management_endpoints/team_endpoints.py b/litellm/types/proxy/management_endpoints/team_endpoints.py index 743b2d7c963..51c48e6ebf9 100644 --- a/litellm/types/proxy/management_endpoints/team_endpoints.py +++ b/litellm/types/proxy/management_endpoints/team_endpoints.py @@ -1,5 +1,5 @@ from collections.abc import Mapping -from typing import Any, Final, Literal +from typing import Final, Literal from pydantic import ConfigDict, Field, field_validator, model_validator @@ -121,8 +121,8 @@ class TeamMemberAddResult(LiteLLMBaseModel): user_email: str | None = None success: bool error: str | None = None - updated_user: dict[str, Any] | None = None - updated_team_membership: dict[str, Any] | None = None + updated_user: dict[str, object] | None = None + updated_team_membership: dict[str, object] | None = None class BulkTeamMemberAddResponse(LiteLLMBaseModel): @@ -133,7 +133,7 @@ class BulkTeamMemberAddResponse(LiteLLMBaseModel): total_requested: int successful_additions: int failed_additions: int - updated_team: dict[str, Any] | None = None + updated_team: dict[str, object] | None = None class TeamMemberRef(MemberDeleteRequest): diff --git a/litellm/types/proxy/prompt_endpoints.py b/litellm/types/proxy/prompt_endpoints.py index 930713663a1..cded9f5fb87 100644 --- a/litellm/types/proxy/prompt_endpoints.py +++ b/litellm/types/proxy/prompt_endpoints.py @@ -1,9 +1,7 @@ -from typing import Any - from litellm.types.llms.base import LiteLLMBaseModel class TestPromptRequest(LiteLLMBaseModel): dotprompt_content: str - prompt_variables: dict[str, Any] | None = None + prompt_variables: dict[str, object] | None = None conversation_history: list[dict[str, str]] | None = None diff --git a/litellm/types/proxy/public_endpoints/public_endpoints.py b/litellm/types/proxy/public_endpoints/public_endpoints.py index c487ce8b6c5..0e8a2bb3401 100644 --- a/litellm/types/proxy/public_endpoints/public_endpoints.py +++ b/litellm/types/proxy/public_endpoints/public_endpoints.py @@ -1,5 +1,5 @@ from collections.abc import Mapping, Sequence -from typing import Any, Literal +from typing import Literal from pydantic import ConfigDict @@ -13,7 +13,7 @@ class PublicModelHubInfo(LiteLLMBaseModel): # Supports both old format (Dict[str, str]) and new format (Dict[str, Dict[str, Any]]) # New format: { "displayName": { "url": "...", "index": 0 } } # Old format: { "displayName": "url" } (for backward compatibility) - useful_links: dict[str, str | dict[str, Any]] | None + useful_links: dict[str, str | dict[str, object]] | None class ProviderCredentialField(LiteLLMBaseModel): diff --git a/litellm/types/proxy/vantage_endpoints.py b/litellm/types/proxy/vantage_endpoints.py index 1833d8b8386..ccf050a8921 100644 --- a/litellm/types/proxy/vantage_endpoints.py +++ b/litellm/types/proxy/vantage_endpoints.py @@ -2,8 +2,8 @@ Vantage endpoint types for LiteLLM Proxy """ +from collections.abc import Mapping from datetime import datetime -from typing import Any from pydantic import Field, field_validator @@ -57,10 +57,10 @@ class VantageExportResponse(LiteLLMBaseModel): message: str status: str - dry_run_data: dict[str, Any] | None = Field( + dry_run_data: Mapping[str, object] | None = Field( None, description="Dry run data including usage data and FOCUS transformed data" ) - summary: dict[str, Any] | None = Field(None, description="Summary statistics for dry run") + summary: Mapping[str, object] | None = Field(None, description="Summary statistics for dry run") class VantageSettingsView(LiteLLMBaseModel): diff --git a/litellm/types/rag.py b/litellm/types/rag.py index 2f912d21587..02c06f242d1 100644 --- a/litellm/types/rag.py +++ b/litellm/types/rag.py @@ -3,7 +3,7 @@ Type definitions for RAG (Retrieval Augmented Generation) Ingest API. """ from collections.abc import Mapping -from typing import Any, Literal +from typing import Literal from pydantic import ConfigDict from typing_extensions import ReadOnly, TypedDict @@ -232,7 +232,7 @@ class RAGIngestRequest(LiteLLMBaseModel): file_url: str | None = None # URL to fetch file from file_id: str | None = None # Existing file ID - ingest_options: dict[str, Any] # RAGIngestOptions as dict for flexibility + ingest_options: dict[str, object] # RAGIngestOptions as dict for flexibility model_config = ConfigDict(extra="allow") # Allow additional fields @@ -260,7 +260,7 @@ class RAGQueryRequest(LiteLLMBaseModel): """Request body for RAG query API.""" model: str - messages: list[Any] + messages: list[object] retrieval_config: RAGRetrievalConfig rerank: RAGRerankConfig | None = None stream: bool | None = False diff --git a/litellm/types/responses/main.py b/litellm/types/responses/main.py index 5588527c141..72dc63a581f 100644 --- a/litellm/types/responses/main.py +++ b/litellm/types/responses/main.py @@ -5,7 +5,7 @@ from typing import Final, Literal, Optional, Union from openai.types.responses.response_function_tool_call import ResponseFunctionToolCall from openai.types.responses.response_function_web_search import ActionSearchSource, ResponseFunctionWebSearch from pydantic import ConfigDict, PrivateAttr -from typing_extensions import Any, TypedDict +from typing_extensions import TypedDict from litellm.types.llms.base import BaseLiteLLMOpenAIResponseObject @@ -100,7 +100,7 @@ class OutputCodeInterpreterCall(BaseLiteLLMOpenAIResponseObject): def build_code_interpreter_log_outputs( - content: Any, + content: object, ) -> list[OutputCodeInterpreterCallLog] | None: """Convert Anthropic bash_code_execution stdout/stderr to log outputs. diff --git a/litellm/types/vector_store_files.py b/litellm/types/vector_store_files.py index 6b4953965d5..3623746216a 100644 --- a/litellm/types/vector_store_files.py +++ b/litellm/types/vector_store_files.py @@ -1,7 +1,9 @@ +import builtins +from collections.abc import Mapping from enum import Enum -from typing import Any, Literal +from typing import Literal -from typing_extensions import TypedDict +from typing_extensions import ReadOnly, TypedDict class VectorStoreFileStatus(str, Enum): @@ -33,7 +35,7 @@ class VectorStoreFileObject(TypedDict, total=False): usage_bytes: int | None vector_store_id: str status: VectorStoreFileStatus - last_error: dict[str, Any] | None + last_error: ReadOnly[Mapping[str, builtins.object] | None] chunking_strategy: VectorStoreFileChunkingStrategy | None attributes: dict[str, str | int | float | bool] | None @@ -83,5 +85,5 @@ class VectorStoreFileContentResponse(TypedDict, total=False): class VectorStoreFileAuthCredentials(TypedDict, total=False): - headers: dict[str, Any] - query_params: dict[str, Any] + headers: ReadOnly[Mapping[str, object]] + query_params: ReadOnly[Mapping[str, object]] diff --git a/litellm/types/vector_stores.py b/litellm/types/vector_stores.py index 5e7d517ece2..2259339e479 100644 --- a/litellm/types/vector_stores.py +++ b/litellm/types/vector_stores.py @@ -31,7 +31,7 @@ class LiteLLM_ManagedVectorStore(TypedDict, total=False): vector_store_name: str | None vector_store_description: str | None - vector_store_metadata: dict[str, Any] | str | None + vector_store_metadata: dict[str, object] | str | None created_at: datetime | None updated_at: datetime | None From 3f93350411ddb50b1beb413665df4d0eb5a065de Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 21 Oct 2026 06:37:24 -0700 Subject: [PATCH 12/14] docs: rename security.md to SECURITY.md and point contact to security@berri.ai (#47581) Co-authored-by: krrish-berri-2 <270687000+krrish-berri-2@users.noreply.github.com> --- security.md => SECURITY.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) rename security.md => SECURITY.md (98%) diff --git a/security.md b/SECURITY.md similarity index 98% rename from security.md rename to SECURITY.md index c73379a9c01..d3603448d46 100644 --- a/security.md +++ b/SECURITY.md @@ -79,7 +79,7 @@ To qualify for a bounty, reports must include clear reproduction steps, a reprod - Audit Logs with retention policy - Control Allowed IP Addresses that can access your Cloud LiteLLM Instance -For security inquiries, please contact us at support@berri.ai +For security inquiries, please contact us at security@berri.ai #### Supported data regions for LiteLLM Cloud From b6f91a221b6b9df6f8a89ec6f8264ac87d63ff5d Mon Sep 17 00:00:00 2001 From: yujonglee Date: Wed, 21 Oct 2026 07:03:18 -0700 Subject: [PATCH 13/14] refactor(inference): rename RouteError to Error (#47771) Every other crate names its error Error. The format crates already re-exported it under that name, so only the inference crate and a few importers change --- .../crates/gateway-inference/src/caching.rs | 8 +++--- .../crates/gateway-inference/src/error.rs | 26 +++++++++-------- litellm-rust/crates/inference-chat/src/lib.rs | 2 +- .../crates/inference-chat/tests/caching.rs | 10 +++---- .../crates/inference-messages/src/lib.rs | 2 +- .../inference-messages/tests/caching.rs | 10 +++---- litellm-rust/crates/inference-ocr/AGENTS.md | 2 +- .../inference-ocr/src/provider_config.rs | 2 +- .../crates/inference-responses/src/lib.rs | 2 +- .../inference-responses/tests/caching.rs | 10 +++---- .../crates/inference-transcription/src/lib.rs | 2 +- litellm-rust/crates/inference/AGENTS.md | 6 ++-- litellm-rust/crates/inference/src/caching.rs | 26 ++++++++--------- litellm-rust/crates/inference/src/context.rs | 10 +++---- litellm-rust/crates/inference/src/error.rs | 24 ++++++++-------- litellm-rust/crates/inference/src/lib.rs | 2 +- litellm-rust/crates/inference/src/provider.rs | 2 +- .../crates/inference/tests/caching.rs | 24 ++++++++-------- .../crates/python-bridge/src/errors.rs | 28 +++++++++---------- .../python-bridge/src/routes/inference.rs | 10 +++---- 20 files changed, 104 insertions(+), 104 deletions(-) diff --git a/litellm-rust/crates/gateway-inference/src/caching.rs b/litellm-rust/crates/gateway-inference/src/caching.rs index 1abea61b164..426b8b91182 100644 --- a/litellm-rust/crates/gateway-inference/src/caching.rs +++ b/litellm-rust/crates/gateway-inference/src/caching.rs @@ -72,26 +72,26 @@ fn duration(seconds: f64) -> Result { #[derive(Clone, Default)] pub(crate) struct CacheHeaders(std::sync::Arc>); -impl litellm_host::interceptors::Interceptors for CacheHeaders { +impl litellm_host::interceptors::Interceptors for CacheHeaders { async fn before_provider_request( &self, wire: litellm_host::interceptors::WireRequest, _: litellm_host::interceptors::RequestContext, - ) -> Result { + ) -> Result { Ok(wire) } async fn after_provider_response( &self, _: litellm_host::interceptors::RawResponse, - ) -> Result<(), litellm_inference::RouteError> { + ) -> Result<(), litellm_inference::Error> { Ok(()) } async fn result_ready( &self, facts: litellm_host::interceptors::ExecutionFacts, - ) -> Result<(), litellm_inference::RouteError> { + ) -> Result<(), litellm_inference::Error> { if let litellm_host::interceptors::ResultSource::Cache { key } = facts.source { let _ = self.0.set(key); } diff --git a/litellm-rust/crates/gateway-inference/src/error.rs b/litellm-rust/crates/gateway-inference/src/error.rs index 39656b933eb..ec82a6ede12 100644 --- a/litellm-rust/crates/gateway-inference/src/error.rs +++ b/litellm-rust/crates/gateway-inference/src/error.rs @@ -4,7 +4,7 @@ use axum::{ response::{IntoResponse, Response}, }; use litellm_http::transport::Error as TransportError; -use litellm_inference::RouteError; +use litellm_inference::Error as InferenceError; use litellm_llms::base_llm::ocr::error::Error as OcrError; use serde_json::{Map, Value, json}; @@ -19,7 +19,7 @@ pub enum Error { )] UnknownModel(String), #[error(transparent)] - Route(#[from] RouteError), + Route(#[from] InferenceError), #[error(transparent)] Ocr(#[from] OcrError), #[error("{0} is not implemented by the Rust gateway")] @@ -36,8 +36,8 @@ impl IntoResponse for Error { } } -impl From> for Error { - fn from(error: litellm_host_http::Error) -> Self { +impl From> for Error { + fn from(error: litellm_host_http::Error) -> Self { match error { litellm_host_http::Error::Call(error) => Self::Route(error), litellm_host_http::Error::Protocol => Self::Internal(error.to_string()), @@ -50,7 +50,7 @@ impl Error { match self { Self::Auth(error) => error.status(), Self::Unsupported(_) - | Self::Route(RouteError::Unsupported(_)) + | Self::Route(InferenceError::Unsupported(_)) | Self::Ocr(OcrError::Unsupported(_)) => StatusCode::NOT_IMPLEMENTED, Self::BodyTooLarge => StatusCode::PAYLOAD_TOO_LARGE, Self::Ocr( @@ -64,10 +64,10 @@ impl Error { .and_then(|status| StatusCode::from_u16(status).ok()) .unwrap_or(StatusCode::INTERNAL_SERVER_ERROR), Self::InvalidBody(_) | Self::UnknownModel(_) => StatusCode::BAD_REQUEST, - Self::Route(RouteError::Transport(TransportError::Http { status, .. })) => { + Self::Route(InferenceError::Transport(TransportError::Http { status, .. })) => { StatusCode::from_u16(*status).unwrap_or(StatusCode::BAD_GATEWAY) } - Self::Route(RouteError::Auth(litellm_auth::Error::MissingApiKey { .. })) => { + Self::Route(InferenceError::Auth(litellm_auth::Error::MissingApiKey { .. })) => { StatusCode::UNAUTHORIZED } Self::Route(error) if error.is_request() => StatusCode::BAD_REQUEST, @@ -97,7 +97,9 @@ impl Error { /// body already in that shape passes through, any other has its message extracted. pub fn body(&self, request_id: Option<&str>) -> Value { let raw = match self { - Self::Route(RouteError::Transport(TransportError::Http { body, .. })) => body.clone(), + Self::Route(InferenceError::Transport(TransportError::Http { body, .. })) => { + body.clone() + } other => other.to_string(), }; let parsed = serde_json::from_str::(&raw).ok(); @@ -175,7 +177,7 @@ mod tests { use super::*; fn upstream(status: u16, body: &str) -> Error { - Error::Route(RouteError::Transport(TransportError::Http { + Error::Route(InferenceError::Transport(TransportError::Http { status, body: body.into(), })) @@ -225,16 +227,16 @@ mod tests { #[rstest] #[case::upstream_status(upstream(429, ""), StatusCode::TOO_MANY_REQUESTS)] - #[case::rejected_request(Error::Route(RouteError::InvalidRequest("top_k".into())), StatusCode::BAD_REQUEST)] + #[case::rejected_request(Error::Route(InferenceError::InvalidRequest("top_k".into())), StatusCode::BAD_REQUEST)] #[case::missing_key( - Error::Route(RouteError::Auth(litellm_auth::Error::MissingApiKey { + Error::Route(InferenceError::Auth(litellm_auth::Error::MissingApiKey { provider: "Anthropic", environment_variable: "ANTHROPIC_API_KEY", })), StatusCode::UNAUTHORIZED, )] #[case::lost_connection( - Error::Route(RouteError::Transport(TransportError::Network("reset".into()))), + Error::Route(InferenceError::Transport(TransportError::Network("reset".into()))), StatusCode::INTERNAL_SERVER_ERROR, )] fn status_follows_who_is_at_fault(#[case] error: Error, #[case] status: StatusCode) { diff --git a/litellm-rust/crates/inference-chat/src/lib.rs b/litellm-rust/crates/inference-chat/src/lib.rs index f6b391e0ba1..3be83840901 100644 --- a/litellm-rust/crates/inference-chat/src/lib.rs +++ b/litellm-rust/crates/inference-chat/src/lib.rs @@ -1,7 +1,7 @@ use litellm_host::observation::ObservationSender; pub mod route; pub mod types; -pub use litellm_inference::RouteError as Error; +pub use litellm_inference::Error; mod common_utils; pub mod constants; pub(crate) mod handler; diff --git a/litellm-rust/crates/inference-chat/tests/caching.rs b/litellm-rust/crates/inference-chat/tests/caching.rs index 3fc432f8b46..70da8bc3255 100644 --- a/litellm-rust/crates/inference-chat/tests/caching.rs +++ b/litellm-rust/crates/inference-chat/tests/caching.rs @@ -21,7 +21,7 @@ use litellm_host::{ lifecycle::{CallEvent, ExecutionEvent}, observation::observation_channel, }; -use litellm_inference::RouteError; +use litellm_inference::Error; use rstest::{fixture, rstest}; use serde_json::{Value, json}; @@ -166,12 +166,12 @@ struct ChangingHooks { facts: std::sync::Mutex>, } -impl Interceptors for ChangingHooks { +impl Interceptors for ChangingHooks { async fn before_provider_request( &self, mut wire: WireRequest, _: RequestContext, - ) -> Result { + ) -> Result { let call = self.calls.fetch_add(1, Ordering::SeqCst); if self.rewrite { wire.body["temperature"] = json!(if call < 2 { 0.1 } else { 0.8 }); @@ -179,11 +179,11 @@ impl Interceptors for ChangingHooks { Ok(wire) } - async fn after_provider_response(&self, _: RawResponse) -> Result<(), RouteError> { + async fn after_provider_response(&self, _: RawResponse) -> Result<(), Error> { Ok(()) } - async fn result_ready(&self, facts: ExecutionFacts) -> Result<(), RouteError> { + async fn result_ready(&self, facts: ExecutionFacts) -> Result<(), Error> { self.facts.lock().unwrap().push(facts); Ok(()) } diff --git a/litellm-rust/crates/inference-messages/src/lib.rs b/litellm-rust/crates/inference-messages/src/lib.rs index f91d70445f3..bfefb9c63b0 100644 --- a/litellm-rust/crates/inference-messages/src/lib.rs +++ b/litellm-rust/crates/inference-messages/src/lib.rs @@ -14,7 +14,7 @@ use std::sync::Arc; use crate::common_utils::messages_provider; -pub use litellm_inference::RouteError as Error; +pub use litellm_inference::Error; pub use types::{ MessagesCall, MessagesCallResponse, MessagesSettings, MessagesShaping, litellm_params, messages_body, diff --git a/litellm-rust/crates/inference-messages/tests/caching.rs b/litellm-rust/crates/inference-messages/tests/caching.rs index 32c962091ad..480b3c4f3b7 100644 --- a/litellm-rust/crates/inference-messages/tests/caching.rs +++ b/litellm-rust/crates/inference-messages/tests/caching.rs @@ -17,7 +17,7 @@ use litellm_host::interceptors::{ ResultSource, WireRequest, }; use litellm_inference::{ - RouteError, + Error, caching::{CacheRequest, execute_unary}, }; use rstest::{fixture, rstest}; @@ -133,12 +133,12 @@ struct ChangingHooks { facts: std::sync::Mutex>, } -impl Interceptors for ChangingHooks { +impl Interceptors for ChangingHooks { async fn before_provider_request( &self, mut wire: WireRequest, _: RequestContext, - ) -> Result { + ) -> Result { let call = self.calls.fetch_add(1, Ordering::SeqCst); if self.rewrite { wire.body["temperature"] = json!(if call < 2 { 0.1 } else { 0.8 }); @@ -146,11 +146,11 @@ impl Interceptors for ChangingHooks { Ok(wire) } - async fn after_provider_response(&self, _: RawResponse) -> Result<(), RouteError> { + async fn after_provider_response(&self, _: RawResponse) -> Result<(), Error> { Ok(()) } - async fn result_ready(&self, facts: ExecutionFacts) -> Result<(), RouteError> { + async fn result_ready(&self, facts: ExecutionFacts) -> Result<(), Error> { self.facts.lock().unwrap().push(facts); Ok(()) } diff --git a/litellm-rust/crates/inference-ocr/AGENTS.md b/litellm-rust/crates/inference-ocr/AGENTS.md index d74b7e7af06..e08f561e75d 100644 --- a/litellm-rust/crates/inference-ocr/AGENTS.md +++ b/litellm-rust/crates/inference-ocr/AGENTS.md @@ -3,4 +3,4 @@ - The route returns its completed response directly - Route closures supply OCR-specific host capabilities, such as the caller's Azure AD token provider (`src/route.rs`) - Provider code reaches the caller's hooks mid-call only through `litellm_llms::base_llm::ocr::handler::CallHooks`, which this crate implements over its host until it folds into `litellm_host::interceptors::Interceptors` -- OCR failures use `base_llm/ocr/error.rs`, the recorded exception to `litellm_inference::RouteError` until OCR folds into `litellm_llms::Error` +- OCR failures use `base_llm/ocr/error.rs`, the recorded exception to `litellm_inference::Error` until OCR folds into `litellm_llms::Error` diff --git a/litellm-rust/crates/inference-ocr/src/provider_config.rs b/litellm-rust/crates/inference-ocr/src/provider_config.rs index 2f0346f6e47..09041e9b72b 100644 --- a/litellm-rust/crates/inference-ocr/src/provider_config.rs +++ b/litellm-rust/crates/inference-ocr/src/provider_config.rs @@ -249,7 +249,7 @@ pub(crate) fn resolve_provider_config<'a>( provider: LlmProviders::Mistral, } } - Err(litellm_inference::RouteError::InvalidProvider(provider)) => { + Err(litellm_inference::Error::InvalidProvider(provider)) => { return Err(Error::InvalidProvider(provider)); } Err(error) => return Err(Error::InvalidProvider(error.to_string())), diff --git a/litellm-rust/crates/inference-responses/src/lib.rs b/litellm-rust/crates/inference-responses/src/lib.rs index 6a9d5372311..b6e151ecec3 100644 --- a/litellm-rust/crates/inference-responses/src/lib.rs +++ b/litellm-rust/crates/inference-responses/src/lib.rs @@ -1,5 +1,5 @@ use litellm_host::observation::ObservationSender; -pub use litellm_inference::RouteError as Error; +pub use litellm_inference::Error; pub mod ids; pub mod route; diff --git a/litellm-rust/crates/inference-responses/tests/caching.rs b/litellm-rust/crates/inference-responses/tests/caching.rs index f45cefdc13b..462461695c3 100644 --- a/litellm-rust/crates/inference-responses/tests/caching.rs +++ b/litellm-rust/crates/inference-responses/tests/caching.rs @@ -18,7 +18,7 @@ use litellm_host::interceptors::{ WireRequest, }; use litellm_inference::{ - RouteError, + Error, caching::{CacheRequest, execute_unary}, }; use rstest::{fixture, rstest}; @@ -124,12 +124,12 @@ struct ChangingHooks { facts: std::sync::Mutex>, } -impl Interceptors for ChangingHooks { +impl Interceptors for ChangingHooks { async fn before_provider_request( &self, mut wire: WireRequest, _: RequestContext, - ) -> Result { + ) -> Result { let call = self.calls.fetch_add(1, Ordering::SeqCst); if self.rewrite { wire.body["temperature"] = json!(if call < 2 { 0.1 } else { 0.8 }); @@ -137,11 +137,11 @@ impl Interceptors for ChangingHooks { Ok(wire) } - async fn after_provider_response(&self, _: RawResponse) -> Result<(), RouteError> { + async fn after_provider_response(&self, _: RawResponse) -> Result<(), Error> { Ok(()) } - async fn result_ready(&self, facts: ExecutionFacts) -> Result<(), RouteError> { + async fn result_ready(&self, facts: ExecutionFacts) -> Result<(), Error> { self.facts.lock().unwrap().push(facts); Ok(()) } diff --git a/litellm-rust/crates/inference-transcription/src/lib.rs b/litellm-rust/crates/inference-transcription/src/lib.rs index 47e7bc70b31..f11e3abcd54 100644 --- a/litellm-rust/crates/inference-transcription/src/lib.rs +++ b/litellm-rust/crates/inference-transcription/src/lib.rs @@ -1,5 +1,5 @@ pub mod types; -pub use litellm_inference::RouteError as Error; +pub use litellm_inference::Error; mod constants; mod handler; diff --git a/litellm-rust/crates/inference/AGENTS.md b/litellm-rust/crates/inference/AGENTS.md index 2c4bd00108b..7f2bd0387be 100644 --- a/litellm-rust/crates/inference/AGENTS.md +++ b/litellm-rust/crates/inference/AGENTS.md @@ -2,7 +2,7 @@ - `gateway-inference` (serving) -> `inference-` (one API's call) -> `inference` (shared base) - This crate is the shared base every format crate builds on - - `RouteError` (`src/error.rs`), `CallOptions` (`src/lib.rs`), `CallContext` (`src/context.rs`) + - `Error` (`src/error.rs`), `CallOptions` (`src/lib.rs`), `CallContext` (`src/context.rs`) - diagnostic spans (`src/diagnostic.rs`), outbound send and signing (`src/outbound.rs`), provider resolution (`src/provider.rs`), `CoreResources` (`src/resources.rs`) - response caching (`src/caching.rs`): `Cachable`, `StreamCachable`, `CacheRequest`, `CallCache`, `execute_unary`, `execute_streaming`, stream capture - Shared test helpers belong in litellm-inference-testing, used only as a dev dependency @@ -49,14 +49,14 @@ - `litellm-inference-` mirrors the route packages (`litellm/ocr/`, `litellm/messages/`, ...) - Provider code never imports from `inference` or `inference-*` - Handlers belong in `inference-*` or `llms`, never in a host crate -- Import every item from its canonical path. Never re-export another crate's items or give an item a second public path. The only allowed re-exports are a private submodule surfacing its item at its module root (`mod error; pub use error::Error;`) and each format crate's `pub use litellm_inference::RouteError as Error;`, which keeps the pre-split `::Error` name +- Import every item from its canonical path. Never re-export another crate's items or give an item a second public path. The only allowed re-exports are a private submodule surfacing its item at its module root (`mod error; pub use error::Error;`) and each format crate's `pub use litellm_inference::Error;`, which keeps the pre-split `::Error` name ## Error placement - The workspace `Error definitions` rules shape each crate's error. This section decides which crate and module a failure belongs to - A failure is declared once, by the lowest crate that raises it - every crate above nests it unchanged (`#[error(transparent)] Auth(#[from] litellm_auth::Error)`) or maps it once at its boundary, as `src/error.rs` does for `litellm_llms::Error` - - `RouteError` collects route failures for every format crate and never re-declares a variant a lower crate raises + - `Error` collects route failures for every format crate and never re-declares a variant a lower crate raises - Scope follows the concept, not the first caller - an error type under `litellm-llms`'s `/` is private to that provider: no other provider and nothing in `base_llm` may import it - a failure two providers or two routes can hit (wire framing, stream event decoding, a malformed provider response) belongs to the crate that owns the concept: `litellm-framer` for framing, `litellm_llms::Error` for the transformation layer diff --git a/litellm-rust/crates/inference/src/caching.rs b/litellm-rust/crates/inference/src/caching.rs index 6dfda20ead4..09d03c06215 100644 --- a/litellm-rust/crates/inference/src/caching.rs +++ b/litellm-rust/crates/inference/src/caching.rs @@ -24,9 +24,9 @@ use serde::{Deserialize, Serialize, de::DeserializeOwned}; use serde_json::Value; use tokio_util::codec::Decoder; -use crate::RouteError; +use crate::Error; -pub trait Cachable: Protocol { +pub trait Cachable: Protocol { const SURFACE: &'static str; fn reusable(_response: &Self::Response) -> bool { @@ -145,15 +145,15 @@ pub async fn execute_unary( request: CacheRequest, cache: Option>, options: Option, - interceptors: &impl Interceptors, + interceptors: &impl Interceptors, observers: Option<&ObservationSender>, provider: F, -) -> Result +) -> Result where P: Cachable, P::Response: Serialize + DeserializeOwned, F: FnOnce() -> Fut, - Fut: Future>, + Fut: Future>, { let identity = request.identity.clone(); crate::diagnostic::provider(&identity.model, &identity.provider); @@ -191,15 +191,15 @@ pub async fn execute_streaming( request: CacheRequest, cache: Option>, options: Option, - interceptors: &impl Interceptors, + interceptors: &impl Interceptors, observers: Option<&ObservationSender>, provider: F, -) -> Result, RouteError> +) -> Result, Error> where P: StreamCachable, P::Response: Serialize + DeserializeOwned, F: FnOnce() -> Fut, - Fut: Future, RouteError>>, + Fut: Future, Error>>, { let identity = request.identity.clone(); crate::diagnostic::provider(&identity.model, &identity.provider); @@ -299,9 +299,9 @@ impl CallCache

{ } fn capture_stream( - chunks: futures_util::stream::BoxStream<'static, Result>, + chunks: futures_util::stream::BoxStream<'static, Result>, session: CacheSession, -) -> futures_util::stream::BoxStream<'static, Result> { +) -> futures_util::stream::BoxStream<'static, Result> { stream::try_unfold( (chunks, Some(Vec::::new()), session), |(mut chunks, captured, session)| async move { @@ -330,7 +330,7 @@ fn capture_stream( { session.store(entry).await; } - Ok::<_, RouteError>(None) + Ok::<_, Error>(None) } } }, @@ -369,9 +369,9 @@ fn successful_stream(text: &str, terminal: &str) -> bool { async fn publish( facts: ExecutionFacts, - interceptors: &impl Interceptors, + interceptors: &impl Interceptors, observers: Option<&ObservationSender>, -) -> Result<(), RouteError> { +) -> Result<(), Error> { if let Some(observers) = observers { observers.emit(CallEvent::Execution(ExecutionEvent::ResultReady { facts: facts.clone(), diff --git a/litellm-rust/crates/inference/src/context.rs b/litellm-rust/crates/inference/src/context.rs index 6fcbb3051b4..655b768dde1 100644 --- a/litellm-rust/crates/inference/src/context.rs +++ b/litellm-rust/crates/inference/src/context.rs @@ -5,7 +5,7 @@ use litellm_host::{ observation::ObservationSender, }; -use crate::{CallOptions, RouteError}; +use crate::{CallOptions, Error}; pub struct CallContext<'a, I> { pub interceptors: &'a I, @@ -13,7 +13,7 @@ pub struct CallContext<'a, I> { pub cache: CachePolicy, } -impl<'a, I: Interceptors> CallContext<'a, I> { +impl<'a, I: Interceptors> CallContext<'a, I> { pub fn new(interceptors: &'a I, options: CallOptions) -> Self { Self { interceptors, @@ -22,7 +22,7 @@ impl<'a, I: Interceptors> CallContext<'a, I> { } } - pub async fn result_ready(&self, facts: ExecutionFacts) -> Result<(), RouteError> { + pub async fn result_ready(&self, facts: ExecutionFacts) -> Result<(), Error> { if let Some(observers) = &self.observers { observers.emit(CallEvent::Execution(ExecutionEvent::ResultReady { facts: facts.clone(), @@ -31,7 +31,7 @@ impl<'a, I: Interceptors> CallContext<'a, I> { self.interceptors.result_ready(facts).await } - pub async fn response_received(&self, body: &str) -> Result<(), RouteError> { + pub async fn response_received(&self, body: &str) -> Result<(), Error> { let raw = RawResponse { body: body.to_owned(), }; @@ -43,6 +43,6 @@ impl<'a, I: Interceptors> CallContext<'a, I> { self.interceptors .after_provider_response(raw) .await - .map_err(RouteError::post_call) + .map_err(Error::post_call) } } diff --git a/litellm-rust/crates/inference/src/error.rs b/litellm-rust/crates/inference/src/error.rs index e5f02754f3d..895c8c55e01 100644 --- a/litellm-rust/crates/inference/src/error.rs +++ b/litellm-rust/crates/inference/src/error.rs @@ -11,7 +11,7 @@ use litellm_http::transport::Error as TransportError; use litellm_llms::Error as LlmError; #[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)] -pub enum RouteError { +pub enum Error { #[error("expected {expected}, got {actual}")] InvalidType { expected: &'static str, @@ -38,10 +38,10 @@ pub enum RouteError { #[error(transparent)] Secret(#[from] SecretError), #[error("post-call hook failed: {0}")] - PostCallHook(#[source] Arc), + PostCallHook(#[source] Arc), } -impl From for RouteError { +impl From for Error { fn from(fault: litellm_host::machine::MachineFault) -> Self { use litellm_host::machine::MachineFault; Self::InvalidRequest(match fault { @@ -51,7 +51,7 @@ impl From for RouteError { } } -impl RouteError { +impl Error { pub fn post_call(error: Self) -> Self { Self::PostCallHook(Arc::new(error)) } @@ -76,7 +76,7 @@ impl RouteError { } } -impl From for RouteError { +impl From for Error { fn from(error: LlmError) -> Self { match error { LlmError::InvalidType { expected, actual } => Self::InvalidType { expected, actual }, @@ -99,7 +99,7 @@ impl SecretError { } } -impl From for RouteError { +impl From for Error { fn from(error: litellm_secrets::Error) -> Self { Self::Secret(SecretError(Arc::new(error))) } @@ -115,22 +115,22 @@ impl Eq for SecretError {} #[cfg(test)] mod tests { - use super::RouteError; + use super::Error; use litellm_llms::{Error as LlmError, ErrorDetail}; use rstest::rstest; #[test] fn a_missing_api_key_is_the_environment_not_the_request() { assert!( - !RouteError::Auth(litellm_auth::Error::MissingApiKey { + !Error::Auth(litellm_auth::Error::MissingApiKey { provider: "Anthropic", environment_variable: "ANTHROPIC_API_KEY", }) .is_request() ); - assert!(RouteError::Auth(litellm_auth::Error::InvalidHeader).is_request()); - assert!(RouteError::InvalidRequest("top_k".into()).is_request()); - assert!(!RouteError::InvalidResponse("bad json".into()).is_request()); + assert!(Error::Auth(litellm_auth::Error::InvalidHeader).is_request()); + assert!(Error::InvalidRequest("top_k".into()).is_request()); + assert!(!Error::InvalidResponse("bad json".into()).is_request()); } #[rstest] #[case::request(true)] @@ -139,7 +139,7 @@ mod tests { let source = serde_json::from_str::("{").unwrap_err(); let source_message = source.to_string(); let detail = ErrorDetail::invalid("test payload", source); - let error = RouteError::from(if request { + let error = Error::from(if request { LlmError::InvalidRequest(detail) } else { LlmError::InvalidResponse(detail) diff --git a/litellm-rust/crates/inference/src/lib.rs b/litellm-rust/crates/inference/src/lib.rs index a168802e7ec..80771e206ca 100644 --- a/litellm-rust/crates/inference/src/lib.rs +++ b/litellm-rust/crates/inference/src/lib.rs @@ -7,7 +7,7 @@ pub mod outbound; pub mod provider; pub mod resources; -pub use error::RouteError; +pub use error::Error; #[derive(Clone, Default)] pub struct CallOptions { diff --git a/litellm-rust/crates/inference/src/provider.rs b/litellm-rust/crates/inference/src/provider.rs index 7e1269db1cf..305aca85f97 100644 --- a/litellm-rust/crates/inference/src/provider.rs +++ b/litellm-rust/crates/inference/src/provider.rs @@ -2,7 +2,7 @@ use litellm_core_utils::get_llm_provider_logic::{ CustomLlmProvider, LlmProviders, get_custom_llm_provider, }; -use crate::error::RouteError as Error; +use crate::error::Error; #[derive(Debug)] pub struct ResolvedProvider<'a> { diff --git a/litellm-rust/crates/inference/tests/caching.rs b/litellm-rust/crates/inference/tests/caching.rs index 1bd53695c30..a0a6a757565 100644 --- a/litellm-rust/crates/inference/tests/caching.rs +++ b/litellm-rust/crates/inference/tests/caching.rs @@ -25,7 +25,7 @@ use litellm_host::{ protocol::Protocol, }; use litellm_inference::{ - RouteError, + Error, caching::{Cachable, CacheRequest, StreamCachable, execute_streaming, execute_unary}, }; use rstest::{fixture, rstest}; @@ -36,7 +36,7 @@ struct TestRoute; impl Protocol for TestRoute { type Request = Value; type Response = Value; - type Error = RouteError; + type Error = Error; type HostCall = Infallible; type Chunk = Bytes; type StreamHead = (); @@ -204,7 +204,7 @@ async fn streamed( .chunks(3) .map(|bytes| Ok(Bytes::copy_from_slice(bytes))) .collect::>(); - let ending = fail.then_some(Err(RouteError::Unsupported("test transport failure"))); + let ending = fail.then_some(Err(Error::Unsupported("test transport failure"))); Ok(CallOutput::Stream { head: (), chunks: stream::iter(chunks.into_iter().chain(ending)).boxed(), @@ -215,7 +215,7 @@ async fn streamed( .unwrap() } -async fn consume(output: OutputOf) -> Result, RouteError> { +async fn consume(output: OutputOf) -> Result, Error> { let CallOutput::Stream { chunks, .. } = output else { panic!("expected a stream"); }; @@ -315,7 +315,7 @@ async fn a_provider_failure_never_populates_the_cache(cache: Arc for CacheHitAccounting { - async fn result_ready(&self, facts: ExecutionFacts) -> Result<(), RouteError> { +impl Interceptors for CacheHitAccounting { + async fn result_ready(&self, facts: ExecutionFacts) -> Result<(), Error> { self.calls.fetch_add(1, Ordering::SeqCst); let ResultSource::Cache { key } = facts.source else { panic!("expected cache source") @@ -482,7 +482,7 @@ impl Interceptors for CacheHitAccounting { ); self.key.set(key).unwrap(); if self.reject { - return Err(RouteError::Unsupported("cache accounting rejected")); + return Err(Error::Unsupported("cache accounting rejected")); } Ok(()) } @@ -491,11 +491,11 @@ impl Interceptors for CacheHitAccounting { &self, wire: WireRequest, _: RequestContext, - ) -> Result { + ) -> Result { Ok(wire) } - async fn after_provider_response(&self, _: RawResponse) -> Result<(), RouteError> { + async fn after_provider_response(&self, _: RawResponse) -> Result<(), Error> { Ok(()) } } @@ -571,7 +571,7 @@ async fn cache_hits_notify_accounting_once_and_propagate_its_failure( if reject { assert!(matches!( result, - Err(RouteError::Unsupported("cache accounting rejected")) + Err(Error::Unsupported("cache accounting rejected")) )); } else { assert_eq!(result.unwrap(), expected); @@ -590,7 +590,7 @@ async fn cache_hits_notify_accounting_once_and_propagate_its_failure( impl Protocol for UnaryTestRoute { type Request = Value; type Response = Value; - type Error = RouteError; + type Error = Error; type HostCall = Infallible; type Chunk = Infallible; type StreamHead = Infallible; diff --git a/litellm-rust/crates/python-bridge/src/errors.rs b/litellm-rust/crates/python-bridge/src/errors.rs index dc63f0d6353..4dab0297e73 100644 --- a/litellm-rust/crates/python-bridge/src/errors.rs +++ b/litellm-rust/crates/python-bridge/src/errors.rs @@ -1,5 +1,5 @@ use litellm_http::transport::Error as TransportError; -use litellm_inference::RouteError; +use litellm_inference::Error; use pyo3::{ exceptions::{PyRuntimeError, PyValueError}, prelude::*, @@ -19,9 +19,9 @@ pyo3::create_exception!( "The provider call was already issued and failed. Args are (status, message); status is 0 when there was no HTTP response." ); -pub(crate) fn route_error_to_pyerr(error: RouteError) -> PyErr { +pub(crate) fn route_error_to_pyerr(error: Error) -> PyErr { match error { - RouteError::Transport(TransportError::Http { status, body }) => { + Error::Transport(TransportError::Http { status, body }) => { RustUpstreamError::new_err((status, body)) } other => by_fault(other.is_request(), other.to_string()), @@ -42,13 +42,13 @@ mod tests { use super::*; #[rstest::rstest] - #[case::unsupported(RouteError::Unsupported("test capability"), true)] - #[case::invalid_provider(RouteError::InvalidProvider("unknown".into()), true)] - #[case::invalid_request(RouteError::InvalidRequest("empty messages".into()), true)] + #[case::unsupported(Error::Unsupported("test capability"), true)] + #[case::invalid_provider(Error::InvalidProvider("unknown".into()), true)] + #[case::invalid_request(Error::InvalidRequest("empty messages".into()), true)] #[case::connection(TransportError::Connect("unreachable".into()).into(), false)] #[case::network(TransportError::Network("timed out".into()).into(), false)] - #[case::invalid_response(RouteError::InvalidResponse("missing usage".into()), false)] - fn route_failures_are_terminal(#[case] error: RouteError, #[case] is_request: bool) { + #[case::invalid_response(Error::InvalidResponse("missing usage".into()), false)] + fn route_failures_are_terminal(#[case] error: Error, #[case] is_request: bool) { Python::initialize(); Python::attach(|py| { let failure = route_error_to_pyerr(error); @@ -86,14 +86,12 @@ mod tests { fn missing_api_key_stays_a_runtime_error_while_other_auth_failures_are_value_errors() { Python::initialize(); Python::attach(|py| { - let missing = - route_error_to_pyerr(RouteError::Auth(litellm_auth::Error::MissingApiKey { - provider: "Anthropic", - environment_variable: "ANTHROPIC_API_KEY", - })); + let missing = route_error_to_pyerr(Error::Auth(litellm_auth::Error::MissingApiKey { + provider: "Anthropic", + environment_variable: "ANTHROPIC_API_KEY", + })); assert!(missing.is_instance_of::(py)); - let invalid = - route_error_to_pyerr(RouteError::Auth(litellm_auth::Error::InvalidHeader)); + let invalid = route_error_to_pyerr(Error::Auth(litellm_auth::Error::InvalidHeader)); assert!(invalid.is_instance_of::(py)); }); } diff --git a/litellm-rust/crates/python-bridge/src/routes/inference.rs b/litellm-rust/crates/python-bridge/src/routes/inference.rs index 3b590a522fe..c9eb60e47ed 100644 --- a/litellm-rust/crates/python-bridge/src/routes/inference.rs +++ b/litellm-rust/crates/python-bridge/src/routes/inference.rs @@ -7,7 +7,7 @@ use litellm_core_utils::get_llm_provider_logic::get_custom_llm_provider; use litellm_host::{call::HostedMachine, protocol::Protocol}; use litellm_host_python::{PythonBinding, PythonHostCalls, from_py}; use litellm_http::transport::Error as TransportError; -use litellm_inference::RouteError; +use litellm_inference::Error; use litellm_secrets::source::SecretSource; use pyo3::{ exceptions::PyValueError, @@ -119,14 +119,14 @@ impl InferenceHost { public_response(py, self.module, response) } - pub fn error(&self, py: Python<'_>, error: RouteError) -> PyResult { - if let RouteError::Secret(source) = &error + pub fn error(&self, py: Python<'_>, error: Error) -> PyResult { + if let Error::Secret(source) = &error && let Some(original) = crate::secrets::python_error(py, source.source_error()) { return Ok(original); } let native = match error { - RouteError::Transport(TransportError::Http { status, body }) => { + Error::Transport(TransportError::Http { status, body }) => { let error = RustUpstreamError::new_err((status, body)); error .value(py) @@ -144,7 +144,7 @@ impl InferenceHost { } pub(super) trait InferenceRoute: Sized + 'static { - type Protocol: Protocol; + type Protocol: Protocol; const OPERATION: LoggingOperation; const SYNC_CALL_TYPE: &'static str; const ASYNC_CALL_TYPE: &'static str; From f2118fc7670ba3fa736d9b8f4deef9077a45a720 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Wed, 21 Oct 2026 19:09:52 -0700 Subject: [PATCH 14/14] fix(security): stop trusting the compromised eMudhra root CA (#47670) Update certifi, Rust trust roots, and the Docker CA bundle to remove emSign ECC Root CA - G3 Require the fixed certifi release in litellm-core and retain libssl3 in Docker runtimes for Prisma compatibility Co-authored-by: shin-berri --- Dockerfile | 6 +++--- backend/Dockerfile | 6 +++--- docker/Dockerfile.database | 6 +++--- docker/Dockerfile.non_root | 6 +++--- gateway/Dockerfile | 6 +++--- litellm-rust/Cargo.lock | 8 ++++---- litellm-rust/Cargo.toml | 2 +- migrations/Dockerfile | 6 +++--- packaging/litellm-core/pyproject.toml | 1 + pyproject.toml | 3 +++ tests/code_coverage_tests/liccheck.ini | 1 + uv.lock | 11 ++++++++--- 12 files changed, 36 insertions(+), 26 deletions(-) diff --git a/Dockerfile b/Dockerfile index 890c90fe449..306c455c9de 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,10 +1,10 @@ # syntax=docker/dockerfile:1.7 # Base image for building -ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:1d95114038f76513a9ace6fca107d5582b08c65981f81f61cb56bf7fd2ef216d +ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:81211e15a86ea1d5357f284482af45730c51de45eaf10ad9ce205d6bfc2df323 # Runtime image -ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:1d95114038f76513a9ace6fca107d5582b08c65981f81f61cb56bf7fd2ef216d +ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:81211e15a86ea1d5357f284482af45730c51de45eaf10ad9ce205d6bfc2df323 ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a # Pinned by digest like the other base images; bump explicitly on Node upgrades. ARG UI_BUILD_IMAGE=node:24.19-alpine3.24@sha256:d32cdf619f63fe0471182d08996dd516c6275bb5fd31ae06e55a570bd9e1ad43 @@ -141,7 +141,7 @@ USER root RUN echo "https://packages.wolfi.dev/os" >> /etc/apk/repositories # node (without npm) is required by the prisma CLI at runtime -RUN apk add --no-cache bash openssl tzdata nodejs python-3.13 libsndfile libevent +RUN apk add --no-cache bash openssl libssl3 tzdata nodejs python-3.13 libsndfile libevent COPY --from=pgbouncer-builder /usr/local/bin/pgbouncer /usr/local/bin/pgbouncer WORKDIR /app diff --git a/backend/Dockerfile b/backend/Dockerfile index 84c90e16f67..772f3d5b351 100644 --- a/backend/Dockerfile +++ b/backend/Dockerfile @@ -1,5 +1,5 @@ -ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:1d95114038f76513a9ace6fca107d5582b08c65981f81f61cb56bf7fd2ef216d -ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:1d95114038f76513a9ace6fca107d5582b08c65981f81f61cb56bf7fd2ef216d +ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:81211e15a86ea1d5357f284482af45730c51de45eaf10ad9ce205d6bfc2df323 +ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:81211e15a86ea1d5357f284482af45730c51de45eaf10ad9ce205d6bfc2df323 ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a FROM $UV_IMAGE AS uvbin @@ -109,7 +109,7 @@ ENV LITELLM_RELEASE_TAG=${LITELLM_RELEASE_TAG} USER root RUN for i in 1 2 3; do \ - apk add --no-cache bash openssl tzdata python-3.13 libsndfile libatomic && break; \ + apk add --no-cache bash openssl libssl3 tzdata python-3.13 libsndfile libatomic && break; \ [ $i = 3 ] && { echo "apk add failed after 3 retries" >&2; exit 1; }; \ sleep 5; \ done diff --git a/docker/Dockerfile.database b/docker/Dockerfile.database index 56210458532..f1fa64be630 100644 --- a/docker/Dockerfile.database +++ b/docker/Dockerfile.database @@ -1,10 +1,10 @@ # syntax=docker/dockerfile:1.7 # Base image for building -ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:1d95114038f76513a9ace6fca107d5582b08c65981f81f61cb56bf7fd2ef216d +ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:81211e15a86ea1d5357f284482af45730c51de45eaf10ad9ce205d6bfc2df323 # Runtime image -ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:1d95114038f76513a9ace6fca107d5582b08c65981f81f61cb56bf7fd2ef216d +ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:81211e15a86ea1d5357f284482af45730c51de45eaf10ad9ce205d6bfc2df323 ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a # Pinned by digest like the other base images; bump explicitly on Node upgrades. ARG UI_BUILD_IMAGE=node:24.19-alpine3.24@sha256:d32cdf619f63fe0471182d08996dd516c6275bb5fd31ae06e55a570bd9e1ad43 @@ -122,7 +122,7 @@ ENV LITELLM_RELEASE_TAG=${LITELLM_RELEASE_TAG} USER root # node (without npm) is required by the prisma CLI at runtime -RUN apk add --no-cache bash openssl tzdata nodejs python-3.13 libsndfile libevent +RUN apk add --no-cache bash openssl libssl3 tzdata nodejs python-3.13 libsndfile libevent COPY --from=pgbouncer-builder /usr/local/bin/pgbouncer /usr/local/bin/pgbouncer WORKDIR /app diff --git a/docker/Dockerfile.non_root b/docker/Dockerfile.non_root index 086c38e0d60..69da9fe149b 100644 --- a/docker/Dockerfile.non_root +++ b/docker/Dockerfile.non_root @@ -1,8 +1,8 @@ # syntax=docker/dockerfile:1.7 # Base images -ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:1d95114038f76513a9ace6fca107d5582b08c65981f81f61cb56bf7fd2ef216d -ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:1d95114038f76513a9ace6fca107d5582b08c65981f81f61cb56bf7fd2ef216d +ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:81211e15a86ea1d5357f284482af45730c51de45eaf10ad9ce205d6bfc2df323 +ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:81211e15a86ea1d5357f284482af45730c51de45eaf10ad9ce205d6bfc2df323 ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a # Pinned by digest like the other base images; bump explicitly on Node upgrades. ARG UI_BUILD_IMAGE=node:24.19-alpine3.24@sha256:d32cdf619f63fe0471182d08996dd516c6275bb5fd31ae06e55a570bd9e1ad43 @@ -134,7 +134,7 @@ RUN for i in 1 2 3; do \ apk upgrade --no-cache && break || sleep 5; \ done && \ for i in 1 2 3; do \ - apk add --no-cache python-3.13 bash openssl tzdata libsndfile nodejs libevent && break || sleep 5; \ + apk add --no-cache python-3.13 bash openssl libssl3 tzdata libsndfile nodejs libevent && break || sleep 5; \ done COPY --from=pgbouncer-builder /usr/local/bin/pgbouncer /usr/local/bin/pgbouncer diff --git a/gateway/Dockerfile b/gateway/Dockerfile index 051b1c30c8f..6af25df0c5d 100644 --- a/gateway/Dockerfile +++ b/gateway/Dockerfile @@ -1,5 +1,5 @@ -ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:1d95114038f76513a9ace6fca107d5582b08c65981f81f61cb56bf7fd2ef216d -ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:1d95114038f76513a9ace6fca107d5582b08c65981f81f61cb56bf7fd2ef216d +ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:81211e15a86ea1d5357f284482af45730c51de45eaf10ad9ce205d6bfc2df323 +ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:81211e15a86ea1d5357f284482af45730c51de45eaf10ad9ce205d6bfc2df323 ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a # Checksum from https://www.pgbouncer.org/downloads/ (the Wolfi repo only carries 1.24.x) ARG PGBOUNCER_VERSION=1.26.0 @@ -125,7 +125,7 @@ FROM $LITELLM_RUNTIME_IMAGE AS runtime USER root RUN for i in 1 2 3; do \ - apk add --no-cache bash openssl tzdata python-3.13 libsndfile libatomic libevent && break; \ + apk add --no-cache bash openssl libssl3 tzdata python-3.13 libsndfile libatomic libevent && break; \ [ $i = 3 ] && { echo "apk add failed after 3 retries" >&2; exit 1; }; \ sleep 5; \ done diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 98f36607fff..c868d3e4b8a 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -8202,18 +8202,18 @@ dependencies = [ [[package]] name = "webpki-root-certs" -version = "1.0.9" +version = "1.0.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b96554aa2acc8ccdb7e1c9a58a7a68dd5d13bccc69cd124cb09406db612a1c9b" +checksum = "58ec80e73fd8baf57bf0d9615554c75348bda990d6d5bbeab964dd068319fecb" dependencies = [ "rustls-pki-types", ] [[package]] name = "webpki-roots" -version = "1.0.9" +version = "1.0.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7dcd9d09a39985f5344844e66b0c530a33843579125f23e21e9f0f220850f22a" +checksum = "1c1c8e20c8a38fc7d65abfe74eb7280e4aa667e8eec8bc84996d2874810748d9" dependencies = [ "rustls-pki-types", ] diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index f04eb115071..9e5e0a96b73 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -124,7 +124,7 @@ moka = { version = "0.12.16", features = ["future"] } strum = { version = "0.28.0", features = ["derive"] } url = "2.5.8" percent-encoding = "2.3" -webpki-roots = "1" +webpki-roots = "1.0.10" time = { version = "0.3.53", features = ["parsing"] } criterion = "0.8.2" fancy-regex = "0.19.2" diff --git a/migrations/Dockerfile b/migrations/Dockerfile index b39a9bfa9bf..55f100d9e02 100644 --- a/migrations/Dockerfile +++ b/migrations/Dockerfile @@ -1,5 +1,5 @@ -ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:1d95114038f76513a9ace6fca107d5582b08c65981f81f61cb56bf7fd2ef216d -ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:1d95114038f76513a9ace6fca107d5582b08c65981f81f61cb56bf7fd2ef216d +ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:81211e15a86ea1d5357f284482af45730c51de45eaf10ad9ce205d6bfc2df323 +ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:81211e15a86ea1d5357f284482af45730c51de45eaf10ad9ce205d6bfc2df323 ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a FROM $UV_IMAGE AS uvbin @@ -118,7 +118,7 @@ FROM $LITELLM_RUNTIME_IMAGE AS runtime USER root RUN for i in 1 2 3; do \ - apk add --no-cache bash openssl tzdata python-3.13 nodejs libsndfile libatomic && break; \ + apk add --no-cache bash openssl libssl3 tzdata python-3.13 nodejs libsndfile libatomic && break; \ [ $i = 3 ] && { echo "apk add failed after 3 retries" >&2; exit 1; }; \ sleep 5; \ done diff --git a/packaging/litellm-core/pyproject.toml b/packaging/litellm-core/pyproject.toml index 2add6868bd9..890fa0fd773 100644 --- a/packaging/litellm-core/pyproject.toml +++ b/packaging/litellm-core/pyproject.toml @@ -30,6 +30,7 @@ dependencies = [ "pydantic>=2.12.0,<3.0.0; python_version >= '3.14'", "pydantic-settings>=2.14.1,<3.0", "jsonschema>=4.0.0,<5.0", + "certifi>=2026.10.19", "typing-extensions>=4.13.0,<5.0", ] diff --git a/pyproject.toml b/pyproject.toml index 9076760153a..3aefd911937 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -36,6 +36,7 @@ dependencies = [ "pydantic-settings>=2.14.1,<3.0", "jsonschema>=4.0.0,<5.0", "boto3>=1.43.1,<2.0", + "certifi>=2026.10.19", "typing-extensions>=4.13.0,<5.0", ] @@ -369,6 +370,8 @@ override-dependencies = [ default-groups = ["dev"] required-version = ">=0.10.9" exclude-newer = "3 days" +# TODO: ryan, remove after 2026-10-23, once certifi 2026.10.19 (drops the compromised eMudhra root) clears the 3-day window +exclude-newer-package = { certifi = "2026-10-20T03:00:00Z" } [tool.uv.sources] litellm-proxy-extras = { workspace = true } diff --git a/tests/code_coverage_tests/liccheck.ini b/tests/code_coverage_tests/liccheck.ini index e533907659b..cd951c6cd28 100644 --- a/tests/code_coverage_tests/liccheck.ini +++ b/tests/code_coverage_tests/liccheck.ini @@ -176,6 +176,7 @@ langgraph: >=1.0.10 # MIT License langgraph-prebuilt: >=1.0.8 # MIT License - https://github.com/langchain-ai/langgraph/blob/main/LICENSE hypothesis: >=6.165.10 # MPL 2.0 license pytest-rerunfailures: >=15.1 # MPL 2.0 license +certifi: >=2026.10.19 # MPL 2.0 license - https://github.com/certifi/python-certifi/blob/master/LICENSE pytest-recording: >=0.13.4 # MIT license expression: >=5.6.0 # MIT License - https://github.com/cognitedata/Expression/blob/main/LICENSE typing-extensions: >=4.13.0 # PSF-2.0 license - https://github.com/python/typing_extensions/blob/main/LICENSE diff --git a/uv.lock b/uv.lock index 03403828225..8d2a38c1806 100644 --- a/uv.lock +++ b/uv.lock @@ -13,6 +13,9 @@ resolution-markers = [ exclude-newer = "0001-01-01T00:00:00Z" # This has no effect and is included for backwards compatibility when using relative exclude-newer values. exclude-newer-span = "P3D" +[options.exclude-newer-package] +certifi = "2026-10-20T03:00:00Z" + [manifest] members = [ "litellm", @@ -998,11 +1001,11 @@ wheels = [ [[package]] name = "certifi" -version = "2026.4.22" +version = "2026.10.19" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/25/ee/6caf7a40c36a1220410afe15a1cc64993a1f864871f698c0f93acb72842a/certifi-2026.4.22.tar.gz", hash = "sha256:8d455352a37b71bf76a79caa83a3d6c25afee4a385d632127b6afb3963f1c580", size = 137077, upload-time = "2026-04-22T11:26:11.191Z" } +sdist = { url = "https://files.pythonhosted.org/packages/a6/bd/1436c50cc46699de0230fbaacc1d1de7bfec413176099949e9d5112e88f2/certifi-2026.10.19.tar.gz", hash = "sha256:035c1e85572b141bf1c3a3eff0e33a8eb74917522f2559075ff8d2fea780a1ae", size = 137598, upload-time = "2026-10-20T02:24:16.604Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/22/30/7cd8fdcdfbc5b869528b079bfb76dcdf6056b1a2097a662e5e8c04f42965/certifi-2026.4.22-py3-none-any.whl", hash = "sha256:3cb2210c8f88ba2318d29b0388d1023c8492ff72ecdde4ebdaddbb13a31b1c4a", size = 135707, upload-time = "2026-04-22T11:26:09.372Z" }, + { url = "https://files.pythonhosted.org/packages/c4/51/29046e48f693c54e5b7412216d447a1d8362c9653148d98c02cdf65fa733/certifi-2026.10.19-py3-none-any.whl", hash = "sha256:42d35340cc5b9ba78e818ddbff0c5b1c2b779de448a340fd1b989f5fa8670e1b", size = 136464, upload-time = "2026-10-20T02:24:15.072Z" }, ] [[package]] @@ -4611,6 +4614,7 @@ dependencies = [ { name = "anyio" }, { name = "async-timeout", marker = "python_full_version < '3.11'" }, { name = "boto3" }, + { name = "certifi" }, { name = "click" }, { name = "fastuuid" }, { name = "filelock" }, @@ -4881,6 +4885,7 @@ requires-dist = [ { name = "backports-zstd", marker = "python_full_version < '3.14' and extra == 'proxy'", specifier = ">=1.0.0,<2.0" }, { name = "boto3", specifier = ">=1.43.1,<2.0" }, { name = "boto3", marker = "extra == 'proxy'", specifier = ">=1.43.1,<2.0" }, + { name = "certifi", specifier = ">=2026.10.19" }, { name = "click", specifier = ">=8.0.0,<9.0" }, { name = "cryptography", marker = "extra == 'proxy'", specifier = ">=49.0.0,<51.0" }, { name = "ddtrace", marker = "extra == 'proxy-runtime'", specifier = ">=4.8.2,<5.0" },