diff --git a/litellm/caching/caching_handler.py b/litellm/caching/caching_handler.py index 501dc266ba0..74c45189622 100644 --- a/litellm/caching/caching_handler.py +++ b/litellm/caching/caching_handler.py @@ -512,7 +512,7 @@ class LLMCachingHandler: ) if not _should_defer_streaming_cache_hit_callbacks(cached_result=cached_result): - logging_obj.handle_sync_success_callbacks_for_async_calls( + logging_obj.handle_sync_cache_hit_success_callbacks( result=cached_result, start_time=start_time, end_time=end_time, diff --git a/litellm/constants.py b/litellm/constants.py index 7879e3bef54..c39558dac36 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -389,7 +389,7 @@ DEEPGRAM_DEFAULT_API_BASE: Final = "https://api.deepgram.com/v1" NADIR_DEFAULT_API_BASE: Final = "https://api.getnadir.com/v1" DEEPGRAM_LISTEN_DEFAULT_MODEL: Final = "nova-3" -BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY: Final = "litellm.bedrock_realtime.pending_session_update" +BEDROCK_REALTIME_PENDING_OPENING_MESSAGE_SCOPE_KEY: Final = "litellm.bedrock_realtime.pending_opening_message" BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY: Final = "litellm.bedrock_realtime.session_committed" BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY: Final = "litellm.bedrock_realtime.committed_failure" BEDROCK_REALTIME_SDK_DISTRIBUTION: Final = "aws-sdk-bedrock-runtime" @@ -1613,6 +1613,7 @@ CLOUDZERO_EXPORT_INTERVAL_MINUTES: Final = get_env_int("CLOUDZERO_EXPORT_INTERVA MCP_TOOL_NAME_PREFIX: Final = "mcp_tool" MAXIMUM_TRACEBACK_LINES_TO_LOG: Final = get_env_int("MAXIMUM_TRACEBACK_LINES_TO_LOG", 100) PASSTHROUGH_UPSTREAM_ERROR_BODY_MAX_LOG_CHARS: Final = 4096 +INTERNAL_CALLBACK_NAMES: Final = frozenset({"cache", "vector_store_pre_call_hook"}) # Headers to control callbacks X_LITELLM_DISABLE_CALLBACKS: Final = "x-litellm-disable-callbacks" diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 0e315e97432..3b0a7fc4cb7 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -7,7 +7,7 @@ from collections.abc import Mapping, Sequence from datetime import datetime from functools import lru_cache from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Literal, cast +from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias, cast from httpx import Response from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError @@ -147,6 +147,7 @@ from litellm.utils import ( TranscriptionResponse, cached_get_model_info_helper, token_counter, + unmapped_model_generalization, ) if TYPE_CHECKING: @@ -863,6 +864,16 @@ def _base_model_pricing_provider(base_model: str, custom_llm_provider: str | Non return own_model_info["litellm_provider"] +def _pricing_provider_for_cost_candidate( + candidate_provider: str | None, + request_provider: str | None, + current_provider: object, +) -> str | None: + if candidate_provider != request_provider: + return candidate_provider + return current_provider if isinstance(current_provider, str) else None + + def select_model_name_for_cost_calc( model: str | None, completion_response: object | None, @@ -1549,14 +1560,20 @@ def completion_cost( ) requested_model: Final = model - potential_model_names: Final = [ - selected_model, - _get_response_model(completion_response), - ] - if model is not None: - potential_model_names.append(model) + request_pricing_provider: Final = custom_llm_provider if isinstance(custom_llm_provider, str) else None + raw_candidates: Final = ( + ( + selected_model, + base_model_provider if base_model_provider is not None else request_pricing_provider, + ), + (_get_response_model(completion_response), request_pricing_provider), + *(((model, request_pricing_provider),) if model is not None else ()), + ) + potential_model_names: Final = ( + raw_candidates if explicit_pricing else _priced_candidates_before_unpriced_rules(raw_candidates) + ) - for idx, model in enumerate(potential_model_names): + for idx, (model, pricing_provider) in enumerate(potential_model_names): try: if verbose_logger.isEnabledFor(logging.DEBUG): verbose_logger.debug("selected model name for cost calculation: %s", model) @@ -1895,13 +1912,6 @@ def completion_cost( if litellm_logging_obj is not None: request_model_for_cost = litellm_logging_obj.model - pricing_provider: str | None = ( - base_model_provider - if base_model_provider is not None and idx == 0 - else custom_llm_provider - if isinstance(custom_llm_provider, str) - else None - ) ( prompt_tokens_cost_usd_dollar, completion_tokens_cost_usd_dollar, @@ -1909,7 +1919,11 @@ def completion_cost( model=model, prompt_tokens=prompt_tokens or 0, completion_tokens=completion_tokens or 0, - custom_llm_provider=pricing_provider, + custom_llm_provider=_pricing_provider_for_cost_candidate( + candidate_provider=pricing_provider, + request_provider=request_pricing_provider, + current_provider=custom_llm_provider, + ), response_time_ms=response_time_ms, region_name=None if explicit_pricing else region_name, custom_cost_per_second=custom_cost_per_second, @@ -2008,7 +2022,11 @@ def completion_cost( if cost_per_token_usage_object is not None and model: _token_type_breakdown = get_token_type_cost_breakdown( model=model, - custom_llm_provider=pricing_provider, + custom_llm_provider=_pricing_provider_for_cost_candidate( + candidate_provider=pricing_provider, + request_provider=request_pricing_provider, + current_provider=custom_llm_provider, + ), usage=cost_per_token_usage_object, service_tier=service_tier, data_residency=data_residency, @@ -2263,6 +2281,25 @@ def _cost_map_model_info(model: str, custom_llm_provider: str | None) -> ModelIn return None +_CostCandidate: TypeAlias = tuple[str | None, str | None] + + +def _is_unpriced_generalized_candidate(candidate: _CostCandidate) -> bool: + model, pricing_provider = candidate + if model is None: + return False + rule_info: Final = unmapped_model_generalization(model, pricing_provider) + return rule_info is not None and not _cost_map_entry_prices_anything(rule_info) + + +def _priced_candidates_before_unpriced_rules(candidates: tuple[_CostCandidate, ...]) -> tuple[_CostCandidate, ...]: + unpriced: Final = tuple(_is_unpriced_generalized_candidate(candidate) for candidate in candidates) + return ( + *(candidate for candidate, is_unpriced in zip(candidates, unpriced, strict=True) if not is_unpriced), + *(candidate for candidate, is_unpriced in zip(candidates, unpriced, strict=True) if is_unpriced), + ) + + def _raw_cost_map_entry(key: str) -> Mapping[str, object] | None: raw_entry: Final = litellm.model_cost.get(key) return raw_entry if isinstance(raw_entry, Mapping) else None diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index a25362c3b87..f74f6d2de66 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -41,6 +41,7 @@ from litellm.constants import ( DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT, DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT, EMPTY_MAPPING, + INTERNAL_CALLBACK_NAMES, PROVIDER_REQUEST_ID_HEADERS, REDACTED_BY_LITELLM, ) @@ -2287,6 +2288,7 @@ class Logging(LiteLLMLoggingBaseClass): start_time=start_time, end_time=end_time, cache_hit=cache_hit, + for_async_request=True, **kwargs, ) @@ -2480,7 +2482,8 @@ class Logging(LiteLLMLoggingBaseClass): start_time, end_time, build_logging_payload: bool = True, - ): + publish_logging_payload: bool = True, + ) -> StandardLoggingPayload | None: """Resolve hidden params, compute response cost, and emit the standard logging payload.""" self._surface_response_headers_from_result(logging_result) hidden_params: Final = getattr(logging_result, HIDDEN_PARAMS_ATTR, {}) @@ -2507,17 +2510,17 @@ class Logging(LiteLLMLoggingBaseClass): self.model_call_details["response_cost"] = self.response_cost_calculator(result=logging_result) if not build_logging_payload: - return + return None - self.model_call_details["standard_logging_object"] = self._build_standard_logging_payload( - logging_result, start_time, end_time - ) - - standard_logging_payload: Final[StandardLoggingPayload | None] = self.model_call_details.get( - "standard_logging_object" + standard_logging_payload: Final = self._build_standard_logging_payload( + logging_result, start_time, end_time, isolate_call_details=not publish_logging_payload ) + if not publish_logging_payload: + return standard_logging_payload + self.model_call_details["standard_logging_object"] = standard_logging_payload if standard_logging_payload is not None: emit_standard_logging_payload(standard_logging_payload) + return standard_logging_payload async def _prepare_baseline_cache_estimate(self, response_obj: object) -> None: if self.baseline_cache_context is None: @@ -2533,12 +2536,16 @@ class Logging(LiteLLMLoggingBaseClass): await invalidate_baseline_cache(self, reason, completed=completed) def _build_standard_logging_payload( - self, init_response_obj: object, start_time: dt_object, end_time: dt_object + self, + init_response_obj: object, + start_time: dt_object, + end_time: dt_object, + isolate_call_details: bool = False, ) -> StandardLoggingPayload | None: """Build StandardLoggingPayload and accumulate its construction time.""" _start: Final = time.time() payload: Final = get_standard_logging_object_payload( - kwargs=self.model_call_details, + kwargs=dict(self.model_call_details) if isolate_call_details else self.model_call_details, init_response_obj=init_response_obj, start_time=start_time, end_time=end_time, @@ -2555,11 +2562,16 @@ class Logging(LiteLLMLoggingBaseClass): emit_standard_logging_payload(stored_payload) def _transform_usage_objects(self, result: object) -> object: + return self._transform_usage_objects_for_payload(result, self.model_call_details.get("standard_logging_object")) + + def _transform_usage_objects_for_payload( + self, result: object, standard_logging_payload: StandardLoggingPayload | None + ) -> object: if isinstance(result, ResponsesAPIResponse): result = result.model_copy() transformed_usage = ResponseAPILoggingUtils.transform_response_api_usage_to_chat_usage(result.usage) setattr(result, "usage", transformed_usage) - if (standard_logging_payload := self.model_call_details.get("standard_logging_object")) is not None: + if standard_logging_payload is not None: response_dict: Final = result.model_dump() if hasattr(result, "model_dump") else dict(result) # Ensure usage is properly included with transformed chat format if transformed_usage is not None: @@ -2568,7 +2580,7 @@ class Logging(LiteLLMLoggingBaseClass): if hasattr(transformed_usage, "model_dump") else dict(transformed_usage) ) - standard_logging_payload["response"] = response_dict + standard_logging_payload["response"] = response_dict # rebind-ok: mirrors the transformed usage elif isinstance(result, TranscriptionResponse): from litellm.litellm_core_utils.llm_cost_calc.usage_object_transformation import ( TranscriptionUsageObjectTransformation, @@ -2599,6 +2611,27 @@ class Logging(LiteLLMLoggingBaseClass): standard_logging_object: StandardLoggingPayload | None = None, build_logging_payload: bool = True, ) -> tuple[datetime.datetime, datetime.datetime, object]: + prepared: Final = self._prepare_success_logging( + result=result, + start_time=start_time, + end_time=end_time, + cache_hit=cache_hit, + standard_logging_object=standard_logging_object, + build_logging_payload=build_logging_payload, + publish_logging_payload=True, + ) + return prepared[0], prepared[1], prepared[2] + + def _prepare_success_logging( + self, + result: object, + start_time: datetime.datetime | None, + end_time: datetime.datetime | None, + cache_hit: bool | None, + standard_logging_object: StandardLoggingPayload | None, + build_logging_payload: bool, + publish_logging_payload: bool, + ) -> tuple[datetime.datetime, datetime.datetime, object, StandardLoggingPayload | None]: try: if start_time is None: start_time = self.start_time @@ -2629,29 +2662,21 @@ class Logging(LiteLLMLoggingBaseClass): ): result = logging_result - if standard_logging_object is None and result is not None and self.stream is not True: - if self._is_recognized_call_type_for_logging(logging_result=logging_result) or isinstance( - logging_result, (dict, list) - ): - self._process_hidden_params_and_response_cost( - logging_result=logging_result, - start_time=start_time, - end_time=end_time, - build_logging_payload=build_logging_payload, - ) - elif standard_logging_object is not None: - self.model_call_details["standard_logging_object"] = standard_logging_object - else: - # Streaming reaches here before its cost is known, so the cost - # is seeded to None, but only when nothing has already - # established one. A stream that assembles into a response - # object recomputes the cost right after this; a pass-through - # stream cannot (its body is opaque) and carries the cost its - # upstream reported in the response headers, which an - # unconditional reset would discard. - self.model_call_details.setdefault("response_cost", None) + standard_logging_payload: Final = self._resolve_success_logging_payload( + result=result, + logging_result=logging_result, + start_time=start_time, + end_time=end_time, + standard_logging_object=standard_logging_object, + build_logging_payload=build_logging_payload, + publish_logging_payload=publish_logging_payload, + ) - result = self._transform_usage_objects(result=result) + result = ( + self._transform_usage_objects(result=result) + if publish_logging_payload + else self._transform_usage_objects_for_payload(result, standard_logging_payload) + ) if ( litellm.max_budget @@ -2670,10 +2695,47 @@ class Logging(LiteLLMLoggingBaseClass): standard_built_in_tools_params=self.standard_built_in_tools_params, ) - return start_time, end_time, result + return start_time, end_time, result, standard_logging_payload except Exception as e: raise Exception(f"[Non-Blocking] LiteLLM.Success_Call Error: {e}") + def _resolve_success_logging_payload( + self, + result: object, + logging_result: object, + start_time: datetime.datetime, + end_time: datetime.datetime, + standard_logging_object: StandardLoggingPayload | None, + build_logging_payload: bool, + publish_logging_payload: bool, + ) -> StandardLoggingPayload | None: + if standard_logging_object is not None: + if not publish_logging_payload: + return standard_logging_object.copy() + self.model_call_details["standard_logging_object"] = standard_logging_object + return standard_logging_object + if result is None or self.stream is True: + # Streaming reaches here before its cost is known, so the cost + # is seeded to None, but only when nothing has already + # established one. A stream that assembles into a response + # object recomputes the cost right after this; a pass-through + # stream cannot (its body is opaque) and carries the cost its + # upstream reported in the response headers, which an + # unconditional reset would discard. + self.model_call_details.setdefault("response_cost", None) + return None + if not self._is_recognized_call_type_for_logging(logging_result=logging_result) and not isinstance( + logging_result, (dict, list) + ): + return None + return self._process_hidden_params_and_response_cost( + logging_result=logging_result, + start_time=start_time, + end_time=end_time, + build_logging_payload=build_logging_payload, + publish_logging_payload=publish_logging_payload, + ) + def _is_recognized_call_type_for_logging( self, logging_result: object, @@ -2855,6 +2917,8 @@ class Logging(LiteLLMLoggingBaseClass): start_time: datetime.datetime | None = None, end_time: datetime.datetime | None = None, cache_hit: bool | None = None, + *, + for_async_request: bool = False, **kwargs: object, # kwargs-ok: forwarded to _success_handler_body ) -> None: """Restores trace_id/session_id contextvars once this attempt's own success @@ -2862,7 +2926,12 @@ class Logging(LiteLLMLoggingBaseClass): try: with post_response_phase(): return self._success_handler_body( - result=result, start_time=start_time, end_time=end_time, cache_hit=cache_hit, **kwargs + result=result, + start_time=start_time, + end_time=end_time, + cache_hit=cache_hit, + for_async_request=for_async_request, + **kwargs, ) finally: self._restore_correlation_context() @@ -2873,17 +2942,20 @@ class Logging(LiteLLMLoggingBaseClass): start_time: datetime.datetime | None = None, end_time: datetime.datetime | None = None, cache_hit: bool | None = None, + for_async_request: bool = False, **kwargs: Any, # kwargs-ok: forwarded from success_handler ) -> None: verbose_logger.debug("Logging Details LiteLLM-Success Call: Cache_hit=%s", cache_hit) if not self.should_run_logging(event_type="sync_success"): # prevent double logging return - start_time, end_time, result = self._success_handler_helper_fn( + start_time, end_time, result, request_logging_payload = self._prepare_success_logging( start_time=start_time, end_time=end_time, result=result, cache_hit=cache_hit, standard_logging_object=kwargs.get("standard_logging_object", None), + build_logging_payload=True, + publish_logging_payload=not for_async_request, ) litellm_params = self.model_call_details.get("litellm_params", {}) is_sync_request: Final = self.is_sync_litellm_request(litellm_params) @@ -2912,16 +2984,45 @@ class Logging(LiteLLMLoggingBaseClass): logged_streaming_response: Final = self._with_client_facing_stream_model(complete_streaming_response) self.model_call_details["complete_streaming_response"] = logged_streaming_response ## STANDARDIZED LOGGING PAYLOAD - self.model_call_details["standard_logging_object"] = self._build_standard_logging_payload( - logged_streaming_response, start_time, end_time + if not for_async_request: + self.model_call_details["standard_logging_object"] = self._build_standard_logging_payload( + logged_streaming_response, start_time, end_time + ) + standard_logging_payload: Final[StandardLoggingPayload | None] = self.model_call_details.get( + "standard_logging_object" + ) + if standard_logging_payload is not None: + # Only emit for sync requests (async_success_handler handles async) + if is_sync_request: + emit_standard_logging_payload(standard_logging_payload) + callback_logging_payload: Final = ( + self._build_standard_logging_payload( + self.model_call_details["complete_streaming_response"], + start_time, + end_time, + isolate_call_details=True, ) - standard_logging_payload: Final[StandardLoggingPayload | None] = self.model_call_details.get( - "standard_logging_object" + if for_async_request and complete_streaming_response is not None + else request_logging_payload + ) + shared_logging_payload: Final[StandardLoggingPayload | None] = self.model_call_details.get( + "standard_logging_object" + ) + isolated_logging_payload: Final = ( + callback_logging_payload + if callback_logging_payload is not None or shared_logging_payload is None + else shared_logging_payload.copy() + ) + call_details: dict[str, Any] # mutable-ok: legacy loggers take a plain dict + call_details = ( # rebind-ok: logging hooks may return a new dict + self.model_call_details + if not for_async_request + else ( + dict(self.model_call_details) + if isolated_logging_payload is None + else {**self.model_call_details, "standard_logging_object": isolated_logging_payload} ) - if standard_logging_payload is not None: - # Only emit for sync requests (async_success_handler handles async) - if is_sync_request: - emit_standard_logging_payload(standard_logging_payload) + ) callbacks: Final = self.get_combined_callback_list( dynamic_success_callbacks=self.dynamic_success_callbacks, global_callbacks=litellm.success_callback, @@ -2929,7 +3030,7 @@ class Logging(LiteLLMLoggingBaseClass): ## REDACT MESSAGES ## result = redact_message_input_output_from_logging( - model_call_details=(self.model_call_details if hasattr(self, "model_call_details") else {}), + model_call_details=call_details, result=result, ) ## LOGGING HOOK ## @@ -2939,7 +3040,7 @@ class Logging(LiteLLMLoggingBaseClass): if ( callback.should_run_guardrail( - data=self.model_call_details, + data=call_details, event_type=GuardrailEventHooks.logging_only, ) is not True @@ -2947,18 +3048,21 @@ class Logging(LiteLLMLoggingBaseClass): continue self.shadow_eval_request_snapshot = None - self.model_call_details, result = callback.logging_hook( - kwargs=self.model_call_details, + call_details, result = callback.logging_hook( # pyright: ignore[reportUnknownVariableType] # logging_hook is typed with a bare dict + kwargs=call_details, result=result, call_type=self.call_type, ) elif isinstance(callback, CustomLogger): - self.model_call_details, result = callback.logging_hook( - kwargs=self.model_call_details, + call_details, result = callback.logging_hook( # pyright: ignore[reportUnknownVariableType] # logging_hook is typed with a bare dict + kwargs=call_details, result=result, call_type=self.call_type, ) + if not for_async_request: + self.model_call_details = call_details + call_details_view: Final = MappingProxyType[str, object](call_details) self.mark_logging_complete(event_type="sync_success") for callback in callbacks: try: @@ -2972,7 +3076,7 @@ class Logging(LiteLLMLoggingBaseClass): if callback == "promptlayer" and promptLayerLogger is not None: print_verbose("reaches promptlayer for logging!") promptLayerLogger.log_event( - kwargs=self.model_call_details, + kwargs=call_details, response_obj=result, start_time=start_time, end_time=end_time, @@ -2980,7 +3084,7 @@ class Logging(LiteLLMLoggingBaseClass): ) if callback == "supabase" and supabaseClient is not None: print_verbose("reaches supabase for logging!") - kwargs = self.model_call_details + kwargs = call_details # this only logs streaming once, complete_streaming_response exists i.e when stream ends if self.stream: @@ -2988,10 +3092,10 @@ class Logging(LiteLLMLoggingBaseClass): continue else: print_verbose("reaches supabase for streaming logging!") - result = self._model_call_details_view["complete_streaming_response"] + result = call_details_view["complete_streaming_response"] - model = self._model_call_details_view["model"] - messages = self._model_call_details_view["messages"] + model = call_details_view["model"] + messages = call_details_view["messages"] optional_params = kwargs.get("optional_params", {}) litellm_params = kwargs.get("litellm_params", {}) supabaseClient.log_event( @@ -3011,7 +3115,7 @@ class Logging(LiteLLMLoggingBaseClass): if callback == "wandb" and weightsBiasesLogger is not None: print_verbose("reaches wandb for logging!") weightsBiasesLogger.log_event( - kwargs=self.model_call_details, + kwargs=call_details, response_obj=result, start_time=start_time, end_time=end_time, @@ -3020,7 +3124,7 @@ class Logging(LiteLLMLoggingBaseClass): if callback == "logfire" and logfireLogger is not None: verbose_logger.debug("reaches logfire for success logging!") kwargs = {} - for k, v in self._model_call_details_view.items(): + for k, v in call_details_view.items(): if k != "original_response": # copy.deepcopy raises errors as this could be a coroutine kwargs[k] = v @@ -3030,10 +3134,10 @@ class Logging(LiteLLMLoggingBaseClass): continue else: print_verbose("reaches logfire for streaming logging!") - result = self._model_call_details_view["complete_streaming_response"] + result = call_details_view["complete_streaming_response"] logfireLogger.log_event( - kwargs=self.model_call_details, + kwargs=call_details, response_obj=result, start_time=start_time, end_time=end_time, @@ -3044,7 +3148,7 @@ class Logging(LiteLLMLoggingBaseClass): if callback == "lunary" and lunaryLogger is not None: print_verbose("reaches lunary for logging!") model = self.model - kwargs = self.model_call_details + kwargs = call_details input = kwargs.get("messages", kwargs.get("input", None)) @@ -3055,7 +3159,7 @@ class Logging(LiteLLMLoggingBaseClass): if "complete_streaming_response" not in kwargs: continue else: - result = self._model_call_details_view["complete_streaming_response"] + result = call_details_view["complete_streaming_response"] lunaryLogger.log_event( type=type, @@ -3065,7 +3169,7 @@ class Logging(LiteLLMLoggingBaseClass): input=input, user_id=kwargs.get("user", None), # user_props=self.model_call_details.get("user_props", None), - extra=self._model_call_details_view.get("optional_params", {}), + extra=call_details_view.get("optional_params", {}), response_obj=result, start_time=start_time, end_time=end_time, @@ -3075,8 +3179,8 @@ class Logging(LiteLLMLoggingBaseClass): if callback == "helicone" and heliconeLogger is not None: print_verbose("reaches helicone for logging!") model = self.model - messages = self._model_call_details_view["input"] - kwargs = self.model_call_details + messages = call_details_view["input"] + kwargs = call_details # this only logs streaming once, complete_streaming_response exists i.e when stream ends if self.stream: @@ -3084,7 +3188,7 @@ class Logging(LiteLLMLoggingBaseClass): continue else: print_verbose("reaches helicone for streaming logging!") - result = self._model_call_details_view["complete_streaming_response"] + result = call_details_view["complete_streaming_response"] heliconeLogger.log_success( model=model, @@ -3099,7 +3203,7 @@ class Logging(LiteLLMLoggingBaseClass): global langFuseLogger print_verbose("reaches langfuse for success logging!") kwargs = {} - for k, v in self._model_call_details_view.items(): + for k, v in call_details_view.items(): if k != "original_response": # copy.deepcopy raises errors as this could be a coroutine kwargs[k] = v # this only logs streaming once, complete_streaming_response exists i.e when stream ends @@ -3112,7 +3216,7 @@ class Logging(LiteLLMLoggingBaseClass): continue else: print_verbose("reaches langfuse for streaming logging!") - result = self._model_call_details_view["complete_streaming_response"] + result = call_details_view["complete_streaming_response"] langfuse_logger_to_use = LangFuseHandler.get_langfuse_logger_for_request( globalLangfuseLogger=langFuseLogger, @@ -3137,7 +3241,7 @@ class Logging(LiteLLMLoggingBaseClass): ) if callback == "greenscale" and greenscaleLogger is not None: kwargs = {} - for k, v in self._model_call_details_view.items(): + for k, v in call_details_view.items(): if k != "original_response": # copy.deepcopy raises errors as this could be a coroutine kwargs[k] = v # this only logs streaming once, complete_streaming_response exists i.e when stream ends @@ -3150,7 +3254,7 @@ class Logging(LiteLLMLoggingBaseClass): continue else: print_verbose("reaches greenscale for streaming logging!") - result = self._model_call_details_view["complete_streaming_response"] + result = call_details_view["complete_streaming_response"] greenscaleLogger.log_event( kwargs=kwargs, @@ -3161,7 +3265,7 @@ class Logging(LiteLLMLoggingBaseClass): ) if callback == "athina" and athinaLogger is not None: deep_copy = {} - for k, v in self._model_call_details_view.items(): + for k, v in call_details_view.items(): deep_copy[k] = v athinaLogger.log_event( kwargs=deep_copy, @@ -3172,7 +3276,7 @@ class Logging(LiteLLMLoggingBaseClass): ) if callback == "traceloop": deep_copy = {} - for k, v in self._model_call_details_view.items(): + for k, v in call_details_view.items(): if k != "original_response": deep_copy[k] = v traceloopLogger.log_event( @@ -3188,11 +3292,11 @@ class Logging(LiteLLMLoggingBaseClass): if s3Logger is None: s3Logger = S3Logger() if self.stream: - if "complete_streaming_response" in self.model_call_details: + if "complete_streaming_response" in call_details: print_verbose("S3Logger Logger: Got Stream Event - Completed Stream Response") s3Logger.log_event( - kwargs=self.model_call_details, - response_obj=self.model_call_details["complete_streaming_response"], + kwargs=call_details, + response_obj=call_details["complete_streaming_response"], start_time=start_time, end_time=end_time, print_verbose=print_verbose, @@ -3201,7 +3305,7 @@ class Logging(LiteLLMLoggingBaseClass): print_verbose("S3Logger Logger: Got Stream Event - No complete stream response as yet") else: s3Logger.log_event( - kwargs=self.model_call_details, + kwargs=call_details, response_obj=result, start_time=start_time, end_time=end_time, @@ -3215,19 +3319,17 @@ class Logging(LiteLLMLoggingBaseClass): openMeterLogger = OpenMeterLogger() if self.stream and complete_streaming_response is None: openMeterLogger.log_stream_event( - kwargs=self.model_call_details, + kwargs=call_details, response_obj=result, start_time=start_time, end_time=end_time, ) else: if self.stream and complete_streaming_response: - self.model_call_details["complete_response"] = self.model_call_details.get( - "complete_streaming_response", {} - ) - result = self._model_call_details_view["complete_response"] + call_details["complete_response"] = call_details.get("complete_streaming_response", {}) + result = call_details_view["complete_response"] openMeterLogger.log_success_event( - kwargs=self.model_call_details, + kwargs=call_details, response_obj=result, start_time=start_time, end_time=end_time, @@ -3240,20 +3342,18 @@ class Logging(LiteLLMLoggingBaseClass): ): # custom logger class if self.stream and complete_streaming_response is None: callback.log_stream_event( - kwargs=self.model_call_details, + kwargs=call_details, response_obj=result, start_time=start_time, end_time=end_time, ) else: if self.stream and complete_streaming_response: - self.model_call_details["complete_response"] = self.model_call_details.get( - "complete_streaming_response", {} - ) - result = self._model_call_details_view["complete_response"] + call_details["complete_response"] = call_details.get("complete_streaming_response", {}) + result = call_details_view["complete_response"] callback.log_success_event( - kwargs=self.model_call_details, + kwargs=call_details, response_obj=result, start_time=start_time, end_time=end_time, @@ -3264,7 +3364,7 @@ class Logging(LiteLLMLoggingBaseClass): print_verbose(f"success callbacks: Running Custom Callback Function - {callback}") customLogger.log_event( - kwargs=self.model_call_details, + kwargs=call_details, response_obj=result, start_time=start_time, end_time=end_time, @@ -4102,8 +4202,18 @@ class Logging(LiteLLMLoggingBaseClass): start_time, end_time, cache_hit, + for_async_request=True, ) + def handle_sync_cache_hit_success_callbacks( + self, + result: object, + start_time: datetime.datetime, + end_time: datetime.datetime, + cache_hit: bool | None = None, + ) -> None: + executor.submit(self.success_handler, result, start_time, end_time, cache_hit) + def _should_run_sync_callbacks_for_async_calls(self) -> bool: """ Returns: @@ -4115,7 +4225,10 @@ class Logging(LiteLLMLoggingBaseClass): ) _filtered_success_callbacks = self._remove_internal_custom_logger_callbacks(_combined_sync_callbacks) _filtered_success_callbacks = self._remove_internal_litellm_callbacks(_filtered_success_callbacks) - return len(_filtered_success_callbacks) > 0 + return any( + not (isinstance(callback, str) and callback in INTERNAL_CALLBACK_NAMES) + for callback in _filtered_success_callbacks # pyright: ignore[reportUnknownVariableType] # callback lists are untyped + ) def _should_run_sync_failure_callbacks_for_async_calls(self) -> bool: """ diff --git a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py index f3c60679c81..8b94c34942c 100644 --- a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py +++ b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py @@ -251,14 +251,7 @@ async def convert_to_streaming_response_async( setattr( model_response_object, "usage", - Usage( - completion_tokens=response_object["usage"].get("completion_tokens", 0), - prompt_tokens=response_object["usage"].get("prompt_tokens", 0), - total_tokens=response_object["usage"].get("total_tokens", 0), - completion_tokens_details=_cached_completion_tokens_details( - _CACHED_USAGE_DETAILS.validate_python(response_object["usage"]) - ), - ), + _replayed_usage(_CACHED_USAGE_DETAILS.validate_python(response_object["usage"])), ) if "id" in response_object: @@ -345,12 +338,10 @@ def convert_to_streaming_response( model_response_object.choices = choice_list if "usage" in response_object and response_object["usage"] is not None: - setattr(model_response_object, "usage", Usage()) - model_response_object.usage.completion_tokens = response_object["usage"].get("completion_tokens", 0) - model_response_object.usage.prompt_tokens = response_object["usage"].get("prompt_tokens", 0) - model_response_object.usage.total_tokens = response_object["usage"].get("total_tokens", 0) - model_response_object.usage.completion_tokens_details = _cached_completion_tokens_details( - _CACHED_USAGE_DETAILS.validate_python(response_object["usage"]) + setattr( + model_response_object, + "usage", + _replayed_usage(_CACHED_USAGE_DETAILS.validate_python(response_object["usage"])), ) if "id" in response_object: @@ -414,13 +405,28 @@ class _CachedMessageProviderFields(TypedDict, total=False): class _CachedUsageDetails(TypedDict, total=False): + completion_tokens: ReadOnly[object] + prompt_tokens: ReadOnly[object] + total_tokens: ReadOnly[object] completion_tokens_details: ReadOnly[object] + prompt_tokens_details: ReadOnly[object] + cache_read_input_tokens: ReadOnly[object] + cache_creation_input_tokens: ReadOnly[object] + + +class _PromptCacheTokenCounts(TypedDict, total=False): + cache_read_input_tokens: ReadOnly[int] + cache_creation_input_tokens: ReadOnly[int] _CACHED_MESSAGE_PROVIDER_FIELDS: Final = TypeAdapter(_CachedMessageProviderFields) _CACHED_USAGE_DETAILS: Final = TypeAdapter(_CachedUsageDetails) _FIELD_MAP: Final = TypeAdapter(dict[str, object]) _COMPLETION_TOKENS_DETAILS: Final = TypeAdapter(CompletionTokensDetailsWrapper) +_PROMPT_TOKENS_DETAILS: Final = TypeAdapter(PromptTokensDetailsWrapper) +_PROMPT_CACHE_TOKEN_COUNTS: Final = TypeAdapter(_PromptCacheTokenCounts) +_TOKEN_COUNT: Final = TypeAdapter(int) +_PROMPT_CACHE_TOKEN_FIELDS: Final = ("cache_read_input_tokens", "cache_creation_input_tokens") def _cached_provider_specific_fields(message: _CachedMessageProviderFields) -> Mapping[str, object] | None: @@ -438,6 +444,44 @@ def _cached_completion_tokens_details(usage: _CachedUsageDetails) -> CompletionT return None +def _cached_prompt_tokens_details(usage: _CachedUsageDetails) -> PromptTokensDetailsWrapper | None: + details: Final = usage.get("prompt_tokens_details") + if not isinstance(details, dict): + return None + try: + return _PROMPT_TOKENS_DETAILS.validate_python(details) + except ValidationError: + return None + + +def _cached_token_count(count: object) -> int | None: + try: + return _TOKEN_COUNT.validate_python(count) + except ValidationError: + return None + + +def _cached_prompt_cache_token_counts(usage: _CachedUsageDetails) -> _PromptCacheTokenCounts: + return _PROMPT_CACHE_TOKEN_COUNTS.validate_python( + { + field: count + for field in _PROMPT_CACHE_TOKEN_FIELDS + if (count := _cached_token_count(usage.get(field))) is not None + } + ) + + +def _replayed_usage(usage: _CachedUsageDetails) -> Usage: + return Usage( + completion_tokens=_cached_token_count(usage.get("completion_tokens")) or 0, + prompt_tokens=_cached_token_count(usage.get("prompt_tokens")) or 0, + total_tokens=_cached_token_count(usage.get("total_tokens")) or 0, + completion_tokens_details=_cached_completion_tokens_details(usage), + prompt_tokens_details=_cached_prompt_tokens_details(usage), + **_cached_prompt_cache_token_counts(usage), + ) + + def handle_invalid_parallel_tool_calls( tool_calls: list[ChatCompletionMessageToolCall | ChatCompletionMessageCustomToolCall], ) -> list[ChatCompletionMessageToolCall | ChatCompletionMessageCustomToolCall] | None: diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index c9b96c0a4c2..6cf0fdb3cbc 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -1,10 +1,10 @@ import asyncio import json import traceback -from collections.abc import Coroutine, Mapping, Sequence +from collections.abc import Coroutine, Mapping, MutableMapping, Sequence from dataclasses import dataclass from enum import Enum, auto -from typing import TYPE_CHECKING, Any, Final, NoReturn, Protocol, TypedDict, cast +from typing import TYPE_CHECKING, Any, Final, NoReturn, Protocol, TypedDict, cast, runtime_checkable from pydantic import TypeAdapter from typing_extensions import ReadOnly @@ -12,6 +12,7 @@ from typing_extensions import ReadOnly import litellm from litellm._logging import redact_internal_details_from_client_message, verbose_logger from litellm.constants import ( + BEDROCK_REALTIME_PENDING_OPENING_MESSAGE_SCOPE_KEY, REALTIME_METERED_USAGE_KEY, REALTIME_SESSION_FAILURE_LOGGED_KEY, REALTIME_SESSION_SUCCESS_LOGGED_KEY, @@ -98,6 +99,16 @@ class ScopedWebSocket(Protocol): def scope(self) -> _ASGIScope: ... +@runtime_checkable +class _ReplayScopeWebSocket(Protocol): + @property + def scope(self) -> MutableMapping[str, object]: ... + + +def client_websocket_scope(websocket: object) -> MutableMapping[str, object] | None: + return websocket.scope if isinstance(websocket, _ReplayScopeWebSocket) else None + + class _ClientWebSocket(ScopedWebSocket, Protocol): async def send_text(self, data: str) -> None: ... async def receive_text(self) -> str: ... @@ -154,6 +165,7 @@ class RealTimeStreaming: logging_worker: _LoggingWorker = GLOBAL_LOGGING_WORKER, ): self.websocket: _ClientWebSocket = websocket + self._replay_scope: Final = client_websocket_scope(websocket) self.backend_ws = backend_ws self.logging_obj = logging_obj self._logging_worker = logging_worker @@ -1346,6 +1358,13 @@ class RealTimeStreaming: return item async def _receive_client_message(self) -> str | None: + refused_opening_message: Final = ( + None + if self._replay_scope is None + else self._replay_scope.pop(BEDROCK_REALTIME_PENDING_OPENING_MESSAGE_SCOPE_KEY, None) + ) + if isinstance(refused_opening_message, str): + return refused_opening_message try: return await self.websocket.receive_text() except Exception as e: # noqa: BLE001 # whatever the client socket raises, the client is gone diff --git a/litellm/llms/bedrock/realtime/handler.py b/litellm/llms/bedrock/realtime/handler.py index c0a382b136c..3e7c2f7eee9 100644 --- a/litellm/llms/bedrock/realtime/handler.py +++ b/litellm/llms/bedrock/realtime/handler.py @@ -8,7 +8,7 @@ import asyncio import contextlib import importlib.metadata import json -from collections.abc import AsyncIterator, Awaitable, Callable, Mapping, MutableMapping +from collections.abc import AsyncIterator, Awaitable, Callable, Mapping, MutableMapping, Sequence from dataclasses import dataclass from types import MappingProxyType from typing import Final, NoReturn, Protocol, runtime_checkable @@ -21,7 +21,7 @@ from litellm._uuid import uuid from litellm.constants import ( BEDROCK_REALTIME_CLOSE_TIMEOUT_SECONDS, BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY, - BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY, + BEDROCK_REALTIME_PENDING_OPENING_MESSAGE_SCOPE_KEY, BEDROCK_REALTIME_SDK_DISTRIBUTION, BEDROCK_REALTIME_SDK_SUPPORTED_RANGE, BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY, @@ -37,9 +37,8 @@ from litellm.types.realtime import RealtimeResponseTransformInput from ..base_aws_llm import BaseAWSLLM, run_aws_signing from ..common_utils import BedrockError -from .transformation import BedrockRealtimeConfig, BedrockSessionUpdateRejection +from .transformation import BedrockRealtimeConfig, BedrockSessionUpdateRejection, requested_modalities -_CLIENT_MODALITIES_ADAPTER: Final[TypeAdapter["list[str] | None"]] = TypeAdapter(list[str] | None) _CLIENT_MESSAGE_ADAPTER: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) _EMPTY_JSON_OBJECT: Final[Mapping[str, JsonValue]] = MappingProxyType({}) @@ -177,25 +176,35 @@ async def _client_messages(client_ws: RealtimeClientWebSocket, initial_message: return -def _pending_session_update(scope: Mapping[str, object]) -> str | None: - """A fallback attempt on the same websocket replays the session.update the failed attempt never acked.""" - if scope.get(BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY) is True: - committed_failure: Final = scope.get(BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY) - raise BedrockError( - status_code=400, - message=( - "Bedrock realtime session already committed to a provider stream; it cannot be replayed" - + (f". The committed stream failed with: {committed_failure}" if committed_failure else "") - ), - ) - pending: Final = scope.get(BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY) +def _is_session_committed(scope: Mapping[str, object]) -> bool: + return scope.get(BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY) is True + + +def raise_if_session_committed(scope: Mapping[str, object] | None) -> None: + """A fallback attempt on a websocket whose Bedrock stream already took input would silently start over.""" + if scope is None or not _is_session_committed(scope): + return + committed_failure: Final = scope.get(BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY) + raise BedrockError( + status_code=400, + message=( + "Bedrock realtime session already committed to a provider stream; it cannot be replayed" + + (f". The committed stream failed with: {committed_failure}" if committed_failure else "") + ), + ) + + +def _pending_opening_message(scope: Mapping[str, object]) -> str | None: + """A fallback attempt on the same websocket replays the opening message Bedrock never accepted a stream for.""" + raise_if_session_committed(scope) + pending: Final = scope.get(BEDROCK_REALTIME_PENDING_OPENING_MESSAGE_SCOPE_KEY) return pending if isinstance(pending, str) else None def _raise_provider_failure(scope: MutableMapping[str, object], failure: BaseException) -> NoReturn: error: Final = _as_bedrock_error(failure) verbose_proxy_logger.error("Bedrock Realtime: provider stream failed: %s", redact_string(str(error))) - if scope.get(BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY) is True: + if _is_session_committed(scope): scope[BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY] = redact_string(str(error)) raise error from failure @@ -207,26 +216,36 @@ def _parse_client_message(message: str) -> Mapping[str, JsonValue]: return _EMPTY_JSON_OBJECT -async def _ack_session_update( +async def _send_opening_message( client_ws: RealtimeClientWebSocket, bedrock_stream: BedrockBidirectionalStream, + message: str, + bedrock_messages: Sequence[str], + send_to_bedrock: Callable[[str], Awaitable[None]], +) -> None: + """Keep the opening message for a fallback to replay until Bedrock accepts the stream it opened.""" + client_ws.scope[BEDROCK_REALTIME_PENDING_OPENING_MESSAGE_SCOPE_KEY] = message # rebind-ok: scope outlives attempts + for bedrock_message in bedrock_messages: + await send_to_bedrock(bedrock_message) + await bedrock_stream.await_output() + client_ws.scope.pop(BEDROCK_REALTIME_PENDING_OPENING_MESSAGE_SCOPE_KEY, None) + client_ws.scope[BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY] = True # rebind-ok: scope outlives the attempt + + +async def _ack_session_update( + client_ws: RealtimeClientWebSocket, transformation_config: BedrockRealtimeConfig, model: str, logging_obj: LiteLLMLogging | None, parsed_client_message: Mapping[str, JsonValue], ) -> bool: - """Ack the client's session.update once Bedrock accepted the stream; False means the client is gone.""" - await bedrock_stream.await_output() - client_ws.scope.pop(BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY, None) - client_ws.scope[BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY] = True # rebind-ok: scope outlives the attempt + """Ack the client's session.update; False means the client is gone.""" if logging_obj is None: return True - requested_modalities: Final = _CLIENT_MODALITIES_ADAPTER.validate_python( - _json_dict(parsed_client_message.get("session")).get("modalities") - ) + modalities: Final = requested_modalities(_json_dict(parsed_client_message.get("session"))) try: await client_ws.send_text( - json.dumps(transformation_config.session_updated_event(model, logging_obj, requested_modalities)) + json.dumps(transformation_config.session_updated_event(model, logging_obj, modalities)) ) except Exception as e: # noqa: BLE001 # any send failure means the client is gone verbose_proxy_logger.debug("Client to Bedrock forwarding ended: %s", e, exc_info=True) @@ -314,7 +333,7 @@ class BedrockRealtime(BaseAWSLLM): except ImportError as e: raise _sdk_import_error(self._sdk_version_lookup(), e) from e - pending_session_update: Final = _pending_session_update(websocket.scope) + pending_opening_message: Final = _pending_opening_message(websocket.scope) # Get AWS region if aws_region_name is None: @@ -376,7 +395,7 @@ class BedrockRealtime(BaseAWSLLM): try: await self._run_session( - websocket, open_bidirectional_stream, model, logging_obj, pending_session_update, drop_params + websocket, open_bidirectional_stream, model, logging_obj, pending_opening_message, drop_params ) finally: await _close_bedrock_client(bedrock_client) @@ -387,7 +406,7 @@ class BedrockRealtime(BaseAWSLLM): open_bidirectional_stream: Callable[[], Awaitable[BedrockBidirectionalStream]], model: str, logging_obj: LiteLLMLogging, - pending_session_update: str | None, + pending_opening_message: str | None, drop_params: bool = False, ) -> None: transformation_config: Final = BedrockRealtimeConfig(drop_params=drop_params) @@ -396,7 +415,7 @@ class BedrockRealtime(BaseAWSLLM): verbose_proxy_logger.debug("Bedrock Realtime: Bidirectional stream established") - if pending_session_update is None: + if pending_opening_message is None: await websocket.send_text(json.dumps(transformation_config.session_created_event(model, logging_obj))) verbose_proxy_logger.debug("Bedrock Realtime: sent session.created to client on connect") @@ -418,7 +437,7 @@ class BedrockRealtime(BaseAWSLLM): model, session_state, logging_obj, - initial_message=pending_session_update, + initial_message=pending_opening_message, ) logged_events: Final = ( @@ -531,18 +550,22 @@ class BedrockRealtime(BaseAWSLLM): if not await _reject_session_update(client_ws, parsed_client_message, rejection): break continue - client_ws.scope[BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY] = message transformed_messages = transformation_config.transform_realtime_request( message=message, model=model, session_configuration_request=session_state.get("session_configuration_request"), ) - for bedrock_message in transformed_messages: - await send_to_bedrock(bedrock_message) + if transformed_messages and not _is_session_committed(client_ws.scope): + await _send_opening_message( + client_ws, bedrock_stream, message, transformed_messages, send_to_bedrock + ) + else: + for bedrock_message in transformed_messages: + await send_to_bedrock(bedrock_message) if is_session_update and not await _ack_session_update( - client_ws, bedrock_stream, transformation_config, model, logging_obj, parsed_client_message + client_ws, transformation_config, model, logging_obj, parsed_client_message ): break finally: diff --git a/litellm/llms/bedrock/realtime/transformation.py b/litellm/llms/bedrock/realtime/transformation.py index e8bc0482a28..9cb38822881 100644 --- a/litellm/llms/bedrock/realtime/transformation.py +++ b/litellm/llms/bedrock/realtime/transformation.py @@ -123,26 +123,81 @@ def _integer_setting(value: JsonValue) -> int | None: return None if number is None else int(number) -def _malformed_number_settings(session_config: Mapping[str, JsonValue]) -> tuple[str, ...]: +def requested_modalities(session_config: Mapping[str, JsonValue]) -> tuple[str, ...] | None: + modalities: Final = session_config.get("modalities") + if not isinstance(modalities, list): + return None + names: Final = tuple(modality for modality in modalities if isinstance(modality, str)) + return names if len(names) == len(modalities) else None + + +_OPENAI_VOICES: Final = frozenset( + ("alloy", "ash", "ballad", "coral", "echo", "sage", "shimmer", "verse", "marin", "cedar") +) + + +def _unplayable_voice(session_config: Mapping[str, JsonValue]) -> str | None: + voice: Final = session_config.get("voice") + if not isinstance(voice, str): + return None + return voice if not voice.strip() or voice.lower() in _OPENAI_VOICES else None + + +def _malformed_settings(session_config: Mapping[str, JsonValue]) -> tuple[tuple[str, str], ...]: + instructions: Final = session_config.get("instructions") + voice: Final = session_config.get("voice") + input_audio_format: Final = session_config.get("input_audio_format") + output_audio_format: Final = session_config.get("output_audio_format") + modalities: Final = session_config.get("modalities") temperature: Final = session_config.get("temperature") max_output_tokens: Final = session_config.get("max_response_output_tokens") + input_rate_hertz: Final = session_config.get("input_sample_rate_hertz") output_rate_hertz: Final = session_config.get("output_sample_rate_hertz") + tools: Final = session_config.get("tools") malformed: Final = ( - ("temperature", temperature is not None and _number_setting(temperature) is None), + ("instructions", "a string", instructions is not None and not isinstance(instructions, str)), + ("voice", "a string", voice is not None and not isinstance(voice, str)), + ("tools", "a list", tools is not None and not isinstance(tools, list)), + ( + "input_audio_format", + "a string", + input_audio_format is not None and not isinstance(input_audio_format, str), + ), + ( + "output_audio_format", + "a string", + output_audio_format is not None and not isinstance(output_audio_format, str), + ), + ( + "modalities", + "a list of strings", + modalities is not None and requested_modalities(session_config) is None, + ), + ("temperature", "a number", temperature is not None and _number_setting(temperature) is None), ( "max_response_output_tokens", + "a number or 'inf'", max_output_tokens is not None and max_output_tokens != "inf" and _integer_setting(max_output_tokens) is None, ), - ("output_sample_rate_hertz", output_rate_hertz is not None and _integer_setting(output_rate_hertz) is None), + ( + "input_sample_rate_hertz", + "a number", + input_rate_hertz is not None and _integer_setting(input_rate_hertz) is None, + ), + ( + "output_sample_rate_hertz", + "a number", + output_rate_hertz is not None and _integer_setting(output_rate_hertz) is None, + ), ) - return tuple(name for name, is_malformed in malformed if is_malformed) + return tuple((name, expected) for name, expected, is_malformed in malformed if is_malformed) @dataclasses.dataclass(frozen=True, slots=True) class BedrockSessionUpdateRejection: - code: Literal["invalid_type", "session_update_not_applied"] + code: Literal["invalid_type", "invalid_value", "session_update_not_applied"] message: str param: str @@ -312,7 +367,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): return 8000 # G.711 typically uses 8kHz return 24000 if is_output else 16000 - def transform_session_update_event(self, json_message: dict) -> list[str]: + def transform_session_update_event(self, json_message: Mapping[str, JsonValue]) -> list[str]: """ Transform session.update event to Bedrock session configuration. @@ -325,7 +380,10 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): verbose_logger.debug("Handling session.update") messages: Final[list[str]] = [] - session_config: Final = json_message.get("session", {}) + requested_session: Final = json_message.get("session") + session_config: Final[Mapping[str, JsonValue]] = ( + requested_session if isinstance(requested_session, dict) else {} + ) locked_settings: Final = self._apply_locked_settings(session_config, self._started_settings) # Settings Bedrock locks at prompt start may only change before it @@ -336,13 +394,14 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): self.output_sample_rate_hertz = locked_settings.output_sample_rate_hertz # Update audio input configuration from session if provided - if "input_audio_format" in session_config: - input_format: Final = session_config["input_audio_format"] + input_format: Final = session_config.get("input_audio_format") + if isinstance(input_format, str): self.input_sample_rate_hertz = self._map_audio_format_to_sample_rate(input_format, is_output=False) # Allow direct override of the input sample rate if provided (custom extension) - if "input_sample_rate_hertz" in session_config: - self.input_sample_rate_hertz = session_config["input_sample_rate_hertz"] + requested_input_rate: Final = _integer_setting(session_config.get("input_sample_rate_hertz")) + if requested_input_rate is not None: + self.input_sample_rate_hertz = requested_input_rate if self.prompt_started: return messages @@ -352,7 +411,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): messages.append(json.dumps(session_start)) # Send prompt start - tools: Final[Sequence[JsonValue] | None] = session_config.get("tools") + tools: Final = session_config.get("tools") prompt_start_config: Final[Mapping[str, object]] = { "promptName": self.prompt_name, "textOutputConfiguration": {"mediaType": self.text_media_type}, @@ -370,7 +429,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): "toolUseOutputConfiguration": {"mediaType": "application/json"}, "toolConfiguration": {"tools": self._transform_tools_to_bedrock_format(tools)}, } - if tools + if isinstance(tools, list) and tools else {} ), } @@ -445,12 +504,14 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): if isinstance(requested_tools, list) else base.tools_json ) - system_prompt: Final = base.system_prompt if instructions is None else str(instructions) - voice_id: Final = base.voice_id if voice is None else str(voice) + system_prompt: Final = instructions if isinstance(instructions, str) else base.system_prompt + voice_id: Final = ( + voice if isinstance(voice, str) and _unplayable_voice(session_config) is None else base.voice_id + ) format_sample_rate: Final = ( - base.output_sample_rate_hertz - if output_audio_format is None - else self._map_audio_format_to_sample_rate(str(output_audio_format), is_output=True) + self._map_audio_format_to_sample_rate(output_audio_format, is_output=True) + if isinstance(output_audio_format, str) + else base.output_sample_rate_hertz ) requested_output_rate: Final = _integer_setting(session_config.get("output_sample_rate_hertz")) output_sample_rate_hertz: Final = ( @@ -476,16 +537,26 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): session_config: Final = json_message.get("session") if not isinstance(session_config, Mapping): return None - malformed: Final = () if self.drop_params else _malformed_number_settings(session_config) + malformed: Final = () if self.drop_params else _malformed_settings(session_config) if malformed: + name, expected = malformed[0] return BedrockSessionUpdateRejection( code="invalid_type", message=( - f"Invalid type for 'session.{malformed[0]}': Bedrock Nova Sonic needs a number for temperature " - "and output_sample_rate_hertz, and a number or 'inf' for max_response_output_tokens. " + f"Invalid type for 'session.{name}': Bedrock Nova Sonic needs {expected} there. " "Set drop_params to drop it instead" ), - param=f"session.{malformed[0]}", + param=f"session.{name}", + ) + unplayable_voice: Final = None if self.drop_params else _unplayable_voice(session_config) + if unplayable_voice is not None: + return BedrockSessionUpdateRejection( + code="invalid_value", + message=( + f"Invalid value for 'session.voice': '{unplayable_voice}' is not a Bedrock Nova Sonic voice. " + "Set drop_params to drop it instead" + ), + param="session.voice", ) if not self.prompt_started: return None @@ -785,7 +856,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): List of Bedrock format messages (JSON strings) """ try: - json_message: Final[dict[str, object]] = json.loads(message) + json_message: Final[dict[str, JsonValue]] = json.loads(message) except json.JSONDecodeError: verbose_logger.warning("Invalid JSON message: %s", message[:200]) return [] @@ -816,11 +887,11 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): self, model: str, logging_obj: LiteLLMLoggingObj, - modalities: list[str] | None = None, + modalities: Sequence[str] | None = None, ) -> OpenAIRealtimeStreamSession: session: Final = OpenAIRealtimeStreamSession( id=logging_obj.litellm_trace_id, - modalities=modalities if modalities is not None else ["text", "audio"], + modalities=list(modalities) if modalities is not None else ["text", "audio"], ) if model is not None and isinstance(model, str): session["model"] = model @@ -842,7 +913,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): self, model: str, logging_obj: LiteLLMLoggingObj, - modalities: list[str] | None = None, + modalities: Sequence[str] | None = None, ) -> OpenAIRealtimeStreamSessionEvents: """Build the OpenAI session.updated ack reflecting the client's requested modalities.""" return OpenAIRealtimeStreamSessionEvents( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index bbd9d96b84b..31692857787 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -75,6 +75,7 @@ from litellm.constants import ( DEFAULT_SHARED_HEALTH_CHECK_LOCK_TTL, DEFAULT_SHARED_HEALTH_CHECK_TTL, DEFAULT_SLACK_ALERTING_THRESHOLD, + INTERNAL_CALLBACK_NAMES, LANGFUSE_SHUTDOWN_FLUSH_TIMEOUT_MILLIS, LITELLM_EMBEDDING_PROVIDERS_SUPPORTING_INPUT_ARRAY_OF_TOKENS, LITELLM_SETTINGS_SAFE_DB_OVERRIDES, @@ -19812,8 +19813,7 @@ def _hidden_runtime_callback_names(configured_callback_names: frozenset[str]) -> if name in CustomLoggerRegistry.CALLBACK_CLASS_STR_TO_CLASS_TYPE ) configured_modules: Final = frozenset(name.rsplit(".", 1)[0] for name in configured_callback_names if "." in name) - internal_callback_names: Final = frozenset({"cache", "vector_store_pre_call_hook"}) - return internal_callback_names | frozenset( + return INTERNAL_CALLBACK_NAMES | frozenset( callback_name for callback_name, callback in litellm.logging_callback_manager.get_callback_objects() if isinstance(callback, CustomGuardrail) diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py index 367d5be7bfc..362191b4642 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -37,9 +37,10 @@ from ..litellm_core_utils.core_helpers import normalize_drop_params from ..litellm_core_utils.credential_accessor import CredentialAccessor from ..litellm_core_utils.get_litellm_params import get_litellm_params from ..litellm_core_utils.litellm_logging import Logging as LiteLLMLogging +from ..litellm_core_utils.realtime_streaming import client_websocket_scope from ..llms.azure.common_utils import get_azure_ad_token from ..llms.azure.realtime.handler import AzureOpenAIRealtime, azure_realtime_protocol_for_client -from ..llms.bedrock.realtime.handler import BedrockRealtime +from ..llms.bedrock.realtime.handler import BedrockRealtime, raise_if_session_committed from ..llms.custom_httpx.http_handler import realtime_ssl_for_url from ..llms.openai.live.handler import OpenAILiveSessions from ..llms.openai.realtime.handler import OpenAIRealtime @@ -451,6 +452,7 @@ async def _arealtime( For PROXY use only. """ + raise_if_session_committed(client_websocket_scope(websocket)) headers = cast(dict | None, kwargs.get("headers")) extra_headers: Final = cast(dict | None, kwargs.get("extra_headers")) if headers is None: diff --git a/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py b/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py index 091472c1ea3..b0da3ed264a 100644 --- a/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py +++ b/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py @@ -4,8 +4,10 @@ Check if prompt caching is valid for a given deployment Route to previously cached model id, if valid """ -from collections.abc import Callable -from typing import Final, cast +from collections.abc import Callable, Mapping +from typing import Final, TypedDict, cast + +from typing_extensions import ReadOnly from litellm import verbose_logger from litellm.caching.dual_cache import DualCache @@ -16,7 +18,7 @@ from litellm.integrations.anthropic_cache_control_hook import ( ) from litellm.integrations.custom_logger import CustomLogger, Span from litellm.litellm_core_utils.token_counter import offload_token_count -from litellm.llms.bedrock.common_utils import supports_bedrock_prompt_cache_ttl +from litellm.llms.bedrock.common_utils import bedrock_cache_point_ttl_supported, bedrock_capability_model from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import CallTypes, StandardLoggingPayload from litellm.utils import get_prompt_cache_min_tokens, is_prompt_caching_valid_prompt @@ -24,6 +26,11 @@ from litellm.utils import get_prompt_cache_min_tokens, is_prompt_caching_valid_p from ..prompt_caching_cache import PromptCachingCache +class _PromptCachingSuccessEventKwargs(TypedDict, total=False): + standard_logging_object: ReadOnly[StandardLoggingPayload | None] + litellm_params: ReadOnly[Mapping[str, object] | None] + + def _get_min_token_count_for_deployments(healthy_deployments: list[dict]) -> int: """ Returns the lowest minimum cacheable prefix across a model group. @@ -49,9 +56,13 @@ def _get_min_token_count_for_deployments(healthy_deployments: list[dict]) -> int ) -def _honors_one_hour_cache_ttl(standard_logging_object: StandardLoggingPayload) -> bool: +def _honors_one_hour_cache_ttl( + standard_logging_object: StandardLoggingPayload, litellm_params: Mapping[str, object] +) -> bool: if standard_logging_object["custom_llm_provider"] == "bedrock": - return supports_bedrock_prompt_cache_ttl(standard_logging_object["model"]) + deployment_model: Final = standard_logging_object["model"] + capability_model: Final = bedrock_capability_model(deployment_model, litellm_params.get("base_model")) + return bedrock_cache_point_ttl_supported(capability_model, deployment_model) model_map_value: Final = standard_logging_object["model_map_information"]["model_map_value"] return model_map_value is not None and (model_map_value.get("cache_creation_input_token_cost_above_1hr") or 0) > 0 @@ -121,7 +132,13 @@ class PromptCachingDeploymentCheck(CustomLogger): return healthy_deployments - async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + async def async_log_success_event( + self, + kwargs: _PromptCachingSuccessEventKwargs, + response_obj: object, + start_time: object, + end_time: object, + ) -> None: standard_logging_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object", None) if standard_logging_object is None: @@ -163,11 +180,15 @@ class PromptCachingDeploymentCheck(CustomLogger): cache: Final = PromptCachingCache( cache=self.cache, ) + raw_litellm_params: Final = kwargs.get("litellm_params") + litellm_params: Final[Mapping[str, object]] = ( + raw_litellm_params if isinstance(raw_litellm_params, Mapping) else {} + ) await cache.async_add_model_id( model_id=model_id, messages=prompt_messages, tools=None, # [TODO]: add tools once standard_logging_object supports it - honors_one_hour_ttl=_honors_one_hour_cache_ttl(standard_logging_object), + honors_one_hour_ttl=_honors_one_hour_cache_ttl(standard_logging_object, litellm_params), ) return diff --git a/litellm/rust_bridge/host/callbacks_legacy_python.py b/litellm/rust_bridge/host/callbacks_legacy_python.py index bfd16b29bf6..a82e734b47d 100644 --- a/litellm/rust_bridge/host/callbacks_legacy_python.py +++ b/litellm/rust_bridge/host/callbacks_legacy_python.py @@ -151,6 +151,8 @@ class LoggingSurface(Protocol): start_time: datetime.datetime | None = None, end_time: datetime.datetime | None = None, cache_hit: bool | None = None, + *, + for_async_request: bool = False, **kwargs: object, ) -> None: ... diff --git a/litellm/utils.py b/litellm/utils.py index 65408663f6f..9dc33637e5b 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -5936,6 +5936,15 @@ def _get_model_info_from_generalization( return None +def unmapped_model_generalization(model: str, custom_llm_provider: str | None) -> Mapping[str, object] | None: + generalization: Final = _get_model_info_from_generalization( + model=model, + potential_model_names=get_potential_model_names(model, custom_llm_provider), + custom_llm_provider=custom_llm_provider, + ) + return generalization[1] if generalization is not None else None + + def _strip_mantle_region_prefix(model: str) -> str: from litellm.llms.bedrock_mantle.common_utils import split_mantle_region_prefix diff --git a/tests/e2e/router/test_reliability_memory_e2e.py b/tests/e2e/router/test_reliability_memory_e2e.py index 9ae6edb3eb2..196784b5345 100644 --- a/tests/e2e/router/test_reliability_memory_e2e.py +++ b/tests/e2e/router/test_reliability_memory_e2e.py @@ -13,19 +13,26 @@ deployment that fails a handful of times in a row is benched (a bad-credential 4 included), the router answers "No deployments available" without retrying, and the retry loop that leaks stops running; cooldown_time 0 keeps it running. -Two identical phases run back to back. The first is the warmup that grows the -proxy's caches and allocator arenas to their steady state, the second is the one -the budget applies to, so a healthy proxy shows the second phase adding roughly -nothing while a leaking one adds a fixed amount per request. RSS is read through -/debug/memory/summary on every configured replica; a burst of failing calls leaves -a transient bulge of garbage that gc reclaims within seconds, so each checkpoint -samples until no new worker has answered for a settle window and keeps the lowest -reading per worker. The growth is judged per worker (by replica address, hostname -and pid, since pods in their own pid namespaces report the same pids) so each -worker is compared with itself, and the two checkpoints must see the same workers: -a single load-balanced address reaches the workers behind it one answer at a time, -and a worker that answered only one checkpoint would otherwise drop out of the -comparison, which is where a leaking worker could hide. +Three identical batches run back to back. The first two are the warmup that grows +the proxy's caches and allocator arenas to their steady state, the third is the one +the budget applies to, so a healthy proxy shows the third batch adding roughly +nothing while a leaking one adds a fixed amount per request. One warmup batch is not +enough: each batch's in-flight retries peak well above the live heap, the allocator +keeps that high-water mark, and on a cold worker the second batch still set a new one +(about 50 MB of RSS with the live Python heap flat to within 1 MB), where the later +batches each settled under 25 MB. RSS is read through /debug/memory/summary on every +configured replica; a burst of failing calls leaves a transient bulge of garbage that +gc reclaims within seconds, so each checkpoint samples until no new worker has +answered for a settle window and keeps the lowest reading per worker. A checkpoint +runs before the warmup and after every batch, and each batch's settled growth per +worker rides out as a repeated `rss_batch_growth` JUnit property as soon as that +batch settles, so a run that fails keeps the curve up to the batch it failed in. The budget is judged per worker +(by replica address, hostname and pid, since pods in their own pid namespaces report +the same pids) so each worker is compared with itself, and the checkpoints around the +measured batch must see the same workers: a single load-balanced address reaches the +workers behind it one answer at a time, and a worker that answered only one +checkpoint would otherwise drop out of the comparison, which is where a leaking +worker could hide. RSS alone is a coarse gauge: on the release stack (spend logs storing prompts, json logs, prometheus and otel callbacks) the same v1.100.0 breadcrumbs grew RSS @@ -58,7 +65,7 @@ from __future__ import annotations import json import time -from collections.abc import Mapping, Sequence +from collections.abc import Callable, Mapping, Sequence from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass from types import MappingProxyType @@ -91,6 +98,8 @@ pytestmark = [pytest.mark.e2e, pytest.mark.quiet_stack] DEPLOYMENTS_PER_GROUP: Final = 2 RSS_SAMPLE_CAP: Final = 4 * MEMORY_RSS_SETTLE_SAMPLES +WARMUP_BATCHES: Final = 2 +MEASURED_BATCH: Final = WARMUP_BATCHES + 1 @dataclass(frozen=True, slots=True) @@ -177,12 +186,52 @@ def _settled_rss_per_worker(proxy: ProxyClient) -> Mapping[WorkerKey, RssReading ) +def _record_batch_growth( + record_property: Callable[[str, object], None], + batch: int, + before: Mapping[WorkerKey, RssReading], + after: Mapping[WorkerKey, RssReading], +) -> None: + for growth in (WorkerGrowth(before[worker], after[worker]) for worker in before.keys() & after.keys()): + record_property( + "rss_batch_growth", + json.dumps( + { + "batch": batch, + "worker": growth.warm.where, + "before_mb": growth.warm.ram_usage_mb, + "after_mb": growth.after.ram_usage_mb, + "growth_mb": round(growth.growth_mb, 2), + } + ), + ) + + +def _settled_rss_around_measured_batch( + proxy: ProxyClient, + key: str, + model: str, + override: RouterSettingsOverride, + fallback: str, + record_property: Callable[[str, object], None], + before: Mapping[WorkerKey, RssReading], + batch: int, +) -> tuple[Mapping[WorkerKey, RssReading], Mapping[WorkerKey, RssReading]]: + calls: Final = _fail_many(proxy, key, model, override) + _assert_every_call_failed_through_fallback(calls, fallback) + after: Final = _settled_rss_per_worker(proxy) + _record_batch_growth(record_property, batch, before, after) + if batch == MEASURED_BATCH: + return before, after + return _settled_rss_around_measured_batch(proxy, key, model, override, fallback, record_property, after, batch + 1) + + def _heaviest_worker_growth( warm: Mapping[WorkerKey, RssReading], after: Mapping[WorkerKey, RssReading] ) -> WorkerGrowth: assert warm.keys() == after.keys(), ( - f"the workers answering /debug/memory/summary changed between the checkpoints, so not every worker can " - f"be compared with itself: gone after the measured batch {sorted(warm.keys() - after.keys())} (a worker " + f"the workers answering /debug/memory/summary changed around measured batch {MEASURED_BATCH}, so not every " + f"worker can be compared with itself: gone after it {sorted(warm.keys() - after.keys())} (a worker " f"that died or was restarted under failing traffic, which is what an OOM kill looks like), first seen " f"after it {sorted(after.keys() - warm.keys())} (the warm window never reached them, so they have no " f"baseline; raise E2E_MEMORY_RSS_SETTLE_SAMPLES if the stack has more workers than the window covers)" @@ -247,6 +296,7 @@ class TestReliabilityMemory: f"leaves the worker little room for traffic" ) + @pytest.mark.filterwarnings("ignore:record_property is incompatible with junit_family") @pytest.mark.covers("reliability.perf.memory.under_slo") @meta( Subject( @@ -257,7 +307,11 @@ class TestReliabilityMemory: ) ) def test_failing_requests_do_not_grow_rss_or_stored_request( - self, client: ComplexityRouterClient, resources: ResourceManager, scoped_key: str + self, + client: ComplexityRouterClient, + resources: ResourceManager, + scoped_key: str, + record_property: Callable[[str, object], None], ) -> None: marker: Final = unique_marker() primary: Final = f"reliability-memory-{marker}" @@ -278,19 +332,23 @@ class TestReliabilityMemory: f"request into the stored snapshot the way the v1.100.0 ones did" ) - warmup: Final = _fail_many(client.proxy, scoped_key, primary, override) - _assert_every_call_failed_through_fallback(warmup, fallback) - warm: Final = _settled_rss_per_worker(client.proxy) - - measured: Final = _fail_many(client.proxy, scoped_key, primary, override) - _assert_every_call_failed_through_fallback(measured, fallback) - after: Final = _settled_rss_per_worker(client.proxy) + warm, after = _settled_rss_around_measured_batch( + client.proxy, + scoped_key, + primary, + override, + fallback, + record_property, + _settled_rss_per_worker(client.proxy), + 1, + ) heaviest: Final = _heaviest_worker_growth(warm, after) assert heaviest.growth_mb <= MEMORY_RSS_BUDGET_MB, ( - f"proxy RSS grew {heaviest.growth_mb:.1f} MB over a second batch of {MEMORY_REQUESTS_PER_PHASE} failing " - f"requests ({MEMORY_RETRIES_PER_REQUEST} retries each plus a fallback) after an identical warmup batch, " - f"past the {MEMORY_RSS_BUDGET_MB:.0f} MB budget: worker pid {heaviest.warm.worker_pid} on " + f"proxy RSS grew {heaviest.growth_mb:.1f} MB over measured batch {MEASURED_BATCH} of " + f"{MEMORY_REQUESTS_PER_PHASE} failing requests ({MEMORY_RETRIES_PER_REQUEST} retries each plus a fallback) " + f"after {WARMUP_BATCHES} identical warmup batches, past the {MEMORY_RSS_BUDGET_MB:.0f} MB budget: " + f"worker pid {heaviest.warm.worker_pid} on " f"{heaviest.warm.hostname or 'an unnamed host'} behind {heaviest.warm.replica} settled at " f"{heaviest.warm.ram_usage_mb:.1f} MB warm and " f"{heaviest.after.ram_usage_mb:.1f} MB after; failing requests are leaking memory the way the " diff --git a/tests/e2e/router/test_reliability_prompt_caching_e2e.py b/tests/e2e/router/test_reliability_prompt_caching_e2e.py index 7bf23bf967f..25e842f362e 100644 --- a/tests/e2e/router/test_reliability_prompt_caching_e2e.py +++ b/tests/e2e/router/test_reliability_prompt_caching_e2e.py @@ -43,6 +43,7 @@ FOLLOW_UPS = 3 class TestReliabilityPromptCachingAffinity: + @pytest.mark.provider_live @pytest.mark.covers("reliability.cache.prompt_caching_model_select.returns_cached") @meta( Subject( diff --git a/tests/integration/_support/bedrock_realtime_peer.py b/tests/integration/_support/bedrock_realtime_peer.py index d30bf1b92f5..ff618a00136 100644 --- a/tests/integration/_support/bedrock_realtime_peer.py +++ b/tests/integration/_support/bedrock_realtime_peer.py @@ -14,6 +14,12 @@ text message (the injected trigger-audio block is USER/AUDIO, so it never trips ``tool_result`` fires on the Nth inbound ``toolResult`` event, which the proxy sends inside a TOOL contentStart/toolResult/contentEnd triple whose ``toolResultInputConfiguration`` ``toolUseId`` names the call being answered. + +A ``Refusal`` peer withholds each stream's initial response until that stream's opening arrived +(its SYSTEM prompt block or its first USER text block closed), then answers the HTTP error Bedrock +sends for an unknown model id. A ``StreamFailure`` peer accepts the stream, plays its turns, and +then sends an event-stream exception frame and ends the stream, as Bedrock does for a mid-call +model error. ``streams()`` groups the inbound events per opened stream. """ from __future__ import annotations @@ -26,22 +32,24 @@ import uuid from collections.abc import Iterator, Mapping from contextlib import contextmanager, suppress from dataclasses import dataclass -from itertools import chain +from itertools import chain, count from pathlib import Path from queue import SimpleQueue -from typing import Final, Literal +from typing import Final, Literal, TypeVar import h2.config import h2.connection import h2.events import h2.exceptions from integration._support.async_server import Handler, LoopbackServer -from integration._support.process import _free_port +from integration._support.ports import allocate_port from integration._support.tls import server_context from integration._support.upstream import aws_event_stream_frame from pydantic import JsonValue, TypeAdapter BedrockEvent = Mapping[str, JsonValue] +OpeningEnd = Literal["system_prompt", "user_input"] +Recorded = TypeVar("Recorded") _JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) @@ -177,6 +185,20 @@ class Turn: events: tuple[BedrockEvent, ...] +@dataclass(frozen=True, slots=True) +class Refusal: + after: OpeningEnd + status: int = 400 + error_type: str = "ValidationException" + message: str = "The provided model identifier is invalid." + + +@dataclass(frozen=True, slots=True) +class StreamFailure: + exception_type: str = "modelStreamErrorException" + message: str = "The system encountered an unexpected error during processing. Try your request again." + + def _event_body(event: BedrockEvent) -> tuple[str, Mapping[str, JsonValue]]: body: Final = event["event"] assert isinstance(body, Mapping) and len(body) == 1, event @@ -214,6 +236,17 @@ def _frame_payload(raw: bytes) -> bytes: return raw[12 + header_length : -4] +def _exception_frame(failure: StreamFailure) -> bytes: + return aws_event_stream_frame( + { + ":message-type": "exception", + ":exception-type": failure.exception_type, + ":content-type": "application/json", + }, + json.dumps({"message": failure.message}).encode(), + ) + + def _decode_inbound(chunk_payload: bytes) -> BedrockEvent | None: """One signed chunk frame's payload -> the inner chunk's {"bytes": ...} event.""" inner_payload: Final = _frame_payload(chunk_payload) @@ -239,6 +272,7 @@ class _Inbound: self.tool_use_ids: Final[dict[str, str]] = {} self.tool_results_seen = 0 self.user_text_closed = False + self.system_text_closed = False def record(self, event: BedrockEvent) -> None: self.events.put(event) @@ -253,6 +287,8 @@ class _Inbound: role_type: Final = self.open_blocks.pop(content_name, ("", "")) if role_type == ("USER", "TEXT"): self.user_text_closed = True + elif role_type == ("SYSTEM", "TEXT"): + self.system_text_closed = True elif name == "toolResult": self.tool_results_seen += 1 @@ -283,10 +319,35 @@ class _Script: return turn.events +def _opening_arrived(after: OpeningEnd, inbound: _Inbound) -> bool: + match after: + case "system_prompt": + return inbound.system_text_closed + case "user_input": + return inbound.user_text_closed + + class _PeerState: - def __init__(self, turns: tuple[Turn, ...]) -> None: + def __init__(self, turns: tuple[Turn, ...], refusal: Refusal | None, failure: StreamFailure | None) -> None: self.inbound: Final = _Inbound() self.script: Final = _Script(turns) + self.refusal: Final = refusal + self.failure: Final = failure + self.ordinals: Final = count() + self.opened: Final = SimpleQueue[int]() + self.stream_events: Final = SimpleQueue[tuple[int, BedrockEvent]]() + + +def _drain(queue: SimpleQueue[Recorded], log: list[Recorded]) -> tuple[Recorded, ...]: + while not queue.empty(): + log.append(queue.get()) # drained log, not a built value + return tuple(log) + + +@dataclass(frozen=True, slots=True) +class _Stream: + ordinal: int + inbound: _Inbound @dataclass(frozen=True, slots=True) @@ -294,11 +355,16 @@ class BedrockRealtimePeer: api_base: str _state: _PeerState _drained: list[BedrockEvent] + _opened: list[int] + _stream_events: list[tuple[int, BedrockEvent]] def received(self) -> tuple[BedrockEvent, ...]: - while not self._state.inbound.events.empty(): - self._drained.append(self._state.inbound.events.get()) # drained log, not a built value - return tuple(self._drained) + return _drain(self._state.inbound.events, self._drained) + + def streams(self) -> tuple[tuple[BedrockEvent, ...], ...]: + opened: Final = _drain(self._state.opened, self._opened) + recorded: Final = _drain(self._state.stream_events, self._stream_events) + return tuple(tuple(event for ordinal, event in recorded if ordinal == stream) for stream in opened) def tool_specs(self) -> tuple[Mapping[str, JsonValue], ...]: return tuple(chain.from_iterable(_tool_specs(event) for event in self.received())) @@ -344,6 +410,7 @@ class _H2Session: self._writer: Final = writer self._h2: Final = h2.connection.H2Connection(config=h2.config.H2Configuration(client_side=False)) self._buffers: Final[dict[int, bytes]] = {} + self._streams: Final[dict[int, _Stream]] = {} self._ended: set[int] = set() self._closed = False @@ -367,10 +434,14 @@ class _H2Session: def _apply(self, event: h2.events.Event) -> None: if isinstance(event, h2.events.RequestReceived): - self._h2.send_headers( - event.stream_id, - [(":status", "200"), ("content-type", "application/vnd.amazon.eventstream")], - ) + ordinal: Final = next(self._state.ordinals) + self._streams[event.stream_id] = _Stream(ordinal, _Inbound()) + self._state.opened.put(ordinal) + if self._state.refusal is None: + self._h2.send_headers( + event.stream_id, + [(":status", "200"), ("content-type", "application/vnd.amazon.eventstream")], + ) elif isinstance(event, h2.events.DataReceived): self._h2.acknowledge_received_data(event.flow_controlled_length, event.stream_id) self._buffers[event.stream_id] = self._buffers.get(event.stream_id, b"") + event.data @@ -380,15 +451,44 @@ class _H2Session: decoded: Final = _decode_inbound(payload) if decoded is None: continue - self._state.inbound.record(decoded) - if event.stream_id not in self._ended: - for outbound in self._state.script.ready_events(self._state.inbound, decoded): - self._h2.send_data(event.stream_id, _outbound_chunk(outbound)) + self._record(event.stream_id, decoded) elif isinstance(event, (h2.events.StreamEnded, h2.events.StreamReset)): self._close_stream(event.stream_id) elif isinstance(event, h2.events.ConnectionTerminated): self._writer.close() + def _record(self, stream_id: int, decoded: BedrockEvent) -> None: + stream: Final = self._streams[stream_id] + self._state.inbound.record(decoded) + stream.inbound.record(decoded) + self._state.stream_events.put((stream.ordinal, decoded)) + if stream_id in self._ended: + return + refusal: Final = self._state.refusal + if refusal is not None: + if _opening_arrived(refusal.after, stream.inbound): + self._refuse(stream_id, refusal) + return + outbound_events: Final = self._state.script.ready_events(self._state.inbound, decoded) + for outbound in outbound_events: + self._h2.send_data(stream_id, _outbound_chunk(outbound)) + failure: Final = self._state.failure + if outbound_events and failure is not None and self._state.script.exhausted: + self._h2.send_data(stream_id, _exception_frame(failure)) + self._end_stream(stream_id) + + def _refuse(self, stream_id: int, refusal: Refusal) -> None: + self._h2.send_headers( + stream_id, + [ + (":status", str(refusal.status)), + ("content-type", "application/json"), + ("x-amzn-errortype", refusal.error_type), + ], + ) + self._h2.send_data(stream_id, json.dumps({"message": refusal.message}).encode()) + self._end_stream(stream_id) + def _close_stream(self, stream_id: int) -> None: self._end_stream(stream_id) self._closed = True @@ -411,19 +511,27 @@ def _session_handler(state: _PeerState) -> Handler: @contextmanager -def bedrock_realtime_peer(turns: tuple[Turn, ...], cert: tuple[Path, Path]) -> Iterator[BedrockRealtimePeer]: +def bedrock_realtime_peer( + turns: tuple[Turn, ...], + cert: tuple[Path, Path], + *, + refusal: Refusal | None = None, + failure: StreamFailure | None = None, +) -> Iterator[BedrockRealtimePeer]: """Serve scripted `turns` on a loopback TLS h2 endpoint; yield the peer. `cert` is the self-signed pair the owned proxy already trusts via SSL_CERT_FILE. SigV4 headers are ignored; the stream closes cleanly when the client ends it. """ - port: Final = _free_port() + port: Final = allocate_port() context: Final = server_context(cert[0], cert[1]) context.set_alpn_protocols(["h2"]) - state: Final = _PeerState(turns) + state: Final = _PeerState(turns, refusal, failure) server: Final = LoopbackServer(port, _session_handler(state), ssl_context=context) server.start() try: - yield BedrockRealtimePeer(api_base=f"https://localhost:{port}", _state=state, _drained=[]) + yield BedrockRealtimePeer( + api_base=f"https://localhost:{port}", _state=state, _drained=[], _opened=[], _stream_events=[] + ) finally: server.stop() diff --git a/tests/integration/_support/process.py b/tests/integration/_support/process.py index fd61ca3aea2..00b9dca5184 100644 --- a/tests/integration/_support/process.py +++ b/tests/integration/_support/process.py @@ -628,8 +628,11 @@ def worker_configuration(config: Path, directory: Path, inherited: Mapping[str, return f"http://127.0.0.1:{inherited['INTEGRATION_PROVIDER_PORT']}" contents: Final = re.sub(r"http://127[.]0[.]0[.]1:(8190|8191)(?![0-9])", endpoint, config.read_text()) - directory.mkdir(parents=True, exist_ok=True) - destination: Final = directory / f"worker-config-{uuid.uuid4().hex}.yaml" + checkouts: Final = (_harness_root(), _proxy_root().resolve()) + inside_checkout: Final = any(config.resolve().is_relative_to(checkout) for checkout in checkouts) + target: Final = directory if inside_checkout else config.parent + target.mkdir(parents=True, exist_ok=True) + destination: Final = target / f"worker-config-{uuid.uuid4().hex}.yaml" loaded: Final[object] = yaml.safe_load(contents) if not isinstance(loaded, dict): destination.write_text(contents) diff --git a/tests/integration/compatibility/test_missing_body_param_status.py b/tests/integration/compatibility/test_missing_body_param_status.py index 42ee063ddc7..dd98ff6dea1 100644 --- a/tests/integration/compatibility/test_missing_body_param_status.py +++ b/tests/integration/compatibility/test_missing_body_param_status.py @@ -1695,15 +1695,19 @@ def test_upstream_pause_and_worker_kill_preserve_required_body_status( record_property("chaos_spend_row_counts", str(tuple(len(rows) for rows in spend_rows))) with owned.replacing_worker(min(workers)) as kill: kill() - post_kill: Final = tuple( - _post(candidate, path, body) for path, body, _expected, _identity in missing_calls[:5] - ) - assert all( - response.status_code == 400 and response.json() == expected - for response, (_path, _body, expected, _identity) in zip(post_kill, missing_calls[:5]) - ), [response.text for response in post_kill] - recovered: Final = _post(candidate, "/v1/chat/completions", _chat_body(chat_model, "recovered", True)) - assert recovered.status_code == 200, recovered.text + with httpx.Client(base_url=candidate.client.base_url, timeout=15, trust_env=False) as fresh: + survivor: Final = Gateway(fresh, candidate.key, candidate.upstream_url) + post_kill: Final = tuple( + _post(survivor, path, body) for path, body, _expected, _identity in missing_calls[:5] + ) + assert all( + response.status_code == 400 and response.json() == expected + for response, (_path, _body, expected, _identity) in zip(post_kill, missing_calls[:5]) + ), [response.text for response in post_kill] + recovered: Final = _post( + survivor, "/v1/chat/completions", _chat_body(chat_model, "recovered", True) + ) + assert recovered.status_code == 200, recovered.text def test_eval_create_and_run_forward_client_metadata(gateway: Gateway) -> None: diff --git a/tests/integration/cost_calculation/cases/misc.py b/tests/integration/cost_calculation/cases/misc.py index c3dc0225d0d..91fd707a9b2 100644 --- a/tests/integration/cost_calculation/cases/misc.py +++ b/tests/integration/cost_calculation/cases/misc.py @@ -702,6 +702,31 @@ BEDROCK_INVOKE_HAIKU_STREAM: Final = CostTrackingTestCase( ), ) +BEDROCK_INVOKE_HAIKU_CHAT_STREAM: Final = CostTrackingTestCase( + name="bedrock-invoke-haiku-chat-stream", + covers="quota_management.spend_tracking.scripted_wire.logs_cost", + model="bedrock/invoke/anthropic.claude-haiku-4-5-20251001-v1:0", + request={ + "model": "$MODEL", + "messages": [{"role": "user", "content": "bedrock-invoke-haiku-chat-stream"}], + "stream": True, + "stream_options": {"include_usage": True}, + "max_tokens": 412, + }, + response=BEDROCK_INVOKE_HAIKU_STREAM.response, + expected=BEDROCK_INVOKE_HAIKU_STREAM.expected, +) + +BEDROCK_INVOKE_HAIKU_RESPONSES: Final = CostTrackingTestCase( + name="bedrock-invoke-haiku-responses", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="bedrock/invoke/anthropic.claude-haiku-4-5-20251001-v1:0", + endpoint="/v1/responses", + request={"model": "$MODEL", "input": "bedrock-invoke-haiku-responses", "max_output_tokens": 412}, + response=BEDROCK_INVOKE_HAIKU_JSON.response, + expected=BEDROCK_INVOKE_HAIKU_JSON.expected, +) + # azure_ai/gpt-5.4-mini-2026-03-17 AZURE_AI_GPT_5_4_MINI_LATEST: Final = CostTrackingTestCase( @@ -1741,6 +1766,8 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( XAI_GROK_5_PROVIDER_REPORTED_COST, BEDROCK_INVOKE_HAIKU_JSON, BEDROCK_INVOKE_HAIKU_STREAM, + BEDROCK_INVOKE_HAIKU_CHAT_STREAM, + BEDROCK_INVOKE_HAIKU_RESPONSES, AZURE_AI_GPT_5_4_MINI_LATEST, AZURE_AI_GPT_5_4_MINI_LATEST_STREAM, AZURE_AI_CLAUDE_HAIKU_4_5_MESSAGES_WEB_SEARCH, diff --git a/tests/integration/management/test_scim_group_bulk_roster.py b/tests/integration/management/test_scim_group_bulk_roster.py index 8387beca958..2af6ae0a83d 100644 --- a/tests/integration/management/test_scim_group_bulk_roster.py +++ b/tests/integration/management/test_scim_group_bulk_roster.py @@ -18,7 +18,6 @@ from pathlib import Path from typing import Final, LiteralString import httpx -import psutil import psycopg import pytest from integration._support.budget import case_timeout_seconds @@ -55,9 +54,6 @@ LOCK_KEY_ROW_SQL: Final = 'SELECT token FROM "LiteLLM_VerificationToken" WHERE t SESSIONS_BLOCKED_BY_SQL: Final = ( "SELECT count(*)::int AS waiting FROM pg_stat_activity WHERE %s::int = ANY(pg_blocking_pids(pid))" ) -CLIENT_PORTS_BLOCKED_BY_SQL: Final = ( - "SELECT client_port FROM pg_stat_activity WHERE %s::int = ANY(pg_blocking_pids(pid)) AND client_port > 0" -) TRANSACTIONS_OPEN_SINCE_SQL: Final = """ SELECT count(*)::int AS open FROM pg_stat_activity WHERE datname = current_database() AND backend_type = 'client backend' AND pid <> pg_backend_pid() @@ -296,22 +292,6 @@ def _transactions_open_since(moment: str) -> int: return _int_value(read_rows(TRANSACTIONS_OPEN_SINCE_SQL, (moment,)), "open") -def _client_ports_blocked_by(backend_pid: int) -> frozenset[int]: - rows: Final = read_rows(CLIENT_PORTS_BLOCKED_BY_SQL, (str(backend_pid),)) - return frozenset(_int_value((row,), "client_port") for row in rows) - - -def _local_tcp_ports(pid: int) -> Iterator[int]: - worker: Final = psutil.Process(pid) - for process in (worker, *worker.children(recursive=True)): - yield from (connection.laddr.port for connection in process.net_connections(kind="tcp")) - - -def _worker_holding(owned: OwnedProxy, client_ports: frozenset[int]) -> int: - (worker,) = tuple(pid for pid in owned.worker_pids() if client_ports & frozenset(_local_tcp_ports(pid))) - return worker - - def _roster_ids(team: str) -> tuple[str, ...]: rows: Final = read_rows(TEAM_ROW_SQL, (team,)) assert len(rows) == 1, rows @@ -508,11 +488,16 @@ def test_a_worker_killed_mid_push_leaves_the_old_roster_and_a_retry_lands_the_ne pinned_key: Final = _team_key_for(scenario, team, max(leaving)) with _holding_row_lock(LOCK_KEY_ROW_SQL, sha256(pinned_key.encode()).hexdigest()) as holder_pid: push: Final = pool.submit(_push_outcome, candidate, team, arriving) - blocked_ports: Final = eventually( - lambda: _client_ports_blocked_by(holder_pid), bool, seconds=UNFIXED_PUSH_SECONDS + _ = eventually( + lambda: _sessions_blocked_by(holder_pid), + lambda waiting: waiting > 0, + seconds=UNFIXED_PUSH_SECONDS, ) killed_at: Final = _database_now() - _ = owned.kill_worker(_worker_holding(owned, blocked_ports)) + victims: Final = tuple(sorted(owned.worker_pids())) + assert victims, "expected at least one serving worker" + for pid in victims: + _ = owned.kill_worker(pid) record_property("interrupted_push", push.result()) _ = eventually(lambda: _transactions_open_since(killed_at), lambda open: open == 0, seconds=ROLLBACK_SECONDS) settled: Final = frozenset(_roster_ids(team)) diff --git a/tests/integration/mcp/test_mcp_oauth_flows.py b/tests/integration/mcp/test_mcp_oauth_flows.py index 9d15ea0afa9..6f59f694352 100644 --- a/tests/integration/mcp/test_mcp_oauth_flows.py +++ b/tests/integration/mcp/test_mcp_oauth_flows.py @@ -860,6 +860,7 @@ def test_spec_client_discovers_the_aggregate_authorization_server_and_signs_in_t "client_origin": "http://127.0.0.1:9", "server_id": None, "server_name": None, + "servers": None, "connected": None, }, described.text completed: Final = gateway.client.post("/authorize/complete", data={"flow": connect_flow}, cookies=flow_cookies) @@ -969,6 +970,7 @@ def test_per_server_resource_scopes_the_session_bearer_to_the_named_server(gatew "client_origin": "http://127.0.0.1:9", "server_id": server_a, "server_name": alias_a, + "servers": None, "connected": True, }, described.text completed: Final = gateway.client.post("/authorize/complete", data={"flow": connect_flow}, cookies=flow_cookies) diff --git a/tests/integration/providers/test_bedrock_mantle_context_overflow_fallbacks_wire.py b/tests/integration/providers/test_bedrock_mantle_context_overflow_fallbacks_wire.py index 24063d8052a..37aa96c2467 100644 --- a/tests/integration/providers/test_bedrock_mantle_context_overflow_fallbacks_wire.py +++ b/tests/integration/providers/test_bedrock_mantle_context_overflow_fallbacks_wire.py @@ -519,10 +519,12 @@ def test_killing_one_worker_leaves_the_sibling_serving_the_fallback(fb_rig: Rig) workers: Final = eventually(fb_rig.owned.worker_pids, lambda pids: len(pids) >= 2, seconds=30) with fb_rig.owned.replacing_worker(min(workers)) as kill: kill() - responses: Final = tuple( - fb_rig.gateway.request("POST", "/v1/chat/completions", _chat(fb_rig.models["fb"], _prompt("after kill"))) - for _ in range(3) - ) + with httpx.Client(base_url=fb_rig.gateway.client.base_url, timeout=15, trust_env=False) as fresh: + survivor: Final = Gateway(fresh, fb_rig.gateway.key, fb_rig.gateway.upstream_url) + responses: Final = tuple( + survivor.request("POST", "/v1/chat/completions", _chat(fb_rig.models["fb"], _prompt("after kill"))) + for _ in range(3) + ) for response in responses: assert response.status_code == 200, response.text assert _FALLBACK_TEXT in response.text, response.text diff --git a/tests/integration/providers/test_bedrock_nova_sonic_realtime_tool_calls.py b/tests/integration/providers/test_bedrock_nova_sonic_realtime_tool_calls.py index ccef839a163..1c8850b9bfa 100644 --- a/tests/integration/providers/test_bedrock_nova_sonic_realtime_tool_calls.py +++ b/tests/integration/providers/test_bedrock_nova_sonic_realtime_tool_calls.py @@ -8,6 +8,11 @@ dials: ``AWSCRTHTTPClient`` -> TLS ALPN h2 -> invoke-with-bidirectional-stream) sides: the OpenAI realtime events the websocket client receives, and the toolSpec / toolResult documents the proxy forwarded to the peer. +The fallback cells cover a websocket that outlives one provider attempt. A refused stream keeps the +client's opening message for the fallback, which gets it once whether it dials Bedrock again or +an OpenAI realtime deployment, and once Bedrock accepted a stream no later attempt may open a second +provider session on that websocket. + The file boots its own proxy instead of the shared gateway so it can pass ``SSL_CERT_FILE`` pointing at the peer's self-signed cert, which is the only way the CRT h2 client trusts it. """ @@ -21,12 +26,19 @@ import uuid from collections.abc import AsyncIterator, Callable, Iterator, Mapping from contextlib import asynccontextmanager from pathlib import Path +from types import MappingProxyType from typing import Final +import httpx import pytest import websockets +import yaml from integration._support.bedrock_realtime_peer import ( + BedrockEvent, BedrockRealtimePeer, + OpeningEnd, + Refusal, + StreamFailure, Turn, assistant_text, bedrock_realtime_peer, @@ -43,12 +55,15 @@ from integration._support.client import ( object_value, string_value, ) -from integration._support.process import owned_proxy +from integration._support.process import graceful_stop_seconds, owned_proxy from integration._support.tls import write_self_signed_cert +from integration._support.upstream import delete_scenario, register_scenario +from integration.cost_calculation.cost_tracking_case import RealtimeResponse from pydantic import JsonValue from websockets.asyncio.client import ClientConnection +from websockets.exceptions import ConnectionClosed -pytestmark: Final = pytest.mark.timeout(180) +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) DEPLOYMENT: Final = "bedrock/amazon.nova-2-sonic-v1:0" WEATHER_PARAMETERS: Final = {"type": "object", "properties": {"location": {"type": "string"}}, "required": ["location"]} @@ -60,6 +75,30 @@ WEATHER_FUNCTION: Final = { NESTED_TOOL: Final = {"type": "function", "function": WEATHER_FUNCTION} FLAT_TOOL: Final = {"type": "function", **WEATHER_FUNCTION} TOOL_RESULT_OUTPUT: Final = '{"temperature_f":72,"conditions":"sunny"}' +QUESTION: Final = "What is the weather in New York right now?" +ANSWER: Final = "It is 72 and sunny" +INSTRUCTIONS: Final = "Answer every weather question in one short sentence." +OPENAI_VOICE_MODEL: Final = "gpt-realtime-2.1" +BEDROCK_REPLAY_GROUPS: Final = MappingProxyType( + { + "system_prompt": ("sonic-refused-session-update-first", "sonic-replay-session-update-first"), + "user_input": ("sonic-refused-text-first", "sonic-replay-text-first"), + } +) +OPENAI_REPLAY_GROUPS: Final = ("sonic-refused-openai-fallback", "openai-replay-target") +OPENAI_COMMITTED_GROUPS: Final = ("sonic-committed-openai-fallback", "openai-committed-target") +FALLBACK_GROUPS: Final = (*BEDROCK_REPLAY_GROUPS.values(), OPENAI_REPLAY_GROUPS, OPENAI_COMMITTED_GROUPS) +SCRIPTED_DONE: Final[dict[str, JsonValue]] = { + "type": "response.done", + "event_id": "evt_$UNIQUE_ID", + "response": { + "id": "resp_$UNIQUE_ID", + "object": "realtime.response", + "status": "completed", + "output": [], + "usage": {"total_tokens": 0, "input_tokens": 0, "output_tokens": 0}, + }, +} @pytest.fixture(scope="module") @@ -67,15 +106,35 @@ def cert(tmp_path_factory: pytest.TempPathFactory) -> tuple[Path, Path]: return write_self_signed_cert(tmp_path_factory.mktemp("bedrock-sonic-cert")) +def _fallback_config(directory: Path) -> Path: + stock: Final = JSON_OBJECT.validate_python( + yaml.safe_load((Path(__file__).resolve().parents[1] / "proxy_config.yaml").read_text()) + ) + config: Final = { + **stock, + "router_settings": { + **object_value(stock["router_settings"]), + "fallbacks": [{source: [target]} for source, target in FALLBACK_GROUPS], + }, + } + path: Final = directory / "bedrock-sonic-fallbacks.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + @pytest.fixture(scope="module") def candidate(tmp_path_factory: pytest.TempPathFactory, cert: tuple[Path, Path]) -> Iterator[Gateway]: + directory: Final = tmp_path_factory.mktemp("bedrock-sonic") with gateway_from_environment() as base: - with owned_proxy(base, tmp_path_factory.mktemp("bedrock-sonic"), {"SSL_CERT_FILE": str(cert[0])}) as proxy: + with owned_proxy(base, directory, {"SSL_CERT_FILE": str(cert[0])}, config=_fallback_config(directory)) as proxy: yield proxy def _deployment(gateway: Gateway, scenario: Scenario, peer: BedrockRealtimePeer) -> str: - name: Final = f"bedrock-sonic-{uuid.uuid4().hex[:12]}" + return _bedrock_deployment(gateway, scenario, peer, f"bedrock-sonic-{uuid.uuid4().hex[:12]}") + + +def _bedrock_deployment(gateway: Gateway, scenario: Scenario, peer: BedrockRealtimePeer, name: str) -> str: created: Final = gateway.post( "/model/new", { @@ -94,6 +153,44 @@ def _deployment(gateway: Gateway, scenario: Scenario, peer: BedrockRealtimePeer) return name +def _openai_deployment(gateway: Gateway, scenario: Scenario, name: str) -> str: + scenario_id: Final = f"bedrock-sonic-fallback-{uuid.uuid4().hex[:12]}" + handle: Final = register_scenario( + scenario_id, RealtimeResponse(content_type="application/x-realtime", events=(SCRIPTED_DONE,)) + ) + scenario.cleanups.callback(delete_scenario, handle) + created: Final = gateway.post( + "/model/new", + { + "model_name": name, + "litellm_params": { + "model": f"openai/{OPENAI_VOICE_MODEL}", + "api_key": scenario_id, + "api_base": f"{gateway.upstream_url}/v1", + }, + "model_info": {"mode": "realtime"}, + }, + ) + scenario.cleanups.callback(scenario.delete_model, string_value(object_value(created["model_info"])["id"])) + return scenario_id + + +def _upstream_requests(gateway: Gateway, scenario_id: str) -> tuple[dict[str, JsonValue], ...]: + with httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream: + observations: Final = JSON_OBJECT.validate_json(upstream.get("/__observations").content) + requests: Final = observations["requests"] + assert isinstance(requests, list), observations + return tuple( + request + for request in map(object_value, requests) + if request["authorization"] == f"Bearer {scenario_id}" or request["api_key"] == scenario_id + ) + + +def _upstream_frames(requests: tuple[dict[str, JsonValue], ...]) -> tuple[dict[str, JsonValue], ...]: + return tuple(object_value(request["body"]) for request in requests if request["method"] == "WEBSOCKET_FRAME") + + def _session_update(tool: Mapping[str, JsonValue]) -> str: return json.dumps( { @@ -113,12 +210,16 @@ def _user_text() -> str: "item": { "type": "message", "role": "user", - "content": [{"type": "input_text", "text": "What is the weather in New York right now?"}], + "content": [{"type": "input_text", "text": QUESTION}], }, } ) +def _instructions_update() -> str: + return json.dumps({"type": "session.update", "session": {"instructions": INSTRUCTIONS}}) + + _RESPONSE_CREATE: Final = json.dumps({"type": "response.create"}) @@ -207,10 +308,13 @@ def test_nested_tool_call_before_assistant_block(candidate: Gateway, cert: tuple assert _tail_types(events) == ( "response.created", "response.output_item.added", + "response.function_call_arguments.delta", "response.function_call_arguments.done", "response.output_item.done", "response.done", ), events + (delta,) = tuple(event for event in events if event.get("type") == "response.function_call_arguments.delta") + assert string_value(delta["delta"]) == '{"location":"New York"}' (call,) = _function_calls(events) assert string_value(call["call_id"]) == "tooluse_A" assert string_value(call["name"]) == "get_weather" @@ -376,3 +480,160 @@ def test_malformed_tool_content(candidate: Gateway, cert: tuple[Path, Path], con events: Final = asyncio.run(_one_call_session(candidate, model, key, NESTED_TOOL)) (call,) = _function_calls(events) assert string_value(call["arguments"]) == (content if content is not None else "{}"), events + + +def _of_type(events: tuple[Event, ...], event_type: str) -> tuple[Event, ...]: + return tuple(event for event in events if event.get("type") == event_type) + + +def _bedrock_fields(events: tuple[BedrockEvent, ...], name: str) -> tuple[dict[str, JsonValue], ...]: + bodies: Final = tuple(object_value(event["event"]) for event in events) + return tuple(object_value(body[name]) for body in bodies if name in body) + + +def _text_inputs(events: tuple[BedrockEvent, ...], role: str) -> tuple[str, ...]: + roles: Final = MappingProxyType( + {string_value(start["contentName"]): start.get("role") for start in _bedrock_fields(events, "contentStart")} + ) + return tuple( + string_value(text["content"]) + for text in _bedrock_fields(events, "textInput") + if roles.get(string_value(text["contentName"])) == role + ) + + +def _answer_turn() -> Turn: + return Turn("user_input", assistant_text(ANSWER) + (usage(30, 12),)) + + +def _answer_text(events: tuple[Event, ...]) -> str: + return "".join(string_value(event["delta"]) for event in _of_type(events, "response.text.delta")) + + +async def _opened_with(socket: ClientConnection, opening: tuple[str, ...], stop: frozenset[str]) -> tuple[Event, ...]: + created: Final = JSON_OBJECT.validate_json(await socket.recv()) + for message in opening: + await socket.send(message) + return (created,) + await _events_until(socket, lambda event, _events: event.get("type") in stop) + + +async def _session( + gateway: Gateway, model: str, key: str, opening: tuple[str, ...], stop: frozenset[str] +) -> tuple[Event, ...]: + async with _realtime(gateway, model, key) as socket: + return await _opened_with(socket, opening, stop) + + +async def _session_until_closed( + gateway: Gateway, model: str, key: str, opening: tuple[str, ...] +) -> tuple[tuple[Event, ...], int | None]: + async with _realtime(gateway, model, key) as socket: + events: Final = await _opened_with(socket, opening, frozenset(("error", "session.created"))) + if events[-1].get("type") != "error": + return events, None + try: + await asyncio.wait_for(socket.recv(), 30) + except ConnectionClosed: + return events, socket.close_code + return events, None + + +@pytest.mark.parametrize("after", ["system_prompt", "user_input"], ids=["session-update-first", "text-first"]) +def test_refused_opening_is_replayed_on_the_bedrock_fallback( + candidate: Gateway, cert: tuple[Path, Path], after: OpeningEnd +) -> None: + source, target = BEDROCK_REPLAY_GROUPS[after] + update_first: Final = after == "system_prompt" + opening: Final = ((_instructions_update(),) if update_first else ()) + (_user_text(), _RESPONSE_CREATE) + opening_role, opening_text = ("SYSTEM", INSTRUCTIONS) if update_first else ("USER", QUESTION) + with ( + bedrock_realtime_peer((), cert, refusal=Refusal(after)) as refusing, + bedrock_realtime_peer((_answer_turn(),), cert) as healthy, + candidate.scenario() as scenario, + ): + _bedrock_deployment(candidate, scenario, refusing, source) + _bedrock_deployment(candidate, scenario, healthy, target) + events: Final = asyncio.run( + _session( + candidate, source, scenario.key(), opening, frozenset(("response.done", "error", "session.created")) + ) + ) + assert len(_of_type(events, "session.created")) == 1, events + assert len(_of_type(events, "session.updated")) == int(update_first), events + assert _of_type(events, "error") == (), events + assert len(_response_dones(events)) == 1, events + assert _answer_text(events) == ANSWER, events + (replayed,) = healthy.streams() + assert _text_inputs(replayed, "USER") == (QUESTION,), replayed + assert _text_inputs(replayed, "SYSTEM").count(INSTRUCTIONS) == int(update_first), replayed + refused: Final = refusing.streams() + assert refused and all(opening_text in _text_inputs(stream, opening_role) for stream in refused), refused + + +def test_refused_text_opening_is_replayed_on_the_openai_fallback(candidate: Gateway, cert: tuple[Path, Path]) -> None: + source, target = OPENAI_REPLAY_GROUPS + with ( + bedrock_realtime_peer((), cert, refusal=Refusal("user_input")) as refusing, + candidate.scenario() as scenario, + ): + _bedrock_deployment(candidate, scenario, refusing, source) + scenario_id: Final = _openai_deployment(candidate, scenario, target) + events: Final = asyncio.run( + _session( + candidate, + source, + scenario.key(), + (_user_text(), _RESPONSE_CREATE), + frozenset(("response.done", "error")), + ) + ) + assert _of_type(events, "error") == (), events + assert len(_response_dones(events)) == 1, events + requests: Final = _upstream_requests(candidate, scenario_id) + assert len(tuple(request for request in requests if request["method"] == "WEBSOCKET")) == 1, requests + assert _upstream_frames(requests) == (json.loads(_user_text()), json.loads(_RESPONSE_CREATE)), requests + + +def test_committed_session_failure_closes_the_socket_without_another_provider_stream( + candidate: Gateway, cert: tuple[Path, Path] +) -> None: + source, target = OPENAI_COMMITTED_GROUPS + with ( + bedrock_realtime_peer((_answer_turn(),), cert, failure=StreamFailure()) as committed, + candidate.scenario() as scenario, + ): + _bedrock_deployment(candidate, scenario, committed, source) + scenario_id: Final = _openai_deployment(candidate, scenario, target) + events, close_code = asyncio.run( + _session_until_closed(candidate, source, scenario.key(), (_user_text(), _RESPONSE_CREATE)) + ) + assert len(_of_type(events, "session.created")) == 1, events + assert len(_response_dones(events)) == 1, events + assert _answer_text(events) == ANSWER, events + assert string_value(events[-1]["type"]) == "error", events + assert close_code == 1011, events + assert _upstream_requests(candidate, scenario_id) == () + assert len(committed.streams()) == 1 + + +def test_voice_object_is_refused_and_never_reaches_bedrock(candidate: Gateway, cert: tuple[Path, Path]) -> None: + voice_update: Final = json.dumps({"type": "session.update", "session": {"voice": {"id": "voice_123"}}}) + with bedrock_realtime_peer((_answer_turn(),), cert) as peer, candidate.scenario() as scenario: + model: Final = _deployment(candidate, scenario, peer) + events: Final = asyncio.run( + _session( + candidate, + model, + scenario.key(), + (voice_update, _user_text(), _RESPONSE_CREATE), + frozenset(("response.done", "session.created")), + ) + ) + errors: Final = tuple(object_value(event["error"]) for event in _of_type(events, "error")) + assert tuple((error.get("code"), error.get("param")) for error in errors) == ( + ("invalid_type", "session.voice"), + ), events + assert _of_type(events, "session.updated") == (), events + assert len(_of_type(events, "session.created")) == 1, events + assert _answer_text(events) == ANSWER, events + assert "voice_123" not in json.dumps(peer.received()) diff --git a/tests/integration/routing/test_prompt_cache_pin_ttl.py b/tests/integration/routing/test_prompt_cache_pin_ttl.py index d4328cd0e31..1ca92505191 100644 --- a/tests/integration/routing/test_prompt_cache_pin_ttl.py +++ b/tests/integration/routing/test_prompt_cache_pin_ttl.py @@ -84,7 +84,7 @@ _FAMILIES: Final = ( _Family("pin-bedrock-converse-1h", "pin-bedrock-converse-1h", one_hour=True, wire_one_hour=True), _Family("pin-bedrock-converse-5m", "pin-bedrock-converse-5m", one_hour=False, wire_one_hour=False), _Family("pin-bedrock-invoke-1h", "pin-bedrock-invoke-1h", one_hour=True, wire_one_hour=True), - _Family("pin-bedrock-arn-base", "pin-bedrock-arn-base", one_hour=False, wire_one_hour=False), + _Family("pin-bedrock-arn-base", "pin-bedrock-arn-base", one_hour=True, wire_one_hour=True), _Family("pin-bedrock-arn", "pin-bedrock-arn", one_hour=False, wire_one_hour=False), _Family("pin-azure-ai", "pin-azure-ai", one_hour=True, wire_one_hour=True), _Family("pin-vertex", "pin-vertex", one_hour=True, wire_one_hour=True), @@ -449,19 +449,39 @@ def test_provider_family_pin_lifetime_agrees_with_the_ttl_the_wire_carries(rig: spend_row_lands(response_id) -def test_messages_through_bedrock_invoke_pins_an_hour_and_sends_the_hour(rig: _Rig) -> None: +@pytest.mark.parametrize("group", ["pin-bedrock-invoke-1h", "pin-bedrock-arn-base"]) +def test_messages_through_bedrock_pins_an_hour_and_sends_the_hour(rig: _Rig, group: str) -> None: before: Final = _before(rig) marker: Final = new_marker() response: Final = rig.owned.gateway.request( - "POST", "/v1/messages", {**_messages_body(new_marker(), marker), "model": "pin-bedrock-invoke-1h"} + "POST", "/v1/messages", {**_messages_body(new_marker(), marker), "model": group} ) assert response.status_code == 200, response.text _, pin = one_new_pin(rig.redis, before) - assert pin.model_id == "pin-bedrock-invoke-1h", pin + assert pin.model_id == group, pin assert lives_one_hour(pin), pin assert [wire_carries_one_hour(request) for request in sent_with((rig.bedrock,), marker)] == [True] +def test_response_cache_hit_on_an_application_profile_rewrites_the_pin_for_its_base_model_hour(rig: _Rig) -> None: + before: Final = _before(rig) + marker: Final = new_marker() + body: Final = chat_body("pin-bedrock-arn-base", _first_turn(new_marker(), marker, ONE_HOUR_MARK)) + first, served = post_chat(rig.owned.gateway, body) + key, pin = one_new_pin(rig.redis, before) + assert served == pin.model_id == "pin-bedrock-arn-base", (served, pin) + assert lives_one_hour(pin), pin + assert [wire_carries_one_hour(request) for request in sent_with(rig.wires(), marker)] == [True] + spend_row_lands(first) + rig.redis.delete(key) + second, _ = post_chat(rig.owned.gateway, body) + assert sent_with(rig.wires(), marker) == () + assert second == first + rewritten: Final = eventually(lambda: pin_at(rig.redis, key), lambda found: found is not None, seconds=30) + assert rewritten is not None and rewritten.model_id == "pin-bedrock-arn-base", rewritten + assert lives_one_hour(rewritten), rewritten + + @pytest.mark.parametrize("ttl", _MALFORMED_TTLS, ids=["int", "2h", "upper-1H", "empty", "null", "list", "5kb"]) def test_malformed_ttl_answers_and_pins_five_minutes_without_a_wire_ttl(rig: _Rig, ttl: JsonValue) -> None: before: Final = _before(rig) diff --git a/tests/integration/spend/test_daily_rollup_retry.py b/tests/integration/spend/test_daily_rollup_retry.py index 5b48d841f75..9ed712cdd06 100644 --- a/tests/integration/spend/test_daily_rollup_retry.py +++ b/tests/integration/spend/test_daily_rollup_retry.py @@ -142,6 +142,7 @@ def test_failed_daily_user_rollup_commit_is_retried_so_spend_report_and_daily_ac assert metrics == { "spend": pytest.approx(0.06), "flat_cost": pytest.approx(0.0), + "ptu_hours": pytest.approx(0.0), "prompt_tokens": 20, "completion_tokens": 20, "cache_read_input_tokens": 0, diff --git a/tests/unit/caching/test_caching_handler.py b/tests/unit/caching/test_caching_handler.py index a9541f9fcb8..b84d1f79dc7 100644 --- a/tests/unit/caching/test_caching_handler.py +++ b/tests/unit/caching/test_caching_handler.py @@ -7,6 +7,7 @@ from dotenv import load_dotenv load_dotenv() import asyncio import hashlib +import threading import random import pytest @@ -64,6 +65,46 @@ def setup_cache(): return cache +@pytest.mark.parametrize( + ("call", "request_kwargs"), + [ + ( + completion, + {"model": "gpt-4o-mini", "messages": [{"role": "user", "content": "hi"}], "mock_response": "x"}, + ), + (embedding, {"model": "text-embedding-3-small", "input": ["hi"], "mock_response": [0.1, 0.2]}), + (litellm.text_completion, {"model": "gpt-3.5-turbo-instruct", "prompt": "hi", "mock_response": "x"}), + ], +) +def test_sync_cache_hits_reach_sync_custom_logger( + monkeypatch: pytest.MonkeyPatch, call: Callable[..., object], request_kwargs: dict[str, object] +) -> None: + from litellm.integrations.custom_logger import CustomLogger + + logged_cache_hits: Final[list[object]] = [] + all_logged: Final = threading.Event() + + class SyncRecorder(CustomLogger): + def log_success_event( + self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + logged_cache_hits.append(kwargs.get("cache_hit")) + if len(logged_cache_hits) == 3: + all_logged.set() + + monkeypatch.setattr(litellm, "success_callback", []) + monkeypatch.setattr(litellm, "_async_success_callback", []) + monkeypatch.setattr(litellm, "input_callback", []) + monkeypatch.setattr(litellm, "callbacks", [SyncRecorder()]) + monkeypatch.setattr(litellm, "cache", Cache(type=LiteLLMCacheType.LOCAL)) + + for _ in range(3): + call(**request_kwargs) + + assert all_logged.wait(timeout=10) + assert logged_cache_hits == [None, True, True] + + @pytest.fixture def copilot_response_cache(monkeypatch: pytest.MonkeyPatch) -> Cache: cache: Final = Cache(type=LiteLLMCacheType.LOCAL) @@ -1071,7 +1112,7 @@ def test_sync_get_cache_does_not_eagerly_log_streaming_responses_hits(): stream=True, start_time=datetime.now(), ) - logging_obj.handle_sync_success_callbacks_for_async_calls = MagicMock() + logging_obj.handle_sync_cache_hit_success_callbacks = MagicMock() kwargs = { "model": original_model, @@ -1095,7 +1136,7 @@ def test_sync_get_cache_does_not_eagerly_log_streaming_responses_hits(): assert isinstance( cached_response.cached_result, CachedResponsesAPIStreamingIterator ) - logging_obj.handle_sync_success_callbacks_for_async_calls.assert_not_called() + logging_obj.handle_sync_cache_hit_success_callbacks.assert_not_called() def test_sync_get_cache_defers_streaming_completion_hit_callbacks(): @@ -1115,7 +1156,7 @@ def test_sync_get_cache_defers_streaming_completion_hit_callbacks(): stream=True, start_time=datetime.now(), ) - logging_obj.handle_sync_success_callbacks_for_async_calls = MagicMock() + logging_obj.handle_sync_cache_hit_success_callbacks = MagicMock() kwargs = { "model": original_model, @@ -1136,7 +1177,7 @@ def test_sync_get_cache_defers_streaming_completion_hit_callbacks(): ) assert cached_response.cached_result is not None - logging_obj.handle_sync_success_callbacks_for_async_calls.assert_not_called() + logging_obj.handle_sync_cache_hit_success_callbacks.assert_not_called() def test_should_defer_streaming_cache_hit_callbacks_for_any_streaming_request(): diff --git a/tests/unit/integration_support/test_process.py b/tests/unit/integration_support/test_process.py index 2106436edb8..4b053fb5a97 100644 --- a/tests/unit/integration_support/test_process.py +++ b/tests/unit/integration_support/test_process.py @@ -20,6 +20,7 @@ import pytest import yaml from pydantic import TypeAdapter +from litellm.proxy.types_utils.utils import get_instance_fn from tests.integration._support.client import Gateway, eventually from tests.integration._support.process import ( OwnedProxy, @@ -36,6 +37,32 @@ from tests.integration._support.process import ( TESTS_DIR: Final = Path(__file__).resolve().parents[2] +@pytest.mark.parametrize("config_contents", ("[]\n", "general_settings: []\n", "general_settings: {}\n")) +def test_worker_configuration_preserves_sibling_auth_modules( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + config_contents: str, +) -> None: + monkeypatch.setenv("INTEGRATION_WORKER_ISOLATION", "1") + monkeypatch.setenv("INTEGRATION_PROVIDER_PORT", "8203") + monkeypatch.setenv("INTEGRATION_UPSTREAM_URL", "http://127.0.0.1:8202") + source: Final = tmp_path / "source.yaml" + source.write_text(config_contents) + auth_state: Final = tmp_path / "sibling_auth.json" + auth_state.write_text("authenticated") + (tmp_path / "sibling_auth.py").write_text( + "from pathlib import Path\n" + "def authenticate():\n" + " return Path(__file__).with_suffix('.json').read_text()\n" + ) + + configured: Final = worker_configuration(source, tmp_path / "worker") + + assert get_instance_fn("sibling_auth.authenticate", config_file_path=str(configured))() == "authenticated" + auth_state.write_text("updated") + assert get_instance_fn("sibling_auth.authenticate", config_file_path=str(configured))() == "updated" + + def test_worker_configuration_routes_stock_edges_to_its_own_servers_and_preserves_other_targets( monkeypatch: pytest.MonkeyPatch, tmp_path: Path ) -> None: @@ -93,6 +120,42 @@ def test_worker_configuration_budgets_implicit_pools_without_changing_connection assert yaml.safe_load(source.read_text()) == original +def test_worker_configuration_keeps_modules_beside_a_test_config_importable( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + monkeypatch.setenv("INTEGRATION_WORKER_ISOLATION", "1") + source: Final = tmp_path / "proxy.yaml" + source.write_text(yaml.safe_dump({"general_settings": {"custom_auth": "sibling_auth.authenticate"}})) + (tmp_path / "sibling_auth.py").write_text("def authenticate() -> str:\n return 'sibling'\n") + configured: Final = worker_configuration(source, tmp_path / "gateway") + assert configured != source + assert get_instance_fn("sibling_auth.authenticate", config_file_path=str(configured))() == "sibling" + + +def test_worker_configuration_writes_the_repository_config_copy_into_the_worker_directory( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + monkeypatch.setenv("INTEGRATION_WORKER_ISOLATION", "1") + monkeypatch.setenv("INTEGRATION_PROVIDER_PORT", "8203") + monkeypatch.setenv("INTEGRATION_UPSTREAM_URL", "http://127.0.0.1:8202") + monkeypatch.setenv("INTEGRATION_PROXY_ROOT", str(tmp_path / "product")) + configured: Final = worker_configuration(TESTS_DIR / "integration/proxy_config.yaml", tmp_path / "worker") + assert configured.parent == tmp_path / "worker" + + +def test_worker_configuration_writes_the_proxy_checkout_config_copy_into_the_worker_directory( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + monkeypatch.setenv("INTEGRATION_WORKER_ISOLATION", "1") + checkout: Final = tmp_path / "checkout" + monkeypatch.setenv("INTEGRATION_PROXY_ROOT", str(checkout)) + source: Final = checkout / "tests/integration/proxy_config.yaml" + source.parent.mkdir(parents=True) + source.write_text("model_list: []\n") + configured: Final = worker_configuration(source, tmp_path / "worker") + assert configured.parent == tmp_path / "worker" + + def test_serial_proxy_keeps_the_requested_configuration(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: monkeypatch.delenv("INTEGRATION_WORKER_ISOLATION", raising=False) source: Final = tmp_path / "source.yaml" diff --git a/tests/unit/litellm_core_utils/llm_response_utils/test_convert_to_streaming_response.py b/tests/unit/litellm_core_utils/llm_response_utils/test_convert_to_streaming_response.py index 6b017f9ff79..348d0105a6c 100644 --- a/tests/unit/litellm_core_utils/llm_response_utils/test_convert_to_streaming_response.py +++ b/tests/unit/litellm_core_utils/llm_response_utils/test_convert_to_streaming_response.py @@ -155,6 +155,89 @@ async def test_a_replayed_stream_skips_cached_completion_tokens_details_that_do_ assert rebuilt.usage.completion_tokens == 7 +def _prompt_cached_payload(content: str, cache_read: int, cache_write: int) -> dict: + cached = _async_payload(content=content) + cached["usage"] = litellm.Usage( + prompt_tokens=42145, + completion_tokens=6, + total_tokens=42151, + cache_read_input_tokens=cache_read, + cache_creation_input_tokens=cache_write, + ).model_dump() + return cached + + +@pytest.mark.parametrize("content", ["", "Hello world! How are you?"], ids=["single_frame", "sliced"]) +@pytest.mark.parametrize("cache_read, cache_write", [(42139, 0), (0, 42139)], ids=["read", "write"]) +@pytest.mark.parametrize("replay", ["async", "sync"]) +@pytest.mark.asyncio +async def test_a_replayed_stream_keeps_the_cached_prompt_cache_tokens(replay, cache_read, cache_write, content): + cached = _prompt_cached_payload(content, cache_read, cache_write) + + chunks = await _collect_async(cached) if replay == "async" else list(convert_to_streaming_response(cached)) + usage: Final = next(c.usage for c in chunks if getattr(c, "usage", None) is not None) + + assert usage.prompt_tokens == 42145 + assert usage.prompt_tokens_details.cached_tokens == cache_read + assert usage.prompt_tokens_details.cache_write_tokens == cache_write + assert usage.cache_read_input_tokens == cache_read + assert usage.cache_creation_input_tokens == cache_write + + +@pytest.mark.parametrize("replay", ["async", "sync"]) +@pytest.mark.asyncio +async def test_a_replayed_stream_keeps_cached_tokens_reported_only_in_the_prompt_tokens_details(replay): + cached = _async_payload() + cached["usage"]["prompt_tokens_details"] = {"cached_tokens": 4, "audio_tokens": 1} + + chunks = await _collect_async(cached) if replay == "async" else list(convert_to_streaming_response(cached)) + rebuilt = litellm.stream_chunk_builder(chunks=chunks) + + assert rebuilt.usage.prompt_tokens_details.cached_tokens == 4 + assert rebuilt.usage.prompt_tokens_details.audio_tokens == 1 + + +@pytest.mark.parametrize("replay", ["async", "sync"]) +@pytest.mark.asyncio +async def test_a_replayed_stream_keeps_integral_counts_cached_as_floats_or_numeric_strings(replay): + cached = _async_payload() + cached["usage"] = { + "prompt_tokens": 100.0, + "completion_tokens": "6", + "total_tokens": 106.0, + "cache_read_input_tokens": 90.0, + } + + chunks = await _collect_async(cached) if replay == "async" else list(convert_to_streaming_response(cached)) + usage: Final = next(c.usage for c in chunks if getattr(c, "usage", None) is not None) + + assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (100, 6, 106) + assert usage.cache_read_input_tokens == 90 + + +@pytest.mark.parametrize( + "field, value", + [ + ("prompt_tokens_details", {"cached_tokens": "abc"}), + ("prompt_tokens_details", [1, 2]), + ("cache_read_input_tokens", "abc"), + ("cache_creation_input_tokens", 1.5), + ], + ids=["details_text", "details_list", "read_text", "write_fraction"], +) +@pytest.mark.parametrize("replay", ["async", "sync"]) +@pytest.mark.asyncio +async def test_a_replayed_stream_skips_cached_prompt_cache_usage_that_does_not_validate(replay, field, value): + cached = _async_payload() + cached["usage"][field] = value + + chunks = await _collect_async(cached) if replay == "async" else list(convert_to_streaming_response(cached)) + rebuilt = litellm.stream_chunk_builder(chunks=chunks) + + assert rebuilt.choices[0].message.content == "Hello world! How are you?" + assert rebuilt.usage.prompt_tokens == 5 + + @pytest.mark.asyncio async def test_async_finish_reason_only_on_last_chunk(): chunks = await _collect_async(_async_payload()) diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index 0edd2801c7a..5094e4156a6 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -7,6 +7,7 @@ import json import logging import os import sys +import threading import time import traceback from collections.abc import AsyncIterator, Callable, Iterator, Mapping, Sequence @@ -7540,6 +7541,252 @@ async def test_streaming_success_callbacks_survive_guardrail_logging_hook_failur releasing.async_log_success_event.assert_awaited_once() +def _async_request_logging_obj(call_id: str, callbacks: list[CustomLogger], stream: bool = False) -> LitellmLogging: + logging_obj: Final = LitellmLogging( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + stream=stream, + call_type="acompletion", + start_time=datetime.datetime.now(), + litellm_call_id=call_id, + function_id=call_id, + dynamic_async_success_callbacks=callbacks, + ) + logging_obj.update_environment_variables(litellm_params={"acompletion": True, "metadata": {}}, optional_params={}) + return logging_obj + + +@pytest.mark.parametrize(("success_callbacks", "expected"), [(["cache"], False), (["cache", "athina"], True)]) +def test_cache_callback_alone_does_not_start_sync_callbacks_for_async_requests( + monkeypatch: pytest.MonkeyPatch, success_callbacks: list[str], expected: bool +) -> None: + monkeypatch.setattr(litellm, "success_callback", success_callbacks) + logging_obj: Final = _async_request_logging_obj("cache-gate", []) + + assert logging_obj._should_run_sync_callbacks_for_async_calls() is expected + + +@pytest.mark.asyncio +@pytest.mark.parametrize("turn_off_message_logging", [False, True]) +@pytest.mark.parametrize("make_response", [ModelResponse, _responses_api_response_with_text]) +async def test_sync_success_handler_landing_after_logging_hook_keeps_hook_payload_edits( + monkeypatch: pytest.MonkeyPatch, make_response: Callable[[], object], turn_off_message_logging: bool +) -> None: + mirrored: Final = {"x-ratelimit-model_per_project-limit-requests": 600} + response: Final = make_response() + seen_edits: Final[list[tuple[object, object]]] = [] + + class MirroringHook(CustomLogger): + async def async_logging_hook( + self, kwargs: dict[str, object], result: object, call_type: str + ) -> tuple[dict[str, object], object]: + payload = kwargs["standard_logging_object"] + assert isinstance(payload, dict) + payload["hidden_params"]["additional_headers"] = dict(mirrored) + payload["response"] = "masked by hook" + logging_obj.success_handler(result=response, for_async_request=True) + return kwargs, result + + class Recorder(CustomLogger): + async def async_log_success_event( + self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + payload = kwargs["standard_logging_object"] + assert isinstance(payload, dict) + seen_edits.append((payload["hidden_params"]["additional_headers"], payload["response"])) + + monkeypatch.setattr(litellm, "_async_success_callback", []) + monkeypatch.setattr(litellm, "success_callback", []) + monkeypatch.setattr(litellm, "turn_off_message_logging", turn_off_message_logging) + logging_obj: Final = _async_request_logging_obj("late-sync-rebuild", [MirroringHook(), Recorder()]) + + await logging_obj.async_success_handler(result=response) + + assert seen_edits == [(mirrored, "masked by hook")] + + +class _ResponseEditingHook(CustomLogger): + def logging_hook( + self, kwargs: dict[str, object], result: object, call_type: str + ) -> tuple[dict[str, object], object]: + payload = kwargs["standard_logging_object"] + assert isinstance(payload, dict) + payload["response"] = "edited by sync hook" + return kwargs, result + + +def test_sync_logging_hook_edits_stay_off_a_caller_supplied_payload_for_async_requests( + monkeypatch: pytest.MonkeyPatch, +) -> None: + sync_payloads: Final[list[object]] = [] + + class AthinaRecorder: + def log_event( + self, + kwargs: dict[str, object], + response_obj: object, + start_time: object, + end_time: object, + print_verbose: object, + ) -> None: + sync_payloads.append(kwargs["standard_logging_object"]) + + monkeypatch.setattr(litellm, "_async_success_callback", []) + monkeypatch.setattr(litellm, "success_callback", [_ResponseEditingHook(), "athina"]) + monkeypatch.setattr(litellm_logging, "athinaLogger", AthinaRecorder()) + logging_obj: Final = _async_request_logging_obj("caller-supplied-payload", []) + supplied_payload: Final = {"id": "caller-supplied-payload", "response": "original"} + logging_obj.model_call_details["standard_logging_object"] = supplied_payload + + logging_obj.success_handler( + result=ModelResponse(), for_async_request=True, standard_logging_object=supplied_payload + ) + + assert len(sync_payloads) == 1 + sync_payload: Final = sync_payloads[0] + assert isinstance(sync_payload, dict) + assert sync_payload["response"] == "edited by sync hook" + assert supplied_payload == {"id": "caller-supplied-payload", "response": "original"} + assert logging_obj.model_call_details["standard_logging_object"] is supplied_payload + + +@pytest.mark.asyncio +async def test_sync_only_logger_keeps_the_async_payload_when_its_result_builds_none( + monkeypatch: pytest.MonkeyPatch, +) -> None: + sync_payloads: Final[list[object]] = [] + + class AthinaRecorder: + def log_event( + self, + kwargs: dict[str, object], + response_obj: object, + start_time: object, + end_time: object, + print_verbose: object, + ) -> None: + sync_payloads.append(kwargs.get("standard_logging_object")) + + monkeypatch.setattr(litellm, "_async_success_callback", []) + monkeypatch.setattr(litellm, "success_callback", ["athina"]) + monkeypatch.setattr(litellm_logging, "athinaLogger", AthinaRecorder()) + logging_obj: Final = LitellmLogging( + model="pass-through-model", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="pass_through_endpoint", + start_time=datetime.datetime.now(), + litellm_call_id="pass-through-sync-logger", + function_id="pass-through-sync-logger", + ) + logging_obj.update_environment_variables(litellm_params={"metadata": {}}, optional_params={}) + raw_body: Final = json.dumps({"answer": "plain text"}) + await logging_obj.async_success_handler(result=raw_body) + async_payload: Final = logging_obj.model_call_details.get("standard_logging_object") + assert isinstance(async_payload, dict) + async_payload_before_sync: Final = dict(async_payload) + monkeypatch.setattr(litellm, "success_callback", [_ResponseEditingHook(), "athina"]) + + logging_obj.success_handler(result=raw_body, for_async_request=True) + + assert sync_payloads == [{**async_payload_before_sync, "response": "edited by sync hook"}] + assert logging_obj.model_call_details["standard_logging_object"] is async_payload + assert async_payload == async_payload_before_sync + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("entrypoint", ["handle_sync_success_callbacks_for_async_calls", "dispatch_success_handlers"]) +@pytest.mark.parametrize( + ("callback_name", "logger_global"), + [ + ("athina", "athinaLogger"), + ("promptlayer", "promptLayerLogger"), + ("wandb", "weightsBiasesLogger"), + ("lunary", "lunaryLogger"), + ("helicone", "heliconeLogger"), + ("greenscale", "greenscaleLogger"), + ("s3", "s3Logger"), + ("traceloop", "traceloopLogger"), + ], +) +async def test_sync_only_logger_gets_its_own_payload_for_async_requests( + monkeypatch: pytest.MonkeyPatch, stream: bool, entrypoint: str, callback_name: str, logger_global: str +) -> None: + sync_logged: Final = threading.Event() + sync_payloads: Final[list[object]] = [] + async_payloads: Final[list[object]] = [] + + class LegacyLoggerRecorder: + def _record(self, **call: object) -> None: + kwargs = call["kwargs"] + assert isinstance(kwargs, dict) + sync_payloads.append(kwargs["standard_logging_object"]) + sync_logged.set() + + log_event = log_success = _record + + class AsyncRecorder(CustomLogger): + async def async_log_success_event( + self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + async_payloads.append(kwargs["standard_logging_object"]) + + monkeypatch.setattr(litellm, "_async_success_callback", []) + monkeypatch.setattr(litellm, "success_callback", [callback_name]) + monkeypatch.setattr(litellm_logging, logger_global, LegacyLoggerRecorder(), raising=False) + response: Final = ModelResponse(id="chatcmpl-sync-only") + logging_obj: Final = _async_request_logging_obj("sync-only-logger", [AsyncRecorder()], stream=stream) + now: Final = datetime.datetime.now() + + if entrypoint == "dispatch_success_handlers": + await logging_obj.dispatch_success_handlers(result=response, start_time=now, end_time=now) + else: + await logging_obj.async_success_handler(result=response, start_time=now, end_time=now) + logging_obj.handle_sync_success_callbacks_for_async_calls(result=response, start_time=now, end_time=now) + + assert sync_logged.wait(timeout=10) + assert len(async_payloads) == 1 + assert len(sync_payloads) == 1 + async_payload: Final = async_payloads[0] + sync_payload: Final = sync_payloads[0] + assert isinstance(async_payload, dict) + assert isinstance(sync_payload, dict) + assert sync_payload["id"] == async_payload["id"] + assert sync_payload is not async_payload + assert logging_obj.model_call_details["standard_logging_object"] is async_payload + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stream", [False, True]) +async def test_supabase_logs_the_request_messages_and_response_for_async_requests( + monkeypatch: pytest.MonkeyPatch, stream: bool +) -> None: + supabase_logged: Final = threading.Event() + supabase_calls: Final[list[tuple[object, object]]] = [] + + class SupabaseRecorder: + def log_event(self, **call: object) -> None: + supabase_calls.append((call["messages"], call["response_obj"])) + supabase_logged.set() + + monkeypatch.setattr(litellm, "_async_success_callback", []) + monkeypatch.setattr(litellm, "success_callback", ["supabase"]) + monkeypatch.setattr(litellm_logging, "supabaseClient", SupabaseRecorder()) + response: Final = ModelResponse(id="chatcmpl-supabase") + logging_obj: Final = _async_request_logging_obj("supabase-logger", [], stream=stream) + now: Final = datetime.datetime.now() + + await logging_obj.dispatch_success_handlers(result=response, start_time=now, end_time=now) + + assert supabase_logged.wait(timeout=10) + assert len(supabase_calls) == 1 + messages, response_obj = supabase_calls[0] + assert messages == [{"role": "user", "content": "hi"}] + assert isinstance(response_obj, ModelResponse) + assert response_obj.id == "chatcmpl-supabase" + + def _resolve(custom_llm_provider, litellm_params, optional_params, model): from litellm.litellm_core_utils.litellm_logging import ( _resolve_vertex_location_for_cost, diff --git a/tests/unit/litellm_core_utils/test_realtime_streaming.py b/tests/unit/litellm_core_utils/test_realtime_streaming.py index 8de5e27d99c..501700e5d57 100644 --- a/tests/unit/litellm_core_utils/test_realtime_streaming.py +++ b/tests/unit/litellm_core_utils/test_realtime_streaming.py @@ -13,6 +13,7 @@ from websockets.frames import Close import litellm from litellm.constants import ( + BEDROCK_REALTIME_PENDING_OPENING_MESSAGE_SCOPE_KEY, REALTIME_METERED_USAGE_KEY, REALTIME_SESSION_FAILURE_LOGGED_KEY, REALTIME_SESSION_SUCCESS_LOGGED_KEY, @@ -543,6 +544,32 @@ async def test_backend_to_client_stores_normalized_events_for_logging(): assert streaming.messages[0]["response"]["usage"]["total_tokens"] == 0 +@pytest.mark.asyncio +async def test_client_ack_messages_replays_the_opening_message_a_refused_bedrock_stream_left(): + opening_message: Final = json.dumps( + { + "type": "conversation.item.create", + "item": {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "What is 17 * 3?"}]}, + } + ) + next_message: Final = json.dumps({"type": "response.create"}) + client_ws = MagicMock() + client_ws.scope = {"headers": [], BEDROCK_REALTIME_PENDING_OPENING_MESSAGE_SCOPE_KEY: opening_message} + client_ws.receive_text = AsyncMock(side_effect=[next_message, Exception("connection closed")]) + backend_ws = MagicMock() + backend_ws.send = AsyncMock() + logging_obj: Final = _mock_logging() + logging_obj.pre_call = MagicMock() + streaming = RealTimeStreaming(client_ws, backend_ws, logging_obj) + + await streaming.client_ack_messages() + + sent_types: Final = [json.loads(call.args[0])["type"] for call in backend_ws.send.call_args_list] + assert sent_types == ["conversation.item.create", "response.create"] + assert json.loads(backend_ws.send.call_args_list[0].args[0])["item"]["content"][0]["text"] == "What is 17 * 3?" + assert BEDROCK_REALTIME_PENDING_OPENING_MESSAGE_SCOPE_KEY not in client_ws.scope + + @pytest.mark.asyncio async def test_client_ack_messages_keeps_beta_session_shape_for_beta_clients(): client_ws = MagicMock() diff --git a/tests/unit/llms/anthropic/pass_through/adapters/test_streaming_iterator_combined_chunk.py b/tests/unit/llms/anthropic/pass_through/adapters/test_streaming_iterator_combined_chunk.py index 9a6940ff12b..9f8c78fdaa9 100644 --- a/tests/unit/llms/anthropic/pass_through/adapters/test_streaming_iterator_combined_chunk.py +++ b/tests/unit/llms/anthropic/pass_through/adapters/test_streaming_iterator_combined_chunk.py @@ -12,8 +12,14 @@ content survives. import asyncio import json from types import SimpleNamespace -from typing import AsyncIterator +from typing import AsyncIterator, Final +import pytest + +from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( + convert_to_streaming_response, + convert_to_streaming_response_async, +) from litellm.llms.anthropic.pass_through.adapters.streaming_iterator import ( AnthropicStreamWrapper, _CombinedChunkSplitter, @@ -139,6 +145,43 @@ def test_delayed_usage_chunk_preserves_cache_tokens(): assert message_delta["usage"]["cache_creation_input_tokens"] == 20 +def _cached_bedrock_response(cache_read: int, cache_write: int) -> dict: + return { + "id": "chatcmpl-cached", + "object": "chat.completion", + "created": 1700000000, + "model": "us.anthropic.claude-opus-4-7", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "pong"}, "finish_reason": "stop"}], + "usage": Usage( + prompt_tokens=42145, + completion_tokens=6, + total_tokens=42151, + cache_read_input_tokens=cache_read, + cache_creation_input_tokens=cache_write, + ).model_dump(), + } + + +async def _replay_async(cached: dict) -> list[ModelResponseStream]: + return [chunk async for chunk in convert_to_streaming_response_async(cached)] + + +@pytest.mark.parametrize("cache_read, cache_write", [(42139, 0), (0, 42139)], ids=["read", "write"]) +@pytest.mark.parametrize("replay", ["async", "sync"]) +def test_a_cache_hit_replayed_as_a_messages_stream_keeps_the_prompt_cache_tokens(replay, cache_read, cache_write): + cached: Final = _cached_bedrock_response(cache_read, cache_write) + chunks: Final = ( + asyncio.run(_replay_async(cached)) if replay == "async" else list(convert_to_streaming_response(cached)) + ) + + events: Final = list(AnthropicStreamWrapper(completion_stream=iter(chunks), model="claude-opus-4-7")) + message_delta: Final = next(event for event in events if event.get("type") == "message_delta") + + assert message_delta["usage"]["input_tokens"] == 6 + assert message_delta["usage"].get("cache_read_input_tokens", 0) == cache_read + assert message_delta["usage"].get("cache_creation_input_tokens", 0) == cache_write + + def test_splitter_passes_through_non_combined_chunks(): """A chunk with content but no finish_reason is not split.""" chunk = ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(content="partial"), finish_reason=None)]) diff --git a/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py b/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py index 8f95a502137..a181c5a5425 100644 --- a/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py +++ b/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py @@ -17,6 +17,7 @@ from litellm.constants import ( from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.bedrock.realtime.handler import BedrockRealtime from litellm.llms.bedrock.realtime.transformation import BedrockRealtimeConfig +from litellm.realtime_api import main as realtime_main @pytest.fixture(autouse=True) @@ -109,6 +110,14 @@ class UnavailableBedrockStream: raise ServiceUnavailableException("fault injected: Bedrock realtime unavailable") +class RecordingRealtimeHandler: + def __init__(self): + self.calls = 0 + + async def async_realtime(self, **kwargs): + self.calls += 1 + + class FakeLogging: def __init__(self, trace_id="trace-nova-sonic"): self.litellm_trace_id = trace_id @@ -678,18 +687,23 @@ class TestBedrockRealtimeSessionLifecycle: @pytest.mark.asyncio @pytest.mark.parametrize( - "drop_params,expected_errors", - [(False, [("invalid_type", "session.temperature")]), (True, [])], + "session,drop_params,expected_errors", + [ + ({"temperature": "0.2", "modalities": ["text"]}, False, [("invalid_type", "session.temperature")]), + ({"temperature": "0.2", "modalities": ["text"]}, True, []), + ({"modalities": "text"}, False, [("invalid_type", "session.modalities")]), + ({"modalities": "text"}, True, []), + ], ) - async def test_malformed_number_setting_is_rejected_or_dropped_and_the_session_still_answers( - self, stub_aws_sdk_client, drop_params, expected_errors + async def test_malformed_setting_is_rejected_or_dropped_and_the_session_still_answers( + self, stub_aws_sdk_client, session, drop_params, expected_errors ): stream: Final = ScriptedBedrockStream([]) stub_aws_sdk_client["streams"] = [stream] user_text: Final = {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "hi"}]} websocket: Final = ConnectedClientWS( [ - json.dumps({"type": "session.update", "session": {"temperature": "0.2", "modalities": ["text"]}}), + json.dumps({"type": "session.update", "session": session}), json.dumps({"type": "conversation.item.create", "item": user_text}), ] ) @@ -777,12 +791,26 @@ class TestBedrockRealtimeProviderFailurePropagation: dispatched["logging_obj"] = RecordingLogging() return dispatched + USER_TEXT: Final = json.dumps( + { + "type": "conversation.item.create", + "item": {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "hi"}]}, + } + ) + REJECTED_SESSION_UPDATE: Final = json.dumps({"type": "session.update", "session": {"temperature": "0.2"}}) + USER_AUDIO: Final = json.dumps({"type": "input_audio_buffer.append", "audio": "AAAA"}) + @pytest.mark.asyncio + @pytest.mark.parametrize( + "client_messages", + [[SESSION_UPDATE], [USER_TEXT], [USER_AUDIO], [REJECTED_SESSION_UPDATE, USER_TEXT]], + ids=["session-update-first", "text-first", "audio-first", "input-after-rejected-update"], + ) async def test_mid_stream_failure_escapes_keeps_partial_spend_and_blocks_replay( - self, stub_aws_sdk_client, spend_dispatch + self, stub_aws_sdk_client, spend_dispatch, monkeypatch, client_messages: list[str] ): handler = BedrockRealtime() - websocket = ConnectedClientWS([self.SESSION_UPDATE]) + websocket = ConnectedClientWS(client_messages) stream = ScriptedBedrockStream(self.TEXT_TURN, receiver_type=BreakingBedrockReceiver) stub_aws_sdk_client["streams"] = [stream] @@ -814,6 +842,37 @@ class TestBedrockRealtimeProviderFailurePropagation: assert not litellm.should_retry(replay.value.status_code), "the router must not retry the replay refusal" assert "Nova Sonic stream broke" in replay.value.message, "the router surfaces the last attempt's error" + other_provider: Final = RecordingRealtimeHandler() + monkeypatch.setattr(realtime_main, "openai_realtime", other_provider) + with pytest.raises(BedrockError) as fallback: + await realtime_main._arealtime.__wrapped__( + model="openai/gpt-realtime", websocket=websocket, litellm_logging_obj=FakeLogging() + ) + + assert fallback.value.status_code == 400, "a fallback to another provider must not start a second session" + assert other_provider.calls == 0 + + @pytest.mark.asyncio + async def test_readiness_failure_on_an_input_first_start_replays_the_input(self, stub_aws_sdk_client): + handler = BedrockRealtime() + websocket = ConnectedClientWS([self.USER_TEXT]) + healthy_stream = ScriptedBedrockStream([]) + stub_aws_sdk_client["streams"] = [UnavailableBedrockStream(), healthy_stream] + + with pytest.raises(BedrockError) as failure: + await handler.async_realtime( + model="amazon.nova-sonic-v1:0", websocket=websocket, logging_obj=FakeLogging(), **self.AWS_PARAMS + ) + assert failure.value.status_code == 503 + + await handler.async_realtime( + model="amazon.nova-sonic-v1:0", websocket=websocket, logging_obj=FakeLogging(), **self.AWS_PARAMS + ) + + assert [json.loads(m)["type"] for m in websocket.sent_to_client] == ["session.created"] + replayed = [json.loads(chunk.value.bytes_.decode("utf-8")) for chunk in healthy_stream.input_stream.sent] + assert any(event["event"].get("textInput", {}).get("content") == "hi" for event in replayed) + @pytest.mark.asyncio async def test_input_side_failure_keeps_spend_for_responses_already_delivered( self, stub_aws_sdk_client, spend_dispatch diff --git a/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_transformation.py b/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_transformation.py index 1cc135a8671..3e10c3bcbdc 100644 --- a/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_transformation.py +++ b/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_transformation.py @@ -1655,9 +1655,17 @@ class TestBedrockRealtimeUserEventsAndUsage: ({"max_response_output_tokens": float("inf")}, "session.max_response_output_tokens"), ({"max_response_output_tokens": True}, "session.max_response_output_tokens"), ({"output_sample_rate_hertz": "16000"}, "session.output_sample_rate_hertz"), + ({"instructions": {"role": "pirate"}}, "session.instructions"), + ({"voice": {"id": "voice_123"}}, "session.voice"), + ({"modalities": "text"}, "session.modalities"), + ({"modalities": ["text", 1]}, "session.modalities"), + ({"input_sample_rate_hertz": "16000"}, "session.input_sample_rate_hertz"), + ({"tools": {"type": "function", "name": "lookup"}}, "session.tools"), + ({"input_audio_format": {"type": "audio/pcm"}}, "session.input_audio_format"), + ({"output_audio_format": {"type": "audio/pcmu"}}, "session.output_audio_format"), ], ) - def test_malformed_number_setting_is_rejected_without_drop_params(self, session, param): + def test_malformed_setting_is_rejected_without_drop_params(self, session, param): rejection: Final = BedrockRealtimeConfig().session_update_rejection( {"type": "session.update", "session": session} ) @@ -1677,6 +1685,59 @@ class TestBedrockRealtimeUserEventsAndUsage: inference: Final = self._inference_configuration(config, session) assert (inference["temperature"], inference["maxTokens"]) == (0.7, 1024) + def test_non_number_input_sample_rate_is_dropped_with_drop_params(self): + config: Final = BedrockRealtimeConfig(drop_params=True) + session: Final = {"input_sample_rate_hertz": "8000"} + assert config.session_update_rejection({"type": "session.update", "session": session}) is None + config.transform_session_update_event({"type": "session.update", "session": session}) + assert config.input_sample_rate_hertz == BedrockRealtimeConfig().input_sample_rate_hertz + + def test_non_string_instructions_and_voice_are_dropped_with_drop_params(self): + config: Final = BedrockRealtimeConfig(drop_params=True) + session: Final = {"instructions": {"role": "pirate"}, "voice": {"id": "voice_123"}} + assert config.session_update_rejection({"type": "session.update", "session": session}) is None + sent: Final = self._voice_and_system_prompt(config, session) + assert sent == self._voice_and_system_prompt(BedrockRealtimeConfig(), {}) + + @pytest.mark.parametrize("voice", ["alloy", "marin", "Alloy", "", " "]) + def test_openai_or_empty_voice_is_rejected_without_drop_params(self, voice): + rejection: Final = BedrockRealtimeConfig().session_update_rejection( + {"type": "session.update", "session": {"voice": voice}} + ) + assert rejection is not None + assert (rejection.code, rejection.param) == ("invalid_value", "session.voice") + assert f"'{voice}'" in rejection.message + + def test_openai_voice_is_dropped_with_drop_params(self): + config: Final = BedrockRealtimeConfig(drop_params=True) + session: Final = {"voice": "alloy", "instructions": "Answer as a pirate."} + assert config.session_update_rejection({"type": "session.update", "session": session}) is None + voice, system_prompt = self._voice_and_system_prompt(config, session) + assert voice == self._voice_and_system_prompt(BedrockRealtimeConfig(), {})[0] + assert system_prompt == "Answer as a pirate." + + def test_nova_voice_is_applied(self): + session: Final = {"voice": "tiffany"} + assert BedrockRealtimeConfig().session_update_rejection({"type": "session.update", "session": session}) is None + assert self._voice_and_system_prompt(BedrockRealtimeConfig(), session)[0] == "tiffany" + + def test_g711_input_audio_format_sets_the_input_sample_rate(self): + config: Final = BedrockRealtimeConfig() + config.transform_session_update_event({"type": "session.update", "session": {"input_audio_format": "g711_ulaw"}}) + assert config.input_sample_rate_hertz == 8000 + + @staticmethod + def _voice_and_system_prompt(config: BedrockRealtimeConfig, session: dict[str, object]) -> tuple[object, object]: + events: Final = [ + json.loads(message)["event"] + for message in config.transform_session_update_event({"type": "session.update", "session": session}) + ] + voice: Final = next( + e["promptStart"]["audioOutputConfiguration"]["voiceId"] for e in events if "promptStart" in e + ) + system_prompt: Final = next(e["textInput"]["content"] for e in events if "textInput" in e) + return voice, system_prompt + def test_null_number_settings_count_as_unset(self): session: Final = {"temperature": None, "max_response_output_tokens": None} config: Final = BedrockRealtimeConfig() diff --git a/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py b/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py index 018fa808f70..e5565a3df16 100644 --- a/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py @@ -15,10 +15,12 @@ from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheCont from litellm.integrations.custom_logger import CustomLogger from litellm.router_utils.pre_call_checks.prompt_caching_deployment_check import ( PromptCachingDeploymentCheck, + _honors_one_hour_cache_ttl, _get_min_token_count_for_deployments, ) from litellm.router_utils.prompt_caching_cache import PROMPT_CACHE_PIN_TTL_SECONDS, PromptCachingCache from litellm.types.llms.openai import AllMessageValues +from litellm.types.utils import StandardLoggingPayload from litellm.utils import get_prompt_cache_min_tokens, is_prompt_caching_valid_prompt, token_counter MODEL_GROUP_ALIAS = "my-claude-group" @@ -126,6 +128,21 @@ def test_get_min_token_count_for_deployments_falls_back_to_default_for_empty_gro assert _get_min_token_count_for_deployments([]) == DEFAULT_MINIMUM_PROMPT_CACHE_TOKEN_COUNT +def test_bedrock_application_profile_uses_its_base_model_for_one_hour_ttl() -> None: + standard_logging_object: Final = cast( + StandardLoggingPayload, + { + "custom_llm_provider": "bedrock", + "model": "bedrock/converse/arn:aws:bedrock:us-east-1:000000000000:application-inference-profile/pinprofile", + }, + ) + + assert _honors_one_hour_cache_ttl( + standard_logging_object, + {"model": None, "base_model": "bedrock/us.anthropic.claude-sonnet-5-5"}, + ) is True + + @pytest.mark.asyncio async def test_async_filter_deployments_does_not_narrow_prompt_below_model_minimum(): """ diff --git a/tests/unit/test_cost_calculator.py b/tests/unit/test_cost_calculator.py index cc2a1d62d52..f6228f242dd 100644 --- a/tests/unit/test_cost_calculator.py +++ b/tests/unit/test_cost_calculator.py @@ -2,7 +2,7 @@ import json import asyncio import datetime import time -from collections.abc import Mapping +from collections.abc import Iterator, Mapping from types import MappingProxyType, SimpleNamespace from typing import Final, Literal, cast @@ -17,6 +17,10 @@ from litellm import TranscriptionResponse, model_cost from litellm.constants import REALTIME_METERED_USAGE_KEY from litellm.cost_calculator import BaseTokenUsageProcessor, RealtimeAPITokenUsageProcessor, ResponsesWebSocketTokenUsageProcessor, batch_cost_calculator, completion_cost, cost_per_token, handle_realtime_stream_cost_calculation, response_cost_calculator from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.fallback_generalizations import ( + get_fallback_generalization_rules, + set_fallback_generalizations, +) from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( @@ -75,6 +79,48 @@ def _local_model_cost_map(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) +@pytest.fixture +def _claude_family_generalization() -> Iterator[None]: + previous_rules: Final = list(get_fallback_generalization_rules()) + try: + set_fallback_generalizations( + [ + { + "name": "claude-family-baseline", + "pattern": r"claude-[a-z]+-\d+(?:[-.]\d+)?(?:-\d{8})?", + "model_info": {"mode": "chat", "max_input_tokens": 200000, "max_output_tokens": 64000}, + } + ] + ) + yield + finally: + set_fallback_generalizations(previous_rules) + + +@pytest.fixture +def _priced_claude_family_generalization() -> Iterator[None]: + previous_rules: Final = list(get_fallback_generalization_rules()) + try: + set_fallback_generalizations( + [ + { + "name": "claude-family-priced", + "pattern": r"claude-[a-z]+-\d+(?:[-.]\d+)?(?:-\d{8})?", + "model_info": { + "mode": "chat", + "max_input_tokens": 200000, + "max_output_tokens": 64000, + "input_cost_per_token": 0.0000321, + "output_cost_per_token": 0.0000765, + }, + } + ] + ) + yield + finally: + set_fallback_generalizations(previous_rules) + + def test_cost_per_token_duplicate_openai_prefix_matches_model_cost(monkeypatch): """ Router/proxy configs may use deployment ids like openai/openai/. Cost lookup must @@ -149,6 +195,265 @@ def test_completion_cost_uses_response_model_for_dynamic_routing(_local_model_co assert cost > 0, "Cost should be calculated using response model" +def test_completion_cost_prefers_exact_route_price_over_capability_generalization( + _local_model_cost_map: None, + _claude_family_generalization: None, + monkeypatch: pytest.MonkeyPatch, +) -> None: + priced_model: Final = "bedrock/invoke/anthropic.claude-haiku-4-5-20251001-v1:0" + priced_entry: Final = { + "litellm_provider": "bedrock", + "mode": "chat", + "input_cost_per_token": 0.0000123, + "output_cost_per_token": 0.0000456, + } + monkeypatch.setattr(litellm, "model_cost", {priced_model: priced_entry}) + usage: Final = Usage(prompt_tokens=17, completion_tokens=9, total_tokens=26) + response: Final = ModelResponse( + model="anthropic.claude-haiku-4-5-20251001-v1:0", + choices=[], + usage=usage, + ) + + cost: Final = completion_cost( + completion_response=response, + model=priced_model, + custom_llm_provider="bedrock", + ) + + input_rate: Final = priced_entry["input_cost_per_token"] + output_rate: Final = priced_entry["output_cost_per_token"] + assert isinstance(input_rate, (int, float)), priced_entry + assert isinstance(output_rate, (int, float)), priced_entry + assert cost == usage.prompt_tokens * input_rate + usage.completion_tokens * output_rate + + +def test_completion_cost_keeps_priced_generalization_before_later_exact_candidate( + _local_model_cost_map: None, + _priced_claude_family_generalization: None, + monkeypatch: pytest.MonkeyPatch, +) -> None: + priced_model: Final = "bedrock/invoke/anthropic.claude-haiku-4-5-20251001-v1:0" + exact_entry: Final = { + "litellm_provider": "bedrock", + "mode": "chat", + "input_cost_per_token": 0.0000123, + "output_cost_per_token": 0.0000456, + } + monkeypatch.setattr(litellm, "model_cost", {priced_model: exact_entry}) + usage: Final = Usage(prompt_tokens=17, completion_tokens=9, total_tokens=26) + response: Final = ModelResponse( + model="anthropic.claude-haiku-4-5-20251001-v1:0", + choices=[], + usage=usage, + ) + + cost: Final = completion_cost( + completion_response=response, + model=priced_model, + custom_llm_provider="bedrock", + ) + + assert cost == pytest.approx(usage.prompt_tokens * 0.0000321 + usage.completion_tokens * 0.0000765) + + +@pytest.fixture +def _per_second_priced_claude_family_generalization() -> Iterator[None]: + previous_rules: Final = list(get_fallback_generalization_rules()) + try: + set_fallback_generalizations( + [ + { + "name": "claude-family-per-second", + "pattern": r"claude-[a-z]+-\d+(?:[-.]\d+)?(?:-\d{8})?", + "model_info": {"cost_per_second": 0.25}, + } + ] + ) + yield + finally: + set_fallback_generalizations(previous_rules) + + +def test_completion_cost_keeps_per_second_priced_generalization_before_later_exact_candidate( + _local_model_cost_map: None, + _per_second_priced_claude_family_generalization: None, + monkeypatch: pytest.MonkeyPatch, +) -> None: + priced_model: Final = "bedrock/invoke/anthropic.claude-haiku-4-5-20251001-v1:0" + monkeypatch.setattr( + litellm, + "model_cost", + { + priced_model: { + "litellm_provider": "bedrock", + "mode": "chat", + "input_cost_per_token": 0.0000123, + "output_cost_per_token": 0.0000456, + } + }, + ) + response: Final = ModelResponse( + model="anthropic.claude-haiku-4-5-20251001-v1:0", + choices=[], + usage=Usage(prompt_tokens=17, completion_tokens=9, total_tokens=26), + ) + + cost: Final = completion_cost( + completion_response=response, + model=priced_model, + custom_llm_provider="bedrock", + total_time=2000, + ) + + assert cost == pytest.approx(0.25 * 2) + + +@pytest.fixture +def _price_free_gpt_family_generalization() -> Iterator[None]: + previous_rules: Final = list(get_fallback_generalization_rules()) + try: + set_fallback_generalizations( + [ + { + "name": "gpt-family-baseline", + "pattern": r"gpt-\d+(?:\.\d+)?", + "model_info": {"mode": "chat", "supports_reasoning": True}, + } + ] + ) + yield + finally: + set_fallback_generalizations(previous_rules) + + +def test_completion_cost_prices_rule_only_response_model_at_the_azure_deployment_rate( + _local_model_cost_map: None, + _price_free_gpt_family_generalization: None, + monkeypatch: pytest.MonkeyPatch, +) -> None: + deployment_model: Final = "azure/gpt-4o-2024-11-20" + monkeypatch.setattr( + litellm, + "model_cost", + { + deployment_model: { + "litellm_provider": "azure", + "mode": "chat", + "input_cost_per_token": 0.0000111, + "output_cost_per_token": 0.0000222, + } + }, + ) + usage: Final = Usage(prompt_tokens=17, completion_tokens=9, total_tokens=26) + response: Final = ModelResponse(model="gpt-9.9", choices=[], usage=usage) + + cost: Final = completion_cost( + completion_response=response, + model=deployment_model, + custom_llm_provider="azure", + ) + + assert cost == pytest.approx(usage.prompt_tokens * 0.0000111 + usage.completion_tokens * 0.0000222) + + +def test_completion_cost_orders_candidates_without_extra_huggingface_config_requests( + _local_model_cost_map: None, +) -> None: + model: Final = "lazy-candidate-regression-org/lazy-candidate-model" + usage: Final = Usage(prompt_tokens=17, completion_tokens=9, total_tokens=26) + response: Final = ModelResponse( + model=model, + choices=[], + usage=usage, + ) + + with respx.mock(assert_all_called=False) as huggingface: + huggingface.get(host="huggingface.co").mock( + return_value=httpx.Response(200, json={"max_position_embeddings": 4096}) + ) + cost: Final = completion_cost( + completion_response=response, + model=f"huggingface/{model}", + custom_llm_provider="huggingface", + ) + + assert len(huggingface.calls) == 1 + assert cost == 0.0 + + +def test_completion_cost_skips_huggingface_config_request_with_custom_pricing( + _local_model_cost_map: None, +) -> None: + model: Final = "unmapped-cost-regression-org/unmapped-model" + usage: Final = Usage(prompt_tokens=17, completion_tokens=9, total_tokens=26) + response: Final = ModelResponse( + model=model, + choices=[], + usage=usage, + ) + + with respx.mock(assert_all_called=False) as huggingface: + huggingface.get(host="huggingface.co").mock( + return_value=httpx.Response(200, json={"max_position_embeddings": 4096}) + ) + cost: Final = completion_cost( + completion_response=response, + model=model, + custom_llm_provider="huggingface", + custom_pricing=True, + custom_cost_per_token={"input_cost_per_token": 0.001, "output_cost_per_token": 0.002}, + ) + + assert len(huggingface.calls) == 0 + assert cost == pytest.approx(usage.prompt_tokens * 0.001 + usage.completion_tokens * 0.002) + + +def test_completion_cost_returns_zero_for_capability_only_generalization( + _local_model_cost_map: None, + _claude_family_generalization: None, + monkeypatch: pytest.MonkeyPatch, +) -> None: + model: Final = "bedrock/invoke/anthropic.claude-haiku-9-9-20991231-v1:0" + monkeypatch.setattr(litellm, "model_cost", {}) + response: Final = ModelResponse( + model="anthropic.claude-haiku-9-9-20991231-v1:0", + choices=[], + usage=Usage(prompt_tokens=17, completion_tokens=9, total_tokens=26), + ) + + cost: Final = completion_cost( + completion_response=response, + model=model, + custom_llm_provider="bedrock", + ) + + assert cost == 0.0 + + +@pytest.mark.parametrize("custom_llm_provider", ["openai", "hosted_vllm"]) +def test_completion_cost_prices_base_model_under_its_own_provider( + custom_llm_provider: str, _local_model_cost_map: None +) -> None: + response: Final = ModelResponse( + model="my-deploy", + choices=[], + usage=Usage(prompt_tokens=1000, completion_tokens=1000, total_tokens=2000), + ) + + cost: Final = completion_cost( + completion_response=response, + model=f"{custom_llm_provider}/my-deploy", + custom_llm_provider=custom_llm_provider, + base_model="claude-sonnet-4-5", + ) + + info: Final = litellm.get_model_info("claude-sonnet-4-5", custom_llm_provider="anthropic") + assert cost == pytest.approx( + 1000 * info["input_cost_per_token"] + 1000 * info["output_cost_per_token"] + ) + + def test_completion_cost_strips_dated_azure_snapshot_model(_local_model_cost_map: None) -> None: dated_response = ModelResponse( model="gpt-5.6-luna-2099-01-01", diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index 14cea8f0c2c..6d9fa1d8cd1 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -4556,7 +4556,7 @@ def test_stream_chunk_builder_skips_stamp_when_cost_is_unpriceable(): from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging logging_obj: Final = LiteLLMLogging( - model="us.anthropic.claude-opus-5", + model="us.anthropic.claude-unpriced-999", messages=[{"role": "user", "content": "hi"}], stream=True, call_type="completion", @@ -4582,6 +4582,41 @@ def test_stream_chunk_builder_skips_stamp_when_cost_is_unpriceable(): assert response._hidden_params.get("response_cost") is None +def test_stream_chunk_builder_prices_alias_response_at_its_deployment_model(): + import time as time_module + + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging + + deployment_model: Final = "us.anthropic.claude-opus-5" + logging_obj: Final = LiteLLMLogging( + model=deployment_model, + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="completion", + start_time=time_module.time(), + litellm_call_id="stream-builder-alias-priced", + function_id="1", + ) + logging_obj.model_call_details["custom_llm_provider"] = "bedrock" + logging_obj.optional_params = {} + usage_chunk: Final = _stream_builder_text_chunk("bedrock-claude-opus-5", "") + usage_chunk.usage = Usage(prompt_tokens=40, completion_tokens=5, total_tokens=45) + chunks: Final = [ + _stream_builder_text_chunk("bedrock-claude-opus-5", "Hello ", finish_reason="stop"), + usage_chunk, + ] + + response: Final = litellm.stream_chunk_builder( + chunks=chunks, messages=[{"role": "user", "content": "hi"}], logging_obj=logging_obj + ) + + entry: Final = litellm.model_cost[deployment_model] + expected_cost: Final = 40 * entry["input_cost_per_token"] + 5 * entry["output_cost_per_token"] + assert response is not None + assert getattr(response.usage, "cost", None) == pytest.approx(expected_cost) + assert response._hidden_params.get("response_cost") == pytest.approx(expected_cost) + + def test_stream_chunk_builder_keeps_provider_reported_usage_cost(): usage_chunk: Final = _stream_builder_text_chunk("gpt-4o", "") usage_chunk.usage = Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15, cost=0.5)