diff --git a/.circleci/config.yml b/.circleci/config.yml index 0d442e32b0a..1edf22578ff 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -1291,7 +1291,7 @@ jobs: name: Install MCP SDK1 peer command: | uv venv --python 3.12 .venv-mcp-peer - uv pip install --python .venv-mcp-peer 'mcp==1.28.1' 'langchain-mcp-adapters==0.2.1' + uv pip install --no-config --exclude-newer "3 days" --python .venv-mcp-peer 'mcp==1.28.1' 'langchain-mcp-adapters==0.2.1' echo "export MCP_TEST_PEER_PYTHON=$PWD/.venv-mcp-peer/bin/python" >> "$BASH_ENV" - when: condition: 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/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/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/harness/endpoint.py b/litellm/harness/endpoint.py index 0abf7c11280..71e2742fa2e 100644 --- a/litellm/harness/endpoint.py +++ b/litellm/harness/endpoint.py @@ -37,6 +37,7 @@ from litellm.harness.context import GatewayTarget from litellm.harness.errors import HarnessError, HarnessInstallFailed from litellm.harness.types import Harness, Usage from litellm.llms.custom_httpx.http_handler import get_async_httpx_client +from litellm.types.litellm_params import STREAM_USAGE_INJECTED_FIELD from litellm.types.llms.custom_http import httpxSpecialProvider if TYPE_CHECKING: @@ -690,18 +691,21 @@ class ModelEndpoint: async def _invoke_sdk( self, route: str, - kwargs: dict[str, Any], # mutable-ok: injects stream_options into the SDK kwargs + kwargs: dict[str, Any], # mutable-ok: SDK call kwargs built by _sdk_kwargs ) -> object: if route == ROUTE_MESSAGES: return await litellm.anthropic.messages.acreate(**kwargs) if route == ROUTE_CHAT: - if kwargs.get("stream"): - stream_options = kwargs.get("stream_options") or {} # mutable-ok: empty default for a JSON field - kwargs["stream_options"] = { # mutable-ok: JSON field sent to litellm.acompletion - "include_usage": True, - **stream_options, + if not kwargs.get("stream"): + return await litellm.acompletion(**kwargs) + stream_options = kwargs.get("stream_options") or {} # mutable-ok: empty default for a JSON field + return await litellm.acompletion( + **{ + **kwargs, + "stream_options": {"include_usage": True, **stream_options}, + STREAM_USAGE_INJECTED_FIELD: "include_usage" not in stream_options, } - return await litellm.acompletion(**kwargs) + ) return await litellm.aresponses(**kwargs) async def _call_sdk(self, route: str, body: Mapping[str, object]) -> Response: 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/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 473666e5ba1..d75d35e4531 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -12,7 +12,7 @@ import subprocess import sys import time import traceback -from collections.abc import Awaitable, Callable, Iterable, Iterator, Mapping, Sequence +from collections.abc import Awaitable, Callable, Coroutine, Iterable, Iterator, Mapping, Sequence from dataclasses import dataclass from datetime import datetime as dt_object from functools import lru_cache @@ -138,6 +138,7 @@ from litellm.types.interactions import ( InteractionsAPIResponse, InteractionsAPIStreamingResponse, ) +from litellm.types.litellm_params import STREAM_USAGE_INJECTED_FIELD from litellm.types.llms.openai import ( AllMessageValues, Batch, @@ -574,6 +575,27 @@ class UncachedPrice: cost_breakdown: CostBreakdown | None +@dataclass(frozen=True, slots=True) +class BilledRoundsSpend: + usage: Usage + cost: float + cost_breakdown: CostBreakdown | None + + +@dataclass(frozen=True, slots=True) +class BilledRoundsDispatch: + """A parked success log for a stream bridged through chat completions that billed several rounds, and the spend + its failure log bills when a post-call hook rejects the response, as it does for a rejected one-round bridged + stream.""" + + dispatch: Coroutine[object, object, None] + spend: BilledRoundsSpend + + +def _fixed_margin(cost_breakdown: CostBreakdown | None) -> float: + return 0.0 if cost_breakdown is None else cost_breakdown.get("margin_fixed_amount", 0.0) + + def _timestamp_seconds(moment: object) -> float | None: if isinstance(moment, datetime.datetime): return moment.timestamp() @@ -582,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, \ @@ -595,6 +621,7 @@ class Logging(LiteLLMLoggingBaseClass): prometheusLogger, \ slack_app custom_pricing: bool = False + replaced_global_callback_keys: frozenset[str] = frozenset() stream_options = None litellm_request_debug: bool = False streamed_anthropic_message_id: str | None = None @@ -602,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: @@ -649,11 +677,11 @@ class Logging(LiteLLMLoggingBaseClass): litellm_call_id: str, function_id: str, litellm_trace_id: str | None = None, - dynamic_input_callbacks: list[str | Callable | CustomLogger] | None = None, + dynamic_input_callbacks: list[str | Callable[..., object] | CustomLogger] | None = None, dynamic_success_callbacks: list[str | Callable[..., object] | CustomLogger] | None = None, dynamic_async_success_callbacks: list[str | Callable[..., object] | CustomLogger] | None = None, - dynamic_failure_callbacks: list[str | Callable | CustomLogger] | None = None, - dynamic_async_failure_callbacks: list[str | Callable | CustomLogger] | None = None, + dynamic_failure_callbacks: list[str | Callable[..., object] | CustomLogger] | None = None, + dynamic_async_failure_callbacks: list[str | Callable[..., object] | CustomLogger] | None = None, applied_guardrails: list[str] | None = None, kwargs: dict | None = None, log_raw_request_response: bool = False, @@ -698,6 +726,7 @@ class Logging(LiteLLMLoggingBaseClass): self._pre_call_session_id: str = session_id_var.get() _sid: Final = kwargs.get("litellm_session_id") if kwargs else None self.litellm_session_id: str = str(_sid) if _sid else "" + self.stream_usage_injected: Final = kwargs is not None and kwargs.get(STREAM_USAGE_INJECTED_FIELD) is True # supports_correlation_logging is False for calls originating from the # sync client entry point (wrapper() in utils.py): a plain OS thread # has no per-call context isolation the way an asyncio Task does, and @@ -726,15 +755,17 @@ class Logging(LiteLLMLoggingBaseClass): self.raw_request_only = raw_request_only # Initialize dynamic callbacks - self.dynamic_input_callbacks: list[str | Callable | CustomLogger] | None = dynamic_input_callbacks + self.dynamic_input_callbacks: list[str | Callable[..., object] | CustomLogger] | None = dynamic_input_callbacks self.dynamic_success_callbacks: list[str | Callable[..., object] | CustomLogger] | None = ( dynamic_success_callbacks ) self.dynamic_async_success_callbacks: list[str | Callable[..., object] | CustomLogger] | None = ( dynamic_async_success_callbacks ) - self.dynamic_failure_callbacks: list[str | Callable | CustomLogger] | None = dynamic_failure_callbacks - self.dynamic_async_failure_callbacks: list[str | Callable | CustomLogger] | None = ( + self.dynamic_failure_callbacks: list[str | Callable[..., object] | CustomLogger] | None = ( + dynamic_failure_callbacks + ) + self.dynamic_async_failure_callbacks: list[str | Callable[..., object] | CustomLogger] | None = ( dynamic_async_failure_callbacks ) @@ -826,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 = [] @@ -851,10 +883,32 @@ class Logging(LiteLLMLoggingBaseClass): self.dynamic_async_failure_callbacks, callback ) + def replace_global_callback(self, callback: CustomLogger) -> None: + self.add_dynamic_callback(callback) + self.replaced_global_callback_keys = self.replaced_global_callback_keys | { + litellm.logging_callback_manager.get_custom_logger_key(callback) + } + + def add_dynamic_function_callbacks( + self, + *, + async_success: Callable[..., object], + failure: Callable[..., object], + async_failure: Callable[..., object], + ) -> None: + self.dynamic_async_success_callbacks = self._with_dynamic_callback( + self.dynamic_async_success_callbacks, async_success + ) + self.dynamic_failure_callbacks = self._with_dynamic_callback(self.dynamic_failure_callbacks, failure) + self.dynamic_async_failure_callbacks = self._with_dynamic_callback( + self.dynamic_async_failure_callbacks, async_failure + ) + @staticmethod def _with_dynamic_callback( - callbacks: Sequence[str | Callable | CustomLogger] | None, callback: CustomLogger - ) -> list[str | Callable | CustomLogger]: + callbacks: Sequence[str | Callable[..., object] | CustomLogger] | None, + callback: Callable[..., object] | CustomLogger, + ) -> list[str | Callable[..., object] | CustomLogger]: existing: Final = tuple(callbacks or ()) return [*existing, *(() if callback in existing else (callback,))] @@ -891,9 +945,9 @@ class Logging(LiteLLMLoggingBaseClass): def _process_dynamic_callback_list( self, - callback_list: list[str | Callable | CustomLogger] | None, + callback_list: list[str | Callable[..., object] | CustomLogger] | None, dynamic_callbacks_type: Literal["input", "success", "failure", "async_success", "async_failure"], - ) -> list[str | Callable | CustomLogger] | None: + ) -> list[str | Callable[..., object] | CustomLogger] | None: """ Helper function to initialize CustomLogger compatible callbacks in self.dynamic_* callbacks @@ -905,7 +959,7 @@ class Logging(LiteLLMLoggingBaseClass): if callback_list is None: return None - processed_list: Final[list[str | Callable | CustomLogger]] = [] + processed_list: Final[list[str | Callable[..., object] | CustomLogger]] = [] for callback in callback_list: if isinstance(callback, str) and callback in litellm._known_custom_logger_compatible_callbacks: for callback_instance in self._resolve_dynamic_callback_string(callback): @@ -1554,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 @@ -1562,11 +1618,12 @@ 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 = litellm.input_callback + (self.dynamic_input_callbacks or []) + callbacks: Final = ( + *self._without_replaced_global_callbacks(litellm.input_callback), + *(self.dynamic_input_callbacks or ()), + ) for callback in callbacks: try: if callback == "supabase" and supabaseClient is not None: @@ -1739,7 +1796,10 @@ class Logging(LiteLLMLoggingBaseClass): ) # Input Integration Logging -> If you want to log the fact that an attempt to call the model was made - callbacks: Final = litellm.input_callback + (self.dynamic_input_callbacks or []) + callbacks: Final = ( + *self._without_replaced_global_callbacks(litellm.input_callback), + *(self.dynamic_input_callbacks or ()), + ) for callback in callbacks: try: if callback == "sentry" and add_breadcrumb: @@ -2388,6 +2448,31 @@ class Logging(LiteLLMLoggingBaseClass): self.model_call_details["combined_usage_object"] = usage self.model_call_details["response_cost"] = response_cost + def record_failure_after_earlier_rounds( + self, earlier: BilledRoundsSpend, failed_usage: Usage, failed_cost: float + ) -> None: + """Bill a failed round on the failure log after the rounds of its request before it were billed on the success + log. + + The fixed margin is a per-request fee, so it is charged once: the failed round's margin is dropped when the + earlier rounds already paid one.""" + from litellm.responses.mcp.round_billing import breakdown_less_repeated_fixed_margin + + failed_breakdown: Final = None if self.cost_breakdown is earlier.cost_breakdown else self.cost_breakdown + repeated_fixed_margin: Final = ( + 0.0 + if failed_breakdown is None + else min(_fixed_margin(earlier.cost_breakdown), _fixed_margin(failed_breakdown), earlier.cost, failed_cost) + ) + self.model_call_details["combined_usage_object"] = failed_usage + self.model_call_details["response_cost"] = failed_cost - repeated_fixed_margin + self.cost_breakdown = breakdown_less_repeated_fixed_margin(failed_breakdown, repeated_fixed_margin) + + def record_billed_rounds_for_failure(self, spend: BilledRoundsSpend) -> None: + """Bill every round of a fully streamed request on the failure log when a post-call hook rejects it.""" + self.record_partial_usage_for_failure(spend.usage, spend.cost) + self.cost_breakdown = spend.cost_breakdown + def attempt_scoped_copy(self) -> "Logging": """A copy with its own model_call_details, so logging one attempt of a request through it leaves the state the request's other attempts share untouched.""" @@ -2395,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), @@ -2514,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( @@ -4331,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, @@ -4516,14 +4631,26 @@ class Logging(LiteLLMLoggingBaseClass): _filtered_failure_callbacks = self._remove_internal_litellm_callbacks(_filtered_failure_callbacks) return len(_filtered_failure_callbacks) > 0 + def _without_replaced_global_callbacks( + self, global_callbacks: Iterable[str | Callable[..., object] | CustomLogger] + ) -> tuple[str | Callable[..., object] | CustomLogger, ...]: + if not self.replaced_global_callback_keys: + return tuple(global_callbacks) + return tuple( + litellm.logging_callback_manager.without_custom_logger_keys( + global_callbacks, self.replaced_global_callback_keys + ) + ) + def get_combined_callback_list( self, dynamic_success_callbacks: list[str | Callable[..., object] | CustomLogger] | None, global_callbacks: list[str | Callable[..., object] | CustomLogger], ) -> list[str | Callable[..., object] | CustomLogger]: + kept_global_callbacks: Final = self._without_replaced_global_callbacks(global_callbacks) if dynamic_success_callbacks is None: - return list(global_callbacks) - return list(dict.fromkeys(dynamic_success_callbacks + global_callbacks)) + return list(kept_global_callbacks) + return list(dict.fromkeys((*dynamic_success_callbacks, *kept_global_callbacks))) def _remove_internal_litellm_callbacks(self, callbacks: list) -> list: """ diff --git a/litellm/litellm_core_utils/logging_callback_manager.py b/litellm/litellm_core_utils/logging_callback_manager.py index 9f521ffc6cc..307353d081b 100644 --- a/litellm/litellm_core_utils/logging_callback_manager.py +++ b/litellm/litellm_core_utils/logging_callback_manager.py @@ -1,4 +1,6 @@ -from collections.abc import Callable +import threading +from collections.abc import Callable, Iterable, Iterator +from types import MethodType from typing import TYPE_CHECKING, Final import litellm @@ -15,6 +17,7 @@ else: _custom_logger_compatible_callbacks_literal = str _generic_api_logger_cache: Final[dict[str, GenericAPILogger]] = {} +_callback_lists_lock: Final = threading.RLock() class LoggingCallbackManager: @@ -43,16 +46,18 @@ class LoggingCallbackManager: Add a input callback to litellm.input_callback. Auto-routes async callbacks to litellm._async_input_callback. """ - if not isinstance(callback, str) and self._is_async_callable(callback): - self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_input_callback) - else: - self._safe_add_callback_to_list(callback=callback, parent_list=litellm.input_callback) + with _callback_lists_lock: + if not isinstance(callback, str) and self._is_async_callable(callback): + self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_input_callback) + else: + self._safe_add_callback_to_list(callback=callback, parent_list=litellm.input_callback) def add_litellm_service_callback(self, callback: CustomLogger | str | Callable[..., object]): """ Add a service callback to litellm.service_callback """ - self._safe_add_callback_to_list(callback=callback, parent_list=litellm.service_callback) + with _callback_lists_lock: + self._safe_add_callback_to_list(callback=callback, parent_list=litellm.service_callback) def add_litellm_callback(self, callback: CustomLogger | str | Callable[..., object]): """ @@ -60,10 +65,11 @@ class LoggingCallbackManager: Ensures no duplicates are added. """ - self._safe_add_callback_to_list( - callback=callback, - parent_list=litellm.callbacks, - ) + with _callback_lists_lock: + self._safe_add_callback_to_list( + callback=callback, + parent_list=litellm.callbacks, + ) def add_litellm_success_callback(self, callback: CustomLogger | str | Callable[..., object]): """ @@ -71,37 +77,41 @@ class LoggingCallbackManager: Auto-routes async callbacks to litellm._async_success_callback. Special-cases 'dynamodb' and 'openmeter' as async callbacks. """ - if ( - isinstance(callback, str) - and callback in ("dynamodb", "openmeter") - or not isinstance(callback, str) - and self._is_async_callable(callback) - ): - self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_success_callback) - else: - self._safe_add_callback_to_list(callback=callback, parent_list=litellm.success_callback) + with _callback_lists_lock: + if ( + isinstance(callback, str) + and callback in ("dynamodb", "openmeter") + or not isinstance(callback, str) + and self._is_async_callable(callback) + ): + self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_success_callback) + else: + self._safe_add_callback_to_list(callback=callback, parent_list=litellm.success_callback) def add_litellm_failure_callback(self, callback: CustomLogger | str | Callable[..., object]): """ Add a failure callback to `litellm.failure_callback`. Auto-routes async callbacks to litellm._async_failure_callback. """ - if not isinstance(callback, str) and self._is_async_callable(callback): - self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_failure_callback) - else: - self._safe_add_callback_to_list(callback=callback, parent_list=litellm.failure_callback) + with _callback_lists_lock: + if not isinstance(callback, str) and self._is_async_callable(callback): + self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_failure_callback) + else: + self._safe_add_callback_to_list(callback=callback, parent_list=litellm.failure_callback) def add_litellm_async_success_callback(self, callback: CustomLogger | Callable[..., object] | str): """ Add a success callback to litellm._async_success_callback """ - self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_success_callback) + with _callback_lists_lock: + self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_success_callback) def add_litellm_async_failure_callback(self, callback: CustomLogger | Callable[..., object] | str): """ Add a failure callback to litellm._async_failure_callback """ - self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_failure_callback) + with _callback_lists_lock: + self._safe_add_callback_to_list(callback=callback, parent_list=litellm._async_failure_callback) def remove_callback_from_list_by_object(self, callback_list, obj, require_self=True): """ @@ -110,13 +120,14 @@ class LoggingCallbackManager: if not isinstance(callback_list, list): # Not list -> do nothing return - if require_self: - remove_list = [c for c in callback_list if hasattr(c, "__self__") and c.__self__ == obj] - else: - remove_list = [c for c in callback_list if c == obj] + with _callback_lists_lock: + if require_self: + remove_list = [c for c in callback_list if hasattr(c, "__self__") and c.__self__ == obj] + else: + remove_list = [c for c in callback_list if c == obj] - for c in remove_list: - callback_list.remove(c) + for c in remove_list: + callback_list.remove(c) def remove_callbacks_by_type(self, callback_list, callback_type): """ @@ -134,10 +145,11 @@ class LoggingCallbackManager: if not isinstance(callback_list, list): return - remove_list: Final = [c for c in callback_list if isinstance(c, callback_type)] + with _callback_lists_lock: + remove_list: Final = [c for c in callback_list if isinstance(c, callback_type)] - for c in remove_list: - callback_list.remove(c) + for c in remove_list: + callback_list.remove(c) def _add_string_callback_to_list( self, callback: str, parent_list: list[CustomLogger | Callable[..., object] | str] @@ -336,6 +348,55 @@ class LoggingCallbackManager: return "-".join(key_parts) + get_custom_logger_key = _get_custom_logger_key + callback_lists_lock: Final = _callback_lists_lock + + def drop_callbacks(self, callbacks: Iterable[object], method_owners: Iterable[object] = ()) -> frozenset[str]: + dropped_ids: Final = frozenset(map(id, callbacks)) + owner_ids: Final = frozenset(map(id, method_owners)) + + def keep(callback: object) -> bool: + if id(callback) in dropped_ids: + return False + return not (isinstance(callback, MethodType) and id(callback.__self__) in owner_ids) + + with _callback_lists_lock: + vacated_keys: Final = frozenset( + self._get_custom_logger_key(callback) + for callback in litellm.callbacks + if isinstance(callback, CustomLogger) and id(callback) in dropped_ids + ) + litellm.callbacks = [callback for callback in litellm.callbacks if keep(callback)] + litellm.input_callback = [callback for callback in litellm.input_callback if keep(callback)] + litellm.success_callback = [callback for callback in litellm.success_callback if keep(callback)] + litellm.failure_callback = [callback for callback in litellm.failure_callback if keep(callback)] + litellm.service_callback = [callback for callback in litellm.service_callback if keep(callback)] + litellm._async_success_callback = [ + callback for callback in litellm._async_success_callback if keep(callback) + ] + litellm._async_failure_callback = [ + callback for callback in litellm._async_failure_callback if keep(callback) + ] + return vacated_keys + + def with_custom_logger_keys( + self, candidates: Iterable[CustomLogger], keys: frozenset[str] + ) -> tuple[CustomLogger, ...]: + if not keys: + return () + class_names: Final = frozenset(key.split("-", 1)[0] for key in keys) + same_class: Final = (candidate for candidate in candidates if type(candidate).__name__ in class_names) + return tuple(candidate for candidate in same_class if self._get_custom_logger_key(candidate) in keys) + + def without_custom_logger_keys( + self, callbacks: Iterable[str | Callable[..., object] | CustomLogger], keys: frozenset[str] + ) -> Iterator[str | Callable[..., object] | CustomLogger]: + return ( + callback + for callback in callbacks + if not (isinstance(callback, CustomLogger) and self._get_custom_logger_key(callback) in keys) + ) + def _reset_all_callbacks(self): """ Reset all callbacks to an empty list @@ -363,19 +424,21 @@ class LoggingCallbackManager: _get_all_callbacks = get_all_callbacks - def remove_callback_from_all_lists(self, obj, require_self=False) -> None: + def remove_callback_from_all_lists(self, obj: object, require_self: bool = False) -> None: """ Remove a callback object from every callback list it may have been promoted into, so a re-initialized callback leaves no stale instance behind. """ - for callback_list in ( - litellm.callbacks, - litellm.success_callback, - litellm.failure_callback, - litellm._async_success_callback, - litellm._async_failure_callback, - ): - self.remove_callback_from_list_by_object(callback_list, obj, require_self=require_self) + with _callback_lists_lock: + for callback_list in ( + litellm.callbacks, + litellm.input_callback, + litellm.success_callback, + litellm.failure_callback, + litellm._async_success_callback, + litellm._async_failure_callback, + ): + self.remove_callback_from_list_by_object(callback_list, obj, require_self=require_self) def get_active_additional_logging_utils_from_custom_logger( self, 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/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 6ec426d6da5..f994baa52b1 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -10,7 +10,7 @@ import traceback from collections.abc import AsyncIterator, Callable, Iterable, Iterator, Mapping, Sequence from dataclasses import dataclass from types import MappingProxyType -from typing import Any, Final, NoReturn, Protocol, TypeVar, cast +from typing import TYPE_CHECKING, Any, Final, NoReturn, Protocol, TypeVar, cast import anyio import httpx @@ -60,6 +60,9 @@ from .exception_mapping_utils import exception_type, with_deployment_api_base from .llm_response_utils.get_api_base import get_api_base from .rules import Rules +if TYPE_CHECKING: + from .litellm_logging import BilledRoundsSpend + # Constants for special delta attribute names AUDIO_ATTRIBUTE: Final = "audio" IMAGE_ATTRIBUTE: Final = "images" @@ -259,6 +262,8 @@ def _provider_hidden_params( class CustomStreamWrapper: + _earlier_rounds_spend: "BilledRoundsSpend | None" = None + @property def hidden_params(self) -> dict[str, object]: # mutable-ok: API requires mutation return self._hidden_params @@ -342,6 +347,10 @@ class CustomStreamWrapper: self.messages = getattr(logging_obj, "messages", None) self.sent_stream_usage = False self.send_stream_usage = True if self.check_send_stream_usage(self.stream_options) else False + self.caller_requested_stream_usage: Final = ( + self.check_send_stream_usage(getattr(logging_obj, "stream_options", None)) + and getattr(logging_obj, "stream_usage_injected", False) is not True + ) self.tool_call = False # keep track of the returned chunks - used for calculating the input/output tokens for stream options self.chunks: list[ModelResponseStream] = [] @@ -2070,7 +2079,7 @@ class CustomStreamWrapper: except StopIteration: if self.sent_last_chunk is True: complete_streaming_response: Final = self._build_complete_streaming_response() - response = self.model_response_creator() + response = self._new_usage_chunk() if complete_streaming_response is not None: self._propagate_usage_cost_to_hidden_params(complete_streaming_response, self.custom_llm_provider) @@ -2222,27 +2231,10 @@ class CustomStreamWrapper: ) if _has_usage: - # Store a copy ONLY when usage stripping below will mutate - # the chunk. For non-usage chunks (vast majority), store - # directly to avoid expensive model_copy() per chunk. - self.chunks.append(processed_chunk.model_copy()) - - # Strip usage from the outgoing chunk so it's not sent twice - # (once in the chunk, once in _hidden_params). - obj_dict = processed_chunk.model_dump() - if "usage" in obj_dict: - del obj_dict["usage"] - processed_chunk = self.model_response_creator( - chunk=obj_dict, - hidden_params=cast( # cast-ok: preserve mapping operations on dynamic storage - Mapping[str, object], getattr(processed_chunk, HIDDEN_PARAMS_ATTR) - ), - ) - is_empty = is_model_response_stream_empty( - model_response=cast(ModelResponseStream, processed_chunk) - ) - if is_empty: + usage_free_chunk = self._record_and_strip_usage_chunk(processed_chunk) + if usage_free_chunk is None: continue + processed_chunk = usage_free_chunk else: # No usage data — safe to store directly without copying self.chunks.append(processed_chunk) @@ -2273,25 +2265,25 @@ class CustomStreamWrapper: if processed_chunk is None: continue - choice = processed_chunk.choices[0] - if isinstance(choice, StreamingChoices): - self._accumulate_streamed_delta(choice.delta) + if processed_chunk.choices: + choice = processed_chunk.choices[0] + if isinstance(choice, StreamingChoices): + self._accumulate_streamed_delta(choice.delta) self.rules.post_call_rules(input=self.response_uptil_now, model=self.model) - # RETURN RESULT - self.chunks.append(processed_chunk) - return processed_chunk + if not self.send_stream_usage or getattr(processed_chunk, "usage", None) is None: + self.chunks.append(processed_chunk) + return processed_chunk + usage_free_chunk = self._record_and_strip_usage_chunk(processed_chunk) + if usage_free_chunk is None: + continue + return usage_free_chunk except (StopAsyncIteration, StopIteration): return await self._finalize_completed_stream(cache_hit=cache_hit) except httpx.TimeoutException as e: # if httpx read timeout error occues 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: @@ -2306,7 +2298,7 @@ class CustomStreamWrapper: async def _finalize_completed_stream(self, cache_hit: bool) -> "ModelResponseStream": if self.sent_last_chunk is True: complete_streaming_response: Final = await self._async_build_complete_streaming_response() - response: Final = self.model_response_creator() + response: Final = self._new_usage_chunk() if complete_streaming_response is not None: self._propagate_usage_cost_to_hidden_params(complete_streaming_response, self.custom_llm_provider) @@ -2392,29 +2384,36 @@ 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 await asyncio.wait((self._failure_logging,), timeout=timeout) + def bill_failure_after_earlier_rounds(self, earlier: "BilledRoundsSpend") -> None: + self._earlier_rounds_spend = earlier + def _record_partial_usage_for_failure(self) -> None: """ 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 @@ -2433,10 +2432,12 @@ class CustomStreamWrapper: if self.model: partial_response.model = self.model backfill_missing_cache_usage_fields(usage) - self.logging_obj.model_call_details["combined_usage_object"] = usage - self.logging_obj.model_call_details["response_cost"] = ( - self.logging_obj.response_cost_calculator(result=partial_response) or 0.0 - ) + partial_cost: Final = self.logging_obj.response_cost_calculator(result=partial_response) or 0.0 + if self._earlier_rounds_spend is None: + self.logging_obj.model_call_details["combined_usage_object"] = usage + self.logging_obj.model_call_details["response_cost"] = partial_cost + return + self.logging_obj.record_failure_after_earlier_rounds(self._earlier_rounds_spend, usage, partial_cost) except Exception as recover_error: verbose_logger.debug( "could not recover partial usage for interrupted stream: %s", @@ -2463,6 +2464,22 @@ class CustomStreamWrapper: return False return any(get(field) for field in cls._CONTINUATION_DISQUALIFYING_DELTA_FIELDS) + def _new_usage_chunk(self) -> ModelResponseStream: + usage_chunk: Final = self.model_response_creator() + if self.caller_requested_stream_usage: + usage_chunk.choices = [] + return usage_chunk + + def _record_and_strip_usage_chunk(self, chunk: ModelResponseStream) -> ModelResponseStream | None: + self.chunks.append(chunk.model_copy()) + usage_free_chunk: Final = self.model_response_creator( + chunk={key: value for key, value in chunk.model_dump().items() if key != "usage"}, + hidden_params=cast( # cast-ok: preserve mapping operations on dynamic storage + Mapping[str, object], getattr(chunk, HIDDEN_PARAMS_ATTR) + ), + ) + return None if is_model_response_stream_empty(model_response=usage_free_chunk) else usage_free_chunk + def _accumulate_streamed_delta(self, delta: object) -> None: # Shared by the sync, async, and non-aiohttp iteration sites so answer # text and the disqualifying-content latch stay in step across all three. 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/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/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/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/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/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/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 176569f5d8d..e36aba00293 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,11 +31,13 @@ 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, 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 @@ -1081,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 [] @@ -1103,6 +1106,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 +1125,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 +1161,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/long_decode.py b/litellm/llms/vertex_ai/gemini/long_decode.py index 81d6c04d0e8..571fa94d4a5 100644 --- a/litellm/llms/vertex_ai/gemini/long_decode.py +++ b/litellm/llms/vertex_ai/gemini/long_decode.py @@ -2,17 +2,20 @@ import json from collections.abc import Awaitable, Callable, Mapping, Sequence from dataclasses import dataclass from itertools import accumulate, chain -from typing import Final, TypeAlias, TypeVar, cast # noqa: TID251 # sum_usage_metadata casts computed keys +from typing import Final, Literal, TypeAlias, TypeVar, cast # noqa: TID251 # sum_usage_metadata casts computed keys import httpx +from pydantic import TypeAdapter, ValidationError from litellm._logging import verbose_logger from litellm.exceptions import Timeout +from litellm.litellm_core_utils.env_utils import get_env_int_or_none from litellm.types.llms.vertex_ai import ( Candidates, Citation, CitationMetadata, GenerateContentResponseBody, + GenerationConfig, GroundingMetadata, GroundingSegment, GroundingSupport, @@ -48,8 +51,14 @@ HOP_FAILURES: Final = (VertexAIError, httpx.HTTPError, Timeout) HOP_END_FINISH_REASON: Final = "CONTINUATION" CUT_OFF_FINISH_REASONS: Final = frozenset({"MAX_TOKENS", HOP_END_FINISH_REASON}) +DEFAULT_OUTPUT_BUDGET_ENV: Final = "LITELLM_GEMINI_DEFAULT_MAX_OUTPUT_TOKENS" +_OUTPUT_LIMIT_KEYS: Final = frozenset({"max_output_tokens", "maxOutputTokens"}) +_REQUEST_KEYS_THAT_SKIP_THE_DEFAULT_BUDGET: Final = frozenset({"generation_config", "continuationToken"}) +_TEXT_ONLY_MODALITIES: Final = TypeAdapter(list[Literal["TEXT"]]) + _VERTEX_V1_PATH: Final = "/v1/projects/" _VERTEX_PATH_ACCEPTING_CONTINUATION_TOKEN: Final = "/v1beta1/projects/" +_CANDIDATE_COUNT_KEYS: Final = ("candidate_count", "candidateCount") LongDecodeHops: TypeAlias = tuple[GenerateContentResponseBody, ...] AsyncHopPost: TypeAlias = Callable[[RequestBody], Awaitable[httpx.Response]] @@ -76,15 +85,53 @@ def _int_or_none(value: object) -> int | None: return value if isinstance(value, int) else None +def _asks_for_one_candidate(generation_config: Mapping[str, object]) -> bool: + return all( + _int_or_none(generation_config[key] or 1) == 1 for key in _CANDIDATE_COUNT_KEYS if key in generation_config + ) + + def continuation_cap_of(request_body: RequestBody) -> int | None: generation_config: Final = request_body.get("generationConfig") if generation_config is None or "max_output_tokens" not in generation_config: return None - if "candidate_count" in generation_config and _int_or_none(generation_config["candidate_count"]) != 1: + if not _asks_for_one_candidate(generation_config): return None return _int_or_none(generation_config["max_output_tokens"]) +def _asks_for_text_only(generation_config: Mapping[str, object]) -> bool: + if "response_modalities" in generation_config: + return False + try: + _ = _TEXT_ONLY_MODALITIES.validate_python(generation_config.get("responseModalities", [])) + except ValidationError: + return False + return True + + +def with_default_output_budget( + request_body: RequestBody, model_max_output_tokens: Callable[[], int | None] +) -> RequestBody: + generation_config: Final = request_body.get("generationConfig", GenerationConfig()) + request_json: Final[Mapping[str, object]] = request_body + budget: Final = get_env_int_or_none(DEFAULT_OUTPUT_BUDGET_ENV) + if budget is None or budget <= 0: + return request_body + if not _REQUEST_KEYS_THAT_SKIP_THE_DEFAULT_BUDGET.isdisjoint(request_body): + return request_body + if not isinstance(request_json.get("generationConfig", {}), dict): + return request_body + if not _OUTPUT_LIMIT_KEYS.isdisjoint(generation_config): + return request_body + if not _asks_for_text_only(generation_config): + return request_body + model_max: Final = model_max_output_tokens() + if model_max is None: + return request_body + return {**request_body, "generationConfig": {**generation_config, "max_output_tokens": min(budget, model_max)}} + + def client_continuation_token(candidate: Candidates) -> str | None: if candidate.get("finishReason") not in CUT_OFF_FINISH_REASONS: return None diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index e5ce395031a..71bb05ee155 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -7,16 +7,19 @@ Why separate file? Make it easy to see how transformation works import os import re from collections.abc import Callable, Mapping +from functools import partial from typing import TYPE_CHECKING, Any, Final, Literal, cast from urllib.parse import quote import httpx -from pydantic import BaseModel, TypeAdapter +from pydantic import BaseModel, PositiveInt, TypeAdapter, ValidationError +from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.core_helpers import drops_unsupported_params +from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_image_mime_type_from_url, ) @@ -65,6 +68,7 @@ from ..common_utils import ( get_supports_response_schema, get_supports_system_message, ) +from .long_decode import with_default_output_budget # VertexBase is imported only for type checking to avoid introducing a # module-load-time cyclic import to vertex_llm_base. The instance is lazily @@ -1304,13 +1308,75 @@ def sync_transform_request_body( vertex_auth_header=vertex_auth_header, ) - return transform_request_body( - messages=messages, - model=model, - custom_llm_provider=custom_llm_provider, - litellm_params=litellm_params, - cached_content=cached_content, - optional_params=optional_params, + return _with_default_output_budget( + transform_request_body( + messages=messages, + model=model, + custom_llm_provider=custom_llm_provider, + litellm_params=litellm_params, + cached_content=cached_content, + optional_params=optional_params, + ), + model, + custom_llm_provider, + litellm_params.get("model_info"), # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # litellm_params is an untyped dict + ) + + +class _CostMapOutputModalities(TypedDict): + supported_output_modalities: NotRequired[ReadOnly[tuple[str, ...]]] + + +_COST_MAP_OUTPUT_MODALITIES: Final = TypeAdapter(_CostMapOutputModalities) + + +def _answers_text(cost_map_key: str) -> bool: + try: + row: Final = _COST_MAP_OUTPUT_MODALITIES.validate_python(litellm.model_cost.get(cost_map_key)) # pyright: ignore[reportUnknownMemberType] # cost map rows are untyped JSON + except ValidationError: + return False + return "text" in row.get("supported_output_modalities", ("text",)) + + +class _OutputLimit(TypedDict): + max_output_tokens: NotRequired[ReadOnly[PositiveInt | None]] + + +_OUTPUT_LIMIT: Final = TypeAdapter(_OutputLimit) + + +def _output_limit(fields: object) -> int | None: + try: + return _OUTPUT_LIMIT.validate_python(fields).get("max_output_tokens") + except ValidationError: + return None + + +def _catalog_capped_output_limit(cost_map_key: str, override: object) -> int | None: + limits: Final = ( + _output_limit(GetModelCostMap.loaded_model_cost_map().get(cost_map_key)), + _output_limit(override), + ) + return min((limit for limit in limits if limit is not None), default=None) + + +def _text_chat_model_max_output_tokens( + model: str, custom_llm_provider: str, deployment_model_info: object +) -> int | None: + info: Final = litellm.utils.lookup_model_info_for_registration(model, custom_llm_provider) + if isinstance(info, litellm.utils.ModelInfoLookupMiss) or info.get("mode") != "chat": + return None + if not _answers_text(info["key"]): + return None + return _catalog_capped_output_limit(info["key"], deployment_model_info or info) + + +def _with_default_output_budget( + request_body: RequestBody, model: str, custom_llm_provider: str, deployment_model_info: object +) -> RequestBody: + return with_default_output_budget( + request_body, + partial(_text_chat_model_max_output_tokens, model, custom_llm_provider, deployment_model_info), ) @@ -1380,21 +1446,8 @@ async def async_transform_request_body( messages, should_inline=should_inline_remote_media(custom_llm_provider) ) - if _openai_messages_may_need_sync_gcs_metadata_fetch(inlined_messages): - # _transform_request_body may issue a sync httpx.get (up to 5s timeout) - # via _get_gcs_object_content_type to fetch GCS object metadata. Run the - # whole sync transformation on a worker thread so it does not block the - # async event loop. - return await asyncify(transform_request_body)( - messages=inlined_messages, - model=model, - custom_llm_provider=custom_llm_provider, - litellm_params=litellm_params, - cached_content=cached_content, - optional_params=optional_params, - ) - - return transform_request_body( + transform: Final = partial( + transform_request_body, messages=inlined_messages, model=model, custom_llm_provider=custom_llm_provider, @@ -1402,6 +1455,21 @@ async def async_transform_request_body( cached_content=cached_content, optional_params=optional_params, ) + # _transform_request_body may issue a sync httpx.get (up to 5s timeout) + # via _get_gcs_object_content_type to fetch GCS object metadata. Run the + # whole sync transformation on a worker thread so it does not block the + # async event loop. + request_body: Final[RequestBody] = ( + await asyncify(transform)() + if _openai_messages_may_need_sync_gcs_metadata_fetch(inlined_messages) + else transform() + ) + return _with_default_output_budget( + request_body, + model, + custom_llm_provider, + litellm_params.get("model_info"), # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # litellm_params is an untyped dict + ) def _default_user_message_when_system_message_passed() -> ChatCompletionUserMessage: 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/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 2fae2a9de1c..c0082bc5ca6 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.url_utils import encode_url_path_segment +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, join_api_base +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): @@ -588,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 d647cab24e3..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 @@ -8846,8 +8838,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/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 3d89af9d691..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": { @@ -76709,14 +76713,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, @@ -77018,5 +77022,33 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": false + }, + "azure_ai/MAI-Code-1.1-Flash": { + "cache_read_input_token_cost": 2e-08, + "deprecation_date": "2026-12-09", + "input_cost_per_token": 2e-07, + "litellm_provider": "azure_ai", + "max_input_tokens": 256000, + "max_output_tokens": 64000, + "max_tokens": 64000, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true } } 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/client/cli/README.md b/litellm/proxy/client/cli/README.md index 842705ad5a7..4038ff6f582 100644 --- a/litellm/proxy/client/cli/README.md +++ b/litellm/proxy/client/cli/README.md @@ -494,7 +494,7 @@ Each command resolves your LiteLLM key (logging in via SSO when none is stored a The right variables are picked per agent. Claude Code gets `ANTHROPIC_BASE_URL` (the proxy root, so it appends `/v1/messages`) and `ANTHROPIC_AUTH_TOKEN`, with any stray `ANTHROPIC_API_KEY` cleared so the proxy token wins, and `ENABLE_TOOL_SEARCH=true` (unless you already set it) so Claude Code keeps tool search on even though the base URL is a proxy rather than a first-party Anthropic host. It also gets `CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY=1` (again unless you already set it) so Claude Code v2.1.129+ fills its `/model` picker from the proxy's `/v1/models`; Claude Code only lists entries whose id contains `claude` or `anthropic`, so the proxy lists every other group to Claude Code as `claude-router-` and marks a group whose input window reaches 1M with `[1m]`, and a request on such an id is served by the group. Older Claude Code versions ignore the variable. Export it as `0` to turn discovery off. Codex and OpenCode get `OPENAI_BASE_URL` (the proxy plus `/v1`) and `OPENAI_API_KEY`. Codex ignores `OPENAI_BASE_URL`, so it is additionally pointed at the proxy through a custom provider passed as `-c` config overrides (HTTP/SSE Responses transport, since the proxy does not speak the Responses WebSocket protocol). OpenCode additionally gets `OPENCODE_CONFIG_CONTENT` holding a generated `litellm` provider (`@ai-sdk/openai-compatible`, the proxy `/v1` URL, `{env:OPENAI_API_KEY}`) with one model entry per chat model your key can see on `/v1/models`, so its model picker mirrors the proxy without a hand-maintained `opencode.json`; OpenCode merges that over your own config files, and if you already export `OPENCODE_CONFIG_CONTENT` yours is left alone. When the list cannot be fetched, `lite opencode` says so on stderr and launches anyway. Codex gets the same list as a catalog file, `$CODEX_HOME/litellm-models.json` (default `~/.codex/`), passed as `-c model_catalog_json=` so `/model` lists exactly the proxy's chat models; a proxy model the installed Codex already knows (`gpt-5.5`, say) keeps that Codex's own entry, reasoning levels and prompt included, and only its place in the picker comes from the proxy, while a model Codex does not know gets the plain entry Codex uses for an unknown `-m` slug. Before launching, `lite codex` asks the installed Codex for its own list and then has it read the written file back (both through `codex debug models`), and when the fetch, either of those or the write fails (Codex releases older than 0.130 have no such command) it says so on stderr and launches with Codex's built-in catalog, leaving a rejected file in place. -pi ignores base-URL environment variables entirely, so `lite pi` (kept out of the `lite --help` command listing for now, but fully functional) wires it up differently: before handoff it fetches the models your key can use from the proxy's `/v1/models` (plus each model's context window and output cap from `/model_group/info`, when available) and syncs them into a `litellm` provider entry in pi's `~/.pi/agent/models.json` (honoring `PI_CODING_AGENT_DIR`), then starts pi on that provider's first model via an injected `--model litellm/`. Only that one provider entry is rewritten; the rest of the file, including any other custom providers, is left alone. The entry references the key as `$LITELLM_PROXY_API_KEY`, which the wrapper exports for the session, so the token itself never lands on disk and plain `pi` outside the wrapper simply shows the litellm models as unavailable. Your own flags come after the injected pin, so `lite pi --model litellm/` wins, and inside the TUI the `/model` picker lists every synced litellm model. +pi ignores base-URL environment variables entirely, so `lite pi` (kept out of the `lite --help` command listing for now, but fully functional) wires it up differently: before handoff it fetches the models your key can use from the proxy's `/v1/models` (plus each model's context window and output cap from `/model_group/info`, when available) and syncs them into a `litellm` provider entry in pi's `~/.pi/agent/models.json` (honoring `PI_CODING_AGENT_DIR`), then starts pi on that provider's first model via an injected `--model litellm/`. Only that one provider entry is rewritten; the rest of the file, including any other custom providers, is left alone. The entry references the key as `$LITELLM_PROXY_API_KEY`, which the wrapper exports for the session, so the token itself never lands on disk and plain `pi` outside the wrapper simply shows the litellm models as unavailable. The entry also turns on pi's session header (`x-session-id`, sent by pi 0.80.7 and later), so every chat turn of one pi session shares one session id in the logs and sticks to one deployment when session affinity is on. A custom `pi --session-id` must be at least 8 letters, digits, `_`, or `-` for the proxy to read it. Your own flags come after the injected pin, so `lite pi --model litellm/` wins, and inside the TUI the `/model` picker lists every synced litellm model. Options (these belong to the wrapper, so put them before the agent's own flags): diff --git a/litellm/proxy/client/cli/commands/pi.py b/litellm/proxy/client/cli/commands/pi.py index b228617699c..357b82bdf67 100644 --- a/litellm/proxy/client/cli/commands/pi.py +++ b/litellm/proxy/client/cli/commands/pi.py @@ -191,7 +191,7 @@ def provider_block( Real contextWindow/maxTokens matter: pi otherwise assumes 128k/16384, which breaks compaction thresholds and over-asks models with smaller output caps. pi sniffs compat from the base URL, and one gateway URL fronts models with - different capabilities, so both flags are pinned off. + different capabilities, so the store and long-cache flags are pinned off. """ return { "baseUrl": base_url.rstrip("/") + "/v1", @@ -199,6 +199,8 @@ def provider_block( "compat": { "supportsStore": False, "supportsLongCacheRetention": False, + "sendSessionAffinityHeaders": True, + "sessionAffinityFormat": "openrouter", }, "apiKey": f"${LITELLM_PROXY_API_KEY_ENV}", "models": [_model_entry(model_id, limits) for model_id in model_ids], diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index b98a8f2184b..831fc11ad95 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -67,6 +67,7 @@ from litellm.litellm_core_utils.get_supported_openai_params import ( get_supported_openai_params, ) from litellm.litellm_core_utils.internal_call_metadata import is_unbilled_non_inference_call_from_params +from litellm.litellm_core_utils.litellm_logging import BilledRoundsDispatch from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import guardrail_information_cost from litellm.litellm_core_utils.llm_response_utils.get_headers import ( @@ -266,6 +267,7 @@ from litellm.proxy.litellm_pre_call_utils import ( sanitize_for_log, ) from litellm.proxy.policy_engine.response_retrieval import attach_post_call_pipelines_to_retrieval +from litellm.types.litellm_params import STREAM_USAGE_INJECTED_FIELD from litellm.types.utils import ( ModelResponse, ModelResponseStream, @@ -552,22 +554,29 @@ def model_deployments_support_stream_options( return all(_litellm_model_supports_stream_options(m) for m in candidate_models) +_STREAM_USAGE_FLAGS: Final = ("_litellm_strip_stream_usage", STREAM_USAGE_INJECTED_FIELD) + + def stream_usage_tracking_updates( data: Mapping[str, object], general_settings: Mapping[str, object], route_type: str, supports_stream_options: Callable[[], bool], ) -> Mapping[str, object]: - scrub: Final = {"_litellm_strip_stream_usage": False} if "_litellm_strip_stream_usage" in data else {} + scrub: Final = {flag: False for flag in _STREAM_USAGE_FLAGS if flag in data} if data.get("stream", False) is not True: return scrub always_include: Final = general_settings.get("always_include_stream_usage") stream_options: Final = data.get("stream_options") if always_include is True: if "stream_options" not in data: - return {**scrub, "stream_options": {"include_usage": True}} + return {**scrub, "stream_options": {"include_usage": True}, STREAM_USAGE_INJECTED_FIELD: True} if isinstance(stream_options, dict) and "include_usage" not in stream_options: - return {**scrub, "stream_options": {**stream_options, "include_usage": True}} + return { + **scrub, + "stream_options": {**stream_options, "include_usage": True}, + STREAM_USAGE_INJECTED_FIELD: True, + } return scrub if always_include is False and route_type == "agenerate_content_stream" and "stream_options" not in data: return {**scrub, "stream_options": {"include_usage": False}} @@ -581,6 +590,7 @@ def stream_usage_tracking_updates( return { "stream_options": {**merged_stream_options, "include_usage": True}, "_litellm_strip_stream_usage": True, + STREAM_USAGE_INJECTED_FIELD: True, } @@ -1976,6 +1986,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: @@ -2425,6 +2436,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(), @@ -2801,6 +2813,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( @@ -2865,6 +2878,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, @@ -3474,6 +3488,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 @@ -3998,6 +4013,8 @@ class ProxyBaseLLMRequestProcessing: async def _on_deferred_native_stream_complete(*args: object) -> None: match args: + case (BilledRoundsDispatch(dispatch=logging_coroutine),): + GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=logging_coroutine) case (logging_coroutine,) if asyncio.iscoroutine(logging_coroutine): GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=logging_coroutine) case (assembled_response, cache_hit): @@ -4141,6 +4158,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/guardrails/guardrail_hooks/llm_shield_proxy/llm_shield_proxy.py b/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/llm_shield_proxy.py index 06ea3ce4b20..7fb87080cd1 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/llm_shield_proxy.py +++ b/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/llm_shield_proxy.py @@ -494,7 +494,7 @@ class LLMShieldProxyGuardrail(CustomGuardrail): sse: Final = AnthropicSSERestorer(step) events: Final = ResponsesStreamRestorer(step, rehydrate) carries: Final[CarryWindows] = {} - last_chunk = None # rebind-ok: tracks the most recent chunk for the final flush. + last_chunk = None # rebind-ok: tracks the most recent chunk with choices for the final flush. async for chunk in response: if isinstance(chunk, (bytes, str)): @@ -506,7 +506,8 @@ class LLMShieldProxyGuardrail(CustomGuardrail): yield event continue restored_chunk = detached(chunk) - last_chunk = restored_chunk + if getattr(restored_chunk, "choices", None): + last_chunk = restored_chunk for choice in attribute_items(restored_chunk, "choices"): await self._restore_choice(choice, carries, session_id) yield restored_chunk diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index 37cef41034b..bc9cee1cb1d 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/litellm/proxy/hooks/model_max_budget_limiter.py @@ -374,6 +374,8 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): self._redis_increment_queue_lock = asyncio.Lock() self._redis_increment_flush_lock = asyncio.Lock() self._detached_increment_operations = None + self._retired = False + self._flush_after_retire = None self.deployment_budget_config = None @with_service_target("model_budgets") diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 4f2f7f9c256..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, @@ -695,6 +696,7 @@ class RequestRateLimiterStash: owner_litellm_call_id: str | None = None rate_limit_response: RateLimitResponse | None = None parallel_slot: ParallelSlotAcquisition | None = None + parallel_slot_reused_at: datetime | None = None parallel_slot_release_lock: asyncio.Lock = field(default_factory=asyncio.Lock, repr=False, compare=False) reserved_tokens: int = 0 reserved_model: RateLimitedModel | None = None @@ -791,6 +793,19 @@ def claim_request_stash_for_data(data: dict) -> RequestRateLimiterStash: return stash +def _ended_before_slot_reuse(ended_at: datetime | None, reused_at: datetime | None) -> bool: + if ended_at is None or reused_at is None or (ended_at.tzinfo is None) != (reused_at.tzinfo is None): + return False + return ended_at <= reused_at + + +def _callback_kwargs_mapping(kwargs: object) -> Mapping[str, object]: + try: + return _REQUEST_RATE_LIMIT_DATA.validate_python(kwargs) + except ValidationError: + return MappingProxyType({}) + + def get_request_stash_for_call(litellm_call_id: str | None) -> RequestRateLimiterStash | None: stash: Final = _request_stash.get() if stash is None: @@ -1233,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)) @@ -1652,6 +1670,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): read_only: bool = False, skip_tpm_check: bool = False, parallel_slot_id: str | None = None, + skip_parallel_check: bool = False, ) -> RateLimitResponse: """ Check if any of the rate limit descriptors should be rate limited. @@ -1668,6 +1687,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ``reserve_tpm_tokens`` reservation path should set this to avoid the +1-per-key Lua / in-memory increment double-charging the tokens counter. + skip_parallel_check: If True, ``max_parallel_requests`` gauges are + neither checked nor acquired, for a call that runs inside a + request already holding its slot. ``max_parallel_requests`` descriptors are enforced by the dedicated concurrency-gauge path (``_check_parallel_request_gauges``), never by @@ -1760,7 +1782,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): if windowed_response["overall_code"] == "OVER_LIMIT": return windowed_response - if not gauges: + if not gauges or skip_parallel_check: return windowed_response gauge_response: Final = await self._check_parallel_request_gauges( @@ -2023,12 +2045,13 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): parent_otel_span: Span | None, *, in_logging_callback: bool = False, + round_ended_at: datetime | None = None, ) -> None: if stash is None: return async with stash.parallel_slot_release_lock: acquisition: Final = stash.parallel_slot - if acquisition is None: + if acquisition is None or _ended_before_slot_reuse(round_ended_at, stash.parallel_slot_reused_at): return deferred: Final = in_logging_callback and await self._defer_parallel_slot_release( acquisition, parent_otel_span @@ -3816,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 @@ -3947,6 +3971,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): tpm_reservation_amount: int, call_type: str | None = None, endpoint_type: EndpointType = EndpointType.GENERIC, + reuses_request_slot: bool = False, ) -> None: """ Reserve project-scoped ITPM/OTPM tokens (Bedrock Mantle-style @@ -4026,7 +4051,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): parent_otel_span=user_api_key_dict.parent_otel_span, ) stash.reservation_released = True - await self._release_stashed_parallel_slot(stash, user_api_key_dict.parent_otel_span) + if not reuses_request_slot: + await self._release_stashed_parallel_slot(stash, user_api_key_dict.parent_otel_span) self._handle_rate_limit_error( response=io_response, descriptors=descriptors, @@ -4187,6 +4213,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): verbose_proxy_logger.debug("Inside Rate Limit Pre-Call Hook") stash: Final = claim_request_stash_for_data(data) + reuses_request_slot: Final = call_type == CallTypes.call_mcp_tool.value and stash.parallel_slot is not None + if reuses_request_slot: + stash.parallel_slot_reused_at = self._get_current_time() ######################################################### # Check if the call type has a specific rate limiter @@ -4231,11 +4260,15 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # then because _reserve_project_io_tokens_or_raise below charges # them unconditionally and counting them here too would # double-charge every request. - parallel_counter_keys: Final = [ - self.create_rate_limit_keys(d["key"], d["value"], "max_parallel_requests") - for d in descriptors - if (d.get("rate_limit") or {}).get("max_parallel_requests") is not None - ] + parallel_counter_keys: Final = ( + [] + if reuses_request_slot + else [ + self.create_rate_limit_keys(d["key"], d["value"], "max_parallel_requests") + for d in descriptors + if (d.get("rate_limit") or {}).get("max_parallel_requests") is not None + ] + ) parallel_slot_id: Final = uuid.uuid4().hex if parallel_counter_keys else None first_pass_descriptors: Final = ( @@ -4250,6 +4283,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): parent_otel_span=user_api_key_dict.parent_otel_span, skip_tpm_check=self.tpm_reservation_enabled, parallel_slot_id=parallel_slot_id, + skip_parallel_check=reuses_request_slot, ) if response["overall_code"] == "OVER_LIMIT": @@ -4259,7 +4293,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): requested_model=requested_model, ) else: - stash.rate_limit_response = response + if not reuses_request_slot: + stash.rate_limit_response = response if parallel_slot_id is not None: stash.parallel_slot = ParallelSlotAcquisition( slot_id=parallel_slot_id, @@ -4357,7 +4392,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): raw_estimate=estimated_tokens, ) if stash.ptu_ceiling is not None and ptu_estimated_tokens > stash.ptu_ceiling.tpm_limit: - await self._release_stashed_parallel_slot(stash, user_api_key_dict.parent_otel_span) + if not reuses_request_slot: + await self._release_stashed_parallel_slot(stash, user_api_key_dict.parent_otel_span) raise PTUShareExceededError(stash.ptu_ceiling, ptu_estimated_tokens) tpm_response: Final = await self.reserve_tpm_tokens( descriptors=descriptors, @@ -4367,7 +4403,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) if tpm_response["overall_code"] == "OVER_LIMIT": - await self._release_stashed_parallel_slot(stash, user_api_key_dict.parent_otel_span) + if not reuses_request_slot: + await self._release_stashed_parallel_slot(stash, user_api_key_dict.parent_otel_span) self._handle_rate_limit_error( response=tpm_response, descriptors=descriptors, @@ -4418,6 +4455,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): tpm_reservation_amount=tpm_reservation_amount, call_type=call_type, endpoint_type=endpoint_type, + reuses_request_slot=reuses_request_slot, ) def _create_pipeline_operations( @@ -5332,8 +5370,16 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): try: verbose_proxy_logger.debug("INSIDE parallel request limiter ASYNC SUCCESS LOGGING") - stash: Final = get_request_stash_for_call(_call_id_from_callback_kwargs(kwargs)) - await self._release_stashed_parallel_slot(stash, litellm_parent_otel_span, in_logging_callback=True) + callback_kwargs: Final = _callback_kwargs_mapping(kwargs) + stash: Final = get_request_stash_for_call(_call_id_from_callback_kwargs(callback_kwargs)) + await self._release_stashed_parallel_slot( + stash, + litellm_parent_otel_span, + in_logging_callback=True, + round_ended_at=end_time + if isinstance(end_time, datetime) and not callback_kwargs.get("stream") + else None, + ) pipeline_operations: Final = self._build_success_event_pipeline_operations( kwargs=kwargs, 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 09f53cecb72..8a45f658f36 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -268,7 +268,7 @@ def sanitize_for_log(value: object) -> str: from litellm.router import Router from litellm.secret_managers.main import get_secret_bool -from litellm.types.litellm_params import AGENTIC_LOOP_KWARG_NAMES +from litellm.types.litellm_params import AGENTIC_LOOP_KWARG_NAMES, STREAM_USAGE_INJECTED_FIELD from litellm.types.llms.anthropic import ANTHROPIC_API_HEADERS from litellm.types.router import BaggageMappedTagsStamp, holds_api_key_secret_reference from litellm.types.services import ServiceTypes @@ -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", @@ -2370,7 +2371,15 @@ def refresh_proxy_server_request_body_snapshot( if not isinstance(proxy_server_request, dict): return _body_snapshot_exclude: Final = ( - frozenset({"secret_fields", "proxy_server_request", "litellm_logging_obj", VECTOR_STORE_CALLER_METADATA_KWARG}) + frozenset( + { + "secret_fields", + "proxy_server_request", + "litellm_logging_obj", + VECTOR_STORE_CALLER_METADATA_KWARG, + STREAM_USAGE_INJECTED_FIELD, + } + ) | _TRANSPORT_ONLY_CREDENTIAL_KEYS | _CALLBACK_CREDENTIAL_KEYS ) 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/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/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index fe7ac30e434..e1e8864678e 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -880,8 +880,8 @@ def _build_passthrough_failure_request_payload( @dataclass(frozen=True, slots=True) class _TeamCallbackWiring: - success_callbacks: "list[str | Callable | CustomLogger] | None" = None # mutable-ok: Logging.__init__ arg - failure_callbacks: "list[str | Callable | CustomLogger] | None" = None # mutable-ok: Logging.__init__ arg + success_callbacks: "list[str | Callable[..., object] | CustomLogger] | None" = None # mutable-ok: Logging init arg + failure_callbacks: "list[str | Callable[..., object] | CustomLogger] | None" = None # mutable-ok: Logging init arg logging_kwargs: dict[str, str | dict[str, str]] | None = None # mutable-ok: Logging.__init__ arg 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/proxy_server.py b/litellm/proxy/proxy_server.py index 8d6c9dd925b..f632317d9a8 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -976,6 +976,7 @@ from litellm.secret_managers.main import ( ) from litellm.tracing.config import is_lens_tracing_enabled from litellm.types.integrations.slack_alerting import AlertType, SlackAlertingArgs +from litellm.types.litellm_params import STREAM_USAGE_INJECTED_FIELD from litellm.types.llms.anthropic import ( AnthropicMessagesRequest, AnthropicResponse, @@ -10383,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 @@ -10598,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( @@ -18219,6 +18222,7 @@ async def async_queue_request( try: data = await request.json() data.pop("_litellm_strip_stream_usage", None) + data.pop(STREAM_USAGE_INJECTED_FIELD, None) # Include original request and headers in the data data["proxy_server_request"] = { 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/route_llm_request.py b/litellm/proxy/route_llm_request.py index 25f9f1dff15..cb226929167 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -1,8 +1,13 @@ import asyncio -from collections.abc import Mapping +import hashlib +import json +from collections import OrderedDict +from collections.abc import Callable, Mapping from dataclasses import dataclass +from datetime import datetime from types import MappingProxyType from typing import TYPE_CHECKING, Annotated, Any, Final, Literal +from uuid import uuid4 import httpx from fastapi import HTTPException, status @@ -10,6 +15,7 @@ from pydantic import BeforeValidator, TypeAdapter, ValidationError import litellm from litellm.litellm_core_utils.core_helpers import normalize_drop_params +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy._types import ProxyException, UserAPIKeyAuth from litellm.router_utils.common_utils import is_proxy_admin_request @@ -54,7 +60,34 @@ _ROUTER_ONLY_SETTINGS: Final = frozenset( ) -def _route_user_config_request(data: dict, route_type: str): +class UserConfigRouters: + def __init__(self, max_size: int) -> None: + self._max_size: Final = max_size + self._routers: Final[OrderedDict[str, LitellmRouter]] = OrderedDict() # mutable-ok: LRU evicts in place + + def get(self, config: Mapping[str, object], build: Callable[[], LitellmRouter]) -> LitellmRouter: + key: Final = hashlib.sha256(json.dumps(config, sort_keys=True, default=str).encode()).hexdigest() + cached: Final = self._routers.get(key) + if cached is not None: + self._routers.move_to_end(key) + return cached + router: Final = build() + router.serve_through_request_logging() + self._routers[key] = router + if len(self._routers) > self._max_size: + self._routers.popitem(last=False)[1].discard() + return router + + def clear(self) -> None: + for router in self._routers.values(): + router.discard() + self._routers.clear() + + +user_config_routers: Final = UserConfigRouters(max_size=32) + + +async def _route_user_config_request(data: dict, route_type: str): """Route a request using the user-provided router config.""" router_config: Final = data.pop("user_config") @@ -63,10 +96,24 @@ def _route_user_config_request(data: dict, route_type: str): valid_args: Final = litellm.Router.get_valid_args() filtered_config: Final = {k: v for k, v in router_config.items() if k in valid_args} - user_router: Final = litellm.Router(**filtered_config) - ret_val: Final = getattr(user_router, f"{route_type}")(**data) - user_router.discard() - return ret_val + logging_obj: Final = _request_logging_obj(JSON_OBJECT_ADAPTER.validate_python(data), route_type) + cached_router: Final = user_config_routers.get(filtered_config, lambda: litellm.Router(**filtered_config)) + cached_router.bind_to_request_logging(logging_obj) + return await getattr(cached_router, f"{route_type}")(**{**data, "litellm_logging_obj": logging_obj}) + + +def _request_logging_obj(data: Mapping[str, object], route_type: str) -> LiteLLMLoggingObj: + arrived_with: Final = data.get("litellm_logging_obj") + if isinstance(arrived_with, LiteLLMLoggingObj): + return arrived_with + logging_obj, _ = litellm.utils.function_setup( + original_function=route_type, + rules_obj=litellm.utils.Rules(), + start_time=datetime.now(), # noqa: DTZ005 # Logging subtracts it from naive end times + is_async_call=True, + **{"litellm_call_id": str(uuid4()), **data}, + ) + return logging_obj def _is_a2a_agent_model(model_name: object) -> bool: 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 1966c32fb07..d423a95d037 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 @@ -158,7 +159,7 @@ from litellm.litellm_core_utils.core_helpers import ( ) from litellm.litellm_core_utils.duration_parser import get_budget_window_start from litellm.litellm_core_utils.env_utils import get_env_float, get_env_int -from litellm.litellm_core_utils.litellm_logging import Logging, deployment_pricing_model_info +from litellm.litellm_core_utils.litellm_logging import BilledRoundsDispatch, Logging, deployment_pricing_model_info from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.safe_json_loads import safe_json_loads @@ -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], ...]) @@ -1015,13 +1018,14 @@ 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, 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 @@ -1037,9 +1041,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: @@ -1048,7 +1049,7 @@ def _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: @@ -1059,13 +1060,14 @@ def _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")) @@ -1082,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, ) @@ -1281,6 +1287,31 @@ 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, + dispatched_call_type: 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, + 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: + 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({}) @@ -1424,7 +1455,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 @@ -1444,13 +1477,14 @@ 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, + 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 ) @@ -3618,6 +3652,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. @@ -3688,7 +3723,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), ) ) @@ -4736,19 +4773,25 @@ class ProxyLogging: @staticmethod def _discard_deferred_stream_logging_for_failure(request_data: Mapping[str, object], error: Exception) -> bool: - """Drop the parked success dispatch for an assembled chat stream that ends in an error - ``post_call_failure_hook`` logs as a failure, billing its usage on the failure row instead. + """Drop the parked success dispatch for an assembled chat stream or a stream that billed several rounds when it + ends in an error ``post_call_failure_hook`` logs as a failure, billing its usage on the failure row instead. Returns False when the parked dispatch should still be flushed by the caller.""" logging_obj: Final = request_data.get("litellm_logging_obj") if not isinstance(logging_obj, Logging): return False _args: Final[tuple[object, ...] | None] = getattr(logging_obj, "_deferred_stream_complete_args", None) - assembled: Final = _args[0] if _args else None - if not isinstance(error, _PROXY_ONLY_LLM_API_ERRORS) or not isinstance(assembled, ModelResponse): + if not isinstance(error, _PROXY_ONLY_LLM_API_ERRORS): return False + match _args: + case (ModelResponse() as assembled, *_): + logging_obj.record_assembled_response_for_failure(assembled) + case (BilledRoundsDispatch() as billed,): + billed.dispatch.close() + logging_obj.record_billed_rounds_for_failure(billed.spend) + case _: + return False logging_obj._on_deferred_stream_complete = None logging_obj._deferred_stream_complete_args = None - logging_obj.record_assembled_response_for_failure(assembled) return True async def arelease_max_parallel_requests_on_disconnect( 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/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 84d4648f75f..c25dd4a405e 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -5,8 +5,10 @@ from itertools import filterfalse from typing import Any, Final, cast import litellm -from litellm.litellm_core_utils.hidden_params import get_or_create_hidden_params +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, @@ -82,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) @@ -145,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 @@ -185,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, @@ -193,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: @@ -247,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)) @@ -255,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: @@ -591,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, @@ -1003,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) @@ -1090,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, @@ -1138,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( @@ -1194,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( @@ -1355,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: @@ -1363,6 +1429,17 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): return False return not any(getattr(part, "text", None) for part in getattr(item, "content", None) or ()) + def finished_round_event(self) -> ResponseCompletedEvent | ResponseIncompleteEvent | None: + if not isinstance(self.litellm_model_response, ModelResponse): + return None + finished: Final = self._emit_terminal_response_event(self.litellm_model_response) + if finished is not None and served_from_cache(self.litellm_custom_stream_wrapper): + finished.response.hidden_params = {**finished.response.hidden_params, "cache_hit": True} + return finished + + def bill_failure_after_earlier_rounds(self, earlier: BilledRoundsSpend) -> None: + self.litellm_custom_stream_wrapper.bill_failure_after_earlier_rounds(earlier) + def partial_response_on_disconnect(self) -> ResponseIncompleteEvent | None: """ The response.incomplete event to bill when the client went away mid-stream. @@ -1376,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/mcp/mcp_streaming_iterator.py b/litellm/responses/mcp/mcp_streaming_iterator.py index 73cfba8dca4..030911238ec 100644 --- a/litellm/responses/mcp/mcp_streaming_iterator.py +++ b/litellm/responses/mcp/mcp_streaming_iterator.py @@ -3,7 +3,9 @@ from typing import TYPE_CHECKING, Any, Final, cast from litellm._logging import verbose_logger from litellm._uuid import uuid +from litellm.responses.litellm_completion_transformation.streaming_iterator import LiteLLMCompletionStreamingIterator from litellm.responses.mcp.request_context import MCPRequestContext +from litellm.responses.mcp.stream_round_billing import StreamedRoundBilling from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator from litellm.types.llms.openai import ( BaseLiteLLMOpenAIResponseObject, @@ -352,6 +354,9 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): self._round_max_output_index = -1 self._composed_output: list[object] = [] # mutable-ok: grows as each round finishes self._pending_mcp_call_items: list[dict[str, object]] = [] # mutable-ok: grows per executed tool + self._iteration_started = False + self._round_billing: StreamedRoundBilling | None = None + self._bridged_round_billed: LiteLLMCompletionStreamingIterator | None = None def _extract_mcp_headers_from_params(self) -> None: """Extract MCP headers from original request params to pass to tool calls""" @@ -436,7 +441,15 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): return self async def __anext__(self) -> ResponsesAPIStreamingResponse: - chunk: Final = await self._anext_impl() + if not self._iteration_started: + self._iteration_started = True + self._round_billing = self._bill_rounds_on_one_row() + try: + chunk: Final = await self._anext_impl() + except Exception: + if self._round_billing is not None: + self._round_billing.finish() + raise sequence_number: Final = getattr(chunk, "sequence_number", None) if isinstance(sequence_number, int): if sequence_number <= self._last_sequence_number and self._last_sequence_number > 0: @@ -446,6 +459,49 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): self._last_sequence_number = max(self._last_sequence_number, sequence_number) return chunk + def __del__(self) -> None: + round_billing: Final[object] = getattr(self, "_round_billing", None) + if isinstance(round_billing, StreamedRoundBilling): + round_billing.release_abandoned() + + def _bill_rounds_on_one_row(self) -> StreamedRoundBilling | None: + if not self.should_auto_execute: + return None + if isinstance(self.base_iterator, LiteLLMCompletionStreamingIterator): + return StreamedRoundBilling(self.base_iterator.litellm_logging_obj) + if isinstance(self.base_iterator, BaseResponsesAPIStreamingIterator): + return StreamedRoundBilling(self.base_iterator.logging_obj) + return None + + def _restore_round_headers(self, logged_event: ResponsesAPIStreamingResponse) -> None: + if isinstance(self.base_iterator, LiteLLMCompletionStreamingIterator): + return + if isinstance(self.base_iterator, BaseResponsesAPIStreamingIterator): + self.base_iterator._restore_provider_response_headers(logged_event) + + def _round_billed_early(self) -> bool: + return self._bridged_round_billed is not None and self._bridged_round_billed is self.base_iterator + + def _bill_finished_bridged_round(self) -> None: + base: Final = self.base_iterator + if self._round_billing is None or not isinstance(base, LiteLLMCompletionStreamingIterator): + return + if self._round_billed_early(): + return + finished: Final = base.finished_round_event() + if finished is None: + return + self._bridged_round_billed = base + round_output: Final = _output_items(finished.response) + self._round_billing.complete_round( + finished, + finished.response, + round_output + if finished.type == ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE + else (*self._composed_output, *round_output), + self._restore_round_headers, + ) + async def _anext_impl(self) -> ResponsesAPIStreamingResponse: """ Phase-based streaming: @@ -617,6 +673,7 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): follow-up rounds past the items already emitted, and lists every round's items on the final response.completed. """ + self._bill_finished_bridged_round() chunk_type: Final[object] = getattr(chunk, "type", None) if self._round_index > 0 and chunk_type in ( ResponsesAPIStreamEvents.RESPONSE_CREATED, @@ -630,12 +687,21 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): if self._output_index_offset: _set_event_field(chunk, "output_index", output_index + self._output_index_offset) + if self.should_auto_execute and chunk_type == ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE: + return self._charge_incomplete_round(chunk) if not (self.should_auto_execute and self._is_response_completed(chunk)): return chunk response_obj: Final[object] = getattr(chunk, "response", None) if isinstance(response_obj, ResponsesAPIResponse): self.collected_response = response_obj + if self._round_billing is not None and not self._round_billed_early(): + self._round_billing.complete_round( + chunk, + response_obj, + (*self._composed_output, *_output_items(response_obj)), + self._restore_round_headers, + ) # Move to tool execution phase after this chunk self.phase = "tool_execution" await self._generate_tool_execution_events() @@ -645,13 +711,33 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): if self._follow_up_pending(): self._absorb_round(response_obj) return None - if self._composed_output: - merged_output: Final[list[object]] = [ # mutable-ok: the response model declares output as a list - *self._composed_output, - *_output_items(response_obj), - ] - merged_response: Final = response_obj.model_copy(update={"output": merged_output}) - _set_event_field(chunk, "response", merged_response) + merged_response: Final = ( + response_obj.model_copy(update={"output": [*self._composed_output, *_output_items(response_obj)]}) + if self._composed_output + else response_obj + ) + client_response: Final = ( + merged_response + if self._round_billing is None + else self._round_billing.charge_final_round(chunk, merged_response, self._restore_round_headers) + ) + if client_response is not response_obj: + _set_event_field(chunk, "response", client_response) + return chunk + + def _charge_incomplete_round(self, chunk: ResponsesAPIStreamingResponse) -> ResponsesAPIStreamingResponse: + response_obj: Final[object] = getattr(chunk, "response", None) + if self._round_billing is None or not isinstance(response_obj, ResponsesAPIResponse): + return chunk + if not self._round_billed_early(): + self._round_billing.complete_round( + chunk, response_obj, _output_items(response_obj), self._restore_round_headers + ) + client_response: Final = self._round_billing.charge_final_round( + chunk, response_obj, self._restore_round_headers + ) + if client_response is not response_obj: + _set_event_field(chunk, "response", client_response) return chunk async def _process_base_iterator_chunk(self) -> ResponsesAPIStreamingResponse: @@ -949,6 +1035,8 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator): self.collected_response = None self._cached_response_id = None self._round_index += 1 + if self._round_billing is not None: + self._round_billing.record_billed_rounds(follow_up_response) except Exception as e: verbose_logger.error("Error creating follow-up iterator: %s", e) diff --git a/litellm/responses/mcp/round_billing.py b/litellm/responses/mcp/round_billing.py index ae34ab31633..93769fbe82a 100644 --- a/litellm/responses/mcp/round_billing.py +++ b/litellm/responses/mcp/round_billing.py @@ -4,7 +4,7 @@ from contextlib import contextmanager from types import MappingProxyType from typing import Final, TypeVar -from pydantic import TypeAdapter +from pydantic import BaseModel, TypeAdapter from pydantic_core import to_jsonable_python from litellm.litellm_core_utils.hidden_params import get_hidden_params, served_from_cache, set_hidden_params @@ -46,15 +46,18 @@ def _summed_mapping(first: Mapping[str, object], final: Mapping[str, object]) -> return MappingProxyType({key: _summed_entry(key, first, final) for key in keys}) +def _summed_usage(first: BaseModel, final: BaseModel) -> object: + return to_jsonable_python(_summed_mapping(first.model_dump(), final.model_dump()), fallback=dict) + + def merged_round_usage(first: ResponseAPIUsage | None, final: ResponseAPIUsage | None) -> ResponseAPIUsage | None: if first is None or final is None: return final if first is None else first - return ResponseAPIUsage.model_validate(_summed_mapping(first.model_dump(), final.model_dump())) + return ResponseAPIUsage.model_validate(_summed_usage(first, final)) -def _merged_chat_usage(first: Usage, final: Usage) -> Usage: - summed: Final = _summed_mapping(first.model_dump(), final.model_dump()) - return Usage.model_validate(to_jsonable_python(summed, fallback=dict)) +def merged_chat_usage(first: Usage, final: Usage) -> Usage: + return Usage.model_validate(_summed_usage(first, final)) def _chat_usage(round_response: ResponsesAPIResponse | ModelResponse) -> Usage | None: @@ -82,7 +85,7 @@ def _usage_of_both_rounds( final_usage: Final = _chat_usage(final) if first_usage is None or final_usage is None: return final_usage if first_usage is None else first_usage - return _merged_chat_usage(first_usage, final_usage) + return merged_chat_usage(first_usage, final_usage) def _repeated_fixed_margin( @@ -105,6 +108,14 @@ def _less_repeated_fixed_margin(summed: Mapping[str, object], repeated: float) - return MappingProxyType({**summed, **charged_once}) +def breakdown_less_repeated_fixed_margin( + breakdown: CostBreakdown | None, repeated_fixed_margin: float +) -> CostBreakdown | None: + if breakdown is None: + return None + return _COST_BREAKDOWN.validate_python(_less_repeated_fixed_margin(breakdown, repeated_fixed_margin)) + + def summed_cost_breakdown( first: ResponsesAPIResponse | ModelResponse, first_breakdown: CostBreakdown | None, diff --git a/litellm/responses/mcp/stream_round_billing.py b/litellm/responses/mcp/stream_round_billing.py new file mode 100644 index 00000000000..443147230bf --- /dev/null +++ b/litellm/responses/mcp/stream_round_billing.py @@ -0,0 +1,376 @@ +import asyncio +from collections.abc import Callable, Coroutine, Sequence +from dataclasses import dataclass +from datetime import datetime +from typing import Final, Protocol, TypeAlias + +from litellm.litellm_core_utils.hidden_params import get_hidden_params, served_from_cache +from litellm.litellm_core_utils.litellm_logging import BilledRoundsDispatch, BilledRoundsSpend, UncachedPrice +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.responses.mcp.round_billing import billed_for_every_round, merged_chat_usage, summed_cost_breakdown +from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator, spawn_logging_task +from litellm.responses.utils import ResponseAPILoggingUtils +from litellm.types.llms.openai import ResponsesAPIResponse, ResponsesAPIStreamingResponse +from litellm.types.utils import CostBreakdown, ModelResponse, Usage + +Parked: TypeAlias = tuple[object, ...] +RestoresHeaders: TypeAlias = Callable[[ResponsesAPIStreamingResponse], None] + +_PARKED_DISPATCH: Final = "_deferred_stream_complete_args" +_DISPATCH_RELEASE: Final = "_on_deferred_stream_complete" +_SERVED_FROM_CACHE: Final = UncachedPrice(cost=0.0, cost_breakdown=None) + + +class _DispatchesSuccessHandlers(Protocol): + async def dispatch_success_handlers( + self, + result: object = None, + start_time: object = None, + end_time: object = None, + cache_hit: object = None, + prefer_async_handlers: bool = False, + ) -> None: ... + + +def _success_dispatcher(logger: LiteLLMLoggingObj) -> _DispatchesSuccessHandlers: + return logger + + +def _event_loop_running() -> bool: + try: + asyncio.get_running_loop() + except RuntimeError: + return False + return True + + +def _priced(response: ResponsesAPIResponse, price: UncachedPrice) -> ResponsesAPIResponse: + priced: Final = response.model_copy() + priced.hidden_params = {**response.hidden_params, "response_cost": price.cost} + return priced + + +def _assembled(parked: Parked | None) -> ModelResponse | None: + match parked: + case (ModelResponse() as assembled, bool()): + return assembled + case _: + return None + + +def _bridged_round_usage(parked: Parked | None) -> Usage | None: + usage: Final[object] = getattr(_assembled(parked), "usage", None) + return usage if isinstance(usage, Usage) else None + + +def _usage_of_bridged_rounds(billed: Usage | None, round_usage: Usage | None, cost: object) -> Usage | None: + if billed is None or round_usage is None: + return None + return merged_chat_usage(billed, round_usage).model_copy(update={"cost": cost}) + + +def _logged_cache_hit(response: ResponsesAPIResponse, chat_usage: Usage | None) -> bool | None: + if chat_usage is not None: + return served_from_cache(response) + return served_from_cache(response) or None + + +@dataclass(frozen=True, slots=True) +class _RoundRow: + logged: ResponsesAPIStreamingResponse | ModelResponse + cache_hit: bool | None + end_time: datetime + + +@dataclass(frozen=True, slots=True) +class _SuccessLog: + logger: LiteLLMLoggingObj + row: _RoundRow + chat_usage: Usage | None + + def carry_bridged_usage(self) -> None: + if self.chat_usage is not None: + self.logger.model_call_details["combined_usage_object"] = self.chat_usage + + +def _round_row( + completed: ResponsesAPIStreamingResponse, + response: ResponsesAPIResponse, + chat_usage: Usage | None, + restore_headers: RestoresHeaders, +) -> _RoundRow: + logged_response: Final = response.model_copy() + logged_response.hidden_params = {**response.hidden_params} + logged: Final = completed.model_copy(update={"response": logged_response}) + restore_headers(logged) + return _RoundRow( + logged, + _logged_cache_hit(response, chat_usage), + datetime.now(), # noqa: DTZ005 # naive to match the logging pipeline start_time + ) + + +@dataclass(frozen=True, slots=True) +class _BilledRounds: + response: ResponsesAPIResponse + breakdown: CostBreakdown | None + chat_usage: Usage | None + rounds: int = 1 + + @property + def stamped_cost(self) -> float | None: + cost: Final = self.response.hidden_params.get("response_cost") + return float(cost) if isinstance(cost, (int, float)) else None + + @property + def cost(self) -> float: + return self.stamped_cost or 0.0 + + @property + def spend(self) -> BilledRoundsSpend: + usage: Final = ( + ResponseAPILoggingUtils.transform_response_api_usage_to_chat_usage(self.response.usage) + if self.chat_usage is None + else self.chat_usage + ) + return BilledRoundsSpend(usage, self.cost, self.breakdown) + + def with_round( + self, priced: ResponsesAPIResponse, breakdown: CostBreakdown | None, chat_usage: Usage | None + ) -> "_BilledRounds": + merged: Final = billed_for_every_round(self.response, self.breakdown, priced, breakdown) + total: Final = merged.hidden_params.get("response_cost") + charged: Final = merged.model_copy( + update={"usage": None if merged.usage is None else merged.usage.model_copy(update={"cost": total})} + ) + charged.hidden_params = { + **merged.hidden_params, + "cache_hit": served_from_cache(self.response) and served_from_cache(priced), + } + return _BilledRounds( + charged, + summed_cost_breakdown(self.response, self.breakdown, priced, breakdown), + _usage_of_bridged_rounds(self.chat_usage, chat_usage, total), + self.rounds + 1, + ) + + def row( + self, + completed: ResponsesAPIStreamingResponse, + client_output: Sequence[object], + restore_headers: RestoresHeaders, + ) -> _RoundRow: + return _round_row( + completed, + self.response.model_copy(update={"output": list(client_output)}), + self.chat_usage, + restore_headers, + ) + + def bridged_row(self, assembled: ModelResponse | None, logged_id: str) -> _RoundRow | None: + if assembled is None or self.chat_usage is None: + return None + logged: Final = assembled.model_copy(update={"id": logged_id, "usage": self.chat_usage}) + logged.hidden_params = { + **(get_hidden_params(assembled) or {}), + "response_cost": self.stamped_cost, + } + return _RoundRow( + logged, + served_from_cache(self.response), + datetime.now(), # noqa: DTZ005 # naive to match the logging pipeline start_time + ) + + +class StreamedRoundBilling: + """One success row per streamed auto-execute request, billed for every round it ran. + + Each round's iterator parks its success dispatch on the shared logging object rather than sending it, either as a + ready dispatch or, for a stream bridged through chat completions, as the assembled response and its cache hit. + When a round completes, before its tools run, it is folded into the rounds before it and a dispatch for every + round so far stays parked, so a client disconnect or a later round that fails still bills them. Once a later round + starts, that dispatch logs from a copy of the logging state taken then, so the later round's failure, logged on the + request's own state where the proxy reads its failure row, never reaches their success row. Only the latest + dispatch is sent, through the proxy's deferred release when the proxy armed one.""" + + def __init__(self, logging_obj: LiteLLMLoggingObj) -> None: + self._logging_obj: Final = logging_obj + self._billed: _BilledRounds | None = None + self._banked: Parked | None = None + self._banked_row: _RoundRow | None = None + if getattr(logging_obj, _DISPATCH_RELEASE, None) is None: + setattr(logging_obj, _DISPATCH_RELEASE, self._run_parked) + + def complete_round( + self, + completed: ResponsesAPIStreamingResponse, + response: ResponsesAPIResponse, + client_output: Sequence[object], + restore_headers: RestoresHeaders, + ) -> None: + round_parked: Final = self._take_parked() + price: Final = self._price(response, round_parked) + priced: Final = _priced(response, price) + chat_usage: Final = _bridged_round_usage(round_parked) + if self._billed is None: + first_round: Final = _BilledRounds(priced, price.cost_breakdown, chat_usage) + self._billed = first_round + self._bank( + round_parked, + first_round.bridged_row(_assembled(round_parked), priced.id) + or first_round.row(completed, client_output, restore_headers), + ) + return + billed: Final = self._billed.with_round(priced, price.cost_breakdown, chat_usage) + self._close(round_parked) + self._close(self._banked) + self._logging_obj.model_call_details.pop("combined_usage_object", None) + self._logging_obj.model_call_details.pop("response_cost", None) + self._billed = billed + self._logging_obj.cost_breakdown = billed.breakdown + row: Final = billed.bridged_row(_assembled(round_parked), billed.response.id) or billed.row( + completed, client_output, restore_headers + ) + self._bank(None if round_parked is None else (self._dispatch(row, billed, self._logging_obj),), row) + + def record_billed_rounds(self, follow_up: object) -> None: + if self._billed is None or not isinstance(follow_up, BaseResponsesAPIStreamingIterator): + return + follow_up.bill_failure_after_earlier_rounds(self._billed.spend) + self._bank_apart_from_later_rounds(self._billed) + + def charge_final_round( + self, completed: ResponsesAPIStreamingResponse, response: ResponsesAPIResponse, restore_headers: RestoresHeaders + ) -> ResponsesAPIResponse: + parked: Final = self._take_parked() + self._banked = None + if self._billed is None or self._billed.rounds == 1: + if parked is not None: + self._release(parked) + return response + self._close(parked) + charged: Final = response.model_copy(update={"usage": self._billed.response.usage}) + charged.hidden_params = {**self._billed.response.hidden_params} + if parked is not None: + final_row: Final = self._billed.bridged_row(self._latest_assembled(), charged.id) or _round_row( + completed, charged, self._billed.chat_usage, restore_headers + ) + self._release((self._dispatch(final_row, self._billed, self._logging_obj),)) + return charged + + def finish(self) -> None: + parked: Final = self._take_parked() + self._banked = None + if parked is not None: + self._release(parked) + return + self._disarm() + + def release_abandoned(self) -> None: + if not self._owns_release(): + return + if _event_loop_running(): + self.finish() + return + self.abandon() + + def abandon(self) -> None: + self._close(self._take_parked()) + self._close(self._banked) + self._banked = None + self._disarm() + + def _latest_assembled(self) -> ModelResponse | None: + row: Final = self._banked_row + return row.logged if row is not None and isinstance(row.logged, ModelResponse) else None + + def _price(self, response: ResponsesAPIResponse, round_parked: Parked | None) -> UncachedPrice: + if served_from_cache(response): + return _SERVED_FROM_CACHE + match round_parked: + case (ModelResponse() as assembled, bool()): + return self._logging_obj.uncached_price(assembled) + case _: + return self._logging_obj.uncached_price(response) + + def _bank(self, parked: Parked | None, row: _RoundRow) -> None: + self._banked = parked + self._banked_row = row + if parked is not None: + setattr(self._logging_obj, _PARKED_DISPATCH, parked) + + def _bank_apart_from_later_rounds(self, billed: _BilledRounds) -> None: + row: Final = self._banked_row + if row is None: + return + parked: Final = self._take_parked() + if parked is None: + return + self._close(parked) + self._bank((self._dispatch(row, billed, self._logging_obj.attempt_scoped_copy()),), row) + + def _dispatch( + self, row: _RoundRow, billed: _BilledRounds, logger: LiteLLMLoggingObj + ) -> BilledRoundsDispatch | Coroutine[object, object, None]: + dispatch: Final = self._log_success(_SuccessLog(logger, row, billed.chat_usage)) + if billed.chat_usage is None: + return dispatch + return BilledRoundsDispatch(dispatch, billed.spend) + + @staticmethod + async def _log_success(success: _SuccessLog) -> None: + success.carry_bridged_usage() + await _success_dispatcher(success.logger).dispatch_success_handlers( + success.row.logged, + start_time=success.logger.start_time, + end_time=success.row.end_time, + cache_hit=success.row.cache_hit, + prefer_async_handlers=True, + ) + + def _take_parked(self) -> Parked | None: + parked: Final[Parked | None] = getattr(self._logging_obj, _PARKED_DISPATCH, None) + if not isinstance(parked, tuple): + return None + setattr(self._logging_obj, _PARKED_DISPATCH, None) + return parked + + async def _run_parked(self, *parked: object) -> None: + match parked: + case (BilledRoundsDispatch(dispatch=dispatch),): + await dispatch + case (dispatch,) if asyncio.iscoroutine(dispatch): + await dispatch + case (assembled, cache_hit): + await _success_dispatcher(self._logging_obj).dispatch_success_handlers( + assembled, cache_hit=cache_hit, start_time=None, end_time=None, prefer_async_handlers=True + ) + case _: + return + + def _owns_release(self) -> bool: + release: Final[object] = getattr(self._logging_obj, _DISPATCH_RELEASE, None) + return release == self._run_parked + + def _release(self, parked: Parked) -> None: + if getattr(self._logging_obj, _DISPATCH_RELEASE, None) is not None and not self._owns_release(): + setattr(self._logging_obj, _PARKED_DISPATCH, parked) + return + setattr(self._logging_obj, _DISPATCH_RELEASE, None) + spawn_logging_task( + asyncio.get_running_loop(), self._run_parked(*parked), task_name="Responses MCP stream success logging" + ) + + def _disarm(self) -> None: + if self._owns_release(): + setattr(self._logging_obj, _DISPATCH_RELEASE, None) + + @staticmethod + def _close(parked: Parked | None) -> None: + match parked: + case (BilledRoundsDispatch(dispatch=dispatch),): + dispatch.close() + case (dispatch,) if asyncio.iscoroutine(dispatch): + dispatch.close() + case _: + return diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index c5dc45566a9..0d59c93786c 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -42,6 +42,7 @@ from litellm.litellm_core_utils.asyncify import run_async_function from litellm.litellm_core_utils.core_helpers import process_response_headers from litellm.litellm_core_utils.exception_mapping_utils import with_deployment_api_base from litellm.litellm_core_utils.internal_call_metadata import without_budget_reservation +from litellm.litellm_core_utils.litellm_logging import BilledRoundsSpend from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.llm_response_utils.get_api_base import get_api_base from litellm.litellm_core_utils.llm_response_utils.get_headers import provider_request_id_headers @@ -363,6 +364,7 @@ class BaseResponsesAPIStreamingIterator: """ _pending_logging_tasks: tuple[asyncio.Task[object], ...] = () + _earlier_rounds_spend: BilledRoundsSpend | None = None def __init__( self, @@ -705,18 +707,22 @@ class BaseResponsesAPIStreamingIterator: if usage_obj is None: return try: - self.logging_obj.model_call_details["combined_usage_object"] = ( - ResponseAPILoggingUtils.transform_response_api_usage_to_chat_usage(usage_obj) - ) + failed_usage: Final = ResponseAPILoggingUtils.transform_response_api_usage_to_chat_usage(usage_obj) except (TypeError, ValueError) as usage_error: verbose_logger.debug( "could not record usage for failed responses stream: %s", usage_error, ) return - self.logging_obj.model_call_details["response_cost"] = ( - self.logging_obj.response_cost_calculator(result=response_obj) or 0.0 - ) + failed_cost: Final = self.logging_obj.response_cost_calculator(result=response_obj) or 0.0 + if self._earlier_rounds_spend is None: + self.logging_obj.model_call_details["combined_usage_object"] = failed_usage + self.logging_obj.model_call_details["response_cost"] = failed_cost + return + self.logging_obj.record_failure_after_earlier_rounds(self._earlier_rounds_spend, failed_usage, failed_cost) + + def bill_failure_after_earlier_rounds(self, earlier: BilledRoundsSpend) -> None: + self._earlier_rounds_spend = earlier def _map_error_event_exception(self, error_obj: object) -> Exception: return _map_stream_error_to_exception( @@ -1701,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 a18ad7c3a97..da04e59f18a 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -38,6 +38,7 @@ from contextvars import ContextVar from dataclasses import dataclass from datetime import datetime, timezone from functools import lru_cache, partial +from itertools import chain from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Generic, Literal, NoReturn, Optional, TypeAlias, TypeVar, Union, cast @@ -259,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, @@ -361,10 +363,11 @@ from litellm.router_utils.routing_groups import ( validate_routing_strategy, ) from litellm.router_utils.routing_read_batch import ROUTER_USAGE_TARGET, RoutingPrefetch, RoutingReadBatch +from litellm.router_utils.stream_usage_injection import router_config_requests_stream_usage from litellm.scheduler import FlowItem, Scheduler from litellm.secret_managers.main import str_to_bool from litellm.types.guardrails import GuardrailEventHooks -from litellm.types.litellm_params import RoutingStrategyName +from litellm.types.litellm_params import STREAM_USAGE_INJECTED_FIELD, RoutingStrategyName from litellm.types.llms.openai import ( AllMessageValues, ChatCompletionToolParam, @@ -552,6 +555,7 @@ _PreRoutingStrategyT = TypeVar("_PreRoutingStrategyT") _CallbackT = TypeVar("_CallbackT") _DeploymentT = TypeVar("_DeploymentT", bound=Mapping[str, object]) + _ALIAS_PARAMS_NEVER_FORWARDED: Final = frozenset( {"model", "api_base", "api_key", "api_version", "responses_api_version"} ) @@ -1046,6 +1050,8 @@ class RoutingArgs(enum.Enum): # as the per-request one built from a caller-supplied user_config, drops out on its # own rather than leaving entries behind that nothing can withdraw. _live_routers: Final["weakref.WeakSet[Router]"] = weakref.WeakSet() +_routers_holding_global_callbacks: Final["weakref.WeakSet[Router]"] = weakref.WeakSet() +_routers_holding_global_callbacks_lock: Final = threading.Lock() def _replay_live_router_model_cost() -> None: @@ -1578,22 +1584,18 @@ class Router: self._init_routing_groups(self._routing_groups_input) self._override_selectors: dict[str, RouterStrategySelector | None] = {} self._override_selectors_lock = threading.Lock() + self._request_scoped_selectors: tuple[CustomLogger, ...] = () + self._request_scoped_policies: tuple[CustomLogger, ...] = () self.access_groups = None + with _routers_holding_global_callbacks_lock: + _routers_holding_global_callbacks.add(self) ## USAGE TRACKING ## - if isinstance(litellm._async_success_callback, list): - litellm.logging_callback_manager.add_litellm_async_success_callback(self.deployment_callback_on_success) - else: - litellm.logging_callback_manager.add_litellm_async_success_callback(self.deployment_callback_on_success) + litellm.logging_callback_manager.add_litellm_async_success_callback(self.deployment_callback_on_success) if isinstance(litellm.success_callback, list): litellm.logging_callback_manager.add_litellm_success_callback(self.sync_deployment_callback_on_success) else: litellm.success_callback = [self.sync_deployment_callback_on_success] - if isinstance(litellm._async_failure_callback, list): - litellm.logging_callback_manager.add_litellm_async_failure_callback( - self.async_deployment_callback_on_failure - ) - else: - litellm._async_failure_callback = [self.async_deployment_callback_on_failure] + litellm.logging_callback_manager.add_litellm_async_failure_callback(self.async_deployment_callback_on_failure) ## COOLDOWNS ## if isinstance(litellm.failure_callback, list): litellm.logging_callback_manager.add_litellm_failure_callback(self.deployment_callback_on_failure) @@ -1689,20 +1691,116 @@ class Router: # Stop contributing to cost-map rebuilds straight away rather than waiting # for this router to be collected. _live_routers.discard(self) - litellm.logging_callback_manager.remove_callback_from_list_by_object(litellm._async_success_callback, self) - litellm.logging_callback_manager.remove_callback_from_list_by_object(litellm.success_callback, self) - litellm.logging_callback_manager.remove_callback_from_list_by_object(litellm._async_failure_callback, self) - litellm.logging_callback_manager.remove_callback_from_list_by_object(litellm.failure_callback, self) - litellm.logging_callback_manager.remove_callback_from_list_by_object(litellm.input_callback, self) - litellm.logging_callback_manager.remove_callback_from_list_by_object(litellm.service_callback, self) - litellm.logging_callback_manager.remove_callback_from_list_by_object(litellm.callbacks, self) + self.retire_routing_state() + self._release_global_callbacks() - # Remove ForwardClientSideHeadersByModelGroup if it exists - if self.optional_callbacks is not None: - for callback in self.optional_callbacks: - litellm.logging_callback_manager.remove_callback_from_list_by_object( - litellm.callbacks, callback, require_self=False - ) + def _release_global_callbacks(self) -> None: + with _routers_holding_global_callbacks_lock: + _routers_holding_global_callbacks.discard(self) + slack_alerting_loggers: Final = self._slack_alerting_loggers() + adaptive_router_hooks: Final = self._adaptive_router_hooks() + vacated_keys: Final = litellm.logging_callback_manager.drop_callbacks( + ( + *self._callback_selectors(), + *self._override_selectors.values(), + *(self.optional_callbacks or ()), + *slack_alerting_loggers, + *adaptive_router_hooks, + ), + method_owners=(self, *slack_alerting_loggers), + ) + with _routers_holding_global_callbacks_lock: + survivors: Final = tuple(_routers_holding_global_callbacks) + candidates: Final = chain.from_iterable(router._global_callback_candidates() for router in survivors) + for candidate in litellm.logging_callback_manager.with_custom_logger_keys(candidates, vacated_keys): + self._register_router_selector(candidate) + adaptive_keys: Final = frozenset( + map(litellm.logging_callback_manager.get_custom_logger_key, adaptive_router_hooks) + ) + adaptive_survivor: Final = next((router for router in survivors if router.adaptive_routers), None) + if adaptive_survivor is not None and vacated_keys & adaptive_keys: + adaptive_survivor._sync_adaptive_router_hooks() + + def serve_through_request_logging(self) -> None: + self._request_scoped_selectors = (*self._callback_selectors(), *self._request_adaptive_router_hooks()) + self._request_scoped_policies = (*self._optional_loggers(), *self._slack_alerting_loggers()) + self._release_global_callbacks() + + def bind_to_request_logging(self, logging_obj: LiteLLMLogging) -> None: + logging_obj.add_dynamic_function_callbacks( + async_success=self.deployment_callback_on_success, + failure=self.deployment_callback_on_failure, + async_failure=self.async_deployment_callback_on_failure, + ) + for selector in self._request_scoped_selectors: + logging_obj.replace_global_callback(selector) + for policy in self._request_policies_not_set_globally(): + logging_obj.add_dynamic_callback(policy) + + def _request_policies_not_set_globally(self) -> tuple[CustomLogger, ...]: + global_keys: Final = frozenset( + map( + litellm.logging_callback_manager.get_custom_logger_key, + (callback for callback in litellm.callbacks if isinstance(callback, CustomLogger)), + ) + ) + return tuple( + policy + for policy in self._request_scoped_policies + if litellm.logging_callback_manager.get_custom_logger_key(policy) not in global_keys + ) + + def _request_adaptive_router_hooks(self) -> tuple[CustomLogger, ...]: + from litellm.router_strategy.adaptive_router.hooks import ( + AdaptiveRouterPostCallHook, + ) + + return tuple( + AdaptiveRouterPostCallHook(adaptive_router=tagged.strategy) + for tagged in chain.from_iterable(self.adaptive_routers.values()) + ) + + def _routing_callbacks(self) -> Iterable[object]: + if not (self._request_scoped_selectors or self._request_scoped_policies): + return litellm.callbacks + replaced_keys: Final = frozenset( + map(litellm.logging_callback_manager.get_custom_logger_key, self._request_scoped_selectors) + ) + return chain( + litellm.logging_callback_manager.without_custom_logger_keys(litellm.callbacks, replaced_keys), + self._request_scoped_selectors, + self._request_policies_not_set_globally(), + ) + + def retire_routing_state(self) -> None: + for strategy in (*self._global_callback_candidates(), *self._override_selectors.values()): + if isinstance(strategy, BaseRoutingStrategy | RouterBudgetLimiting): + strategy.retire() + + def _slack_alerting_loggers(self) -> tuple[CustomLogger, ...]: + slack_alerting_logger: Final = getattr(self, "slack_alerting_logger", None) + return (slack_alerting_logger,) if isinstance(slack_alerting_logger, CustomLogger) else () + + def _adaptive_router_hooks(self) -> tuple[CustomLogger, ...]: + from litellm.router_strategy.adaptive_router.hooks import ( + AdaptiveRouterPostCallHook, + ) + + own_strategy_ids: Final = frozenset( + id(tagged.strategy) for tagged in chain.from_iterable(self.adaptive_routers.values()) + ) + hooks: Final = litellm.logging_callback_manager.get_custom_loggers_for_type(AdaptiveRouterPostCallHook) + return tuple( + hook + for hook in hooks + if isinstance(hook, AdaptiveRouterPostCallHook) and id(hook.adaptive_router) in own_strategy_ids + ) + + def _optional_loggers(self) -> tuple[CustomLogger, ...]: + return tuple(callback for callback in self.optional_callbacks or () if isinstance(callback, CustomLogger)) + + def _global_callback_candidates(self) -> tuple[CustomLogger, ...]: + return (*self._callback_selectors(), *self._optional_loggers(), *self._slack_alerting_loggers()) @staticmethod def _create_redis_cache( @@ -1808,7 +1906,7 @@ class Router: return selector @staticmethod - def _register_router_selector(selector: RouterStrategySelector) -> None: + def _register_router_selector(selector: CustomLogger) -> None: if isinstance(selector, LeastBusyLoggingHandler): if isinstance(litellm.input_callback, list): litellm.logging_callback_manager.add_litellm_input_callback(selector) @@ -1817,23 +1915,28 @@ class Router: if isinstance(litellm.callbacks, list): litellm.logging_callback_manager.add_litellm_callback(selector) + def _callback_selectors(self) -> tuple[RouterStrategySelector, ...]: + return tuple( + selector + for selector in ( + *(getattr(self, attr, None) for attr in self._DEFAULT_SELECTOR_ATTR_BY_STRATEGY.values()), + *chain.from_iterable(selectors.values() for selectors in self._group_selectors.values()), + ) + if selector is not None + ) + def _unregister_router_selectors(self, selectors: Sequence[object]) -> None: """ Drop router-owned strategy selectors from litellm's global callback - lists by identity. Used before re-init (`routing_strategy_init` / - `_init_routing_groups`) so repeated `update_settings` calls don't - accumulate dead selectors that keep receiving callback events. + lists. Used before re-init (`routing_strategy_init` / + `_init_routing_groups`) and by `discard` so dead selectors stop + receiving callback events. `function_setup` copies `litellm.callbacks` + into the per-event lists on every request, so those are purged too. """ for selector in selectors: if isinstance(selector, BaseRoutingStrategy): selector.retire() - selector_ids: Final = {id(s) for s in selectors if s is not None} - if not selector_ids: - return - if isinstance(litellm.callbacks, list): - litellm.callbacks = [c for c in litellm.callbacks if id(c) not in selector_ids] - if isinstance(litellm.input_callback, list): - litellm.input_callback = [c for c in litellm.input_callback if id(c) not in selector_ids] + litellm.logging_callback_manager.drop_callbacks(selectors) def _apply_updated_routing_strategy_args(self) -> None: """ @@ -4673,7 +4776,8 @@ class Router: deployment=deployment, forwarded_keys=kwargs.pop(_ALIAS_MARKER_FORWARDED_PARAMS_KWARG, ()) ): kwargs.pop(key, None) - self._merge_tools_from_deployment(deployment=deployment, kwargs=kwargs, defaults=self.default_litellm_params) + defaults: Final = _LITELLM_PARAMS_ADAPTER.validate_python(self.default_litellm_params) + self._merge_tools_from_deployment(deployment=deployment, kwargs=kwargs, defaults=defaults) model_info = deployment.get("model_info", {}).copy() deployment_litellm_model_name = deployment["litellm_params"]["model"] @@ -4817,6 +4921,19 @@ class Router: data=_DEPLOYMENT_PARAMS_ADAPTER.validate_python(deployment["litellm_params"]), ) + deployment_litellm_params: Final = _DEPLOYMENT_PARAMS_ADAPTER.validate_python(deployment["litellm_params"]) + if router_config_requests_stream_usage( + _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), + ( + _LITELLM_PARAMS_ADAPTER.validate_python( + deployment_litellm_params + if isinstance(deployment_litellm_params, Mapping) + else deployment_litellm_params.model_dump(exclude_none=True) + ), + defaults, + ), + ): + kwargs[STREAM_USAGE_INJECTED_FIELD] = True self._update_kwargs_with_default_litellm_params(kwargs=kwargs, metadata_variable_name=metadata_variable_name) hop_kwargs: Final = _FALLBACK_HOP_ADAPTER.validate_python(kwargs) if _is_fallback_hop(hop_kwargs): @@ -6605,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 @@ -9789,7 +9907,7 @@ class Router: Raises: - Rate Limit Exception - If the deployment is over it's tpm/rpm limits """ - for _callback in litellm.callbacks: + for _callback in self._routing_callbacks(): if isinstance(_callback, CustomLogger): _callback.pre_call_check(deployment) @@ -9810,7 +9928,7 @@ class Router: Raises: - Rate Limit Exception - If the deployment is over it's tpm/rpm limits """ - for _callback in litellm.callbacks: + for _callback in self._routing_callbacks(): if isinstance(_callback, CustomLogger): try: await _callback.async_pre_call_check(deployment, parent_otel_span) @@ -9947,7 +10065,7 @@ class Router: - Rate Limit Exception - If the deployment is over it's tpm/rpm limits """ returned_healthy_deployments = healthy_deployments - for _callback in litellm.callbacks: + for _callback in self._routing_callbacks(): if isinstance(_callback, CustomLogger): try: returned_healthy_deployments = await _callback.async_filter_deployments( @@ -10602,6 +10720,7 @@ class Router: changes, otherwise a released router keeps recording turns through its hook.""" from litellm.router_strategy.adaptive_router.hooks import ( AdaptiveRouterPostCallHook, + AdaptiveRouterResponseHeaderHook, ) for callback in litellm.logging_callback_manager.get_custom_loggers_for_type(AdaptiveRouterPostCallHook): @@ -10611,6 +10730,8 @@ class Router: litellm.logging_callback_manager.add_litellm_callback( AdaptiveRouterPostCallHook(adaptive_router=tagged.strategy) ) + if self.adaptive_routers: + litellm.logging_callback_manager.add_litellm_callback(AdaptiveRouterResponseHeaderHook()) def init_adaptive_router_deployment(self, deployment: Deployment) -> None: """ @@ -13948,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/router_strategy/adaptive_router/__init__.py b/litellm/router_strategy/adaptive_router/__init__.py index d7f55ebced9..1f7e6b93169 100644 --- a/litellm/router_strategy/adaptive_router/__init__.py +++ b/litellm/router_strategy/adaptive_router/__init__.py @@ -1,6 +1,9 @@ """Adaptive router strategy. See README.md for design overview.""" from litellm.router_strategy.adaptive_router.adaptive_router import AdaptiveRouter -from litellm.router_strategy.adaptive_router.hooks import AdaptiveRouterPostCallHook +from litellm.router_strategy.adaptive_router.hooks import ( + AdaptiveRouterPostCallHook, + AdaptiveRouterResponseHeaderHook, +) -__all__ = ["AdaptiveRouter", "AdaptiveRouterPostCallHook"] +__all__ = ["AdaptiveRouter", "AdaptiveRouterPostCallHook", "AdaptiveRouterResponseHeaderHook"] diff --git a/litellm/router_strategy/adaptive_router/hooks.py b/litellm/router_strategy/adaptive_router/hooks.py index 92207e07c52..3427b81be68 100644 --- a/litellm/router_strategy/adaptive_router/hooks.py +++ b/litellm/router_strategy/adaptive_router/hooks.py @@ -170,12 +170,7 @@ def _assistant_content_and_tool_calls(response_obj: Any) -> tuple[object, Sequen return content, tool_calls -class AdaptiveRouterPostCallHook(CustomLogger): - """One hook instance per AdaptiveRouter. Registered into litellm.callbacks.""" - - def __init__(self, adaptive_router: AdaptiveRouter) -> None: - self.adaptive_router = adaptive_router - +class AdaptiveRouterResponseHeaderHook(CustomLogger): async def async_post_call_response_headers_hook( self, data: Mapping[str, object], @@ -200,6 +195,13 @@ class AdaptiveRouterPostCallHook(CustomLogger): return None return {ADAPTIVE_ROUTER_RESPONSE_HEADER: chosen} + +class AdaptiveRouterPostCallHook(CustomLogger): + """One hook instance per AdaptiveRouter. Registered into litellm.callbacks.""" + + def __init__(self, adaptive_router: AdaptiveRouter) -> None: + self.adaptive_router = adaptive_router + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): await self._record(kwargs, response_obj, response_status=200) diff --git a/litellm/router_strategy/base_routing_strategy.py b/litellm/router_strategy/base_routing_strategy.py index 51400e6da24..2cea7389808 100644 --- a/litellm/router_strategy/base_routing_strategy.py +++ b/litellm/router_strategy/base_routing_strategy.py @@ -24,6 +24,8 @@ class BaseRoutingStrategy(ABC): self.dual_cache = dual_cache self.redis_increment_operation_queue: list[RedisPipelineIncrementOperation] = [] self._sync_task: asyncio.Task[None] | None = None + self._retired = False + self._flush_after_retire: asyncio.Task[None] | None = None if should_batch_redis_writes: self.setup_sync_task(default_sync_interval) @@ -47,13 +49,25 @@ class BaseRoutingStrategy(ABC): def retire(self) -> None: self.cancel_sync_task() - if not self.redis_increment_operation_queue: + self._retired = True + self._schedule_flush_after_retire() + + def _schedule_flush_after_retire(self) -> None: + if self._flush_after_retire is not None and not self._flush_after_retire.done(): + return + if not self.redis_increment_operation_queue or self.dual_cache.redis_cache is None: return try: loop: Final = asyncio.get_running_loop() except RuntimeError: return - loop.create_task(self._push_in_memory_increments_to_redis()) + self._flush_after_retire = loop.create_task(self._drain_increments_after_retire()) + + async def _drain_increments_after_retire(self) -> None: + if not self.redis_increment_operation_queue: + return + await self._push_in_memory_increments_to_redis() + await self._drain_increments_after_retire() async def cleanup(self): """Cleanup method to be called when shutting down""" @@ -98,6 +112,8 @@ class BaseRoutingStrategy(ABC): self.redis_increment_operation_queue.append(increment_op) self.add_to_in_memory_keys_to_update(key=key) + if self._retired: + self._schedule_flush_after_retire() return result @with_service_target("router_usage") diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index 147ce60a99b..eab232eae18 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/litellm/router_strategy/budget_limiter.py @@ -128,7 +128,9 @@ class RouterBudgetLimiting(CustomLogger): self._redis_increment_queue_lock = asyncio.Lock() self._redis_increment_flush_lock = asyncio.Lock() self._detached_increment_operations: tuple[RedisPipelineIncrementOperation, ...] | None = None - asyncio.create_task(self.periodic_sync_in_memory_spend_with_redis()) + self._sync_task: asyncio.Task[None] = asyncio.create_task(self.periodic_sync_in_memory_spend_with_redis()) + self._retired = False + self._flush_after_retire: asyncio.Task[None] | None = None self.provider_budget_config: GenericBudgetConfigType | None = provider_budget_config self.deployment_budget_config: GenericBudgetConfigType | None = None self.tag_budget_config: GenericBudgetConfigType | None = None @@ -140,6 +142,30 @@ class RouterBudgetLimiting(CustomLogger): if isinstance(litellm.callbacks, list): litellm.logging_callback_manager.add_litellm_callback(self) + def retire(self) -> None: + self._sync_task.cancel() + self._retired = True + self._schedule_flush_after_retire() + + def _schedule_flush_after_retire(self) -> None: + if self._flush_after_retire is not None and not self._flush_after_retire.done(): + return + if not self._has_unwritten_increments(): + return + try: + loop: Final = asyncio.get_running_loop() + except RuntimeError: + return + self._flush_after_retire = loop.create_task(self._drain_increments_after_retire()) + + def _has_unwritten_increments(self) -> bool: + return bool(self.redis_increment_operation_queue) or self._detached_increment_operations is not None + + async def _drain_increments_after_retire(self) -> None: + while self._has_unwritten_increments(): + if not await self._push_in_memory_increments_to_redis(): + await asyncio.sleep(DEFAULT_REDIS_SYNC_INTERVAL) + @with_service_target("router_budgets") async def async_filter_deployments( self, @@ -427,6 +453,8 @@ class RouterBudgetLimiting(CustomLogger): if self.dual_cache.redis_cache is None: return self.redis_increment_operation_queue.append(increment_op) + if self._retired: + self._schedule_flush_after_retire() def _get_redis_increment_queue_lock(self) -> asyncio.Lock: return self._redis_increment_queue_lock 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/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/litellm/router_utils/stream_usage_injection.py b/litellm/router_utils/stream_usage_injection.py new file mode 100644 index 00000000000..946543d9b83 --- /dev/null +++ b/litellm/router_utils/stream_usage_injection.py @@ -0,0 +1,17 @@ +from collections.abc import Mapping, Sequence + + +def router_config_requests_stream_usage( + request_kwargs: Mapping[str, object], configured_params: Sequence[Mapping[str, object]] +) -> bool: + if request_kwargs.get("stream") is not True or "stream_options" in request_kwargs: + return False + return any(_includes_usage(params.get("stream_options")) for params in configured_params) + + +def _includes_usage(stream_options: object) -> bool: + match stream_options: + case {"include_usage": True}: + return True + case _: + return False diff --git a/litellm/types/litellm_params.py b/litellm/types/litellm_params.py index 5f318cc3c69..089f21a33c1 100644 --- a/litellm/types/litellm_params.py +++ b/litellm/types/litellm_params.py @@ -54,6 +54,7 @@ RoutingStrategyName: TypeAlias = Literal[ TRUSTED_CALLBACK_VARS_FIELD: Final = "litellm_trusted_callback_vars" ADDRESSED_RESPONSE_ID_FIELD: Final = "_litellm_addressed_response_id" +STREAM_USAGE_INJECTED_FIELD: Final = "_litellm_stream_usage_injected" WIRE_NAME: Final = "wire_name" @@ -385,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 @@ -397,6 +401,7 @@ class ProxyRequestState: trusted_callback_vars: Mapping[str, str] | None = field(default=None, metadata=wire(TRUSTED_CALLBACK_VARS_FIELD)) addressed_response_id: str | None = field(default=None, metadata=wire(ADDRESSED_RESPONSE_ID_FIELD)) strip_stream_usage: bool | None = field(default=None, metadata=wire("_litellm_strip_stream_usage")) + stream_usage_injected: bool | None = field(default=None, metadata=wire(STREAM_USAGE_INJECTED_FIELD)) client_side_timeout: bool | None = None model_file_id_mapping: Mapping[str, Mapping[str, str]] | 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/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 379eddb3155..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. @@ -1126,83 +1131,90 @@ def function_setup( dynamic_callbacks: Final[list[str | Callable | CustomLogger] | None] = kwargs.pop("callbacks", None) all_callbacks: Final = get_dynamic_callbacks(dynamic_callbacks=dynamic_callbacks) - if len(all_callbacks) > 0: - for callback in all_callbacks: - # check if callback is a string - e.g. "lago", "openmeter" - if isinstance(callback, str): - callback = litellm.litellm_core_utils.litellm_logging._init_custom_logger_compatible_class( - callback, - internal_usage_cache=None, - llm_router=None, - ) - if callback is None or any( - type(cb) is type(callback) for cb in litellm._async_success_callback - ): # don't double add a callback - continue - if callback not in litellm.input_callback: - litellm.input_callback.append(callback) - if callback not in litellm.success_callback: - litellm.logging_callback_manager.add_litellm_success_callback(callback) - if callback not in litellm.failure_callback: - litellm.logging_callback_manager.add_litellm_failure_callback(callback) - if callback not in litellm._async_success_callback: - litellm.logging_callback_manager.add_litellm_async_success_callback(callback) - if callback not in litellm._async_failure_callback: - litellm.logging_callback_manager.add_litellm_async_failure_callback(callback) - print_verbose(f"Initialized litellm callbacks, Async Success Callbacks: {litellm._async_success_callback}") + with litellm.logging_callback_manager.callback_lists_lock: + if len(all_callbacks) > 0: + for callback in all_callbacks: + # check if callback is a string - e.g. "lago", "openmeter" + if isinstance(callback, str): + callback = litellm.litellm_core_utils.litellm_logging._init_custom_logger_compatible_class( + callback, + internal_usage_cache=None, + llm_router=None, + ) + if callback is None or any( + type(cb) is type(callback) for cb in litellm._async_success_callback + ): # don't double add a callback + continue + if callback not in litellm.input_callback: + litellm.input_callback.append(callback) + if callback not in litellm.success_callback: + litellm.logging_callback_manager.add_litellm_success_callback(callback) + if callback not in litellm.failure_callback: + litellm.logging_callback_manager.add_litellm_failure_callback(callback) + if callback not in litellm._async_success_callback: + litellm.logging_callback_manager.add_litellm_async_success_callback(callback) + if callback not in litellm._async_failure_callback: + litellm.logging_callback_manager.add_litellm_async_failure_callback(callback) + print_verbose( + f"Initialized litellm callbacks, Async Success Callbacks: {litellm._async_success_callback}" + ) - if ( - len(litellm.input_callback) > 0 or len(litellm.success_callback) > 0 or len(litellm.failure_callback) > 0 - ) and len(callback_list) == 0: - callback_list = list(set(litellm.input_callback + litellm.success_callback + litellm.failure_callback)) + callbacks_to_set_up: Final = ( + list(set(litellm.input_callback + litellm.success_callback + litellm.failure_callback)) + if len(callback_list) == 0 + else [] + ) + if callbacks_to_set_up: + callback_list = callbacks_to_set_up + ## ASYNC CALLBACKS - safety net for callbacks added via direct append + if len(litellm.input_callback) > 0: + removed_async_items = [] + for index, callback in enumerate(litellm.input_callback): + if coroutine_checker.is_async_callable(callback): + litellm._async_input_callback.append(callback) + removed_async_items.append(index) + + # Pop the async items from input_callback in reverse order to avoid index issues + for index in reversed(removed_async_items): + litellm.input_callback.pop(index) + if len(litellm.success_callback) > 0: + removed_async_items = [] + for index, callback in enumerate(litellm.success_callback): + if coroutine_checker.is_async_callable(callback): + litellm.logging_callback_manager.add_litellm_async_success_callback(callback) + removed_async_items.append(index) + elif callback == "dynamodb" or callback == "openmeter": + # dynamo is an async callback, it's used for the proxy and needs to be async + # we only support async dynamo db logging for acompletion/aembedding since that's used on proxy + litellm.logging_callback_manager.add_litellm_async_success_callback(callback) + removed_async_items.append(index) + elif callback in litellm._known_custom_logger_compatible_callbacks and isinstance(callback, str): + add_custom_logger_callback_to_specific_event(callback, "success") + + # Pop the async items from success_callback in reverse order to avoid index issues + for index in reversed(removed_async_items): + litellm.success_callback.pop(index) + + if len(litellm.failure_callback) > 0: + removed_async_items = [] + for index, callback in enumerate(litellm.failure_callback): + if coroutine_checker.is_async_callable(callback): + litellm.logging_callback_manager.add_litellm_async_failure_callback(callback) + removed_async_items.append(index) + elif callback in litellm._known_custom_logger_compatible_callbacks and isinstance(callback, str): + add_custom_logger_callback_to_specific_event(callback, "failure") + + # Pop the async items from failure_callback in reverse order to avoid index issues + for index in reversed(removed_async_items): + litellm.failure_callback.pop(index) + if callbacks_to_set_up: get_set_callbacks: Final[Callable[[], Callable[..., None]]] = litellm_utils.get_set_callbacks - get_set_callbacks()(callback_list=callback_list, function_id=function_id) - ## ASYNC CALLBACKS - safety net for callbacks added via direct append - if len(litellm.input_callback) > 0: - removed_async_items = [] - for index, callback in enumerate(litellm.input_callback): - if coroutine_checker.is_async_callable(callback): - litellm._async_input_callback.append(callback) - removed_async_items.append(index) - - # Pop the async items from input_callback in reverse order to avoid index issues - for index in reversed(removed_async_items): - litellm.input_callback.pop(index) - if len(litellm.success_callback) > 0: - removed_async_items = [] - for index, callback in enumerate(litellm.success_callback): - if coroutine_checker.is_async_callable(callback): - litellm.logging_callback_manager.add_litellm_async_success_callback(callback) - removed_async_items.append(index) - elif callback == "dynamodb" or callback == "openmeter": - # dynamo is an async callback, it's used for the proxy and needs to be async - # we only support async dynamo db logging for acompletion/aembedding since that's used on proxy - litellm.logging_callback_manager.add_litellm_async_success_callback(callback) - removed_async_items.append(index) - elif callback in litellm._known_custom_logger_compatible_callbacks and isinstance(callback, str): - add_custom_logger_callback_to_specific_event(callback, "success") - - # Pop the async items from success_callback in reverse order to avoid index issues - for index in reversed(removed_async_items): - litellm.success_callback.pop(index) - - if len(litellm.failure_callback) > 0: - removed_async_items = [] - for index, callback in enumerate(litellm.failure_callback): - if coroutine_checker.is_async_callable(callback): - litellm.logging_callback_manager.add_litellm_async_failure_callback(callback) - removed_async_items.append(index) - elif callback in litellm._known_custom_logger_compatible_callbacks and isinstance(callback, str): - add_custom_logger_callback_to_specific_event(callback, "failure") - - # Pop the async items from failure_callback in reverse order to avoid index issues - for index in reversed(removed_async_items): - litellm.failure_callback.pop(index) + get_set_callbacks()(callback_list=callbacks_to_set_up, function_id=function_id) ### DYNAMIC CALLBACKS ### - dynamic_success_callbacks: list[str | Callable | CustomLogger] | None = None - dynamic_async_success_callbacks: list[str | Callable | CustomLogger] | None = None - dynamic_failure_callbacks: list[str | Callable | CustomLogger] | None = None - dynamic_async_failure_callbacks: Final[list[str | Callable | CustomLogger] | None] = None + dynamic_success_callbacks: list[str | Callable[..., object] | CustomLogger] | None = None + dynamic_async_success_callbacks: list[str | Callable[..., object] | CustomLogger] | None = None + dynamic_failure_callbacks: list[str | Callable[..., object] | CustomLogger] | None = None + dynamic_async_failure_callbacks: Final[list[str | Callable[..., object] | CustomLogger] | None] = None if kwargs.get("success_callback", None) is not None and isinstance(kwargs["success_callback"], list): removed_async_items = [] for index, callback in enumerate(kwargs["success_callback"]): @@ -7962,12 +7974,16 @@ class TextCompletionStreamWrapper: text_choices: Final = TextChoices() if isinstance(chunk, Choices): # chunk should always be of type StreamingChoices raise Exception - delta: Final = chunk["choices"][0]["delta"] - text_choices["text"] = delta["content"] - text_choices["reasoning_content"] = delta.get("reasoning_content") - text_choices["index"] = chunk["choices"][0]["index"] - text_choices["finish_reason"] = chunk["choices"][0]["finish_reason"] - response["choices"] = [text_choices] + + if chunk["choices"]: + delta: Final = chunk["choices"][0]["delta"] + text_choices["text"] = delta["content"] + text_choices["reasoning_content"] = delta.get("reasoning_content") + text_choices["index"] = chunk["choices"][0]["index"] + text_choices["finish_reason"] = chunk["choices"][0]["finish_reason"] + response["choices"] = [text_choices] + else: + response["choices"] = [] # only pass usage when stream_options["include_usage"] is True if self.stream_options and self.stream_options.get("include_usage", False) is True: diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 3d89af9d691..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": { @@ -76709,14 +76713,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, @@ -77018,5 +77022,33 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": false + }, + "azure_ai/MAI-Code-1.1-Flash": { + "cache_read_input_token_cost": 2e-08, + "deprecation_date": "2026-12-09", + "input_cost_per_token": 2e-07, + "litellm_provider": "azure_ai", + "max_input_tokens": 256000, + "max_output_tokens": 64000, + "max_tokens": 64000, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://prices.azure.com/api/retail/prices?$filter=serviceName%20eq%20'Foundry%20Models'%20and%20armRegionName%20eq%20'eastus'%20and%20priceType%20eq%20'Consumption'", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "supports_vision": true } } diff --git a/tests/audio_tests/gemini_tts_speech.mp3 b/tests/audio_tests/gemini_tts_speech.mp3 new file mode 100644 index 00000000000..71a916ad6c0 Binary files /dev/null and b/tests/audio_tests/gemini_tts_speech.mp3 differ diff --git a/tests/audio_tests/gemini_tts_speech.ogg b/tests/audio_tests/gemini_tts_speech.ogg new file mode 100644 index 00000000000..13397fb0c9d Binary files /dev/null and b/tests/audio_tests/gemini_tts_speech.ogg differ 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/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py index 55435c51648..af2f386a362 100644 --- a/tests/code_coverage_tests/router_code_coverage.py +++ b/tests/code_coverage_tests/router_code_coverage.py @@ -146,6 +146,18 @@ ignored_function_names = [ "async_admit_deployment_turn", # Bound as the WebSocket turn_admission partial, tested through it in test_router_helper_utils.py "_deployment_as_configured_now", # Read by async_admit_deployment_turn, tested through turn_admission in test_router_helper_utils.py "_admit_turn_through_routing_strategy", # Run by async_admit_deployment_turn, tested through turn_admission in test_router_helper_utils.py + "_callback_selectors", # Run by discard, tested through the discard tests in test_router_helper_utils.py + "bind_to_request_logging", # Run by the proxy user_config route, tested through the user_config tests in test_route_llm_request.py + "serve_through_request_logging", # Run by the proxy user_config route, tested through the user_config tests in test_route_llm_request.py + "_request_adaptive_router_hooks", # Run by serve_through_request_logging, tested through the user_config adaptive test + "_release_global_callbacks", # Run by discard, tested through the discard tests in test_router_helper_utils.py + "_routing_callbacks", # Read by the pre-call checks, tested through the user_config limit tests in test_route_llm_request.py + "_request_policies_not_set_globally", # Read by bind_to_request_logging, tested through the user_config budget test in test_route_llm_request.py + "_optional_loggers", # Read by discard and serve_through_request_logging, tested through the budget limiter discard test + "retire_routing_state", # Run by discard, tested through the budget limiter discard test + "_slack_alerting_loggers", # Run by discard, tested through the Slack alerting discard test in test_router_helper_utils.py + "_adaptive_router_hooks", # Run by discard, tested through the adaptive router discard tests in test_router_helper_utils.py + "_global_callback_candidates", # Run by discard, tested through the survivor discard tests in test_router_helper_utils.py "without_deployments_that_cannot_batch", # Tested through the Azure and OpenAI mixed-group acreate_batch tests in test_router.py "batch_create_routing", # Tested through the mixed-group acreate_batch tests in test_router.py "is_batch_create_routing", # Tested through the chat-after-batch-create test in test_router.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/gateway/perf_ci_config.yml b/tests/e2e/gateway/perf_ci_config.yml index 7ab6637e985..767db313847 100644 --- a/tests/e2e/gateway/perf_ci_config.yml +++ b/tests/e2e/gateway/perf_ci_config.yml @@ -46,6 +46,11 @@ model_list: model: openai/gpt-6-luna api_base: os.environ/E2E_PERF_MOCK_OPENAI_URL api_key: perf-mock-not-checked + - model_name: gpt-6-luna-responses + litellm_params: + model: openai/responses/gpt-6-luna + api_base: os.environ/E2E_PERF_MOCK_OPENAI_URL + api_key: perf-mock-not-checked - model_name: claude-sonnet-5 litellm_params: model: anthropic/claude-sonnet-5 @@ -58,6 +63,13 @@ model_list: aws_region_name: us-east-1 aws_access_key_id: perf-mock-not-checked aws_secret_access_key: perf-mock-not-checked + - model_name: claude-sonnet-5-bedrock-converse + litellm_params: + model: bedrock/converse/anthropic.claude-sonnet-5 + aws_bedrock_runtime_endpoint: os.environ/E2E_PERF_MOCK_BEDROCK_URL + aws_region_name: us-east-1 + aws_access_key_id: perf-mock-not-checked + aws_secret_access_key: perf-mock-not-checked guardrails: - guardrail_name: perf-pre litellm_params: diff --git a/tests/e2e/perf/cases/chat_completions_2rps.py b/tests/e2e/perf/cases/chat_completions_2rps.py new file mode 100644 index 00000000000..9792d99c14f --- /dev/null +++ b/tests/e2e/perf/cases/chat_completions_2rps.py @@ -0,0 +1,15 @@ +"""`/v1/chat/completions` at 2 rps: `gpt-6-luna` is served by the OpenAI mock's chat completions +route, `gpt-6-luna-responses` (`openai/responses/gpt-6-luna`) is bridged to its Responses route.""" + +from __future__ import annotations + +from typing import Final + +from .direct_vs_bridge import direct_vs_bridge + +CASE: Final = direct_vs_bridge( + "chat_completions_2rps", + "chat_completions", + direct="gpt-6-luna", + bridge="gpt-6-luna-responses", +) diff --git a/tests/e2e/perf/cases/direct_vs_bridge.py b/tests/e2e/perf/cases/direct_vs_bridge.py new file mode 100644 index 00000000000..13f609c8299 --- /dev/null +++ b/tests/e2e/perf/cases/direct_vs_bridge.py @@ -0,0 +1,31 @@ +"""The mixed case's traffic on one public API, half of it on a deployment the gateway serves on +that provider's own API (direct) and half on a deployment it has to translate to another provider +API (bridge). Same 2 rps, request count, shapes, stream share, callers, guardrails and limits as +`mixed_size_requests_2rps`, so a gap between the two halves and the mixed baseline is the +translation's cost.""" + +from __future__ import annotations + +from dataclasses import replace + +from setup.performance_test_case import Endpoint, PerformanceTestCase, VirtualKey + +from .mixed_size_requests_2rps import CASE as MIXED + + +def _scoped(key: VirtualKey, direct: str) -> VirtualKey: + return replace(key, models=(direct,)) if key.models else key + + +def direct_vs_bridge(name: str, endpoint: Endpoint, *, direct: str, bridge: str) -> PerformanceTestCase: + return replace( + MIXED, + name=name, + teams=( + replace(MIXED.team("team-a"), models=(direct, bridge)), + replace(MIXED.team("team-b"), models=(bridge,)), + ), + virtual_keys=tuple(_scoped(key, direct) for key in MIXED.virtual_keys), + model_groups=(direct, bridge), + traffic=replace(MIXED.traffic, endpoints={endpoint: 1.0}), + ) diff --git a/tests/e2e/perf/cases/event_loop_stalls.py b/tests/e2e/perf/cases/event_loop_stalls.py index 0d67a98257f..8a5654cbf7b 100644 --- a/tests/e2e/perf/cases/event_loop_stalls.py +++ b/tests/e2e/perf/cases/event_loop_stalls.py @@ -8,8 +8,8 @@ new hook that deep-copies the request inline all land here as a new row. Runs on `PYTHONASYNCIODEBUG=1` and `LITELLM_EVENT_LOOP=asyncio` (uvloop reports nothing and crashed under debug locally), and the runner `E2E_PERF_EVENT_LOOP=1`, so the nightly build, which sets neither, skips it. The allowlist counts are twice the highest count any of ten back-to-back local runs showed on that row (rows that showed nothing in -those ten but did in the earlier before/after legs for 46925 and 46967 get 2), so each run sits at or under half -of every cap and a listed location doubling in frequency is a failure, not a flake. The +those ten but did in the earlier before/after legs for 46925 and 46967, or in the first scheduled-lane run, get 2), +so each run sits at or under half of every cap and a listed location doubling in frequency is a failure, not a flake. The `garbage collection` bucket is the gen-2 passes the collector logs in debug mode (0.6 to 0.7 s each, one every minute or two of this traffic, plus the ramp while the heap fills on a local whole-file run), which asyncio would otherwise charge to a random task; its count is the signal for a change that makes much more garbage.""" @@ -44,6 +44,7 @@ CASE: Final = replace( KnownStall("run_coroutine_job", max_count=2), KnownStall("user_api_key_auth.py", max_count=2), KnownStall("route_streaming_logging_to_handler", max_count=2), + KnownStall("TaskHandle._run_coro", max_count=2), ), ), ) diff --git a/tests/e2e/perf/cases/messages_2rps.py b/tests/e2e/perf/cases/messages_2rps.py new file mode 100644 index 00000000000..df5d221d430 --- /dev/null +++ b/tests/e2e/perf/cases/messages_2rps.py @@ -0,0 +1,16 @@ +"""`/v1/messages` at 2 rps: `claude-sonnet-5` is served by the Anthropic mock's Messages route, +`claude-sonnet-5-bedrock-converse` (`bedrock/converse/anthropic.claude-sonnet-5`) is bridged to the +Bedrock mock's Converse route.""" + +from __future__ import annotations + +from typing import Final + +from .direct_vs_bridge import direct_vs_bridge + +CASE: Final = direct_vs_bridge( + "messages_2rps", + "messages", + direct="claude-sonnet-5", + bridge="claude-sonnet-5-bedrock-converse", +) 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/responses_2rps.py b/tests/e2e/perf/cases/responses_2rps.py new file mode 100644 index 00000000000..dd55d536eb6 --- /dev/null +++ b/tests/e2e/perf/cases/responses_2rps.py @@ -0,0 +1,15 @@ +"""`/v1/responses` at 2 rps: `gpt-6-luna` is served by the OpenAI mock's Responses route, +`claude-sonnet-5` is bridged to the Anthropic mock's Messages route.""" + +from __future__ import annotations + +from typing import Final + +from .direct_vs_bridge import direct_vs_bridge + +CASE: Final = direct_vs_bridge( + "responses_2rps", + "responses", + direct="gpt-6-luna", + bridge="claude-sonnet-5", +) 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/event_loop.py b/tests/e2e/perf/setup/event_loop.py index 1c582e24fd8..58144424544 100644 --- a/tests/e2e/perf/setup/event_loop.py +++ b/tests/e2e/perf/setup/event_loop.py @@ -18,7 +18,7 @@ import json import os import re from collections.abc import Iterable -from dataclasses import dataclass, replace +from dataclasses import asdict, dataclass, replace from datetime import datetime from itertools import groupby from pathlib import Path @@ -80,11 +80,24 @@ class StallRow: total_seconds: float max_seconds: float allowed_by: KnownStall | None + """Every hold in the row, longest first.""" + holds: tuple[float, ...] = () @property def offending(self) -> bool: return self.allowed_by is None or self.count > self.allowed_by.max_count + @property + def over_allowance(self) -> int: + return self.count if self.allowed_by is None else max(0, self.count - self.allowed_by.max_count) + + @property + def note(self) -> str: + allowed: Final = ( + f"{self.count} of {self.allowed_by.max_count} allowed" if self.allowed_by else "not on the allowlist" + ) + return f"{allowed}, max {self.max_seconds:.2f} s" + def _short(path: str) -> str: if "site-packages/" in path: @@ -161,6 +174,7 @@ def summarize(spec: EventLoopStalls, stalls: Iterable[Stall]) -> tuple[StallRow, total_seconds=sum(stall.seconds for stall in group), max_seconds=max(stall.seconds for stall in group), allowed_by=_allowed_by(spec, coroutine, where), + holds=tuple(sorted((stall.seconds for stall in group), reverse=True)), ) for (coroutine, where), group in ( (key, tuple(items)) for key, items in groupby(over, key=lambda s: (s.coroutine, s.where)) @@ -182,28 +196,65 @@ def offenders(rows: Iterable[StallRow]) -> tuple[str, ...]: ) -def _over_allowance(row: StallRow) -> int: - return row.count if row.allowed_by is None else max(0, row.count - row.allowed_by.max_count) +TOP_HOLDS: Final = 10 + + +@dataclass(frozen=True, slots=True) +class Slice: + """One `_breakdown` row in the shape the other metrics use (label, value, share) plus the allowlist verdict the + dashboard shows under the label and colours by. `share` is None where the rows are not parts of a total.""" + + label: str + value: float + share: float | None + note: str | None + over: bool + + +def _label(row: StallRow) -> str: + return f"{row.coroutine} {row.where}".strip() + + +def _slice(row: StallRow, value: float, total: float) -> Slice: + return Slice(label=_label(row), value=value, share=value / total, note=row.note, over=row.offending) + + +def _breakdown(name: str, slices: Iterable[Slice]) -> tuple[tuple[str, str], ...]: + rows: Final = tuple(asdict(row) for row in slices) + return ((f"perf_event_loop_{name}_breakdown", json.dumps(rows)),) if rows else () + + +def _longest_holds(table: tuple[StallRow, ...]) -> tuple[Slice, ...]: + holds: Final = sorted(((seconds, row) for row in table for seconds in row.holds), key=lambda hold: -hold[0]) + return tuple( + Slice(label=_label(row), value=seconds, share=None, note=None, over=row.offending) + for seconds, row in holds[:TOP_HOLDS] + ) def junit_properties(spec: EventLoopStalls, rows: Iterable[StallRow]) -> tuple[tuple[str, float | str], ...]: """The stall table as `perf_event_loop_*` properties for the dashboard: every hold over the threshold with the allowlist's total as its limit, the holds past their allowance (limit 0, the case's one assertion), the longest - and the summed hold, and a `_breakdown` of holds per location in the shape the other metrics use.""" + and the summed hold, each with a `_breakdown` the dashboard opens on click: holds and seconds per location, the + locations past their cap by how far over they are, and the ten longest single holds with where they landed.""" table: Final = tuple(rows) holds: Final = sum(row.count for row in table) - breakdown: Final = tuple( - {"label": f"{row.coroutine} {row.where}".strip(), "value": row.count, "share": row.count / holds} - for row in table - ) + offending: Final = sum(row.over_allowance for row in table) + total_seconds: Final = sum(row.total_seconds for row in table) return ( ("perf_event_loop_holds", holds), ("perf_event_loop_holds_limit", sum(known.max_count for known in spec.allowlist)), - ("perf_event_loop_offending_holds", sum(map(_over_allowance, table))), + ("perf_event_loop_offending_holds", offending), ("perf_event_loop_offending_holds_limit", 0), ("perf_event_loop_hold_seconds_max", max((row.max_seconds for row in table), default=0.0)), - ("perf_event_loop_hold_seconds_total", sum(row.total_seconds for row in table)), - *((("perf_event_loop_holds_breakdown", json.dumps(breakdown)),) if breakdown else ()), + ("perf_event_loop_hold_seconds_total", total_seconds), + *_breakdown("holds", (_slice(row, row.count, holds) for row in table)), + *_breakdown( + "offending_holds", + (_slice(row, row.over_allowance, offending) for row in table if row.over_allowance), + ), + *_breakdown("hold_seconds_total", (_slice(row, row.total_seconds, total_seconds) for row in table)), + *_breakdown("hold_seconds_max", _longest_holds(table)), ) 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 ccdae6398a9..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: @@ -565,6 +662,7 @@ def _diagnostics(observation: Observation) -> list[str]: return [ " diagnostics (not limits):", f" bodies off reply {observation.bodies_off_reply()}", + f" mock routes {_mock_routes(observation)}", f" guardrail calls per request {len(observation.guardrail_calls) / count:.3f}", f" guardrail missed {observation.guardrail_missed() if observation.guardrail_calls else 'n/a'}", f" retried share {len(observation.retried_markers) / count:.4f}", @@ -575,6 +673,12 @@ def _diagnostics(observation: Observation) -> list[str]: ] +def _mock_routes(observation: Observation) -> str: + """Which provider route each request landed on, so a direct-vs-bridge case shows its split.""" + counts: Final = Counter(f"{seen.provider}/{seen.route}" for seen in observation.mock_requests) + return ", ".join(f"{route} {count}" for route, count in sorted(counts.items())) or "n/a" + + def _datastore_lines(observation: Observation) -> list[str]: """What the per-request Postgres and Redis counts are made of, so a red run names the statement or command that grew. `plumbing` (BEGIN, COMMIT, session settings) and `unclassified` count in neither limit.""" 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_event_loop.py b/tests/e2e_harness/perf/setup/test_event_loop.py index e03e985709b..75989b38275 100644 --- a/tests/e2e_harness/perf/setup/test_event_loop.py +++ b/tests/e2e_harness/perf/setup/test_event_loop.py @@ -37,8 +37,14 @@ class _Slice(BaseModel): model_config = ConfigDict(frozen=True) label: str - value: int - share: float + value: float + share: float | None + note: str | None + over: bool + + +def _breakdown(properties: dict[str, float | str], name: str) -> list[_Slice]: + return TypeAdapter(list[_Slice]).validate_json(str(properties[f"perf_event_loop_{name}_breakdown"])) def test_parse_keeps_the_coroutine_its_location_and_the_hold() -> None: @@ -116,13 +122,56 @@ def test_junit_properties_count_every_hold_and_only_the_holds_past_their_allowan assert properties["perf_event_loop_offending_holds_limit"] == 0 assert properties["perf_event_loop_hold_seconds_max"] == 0.731 assert math.isclose(float(properties["perf_event_loop_hold_seconds_total"]), 0.302 + 2 * 0.412 + 3 * 0.731) - breakdown: Final = TypeAdapter(list[_Slice]).validate_json(str(properties["perf_event_loop_holds_breakdown"])) - assert [(slice.label, slice.value) for slice in breakdown] == [ - ("RequestResponseCycle.run_asgi() uvicorn/protocols/http/h11_impl.py:416", 3), - ("TaskHandle._run_coro() anyio/_core/_tasks.py:275", 2), - ("LifespanOn.main() uvicorn/lifespan/on.py:86", 1), + breakdown: Final = _breakdown(properties, "holds") + assert [(slice.label, slice.value, slice.note, slice.over) for slice in breakdown] == [ + ( + "RequestResponseCycle.run_asgi() uvicorn/protocols/http/h11_impl.py:416", + 3, + "3 of 2 allowed, max 0.73 s", + True, + ), + ("TaskHandle._run_coro() anyio/_core/_tasks.py:275", 2, "not on the allowlist, max 0.41 s", True), + ("LifespanOn.main() uvicorn/lifespan/on.py:86", 1, "1 of 3 allowed, max 0.30 s", False), ] - assert math.isclose(sum(slice.share for slice in breakdown), 1.0) + assert math.isclose(sum(slice.share or 0.0 for slice in breakdown), 1.0) + + +def test_the_offending_breakdown_lists_only_the_locations_past_their_cap_by_how_far_over_they_are() -> None: + rows: Final = summarize(_SPEC, parse_stalls((_STARTUP, _STREAM, _STREAM, _HANDLER, _HANDLER, _HANDLER))) + offending: Final = _breakdown(dict(junit_properties(_SPEC, rows)), "offending_holds") + assert [(slice.label, slice.value, slice.share, slice.over) for slice in offending] == [ + ("RequestResponseCycle.run_asgi() uvicorn/protocols/http/h11_impl.py:416", 1, 1 / 3, True), + ("TaskHandle._run_coro() anyio/_core/_tasks.py:275", 2, 2 / 3, True), + ] + clean: Final = dict(junit_properties(_SPEC, summarize(_SPEC, parse_stalls((_HANDLER, _HANDLER))))) + assert clean["perf_event_loop_offending_holds"] == 0 + assert "perf_event_loop_offending_holds_breakdown" not in clean + + +def test_the_seconds_breakdowns_carry_the_time_per_location_and_the_longest_single_holds() -> None: + long_handler: Final = _HANDLER.replace("0.731", "0.950") + rows: Final = summarize(_SPEC, parse_stalls((_STARTUP, _STREAM, long_handler, _HANDLER))) + properties: Final = dict(junit_properties(_SPEC, rows)) + per_location: Final = _breakdown(properties, "hold_seconds_total") + assert [(slice.label.split()[0], slice.value) for slice in per_location] == [ + ("RequestResponseCycle.run_asgi()", 0.950 + 0.731), + ("TaskHandle._run_coro()", 0.412), + ("LifespanOn.main()", 0.302), + ] + assert math.isclose(sum(slice.share or 0.0 for slice in per_location), 1.0) + longest: Final = _breakdown(properties, "hold_seconds_max") + assert [(slice.label.split()[0], slice.value, slice.share, slice.over) for slice in longest] == [ + ("RequestResponseCycle.run_asgi()", 0.950, None, False), + ("RequestResponseCycle.run_asgi()", 0.731, None, False), + ("TaskHandle._run_coro()", 0.412, None, True), + ("LifespanOn.main()", 0.302, None, False), + ] + + +def test_the_longest_holds_breakdown_stops_at_ten() -> None: + rows: Final = summarize(_SPEC, parse_stalls(tuple(_STARTUP for _ in range(12)))) + longest: Final = _breakdown(dict(junit_properties(_SPEC, rows)), "hold_seconds_max") + assert len(longest) == 10 def test_junit_properties_without_holds_report_zeros_and_no_breakdown() -> None: @@ -130,4 +179,4 @@ def test_junit_properties_without_holds_report_zeros_and_no_breakdown() -> None: assert properties["perf_event_loop_holds"] == 0 assert properties["perf_event_loop_offending_holds"] == 0 assert properties["perf_event_loop_hold_seconds_max"] == 0 - assert "perf_event_loop_holds_breakdown" not in properties + assert not [name for name in properties if name.endswith("_breakdown")] 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 ) diff --git a/tests/integration/compatibility/test_openai_consumer.py b/tests/integration/compatibility/test_openai_consumer.py index a49f20f4711..6150995723b 100644 --- a/tests/integration/compatibility/test_openai_consumer.py +++ b/tests/integration/compatibility/test_openai_consumer.py @@ -459,7 +459,6 @@ async def test_openai_clients_parse_text_completion_streams_and_choices(gateway: def test_text_completion_usage_chunk_has_empty_choices(gateway: Gateway) -> None: - pytest.skip("BUG: the /v1/completions include_usage chunk carries choices [{index: 0}] instead of []") identity: Final = f"cmpl-usage-{uuid.uuid4().hex}" prompt: Final = f"raw-stream {identity}" with wire_server(_text_provider(identity, {prompt})) as wire, gateway.scenario() as scenario: diff --git a/tests/integration/database/test_cleanup_transactions.py b/tests/integration/database/test_cleanup_transactions.py index 72da2bb48ee..202657d180e 100644 --- a/tests/integration/database/test_cleanup_transactions.py +++ b/tests/integration/database/test_cleanup_transactions.py @@ -113,7 +113,7 @@ async def test_size_rule_deletes_the_oldest_rows_until_under_the_target_and_does """ Deleting rows never shrinks the table file, so the size pass must stop on live bytes, not on pg_total_relation_size, taking the oldest rows first and keeping the tool index on the same cutoff. Once - autovacuum has turned the deleted rows into free space the catalog no longer shows them as dead, and the + VACUUM has turned the deleted rows into free space the catalog no longer shows them as dead, and the files still measure the same; the next check must read what the last run freed from its saved count and leave the table alone instead of calling the whole file live again. """ @@ -140,13 +140,15 @@ async def test_size_rule_deletes_the_oldest_rows_until_under_the_target_and_does setup.execute( sql.SQL( 'CREATE TABLE {} (request_id text, "startTime" timestamp NOT NULL, payload text, ' - 'PRIMARY KEY (request_id, "startTime"))' + 'PRIMARY KEY (request_id, "startTime")) ' + 'WITH (autovacuum_enabled = false, toast.autovacuum_enabled = false)' ).format(logs) ) setup.execute( sql.SQL( "CREATE TABLE {} (request_id text, tool_name text, start_time timestamp NOT NULL, " - "PRIMARY KEY (request_id, tool_name))" + "PRIMARY KEY (request_id, tool_name)) " + "WITH (autovacuum_enabled = false, toast.autovacuum_enabled = false)" ).format(index) ) setup.execute(sql.SQL("CREATE TABLE {} (param_name text PRIMARY KEY, param_value jsonb)").format(config)) 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/integration/observability/_openinference_support.py b/tests/integration/observability/_openinference_support.py index b46503703c7..58785c89033 100644 --- a/tests/integration/observability/_openinference_support.py +++ b/tests/integration/observability/_openinference_support.py @@ -312,7 +312,7 @@ def _chat_caller_stream_usage_event(event: dict[str, JsonValue], model: str) -> "object": event["object"], "created": 1, "model": model, - "choices": [{"index": 0, "delta": {}}], + "choices": [], "usage": {**_json_object_value(event["usage"]), "cost": 4.05e-6}, } 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/integration/providers/test_bedrock_mantle_claude_chat_wire.py b/tests/integration/providers/test_bedrock_mantle_claude_chat_wire.py index 2badf810b6a..d4f9533fcdc 100644 --- a/tests/integration/providers/test_bedrock_mantle_claude_chat_wire.py +++ b/tests/integration/providers/test_bedrock_mantle_claude_chat_wire.py @@ -1,6 +1,7 @@ import json import re from collections.abc import Callable, Mapping +from itertools import chain from typing import Final from uuid import uuid4 @@ -263,13 +264,17 @@ def _assert_chunk_envelope(chunks: tuple[dict[str, JsonValue], ...], model: str, return next(iter(identities)) +def _content_chunks(chunks: tuple[dict[str, JsonValue], ...]) -> tuple[dict[str, JsonValue], ...]: + return tuple(chunk for chunk in chunks if chunk.get("usage") is None) + + def _finish_reasons(chunks: tuple[dict[str, JsonValue], ...]) -> list[JsonValue]: - reasons: Final = (_only_choice(chunk).get("finish_reason") for chunk in chunks) + reasons: Final = (_only_choice(chunk).get("finish_reason") for chunk in _content_chunks(chunks)) return [reason for reason in reasons if reason is not None] def _deltas(chunks: tuple[dict[str, JsonValue], ...]) -> tuple[dict[str, JsonValue], ...]: - return tuple(object_value(_only_choice(chunk)["delta"]) for chunk in chunks) + return tuple(object_value(_only_choice(chunk)["delta"]) for chunk in _content_chunks(chunks)) def _sdk_base(gateway: Gateway) -> str: @@ -301,12 +306,12 @@ def test_chat_completion_on_a_mantle_claude_id_posts_native_messages_and_spends_ @pytest.mark.parametrize( ("stream_options", "usage_chunks"), [ - pytest.param({"stream_options": {"include_usage": True}}, [_USAGE], id="include_usage"), + pytest.param({"stream_options": {"include_usage": True}}, [(True, *_USAGE)], id="include_usage"), pytest.param({}, [], id="no_usage"), ], ) def test_chat_stream_on_a_mantle_claude_id_relays_openai_chunks_from_the_native_stream( - gateway: Gateway, stream_options: dict[str, JsonValue], usage_chunks: list[tuple[int, int, int]] + gateway: Gateway, stream_options: dict[str, JsonValue], usage_chunks: list[tuple[JsonValue, ...]] ) -> None: marker: Final = uuid4().hex expected: Final = _outbound(_HAIKU, [_user_turn(marker)], stream=True) @@ -322,9 +327,9 @@ def test_chat_stream_on_a_mantle_claude_id_relays_openai_chunks_from_the_native_ identity: Final = _assert_chunk_envelope(chunks, model, response.text) assert "".join(str(delta.get("content") or "") for delta in _deltas(chunks)) == _answer(marker), response.text assert _finish_reasons(chunks) == ["stop"], response.text - assert [_token_usage(chunk) for chunk in chunks if chunk.get("usage") is not None] == usage_chunks, ( - response.text - ) + assert [ + (chunk["choices"] == [], *_token_usage(chunk)) for chunk in chunks if chunk.get("usage") is not None + ] == usage_chunks, response.text _assert_success_row( _spend_row(_BY_REQUEST_ID, identity), model, _HAIKU, "acompletion", _mantle_cost(gateway, _HAIKU) ) @@ -620,15 +625,14 @@ def test_openai_sdk_chat_stream_on_a_mantle_claude_id(gateway: Gateway) -> None: ) identities: Final = {chunk.id for chunk in chunks} assert len(identities) == 1 and {chunk.model for chunk in chunks} == {model}, chunks - assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks) == _answer(marker), chunks - assert [chunk.choices[0].finish_reason for chunk in chunks if chunk.choices[0].finish_reason] == ["stop"], ( - chunks - ) + choices: Final = tuple(chain.from_iterable(chunk.choices for chunk in chunks)) + assert "".join(choice.delta.content or "" for choice in choices) == _answer(marker), chunks + assert [choice.finish_reason for choice in choices if choice.finish_reason] == ["stop"], chunks assert [ - (chunk.usage.prompt_tokens, chunk.usage.completion_tokens, chunk.usage.total_tokens) + (len(chunk.choices), chunk.usage.prompt_tokens, chunk.usage.completion_tokens, chunk.usage.total_tokens) for chunk in chunks if chunk.usage is not None - ] == [_USAGE], chunks + ] == [(0, *_USAGE)], chunks _assert_success_row( _spend_row(_BY_REQUEST_ID, next(iter(identities))), model, @@ -688,15 +692,14 @@ async def test_async_openai_sdk_chat_stream_on_a_mantle_claude_id(gateway: Gatew chunks: Final = tuple([chunk async for chunk in stream]) identities: Final = {chunk.id for chunk in chunks} assert len(identities) == 1 and {chunk.model for chunk in chunks} == {model}, chunks - assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks) == _answer(marker), chunks - assert [chunk.choices[0].finish_reason for chunk in chunks if chunk.choices[0].finish_reason] == ["stop"], ( - chunks - ) + choices: Final = tuple(chain.from_iterable(chunk.choices for chunk in chunks)) + assert "".join(choice.delta.content or "" for choice in choices) == _answer(marker), chunks + assert [choice.finish_reason for choice in choices if choice.finish_reason] == ["stop"], chunks assert [ - (chunk.usage.prompt_tokens, chunk.usage.completion_tokens, chunk.usage.total_tokens) + (len(chunk.choices), chunk.usage.prompt_tokens, chunk.usage.completion_tokens, chunk.usage.total_tokens) for chunk in chunks if chunk.usage is not None - ] == [_USAGE], chunks + ] == [(0, *_USAGE)], chunks _assert_success_row( _spend_row(_BY_REQUEST_ID, next(iter(identities))), model, 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/integration/spend/test_spend_log_tool_payload_content.py b/tests/integration/spend/test_spend_log_tool_payload_content.py index 9055f6f06c4..226accdbf54 100644 --- a/tests/integration/spend/test_spend_log_tool_payload_content.py +++ b/tests/integration/spend/test_spend_log_tool_payload_content.py @@ -15,7 +15,7 @@ import openai import pytest from integration._support.client import Gateway, eventually, object_value, string_value from integration._support.database import read_rows -from integration._support.process import owned_proxy +from integration._support.process import graceful_stop_seconds, owned_proxy from integration._support.wire import Reply, Request, wire_server from pydantic import JsonValue, TypeAdapter @@ -723,6 +723,33 @@ def test_chat_sdk_stream_include_usage_masks_request_fields(gateway: Gateway, tm assert len(post_requests) == 1 +@pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +def test_chat_stream_usage_the_proxy_adds_keeps_its_marker_out_of_the_stored_request( + gateway: Gateway, tmp_path: Path +) -> None: + response_id: Final = f"chatcmpl-proxy-usage-{uuid4()}" + with ( + wire_server( + _answering_model_listing(lambda request: _chat_stream(response_id, "sort_key", include_usage=True)) + ) as wire, + owned_proxy(gateway, tmp_path, {}, config=_prompt_storage_config(tmp_path)) as isolated, + isolated.scenario() as scenario, + ): + model: Final = scenario.model(model="openai/gpt-4o-mini", api_base=wire.url, api_key="synthetic-openai-key") + response: Final = isolated.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": "stream marker"}], "stream": True}, + ) + assert response.status_code == 200, response.text + assert {string_value(event["id"]) for event in _sse_events(response.text)} == {response_id}, response.text + stored_request: Final = object_value(_stored_row(response_id)["proxy_server_request"]) + assert stored_request["messages"] == [{"role": "user", "content": "stream marker"}], stored_request + assert "_litellm_stream_usage_injected" not in stored_request, stored_request + observed: Final = _provider_calls(wire.drain()) + assert [_json_object(request.body)["stream_options"] for request in observed] == [{"include_usage": True}] + + def test_chat_history_keeps_string_tool_arguments_and_tool_content(gateway: Gateway, tmp_path: Path) -> None: arguments: Final = '{"sort_key":"created_at"}' tool_content: Final = "tool-result-created_at" diff --git a/tests/integration/streaming/test_stream_contracts.py b/tests/integration/streaming/test_stream_contracts.py index d7537a78db8..1d2a8b25792 100644 --- a/tests/integration/streaming/test_stream_contracts.py +++ b/tests/integration/streaming/test_stream_contracts.py @@ -3,6 +3,7 @@ import json import threading import uuid from collections.abc import Callable +from itertools import chain from pathlib import Path from queue import SimpleQueue from typing import Final @@ -13,9 +14,10 @@ from hypothesis import Phase, example, given, settings from hypothesis import strategies as st from integration._support.client import Gateway, eventually, object_value, string_value from integration._support.database import read_rows -from integration._support.process import owned_proxy +from integration._support.process import graceful_stop_seconds, owned_proxy from integration._support.wire import Reply, Request, wire_server -from openai import OpenAI +from openai import OpenAI, omit +from openai.types.chat import ChatCompletionChunk from pydantic import JsonValue, TypeAdapter pytestmark: Final = pytest.mark.shares_worker_state @@ -197,6 +199,7 @@ def test_proxy_stream_usage_visibility_keeps_exact_persisted_charge(gateway: Gat assert {chunk.id for chunk in chunks} == {identity} usages: Final = tuple(chunk.usage for chunk in chunks if chunk.usage is not None) assert len(usages) == (1 if include else 0) + assert [len(chunk.choices) for chunk in chunks if chunk.usage is not None] == ([0] if include else []) if include: assert usages[0].prompt_tokens == 11 and usages[0].completion_tokens == 4 requests: Final = wire.drain() @@ -214,6 +217,119 @@ def test_proxy_stream_usage_visibility_keeps_exact_persisted_charge(gateway: Gat assert float(rows[0]["spend"]) == pytest.approx(0.019) +_STREAM_USAGE_INJECTED_FLAG: Final = "_litellm_stream_usage_injected" + + +def _registered_model_names(gateway: Gateway) -> tuple[JsonValue, ...]: + entries: Final = gateway.get("/model/info")["data"] + assert isinstance(entries, list), entries + return tuple(object_value(entry).get("model_name") for entry in entries) + + +def _registered(gateway: Gateway, model: str) -> None: + eventually(lambda: _registered_model_names(gateway), lambda model_names: model in model_names, seconds=60) + + +def _sdk_stream( + gateway: Gateway, + route_prefix: str, + model: str, + identity: str, + *, + include_usage: bool | None, + extra_body: dict[str, JsonValue], +) -> tuple[ChatCompletionChunk, ...]: + with OpenAI( + api_key=gateway.key, + base_url=str(gateway.client.base_url).rstrip("/") + route_prefix, + timeout=30, + max_retries=0, + ) as client: + stream: Final = client.chat.completions.create( + model=model, + messages=[{"role": "user", "content": identity}], + stream=True, + stream_options=omit if include_usage is None else {"include_usage": include_usage}, + extra_body=extra_body, + ) + with stream: + return tuple(stream) + + +def _streamed_text(chunks: tuple[ChatCompletionChunk, ...]) -> str: + choices: Final = chain.from_iterable(chunk.choices for chunk in chunks) + return "".join(choice.delta.content or "" for choice in choices) + + +def _usage_chunk_shapes(chunks: tuple[ChatCompletionChunk, ...]) -> list[tuple[int, int]]: + return [(len(chunk.choices), chunk.usage.total_tokens) for chunk in chunks if chunk.usage is not None] + + +@pytest.mark.parametrize( + ("route_prefix", "sends", "extra_body"), + [pytest.param("/v1", 2, {}, id="chat_miss_then_cache_hit"), pytest.param("/queue", 1, {"priority": 0}, id="queue")], +) +def test_caller_requested_stream_usage_has_no_choices_even_with_a_client_sent_injection_flag( + gateway: Gateway, route_prefix: str, sends: int, extra_body: dict[str, JsonValue] +) -> None: + identity: Final = "stream-usage-flag-" + uuid.uuid4().hex + with ( + wire_server(lambda request: Reply(content_type="text/event-stream", chunks=text_stream(identity))) as wire, + gateway.scenario() as scenario, + ): + model: Final = scenario.model(api_base=wire.url + "/v1") + _registered(gateway, model) + streams: Final = tuple( + _sdk_stream( + gateway, + route_prefix, + model, + identity, + include_usage=True, + extra_body={**extra_body, _STREAM_USAGE_INJECTED_FLAG: True}, + ) + for _ in range(sends) + ) + assert [_streamed_text(chunks) for chunks in streams] == ["Hello 雪 café"] * sends + assert [_usage_chunk_shapes(chunks) for chunks in streams] == [[(0, 15)]] * sends + requests: Final = wire.drain() + assert len(requests) == 1 + assert _STREAM_USAGE_INJECTED_FLAG not in json.loads(requests[0].body) + + +@pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +def test_stream_usage_the_deployment_adds_keeps_one_empty_choice(gateway: Gateway, tmp_path: Path) -> None: + identity: Final = "deployment-stream-usage-" + uuid.uuid4().hex + with wire_server(lambda request: Reply(content_type="text/event-stream", chunks=text_stream(identity))) as wire: + shared_config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config: Final = { + **shared_config, + "model_list": [ + { + "model_name": "deployment-usage", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "synthetic-deployment-usage-key", + "api_base": wire.url + "/v1", + "stream_options": {"include_usage": True}, + }, + } + ], + "general_settings": {**shared_config["general_settings"], "always_include_stream_usage": False}, + } + path: Final = tmp_path / "deployment-stream-usage.yaml" + path.write_text(yaml.safe_dump(config)) + with owned_proxy(gateway, tmp_path, {}, config=path) as candidate: + chunks: Final = _sdk_stream( + candidate, "/v1", "deployment-usage", identity, include_usage=None, extra_body={} + ) + assert _streamed_text(chunks) == "Hello 雪 café" + assert _usage_chunk_shapes(chunks) == [(1, 15)] + requests: Final = wire.drain() + assert len(requests) == 1 + assert json.loads(requests[0].body)["stream_options"]["include_usage"] is True + + @pytest.mark.covers("other.streaming.messages_bridge.empty_choices_usage_chunk_completes_stream") def test_messages_stream_completes_through_trailing_empty_choices_usage_chunk(gateway: Gateway) -> None: identity: Final = "messages-empty-choices-" + uuid.uuid4().hex diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index c96e3e2c031..e64f5cca5b8 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -252,6 +252,7 @@ def isolate_host_environment(isolated_aws_config_files: tuple[Path, Path]) -> It @pytest.fixture(autouse=True) def isolate_litellm_globals() -> Iterator[None]: original_callbacks: Final = {name: list(getattr(litellm, name) or []) for name in CALLBACK_LISTS} + original_callback_holders: Final = frozenset(litellm_router_module._routers_holding_global_callbacks) original_reset: Final = {name: getattr(litellm, name) for name in RESET_TO_NONE_GLOBALS} original_restored: Final = {name: getattr(litellm, name) for name in RESTORED_GLOBALS if hasattr(litellm, name)} original_clients: Final = {name: litellm.__dict__[name] for name in MODULE_LEVEL_CLIENTS if name in litellm.__dict__} @@ -263,6 +264,7 @@ def isolate_litellm_globals() -> Iterator[None]: _flush_client_caches() for name in CALLBACK_LISTS: setattr(litellm, name, []) + litellm_router_module._routers_holding_global_callbacks.clear() for name in RESET_TO_NONE_GLOBALS: setattr(litellm, name, None) for name in MODULE_LEVEL_CLIENTS: @@ -277,6 +279,8 @@ def isolate_litellm_globals() -> Iterator[None]: litellm.__dict__.update(original_clients) for name, value in (original_callbacks | original_reset | original_restored).items(): setattr(litellm, name, value) + litellm_router_module._routers_holding_global_callbacks.clear() + litellm_router_module._routers_holding_global_callbacks.update(original_callback_holders) for logger, (level, disabled, propagate, handlers, filters) in original_loggers.items(): logger.setLevel(level) logger.disabled = disabled 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/fine_tuning/test_main.py b/tests/unit/fine_tuning/test_main.py index 57c845c6e1a..89f51d2e74a 100644 --- a/tests/unit/fine_tuning/test_main.py +++ b/tests/unit/fine_tuning/test_main.py @@ -1,11 +1,12 @@ import json -from typing import Final +from typing import Final, Literal import httpx import pytest import respx import litellm +from litellm.types.utils import LiteLLMFineTuningJob API_BASE: Final = "https://fine-tuning.test" BASE_MODEL: Final = "gpt-4.1-nano-2025-04-14" @@ -52,3 +53,42 @@ async def test_create_sends_the_model_without_the_provider_prefix( ) assert json.loads(sent.calls.last.request.content)["model"] == BASE_MODEL + + +@pytest.mark.asyncio +@respx.mock +@pytest.mark.parametrize( + ("custom_llm_provider", "api_base", "route", "sent_api_version"), + [ + ("openai", f"{API_BASE}/v1", f"{API_BASE}/v1/fine_tuning/jobs", None), + ("azure", API_BASE, f"{API_BASE}/openai/fine_tuning/jobs", "2024-10-21"), + ], +) +@pytest.mark.parametrize("is_async", [True, False], ids=["async", "sync"]) +async def test_create_with_api_version_reaches_the_provider_and_only_azure_gets_it( + monkeypatch: pytest.MonkeyPatch, + custom_llm_provider: Literal["openai", "azure"], + api_base: str, + route: str, + sent_api_version: str | None, + is_async: bool, +) -> None: + sent: Final = respx.post(url__startswith=route).mock(return_value=httpx.Response(200, json=PROVIDER_JOB)) + request: Final = { + "model": BASE_MODEL, + "training_file": "file-abc123", + "custom_llm_provider": custom_llm_provider, + "api_key": "sk-test", + "api_base": api_base, + "api_version": "2024-10-21", + } + + async with httpx.AsyncClient() as session: + monkeypatch.setattr(litellm, "aclient_session", session) + job: Final = ( + await litellm.acreate_fine_tuning_job(**request) if is_async else litellm.create_fine_tuning_job(**request) + ) + + assert isinstance(job, LiteLLMFineTuningJob) + assert job.id == PROVIDER_JOB["id"] + assert sent.calls.last.request.url.params.get("api-version") == sent_api_version diff --git a/tests/unit/harness/test_endpoint.py b/tests/unit/harness/test_endpoint.py index 19bd6d0450e..d1447fd40cc 100644 --- a/tests/unit/harness/test_endpoint.py +++ b/tests/unit/harness/test_endpoint.py @@ -272,6 +272,7 @@ async def test_sdk_chat_stream(monkeypatch: pytest.MonkeyPatch) -> None: assert lines[-1] == "data: [DONE]" assert json.loads(lines[0][6:])["choices"][0]["delta"]["content"] == "he" assert calls[0]["stream_options"] == {"include_usage": True} + assert calls[0]["_litellm_stream_usage_injected"] is True assert (ep.usage.input_tokens, ep.usage.output_tokens) == (6, 2) assert ep.usage.cost == pytest.approx(6 * 0.001 + 2 * 0.002) 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 diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index 79645ffe8f6..bbfae2229ce 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -42,6 +42,7 @@ from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils import litellm_logging from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging from litellm.litellm_core_utils.litellm_logging import ( + BilledRoundsSpend, Logging, StandardLoggingPayloadSetup, _extract_response_obj_and_hidden_params, @@ -402,6 +403,32 @@ def test_get_combined_callback_list_preserves_insertion_order(logging_obj): ) == ["prometheus", "langfuse", "datadog", "otel", "s3", "gcs_bucket", "arize", "logfire"] +class _InputEventRecorder(CustomLogger): + def __init__(self, label: str, events: list[str]) -> None: + self._label = label + self.events = events + + def log_pre_api_call(self, model: object, messages: object, kwargs: object) -> None: + self.events.append(f"{self._label} pre") + + def log_post_api_call(self, kwargs: object, response_obj: object, start_time: object, end_time: object) -> None: + self.events.append(f"{self._label} post") + + +def test_input_events_reach_the_replacement_and_not_the_global_callback_it_replaced( + logging_obj: LitellmLogging, monkeypatch: pytest.MonkeyPatch +) -> None: + events: Final[list[str]] = [] + monkeypatch.setattr(litellm, "input_callback", [_InputEventRecorder("global", events)]) + logging_obj.replace_global_callback(_InputEventRecorder("replacement", events)) + logging_obj.update_environment_variables(litellm_params={}, optional_params={}) + + logging_obj.pre_call(input="", api_key="") + logging_obj.post_call(original_response="ok", input="", api_key="") + + assert events == ["replacement pre", "replacement post"] + + def test_get_masked_api_base(logging_obj): api_base = "https://api.openai.com/v1" masked_api_base = logging_obj._get_masked_api_base(api_base) @@ -1822,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]) @@ -6811,6 +6844,26 @@ def test_failure_handler_zeroes_spend_without_recovered_usage(logging_obj): assert payload["total_tokens"] == 0 +def test_failure_after_earlier_rounds_bills_only_the_failed_round_when_it_priced_no_breakdown( + logging_obj: LitellmLogging, +) -> None: + earlier_breakdown: Final = CostBreakdown(total_cost=0.3, margin_fixed_amount=0.1) + logging_obj.cost_breakdown = earlier_breakdown + logging_obj.model_call_details["litellm_params"] = {} + + logging_obj.record_failure_after_earlier_rounds( + BilledRoundsSpend(Usage(prompt_tokens=5, completion_tokens=2, total_tokens=7), 0.3, earlier_breakdown), + Usage(prompt_tokens=7, completion_tokens=0, total_tokens=7), + 0.0, + ) + logging_obj.failure_handler(litellm.ServiceUnavailableError("503", "anthropic", "claude"), "tb") + + row: Final = logging_obj.model_call_details["standard_logging_object"] + assert (row["prompt_tokens"], row["completion_tokens"]) == (7, 0) + assert row["response_cost"] == pytest.approx(0.0) + assert row.get("cost_breakdown") is None or row["cost_breakdown"]["total_cost"] == pytest.approx(0.0) + + def test_set_cost_breakdown_stores_reasoning_cost(): """reasoning_cost is stored only when positive, mirroring the cache-cost fields.""" from datetime import datetime diff --git a/tests/unit/litellm_core_utils/test_logging_callback_manager.py b/tests/unit/litellm_core_utils/test_logging_callback_manager.py index f8260702068..25a52137d65 100644 --- a/tests/unit/litellm_core_utils/test_logging_callback_manager.py +++ b/tests/unit/litellm_core_utils/test_logging_callback_manager.py @@ -145,6 +145,7 @@ def test_remove_callback_from_all_lists(): obj = TestLogger() manager.add_litellm_callback(obj) + manager.add_litellm_input_callback(obj) manager.add_litellm_success_callback(obj) manager.add_litellm_failure_callback(obj) manager.add_litellm_async_success_callback(obj) @@ -153,6 +154,7 @@ def test_remove_callback_from_all_lists(): manager.remove_callback_from_all_lists(obj) assert obj not in litellm.callbacks + assert obj not in litellm.input_callback assert obj not in litellm.success_callback assert obj not in litellm.failure_callback assert obj not in litellm._async_success_callback 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_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/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index 68aabc3c01e..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 @@ -603,6 +603,91 @@ async def test_streaming_handler_with_usage( assert chunk_has_usage +def _usage_test_logging( + caller_stream_options: dict[str, bool] | None, stream_usage_injected: bool = False +) -> Logging: + logging_obj: Final = Logging( + model="hosted_vllm/m", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="completion", + start_time=1_790_000_000.0, + litellm_call_id="12345", + function_id="1245", + kwargs={"_litellm_stream_usage_injected": True} if stream_usage_injected else {}, + ) + logging_obj.update_environment_variables( + litellm_params={}, optional_params={}, model="hosted_vllm/m", stream_options=caller_stream_options + ) + return logging_obj + + +@pytest.mark.parametrize( + ("caller_stream_options", "stream_usage_injected", "usage_chunk_choice_count"), + [ + ({"include_usage": True}, False, 0), + (None, False, 1), + ({"include_usage": True}, True, 1), + ], + ids=["caller-asked", "provider-handler-injected", "proxy-injected"], +) +@pytest.mark.parametrize("plain_iterator", [False, True]) +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +async def test_usage_chunk_has_empty_choices_only_for_callers_that_asked( + sync_mode: bool, + plain_iterator: bool, + caller_stream_options: dict[str, bool] | None, + stream_usage_injected: bool, + usage_chunk_choice_count: int, +): + usage: Final = Usage(prompt_tokens=5, completion_tokens=1, total_tokens=6) + provider_chunks: Final = [ + ModelResponseStream(id="c1", choices=[StreamingChoices(index=0, delta=Delta(role="assistant", content="ok"))]), + ModelResponseStream(id="c1", choices=[StreamingChoices(index=0, delta=Delta(), finish_reason="stop")]), + ModelResponseStream(id="c1", choices=[], provider_specific_fields={"citations": ["a"]}), + ModelResponseStream(id="c1", choices=[], usage=usage), + ] + wrapper: Final = CustomStreamWrapper( + completion_stream=( + iter(provider_chunks) if plain_iterator else ModelResponseListIterator(model_responses=provider_chunks) + ), + model="hosted_vllm/m", + custom_llm_provider="hosted_vllm", + logging_obj=_usage_test_logging(caller_stream_options, stream_usage_injected), + stream_options={"include_usage": True}, + ) + + chunks: Final = list(wrapper) if sync_mode else [chunk async for chunk in wrapper] + + assert [len(chunk.choices) for chunk in chunks] == [1, 1, 0, usage_chunk_choice_count] + assert chunks[2].provider_specific_fields == {"citations": ["a"]} + assert [getattr(chunk, "usage", None) is not None for chunk in chunks] == [False, False, False, True] + assert chunks[-1].usage.total_tokens == 6 + + +@pytest.mark.asyncio +async def test_plain_iterator_stream_keeps_usage_on_its_chunk_when_usage_was_not_requested(): + usage: Final = Usage(prompt_tokens=5, completion_tokens=1, total_tokens=6) + provider_chunks: Final = [ + ModelResponseStream(id="c1", choices=[StreamingChoices(index=0, delta=Delta(role="assistant", content="ok"))]), + ModelResponseStream( + id="c1", choices=[StreamingChoices(index=0, delta=Delta(), finish_reason="stop")], usage=usage + ), + ] + wrapper: Final = CustomStreamWrapper( + completion_stream=iter(provider_chunks), + model="hosted_vllm/m", + custom_llm_provider="hosted_vllm", + logging_obj=_usage_test_logging(None), + ) + + chunks: Final = [chunk async for chunk in wrapper] + + assert [chunk.choices[0].finish_reason for chunk in chunks] == [None, "stop"] + assert chunks[-1].usage.total_tokens == 6 + + @pytest.mark.parametrize("sync_mode", [False]) @pytest.mark.asyncio @pytest.mark.flaky(reruns=3) @@ -4180,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/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/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/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/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/openai/completion/test_completion_handler.py b/tests/unit/llms/openai/completion/test_completion_handler.py index f465bbea663..9e2754eef68 100644 --- a/tests/unit/llms/openai/completion/test_completion_handler.py +++ b/tests/unit/llms/openai/completion/test_completion_handler.py @@ -364,7 +364,8 @@ def test_text_completion_stream_forwards_stream_options_and_usage(monkeypatch): ) ) - assert tuple(chunk.choices[0].text for chunk in chunks) == ("hello", None, None) + assert tuple(chunk.choices[0].text for chunk in chunks[:-1]) == ("hello", None) + assert chunks[-1].choices == [] assert tuple(chunk.usage for chunk in chunks[:-1]) == (None, None) assert (chunks[-1].usage.prompt_tokens, chunks[-1].usage.completion_tokens, chunks[-1].usage.total_tokens) == ( 2, 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/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/llms/sagemaker/test_sagemaker_nova_transformation.py b/tests/unit/llms/sagemaker/test_sagemaker_nova_transformation.py index 12690d5822b..56c0aee3efd 100644 --- a/tests/unit/llms/sagemaker/test_sagemaker_nova_transformation.py +++ b/tests/unit/llms/sagemaker/test_sagemaker_nova_transformation.py @@ -543,8 +543,9 @@ def test_nova_streaming_sends_stream_flag_and_parses_chunks_and_usage(respx_mock "stream": True, "stream_options": {"include_usage": True}, } - assert tuple(chunk.choices[0].delta.content for chunk in chunks) == ("Nova ", "stream", None, None) - assert tuple(chunk.choices[0].finish_reason for chunk in chunks) == (None, None, "stop", None) + assert tuple(chunk.choices[0].delta.content for chunk in chunks[:-1]) == ("Nova ", "stream", None) + assert chunks[-1].choices == [] + assert tuple(chunk.choices[0].finish_reason for chunk in chunks[:-1]) == (None, None, "stop") assert (chunks[-1].usage.prompt_tokens, chunks[-1].usage.completion_tokens, chunks[-1].usage.total_tokens) == ( 2, 1, 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/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_long_decode.py b/tests/unit/llms/vertex_ai/gemini/test_long_decode.py index 4465ad7f3df..3fe9a5c7a92 100644 --- a/tests/unit/llms/vertex_ai/gemini/test_long_decode.py +++ b/tests/unit/llms/vertex_ai/gemini/test_long_decode.py @@ -516,9 +516,16 @@ async def test_unary_answer_cut_at_the_callers_cap_is_returned_with_its_token_in assert response.choices[0].message.provider_specific_fields["continuation_token"] == "dG9rLTE=" +@pytest.mark.parametrize( + "choice_count_request", + [{"max_tokens": 100_000, "n": 2}, {"extra_body": {"generationConfig": {"candidateCount": 2}}}], +) @pytest.mark.parametrize("is_async", [True, False]) @pytest.mark.asyncio -async def test_unary_request_for_more_than_one_choice_is_never_continued(is_async: bool): +async def test_unary_request_for_more_than_one_choice_is_never_continued( + monkeypatch: pytest.MonkeyPatch, is_async: bool, choice_count_request: Mapping[str, object] +): + monkeypatch.setenv(BUDGET_ENV, "100000") sent: Final[list[dict[str, object]]] = [] # Vertex AI answered 400 "Candidate count greater than 1 is not supported when continuation_token is set" # on 2026-10-13 @@ -526,9 +533,8 @@ async def test_unary_request_for_more_than_one_choice_is_never_continued(is_asyn kwargs: Final = { "model": MODEL, "messages": list(MESSAGES), - "max_tokens": 100_000, - "n": 2, "api_key": "fake-gemini-key", + **choice_count_request, } response: Final = ( @@ -548,6 +554,42 @@ async def test_unary_request_for_more_than_one_choice_is_never_continued(is_asyn ] +@pytest.mark.parametrize( + "one_choice_request", + [ + {"max_tokens": 100_000, "n": 0}, + {"extra_body": {"generationConfig": {"candidateCount": None}}}, + {"extra_body": {"generationConfig": {"candidateCount": 0}}}, + ], +) +@pytest.mark.parametrize("is_async", [True, False]) +@pytest.mark.asyncio +async def test_unary_request_whose_candidate_count_google_reads_as_one_is_continued( + monkeypatch: pytest.MonkeyPatch, is_async: bool, one_choice_request: Mapping[str, object] +): + monkeypatch.setenv(BUDGET_ENV, "100000") + sent: Final[list[dict[str, object]]] = [] + hops: Final = (httpx.Response(200, json=HOP_1), httpx.Response(200, json=HOP_2)) + kwargs: Final = { + "model": MODEL, + "messages": list(MESSAGES), + "api_key": "fake-gemini-key", + **one_choice_request, + } + + response: Final = ( + await litellm.acompletion(**kwargs, client=AsyncHTTPHandler(transport=httpx.MockTransport(_replay(hops, sent)))) + if is_async + else litellm.completion( + **kwargs, client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(_replay(hops, sent)))) + ) + ) + + assert sent[1] == {**sent[0], "continuationToken": "dG9rLTE="} + assert response.choices[0].message.content == "hello world" + assert response.choices[0].finish_reason == "stop" + + @pytest.mark.asyncio async def test_unary_failed_continuation_hop_returns_the_answer_so_far_with_its_token(): sent: Final[list[dict[str, object]]] = [] @@ -1295,6 +1337,468 @@ async def test_vertex_long_decode_keeps_the_callers_own_authorization_header_on_ assert seen == ["Bearer mine", "Bearer mine"] +BUDGET_ENV: Final = "LITELLM_GEMINI_DEFAULT_MAX_OUTPUT_TOKENS" +ARGON: Final = "gemini/gemini-4-argon" +UNMAPPED_GEMINI: Final = "gemini/gemini-99-unreleased" +OUTPUT_LIMIT_KEYS: Final = ("max_output_tokens", "maxOutputTokens") + + +def _output_limits(request: Mapping[str, object]) -> dict[str, object]: + generation_config: Final = request.get("generationConfig", {}) + assert isinstance(generation_config, dict) + return {key: value for key, value in generation_config.items() if key in OUTPUT_LIMIT_KEYS} + + +def _cost_map_max_output_tokens(model: str) -> int: + provider, name = model.split("/", 1) + max_output_tokens: Final = litellm.get_model_info(model=name, custom_llm_provider=provider)["max_output_tokens"] + assert isinstance(max_output_tokens, int) + return max_output_tokens + + +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("is_async", [True, False]) +@pytest.mark.asyncio +async def test_request_with_no_output_limit_keeps_continuing_until_the_env_budget_is_spent( + monkeypatch: pytest.MonkeyPatch, is_async: bool, stream: bool +): + monkeypatch.setenv(BUDGET_ENV, "50") + sent: Final[list[dict[str, object]]] = [] + hops: Final = ( + (_sse(STREAM_FIRST_CHUNK, STREAM_HOP_1_END), _sse(HOP_2_TRUNCATED), _sse(HOP_2)) + if stream + else ( + httpx.Response(200, json=HOP_1), + httpx.Response(200, json=HOP_2_TRUNCATED), + httpx.Response(200, json=HOP_2), + ) + ) + kwargs: Final = {"model": ARGON, "messages": list(MESSAGES), "stream": stream, "api_key": "fake-gemini-key"} + transport: Final = httpx.MockTransport(_replay(hops, sent)) + + response: Final = ( + await litellm.acompletion(**kwargs, client=AsyncHTTPHandler(transport=transport)) + if is_async + else litellm.completion(**kwargs, client=HTTPHandler(client=httpx.Client(transport=transport))) + ) + chunks: Final = ([chunk async for chunk in response] if is_async else list(response)) if stream else [] + + assert len(sent) == 2 + assert _output_limits(sent[0]) == {"max_output_tokens": 50} + if stream: + assert _stream_text(chunks) == "hello world" + assert _stream_finish_reasons(chunks) == ["length"] + assert {"continuation_token": "dG9rLTI="} in _delta_fields(chunks) + else: + assert response.choices[0].message.content == "hello world" + assert response.choices[0].finish_reason == "length" + assert response.choices[0].message.provider_specific_fields["continuation_token"] == "dG9rLTI=" + + +@pytest.mark.parametrize( + "budget, model, request_params, expected_limits", + [ + ("50", ARGON, {"max_tokens": 30}, {"max_output_tokens": 30}), + ("50", ARGON, {"max_completion_tokens": 30}, {"max_output_tokens": 30}), + ( + "50", + ARGON, + {"extra_body": {"generationConfig": {"maxOutputTokens": 30}}}, + {"maxOutputTokens": 30}, + ), + (None, ARGON, {}, {}), + ("", ARGON, {}, {}), + ("a lot", ARGON, {}, {}), + ("0", ARGON, {}, {}), + ("-5", ARGON, {}, {}), + ("50", UNMAPPED_GEMINI, {}, {}), + ], +) +@pytest.mark.asyncio +async def test_env_budget_applies_only_to_a_request_that_sets_no_output_limit_of_its_own( + monkeypatch: pytest.MonkeyPatch, + budget: str | None, + model: str, + request_params: Mapping[str, object], + expected_limits: Mapping[str, object], +): + if budget is None: + monkeypatch.delenv(BUDGET_ENV, raising=False) + else: + monkeypatch.setenv(BUDGET_ENV, budget) + sent: Final[list[dict[str, object]]] = [] + + await litellm.acompletion( + model=model, + messages=list(MESSAGES), + api_key="fake-gemini-key", + client=AsyncHTTPHandler(transport=httpx.MockTransport(_replay((httpx.Response(200, json=HOP_2),), sent))), + **request_params, + ) + + assert _output_limits(sent[0]) == expected_limits + + +@pytest.mark.parametrize( + "modalities, expected_limits", + [ + (["audio"], {}), + (["text", "audio"], {}), + (["text", "image"], {}), + (["text"], {"max_output_tokens": 50}), + ], +) +@pytest.mark.asyncio +async def test_env_budget_applies_only_to_a_request_asking_for_text_output( + monkeypatch: pytest.MonkeyPatch, modalities: list[str], expected_limits: Mapping[str, object] +): + monkeypatch.setenv(BUDGET_ENV, "50") + sent: Final[list[dict[str, object]]] = [] + + await litellm.acompletion( + model=ARGON, + messages=list(MESSAGES), + modalities=modalities, + api_key="fake-gemini-key", + client=AsyncHTTPHandler(transport=httpx.MockTransport(_replay((httpx.Response(200, json=HOP_2),), sent))), + ) + + assert _output_limits(sent[0]) == expected_limits + + +def _ai_studio_model_in_mode(mode: str) -> str: + return next( + name + for name, row in litellm.model_cost.items() + if name.startswith("gemini/") and row.get("mode") == mode and row.get("max_output_tokens") + ) + + +@pytest.mark.parametrize("mode", ["audio_speech", "image_generation"]) +@pytest.mark.asyncio +async def test_env_budget_skips_a_model_the_cost_map_does_not_list_as_a_chat_model( + monkeypatch: pytest.MonkeyPatch, mode: str +): + monkeypatch.setenv(BUDGET_ENV, "50") + sent: Final[list[dict[str, object]]] = [] + + await litellm.acompletion( + model=_ai_studio_model_in_mode(mode), + messages=list(MESSAGES), + api_key="fake-gemini-key", + client=AsyncHTTPHandler(transport=httpx.MockTransport(_replay((httpx.Response(200, json=HOP_2),), sent))), + ) + + assert _output_limits(sent[0]) == {} + + +@pytest.mark.parametrize( + "output_modalities, expected_limits", + [ + ({}, {"max_output_tokens": 50}), + ({"supported_output_modalities": ["text"]}, {"max_output_tokens": 50}), + ({"supported_output_modalities": ["text", "audio"]}, {"max_output_tokens": 50}), + ({"supported_output_modalities": ["audio"]}, {}), + ({"supported_output_modalities": "audio"}, {}), + ], +) +@pytest.mark.asyncio +async def test_env_budget_applies_only_to_a_chat_model_the_cost_map_lists_as_answering_text( + monkeypatch: pytest.MonkeyPatch, output_modalities: Mapping[str, object], expected_limits: Mapping[str, object] +): + monkeypatch.setenv(BUDGET_ENV, "50") + litellm.register_model( + { + UNMAPPED_GEMINI: { + "litellm_provider": "gemini", + "mode": "chat", + "max_output_tokens": 1000, + **output_modalities, + } + } + ) + sent: Final[list[dict[str, object]]] = [] + + await litellm.acompletion( + model=UNMAPPED_GEMINI, + messages=list(MESSAGES), + api_key="fake-gemini-key", + client=AsyncHTTPHandler(transport=httpx.MockTransport(_replay((httpx.Response(200, json=HOP_2),), sent))), + ) + + assert _output_limits(sent[0]) == expected_limits + + +@pytest.mark.parametrize( + "extra_body", + [ + {"generation_config": {"temperature": 0.1}}, + {"generation_config": {"max_output_tokens": 30}}, + {"generationConfig": {"response_modalities": ["AUDIO"]}}, + ], +) +@pytest.mark.asyncio +async def test_env_budget_leaves_a_snake_case_generation_config_from_extra_body_alone( + monkeypatch: pytest.MonkeyPatch, extra_body: Mapping[str, object] +): + monkeypatch.setenv(BUDGET_ENV, "50") + sent: Final[list[dict[str, object]]] = [] + + await litellm.acompletion( + model=ARGON, + messages=list(MESSAGES), + extra_body=dict(extra_body), + api_key="fake-gemini-key", + client=AsyncHTTPHandler(transport=httpx.MockTransport(_replay((httpx.Response(200, json=HOP_2),), sent))), + ) + + assert _output_limits(sent[0]) == {} + + +def _gemini_chat_models_answering_no_text() -> tuple[str, ...]: + rows: Final = ( + (name, row) + for name, row in litellm.model_cost.items() + if row.get("mode") == "chat" and "text" not in row.get("supported_output_modalities", ("text",)) + ) + return tuple( + name if name.startswith("gemini/") else f"vertex_ai/{name}" + for name, row in rows + if row.get("litellm_provider") in ("gemini", "vertex_ai-language-models") + ) + + +@pytest.mark.asyncio +async def test_env_budget_skips_every_chat_model_the_cost_map_lists_as_answering_no_text( + monkeypatch: pytest.MonkeyPatch, token_endpoint: _TokenEndpoint +): + monkeypatch.setenv(BUDGET_ENV, "50") + models: Final = _gemini_chat_models_answering_no_text() + sent: Final[list[dict[str, object]]] = [] + replies: Final = tuple(httpx.Response(200, json=HOP_2) for _ in models) + client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(_replay(replies, sent))) + + for model in models: + await litellm.acompletion( + model=model, + messages=list(MESSAGES), + api_key="fake-gemini-key", + vertex_project="p", + vertex_location="us-central1", + vertex_credentials=token_endpoint.credentials, + client=client, + ) + + assert models + assert {model: _output_limits(request) for model, request in zip(models, sent, strict=True)} == dict.fromkeys( + models, {} + ) + + +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.asyncio +async def test_env_budget_leaves_a_continuation_token_resume_with_no_limit_alone( + monkeypatch: pytest.MonkeyPatch, stream: bool +): + monkeypatch.setenv(BUDGET_ENV, "50") + sent: Final[list[dict[str, object]]] = [] + reply: Final = _sse(HOP_2) if stream else httpx.Response(200, json=HOP_2) + + response: Final = await litellm.acompletion( + model=ARGON, + messages=list(MESSAGES), + continuation_token="dG9rLTE=", + stream=stream, + api_key="fake-gemini-key", + client=AsyncHTTPHandler(transport=httpx.MockTransport(_replay((reply,), sent))), + ) + if stream: + _ = [chunk async for chunk in response] + + assert sent[0]["continuationToken"] == "dG9rLTE=" + assert _output_limits(sent[0]) == {} + + +@pytest.mark.parametrize("model", ["gemini/gemini-2.5-flash", "vertex_ai/gemini-2.5-flash", ARGON]) +@pytest.mark.asyncio +async def test_env_budget_above_a_models_output_limit_sends_that_limit_instead( + monkeypatch: pytest.MonkeyPatch, token_endpoint: _TokenEndpoint, model: str +): + budget: Final = 2 * _cost_map_max_output_tokens(model) + monkeypatch.setenv(BUDGET_ENV, str(budget)) + sent: Final[list[dict[str, object]]] = [] + + await litellm.acompletion( + model=model, + messages=list(MESSAGES), + api_key="fake-gemini-key", + vertex_project="p", + vertex_location="us-central1", + vertex_credentials=token_endpoint.credentials, + client=AsyncHTTPHandler(transport=httpx.MockTransport(_replay((httpx.Response(200, json=HOP_2),), sent))), + ) + + assert _output_limits(sent[0]) == {"max_output_tokens": _cost_map_max_output_tokens(model)} + + +@pytest.mark.asyncio +async def test_env_budget_never_goes_above_the_cost_map_limit_when_a_deployment_overrides_it( + monkeypatch: pytest.MonkeyPatch, +): + model: Final = "gemini/gemini-2.5-flash" + model_limit: Final = _cost_map_max_output_tokens(model) + monkeypatch.setenv(BUDGET_ENV, str(2 * model_limit)) + router: Final = litellm.Router( + model_list=[ + {"model_name": "plain", "litellm_params": {"model": model, "api_key": "fake-gemini-key"}}, + { + "model_name": "raised", + "litellm_params": {"model": model, "api_key": "fake-gemini-key"}, + "model_info": {"max_output_tokens": 4 * model_limit}, + }, + ] + ) + sent: Final[list[dict[str, object]]] = [] + replies: Final = (httpx.Response(200, json=HOP_2), httpx.Response(200, json=HOP_2)) + client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(_replay(replies, sent))) + + for model_name in ("plain", "raised"): + await router.acompletion(model=model_name, messages=list(MESSAGES), client=client) + + sent_limits: Final = [_output_limits(request).get("max_output_tokens", 0) for request in sent] + assert len(sent_limits) == 2 + assert all(isinstance(limit, int) and limit <= model_limit for limit in sent_limits) + + +@pytest.mark.asyncio +async def test_env_budget_caps_each_deployment_at_its_own_model_info_limit_and_never_at_a_siblings( + monkeypatch: pytest.MonkeyPatch, +): + model: Final = "gemini/gemini-2.5-flash" + budget: Final = _cost_map_max_output_tokens(model) // 2 + low_limit: Final = budget // 4 + monkeypatch.setenv(BUDGET_ENV, str(budget)) + router: Final = litellm.Router( + model_list=[ + { + "model_name": "low", + "litellm_params": {"model": model, "api_key": "fake-gemini-key"}, + "model_info": {"max_output_tokens": low_limit}, + }, + {"model_name": "plain", "litellm_params": {"model": model, "api_key": "fake-gemini-key"}}, + ] + ) + sent: Final[list[dict[str, object]]] = [] + replies: Final = (httpx.Response(200, json=HOP_2), httpx.Response(200, json=HOP_2)) + client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(_replay(replies, sent))) + + for model_name in ("plain", "low"): + await router.acompletion(model=model_name, messages=list(MESSAGES), client=client) + + assert [_output_limits(request) for request in sent] == [ + {"max_output_tokens": budget}, + {"max_output_tokens": low_limit}, + ] + + +@pytest.mark.parametrize("deployment_limit", [0, -1]) +@pytest.mark.asyncio +async def test_env_budget_ignores_a_deployment_model_info_limit_that_is_not_positive( + monkeypatch: pytest.MonkeyPatch, deployment_limit: int +): + model: Final = "gemini/gemini-2.5-flash" + budget: Final = _cost_map_max_output_tokens(model) // 2 + monkeypatch.setenv(BUDGET_ENV, str(budget)) + router: Final = litellm.Router( + model_list=[ + { + "model_name": "odd", + "litellm_params": {"model": model, "api_key": "fake-gemini-key"}, + "model_info": {"max_output_tokens": deployment_limit}, + } + ] + ) + sent: Final[list[dict[str, object]]] = [] + + await router.acompletion( + model="odd", + messages=list(MESSAGES), + client=AsyncHTTPHandler(transport=httpx.MockTransport(_replay((httpx.Response(200, json=HOP_2),), sent))), + ) + + assert _output_limits(sent[0]) == {"max_output_tokens": budget} + + +@pytest.mark.asyncio +async def test_env_budget_on_a_direct_call_never_goes_above_the_cost_map_limit_after_a_router_raised_it( + monkeypatch: pytest.MonkeyPatch, +): + model: Final = "gemini/gemini-2.5-flash" + model_limit: Final = _cost_map_max_output_tokens(model) + monkeypatch.setenv(BUDGET_ENV, str(2 * model_limit)) + litellm.Router( + model_list=[ + { + "model_name": "raised", + "litellm_params": {"model": model, "api_key": "fake-gemini-key"}, + "model_info": {"max_output_tokens": 4 * model_limit}, + } + ] + ) + sent: Final[list[dict[str, object]]] = [] + + await litellm.acompletion( + model=model, + messages=list(MESSAGES), + api_key="fake-gemini-key", + client=AsyncHTTPHandler(transport=httpx.MockTransport(_replay((httpx.Response(200, json=HOP_2),), sent))), + ) + + assert _output_limits(sent[0]) == {"max_output_tokens": model_limit} + + +@pytest.mark.asyncio +async def test_env_budget_on_a_direct_call_follows_a_lower_limit_set_with_register_model( + monkeypatch: pytest.MonkeyPatch, +): + model: Final = "gemini/gemini-2.5-flash" + lowered_limit: Final = _cost_map_max_output_tokens(model) // 4 + monkeypatch.setenv(BUDGET_ENV, str(2 * lowered_limit)) + litellm.register_model({model: {"litellm_provider": "gemini", "max_output_tokens": lowered_limit}}) + sent: Final[list[dict[str, object]]] = [] + + await litellm.acompletion( + model=model, + messages=list(MESSAGES), + api_key="fake-gemini-key", + client=AsyncHTTPHandler(transport=httpx.MockTransport(_replay((httpx.Response(200, json=HOP_2),), sent))), + ) + + assert _output_limits(sent[0]) == {"max_output_tokens": lowered_limit} + + +@pytest.mark.parametrize( + "generation_config", + [None, {"responseModalities": None}, {"responseModalities": {"TEXT": 1}}], +) +@pytest.mark.asyncio +async def test_env_budget_sends_a_generation_config_of_an_unexpected_shape_as_it_came( + monkeypatch: pytest.MonkeyPatch, generation_config: object +): + monkeypatch.setenv(BUDGET_ENV, "50") + sent: Final[list[dict[str, object]]] = [] + + await litellm.acompletion( + model=ARGON, + messages=list(MESSAGES), + extra_body={"generationConfig": generation_config}, + api_key="fake-gemini-key", + client=AsyncHTTPHandler(transport=httpx.MockTransport(_replay((httpx.Response(200, json=HOP_2),), sent))), + ) + + assert sent[0]["generationConfig"] == generation_config + + def test_initial_prompt_tokens_is_absent_on_a_single_hop_and_the_first_hops_count_otherwise(): assert initial_prompt_tokens_of((HOP_1,)) is None assert initial_prompt_tokens_of((HOP_1, HOP_2)) == 10 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 bb624add673..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 @@ -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, @@ -1894,7 +2133,9 @@ def test_gemini_tool_use(respx_mock: respx.MockRouter) -> None: assert [(call.function.name, json.loads(call.function.arguments)) for call in tool_calls] == [ ("get_weather", {"location": "Lima, Peru"}) ] - assert [chunk.choices[0].finish_reason for chunk in chunks if chunk.choices[0].finish_reason] == ["tool_calls"] + assert [chunk.choices[0].finish_reason for chunk in chunks if chunk.choices and chunk.choices[0].finish_reason] == [ + "tool_calls" + ] @pytest.mark.respx(assert_all_called=True) @@ -3037,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/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() 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/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_transformation.py b/tests/unit/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_transformation.py index fb1cf4ced2d..0726ebb5ee1 100644 --- a/tests/unit/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_transformation.py +++ b/tests/unit/llms/vertex_ai/vertex_gemma_models/test_vertex_gemma_transformation.py @@ -632,7 +632,7 @@ class TestVertexGemmaCompletion: assert chunks[0].choices[0].delta.content == "READY" assert chunks[1].choices[0].finish_reason == "stop" if stream_options and stream_options["include_usage"]: - assert chunks[-1].choices[0].delta.content is None + assert chunks[-1].choices == [] assert chunks[-1].usage.total_tokens == 114 assert span.get_metric("_dd.llmobs.total_tokens") == 114 else: diff --git a/tests/unit/proxy/client/cli/test_pi.py b/tests/unit/proxy/client/cli/test_pi.py index ff4dc25cd0c..e6669564358 100644 --- a/tests/unit/proxy/client/cli/test_pi.py +++ b/tests/unit/proxy/client/cli/test_pi.py @@ -195,7 +195,12 @@ class TestProviderBlock: assert block == { "baseUrl": "http://localhost:4000/v1", "api": "openai-completions", - "compat": {"supportsStore": False, "supportsLongCacheRetention": False}, + "compat": { + "supportsStore": False, + "supportsLongCacheRetention": False, + "sendSessionAffinityHeaders": True, + "sessionAffinityFormat": "openrouter", + }, "apiKey": "$LITELLM_PROXY_API_KEY", "models": [{"id": "m-1"}, {"id": "m-2"}], } diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py index 63fcdfbb0c5..df2fc8e1da2 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py @@ -1,6 +1,7 @@ import asyncio import json from types import SimpleNamespace +from typing import Final from unittest.mock import AsyncMock import pytest @@ -1975,6 +1976,20 @@ class TestProxyWiring: class TestStreamUsage: + @pytest.mark.asyncio + async def test_trailing_flush_survives_a_usage_chunk_with_no_choices(self): + guardrail, _ = _shielded({"[EMAIL_1]": "a@example.com"}) + held: Final = ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(content="Mail [EMAI"))]) + usage_chunk: Final = ModelResponseStream( + choices=[], usage=Usage(prompt_tokens=5, completion_tokens=7, total_tokens=12) + ) + + out: Final = await _restore_stream(guardrail, [held, usage_chunk]) + + assert out[1] is not None and out[1].choices == [] + assert [chunk.choices[0].delta.content for chunk in out[2:]] == ["[EMAI"] + assert all(getattr(chunk, "usage", None) is None for chunk in out[2:]) + @pytest.mark.asyncio async def test_trailing_flush_does_not_repeat_usage(self): """With n>=2 and include_usage, the last chunk carries usage; a flush copied from it must not. 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_parallel_request_limiter_v3.py b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py index 8a74f2d50b4..1ec505362dd 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py @@ -3,6 +3,7 @@ Unit Tests for the max parallel request limiter v3 for the proxy """ import asyncio +import contextvars import logging import sys import time @@ -54,6 +55,7 @@ from litellm.types.mcp import MCPPreCallRequestObject, MCPTransport from litellm.types.mcp_server.mcp_server_manager import MCPServer from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem from litellm.types.utils import ( + CallTypes, EmbeddingResponse, ModelResponse, PromptTokensDetailsWrapper, @@ -4206,6 +4208,175 @@ async def test_completed_responses_post_call_releases_parallel_slot() -> None: ) == 0 +@pytest.mark.asyncio +async def test_mcp_tool_call_inside_admitted_request_reuses_its_parallel_slot() -> None: + api_key: Final = hash_token("sk-mcp-auto-execute") + local_cache: Final = DualCache() + handler: Final = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + user_api_key_dict: Final = UserAPIKeyAuth(api_key=api_key, max_parallel_requests=1) + parallel_key: Final = f"{{api_key:{api_key}}}:max_parallel_requests" + owner_call_id: Final = "mcp-auto-execute-owner" + + async def in_flight() -> int: + return handler._gauge_in_flight_from_cache_value( + await local_cache.async_get_cache(key=parallel_key) + ) + + def mcp_tool_call() -> dict[str, object]: + return { + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "Tool: add"}], + "mcp_server_name": "math", + } + + async def pre_call(data: dict[str, object], call_type: str) -> None: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data=data, + call_type=call_type, + ) + + async def pre_call_from_another_request( + data: dict[str, object], call_type: str + ) -> None: + await contextvars.Context().run(asyncio.create_task, pre_call(data, call_type)) + + await pre_call( + {"model": "gpt-4o-mini", "input": "add 2 and 3", "litellm_call_id": owner_call_id}, + CallTypes.aresponses.value, + ) + await pre_call(mcp_tool_call(), CallTypes.call_mcp_tool.value) + await pre_call(mcp_tool_call(), CallTypes.call_mcp_tool.value) + assert await in_flight() == 1 + + with pytest.raises(HTTPException): + await pre_call_from_another_request( + mcp_tool_call(), CallTypes.call_mcp_tool.value + ) + + await handler.async_log_success_event( + kwargs={"litellm_call_id": owner_call_id}, + response_obj=None, + start_time=None, + end_time=None, + ) + assert await in_flight() == 0 + await pre_call_from_another_request( + {"model": "gpt-4o-mini", "input": "next request"}, + CallTypes.aresponses.value, + ) + assert await in_flight() == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("stream", "in_flight_after_first_round_logs"), + [(False, 1), (True, 0)], + ids=["non_stream_round_logged_late_keeps_slot", "stream_logged_at_its_end_releases_slot"], +) +async def test_first_round_logged_after_tool_call_reuse_releases_only_at_request_end( + stream: bool, in_flight_after_first_round_logs: int +) -> None: + reused_at: Final = datetime(2026, 1, 5, 12, 0, 0) + api_key: Final = hash_token(f"sk-mcp-late-round-log-{stream}") + local_cache: Final = DualCache() + handler: Final = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache), + time_provider=lambda: reused_at, + ) + user_api_key_dict: Final = UserAPIKeyAuth(api_key=api_key, max_parallel_requests=1) + parallel_key: Final = f"{{api_key:{api_key}}}:max_parallel_requests" + owner_call_id: Final = f"mcp-late-round-log-{stream}" + + async def in_flight() -> int: + return handler._gauge_in_flight_from_cache_value( + await local_cache.async_get_cache(key=parallel_key) + ) + + async def pre_call(data: dict[str, object], call_type: str) -> None: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data=data, + call_type=call_type, + ) + + async def round_logged(ended_at: datetime) -> None: + await handler.async_log_success_event( + kwargs={"litellm_call_id": owner_call_id, "stream": stream}, + response_obj=None, + start_time=None, + end_time=ended_at, + ) + + await pre_call( + { + "model": "claude-haiku-5-5", + "messages": [{"role": "user", "content": "add 2 and 3"}], + "litellm_call_id": owner_call_id, + }, + CallTypes.anthropic_messages.value, + ) + await pre_call( + {"model": "MCP: add", "messages": [{"role": "user", "content": "Tool: add"}], "mcp_server_name": "math"}, + CallTypes.call_mcp_tool.value, + ) + + await round_logged(reused_at - timedelta(seconds=1)) + assert await in_flight() == in_flight_after_first_round_logs + + await round_logged(reused_at + timedelta(seconds=1)) + assert await in_flight() == 0 + + +@pytest.mark.asyncio +async def test_mcp_tool_call_reusing_request_slot_keeps_request_max_parallel_headers() -> None: + api_key: Final = hash_token("sk-mcp-reuse-headers") + local_cache: Final = DualCache() + handler: Final = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + user_api_key_dict: Final = UserAPIKeyAuth(api_key=api_key, max_parallel_requests=2) + data: Final[dict[str, object]] = { + "model": "gpt-4o-mini", + "input": "add 2 and 3", + "litellm_call_id": "mcp-reuse-headers-owner", + } + for request_data, call_type in ( + (data, CallTypes.aresponses.value), + ( + {"model": "MCP: add", "messages": [{"role": "user", "content": "Tool: add"}], "mcp_server_name": "math"}, + CallTypes.call_mcp_tool.value, + ), + ): + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data=request_data, + call_type=call_type, + ) + response: Final = ModelResponse( + id="mcp-reuse-headers", + object="chat.completion", + created=0, + model="gpt-4o-mini", + usage=Usage(prompt_tokens=1, completion_tokens=1, total_tokens=2), + choices=[], + ) + response._hidden_params = {} + + await handler.async_post_call_success_hook( + data=data, user_api_key_dict=user_api_key_dict, response=response + ) + + headers: Final = response._hidden_params["additional_headers"] + assert headers["x-ratelimit-api_key-limit-max_parallel_requests"] == 2 + assert headers["x-ratelimit-api_key-remaining-max_parallel_requests"] == 1 + + @pytest.mark.asyncio async def test_concurrent_success_callbacks_release_parallel_slot_once_when_redis_fails() -> None: from unittest.mock import AsyncMock 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/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/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_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index f0baceecdaa..db0a74d360a 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -8536,6 +8536,7 @@ PROXY_STAMPED_NAMES: Final = frozenset( "litellm_trusted_callback_vars", "_litellm_addressed_response_id", "_litellm_strip_stream_usage", + "_litellm_stream_usage_injected", "client_side_timeout", "model_file_id_mapping", ) 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/proxy_server/test_routes_queue.py b/tests/unit/proxy/proxy_server/test_routes_queue.py index 07a61c8587b..a86314e3063 100644 --- a/tests/unit/proxy/proxy_server/test_routes_queue.py +++ b/tests/unit/proxy/proxy_server/test_routes_queue.py @@ -6,10 +6,13 @@ Pins (PR2): from __future__ import annotations +from collections.abc import Callable +from contextlib import AbstractContextManager from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest +from fastapi.testclient import TestClient from litellm import Router from litellm.proxy import proxy_server @@ -85,6 +88,25 @@ def test_queue_chat_completions_happy(client, auth_as, patched_queue): } +def test_queue_chat_completions_drops_client_sent_stream_usage_flags( + client: TestClient, auth_as: Callable[..., AbstractContextManager[None]], patched_queue: MagicMock +) -> None: + payload: Final = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "hi"}], + "priority": 0, + "_litellm_strip_stream_usage": True, + "_litellm_stream_usage_injected": True, + } + with auth_as(): + response: Final = client.post("/queue/chat/completions", json=payload) + assert response.status_code == 200 + forwarded: Final = patched_queue.schedule_acompletion.call_args.kwargs + assert "_litellm_strip_stream_usage" not in forwarded + assert "_litellm_stream_usage_injected" not in forwarded + assert forwarded["messages"] == [{"role": "user", "content": "hi"}] + + def test_queue_chat_completions_no_router_error(client, auth_as, queue_no_router): """Pins ``POST /queue/chat/completions`` (error: no llm_router).""" payload = {"model": "gpt-4", "messages": [{"role": "user", "content": "hi"}]} 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 5dee9b6c614..700c4958a69 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -4123,6 +4123,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: """ @@ -9238,6 +9304,7 @@ class TestApplyStreamUsageTracking: assert data["stream_options"] == {"include_usage": True} assert data["_litellm_strip_stream_usage"] is True + assert data["_litellm_stream_usage_injected"] is True def test_default_preserves_other_client_stream_options_keys(self): data = {"stream": True, "stream_options": {"include_obfuscation": True}} @@ -9246,6 +9313,7 @@ class TestApplyStreamUsageTracking: assert data["stream_options"] == {"include_obfuscation": True, "include_usage": True} assert data["_litellm_strip_stream_usage"] is True + assert data["_litellm_stream_usage_injected"] is True def test_client_requested_usage_is_left_untouched_and_not_stripped(self): data = {"stream": True, "stream_options": {"include_usage": True}} @@ -9262,6 +9330,7 @@ class TestApplyStreamUsageTracking: assert data["stream_options"]["include_usage"] is True assert data["_litellm_strip_stream_usage"] is True + assert data["_litellm_stream_usage_injected"] is True def test_explicit_false_flag_disables_injection_entirely(self): data = {"stream": True} @@ -9287,6 +9356,36 @@ class TestApplyStreamUsageTracking: assert data["stream_options"] == {"include_usage": True} assert "_litellm_strip_stream_usage" not in data + @pytest.mark.parametrize("client_stream_options", [None, {"include_obfuscation": True}]) + def test_flag_true_marks_usage_injected_for_clients_that_did_not_ask( + self, client_stream_options: dict[str, bool] | None + ): + data: Final = { + "stream": True, + **({} if client_stream_options is None else {"stream_options": client_stream_options}), + } + + _apply_stream_usage_tracking( + data=data, + general_settings={"always_include_stream_usage": True}, + route_type="acompletion", + ) + + assert data["stream_options"]["include_usage"] is True + assert data["_litellm_stream_usage_injected"] is True + + def test_flag_true_leaves_usage_unmarked_when_client_asked_and_neutralizes_client_marker(self): + data: Final = {"stream": True, "stream_options": {"include_usage": True}, "_litellm_stream_usage_injected": True} + + _apply_stream_usage_tracking( + data=data, + general_settings={"always_include_stream_usage": True}, + route_type="acompletion", + ) + + assert data["stream_options"] == {"include_usage": True} + assert data["_litellm_stream_usage_injected"] is False + def test_flag_true_respects_client_explicit_include_usage_false(self): data = {"stream": True, "stream_options": {"include_usage": False}} @@ -10841,6 +10940,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 c2a5b01fe4e..d9def23bde8 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -936,6 +936,26 @@ def test_body_snapshot_excludes_team_callback_credentials() -> None: }, proxy_request +def test_body_snapshot_keeps_the_stream_usage_injection_marker_out_of_the_stored_request() -> None: + from litellm.proxy.common_request_processing import stream_usage_tracking_updates + from litellm.proxy.litellm_pre_call_utils import refresh_proxy_server_request_body_snapshot + from litellm.types.litellm_params import STREAM_USAGE_INJECTED_FIELD + + request: Final = {"model": "gpt-5-mini", "messages": [{"role": "user", "content": "hi"}], "stream": True} + proxy_request: Final = {"body": dict(request)} + data: Final = { + **request, + **stream_usage_tracking_updates(request, {}, "acompletion", lambda: True), + "proxy_server_request": proxy_request, + } + + refresh_proxy_server_request_body_snapshot(data) + + assert data[STREAM_USAGE_INJECTED_FIELD] is True + assert proxy_request["body"]["stream_options"] == {"include_usage": True} + assert STREAM_USAGE_INJECTED_FIELD not in proxy_request["body"] + + @pytest.mark.asyncio @pytest.mark.parametrize("pre_call_ran", [False, True]) async def test_post_guardrail_snapshot_preserves_logging_only_masking_in_spend_logs( @@ -1209,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( @@ -1231,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 7b18cb3350c..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,9 @@ +import asyncio import datetime as real_datetime +import json import smtplib +import traceback +from collections.abc import Mapping from typing import Final import httpx @@ -11,8 +15,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 @@ -735,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 @@ -746,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: @@ -763,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 @@ -2154,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 @@ -2440,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) @@ -2472,10 +2518,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 +3372,124 @@ 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], + 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/proxy/test_route_llm_request.py b/tests/unit/proxy/test_route_llm_request.py index 8928b79bcd4..6044475652c 100644 --- a/tests/unit/proxy/test_route_llm_request.py +++ b/tests/unit/proxy/test_route_llm_request.py @@ -1,19 +1,42 @@ +import asyncio import contextlib +import gc +import threading +import weakref import httpx import json import litellm import pytest +import redis import respx +from collections.abc import Awaitable, Callable, Coroutine, Iterator, Mapping +from datetime import datetime from types import SimpleNamespace +from uuid import uuid4 from typing import Final from unittest.mock import MagicMock +from fakeredis import TcpFakeServer from fastapi import HTTPException -from litellm.proxy.route_llm_request import JSON_OBJECT_ADAPTER, ProxyModelNotFoundError, route_request -from litellm.types.router import RouterGeneralSettings +from litellm.caching.caching import DualCache +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.route_llm_request import ( + JSON_OBJECT_ADAPTER, + ProxyModelNotFoundError, + UserConfigRouters, + route_request, + user_config_routers, +) +from litellm.proxy.utils import ProxyLogging +from litellm.router_strategy.adaptive_router.config import ( + ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY, + ADAPTIVE_ROUTER_RESPONSE_HEADER, +) +from litellm.types.router import RouterErrors, RouterGeneralSettings, RouterRateLimitError @pytest.mark.parametrize( @@ -2434,3 +2457,571 @@ async def test_team_key_search_of_registered_bedrock_store_reaches_bedrock_with_ posted_url ) assert "Credential=AKIAREGISTRYKEY00000/" in signed_auth, signed_auth + + +@pytest.fixture(autouse=True) +def discard_cached_user_config_routers() -> Iterator[None]: + yield + user_config_routers.clear() + + +@pytest.fixture +def fake_redis_port() -> Iterator[int]: + server: Final = TcpFakeServer(("127.0.0.1", 0), server_type="redis") + thread: Final = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + yield server.server_address[1] + server.shutdown() + server.server_close() + + +async def _wait_out_a_rate_limit_minute_rollover() -> None: + seconds_left: Final = 60 - datetime.now().second + if seconds_left < 10: + await asyncio.sleep(seconds_left) + + +def _proxy_logging_obj(model: str = "rate-limited") -> LiteLLMLoggingObj: + logging_obj, _ = litellm.utils.function_setup( + original_function="acompletion", + rules_obj=litellm.utils.Rules(), + start_time=datetime.now(), + model=model, + messages=[{"role": "user", "content": "hi"}], + litellm_call_id=str(uuid4()), + ) + assert isinstance(logging_obj, LiteLLMLoggingObj) + return logging_obj + + +def _mock_deployment(model_name: str, **deployment_limits: int) -> dict[str, object]: + return { + "model_name": model_name, + **deployment_limits, + "litellm_params": {"model": "openai/gpt-5.5", "api_key": "sk-test", "mock_response": "ok"}, + } + + +async def _user_config_request_outcome( + redis_port: int | None, + routing_settings: Mapping[str, object], + deployment_limits: Mapping[str, int], + logging_obj: LiteLLMLoggingObj | None, +) -> str: + data: Final = { + "model": "rate-limited", + "messages": [{"role": "user", "content": "hi"}], + "num_retries": 0, + **({"litellm_logging_obj": logging_obj} if logging_obj is not None else {}), + "user_config": { + **routing_settings, + **({"redis_host": "127.0.0.1", "redis_port": redis_port} if redis_port is not None else {}), + "model_list": [_mock_deployment("rate-limited", **deployment_limits)], + }, + } + try: + await (await route_request(data, None, None, "acompletion")) + except (litellm.RateLimitError, RouterRateLimitError): + return "429" + finally: + await asyncio.sleep(0.3) + return "200" + + +async def _tpm_for_three_requests() -> int: + messages: Final = [{"role": "user", "content": "hi"}] + probe: Final = await litellm.acompletion(model="openai/gpt-5.5", messages=messages, mock_response="ok") + assert isinstance(probe, litellm.ModelResponse) + probe_usage: Final = getattr(probe, "usage", None) + assert isinstance(probe_usage, litellm.Usage) + return 2 * probe_usage.total_tokens + litellm.token_counter(messages=messages) + + +async def _request_cost() -> float: + probe: Final = await litellm.acompletion( + model="openai/gpt-5.5", messages=[{"role": "user", "content": "hi"}], mock_response="ok" + ) + return litellm.completion_cost(completion_response=probe) + + +async def _proxy_router_outcomes( + redis_port: int | None, + routing_settings: Mapping[str, object], + deployment_limits: Mapping[str, int], + requests: int, +) -> list[str]: + router: Final = litellm.Router( + **routing_settings, + num_retries=0, + model_list=[_mock_deployment("proxy-rate-limited", **deployment_limits)], + **({"redis_host": "127.0.0.1", "redis_port": redis_port} if redis_port is not None else {}), + ) + + async def outcome() -> str: + try: + await router.acompletion(model="proxy-rate-limited", messages=[{"role": "user", "content": "hi"}]) + except (litellm.RateLimitError, RouterRateLimitError): + return "429" + finally: + await asyncio.sleep(0.3) + return "200" + + try: + return [await outcome() for _ in range(requests)] + finally: + router.discard() + + +@pytest.mark.asyncio +async def test_a_user_config_adaptive_request_answers_with_the_model_it_picked() -> None: + data: Final = { + "model": "smart-cheap-router", + "messages": [{"role": "user", "content": "Write a one line python function that adds two numbers"}], + "metadata": {}, + "user_config": { + "model_list": [ + { + "model_name": "smart-cheap-router", + "litellm_params": { + "model": "auto_router/adaptive_router", + "adaptive_router_config": {"available_models": ["fast", "smart"]}, + }, + }, + { + **_mock_deployment("fast"), + "model_info": {"adaptive_router_preferences": {"quality_tier": 2, "strengths": []}}, + }, + { + **_mock_deployment("smart"), + "model_info": {"adaptive_router_preferences": {"quality_tier": 3, "strengths": ["code_generation"]}}, + }, + ] + }, + } + + await (await route_request(data, None, None, "acompletion")) + headers: Final = await ProxyLogging(user_api_key_cache=DualCache()).post_call_response_headers_hook( + data=data, user_api_key_dict=UserAPIKeyAuth(), response=None + ) + + assert headers == {ADAPTIVE_ROUTER_RESPONSE_HEADER: data["metadata"][ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY]} + + +def _global_router_callbacks() -> tuple[list[object], list[object]]: + return list(litellm.callbacks), list(litellm.input_callback) + + +@pytest.mark.parametrize( + "routing_settings, deployment_limits", + [ + ({"routing_strategy": "usage-based-routing-v2"}, {"rpm": 2}), + ({"routing_strategy": "usage-based-routing-v2"}, {"tpm": 30}), + ({"routing_strategy": "usage-based-routing"}, {"rpm": 2}), + ({"optional_pre_call_checks": ["enforce_model_rate_limits"]}, {"rpm": 2}), + ], +) +@pytest.mark.parametrize("arrives_with_logging_obj", [True, False]) +@pytest.mark.parametrize("with_redis", [True, False]) +@pytest.mark.asyncio +async def test_user_config_router_enforces_its_limit_across_proxy_requests( + fake_redis_port: int, + routing_settings: Mapping[str, object], + deployment_limits: Mapping[str, int], + arrives_with_logging_obj: bool, + with_redis: bool, +) -> None: + await _wait_out_a_rate_limit_minute_rollover() + callbacks_before: Final = _global_router_callbacks() + outcomes: Final = [ + await _user_config_request_outcome( + fake_redis_port if with_redis else None, + routing_settings, + deployment_limits, + _proxy_logging_obj() if arrives_with_logging_obj else None, + ) + for _ in range(6) + ] + + assert "429" in outcomes, outcomes + assert _global_router_callbacks() == callbacks_before + assert outcomes == await _proxy_router_outcomes( + fake_redis_port if with_redis else None, routing_settings, deployment_limits, len(outcomes) + ) + + +@pytest.mark.parametrize("arrives_with_logging_obj", [True, False]) +@pytest.mark.asyncio +async def test_concurrent_user_config_requests_each_count_once_against_the_tpm_limit( + fake_redis_port: int, arrives_with_logging_obj: bool +) -> None: + tpm_for_three_requests: Final = await _tpm_for_three_requests() + + def request() -> Coroutine[None, None, str]: + return _user_config_request_outcome( + fake_redis_port, + {"routing_strategy": "usage-based-routing-v2"}, + {"tpm": tpm_for_three_requests}, + _proxy_logging_obj() if arrives_with_logging_obj else None, + ) + + await _wait_out_a_rate_limit_minute_rollover() + concurrent_outcomes: Final = await asyncio.gather(request(), request()) + later_outcomes: Final = [await request() for _ in range(2)] + + assert [*concurrent_outcomes, *later_outcomes] == ["200", "200", "200", "429"] + + +@pytest.mark.asyncio +async def test_overlapping_user_config_requests_share_one_rpm_limit(fake_redis_port: int) -> None: + def request() -> Coroutine[None, None, str]: + return _user_config_request_outcome( + fake_redis_port, {"routing_strategy": "usage-based-routing-v2"}, {"rpm": 2}, _proxy_logging_obj() + ) + + await _wait_out_a_rate_limit_minute_rollover() + outcomes: Final = await asyncio.gather(*(request() for _ in range(4))) + + assert outcomes.count("200") <= 2, outcomes + assert "429" in outcomes, outcomes + + +@pytest.mark.parametrize( + "routing_strategy, limited_by", + [("usage-based-routing-v2", "rpm"), ("usage-based-routing-v2", "tpm"), ("usage-based-routing", "rpm")], +) +@pytest.mark.parametrize("proxy_router_shares_redis", [True, False]) +@pytest.mark.asyncio +async def test_a_proxy_router_on_the_same_strategy_leaves_the_user_config_limit_unchanged( + fake_redis_port: int, routing_strategy: str, limited_by: str, proxy_router_shares_redis: bool +) -> None: + deployment_limits: Final = {"rpm": 3} if limited_by == "rpm" else {"tpm": await _tpm_for_three_requests()} + + def request() -> Coroutine[None, None, str]: + return _user_config_request_outcome( + fake_redis_port, {"routing_strategy": routing_strategy}, deployment_limits, _proxy_logging_obj() + ) + + async def outcomes() -> list[str]: + first_instance: Final = [await request() for _ in range(2)] + user_config_routers.clear() + second_instance: Final = [await request() for _ in range(3)] + return [*first_instance, *second_instance] + + await _wait_out_a_rate_limit_minute_rollover() + alone: Final = await outcomes() + user_config_routers.clear() + await asyncio.sleep(0.3) + redis.Redis(host="127.0.0.1", port=fake_redis_port).flushall() + proxy_router: Final = litellm.Router( + routing_strategy=routing_strategy, + model_list=[_mock_deployment("proxy-model")], + redis_host="127.0.0.1" if proxy_router_shares_redis else None, + redis_port=fake_redis_port if proxy_router_shares_redis else None, + ) + try: + beside_proxy_router: Final = await outcomes() + finally: + proxy_router.discard() + + assert "429" in alone, alone + assert beside_proxy_router == alone + + +@pytest.mark.parametrize("user_config_settings", [{}, {"optional_pre_call_checks": ["router_budget_limiting"]}]) +@pytest.mark.asyncio +async def test_user_config_requests_stay_under_the_proxy_provider_budget( + user_config_settings: Mapping[str, object], +) -> None: + proxy_router: Final = litellm.Router( + model_list=[_mock_deployment("proxy-model")], + provider_budget_config={"openai": {"budget_limit": await _request_cost() / 2, "time_period": "1d"}}, + ) + try: + await proxy_router.acompletion(model="proxy-model", messages=[{"role": "user", "content": "hi"}]) + await asyncio.sleep(0.3) + with pytest.raises(ValueError, match=RouterErrors.no_deployments_with_provider_budget_routing.value): + await _user_config_request_outcome(None, user_config_settings, {}, _proxy_logging_obj()) + finally: + proxy_router.discard() + + +@pytest.mark.asyncio +async def test_a_user_config_provider_budget_limits_its_own_requests() -> None: + budget_limit: Final = await _request_cost() / 2 + + def request() -> Coroutine[None, None, str]: + budget: Final = {"openai": {"budget_limit": budget_limit, "time_period": "1d"}} + return _user_config_request_outcome(None, {"provider_budget_config": budget}, {}, _proxy_logging_obj()) + + first: Final = await request() + with pytest.raises(ValueError, match=RouterErrors.no_deployments_with_provider_budget_routing.value): + await request() + assert first == "200" + + +@pytest.mark.asyncio +async def test_a_request_in_flight_when_its_user_config_router_is_dropped_still_counts_against_its_budget( + fake_redis_port: int, +) -> None: + budget_limit: Final = 1.5 * await _request_cost() + + async def request() -> None: + await ( + await route_request( + { + "model": "budgeted", + "messages": [{"role": "user", "content": "hi"}], + "litellm_logging_obj": _proxy_logging_obj("budgeted"), + "user_config": { + "redis_host": "127.0.0.1", + "redis_port": fake_redis_port, + "provider_budget_config": {"openai": {"budget_limit": budget_limit, "time_period": "1d"}}, + "model_list": [ + { + "model_name": "budgeted", + "litellm_params": { + "model": "openai/gpt-5.5", + "api_key": "sk-test", + "mock_response": "ok", + "mock_delay": 0.5, + }, + } + ], + }, + }, + None, + None, + "acompletion", + ) + ) + + await request() + in_flight: Final = asyncio.create_task(request()) + await asyncio.sleep(0.2) + user_config_routers.clear() + await in_flight + await asyncio.sleep(0.3) + + with pytest.raises(ValueError, match=RouterErrors.no_deployments_with_provider_budget_routing.value): + await request() + + +async def _user_config_request(model: str, user_config: Mapping[str, object], arrives_with_logging_obj: bool) -> object: + return await ( + await route_request( + { + "model": model, + "messages": list(_HELLO_MESSAGES), + "num_retries": 0, + **({"litellm_logging_obj": _proxy_logging_obj(model)} if arrives_with_logging_obj else {}), + "user_config": user_config, + }, + None, + None, + "acompletion", + ) + ) + + +def _ordered_pair(model_name: str, first_base: str, second_base: str) -> list[dict[str, object]]: + return [ + { + "model_name": model_name, + "litellm_params": {"model": "openai/gpt-5.5", "api_key": "sk-test", "api_base": api_base, "order": order}, + } + for order, api_base in ((1, first_base), (2, second_base)) + ] + + +@pytest.mark.parametrize("arrives_with_logging_obj", [True, False]) +@pytest.mark.asyncio +async def test_a_rate_limited_user_config_deployment_cools_down_for_the_next_requests( + monkeypatch: pytest.MonkeyPatch, arrives_with_logging_obj: bool +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + rate_limited_base: Final = "https://rate-limited.example/v1" + healthy_base: Final = "https://healthy.example/v1" + user_config: Final = {"model_list": _ordered_pair("flaky", rate_limited_base, healthy_base)} + + with respx.mock(assert_all_called=False) as respx_mock: + rate_limited: Final = respx_mock.post(f"{rate_limited_base}/chat/completions").mock( + return_value=httpx.Response(429, json={"error": {"message": "slow down"}}) + ) + respx_mock.post(f"{healthy_base}/chat/completions").mock(return_value=_chat_completion_response("gpt-5.5")) + for _ in range(3): + await _user_config_request("flaky", user_config, arrives_with_logging_obj) + await asyncio.sleep(0.3) + + assert rate_limited.call_count == 1 + + +async def _remaining_usage_headers_after_429_then_200s( + send: Callable[[], Awaitable[object]], api_base: str +) -> list[tuple[object, object] | str]: + async def remaining_usage() -> tuple[object, object] | str: + try: + response: Final = await send() + except litellm.RateLimitError: + return "429" + finally: + await asyncio.sleep(0.3) + assert isinstance(response, litellm.ModelResponse) + headers: Final = response._hidden_params["additional_headers"] + return headers.get("x-ratelimit-remaining-requests"), headers.get("x-ratelimit-remaining-tokens") + + with respx.mock(assert_all_called=False) as respx_mock: + respx_mock.post(f"{api_base}/chat/completions").mock( + side_effect=[ + httpx.Response(429, json={"error": {"message": "slow down"}}), + _chat_completion_response("gpt-5.5"), + _chat_completion_response("gpt-5.5"), + ] + ) + return [await remaining_usage() for _ in range(3)] + + +@pytest.mark.parametrize("arrives_with_logging_obj", [True, False]) +@pytest.mark.asyncio +async def test_user_config_remaining_usage_headers_match_a_proxy_router_with_the_same_deployment( + monkeypatch: pytest.MonkeyPatch, arrives_with_logging_obj: bool +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + await _wait_out_a_rate_limit_minute_rollover() + proxy_base: Final = "https://proxy-limited.example/v1" + user_config_base: Final = "https://user-config-limited.example/v1" + + def deployment(model_name: str, api_base: str) -> dict[str, object]: + return { + "model_name": model_name, + "litellm_params": { + "model": "openai/gpt-5.5", + "api_key": "sk-test", + "api_base": api_base, + "rpm": 10, + "tpm": 1000, + }, + } + + proxy_router: Final = litellm.Router(model_list=[deployment("proxy-limited", proxy_base)]) + try: + proxy_router_headers: Final = await _remaining_usage_headers_after_429_then_200s( + lambda: proxy_router.acompletion(model="proxy-limited", messages=list(_HELLO_MESSAGES), num_retries=0), + proxy_base, + ) + finally: + proxy_router.discard() + + user_config: Final = {"model_list": [deployment("user-config-limited", user_config_base)]} + user_config_headers: Final = await _remaining_usage_headers_after_429_then_200s( + lambda: _user_config_request("user-config-limited", user_config, arrives_with_logging_obj), user_config_base + ) + + assert user_config_headers == proxy_router_headers + + +async def _first_deployment_calls_over_successes_then_failures( + send: Callable[[], Awaitable[object]], first_base: str, second_base: str +) -> int: + successes: Final = [_chat_completion_response("gpt-5.5") for _ in range(6)] + failures: Final = [httpx.Response(500, json={"error": {"message": "boom"}}) for _ in range(6)] + with respx.mock(assert_all_called=False) as respx_mock: + first: Final = respx_mock.post(f"{first_base}/chat/completions").mock(side_effect=[*successes, *failures]) + respx_mock.post(f"{second_base}/chat/completions").mock(return_value=_chat_completion_response("gpt-5.5")) + for _ in range(len(successes) + len(failures)): + await send() + await asyncio.sleep(0.3) + return first.call_count + + +@pytest.mark.parametrize("arrives_with_logging_obj", [True, False]) +@pytest.mark.asyncio +async def test_user_config_successes_keep_a_deployment_out_of_cooldown_like_a_proxy_router( + monkeypatch: pytest.MonkeyPatch, arrives_with_logging_obj: bool +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + await _wait_out_a_rate_limit_minute_rollover() + proxy_router: Final = litellm.Router( + model_list=_ordered_pair("proxy-pair", "https://proxy-first.example/v1", "https://proxy-second.example/v1") + ) + try: + proxy_router_calls: Final = await _first_deployment_calls_over_successes_then_failures( + lambda: proxy_router.acompletion(model="proxy-pair", messages=list(_HELLO_MESSAGES), num_retries=0), + "https://proxy-first.example/v1", + "https://proxy-second.example/v1", + ) + finally: + proxy_router.discard() + + user_config: Final = { + "model_list": _ordered_pair("user-pair", "https://user-first.example/v1", "https://user-second.example/v1") + } + + user_config_calls: Final = await _first_deployment_calls_over_successes_then_failures( + lambda: _user_config_request("user-pair", user_config, arrives_with_logging_obj), + "https://user-first.example/v1", + "https://user-second.example/v1", + ) + + assert user_config_calls == proxy_router_calls + + +@pytest.mark.asyncio +async def test_user_config_adaptive_router_learns_from_proxy_requests() -> None: + user_config: Final = { + "model_list": [ + _mock_deployment("smart"), + _mock_deployment("fast"), + { + "model_name": "adaptive", + "litellm_params": { + "model": "auto_router/adaptive_router", + "adaptive_router_config": {"available_models": ["smart", "fast"]}, + }, + }, + ] + } + await ( + await route_request( + { + "model": "adaptive", + "messages": [{"role": "user", "content": "write a sorting function"}], + "metadata": {"session_id": "session-1"}, + "litellm_logging_obj": _proxy_logging_obj("adaptive"), + "user_config": user_config, + }, + None, + None, + "acompletion", + ) + ) + await asyncio.sleep(0.3) + + router: Final = user_config_routers.get(user_config, lambda: pytest.fail("the user_config router was not reused")) + snapshot: Final = await router.adaptive_routers["adaptive"][0].strategy.get_state_snapshot() + assert snapshot["feedback_contexts_live"] == 1, snapshot + + +@pytest.mark.asyncio +async def test_user_config_routers_discard_the_least_recently_used_router_past_capacity() -> None: + routers: Final = UserConfigRouters(max_size=2) + + def config(model_name: str) -> dict[str, object]: + return {"routing_strategy": "usage-based-routing-v2", "model_list": [_mock_deployment(model_name)]} + + def build(model_name: str) -> litellm.Router: + return litellm.Router(**config(model_name)) + + def must_be_cached() -> litellm.Router: + pytest.fail("a cached router was rebuilt") + + first_handler: Final = weakref.ref(routers.get(config("first"), lambda: build("first")).lowesttpm_logger_v2) + second_handler: Final = weakref.ref(routers.get(config("second"), lambda: build("second")).lowesttpm_logger_v2) + routers.get(config("first"), must_be_cached) + routers.get(config("third"), lambda: build("third")) + await asyncio.sleep(0.2) + gc.collect() + + assert routers.get(config("first"), must_be_cached) is not None + assert first_handler() is not None + assert second_handler() is None + routers.clear() 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/mcp/test_aresponses_api_with_mcp.py b/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py index 60d4d2e89b5..e88d951ff09 100644 --- a/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py +++ b/tests/unit/responses/mcp/test_aresponses_api_with_mcp.py @@ -1,7 +1,9 @@ import asyncio +import gc import json import logging import sys +import warnings from collections.abc import Callable, Iterator, Mapping, Sequence from contextlib import contextmanager, nullcontext from datetime import datetime @@ -18,12 +20,14 @@ from mcp.types import Tool as MCPTool # Import required modules import litellm +from litellm.exceptions import GuardrailRaisedException, MidStreamFallbackError from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.langfuse.langfuse import langfuse_trace_id, langfuse_trace_metadata 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, @@ -31,7 +35,7 @@ from litellm.types.llms.openai import ( OpenAIMcpServerTool, ToolParam, ) -from litellm.types.utils import StandardLoggingPayload, Usage +from litellm.types.utils import ModelResponse, StandardLoggingPayload, Usage class MockUserAPIKeyAuth: @@ -1518,25 +1522,39 @@ def _openai_answer_round() -> bytes: ) -def _anthropic_round(message_id: str, block: Mapping[str, object], delta: Mapping[str, object], stop: str) -> bytes: +def _anthropic_message_start( + message_id: str, model: str, input_tokens: int, cache_usage: Mapping[str, object] = MappingProxyType({}) +) -> dict[str, object]: + return { + "type": "message_start", + "message": { + "id": message_id, + "type": "message", + "role": "assistant", + "model": model, + "content": [], + "stop_reason": None, + "usage": {"input_tokens": input_tokens, "output_tokens": 1, **cache_usage}, + }, + } + + +def _anthropic_round( + message_id: str, + block: Mapping[str, object], + delta: Mapping[str, object], + stop: str, + model: str = "claude-haiku-4-5", + usage: tuple[int, int] = (5, 5), + cache_usage: Mapping[str, object] = MappingProxyType({}), +) -> bytes: return _sse( ( - { - "type": "message_start", - "message": { - "id": message_id, - "type": "message", - "role": "assistant", - "model": "claude-haiku-4-5", - "content": [], - "stop_reason": None, - "usage": {"input_tokens": 5, "output_tokens": 1}, - }, - }, + _anthropic_message_start(message_id, model, usage[0], cache_usage), {"type": "content_block_start", "index": 0, "content_block": dict(block)}, {"type": "content_block_delta", "index": 0, "delta": dict(delta)}, {"type": "content_block_stop", "index": 0}, - {"type": "message_delta", "delta": {"stop_reason": stop}, "usage": {"output_tokens": 5}}, + {"type": "message_delta", "delta": {"stop_reason": stop}, "usage": {"output_tokens": usage[1]}}, {"type": "message_stop"}, ) ) @@ -1586,6 +1604,7 @@ class _RoundRecorder(CustomLogger): def __init__(self) -> None: super().__init__() self.logged: tuple[StandardLoggingPayload, ...] = () + self.results: tuple[object, ...] = () self.failed: tuple[StandardLoggingPayload, ...] = () self.hooked: tuple[object, ...] = () @@ -1593,6 +1612,7 @@ class _RoundRecorder(CustomLogger): self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime ) -> None: self.logged = (*self.logged, cast(StandardLoggingPayload, kwargs["standard_logging_object"])) + self.results = (*self.results, response_obj) async def async_log_failure_event( self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime @@ -1624,12 +1644,15 @@ _ECHO_MCP_TOOL: Final = MappingProxyType( def _serve_responses( - respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, responses: Sequence[httpx.Response] + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + responses: Sequence[httpx.Response], + url_regex: str = r"https://mcp-rounds\.example\.com/.*responses", ) -> respx.Route: monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") litellm.in_memory_llm_clients_cache.flush_cache() - return respx_mock.post(url__regex=r"https://mcp-rounds\.example\.com/.*responses").mock(side_effect=list(responses)) + return respx_mock.post(url__regex=url_regex).mock(side_effect=list(responses)) def _serve_rounds( @@ -1957,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) ) @@ -2424,3 +2462,894 @@ async def test_streaming_mcp_event_order_and_response_id_consistency( assert lifecycle_ids[0] == lifecycle_ids[1] assert lifecycle_ids[2] != lifecycle_ids[0] assert not [record.getMessage() for record in caplog.records if record.levelno >= logging.ERROR] + + +_FIXED_MARGIN: Final = 0.001 + + +def _streamed(body: Mapping[str, object], **terminal: object) -> httpx.Response: + completed: Final = json.loads(json.dumps({**body, **terminal}, default=dict)) + opened: Final = {**completed, "status": "in_progress", "output": [], "usage": None} + items: Final = tuple(enumerate(completed["output"])) + events: Final = ( + {"type": "response.created", "response": opened}, + {"type": "response.in_progress", "response": opened}, + *({"type": "response.output_item.added", "output_index": index, "item": item} for index, item in items), + *({"type": "response.output_item.done", "output_index": index, "item": item} for index, item in items), + {"type": f"response.{completed['status']}", "response": completed}, + ) + return httpx.Response( + 200, + content=_sse(tuple({**event, "sequence_number": number} for number, event in enumerate(events))), + headers=SSE_HEADERS, + ) + + +class _CutAfterTheFirstEvent(httpx.AsyncByteStream): + async def __aiter__(self): + opened: Final = {**json.loads(json.dumps(_ROUND_TWO_BODY, default=dict)), "status": "in_progress"} + yield _sse(({"type": "response.created", "response": opened, "sequence_number": 0},)) + raise httpx.ReadError("connection reset mid-stream") + + +async def _open_stream_auto_execute(logging_obj: Logging) -> object: + return await litellm.aresponses( + model="openai/gpt-5.6", + input="Call echo with PING", + tools=[dict(_ECHO_MCP_TOOL)], + api_key="sk-mcp-rounds", + api_base="https://mcp-rounds.example.com/v1", + stream=True, + litellm_logging_obj=logging_obj, + ) + + +async def _drain(stream: object) -> tuple[ResponsesAPIStreamingResponse, ...]: + return tuple([event async for event in cast(Any, stream)]) + + +async def _stream_auto_execute(logging_obj: Logging) -> tuple[ResponsesAPIStreamingResponse, ...]: + with _echo_tool_served("PING"): + return await _drain(await _open_stream_auto_execute(logging_obj)) + + +def _final_response( + events: Sequence[ResponsesAPIStreamingResponse], terminal: str = "response.completed" +) -> ResponsesAPIResponse: + assert _event_type(events[-1]) == terminal + response: Final = getattr(events[-1], "response", None) + assert isinstance(response, ResponsesAPIResponse) + return response + + +def _assert_one_stream_row_bills_both_rounds_once( + logged: tuple[StandardLoggingPayload, ...], client: ResponsesAPIResponse +) -> None: + charged_once: Final = _round_cost(_ROUND_ONE_BODY) + _round_cost(_ROUND_TWO_BODY) - _FIXED_MARGIN + assert len(logged) == 1 + row: Final = logged[0] + assert _output_types(row["response"]["output"]) == _output_types(client.output) + assert "message" in _output_types(client.output) + assert "function_call" not in _output_types(client.output) + assert (row["prompt_tokens"], row["completion_tokens"]) == (71 + 149, 19 + 5) + assert row["cache_hit"] is None + assert row["metadata"]["usage_object"]["prompt_tokens_details"]["cached_tokens"] == 32 + 64 + assert row["hidden_params"]["additional_headers"]["llm_provider-content-type"] == SSE_HEADERS["content-type"] + assert row["response_cost"] == pytest.approx(charged_once) + assert row["cost_breakdown"]["total_cost"] == pytest.approx(charged_once) + assert client.usage is not None + assert (client.usage.input_tokens, client.usage.output_tokens) == (71 + 149, 19 + 5) + assert client.usage.cost == pytest.approx(charged_once) + + +def _assert_breakdown_agrees_with_the_billed_cost(row: StandardLoggingPayload) -> None: + breakdown: Final = row.get("cost_breakdown") + assert breakdown is None or breakdown["total_cost"] == pytest.approx(row["response_cost"]) + + +def _request_failure_row(logging_obj: Logging, success_row: StandardLoggingPayload) -> Mapping[str, object]: + failure_row: Final = _failure_fields_to_lift(MappingProxyType({"litellm_logging_obj": logging_obj})) + failure_payload: Final = cast(StandardLoggingPayload, failure_row["standard_logging_object"]) + assert failure_payload["status"] == "failure" + assert failure_payload["id"] != success_row["id"] + return failure_row + + +def _assert_billed_tokens(row: Mapping[str, object], tokens: tuple[int, int]) -> None: + usage: Final = row["combined_usage_object"] + assert isinstance(usage, Usage) + assert (usage.prompt_tokens, usage.completion_tokens) == tokens + + +def _assert_one_row_bills(logged: tuple[StandardLoggingPayload, ...], body: Mapping[str, object]) -> None: + usage: Final = cast(Mapping[str, int], body["usage"]) + assert len(logged) == 1 + assert (logged[0]["prompt_tokens"], logged[0]["completion_tokens"]) == ( + usage["input_tokens"], + usage["output_tokens"], + ) + assert logged[0]["response_cost"] == pytest.approx(_round_cost(body)) + assert logged[0]["response_cost"] > 0 + + +def _assert_one_row_bills_round_one_twice_with_one_fixed_margin(logged: tuple[StandardLoggingPayload, ...]) -> None: + charged_once: Final = 2 * _round_cost(_ROUND_ONE_BODY) - _FIXED_MARGIN + assert len(logged) == 1 + assert (logged[0]["prompt_tokens"], logged[0]["completion_tokens"]) == (2 * 71, 2 * 19) + assert logged[0]["response_cost"] == pytest.approx(charged_once) + assert logged[0]["cost_breakdown"]["total_cost"] == pytest.approx(charged_once) + + +@pytest.mark.asyncio +async def test_stream_auto_execute_logs_the_final_response_billed_for_both_rounds_with_one_fixed_margin( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + _serve_responses(respx_mock, monkeypatch, (_streamed(_ROUND_ONE_BODY), _streamed(_ROUND_TWO_BODY))) + recorder: Final = _RoundRecorder() + + events: Final = await _stream_auto_execute(_rounds_logging_obj(recorder, "mcp-stream-rounds")) + await _settle(lambda: len(recorder.logged) > 0) + + _assert_one_stream_row_bills_both_rounds_once(recorder.logged, _final_response(events)) + + +@pytest.mark.asyncio +async def test_stream_auto_execute_deferred_for_a_post_call_guardrail_logs_both_rounds_once_when_the_proxy_releases( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.utils import ProxyLogging + + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + _serve_responses(respx_mock, monkeypatch, (_streamed(_ROUND_ONE_BODY), _streamed(_ROUND_TWO_BODY))) + recorder: Final = _RoundRecorder() + logging_obj: Final = _rounds_logging_obj(recorder, "mcp-stream-rounds-deferred") + + with _echo_tool_served("PING"): + stream: Final = await _open_stream_auto_execute(logging_obj) + ProxyBaseLLMRequestProcessing(data={})._arm_deferred_stream_dispatch( + response=stream, route_type="aresponses", user_api_key_dict=MagicMock(), logging_obj=logging_obj + ) + events: Final = await _drain(stream) + await asyncio.sleep(0.2) + logged_before_release: Final = recorder.logged + ProxyLogging.fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) + await _settle(lambda: len(recorder.logged) > 0) + + assert logged_before_release == () + _assert_one_stream_row_bills_both_rounds_once(recorder.logged, _final_response(events)) + + +@pytest.mark.asyncio +async def test_stream_auto_execute_bills_the_paid_first_round_on_its_success_row_when_the_follow_up_stream_breaks( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + from litellm.proxy.utils import ProxyLogging + + _serve_responses( + respx_mock, + monkeypatch, + (_streamed(_ROUND_ONE_BODY), httpx.Response(200, stream=_CutAfterTheFirstEvent(), headers=SSE_HEADERS)), + ) + recorder: Final = _RoundRecorder() + logging_obj: Final = _rounds_logging_obj(recorder, "mcp-stream-follow-up-breaks") + + with pytest.raises(litellm.ServiceUnavailableError): + await _stream_auto_execute(logging_obj) + ProxyLogging.fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) + await _settle(lambda: len(recorder.logged) > 0 and len(recorder.failed) > 0) + + _assert_one_row_bills(recorder.logged, _ROUND_ONE_BODY) + assert [row["response_cost"] for row in recorder.failed] == [0.0] + assert _request_failure_row(logging_obj, recorder.logged[0]).get("response_cost", 0.0) == 0.0 + + +@pytest.mark.asyncio +async def test_stream_auto_execute_blocked_by_a_post_call_guardrail_keeps_both_rounds_on_the_success_row( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.utils import ProxyLogging + + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + _serve_responses(respx_mock, monkeypatch, (_streamed(_ROUND_ONE_BODY), _streamed(_ROUND_TWO_BODY))) + recorder: Final = _RoundRecorder() + logging_obj: Final = _rounds_logging_obj(recorder, "mcp-stream-rounds-blocked") + blocked: Final = GuardrailRaisedException(guardrail_name="output-policy", message="blocked") + + with _echo_tool_served("PING"): + stream: Final = await _open_stream_auto_execute(logging_obj) + ProxyBaseLLMRequestProcessing(data={})._arm_deferred_stream_dispatch( + response=stream, route_type="aresponses", user_api_key_dict=MagicMock(), logging_obj=logging_obj + ) + events: Final = await _drain(stream) + if not ProxyLogging._discard_deferred_stream_logging_for_failure({"litellm_logging_obj": logging_obj}, blocked): + ProxyLogging.fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) + await logging_obj.async_failure_handler(blocked, "blocked by output-policy") + await _settle(lambda: len(recorder.logged) > 0 and len(recorder.failed) > 0) + + _assert_one_stream_row_bills_both_rounds_once(recorder.logged, _final_response(events)) + assert [row["response_cost"] for row in recorder.failed] == [0.0] + + +@pytest.mark.asyncio +async def test_stream_auto_execute_bills_the_paid_first_round_once_when_the_follow_up_is_rejected( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + context_window_error: Final = { + "error": { + "message": "Your input exceeds the context window of this model.", + "type": "invalid_request_error", + "param": "input", + "code": "context_length_exceeded", + } + } + _serve_responses( + respx_mock, monkeypatch, (_streamed(_ROUND_ONE_BODY), httpx.Response(400, json=context_window_error)) + ) + recorder: Final = _RoundRecorder() + + events: Final = await _stream_auto_execute(_rounds_logging_obj(recorder, "mcp-stream-follow-up-rejected")) + await _settle(lambda: len(recorder.logged) > 0 and len(recorder.failed) > 0) + + assert _event_type(events[-1]) == "error" + _assert_one_row_bills(recorder.logged, _ROUND_ONE_BODY) + assert [row["response_cost"] for row in recorder.failed] == [0.0] + + +@pytest.mark.asyncio +async def test_stream_auto_execute_bills_every_paid_round_with_one_fixed_margin_when_the_third_round_is_rejected( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + _serve_responses( + respx_mock, + monkeypatch, + ( + _streamed(_ROUND_ONE_BODY), + _streamed(_ROUND_ONE_BODY), + httpx.Response(400, json={"error": {"message": "Your input exceeds the context window of this model."}}), + ), + ) + recorder: Final = _RoundRecorder() + + await _stream_auto_execute(_rounds_logging_obj(recorder, "mcp-stream-third-round-rejected")) + await _settle(lambda: len(recorder.logged) > 0) + + _assert_one_row_bills_round_one_twice_with_one_fixed_margin(recorder.logged) + + +@pytest.mark.asyncio +async def test_stream_auto_execute_bills_every_finished_round_when_the_client_disconnects_during_a_later_tool_call( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + from litellm.proxy.utils import ProxyLogging + + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + _serve_responses(respx_mock, monkeypatch, (_streamed(_ROUND_ONE_BODY), _streamed(_ROUND_ONE_BODY))) + recorder: Final = _RoundRecorder() + logging_obj: Final = _rounds_logging_obj(recorder, "mcp-stream-disconnect-in-tool-call") + echo_then_disconnect: Final = AsyncMock( + side_effect=[[{"tool_call_id": "call_echo", "result": "PING", "name": "echo"}], asyncio.CancelledError()] + ) + + with ( + _echo_tool_served("PING"), + patch.object(LiteLLM_Proxy_MCP_Handler, "execute_tool_calls", echo_then_disconnect), + pytest.raises(asyncio.CancelledError), + ): + await _drain(await _open_stream_auto_execute(logging_obj)) + ProxyLogging.fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) + await _settle(lambda: len(recorder.logged) > 0) + + _assert_one_row_bills_round_one_twice_with_one_fixed_margin(recorder.logged) + + +@pytest.mark.asyncio +async def test_stream_auto_execute_answered_without_a_tool_call_logs_its_one_round( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + _serve_responses(respx_mock, monkeypatch, (_streamed(_ROUND_TWO_BODY),)) + recorder: Final = _RoundRecorder() + + events: Final = await _stream_auto_execute(_rounds_logging_obj(recorder, "mcp-stream-one-round")) + await _settle(lambda: len(recorder.logged) > 0) + + assert _output_types(_final_response(events).output) == ("message",) + _assert_one_row_bills(recorder.logged, _ROUND_TWO_BODY) + + +@pytest.mark.asyncio +async def test_stream_auto_execute_bills_both_rounds_with_one_fixed_margin_when_the_final_round_is_incomplete( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + truncated: Final = _streamed( + _ROUND_TWO_BODY, status="incomplete", incomplete_details={"reason": "max_output_tokens"} + ) + _serve_responses(respx_mock, monkeypatch, (_streamed(_ROUND_ONE_BODY), truncated)) + recorder: Final = _RoundRecorder() + + events: Final = await _stream_auto_execute(_rounds_logging_obj(recorder, "mcp-stream-final-round-incomplete")) + await _settle(lambda: len(recorder.logged) > 0) + + _assert_one_stream_row_bills_both_rounds_once(recorder.logged, _final_response(events, "response.incomplete")) + + +@pytest.mark.asyncio +async def test_stream_auto_execute_bills_a_failed_follow_up_on_its_own_failure_row_with_no_second_fixed_margin( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "cost_margin_config", {"openai": {"fixed_amount": _FIXED_MARGIN}}) + server_error: Final = {"code": "server_error", "message": "The server had an error processing your request."} + _serve_responses( + respx_mock, + monkeypatch, + (_streamed(_ROUND_ONE_BODY), _streamed(_ROUND_TWO_BODY, status="failed", error=server_error)), + ) + recorder: Final = _RoundRecorder() + logging_obj: Final = _rounds_logging_obj(recorder, "mcp-stream-follow-up-failed") + + with pytest.raises(litellm.ServiceUnavailableError): + await _stream_auto_execute(logging_obj) + await _settle(lambda: len(recorder.logged) > 0 and len(recorder.failed) > 0) + + _assert_one_row_bills(recorder.logged, _ROUND_ONE_BODY) + assert len(recorder.failed) == 1 + assert (recorder.failed[0]["prompt_tokens"], recorder.failed[0]["completion_tokens"]) == (149, 5) + assert recorder.failed[0]["response_cost"] == pytest.approx(_round_cost(_ROUND_TWO_BODY) - _FIXED_MARGIN) + _assert_breakdown_agrees_with_the_billed_cost(recorder.failed[0]) + failure_row: Final = _request_failure_row(logging_obj, recorder.logged[0]) + _assert_billed_tokens(failure_row, (149, 5)) + assert failure_row["response_cost"] == pytest.approx(_round_cost(_ROUND_TWO_BODY) - _FIXED_MARGIN) + + +_XAI_MODEL: Final = "xai/grok-4.7" +_XAI_RESPONSES_URL: Final = r"https://api\.x\.ai/v1/responses" + + +def _xai_round(body: Mapping[str, object], cost_in_usd_ticks: int) -> Mapping[str, object]: + usage: Final = cast(Mapping[str, object], body["usage"]) + stated: Final = {**usage, "cost_in_usd_ticks": cost_in_usd_ticks, "context_details": {"cached_tokens": 3}} + return MappingProxyType({**body, "model": "grok-4.7", "usage": stated}) + + +def _xai_round_cost(body: Mapping[str, object]) -> float: + usage: Final = cast(Mapping[str, object], body["usage"]) + stated: Final = {**usage, "cost": xai_reported_cost_in_usd(usage["cost_in_usd_ticks"])} + return litellm.completion_cost( + completion_response=ResponsesAPIResponse.model_validate({**body, "usage": stated}), + model=_XAI_MODEL, + custom_llm_provider="xai", + call_type="aresponses", + ) + + +@pytest.mark.asyncio +async def test_xai_stream_auto_execute_bills_each_round_at_its_own_price_with_one_fixed_margin( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "cost_margin_config", {"xai": {"fixed_amount": _FIXED_MARGIN}}) + rounds: Final = (_xai_round(_ROUND_ONE_BODY, 12_400_000), _xai_round(_ROUND_TWO_BODY, 8_800_000)) + _serve_responses(respx_mock, monkeypatch, tuple(_streamed(body) for body in rounds), _XAI_RESPONSES_URL) + recorder: Final = _RoundRecorder() + + with _echo_tool_served("PING"): + events: Final = await _drain( + await litellm.aresponses( + model=_XAI_MODEL, + input="Call echo with PING", + tools=[dict(_ECHO_MCP_TOOL)], + api_key="xai-mcp-rounds", + stream=True, + litellm_logging_obj=_rounds_logging_obj(recorder, "mcp-xai-stream-rounds"), + ) + ) + await _settle(lambda: len(recorder.logged) > 0) + + charged_once: Final = _xai_round_cost(rounds[0]) + _xai_round_cost(rounds[1]) - _FIXED_MARGIN + client: Final = _final_response(events) + assert len(recorder.logged) == 1 + assert recorder.logged[0]["response_cost"] == pytest.approx(charged_once) + assert recorder.logged[0]["cost_breakdown"]["total_cost"] == pytest.approx(charged_once) + assert client.usage is not None + assert client.usage.cost == pytest.approx(charged_once) + assert json.loads(events[-1].model_dump_json())["response"]["usage"]["context_details"] == {"cached_tokens": 6} + + +_UNPRICED_MODEL: Final = "litellm-unpriced-mcp-rounds" + + +@pytest.mark.asyncio +async def test_stream_auto_execute_on_an_unpriced_model_logs_both_rounds_as_a_failed_cost_lookup_not_a_free_call( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + rounds: Final = tuple( + MappingProxyType({**body, "model": _UNPRICED_MODEL}) for body in (_ROUND_ONE_BODY, _ROUND_TWO_BODY) + ) + _serve_rounds(respx_mock, monkeypatch, rounds) + recorder: Final = _RoundRecorder() + + with _echo_tool_served("PING"): + events: Final = await _drain( + await litellm.aresponses( + model=f"openai/{_UNPRICED_MODEL}", + input="Call echo with PING", + tools=[dict(_ECHO_MCP_TOOL)], + api_key="sk-mcp-rounds", + api_base="https://mcp-rounds.example.com/v1", + stream=True, + litellm_logging_obj=_rounds_logging_obj(recorder, "mcp-stream-unpriced-rounds"), + ) + ) + await _settle(lambda: len(recorder.logged) > 0) + + client: Final = _final_response(events) + assert len(recorder.logged) == 1 + row: Final = recorder.logged[0] + assert (row["prompt_tokens"], row["completion_tokens"]) == (71 + 149, 19 + 5) + assert row["response_cost_failure_debug_info"] is not None + assert client.usage is not None + assert client.usage.cost is None + + +async def _stop_reading_once_the_follow_up_answers(logging_obj: Logging) -> None: + stream: Final = await _open_stream_auto_execute(logging_obj) + async for event in cast(Any, stream): + if _event_type(event) == "response.output_item.added" and _output_types((event.item,)) == ("message",): + return + + +@pytest.mark.asyncio +async def test_stream_auto_execute_bills_the_finished_first_round_when_the_caller_stops_reading_the_follow_up( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + _serve_responses(respx_mock, monkeypatch, (_streamed(_ROUND_ONE_BODY), _streamed(_ROUND_TWO_BODY))) + recorder: Final = _RoundRecorder() + + with _echo_tool_served("PING"): + await _stop_reading_once_the_follow_up_answers(_rounds_logging_obj(recorder, "mcp-stream-caller-stops")) + await _settle(lambda: len(recorder.logged) > 0) + + _assert_one_row_bills(recorder.logged, _ROUND_ONE_BODY) + + +async def _stream_read_until_the_follow_up_answers(logging_obj: Logging) -> object: + stream: Final = await _open_stream_auto_execute(logging_obj) + async for event in cast(Any, stream): + if _event_type(event) == "response.output_item.added" and _output_types((event.item,)) == ("message",): + return stream + pytest.fail("the follow-up round never answered") + + +_STREAM_BILLING_NAMES: Final = ("StreamedRoundBilling", "dispatch_success_handlers", "StreamingIterator") + + +def _names_stream_billing(text: str) -> bool: + return any(name in text for name in _STREAM_BILLING_NAMES) + + +class _UnraisableRecorder: + def __init__(self) -> None: + self.raised: tuple[str, ...] = () + + def __call__(self, unraisable: "sys.UnraisableHookArgs") -> None: + self.raised = (*self.raised, f"{unraisable.object!r}: {unraisable.exc_value!r}") + + +def test_stream_auto_execute_dropped_after_its_event_loop_closed_discards_its_parked_log_cleanly( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + _serve_responses(respx_mock, monkeypatch, (_streamed(_ROUND_ONE_BODY), _streamed(_ROUND_TWO_BODY))) + recorder: Final = _RoundRecorder() + unraisable: Final = _UnraisableRecorder() + gc.collect() + monkeypatch.setattr(sys, "unraisablehook", unraisable) + + with _echo_tool_served("PING"), warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + asyncio.run(_stream_read_until_the_follow_up_answers(_rounds_logging_obj(recorder, "mcp-stream-loop-closed"))) + gc.collect() + + assert recorder.logged == () + assert [raised for raised in unraisable.raised if _names_stream_billing(raised)] == [] + assert [str(warning.message) for warning in caught if _names_stream_billing(str(warning.message))] == [] + + +_BRIDGED_MODEL: Final = "anthropic/claude-haiku-5-5" +_ANTHROPIC_MESSAGES_URL: Final = r"https://api\.anthropic\.com/v1/messages" + + +def _bridged_round(content: bytes) -> httpx.Response: + return httpx.Response(200, content=content, headers=SSE_HEADERS) + + +def _bridged_echo_call_round(cache_usage: Mapping[str, object] = MappingProxyType({})) -> httpx.Response: + return _bridged_round( + _anthropic_round( + "msg_echo_call", + {"type": "tool_use", "id": "call_echo", "name": "echo", "input": {}}, + {"type": "input_json_delta", "partial_json": json.dumps({"value": "PING"})}, + "tool_use", + model="claude-haiku-5-5", + usage=(533, 40), + cache_usage=cache_usage, + ) + ) + + +def _bridged_echo_answer_round( + stop: str = "end_turn", cache_usage: Mapping[str, object] = MappingProxyType({}) +) -> httpx.Response: + return _bridged_round( + _anthropic_round( + "msg_echo_answer", + {"type": "text", "text": ""}, + {"type": "text_delta", "text": "PING"}, + stop, + model="claude-haiku-5-5", + usage=(470, 4), + cache_usage=cache_usage, + ) + ) + + +def _anthropic_prompt_cache_usage(written: int, read: int) -> Mapping[str, object]: + return MappingProxyType( + { + "cache_creation_input_tokens": written, + "cache_read_input_tokens": read, + "cache_creation": {"ephemeral_5m_input_tokens": written, "ephemeral_1h_input_tokens": 0}, + "inference_geo": "us", + "speed": "standard", + } + ) + + +class _CutAfterTheMessageStart(httpx.AsyncByteStream): + async def __aiter__(self): + yield _sse((_anthropic_message_start("msg_echo_answer", "claude-haiku-5-5", 470),)) + raise httpx.ReadError("connection reset mid-stream") + + +def _bridged_round_cost(prompt_tokens: int, completion_tokens: int) -> float: + return litellm.completion_cost( + completion_response=ModelResponse( + model="claude-haiku-5-5", + usage=Usage( + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=prompt_tokens + completion_tokens, + ), + ), + model=_BRIDGED_MODEL, + ) + + +async def _open_bridged_stream_auto_execute(logging_obj: Logging) -> object: + return await litellm.aresponses( + model=_BRIDGED_MODEL, + input="Call echo with PING", + tools=[dict(_ECHO_MCP_TOOL)], + api_key="sk-ant-mcp-rounds", + stream=True, + litellm_logging_obj=logging_obj, + ) + + +def _assert_logged_as_the_final_chat_answer_with_both_rounds_usage(recorder: _RoundRecorder, answer: str) -> None: + assert len(recorder.results) == 1 + logged: Final = recorder.results[0] + assert isinstance(logged, ModelResponse) + assert logged.choices[0].message.content == answer + usage: Final[object] = getattr(logged, "usage", None) + assert isinstance(usage, Usage) + assert (usage.prompt_tokens, usage.completion_tokens) == (533 + 470, 40 + 4) + + +def _assert_one_bridged_row_bills_both_rounds_once(recorder: _RoundRecorder, client: ResponsesAPIResponse) -> None: + charged_once: Final = _bridged_round_cost(533, 40) + _bridged_round_cost(470, 4) - _FIXED_MARGIN + assert len(recorder.logged) == 1 + row: Final = recorder.logged[0] + _assert_logged_as_the_final_chat_answer_with_both_rounds_usage(recorder, "PING") + assert row["id"] == client.id + assert "message" in _output_types(client.output) + assert (row["prompt_tokens"], row["completion_tokens"]) == (533 + 470, 40 + 4) + assert row["cache_hit"] is False + assert row["response_cost"] == pytest.approx(charged_once) + assert row["cost_breakdown"]["total_cost"] == pytest.approx(charged_once) + assert client.usage is not None + assert (client.usage.input_tokens, client.usage.output_tokens) == (533 + 470, 40 + 4) + assert client.usage.cost == pytest.approx(charged_once) + + +@pytest.mark.asyncio +async def test_bridged_stream_auto_execute_logs_the_final_response_billed_for_both_rounds_with_one_fixed_margin( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "cost_margin_config", {"anthropic": {"fixed_amount": _FIXED_MARGIN}}) + _serve_responses( + respx_mock, monkeypatch, (_bridged_echo_call_round(), _bridged_echo_answer_round()), _ANTHROPIC_MESSAGES_URL + ) + recorder: Final = _RoundRecorder() + + with _echo_tool_served("PING"): + events: Final = await _drain( + await _open_bridged_stream_auto_execute(_rounds_logging_obj(recorder, "mcp-bridged-stream-rounds")) + ) + await _settle(lambda: len(recorder.logged) > 0) + + assert _output_types(_final_response(events).output) == ("mcp_call", "message") + _assert_one_bridged_row_bills_both_rounds_once(recorder, _final_response(events)) + + +@pytest.mark.asyncio +async def test_bridged_stream_auto_execute_logs_the_provider_usage_of_both_rounds_summed( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "cost_margin_config", {"anthropic": {"fixed_amount": _FIXED_MARGIN}}) + _serve_responses( + respx_mock, + monkeypatch, + ( + _bridged_echo_call_round(_anthropic_prompt_cache_usage(written=10, read=5)), + _bridged_echo_answer_round(cache_usage=_anthropic_prompt_cache_usage(written=3, read=7)), + ), + _ANTHROPIC_MESSAGES_URL, + ) + recorder: Final = _RoundRecorder() + + with _echo_tool_served("PING"): + await _drain( + await _open_bridged_stream_auto_execute(_rounds_logging_obj(recorder, "mcp-bridged-provider-usage")) + ) + await _settle(lambda: len(recorder.logged) > 0) + + assert len(recorder.logged) == 1 + row: Final = recorder.logged[0] + usage: Final = row["metadata"]["usage_object"] + assert "choices" in row["response"] + assert (row["prompt_tokens"], row["completion_tokens"]) == (533 + 10 + 5 + 470 + 3 + 7, 40 + 4) + assert (usage["cache_creation_input_tokens"], usage["cache_read_input_tokens"]) == (10 + 3, 5 + 7) + assert usage["prompt_tokens_details"]["cache_creation_token_details"]["ephemeral_5m_input_tokens"] == 10 + 3 + assert (usage["inference_geo"], usage["speed"]) == ("us", "standard") + assert usage["cost"] == pytest.approx(row["response_cost"]) + + +@pytest.mark.asyncio +async def test_bridged_stream_auto_execute_bills_both_rounds_with_one_fixed_margin_when_the_final_round_is_incomplete( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "cost_margin_config", {"anthropic": {"fixed_amount": _FIXED_MARGIN}}) + _serve_responses( + respx_mock, + monkeypatch, + (_bridged_echo_call_round(), _bridged_echo_answer_round("max_tokens")), + _ANTHROPIC_MESSAGES_URL, + ) + recorder: Final = _RoundRecorder() + + with _echo_tool_served("PING"): + events: Final = await _drain( + await _open_bridged_stream_auto_execute(_rounds_logging_obj(recorder, "mcp-bridged-stream-incomplete")) + ) + await _settle(lambda: len(recorder.logged) > 0) + + _assert_one_bridged_row_bills_both_rounds_once(recorder, _final_response(events, "response.incomplete")) + + +@pytest.mark.asyncio +async def test_bridged_stream_auto_execute_deferred_for_a_post_call_guardrail_logs_both_rounds_once_on_release( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.utils import ProxyLogging + + monkeypatch.setattr(litellm, "cost_margin_config", {"anthropic": {"fixed_amount": _FIXED_MARGIN}}) + _serve_responses( + respx_mock, monkeypatch, (_bridged_echo_call_round(), _bridged_echo_answer_round()), _ANTHROPIC_MESSAGES_URL + ) + recorder: Final = _RoundRecorder() + logging_obj: Final = _rounds_logging_obj(recorder, "mcp-bridged-stream-rounds-deferred") + + with _echo_tool_served("PING"): + stream: Final = await _open_bridged_stream_auto_execute(logging_obj) + ProxyBaseLLMRequestProcessing(data={})._arm_deferred_stream_dispatch( + response=stream, route_type="aresponses", user_api_key_dict=MagicMock(), logging_obj=logging_obj + ) + events: Final = await _drain(stream) + await asyncio.sleep(0.2) + logged_before_release: Final = recorder.logged + ProxyLogging.fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) + await _settle(lambda: len(recorder.logged) > 0) + + assert logged_before_release == () + _assert_one_bridged_row_bills_both_rounds_once(recorder, _final_response(events)) + + +@pytest.mark.asyncio +async def test_bridged_stream_auto_execute_blocked_by_a_post_call_guardrail_bills_both_rounds_on_the_failure_row( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.utils import ProxyLogging + + monkeypatch.setattr(litellm, "cost_margin_config", {"anthropic": {"fixed_amount": _FIXED_MARGIN}}) + _serve_responses( + respx_mock, monkeypatch, (_bridged_echo_call_round(), _bridged_echo_answer_round()), _ANTHROPIC_MESSAGES_URL + ) + recorder: Final = _RoundRecorder() + logging_obj: Final = _rounds_logging_obj(recorder, "mcp-bridged-stream-rounds-blocked") + blocked: Final = GuardrailRaisedException(guardrail_name="output-policy", message="blocked") + + with _echo_tool_served("PING"): + stream: Final = await _open_bridged_stream_auto_execute(logging_obj) + ProxyBaseLLMRequestProcessing(data={})._arm_deferred_stream_dispatch( + response=stream, route_type="aresponses", user_api_key_dict=MagicMock(), logging_obj=logging_obj + ) + await _drain(stream) + ProxyLogging._discard_deferred_stream_logging_for_failure({"litellm_logging_obj": logging_obj}, blocked) + ProxyLogging.fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) + await logging_obj.async_failure_handler(blocked, "blocked by output-policy") + await _settle(lambda: len(recorder.failed) > 0) + + charged_once: Final = _bridged_round_cost(533, 40) + _bridged_round_cost(470, 4) - _FIXED_MARGIN + assert recorder.logged == () + assert len(recorder.failed) == 1 + assert (recorder.failed[0]["prompt_tokens"], recorder.failed[0]["completion_tokens"]) == (533 + 470, 40 + 4) + assert recorder.failed[0]["response_cost"] == pytest.approx(charged_once) + assert recorder.failed[0]["cost_breakdown"]["total_cost"] == pytest.approx(charged_once) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stop", ("end_turn", "max_tokens"), ids=("completed", "incomplete")) +async def test_bridged_stream_auto_execute_bills_both_rounds_when_the_caller_stops_reading_after_the_answer_text( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, stop: str +): + _serve_responses( + respx_mock, monkeypatch, (_bridged_echo_call_round(), _bridged_echo_answer_round(stop)), _ANTHROPIC_MESSAGES_URL + ) + recorder: Final = _RoundRecorder() + + async def stop_reading_once_the_answer_text_is_done() -> None: + stream: Final = await _open_bridged_stream_auto_execute( + _rounds_logging_obj(recorder, "mcp-bridged-stream-caller-stops") + ) + async for event in cast(Any, stream): + if _event_type(event) == "response.output_text.done": + return + + with _echo_tool_served("PING"): + await stop_reading_once_the_answer_text_is_done() + await _settle(lambda: len(recorder.logged) > 0) + + assert len(recorder.logged) == 1 + _assert_logged_as_the_final_chat_answer_with_both_rounds_usage(recorder, "PING") + assert (recorder.logged[0]["prompt_tokens"], recorder.logged[0]["completion_tokens"]) == (533 + 470, 40 + 4) + assert recorder.logged[0]["response_cost"] == pytest.approx( + _bridged_round_cost(533, 40) + _bridged_round_cost(470, 4) + ) + + +@pytest.mark.asyncio +async def test_bridged_stream_auto_execute_bills_a_broken_follow_up_on_its_own_failure_row_with_no_second_fixed_margin( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "cost_margin_config", {"anthropic": {"fixed_amount": _FIXED_MARGIN}}) + _serve_responses( + respx_mock, + monkeypatch, + (_bridged_echo_call_round(), httpx.Response(200, stream=_CutAfterTheMessageStart(), headers=SSE_HEADERS)), + _ANTHROPIC_MESSAGES_URL, + ) + recorder: Final = _RoundRecorder() + logging_obj: Final = _rounds_logging_obj(recorder, "mcp-bridged-stream-breaks") + + with _echo_tool_served("PING"), pytest.raises(MidStreamFallbackError): + await _drain(await _open_bridged_stream_auto_execute(logging_obj)) + await _settle(lambda: len(recorder.logged) > 0 and len(recorder.failed) > 0) + + assert len(recorder.logged) == 1 + assert (recorder.logged[0]["prompt_tokens"], recorder.logged[0]["completion_tokens"]) == (533, 40) + assert recorder.logged[0]["response_cost"] == pytest.approx(_bridged_round_cost(533, 40)) + assert len(recorder.failed) == 1 + assert (recorder.failed[0]["prompt_tokens"], recorder.failed[0]["completion_tokens"]) == (470, 0) + assert recorder.failed[0]["response_cost"] == pytest.approx(_bridged_round_cost(470, 0) - _FIXED_MARGIN) + _assert_breakdown_agrees_with_the_billed_cost(recorder.failed[0]) + failure_row: Final = _request_failure_row(logging_obj, recorder.logged[0]) + _assert_billed_tokens(failure_row, (470, 0)) + assert failure_row["response_cost"] == pytest.approx(_bridged_round_cost(470, 0) - _FIXED_MARGIN) + + +@pytest.mark.asyncio +async def test_bridged_stream_auto_execute_bills_a_broken_follow_up_its_own_fixed_margin_after_a_cached_first_round( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "cost_margin_config", {"anthropic": {"fixed_amount": _FIXED_MARGIN}}) + monkeypatch.setattr(litellm, "cache", litellm.Cache(type="local")) + _serve_responses( + respx_mock, + monkeypatch, + ( + _bridged_echo_call_round(), + _bridged_echo_answer_round(), + httpx.Response(200, stream=_CutAfterTheMessageStart(), headers=SSE_HEADERS), + ), + _ANTHROPIC_MESSAGES_URL, + ) + warm: Final = _RoundRecorder() + with _echo_tool_served("PING"): + await _drain( + await _open_bridged_stream_auto_execute(_rounds_logging_obj(warm, "mcp-bridged-cache-then-break-warm")) + ) + await _settle(lambda: len(warm.logged) > 0) + + recorder: Final = _RoundRecorder() + with _echo_tool_served("PONG"), pytest.raises(MidStreamFallbackError): + await _drain( + await _open_bridged_stream_auto_execute(_rounds_logging_obj(recorder, "mcp-bridged-cache-then-break")) + ) + await _settle(lambda: len(recorder.logged) > 0 and len(recorder.failed) > 0) + + assert [row["response_cost"] for row in recorder.logged] == [0.0] + assert len(recorder.failed) == 1 + assert recorder.failed[0]["response_cost"] == pytest.approx(_bridged_round_cost(470, 0)) + + +@pytest.mark.asyncio +async def test_bridged_stream_auto_execute_answered_without_a_tool_call_logs_its_one_round( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + _serve_responses(respx_mock, monkeypatch, (_bridged_echo_answer_round(),), _ANTHROPIC_MESSAGES_URL) + recorder: Final = _RoundRecorder() + + with _echo_tool_served("PING"): + await _drain(await _open_bridged_stream_auto_execute(_rounds_logging_obj(recorder, "mcp-bridged-one-round"))) + await _settle(lambda: len(recorder.logged) > 0) + + assert len(recorder.logged) == 1 + assert "choices" in recorder.logged[0]["response"] + assert (recorder.logged[0]["prompt_tokens"], recorder.logged[0]["completion_tokens"]) == (470, 4) + assert recorder.logged[0]["response_cost"] == pytest.approx(_bridged_round_cost(470, 4)) + + +@pytest.mark.asyncio +async def test_bridged_stream_auto_execute_bills_nothing_for_a_round_served_from_the_response_cache( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "cache", litellm.Cache(type="local")) + provider: Final = _serve_responses( + respx_mock, + monkeypatch, + (_bridged_echo_call_round(), _bridged_echo_answer_round(), _bridged_echo_answer_round()), + _ANTHROPIC_MESSAGES_URL, + ) + warm: Final = _RoundRecorder() + with _echo_tool_served("PING"): + await _drain(await _open_bridged_stream_auto_execute(_rounds_logging_obj(warm, "mcp-bridged-cache-warm"))) + await _settle(lambda: len(warm.logged) > 0) + + first_cached: Final = _RoundRecorder() + with _echo_tool_served("PONG"): + events: Final = await _drain( + await _open_bridged_stream_auto_execute(_rounds_logging_obj(first_cached, "mcp-bridged-cache-first-round")) + ) + await _settle(lambda: len(first_cached.logged) > 0) + + both_cached: Final = _RoundRecorder() + with _echo_tool_served("PING"): + await _drain( + await _open_bridged_stream_auto_execute(_rounds_logging_obj(both_cached, "mcp-bridged-cache-both-rounds")) + ) + await _settle(lambda: len(both_cached.logged) > 0) + + client: Final = _final_response(events) + assert provider.call_count == 3 + assert len(first_cached.logged) == 1 + assert (first_cached.logged[0]["prompt_tokens"], first_cached.logged[0]["completion_tokens"]) == (533 + 470, 40 + 4) + assert first_cached.logged[0]["response_cost"] == pytest.approx(_bridged_round_cost(470, 4)) + assert first_cached.logged[0]["cache_hit"] is False + assert [(row["response_cost"], row["cache_hit"]) for row in both_cached.logged] == [(0.0, True)] + assert client.usage is not None + assert client.usage.cost == pytest.approx(_bridged_round_cost(470, 4)) 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/responses/test_streaming_iterator_error_events.py b/tests/unit/responses/test_streaming_iterator_error_events.py index 04d2b9ddf81..7445070bd3a 100644 --- a/tests/unit/responses/test_streaming_iterator_error_events.py +++ b/tests/unit/responses/test_streaming_iterator_error_events.py @@ -23,7 +23,8 @@ import asyncio import json from datetime import datetime from importlib import import_module -from typing import Final +from collections.abc import Mapping +from typing import Final, cast from unittest.mock import Mock, patch import httpx @@ -31,6 +32,7 @@ import pytest import litellm 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.get_litellm_params import get_litellm_params from litellm.litellm_core_utils.litellm_logging import _get_provider_request_id @@ -44,6 +46,7 @@ from litellm.responses.streaming_iterator import ( _status_code_for_error_fields, ) from litellm.utils import ProviderConfigManager +from litellm.types.utils import StandardLoggingPayload from litellm.types.llms.openai import ( ErrorEvent, ErrorEventError, @@ -523,6 +526,82 @@ def test_handle_logging_failed_response_records_usage_and_cost(): iterator.logging_obj.response_cost_calculator.assert_called_once_with(result=chunk.response) +class _FailureRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.failed: tuple[StandardLoggingPayload, ...] = () + + async def async_log_failure_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + self.failed = (*self.failed, cast(StandardLoggingPayload, kwargs["standard_logging_object"])) + + +async def _first_failure_row(recorder: _FailureRecorder) -> StandardLoggingPayload: + for _ in range(200): + if recorder.failed: + return recorder.failed[0] + await asyncio.sleep(0.01) + pytest.fail("the failed stream logged no failure row") + + +@pytest.mark.asyncio +async def test_failed_stream_bills_only_its_own_usage_over_an_earlier_attempts_stash(): + recorder: Final = _FailureRecorder() + logging_obj: Final = LiteLLMLoggingObj( + model="gpt-5.6", + messages=[], + stream=True, + call_type="aresponses", + start_time=datetime(2026, 1, 1), + litellm_call_id="responses-failed-over-stash", + function_id="responses-failed-over-stash", + dynamic_async_failure_callbacks=[recorder], + ) + logging_obj.update_environment_variables( + model="gpt-5.6", optional_params={}, litellm_params=get_litellm_params(custom_llm_provider="openai") + ) + logging_obj.model_call_details["combined_usage_object"] = litellm.Usage( + prompt_tokens=100, completion_tokens=50, total_tokens=150 + ) + logging_obj.model_call_details["response_cost"] = 0.5 + failed: Final = { + "type": "response.failed", + "sequence_number": 1, + "response": { + "id": "resp_failed_over_stash", + "object": "response", + "created_at": 1790000000, + "status": "failed", + "model": "gpt-5.6", + "output": [], + "error": {"code": "server_error", "message": "boom"}, + "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + }, + } + iterator: Final = ResponsesAPIStreamingIterator( + response=httpx.Response( + 200, headers={"content-type": "text/event-stream"}, content=f"data: {json.dumps(failed)}\n\n".encode() + ), + model="gpt-5.6", + responses_api_provider_config=ProviderConfigManager.get_provider_responses_api_config( + provider="openai", model="gpt-5.6" + ), + logging_obj=logging_obj, + custom_llm_provider="openai", + ) + + with pytest.raises(MidStreamFallbackError): + await anext(aiter(iterator)) + + row: Final = await _first_failure_row(recorder) + prompt_cost, completion_cost = litellm.cost_per_token( + model="gpt-5.6", prompt_tokens=10, completion_tokens=5, custom_llm_provider="openai" + ) + assert (row["prompt_tokens"], row["completion_tokens"]) == (10, 5) + assert row["response_cost"] == pytest.approx(prompt_cost + completion_cost) + + def test_handle_logging_failed_response_without_usage_skips_recording(): iterator = _make_iterator() iterator.completed_response = _make_failed_chunk( diff --git a/tests/unit/router_strategy/adaptive_router/test_hooks.py b/tests/unit/router_strategy/adaptive_router/test_hooks.py index ad61f43c5a0..0ed22658144 100644 --- a/tests/unit/router_strategy/adaptive_router/test_hooks.py +++ b/tests/unit/router_strategy/adaptive_router/test_hooks.py @@ -10,6 +10,7 @@ from litellm.router_strategy.adaptive_router.config import ( ) from litellm.router_strategy.adaptive_router.hooks import ( AdaptiveRouterPostCallHook, + AdaptiveRouterResponseHeaderHook, _recent_tool_results, _resolve_session_key, ) @@ -337,8 +338,7 @@ async def test_post_call_response_headers_hook_returns_chosen_model_header(): """The header hook returns the `x-litellm-adaptive-router-model` header so proxy header construction picks it up (works for both streaming and non-streaming; `async_post_call_success_hook` is too late for streaming).""" - hook = _make_hook() - headers = await hook.async_post_call_response_headers_hook( + headers = await AdaptiveRouterResponseHeaderHook().async_post_call_response_headers_hook( data={"metadata": {"adaptive_router_chosen_model": "smart"}}, user_api_key_dict=MagicMock(), response=MagicMock(), @@ -348,8 +348,7 @@ async def test_post_call_response_headers_hook_returns_chosen_model_header(): @pytest.mark.asyncio async def test_post_call_response_headers_hook_noop_when_metadata_missing_key(): - hook = _make_hook() - headers = await hook.async_post_call_response_headers_hook( + headers = await AdaptiveRouterResponseHeaderHook().async_post_call_response_headers_hook( data={"metadata": {"litellm_session_id": "sess-A"}}, user_api_key_dict=MagicMock(), response=MagicMock(), @@ -359,8 +358,7 @@ async def test_post_call_response_headers_hook_noop_when_metadata_missing_key(): @pytest.mark.asyncio async def test_post_call_response_headers_hook_noop_when_no_metadata(): - hook = _make_hook() - headers = await hook.async_post_call_response_headers_hook( + headers = await AdaptiveRouterResponseHeaderHook().async_post_call_response_headers_hook( data={}, user_api_key_dict=MagicMock(), response=MagicMock(), @@ -370,8 +368,7 @@ async def test_post_call_response_headers_hook_noop_when_no_metadata(): @pytest.mark.asyncio async def test_post_call_response_headers_hook_noop_when_metadata_not_dict(): - hook = _make_hook() - headers = await hook.async_post_call_response_headers_hook( + headers = await AdaptiveRouterResponseHeaderHook().async_post_call_response_headers_hook( data={"metadata": "not-a-dict"}, user_api_key_dict=MagicMock(), response=MagicMock(), diff --git a/tests/unit/router_strategy/test_base_routing_strategy.py b/tests/unit/router_strategy/test_base_routing_strategy.py index dc75c3d2916..19ff11dd3f6 100644 --- a/tests/unit/router_strategy/test_base_routing_strategy.py +++ b/tests/unit/router_strategy/test_base_routing_strategy.py @@ -1,15 +1,20 @@ import json import logging -from typing import Any, Dict, List, Optional, Set, Union +from typing import Any, Dict, Final, List, Optional, Set, Union import pytest import asyncio +import threading from unittest.mock import MagicMock, patch +from fakeredis import TcpFakeServer + from litellm.caching.caching import DualCache +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.caching.redis_cache import RedisCache from litellm.caching.redis_cache import RedisCircuitBreakerOpenError, RedisPipelineIncrementOperation from litellm.router_strategy.base_routing_strategy import BaseRoutingStrategy @@ -162,3 +167,24 @@ async def test_push_refused_by_the_open_circuit_breaker_is_not_logged_as_an_erro assert caplog.records == [] assert base_strategy.redis_increment_operation_queue == [] + + +@pytest.mark.asyncio +async def test_an_increment_counted_after_retire_still_reaches_redis(): + server: Final = TcpFakeServer(("127.0.0.1", 0), server_type="redis") + threading.Thread(target=server.serve_forever, daemon=True).start() + redis_cache: Final = RedisCache(host="127.0.0.1", port=server.server_address[1]) + strategy: Final = BaseRoutingStrategy( + dual_cache=DualCache(redis_cache=redis_cache, in_memory_cache=InMemoryCache()), + should_batch_redis_writes=True, + default_sync_interval=60, + ) + try: + strategy.retire() + await strategy._increment_value_in_current_window("counted-after-retire", 3, 60) + await asyncio.sleep(0.2) + + assert float(await redis_cache.async_get_cache("counted-after-retire")) == 3 + finally: + server.shutdown() + server.server_close() diff --git a/tests/unit/router_strategy/test_budget_limiter_hotpath.py b/tests/unit/router_strategy/test_budget_limiter_hotpath.py index 4667ac67efa..10704a29347 100644 --- a/tests/unit/router_strategy/test_budget_limiter_hotpath.py +++ b/tests/unit/router_strategy/test_budget_limiter_hotpath.py @@ -1,4 +1,5 @@ import asyncio +import contextlib import gc import logging from types import SimpleNamespace @@ -492,6 +493,8 @@ def _new_router_budget_limiter( budget_limiter._redis_increment_queue_lock = queue_lock if queue_lock is not None else asyncio.Lock() budget_limiter._redis_increment_flush_lock = asyncio.Lock() budget_limiter._detached_increment_operations = None + budget_limiter._retired = False + budget_limiter._flush_after_retire = None return budget_limiter @@ -784,3 +787,57 @@ async def test_cancelled_flush_does_not_requeue_an_applied_batch(cancellations: await limiter._push_in_memory_increments_to_redis() assert redis_cache.values[_SPEND_KEY] == 10.0 assert limiter.redis_increment_operation_queue == [] + + +@pytest.mark.asyncio +async def test_spend_queued_when_a_limiter_is_retired_reaches_redis_after_a_failed_write() -> None: + pipeline_started: Final = asyncio.Event() + pipeline_completed: Final = asyncio.Event() + redis_cache: Final = _MockRedisCache( + initial_values={}, + pipeline_started=pipeline_started, + pipeline_completed=pipeline_completed, + should_fail_pipeline=True, + ) + budget_limiter: Final = _new_router_budget_limiter( + redis_cache=redis_cache, redis_increment_operation_queue=[_increment(10.0)] + ) + budget_limiter._sync_task = asyncio.create_task(asyncio.sleep(3600)) + + budget_limiter.retire() + await asyncio.wait_for(pipeline_started.wait(), timeout=1) + assert redis_cache.values == {} + redis_cache.should_fail_pipeline = False + await asyncio.wait_for(pipeline_completed.wait(), timeout=5) + + assert redis_cache.values == {_SPEND_KEY: 10.0} + assert budget_limiter.redis_increment_operation_queue == [] + + +@pytest.mark.asyncio +async def test_spend_in_flight_when_a_limiter_is_retired_reaches_redis_after_that_write_fails() -> None: + pipeline_started: Final = asyncio.Event() + allow_pipeline_to_complete: Final = asyncio.Event() + pipeline_completed: Final = asyncio.Event() + redis_cache: Final = _MockRedisCache( + initial_values={}, + pipeline_started=pipeline_started, + allow_pipeline_to_complete=allow_pipeline_to_complete, + should_fail_pipeline=True, + pipeline_completed=pipeline_completed, + ) + budget_limiter: Final = _new_router_budget_limiter( + redis_cache=redis_cache, redis_increment_operation_queue=[_increment(10.0)] + ) + budget_limiter._sync_task = asyncio.create_task(budget_limiter.periodic_sync_in_memory_spend_with_redis()) + await asyncio.wait_for(pipeline_started.wait(), timeout=1) + + budget_limiter.retire() + allow_pipeline_to_complete.set() + with contextlib.suppress(asyncio.CancelledError): + await budget_limiter._sync_task + redis_cache.should_fail_pipeline = False + await asyncio.wait_for(pipeline_completed.wait(), timeout=5) + + assert redis_cache.values == {_SPEND_KEY: 10.0} + assert budget_limiter.redis_increment_operation_queue == [] 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 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_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( diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index f5d13aecb6f..c3d3575b8b3 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -3868,6 +3868,9 @@ def test_stream_chunk_builder_text_completion_combines_text_and_usage(): model="gpt-3.5-turbo-instruct", choices=[{"text": " world", "index": 0, "logprobs": None, "finish_reason": "stop"}], ), + TextCompletionResponse( + id="cmpl-1", object="text_completion", created=1, model="gpt-3.5-turbo-instruct", choices=[] + ), ] response = stream_chunk_builder_text_completion(chunks=chunks, messages=[{"role": "user", "content": "say hello"}]) @@ -8853,7 +8856,8 @@ async def test_text_completion_stream_include_usage_returns_usage_on_final_chunk "stream": True, "stream_options": {"include_usage": True}, } - assert "".join(chunk.choices[0].text or "" for chunk in chunks) == "reply" + assert "".join(chunk.choices[0].text or "" for chunk in chunks if chunk.choices) == "reply" + assert chunks[-1].choices == [] final_usage: Final = chunks[-1].usage assert (final_usage.prompt_tokens, final_usage.completion_tokens, final_usage.total_tokens) == (3, 2, 5) diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 13ac55fa937..9cc1bf6ce79 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 @@ -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), @@ -16726,6 +16749,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 +16760,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 +18004,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 +18071,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 @@ -29670,6 +29741,73 @@ def test_pre_call_checks_ignores_an_image_response_format_string(monkeypatch: py assert [d["model_info"]["id"] for d in filtered] == ["only"] +@pytest.mark.parametrize( + ("deployment_stream_options", "default_stream_options", "request_stream_options", "usage_chunk_choice_count"), + [ + ({"include_usage": True}, None, None, 1), + (None, {"include_usage": True}, None, 1), + ({"include_usage": True}, None, {"include_usage": True}, 0), + (None, None, {"include_usage": True}, 0), + ], + ids=["deployment-config", "router-default", "request-over-deployment", "request-only"], +) +@pytest.mark.asyncio +async def test_router_configured_stream_usage_keeps_a_choice_on_the_usage_chunk( + deployment_stream_options: dict[str, bool] | None, + default_stream_options: dict[str, bool] | None, + request_stream_options: dict[str, bool] | None, + usage_chunk_choice_count: int, +): + deployment_params: Final = {"model": "openai/gpt-5-mini", "api_key": "fake", "mock_response": "hi there"} + router: Final = Router( + model_list=[ + { + "model_name": "m", + "litellm_params": ( + deployment_params + if deployment_stream_options is None + else {**deployment_params, "stream_options": deployment_stream_options} + ), + } + ], + default_litellm_params=None if default_stream_options is None else {"stream_options": default_stream_options}, + ) + request_kwargs: Final = {} if request_stream_options is None else {"stream_options": request_stream_options} + + response: Final = await router.acompletion( + model="m", messages=[{"role": "user", "content": "hi"}], stream=True, **request_kwargs + ) + chunks: Final = [chunk async for chunk in response] + + assert getattr(chunks[-1], "usage", None) is not None + assert len(chunks[-1].choices) == usage_chunk_choice_count + assert all(len(chunk.choices) == 1 for chunk in chunks[:-1]) + + +@pytest.mark.parametrize( + ("stream_options", "marked_injected"), + [({"include_usage": True}, True), (None, False)], + ids=["include-usage", "no-stream-options"], +) +def test_update_kwargs_with_deployment_reads_stream_options_from_pydantic_litellm_params( + stream_options: dict[str, bool] | None, marked_injected: bool +): + from litellm.types.litellm_params import STREAM_USAGE_INJECTED_FIELD + from litellm.types.router import LiteLLM_Params + + kwargs: Final[dict[str, object]] = {"stream": True, "metadata": {}} + Router()._update_kwargs_with_deployment( + deployment={ + "model_name": "gpt-5-mini", + "litellm_params": LiteLLM_Params(model="openai/gpt-5-mini", api_key="fake-key", stream_options=stream_options), + "model_info": {"id": "pydantic-params"}, + }, + kwargs=kwargs, + ) + + assert (kwargs.get(STREAM_USAGE_INJECTED_FIELD) is True) is marked_injected + + @pytest.mark.asyncio async def test_responses_fallback_wrapper_disconnect_billing_keeps_the_primary_streams_tokens(): """After a mid-stream hop, a disconnect bills the primary's relayed text plus the fallback's, not the fallback alone.""" diff --git a/tests/unit/test_router/test_router_helper_utils.py b/tests/unit/test_router/test_router_helper_utils.py index 7088818aa01..af8e7f53462 100644 --- a/tests/unit/test_router/test_router_helper_utils.py +++ b/tests/unit/test_router/test_router_helper_utils.py @@ -1,8 +1,12 @@ +import asyncio import functools +import gc import json import os +import weakref from datetime import datetime, timezone -from collections.abc import Mapping +from collections.abc import Iterator, Mapping +from types import MethodType from typing import Final from unittest.mock import AsyncMock, MagicMock, patch @@ -12,8 +16,12 @@ from litellm import Router from litellm.caching.dual_cache import DualCache from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS, ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY +from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting +from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper -from litellm.types.router import Deployment, DeploymentTypedDict, LiteLLM_Params, ModelInfo +from litellm.router_strategy.adaptive_router.hooks import AdaptiveRouterPostCallHook +from litellm.types.litellm_params import RoutingStrategyName +from litellm.types.router import AlertingConfig, Deployment, DeploymentTypedDict, LiteLLM_Params, ModelInfo from litellm.types.utils import ( CallTypes, ModelResponse, @@ -1852,6 +1860,315 @@ def test_discard(model_list): assert len(litellm.service_callback) == 0 +_CALLBACK_LIST_NAMES: Final = ( + "callbacks", + "input_callback", + "success_callback", + "failure_callback", + "_async_success_callback", + "_async_failure_callback", +) + + +def _callback_ids(list_name: str) -> list[int]: + return [id(callback) for callback in getattr(litellm, list_name)] + + +@pytest.mark.parametrize( + "routing_strategy, selector_attr", + [ + ("least-busy", "leastbusy_logger"), + ("usage-based-routing", "lowesttpm_logger"), + ("usage-based-routing-v2", "lowesttpm_logger_v2"), + ("latency-based-routing", "lowestlatency_logger"), + ("cost-based-routing", "lowestcost_logger"), + ], +) +@pytest.mark.asyncio +async def test_discard_releases_strategy_selector_across_router_lifecycles( + routing_strategy: RoutingStrategyName, selector_attr: str +): + def retire() -> weakref.ReferenceType[object]: + router: Final = Router(model_list=[], routing_strategy=routing_strategy) + selector: Final = getattr(router, selector_attr) + assert any(callback is selector for callback in litellm.callbacks) + router.discard() + return weakref.ref(selector) + + selector_refs: Final = tuple(retire() for _ in range(3)) + await asyncio.sleep(0) + gc.collect() + + assert all(getattr(litellm, list_name) == [] for list_name in _CALLBACK_LIST_NAMES) + assert all(selector_ref() is None for selector_ref in selector_refs) + + +@pytest.mark.parametrize( + "routing_strategy, selector_attr", + [ + ("least-busy", "leastbusy_logger"), + ("usage-based-routing", "lowesttpm_logger"), + ("usage-based-routing-v2", "lowesttpm_logger_v2"), + ("latency-based-routing", "lowestlatency_logger"), + ("cost-based-routing", "lowestcost_logger"), + ], +) +@pytest.mark.asyncio +async def test_discard_purges_selector_from_event_lists_after_serving_a_request( + routing_strategy: RoutingStrategyName, selector_attr: str +): + router: Final = Router( + model_list=[{"model_name": "gpt-5.5", "litellm_params": {"model": "openai/gpt-5.5", "api_key": "sk-test"}}], + routing_strategy=routing_strategy, + ) + selector: Final = getattr(router, selector_attr) + await router.acompletion(model="gpt-5.5", messages=[{"role": "user", "content": "hi"}], mock_response="ok") + assert any(callback is selector for callback in litellm.success_callback) + + router.discard() + + assert all(getattr(litellm, list_name) == [] for list_name in _CALLBACK_LIST_NAMES) + + +def test_discard_keeps_callbacks_owned_by_other_live_routers(): + survivor: Final = Router(model_list=[], routing_strategy="latency-based-routing") + unrelated: Final = CustomLogger() + litellm.logging_callback_manager.add_litellm_callback(unrelated) + before_ids: Final = {list_name: _callback_ids(list_name) for list_name in _CALLBACK_LIST_NAMES} + retired: Final = Router(model_list=[], routing_strategy="least-busy") + assert any(callback is retired.leastbusy_logger for callback in litellm.callbacks) + + retired.discard() + + assert {list_name: _callback_ids(list_name) for list_name in before_ids} == before_ids + assert any(callback is survivor.lowestlatency_logger for callback in litellm.callbacks) + survivor.discard() + + +@pytest.mark.parametrize( + "routing_strategy, selector_attr", + [ + ("usage-based-routing", "lowesttpm_logger"), + ("usage-based-routing-v2", "lowesttpm_logger_v2"), + ("latency-based-routing", "lowestlatency_logger"), + ("cost-based-routing", "lowestcost_logger"), + ], +) +@pytest.mark.asyncio +async def test_discard_leaves_same_strategy_survivor_receiving_events( + routing_strategy: RoutingStrategyName, selector_attr: str +): + model_list: Final[list[DeploymentTypedDict]] = [ + {"model_name": "gpt-5.5", "litellm_params": {"model": "openai/gpt-5.5", "api_key": "sk-test"}} + ] + retired: Final = Router(model_list=model_list, routing_strategy=routing_strategy) + survivor: Final = Router(model_list=model_list, routing_strategy=routing_strategy) + + retired.discard() + await survivor.acompletion(model="gpt-5.5", messages=[{"role": "user", "content": "hi"}], mock_response="ok") + + assert any(callback is getattr(survivor, selector_attr) for callback in litellm.success_callback) + survivor.discard() + + +def test_discard_drops_routing_group_selectors(): + router: Final = Router( + model_list=[ + {"model_name": "grouped-model", "litellm_params": {"model": "openai/gpt-5.5", "api_key": "sk-test"}} + ], + routing_groups=[ + {"group_name": "fast", "models": ["grouped-model"], "routing_strategy": "latency-based-routing"} + ], + ) + _, group_selector = router._get_routing_context("grouped-model") + assert any(callback is group_selector for callback in litellm.callbacks) + + router.discard() + + assert all(getattr(litellm, list_name) == [] for list_name in _CALLBACK_LIST_NAMES) + + +_GROUPED_MODEL_LIST: Final[list[DeploymentTypedDict]] = [ + {"model_name": "grouped-model", "litellm_params": {"model": "openai/gpt-5.5", "api_key": "sk-test"}} +] + + +class _LabelledHook(CustomLogger): + def __init__(self, label: str) -> None: + super().__init__() + self.label = label + + +def _visit_hooks_reloading_routing_groups_mid_loop( + router: Router, hooks: tuple[_LabelledHook, ...] +) -> Iterator[CustomLogger]: + for callback in litellm.callbacks: + if callback is hooks[0]: + router.update_settings( + routing_groups=[{"group_name": "fast", "models": ["grouped-model"], "routing_strategy": "least-busy"}] + ) + if callback in hooks: + yield callback + + +def test_routing_group_reload_keeps_every_hook_of_an_in_flight_callback_loop(): + router: Final = Router( + model_list=_GROUPED_MODEL_LIST, + routing_groups=[{"group_name": "fast", "models": ["grouped-model"], "routing_strategy": "least-busy"}], + ) + hooks: Final = tuple(_LabelledHook(label) for label in ("guardrail_a", "guardrail_b", "guardrail_c")) + for hook in hooks: + litellm.logging_callback_manager.add_litellm_callback(hook) + + visited: Final = tuple(_visit_hooks_reloading_routing_groups_mid_loop(router, hooks)) + + assert visited == hooks + router.discard() + + +def _routers_seen_in_success_loop_discarding_first(first: Router, second: Router) -> Iterator[Router]: + for callback in litellm.success_callback: + if not isinstance(callback, MethodType): + continue + if callback.__self__ is first: + first.discard() + yield first + if callback.__self__ is second: + yield second + + +def test_discard_mid_success_loop_keeps_the_next_router_callback(): + first: Final = Router(model_list=_GROUPED_MODEL_LIST) + second: Final = Router(model_list=_GROUPED_MODEL_LIST) + + seen: Final = tuple(_routers_seen_in_success_loop_discarding_first(first, second)) + + assert seen == (first, second) + second.discard() + + +@pytest.mark.asyncio +async def test_discard_purges_optional_pre_call_checks_from_event_lists_after_serving_a_request(): + router: Final = Router( + model_list=_GROUPED_MODEL_LIST, + optional_pre_call_checks=["enforce_model_rate_limits", "prompt_caching", "deployment_affinity"], + ) + await router.acompletion(model="grouped-model", messages=[{"role": "user", "content": "hi"}], mock_response="ok") + + router.discard() + + assert all(getattr(litellm, list_name) == [] for list_name in _CALLBACK_LIST_NAMES) + + +def test_discard_registers_same_config_survivor_optional_pre_call_checks(): + retired: Final = Router(model_list=_GROUPED_MODEL_LIST, optional_pre_call_checks=["enforce_model_rate_limits"]) + survivor: Final = Router(model_list=_GROUPED_MODEL_LIST, optional_pre_call_checks=["enforce_model_rate_limits"]) + + retired.discard() + + assert survivor.optional_callbacks + assert all(_is_registered(check) for check in survivor.optional_callbacks) + survivor.discard() + + +def _is_registered(callback: object) -> bool: + return any(registered is callback for registered in litellm.callbacks) + + +@pytest.mark.asyncio +async def test_discard_releases_per_request_override_selector(): + def build_override_then_discard() -> weakref.ReferenceType[object]: + router: Final = Router(model_list=_GROUPED_MODEL_LIST) + selector: Final = router._get_override_strategy_selector("usage-based-routing-v2") + assert selector is not None + router.discard() + return weakref.ref(selector) + + selector_ref: Final = build_override_then_discard() + await asyncio.sleep(0) + gc.collect() + + assert selector_ref() is None + + + +def test_discard_registers_the_selector_of_a_same_strategy_survivor_built_without_a_model_list(): + retired: Final = Router(model_list=[], routing_strategy="latency-based-routing") + survivor: Final = Router(routing_strategy="latency-based-routing") + assert not _is_registered(survivor.lowestlatency_logger) + + retired.discard() + + assert _is_registered(survivor.lowestlatency_logger) + survivor.discard() + + +def test_discard_drops_its_slack_alerting_callbacks(): + router: Final = Router(model_list=_GROUPED_MODEL_LIST, alerting_config=AlertingConfig(webhook_url="test")) + assert any(isinstance(callback, SlackAlerting) for callback in litellm.callbacks) + + router.discard() + + assert all(getattr(litellm, list_name) == [] for list_name in _CALLBACK_LIST_NAMES) + + +def _adaptive_router_model_list(name: str) -> list[DeploymentTypedDict]: + return [ + {"model_name": "gpt-5.5", "litellm_params": {"model": "openai/gpt-5.5", "api_key": "sk-test"}}, + { + "model_name": name, + "litellm_params": { + "model": "auto_router/adaptive_router", + "adaptive_router_config": {"available_models": ["gpt-5.5"]}, + }, + }, + ] + + +def _registered_adaptive_router_strategies() -> tuple[object, ...]: + return tuple( + callback.adaptive_router for callback in litellm.callbacks if isinstance(callback, AdaptiveRouterPostCallHook) + ) + + +def test_discard_drops_its_adaptive_router_hook(): + router: Final = Router(model_list=_adaptive_router_model_list("adaptive"), ignore_invalid_deployments=True) + assert _registered_adaptive_router_strategies() + + router.discard() + + assert _registered_adaptive_router_strategies() == () + + +def test_discard_registers_the_adaptive_router_hook_of_a_survivor(): + survivor: Final = Router(model_list=_adaptive_router_model_list("adaptive"), ignore_invalid_deployments=True) + retired: Final = Router(model_list=_adaptive_router_model_list("adaptive"), ignore_invalid_deployments=True) + + retired.discard() + + survivor_strategies: Final = tuple(tagged.strategy for tagged in survivor.adaptive_routers["adaptive"]) + assert _registered_adaptive_router_strategies() == survivor_strategies + survivor.discard() + + +@pytest.mark.asyncio +async def test_discard_releases_its_router_budget_limiter(): + def build_budget_limiter_then_discard() -> weakref.ReferenceType[object]: + router: Final = Router( + model_list=_GROUPED_MODEL_LIST, + provider_budget_config={"openai": {"budget_limit": 10.0, "time_period": "1d"}}, + ) + budget_limiter: Final = router.get_router_deployment_budget_limiter() + assert budget_limiter is not None + router.discard() + return weakref.ref(budget_limiter) + + budget_limiter_ref: Final = build_budget_limiter_then_discard() + await asyncio.sleep(0) + gc.collect() + + assert budget_limiter_ref() is None + def test_initialize_assistants_endpoint(model_list): """Test if the 'initialize_assistants_endpoint' function is working correctly""" router = Router(model_list=model_list) diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 9b6700e9d39..cdff03af3b4 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -3891,6 +3891,114 @@ class TestCallbackAsyncSyncSeparation: assert my_sync_cb not in litellm._async_input_callback +class _DropsACallbackFromAnotherThreadWhenComparedToDynamodb: + def __init__(self, dropped: Callable[..., object]) -> None: + self.dropped = dropped + self.drop_thread: threading.Thread | None = None + + def __eq__(self, other: object) -> bool: + if isinstance(other, str) and other == "dynamodb" and self.drop_thread is None: + self.drop_thread = threading.Thread( + target=litellm.logging_callback_manager.drop_callbacks, args=((self.dropped,),) + ) + self.drop_thread.start() + self.drop_thread.join(timeout=0.2) + return self is other + + __hash__ = object.__hash__ + + +def test_function_setup_moves_async_callbacks_while_another_thread_drops_a_callback( + monkeypatch: pytest.MonkeyPatch, +) -> None: + def dropped_callback(*args: object, **kwargs: object) -> None: + return None + + async def async_callback(*args: object, **kwargs: object) -> None: + return None + + drops_on_compare: Final = _DropsACallbackFromAnotherThreadWhenComparedToDynamodb(dropped_callback) + monkeypatch.setattr(litellm.utils, "callback_list", ["initialized"]) + for list_name in ( + "callbacks", + "input_callback", + "failure_callback", + "service_callback", + "_async_input_callback", + "_async_success_callback", + "_async_failure_callback", + ): + monkeypatch.setattr(litellm, list_name, []) + monkeypatch.setattr(litellm, "success_callback", [dropped_callback, drops_on_compare, async_callback]) + + litellm.utils.function_setup( + "completion", + litellm.utils.Rules(), + datetime.now(), + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + litellm_call_id="callback-drop-race", + ) + assert drops_on_compare.drop_thread is not None + drops_on_compare.drop_thread.join() + + assert litellm.success_callback == [drops_on_compare] + assert litellm._async_success_callback == [async_callback] + + +class _DropsACallbackFromAnotherThreadWhenComparedToSentry: + def __init__(self, dropped: Callable[..., object]) -> None: + self.dropped = dropped + self.drop_thread: threading.Thread | None = None + self.drop_finished_during_logger_setup = False + + def __eq__(self, other: object) -> bool: + if isinstance(other, str) and other == "sentry" and self.drop_thread is None: + self.drop_thread = threading.Thread( + target=litellm.logging_callback_manager.drop_callbacks, args=((self.dropped,),) + ) + self.drop_thread.start() + self.drop_thread.join(timeout=5) + self.drop_finished_during_logger_setup = not self.drop_thread.is_alive() + return self is other + + __hash__ = object.__hash__ + + +def test_function_setup_lets_other_threads_drop_callbacks_while_it_sets_up_loggers( + monkeypatch: pytest.MonkeyPatch, +) -> None: + def dropped_callback(*args: object, **kwargs: object) -> None: + return None + + drops_on_compare: Final = _DropsACallbackFromAnotherThreadWhenComparedToSentry(dropped_callback) + monkeypatch.setattr(litellm.utils, "callback_list", []) + for list_name in ( + "callbacks", + "input_callback", + "failure_callback", + "service_callback", + "_async_input_callback", + "_async_success_callback", + "_async_failure_callback", + ): + monkeypatch.setattr(litellm, list_name, []) + monkeypatch.setattr(litellm, "success_callback", [drops_on_compare]) + + litellm.utils.function_setup( + "completion", + litellm.utils.Rules(), + datetime.now(), + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + litellm_call_id="logger-setup-outside-lock", + ) + assert drops_on_compare.drop_thread is not None + drops_on_compare.drop_thread.join() + + assert drops_on_compare.drop_finished_during_logger_setup + + class TestMetadataNoneHandling: """ Test that metadata=None in kwargs doesn't cause TypeError. diff --git a/tests/unit/types/test_litellm_params.py b/tests/unit/types/test_litellm_params.py index 788d405f99f..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", @@ -263,6 +264,7 @@ INTERNAL_STATE_NAMES: Final = ( "litellm_trusted_callback_vars", "_litellm_addressed_response_id", "_litellm_strip_stream_usage", + "_litellm_stream_usage_injected", "client_side_timeout", "model_file_id_mapping", "acompletion", diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailDeleteModal.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailDeleteModal.test.tsx new file mode 100644 index 00000000000..0d74b714ee6 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailDeleteModal.test.tsx @@ -0,0 +1,78 @@ +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { GuardrailDeleteModal } from "./GuardrailDeleteModal"; +import { deleteGuardrailCall } from "@/components/networking"; +import { toast } from "@/lib/toast"; +import { fireEvent, renderWithProviders, screen, waitFor } from "@/../tests/test-utils"; + +vi.mock("@/components/networking", () => ({ deleteGuardrailCall: vi.fn() })); +vi.mock("@/lib/toast", () => ({ toast: { success: vi.fn(), fromError: vi.fn() } })); + +const GUARDRAIL = { + guardrail_id: "gr-1", + guardrail_name: "word-block", + litellm_params: { guardrail: "litellm_content_filter", mode: "pre_call", default_on: true }, +}; + +const renderModal = () => { + const onClose = vi.fn(); + const onDeleted = vi.fn(); + renderWithProviders( + , + ); + return { onClose, onDeleted }; +}; + +const confirmDelete = () => { + fireEvent.change(screen.getByPlaceholderText(GUARDRAIL.guardrail_name), { + target: { value: GUARDRAIL.guardrail_name }, + }); + fireEvent.click(screen.getByRole("button", { name: "Delete" })); +}; + +describe("GuardrailDeleteModal", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("warns about name references and only enables Delete once the exact name is typed", () => { + renderModal(); + + const warning = screen.getByText(/silently stop running it/); + expect(warning).toHaveTextContent("Keys, teams, and policies that list this guardrail"); + expect(warning).toHaveTextContent("Policy pipeline steps that use it will block every request with a 400 error"); + expect(warning).toHaveTextContent("unless the step's on_error is next or allow"); + const confirmButton = screen.getByRole("button", { name: "Delete" }); + expect(confirmButton).toBeDisabled(); + fireEvent.change(screen.getByPlaceholderText(GUARDRAIL.guardrail_name), { target: { value: "word" } }); + expect(confirmButton).toBeDisabled(); + fireEvent.change(screen.getByPlaceholderText(GUARDRAIL.guardrail_name), { + target: { value: GUARDRAIL.guardrail_name }, + }); + expect(confirmButton).toBeEnabled(); + }); + + it("shows a success toast naming the guardrail and calls onDeleted after a successful delete", async () => { + vi.mocked(deleteGuardrailCall).mockResolvedValue({}); + const { onClose, onDeleted } = renderModal(); + + confirmDelete(); + + await waitFor(() => expect(onDeleted).toHaveBeenCalledTimes(1)); + expect(deleteGuardrailCall).toHaveBeenCalledWith("tok", GUARDRAIL.guardrail_id); + expect(toast.success).toHaveBeenCalledWith('Guardrail "word-block" deleted successfully'); + expect(toast.fromError).not.toHaveBeenCalled(); + expect(onClose).toHaveBeenCalledTimes(1); + }); + + it("shows an error toast and does not call onDeleted when the delete fails", async () => { + vi.mocked(deleteGuardrailCall).mockRejectedValue(new Error("boom")); + const { onClose, onDeleted } = renderModal(); + + confirmDelete(); + + await waitFor(() => expect(onClose).toHaveBeenCalledTimes(1)); + expect(toast.fromError).toHaveBeenCalledWith("Failed to delete guardrail"); + expect(toast.success).not.toHaveBeenCalled(); + expect(onDeleted).not.toHaveBeenCalled(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailDeleteModal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailDeleteModal.tsx new file mode 100644 index 00000000000..a3ed15b630f --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailDeleteModal.tsx @@ -0,0 +1,68 @@ +import React, { useState } from "react"; +import DeleteResourceModal from "@/components/common_components/DeleteResourceModal"; +import { deleteGuardrailCall } from "@/components/networking"; +import { toast } from "@/lib/toast"; +import { formatGuardrailMode, getGuardrailLogoAndName } from "./guardrail_info_helpers"; + +const DANGLING_REFERENCE_WARNING = + "Keys, teams, and policies that list this guardrail will silently stop running it. Policy pipeline steps that use it will block every request with a 400 error, unless the step's on_error is next or allow (or its on_fail, when on_error is not set)."; + +export interface DeletableGuardrail { + guardrail_id: string; + guardrail_name?: string | null; + litellm_params?: { + guardrail?: string; + mode?: unknown; + default_on?: boolean | null; + } | null; +} + +interface GuardrailDeleteModalProps { + guardrail: DeletableGuardrail | null; + accessToken: string | null; + onClose: () => void; + onDeleted: () => void; +} + +export function GuardrailDeleteModal({ guardrail, accessToken, onClose, onDeleted }: GuardrailDeleteModalProps) { + const [isDeleting, setIsDeleting] = useState(false); + const guardrailName = guardrail ? guardrail.guardrail_name || guardrail.guardrail_id : ""; + const providerDisplayName = getGuardrailLogoAndName(guardrail?.litellm_params?.guardrail || "").displayName; + + const handleDelete = async () => { + if (!guardrail || !accessToken) return; + setIsDeleting(true); + try { + await deleteGuardrailCall(accessToken, guardrail.guardrail_id); + toast.success(`Guardrail "${guardrailName}" deleted successfully`); + onClose(); + onDeleted(); + } catch { + toast.fromError("Failed to delete guardrail"); + onClose(); + } finally { + setIsDeleting(false); + } + }; + + return ( + + ); +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailSettingsActions.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailSettingsActions.tsx new file mode 100644 index 00000000000..46bc09e7383 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailSettingsActions.tsx @@ -0,0 +1,51 @@ +import { Code, Trash2 } from "lucide-react"; +import React, { useState } from "react"; +import { Button } from "@/components/ui/button"; +import { DeletableGuardrail, GuardrailDeleteModal } from "./GuardrailDeleteModal"; + +interface GuardrailSettingsActionsProps { + guardrail: DeletableGuardrail & { guardrail_definition_location?: string }; + isEditing: boolean; + onEdit: () => void; + onEditCode: () => void; + accessToken: string | null; + onDeleted: () => void; +} + +export function GuardrailSettingsActions({ + guardrail, + isEditing, + onEdit, + onEditCode, + accessToken, + onDeleted, +}: GuardrailSettingsActionsProps) { + const [isModalOpen, setIsModalOpen] = useState(false); + + if (isEditing || guardrail.guardrail_definition_location === "config") return null; + + return ( + <> + {guardrail.litellm_params?.guardrail === "custom_code" ? ( + + ) : ( + + )} + + setIsModalOpen(false)} + onDeleted={onDeleted} + /> + + ); +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.integration.test.tsx new file mode 100644 index 00000000000..19eb838da6b --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.integration.test.tsx @@ -0,0 +1,135 @@ +import { beforeAll, beforeEach, describe, expect, it, vi } from "vitest"; +import GuardrailsPanel from "./GuardrailsPanel"; +import { + deleteGuardrailCall, + getGuardrailInfo, + getGuardrailProviderSpecificParams, + getGuardrailsList, + getGuardrailUISettings, +} from "@/components/networking"; +import { fireEvent, renderWithProviders, screen, waitFor, within } from "@/../tests/test-utils"; + +vi.mock("@/components/networking", () => ({ + getGuardrailsList: vi.fn(), + deleteGuardrailCall: vi.fn(), + getGuardrailInfo: vi.fn(), + getGuardrailUISettings: vi.fn(), + getGuardrailProviderSpecificParams: vi.fn(), + updateGuardrailCall: vi.fn(), +})); + +vi.mock("./add_guardrail_form", () => ({ __esModule: true, default: () => null })); +vi.mock("./guardrail_garden", () => ({ __esModule: true, default: () => null })); +vi.mock("./GuardrailTestPlayground", () => ({ __esModule: true, default: () => null })); +vi.mock("./TeamGuardrailsTab", () => ({ TeamGuardrailsTab: () => null })); + +const DB_GUARDRAIL = { + guardrail_id: "db-guardrail-1", + guardrail_name: "db-word-block", + litellm_params: { guardrail: "litellm_content_filter", mode: "pre_call", default_on: true }, + guardrail_info: null, + created_at: "2024-01-01T00:00:00Z", + updated_at: "2024-01-01T00:00:00Z", + guardrail_definition_location: "db", +}; + +const CONFIG_GUARDRAIL = { + ...DB_GUARDRAIL, + guardrail_id: "config-guardrail-1", + guardrail_name: "cfg-word-block", + guardrail_definition_location: "config", +}; + +beforeAll(() => { + Object.defineProperty(window, "matchMedia", { + writable: true, + value: vi.fn().mockImplementation((query: string) => ({ + matches: false, + media: query, + onchange: null, + addListener: vi.fn(), + removeListener: vi.fn(), + addEventListener: vi.fn(), + removeEventListener: vi.fn(), + dispatchEvent: vi.fn(), + })), + }); +}); + +const UI_SETTINGS = { + supported_entities: [], + supported_actions: [], + pii_entity_categories: [], + supported_modes: ["pre_call", "post_call"], +}; + +describe("GuardrailsPanel delete from the guardrail detail Settings card", () => { + beforeEach(() => { + vi.clearAllMocks(); + vi.mocked(getGuardrailsList).mockResolvedValue({ guardrails: [DB_GUARDRAIL, CONFIG_GUARDRAIL] }); + vi.mocked(getGuardrailInfo).mockImplementation(async (_token: string, id: string) => + id === CONFIG_GUARDRAIL.guardrail_id ? CONFIG_GUARDRAIL : DB_GUARDRAIL, + ); + vi.mocked(getGuardrailUISettings).mockResolvedValue(UI_SETTINGS); + vi.mocked(getGuardrailProviderSpecificParams).mockResolvedValue({}); + }); + + const renderDetail = (guardrailId: string) => + renderWithProviders(, { + searchParams: `?guardrail=${guardrailId}`, + }); + + it("deletes a DB guardrail only after the admin types its name, then returns to the refreshed table", async () => { + vi.mocked(deleteGuardrailCall).mockResolvedValue({}); + renderDetail(DB_GUARDRAIL.guardrail_id); + + fireEvent.click(await screen.findByRole("tab", { name: "Settings" })); + fireEvent.click(await screen.findByRole("button", { name: "Delete" })); + + const modal = within(await screen.findByRole("dialog")); + expect( + modal.getByText(/pipeline steps that use it will block every request with a 400 error/i), + ).toBeInTheDocument(); + const confirmButton = modal.getByRole("button", { name: "Delete" }); + expect(confirmButton).toBeDisabled(); + + const confirmationInput = modal.getByPlaceholderText(DB_GUARDRAIL.guardrail_name); + fireEvent.change(confirmationInput, { target: { value: "db-word" } }); + expect(confirmButton).toBeDisabled(); + fireEvent.change(confirmationInput, { target: { value: DB_GUARDRAIL.guardrail_name } }); + expect(confirmButton).toBeEnabled(); + + fireEvent.click(confirmButton); + + await waitFor(() => expect(deleteGuardrailCall).toHaveBeenCalledWith("test-token", DB_GUARDRAIL.guardrail_id)); + expect(await screen.findByRole("table")).toBeInTheDocument(); + expect(screen.queryByRole("button", { name: "Edit Settings" })).not.toBeInTheDocument(); + expect(getGuardrailsList).toHaveBeenCalledTimes(2); + }); + + it("keeps the admin on the detail page when the delete request fails", async () => { + vi.mocked(deleteGuardrailCall).mockRejectedValue(new Error("boom")); + renderDetail(DB_GUARDRAIL.guardrail_id); + + fireEvent.click(await screen.findByRole("tab", { name: "Settings" })); + fireEvent.click(await screen.findByRole("button", { name: "Delete" })); + const modal = within(await screen.findByRole("dialog")); + fireEvent.change(modal.getByPlaceholderText(DB_GUARDRAIL.guardrail_name), { + target: { value: DB_GUARDRAIL.guardrail_name }, + }); + fireEvent.click(modal.getByRole("button", { name: "Delete" })); + + await waitFor(() => expect(deleteGuardrailCall).toHaveBeenCalledTimes(1)); + await waitFor(() => expect(screen.queryByRole("dialog")).not.toBeInTheDocument()); + expect(screen.getByRole("button", { name: "Delete" })).toBeInTheDocument(); + expect(screen.queryByRole("table")).not.toBeInTheDocument(); + }); + + it("does not offer Delete on a config guardrail's detail page", async () => { + renderDetail(CONFIG_GUARDRAIL.guardrail_id); + + fireEvent.click(await screen.findByRole("tab", { name: "Settings" })); + expect(await screen.findByRole("img", { name: "Config guardrail details" })).toBeInTheDocument(); + expect(screen.queryByRole("button", { name: "Delete" })).not.toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.test.tsx index 7be120a4cb7..ac824f63be9 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.test.tsx @@ -2,6 +2,7 @@ import { type UrlUpdateEvent } from "nuqs/adapters/testing"; import { beforeAll, beforeEach, describe, expect, it, vi } from "vitest"; import GuardrailsPanel from "./GuardrailsPanel"; import { getGuardrailsList, deleteGuardrailCall } from "@/components/networking"; +import { toast } from "@/lib/toast"; import { fireEvent, renderWithProviders, screen, waitFor, within } from "@/../tests/test-utils"; vi.mock("@/components/networking", () => ({ @@ -9,6 +10,8 @@ vi.mock("@/components/networking", () => ({ deleteGuardrailCall: vi.fn(), })); +vi.mock("@/lib/toast", () => ({ toast: { success: vi.fn(), fromError: vi.fn() } })); + vi.mock("./add_guardrail_form", () => ({ __esModule: true, default: () =>
Mock Add Guardrail Form
, @@ -140,13 +143,36 @@ describe("GuardrailsPanel", () => { expect(modal.getByText("Delete Guardrail")).toBeInTheDocument(); expect(modal.getByText("test-guardrail-1")).toBeInTheDocument(); expect(modal.getByText("Test Provider")).toBeInTheDocument(); + expect( + modal.getByText(/pipeline steps that use it will block every request with a 400 error/i), + ).toBeInTheDocument(); - fireEvent.click(modal.getByRole("button", { name: "Delete" })); + const confirmButton = modal.getByRole("button", { name: "Delete" }); + expect(confirmButton).toBeDisabled(); + fireEvent.change(modal.getByPlaceholderText("Test Guardrail"), { target: { value: "Test Guardrail" } }); + fireEvent.click(confirmButton); await waitFor(() => { expect(mockDeleteGuardrailCall).toHaveBeenCalledWith("test-token", "test-guardrail-1"); }); - expect(mockGetGuardrailsList).toHaveBeenCalledTimes(2); + await waitFor(() => expect(mockGetGuardrailsList).toHaveBeenCalledTimes(2)); + expect(toast.success).toHaveBeenCalledWith('Guardrail "Test Guardrail" deleted successfully'); + }); + + it("should show an error toast and not refetch when the row delete fails", async () => { + mockDeleteGuardrailCall.mockRejectedValue(new Error("boom")); + renderWithProviders(); + fireEvent.click(screen.getByText("Guardrails")); + + fireEvent.click(await screen.findByTestId("delete-button")); + const modal = within(await screen.findByRole("dialog")); + fireEvent.change(modal.getByPlaceholderText("Test Guardrail"), { target: { value: "Test Guardrail" } }); + fireEvent.click(modal.getByRole("button", { name: "Delete" })); + + await waitFor(() => expect(screen.queryByRole("dialog")).not.toBeInTheDocument()); + expect(toast.fromError).toHaveBeenCalledWith("Failed to delete guardrail"); + expect(toast.success).not.toHaveBeenCalled(); + expect(mockGetGuardrailsList).toHaveBeenCalledTimes(1); }); it("should mount every tab panel up front so panel state survives tab switches", async () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.tsx index 901e39004f3..b53d016e950 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.tsx @@ -2,7 +2,7 @@ import { parseAsString, useQueryState } from "nuqs"; import React, { useState, useEffect } from "react"; import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; import { ChevronDown, Code, Plus } from "lucide-react"; -import { getGuardrailsList, deleteGuardrailCall } from "@/components/networking"; +import { getGuardrailsList } from "@/components/networking"; import { buttonVariants } from "@/components/ui/button"; import { DropdownMenu, @@ -16,10 +16,8 @@ import GuardrailTable from "./guardrail_table"; import { isAdminRole } from "@/utils/roles"; import GuardrailInfoView from "./guardrail_info"; import GuardrailTestPlayground from "./GuardrailTestPlayground"; -import { toast } from "@/lib/toast"; import { Guardrail } from "@/components/guardrails/types"; -import DeleteResourceModal from "@/components/common_components/DeleteResourceModal"; -import { formatGuardrailMode, getGuardrailLogoAndName } from "./guardrail_info_helpers"; +import { GuardrailDeleteModal } from "./GuardrailDeleteModal"; import { CustomCodeModal } from "./custom_code"; import GuardrailGarden from "./guardrail_garden"; import { TeamGuardrailsTab } from "./TeamGuardrailsTab"; @@ -38,9 +36,7 @@ const GuardrailsPanel: React.FC = ({ accessToken, userRole const [isAddModalVisible, setIsAddModalVisible] = useState(false); const [isCustomCodeModalVisible, setIsCustomCodeModalVisible] = useState(false); const [isLoading, setIsLoading] = useState(false); - const [isDeleting, setIsDeleting] = useState(false); const [guardrailToDelete, setGuardrailToDelete] = useState(null); - const [isDeleteModalOpen, setIsDeleteModalOpen] = useState(false); const [selectedGuardrailId, setSelectedGuardrailId] = useQueryState( "guardrail", parseAsString.withOptions({ history: "push" }), @@ -97,40 +93,10 @@ const GuardrailsPanel: React.FC = ({ accessToken, userRole fetchGuardrails(); }; - const handleDeleteClick = (guardrailId: string, guardrailName: string) => { - const guardrail = guardrailsList.find((g) => g.guardrail_id === guardrailId) || null; - setGuardrailToDelete(guardrail); - setIsDeleteModalOpen(true); + const handleDeleteClick = (guardrailId: string) => { + setGuardrailToDelete(guardrailsList.find((g) => g.guardrail_id === guardrailId) || null); }; - const handleDeleteConfirm = async () => { - if (!guardrailToDelete || !accessToken) return; - - setIsDeleting(true); - try { - await deleteGuardrailCall(accessToken, guardrailToDelete.guardrail_id); - toast.success(`Guardrail "${guardrailToDelete.guardrail_name}" deleted successfully`); - await fetchGuardrails(); - } catch (error) { - console.error("Error deleting guardrail:", error); - toast.fromError("Failed to delete guardrail"); - } finally { - setIsDeleting(false); - setIsDeleteModalOpen(false); - setGuardrailToDelete(null); - } - }; - - const handleDeleteCancel = () => { - setIsDeleteModalOpen(false); - setGuardrailToDelete(null); - }; - - const providerDisplayName = - guardrailToDelete && guardrailToDelete.litellm_params - ? getGuardrailLogoAndName(guardrailToDelete.litellm_params.guardrail).displayName - : undefined; - return (
@@ -184,6 +150,10 @@ const GuardrailsPanel: React.FC = ({ accessToken, userRole { + closeGuardrailDetail(); + fetchGuardrails(); + }} accessToken={accessToken} isAdmin={isAdmin} /> @@ -210,24 +180,11 @@ const GuardrailsPanel: React.FC = ({ accessToken, userRole onSuccess={handleSuccess} /> - setGuardrailToDelete(null)} + onDeleted={fetchGuardrails} /> diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.characterization.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.characterization.test.tsx index 1689cf93289..732c50647b1 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.characterization.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.characterization.test.tsx @@ -84,7 +84,8 @@ const saveChanges = async (user: ReturnType) => { await user.click(screen.getByText("Save Changes")); }; -const renderView = () => render( {}} accessToken="123" isAdmin />); +const renderView = () => + render( {}} onDeleted={() => {}} accessToken="123" isAdmin />); const lastPayload = () => vi.mocked(networking.updateGuardrailCall).mock.calls.at(-1)?.[2]; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.test.tsx index 836921104ff..ae057689b37 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.test.tsx @@ -10,6 +10,7 @@ vi.mock("@/components/networking", () => ({ getGuardrailUISettings: vi.fn(), getGuardrailProviderSpecificParams: vi.fn(), updateGuardrailCall: vi.fn(), + deleteGuardrailCall: vi.fn(), })); // Mock ContentFilterManager @@ -65,7 +66,9 @@ describe("Guardrail Info", () => { vi.mocked(networking.getGuardrailProviderSpecificParams).mockResolvedValue({}); - render( {}} accessToken="123" isAdmin={true} />); + render( + {}} onDeleted={() => {}} accessToken="123" isAdmin={true} />, + ); // Wait for the loading to complete and data to be rendered await waitFor(() => { @@ -103,7 +106,9 @@ describe("Guardrail Info", () => { vi.mocked(networking.getGuardrailProviderSpecificParams).mockResolvedValue({}); - render( {}} accessToken="123" isAdmin={true} />); + render( + {}} onDeleted={() => {}} accessToken="123" isAdmin={true} />, + ); expect(await screen.findAllByText("pre_call, post_call (tag-based)")).not.toHaveLength(0); }); @@ -131,7 +136,9 @@ describe("Guardrail Info", () => { vi.mocked(networking.getGuardrailProviderSpecificParams).mockResolvedValue({}); - render( {}} accessToken="123" isAdmin={true} />); + render( + {}} onDeleted={() => {}} accessToken="123" isAdmin={true} />, + ); const logo = await screen.findByAltText("Presidio PII logo"); expect(logo).toHaveAttribute("src", expect.stringContaining("microsoft_azure.svg")); @@ -162,7 +169,7 @@ describe("Guardrail Info", () => { vi.mocked(networking.getGuardrailProviderSpecificParams).mockResolvedValue({}); const { container } = render( - {}} accessToken="123" isAdmin={true} />, + {}} onDeleted={() => {}} accessToken="123" isAdmin={true} />, ); await waitFor(() => { @@ -212,7 +219,9 @@ describe("Guardrail Info", () => { vi.mocked(networking.getGuardrailProviderSpecificParams).mockResolvedValue({}); - render( {}} accessToken="123" isAdmin={true} />); + render( + {}} onDeleted={() => {}} accessToken="123" isAdmin={true} />, + ); await waitFor(() => { expect(screen.getByText("PII Entity Configuration")).toBeInTheDocument(); @@ -245,7 +254,9 @@ describe("Guardrail Info", () => { vi.mocked(networking.getGuardrailProviderSpecificParams).mockResolvedValue({}); vi.mocked(networking.updateGuardrailCall).mockResolvedValue({ status: "success" }); - render( {}} accessToken="123" isAdmin={true} />); + render( + {}} onDeleted={() => {}} accessToken="123" isAdmin={true} />, + ); await waitFor(() => { expect(screen.getByText("Settings")).toBeInTheDocument(); @@ -331,7 +342,9 @@ describe("Guardrail Info", () => { }); vi.mocked(networking.getGuardrailProviderSpecificParams).mockResolvedValue({}); - render( {}} accessToken="123" isAdmin={true} />); + render( + {}} onDeleted={() => {}} accessToken="123" isAdmin={true} />, + ); expect(await screen.findByRole("tab", { name: "Overview" })).toHaveAttribute("aria-selected", "true"); expect(screen.getByRole("tab", { name: "Settings" })).toHaveAttribute("aria-selected", "false"); @@ -355,10 +368,90 @@ describe("Guardrail Info when the guardrail cannot be loaded", () => { vi.mocked(networking.getGuardrailProviderSpecificParams).mockResolvedValue({}); const onClose = vi.fn(); - render(); + render( + {}} + accessToken="123" + isAdmin={true} + />, + ); expect(await screen.findByText("Guardrail not found")).toBeInTheDocument(); fireEvent.click(screen.getByRole("button", { name: /back to guardrails/i })); expect(onClose).toHaveBeenCalledTimes(1); }); + + describe("Delete from the Settings card", () => { + const dbGuardrail = { + guardrail_id: "123", + guardrail_name: "Test Guardrail", + litellm_params: { guardrail: "presidio", mode: "pre_call", default_on: true }, + guardrail_definition_location: "db", + }; + const uiSettings = { + supported_entities: [], + supported_actions: [], + pii_entity_categories: [], + supported_modes: ["pre_call"], + }; + const mockDbGuardrail = () => { + vi.mocked(networking.getGuardrailInfo).mockResolvedValue(dbGuardrail); + vi.mocked(networking.getGuardrailUISettings).mockResolvedValue(uiSettings); + vi.mocked(networking.getGuardrailProviderSpecificParams).mockResolvedValue({}); + }; + + const openDeleteModalAndConfirm = async () => { + fireEvent.click(await screen.findByRole("tab", { name: "Settings" })); + fireEvent.click(await screen.findByRole("button", { name: "Delete" })); + const modal = within(await screen.findByRole("dialog")); + expect(modal.getByText(/silently stop running it/)).toBeInTheDocument(); + const confirmButton = modal.getByRole("button", { name: "Delete" }); + expect(confirmButton).toBeDisabled(); + fireEvent.change(modal.getByPlaceholderText("Test Guardrail"), { target: { value: "Test Guardrail" } }); + fireEvent.click(confirmButton); + }; + + it("deletes the guardrail and calls onDeleted once the admin types its name", async () => { + mockDbGuardrail(); + vi.mocked(networking.deleteGuardrailCall).mockResolvedValue({}); + const onDeleted = vi.fn(); + render( + {}} onDeleted={onDeleted} accessToken="tok" isAdmin />, + ); + + await openDeleteModalAndConfirm(); + + await waitFor(() => expect(onDeleted).toHaveBeenCalledTimes(1)); + expect(networking.deleteGuardrailCall).toHaveBeenCalledWith("tok", "123"); + }); + + it("hides Delete while the settings are being edited", async () => { + mockDbGuardrail(); + render( {}} onDeleted={() => {}} accessToken="tok" isAdmin />); + + fireEvent.click(await screen.findByRole("tab", { name: "Settings" })); + fireEvent.click(await screen.findByRole("button", { name: "Edit Settings" })); + + expect(await screen.findByRole("button", { name: "Save Changes" })).toBeInTheDocument(); + expect(screen.queryByRole("button", { name: "Delete" })).not.toBeInTheDocument(); + expect(screen.queryByRole("button", { name: "Edit Settings" })).not.toBeInTheDocument(); + }); + + it("does not call onDeleted when the delete request fails", async () => { + mockDbGuardrail(); + vi.mocked(networking.deleteGuardrailCall).mockRejectedValue(new Error("boom")); + const onDeleted = vi.fn(); + render( + {}} onDeleted={onDeleted} accessToken="tok" isAdmin />, + ); + + await openDeleteModalAndConfirm(); + + await waitFor(() => expect(screen.queryByRole("dialog")).not.toBeInTheDocument()); + expect(networking.deleteGuardrailCall).toHaveBeenCalledTimes(1); + expect(onDeleted).not.toHaveBeenCalled(); + }); + }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.tsx index 2cfe2b79886..eacacf5e30e 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_info.tsx @@ -33,6 +33,7 @@ import { } from "./GuardrailFormField"; import ContentFilterManager, { formatContentFilterDataForAPI } from "./content_filter/ContentFilterManager"; import CustomCodeModal, { EditGuardrailData } from "./custom_code/CustomCodeModal"; +import { GuardrailSettingsActions } from "./GuardrailSettingsActions"; import { GuardrailModeCard } from "./GuardrailModeDisplay"; import { GuardrailReadOnlyDetails } from "./GuardrailReadOnlyDetails"; import { @@ -72,11 +73,12 @@ const SectionHeading: React.FC<{ children: React.ReactNode }> = ({ children }) = export interface GuardrailInfoProps { guardrailId: string; onClose: () => void; + onDeleted: () => void; accessToken: string | null; isAdmin: boolean; } -const GuardrailInfoView: React.FC = ({ guardrailId, onClose, accessToken, isAdmin }) => { +const GuardrailInfoView: React.FC = ({ guardrailId, onClose, onDeleted, accessToken, isAdmin }) => { const [guardrailData, setGuardrailData] = useState(null); const [guardrailProviderSpecificParams, setGuardrailProviderSpecificParams] = useState(null); const [loading, setLoading] = useState(true); @@ -686,25 +688,21 @@ const GuardrailInfoView: React.FC = ({ guardrailId, onClose, {isAdmin && ( -
-

Guardrail Settings

+
+

Guardrail Settings

{isConfigGuardrail && ( )} - {!isEditing && - !isConfigGuardrail && - (guardrailData.litellm_params?.guardrail === "custom_code" ? ( - - ) : ( - - ))} + setIsEditing(true)} + onEditCode={() => setCustomCodeModalVisible(true)} + accessToken={accessToken} + onDeleted={onDeleted} + />
{isEditing ? ( 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 */