diff --git a/litellm/integrations/s3.py b/litellm/integrations/s3.py index 53a982cd2c4..e8252d87572 100644 --- a/litellm/integrations/s3.py +++ b/litellm/integrations/s3.py @@ -161,8 +161,13 @@ def get_s3_object_key( start_time: datetime, s3_file_name: str, ) -> str: + sanitized_s3_file_name = s3_file_name.replace("/", "_") s3_object_key = ( - (s3_path.rstrip("/") + "/" if s3_path else "") + prefix + start_time.strftime("%Y-%m-%d") + "/" + s3_file_name + (s3_path.rstrip("/") + "/" if s3_path else "") + + prefix + + start_time.strftime("%Y-%m-%d") + + "/" + + sanitized_s3_file_name ) # we need the s3 key to include the time, so we log cache hits too s3_object_key += ".json" return s3_object_key diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index c039f0f43ee..33bf546c239 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -445,6 +445,7 @@ class PromptTokensDetailsResult(TypedDict): text_tokens: int audio_tokens: int image_tokens: int + video_tokens: int character_count: int image_count: int video_length_seconds: float @@ -473,6 +474,7 @@ def _parse_prompt_tokens_details(usage: Usage) -> PromptTokensDetailsResult: ) audio_tokens = cast(Optional[int], getattr(usage.prompt_tokens_details, "audio_tokens", 0)) or 0 image_tokens = cast(Optional[int], getattr(usage.prompt_tokens_details, "image_tokens", 0)) or 0 + video_tokens = _coerce_token_count(getattr(usage.prompt_tokens_details, "video_tokens", 0)) character_count = ( cast( Optional[int], @@ -503,6 +505,7 @@ def _parse_prompt_tokens_details(usage: Usage) -> PromptTokensDetailsResult: text_tokens=text_tokens, audio_tokens=audio_tokens, image_tokens=image_tokens, + video_tokens=video_tokens, character_count=character_count, image_count=image_count, video_length_seconds=float(video_length_seconds), @@ -515,6 +518,7 @@ class CompletionTokensDetailsResult(TypedDict): text_tokens: int reasoning_tokens: int image_tokens: int + video_tokens: int def _parse_completion_tokens_details(usage: Usage) -> CompletionTokensDetailsResult: @@ -546,12 +550,14 @@ def _parse_completion_tokens_details(usage: Usage) -> CompletionTokensDetailsRes ) or 0 ) + video_tokens = _coerce_token_count(getattr(usage.completion_tokens_details, "video_tokens", 0)) return CompletionTokensDetailsResult( audio_tokens=audio_tokens, text_tokens=text_tokens, reasoning_tokens=reasoning_tokens, image_tokens=image_tokens, + video_tokens=video_tokens, ) @@ -586,6 +592,13 @@ def _calculate_input_cost( image_token_cost_key = "input_cost_per_token" prompt_cost += calculate_cost_component(model_info, image_token_cost_key, prompt_tokens_details["image_tokens"]) + ### VIDEO TOKEN COST + if prompt_tokens_details["video_tokens"]: + video_token_cost_key = "input_cost_per_video_token" + if model_info.get(video_token_cost_key) is None: + video_token_cost_key = "input_cost_per_token" + prompt_cost += calculate_cost_component(model_info, video_token_cost_key, prompt_tokens_details["video_tokens"]) + ### CACHE WRITING COST - Now uses tiered pricing if ( prompt_tokens_details["cache_creation_tokens"] @@ -698,6 +711,7 @@ def generic_cost_per_token( text_tokens=usage.prompt_tokens, audio_tokens=0, image_tokens=0, + video_tokens=0, character_count=0, image_count=0, video_length_seconds=0.0, @@ -716,13 +730,14 @@ def generic_cost_per_token( audio_tokens = prompt_tokens_details["audio_tokens"] cache_creation = prompt_tokens_details["cache_creation_tokens"] image_tokens = prompt_tokens_details["image_tokens"] + video_tokens = prompt_tokens_details["video_tokens"] # Check for double-counting: sum of details > prompt_tokens means overlap - total_details = text_tokens + cache_hit + audio_tokens + cache_creation + image_tokens + total_details = text_tokens + cache_hit + audio_tokens + cache_creation + image_tokens + video_tokens has_double_counting = cache_hit > 0 and total_details > usage.prompt_tokens if (text_tokens == 0 and prompt_tokens_details["image_count"] == 0) or has_double_counting: - text_tokens = usage.prompt_tokens - cache_hit - audio_tokens - cache_creation - image_tokens + text_tokens = usage.prompt_tokens - cache_hit - audio_tokens - cache_creation - image_tokens - video_tokens # Clamp to zero: inconsistent streaming usage if text_tokens < 0: text_tokens = 0 @@ -751,6 +766,7 @@ def generic_cost_per_token( audio_tokens = 0 reasoning_tokens = 0 image_tokens = 0 + video_tokens = 0 is_text_tokens_total = False if usage.completion_tokens_details is not None: completion_tokens_details = _parse_completion_tokens_details(usage) @@ -758,19 +774,20 @@ def generic_cost_per_token( text_tokens = completion_tokens_details["text_tokens"] reasoning_tokens = completion_tokens_details["reasoning_tokens"] image_tokens = completion_tokens_details["image_tokens"] + video_tokens = completion_tokens_details["video_tokens"] # Handle text_tokens calculation: # 1. If text_tokens is explicitly provided and > 0, use it - # 2. If there's a breakdown (reasoning/audio/image tokens), calculate text_tokens as the remainder + # 2. If there's a breakdown (reasoning/audio/image/video tokens), calculate text_tokens as the remainder # 3. If no breakdown at all, assume all completion_tokens are text_tokens - has_token_breakdown = image_tokens > 0 or audio_tokens > 0 or reasoning_tokens > 0 + has_token_breakdown = image_tokens > 0 or audio_tokens > 0 or reasoning_tokens > 0 or video_tokens > 0 if text_tokens == 0: if has_token_breakdown: # Calculate text tokens as remainder when we have a breakdown # This handles cases like OpenAI's reasoning models where text_tokens isn't provided text_tokens = max( 0, - usage.completion_tokens - reasoning_tokens - audio_tokens - image_tokens, + usage.completion_tokens - reasoning_tokens - audio_tokens - image_tokens - video_tokens, ) else: # No breakdown at all, all tokens are text tokens @@ -803,6 +820,14 @@ def generic_cost_per_token( ) completion_cost += float(image_tokens) * _output_cost_per_image_token + ## VIDEO COST + if not is_text_tokens_total and video_tokens and video_tokens > 0: + _output_cost_per_video_token = _get_cost_per_unit(model_info, "output_cost_per_video_token", None) + _output_cost_per_video_token = ( + _output_cost_per_video_token if _output_cost_per_video_token is not None else completion_base_cost + ) + completion_cost += float(video_tokens) * _output_cost_per_video_token + ## REGIONAL DATA-RESIDENCY UPLIFT # Applied as a flat multiplier across all token costs for the request # when the upstream is a regionalized OpenAI host (eu./us.api.openai.com). diff --git a/litellm/llms/base_llm/guardrail_translation/base_translation.py b/litellm/llms/base_llm/guardrail_translation/base_translation.py index 6c41b46cfa0..bed06832386 100644 --- a/litellm/llms/base_llm/guardrail_translation/base_translation.py +++ b/litellm/llms/base_llm/guardrail_translation/base_translation.py @@ -1,4 +1,5 @@ from abc import ABC, abstractmethod +from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, Dict, List, Optional if TYPE_CHECKING: @@ -11,10 +12,30 @@ if TYPE_CHECKING: from litellm.types.llms.openai import AllMessageValues +@dataclass(slots=True) +class StreamTransformSink: + """Out-parameter used by ``process_output_streaming_response`` to hand the + guardrailed streaming state back to the caller. + + The streaming text-transform path must not mutate ``responses_so_far`` (it is + the raw accumulator the guardrail re-reads every round), so the guardrailed + accumulated text per choice (``mutated_text_per_choice``, keyed by + ``StreamingChoices.index``) and the per-choice trailing holdback the guardrail + requested (``holdback_per_choice``, from ``stream_holdback_chars``) are + reported here instead of in place. Only the OpenAI chat handler populates this + today; the hook passes a fresh sink per round and reads it afterwards. A + mutable dataclass is deliberate: it is a write-once output parameter for a + single call, not shared state. + """ + + mutated_text_per_choice: dict[int, str] = field(default_factory=dict) + holdback_per_choice: dict[int, int] = field(default_factory=dict) + + class BaseTranslation(ABC): @staticmethod def transform_user_api_key_dict_to_metadata( - user_api_key_dict: Optional[Any], + user_api_key_dict: Any | None, ) -> Dict[str, Any]: """ Transform user_api_key_dict to a metadata dict with prefixed keys. @@ -73,7 +94,7 @@ class BaseTranslation(ABC): guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: Optional["LiteLLMLoggingObj"] = None, user_api_key_dict: Optional["UserAPIKeyAuth"] = None, - request_data: Optional[dict] = None, + request_data: dict | None = None, ) -> Any: """ Process output response with guardrails. @@ -92,12 +113,15 @@ class BaseTranslation(ABC): guardrail_to_apply: "CustomGuardrail", litellm_logging_obj: Optional["LiteLLMLoggingObj"] = None, user_api_key_dict: Optional["UserAPIKeyAuth"] = None, - request_data: Optional[dict] = None, + request_data: dict | None = None, + stream_transform_sink: StreamTransformSink | None = None, ) -> Any: """ Process output streaming response with guardrails. - Optional to override in subclasses. + Optional to override in subclasses. ``stream_transform_sink`` is the + out-parameter used by handlers that support streaming text + transformations (see ``StreamTransformSink``); base handlers ignore it. """ return responses_so_far @@ -105,8 +129,8 @@ class BaseTranslation(ABC): self, exc: "ModifyResponseException", stream_started: bool = False, - responses_so_far: Optional[list[Any]] = None, - ) -> Optional[list[bytes]]: + responses_so_far: list[Any] | None = None, + ) -> list[bytes] | None: """ Build the streaming chunks that deliver a guardrail block message and cleanly terminate the stream in this provider's wire format. @@ -125,7 +149,7 @@ class BaseTranslation(ABC): """ return None - def get_structured_messages(self, data: dict) -> Optional[List["AllMessageValues"]]: + def get_structured_messages(self, data: dict) -> List["AllMessageValues"] | None: """ Convert request data to OpenAI-spec structured messages. diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index a6b6a6267c3..fd2c9339248 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -14,11 +14,14 @@ Pattern Overview: This pattern can be replicated for other message formats (e.g., Anthropic). """ -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast +from typing import TYPE_CHECKING, Any, Dict, List, Tuple, Union, cast import litellm from litellm._logging import verbose_proxy_logger -from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation +from litellm.llms.base_llm.guardrail_translation.base_translation import ( + BaseTranslation, + StreamTransformSink, +) from litellm.llms.base_llm.guardrail_translation.utils import ( effective_skip_system_message_for_guardrail, effective_skip_tool_message_for_guardrail, @@ -27,6 +30,9 @@ from litellm.llms.base_llm.guardrail_translation.utils import ( ) from litellm.main import stream_chunk_builder from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam +from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import ( + coerce_stream_holdback_value, +) from litellm.types.utils import ( Choices, GenericGuardrailAPIInputs, @@ -50,7 +56,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): Methods can be overridden to customize behavior for different message formats. """ - def get_structured_messages(self, data: dict) -> Optional[List[AllMessageValues]]: + def get_structured_messages(self, data: dict) -> List[AllMessageValues] | None: """ Convert chat completions request data to OpenAI-spec structured messages. @@ -65,7 +71,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): self, data: dict, guardrail_to_apply: "CustomGuardrail", - litellm_logging_obj: Optional[Any] = None, + litellm_logging_obj: Any | None = None, ) -> Any: """ Process input messages by applying guardrails to text content. @@ -80,7 +86,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): texts_to_check: List[str] = [] images_to_check: List[str] = [] tool_calls_to_check: List[ChatCompletionToolParam] = [] - text_task_mappings: List[Tuple[int, Optional[int]]] = [] + text_task_mappings: List[Tuple[int, int | None]] = [] tool_call_task_mappings: List[Tuple[int, int]] = [] # Step 1: Extract all text content, images, and tool calls @@ -184,7 +190,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): texts_to_check: List[str], images_to_check: List[str], tool_calls_to_check: List[ChatCompletionToolParam], - text_task_mappings: List[Tuple[int, Optional[int]]], + text_task_mappings: List[Tuple[int, int | None]], tool_call_task_mappings: List[Tuple[int, int]], skip_system_message: bool = False, skip_tool_message: bool = False, @@ -239,7 +245,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): self, messages: List[Dict[str, Any]], responses: List[str], - task_mappings: List[Tuple[int, Optional[int]]], + task_mappings: List[Tuple[int, int | None]], ) -> None: """ Apply guardrail responses back to input message text content. @@ -249,7 +255,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): for task_idx, guardrail_response in enumerate(responses): mapping = task_mappings[task_idx] msg_idx = cast(int, mapping[0]) - content_idx_optional = cast(Optional[int], mapping[1]) + content_idx_optional = cast(int | None, mapping[1]) # Handle content content = messages[msg_idx].get("content", None) @@ -291,9 +297,9 @@ class OpenAIChatCompletionsHandler(BaseTranslation): self, response: "ModelResponse", guardrail_to_apply: "CustomGuardrail", - litellm_logging_obj: Optional[Any] = None, - user_api_key_dict: Optional[Any] = None, - request_data: Optional[dict] = None, + litellm_logging_obj: Any | None = None, + user_api_key_dict: Any | None = None, + request_data: dict | None = None, ) -> Any: """ Process output response by applying guardrails to text content. @@ -320,7 +326,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): texts_to_check: List[str] = [] images_to_check: List[str] = [] tool_calls_to_check: List[Dict[str, Any]] = [] - text_task_mappings: List[Tuple[int, Optional[int]]] = [] + text_task_mappings: List[Tuple[int, int | None]] = [] tool_call_task_mappings: List[Tuple[int, int]] = [] # text_task_mappings: Track (choice_index, content_index) for each text # content_index is None for string content, int for list content @@ -402,9 +408,10 @@ class OpenAIChatCompletionsHandler(BaseTranslation): self, responses_so_far: List["ModelResponseStream"], guardrail_to_apply: "CustomGuardrail", - litellm_logging_obj: Optional[Any] = None, - user_api_key_dict: Optional[Any] = None, - request_data: Optional[dict] = None, + litellm_logging_obj: Any | None = None, + user_api_key_dict: Any | None = None, + request_data: dict | None = None, + stream_transform_sink: StreamTransformSink | None = None, ) -> List["ModelResponseStream"]: """ Process output streaming responses by applying guardrails to text content. @@ -414,14 +421,50 @@ class OpenAIChatCompletionsHandler(BaseTranslation): guardrail_to_apply: The guardrail instance to apply litellm_logging_obj: Optional logging object user_api_key_dict: User API key metadata to pass to guardrails + stream_transform_sink: Optional out-parameter for the streaming text + transformation path. When provided, the guardrail runs over the raw + accumulated text (``responses_so_far`` is left untouched so it stays + a correct raw accumulator across rounds) and the guardrailed text + plus requested holdback are reported per choice on the sink. Returns: - Modified list of responses with guardrail applied to content + The (unmodified) list of responses. Response Format Support: - String content: choice.message.content = "text here" - List content: choice.message.content = [{"type": "text", "text": "text here"}, ...] """ + if stream_transform_sink is not None: + await self._process_streaming_transform( + responses_so_far=responses_so_far, + guardrail_to_apply=guardrail_to_apply, + litellm_logging_obj=litellm_logging_obj, + user_api_key_dict=user_api_key_dict, + request_data=request_data, + sink=stream_transform_sink, + ) + return responses_so_far + + return await self._process_streaming_block_only( + responses_so_far=responses_so_far, + guardrail_to_apply=guardrail_to_apply, + litellm_logging_obj=litellm_logging_obj, + user_api_key_dict=user_api_key_dict, + request_data=request_data, + ) + + async def _process_streaming_block_only( + self, + *, + responses_so_far: list["ModelResponseStream"], + guardrail_to_apply: "CustomGuardrail", + litellm_logging_obj: Any | None, + user_api_key_dict: Any | None, + request_data: dict | None, + ) -> list["ModelResponseStream"]: + """Block-only streaming path: run the guardrail so an in-flight BLOCK can + terminate the stream. Text rewrites are not propagated to the client here + (see ``_process_streaming_transform`` for the incremental_diff path).""" # check if the stream has ended has_stream_ended = False for chunk in responses_so_far: @@ -467,7 +510,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): # Step 2: Create lists for guardrail processing texts_to_check: List[str] = [] images_to_check: List[str] = [] - task_mappings: List[Tuple[int, Optional[int]]] = [] + task_mappings: List[Tuple[int, int | None]] = [] # Track (choice_index, content_index) for each combined text for (map_choice_idx, map_content_idx), combined_text in combined_texts.items(): @@ -520,9 +563,109 @@ class OpenAIChatCompletionsHandler(BaseTranslation): return responses_so_far + @staticmethod + def _accumulate_string_content_by_choice_index( + responses_so_far: list["ModelResponseStream"], + ) -> dict[int, str]: + """Accumulate raw string ``delta.content`` per choice, keyed by + ``StreamingChoices.index`` (not enumerate position, which collapses to 0 + when each chunk carries a single non-zero-indexed choice for ``n > 1``). + + Only string content participates; list-of-blocks content is out of scope + for the incremental transform path. Reads ``responses_so_far`` without + mutating it so it stays a correct raw accumulator across rounds. + """ + accumulated: dict[int, str] = {} + for response in responses_so_far: + for choice in response.choices: + if isinstance(choice, litellm.StreamingChoices): + content = choice.delta.content + elif isinstance(choice, litellm.Choices): + content = choice.message.content + else: + continue + if isinstance(content, str) and content: + idx = getattr(choice, "index", 0) or 0 + accumulated[idx] = accumulated.get(idx, "") + content + return accumulated + + async def _process_streaming_transform( + self, + *, + responses_so_far: list["ModelResponseStream"], + guardrail_to_apply: "CustomGuardrail", + litellm_logging_obj: Any | None, + user_api_key_dict: Any | None, + request_data: dict | None, + sink: StreamTransformSink, + ) -> None: + """Run the guardrail over the raw accumulated text and report the + guardrailed text plus requested holdback per choice on ``sink``. + + Unlike the block-only path this never mutates ``responses_so_far``: it + re-derives the raw accumulated text every round (so a rewrite guardrail + always sees consistent input) and hands the result back out of band. + """ + raw_by_index = self._accumulate_string_content_by_choice_index(responses_so_far) + if not raw_by_index: + sink.mutated_text_per_choice = {} + sink.holdback_per_choice = {} + return + + # Fix #2 — sort by StreamingChoices.index so an n>1 stream that emits + # choice 1 before choice 0 still hands the guardrail texts in a + # deterministic index order. Without this, the guardrail's returned + # texts (aligned to the input order it received) would map back to the + # wrong choice indices when we rebuild the sink dicts by + # ``enumerate(indices)``. + indices = sorted(raw_by_index.keys()) + texts_to_check = [raw_by_index[i] for i in indices] + + if request_data is None: + request_data = {"responses": responses_so_far} + elif "responses" not in request_data: + request_data["responses"] = responses_so_far + if "litellm_metadata" not in request_data: + user_metadata = self.transform_user_api_key_dict_to_metadata(user_api_key_dict) + if user_metadata: + request_data["litellm_metadata"] = user_metadata + + inputs = GenericGuardrailAPIInputs(texts=texts_to_check) + if responses_so_far and getattr(responses_so_far[0], "model", None): + inputs["model"] = responses_so_far[0].model + guardrailed_inputs = await guardrail_to_apply.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="response", + logging_obj=litellm_logging_obj, + ) + + returned_texts = guardrailed_inputs.get("texts") + # No "texts" key means the guardrail made no change (action NONE): the raw + # accumulated text is the guardrailed text. A present-but-shorter list is a + # guardrail contract violation; those choices are omitted below (withheld, + # not emitted raw) so a malformed response fails closed instead of leaking. + if returned_texts is None: + returned_texts = texts_to_check + elif len(returned_texts) < len(texts_to_check): + verbose_proxy_logger.warning( + "OpenAI Chat Completions: guardrail returned %s transformed texts for %s inputs on the " + "streaming transform path; withholding the unmatched choices to fail closed.", + len(returned_texts), + len(texts_to_check), + ) + + holdback = guardrailed_inputs.get("stream_holdback_chars") or [] + sink.mutated_text_per_choice = { + idx: returned_texts[i] for i, idx in enumerate(indices) if i < len(returned_texts) + } + sink.holdback_per_choice = { + indices[i]: coerce_stream_holdback_value(holdback[i]) for i in range(len(indices)) if i < len(holdback) + } + def _combine_streaming_texts( self, responses_so_far: List["ModelResponseStream"] - ) -> Dict[Tuple[int, Optional[int]], str]: + ) -> Dict[Tuple[int, int | None], str]: """ Combine all streaming chunks into complete text per choice. @@ -534,7 +677,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): Returns: Dict mapping (choice_idx, content_idx) to combined text string """ - combined_texts: Dict[Tuple[int, Optional[int]], str] = {} + combined_texts: Dict[Tuple[int, int | None], str] = {} for response_idx, response in enumerate(responses_so_far): for choice_idx, choice in enumerate(response.choices): @@ -550,7 +693,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): if isinstance(content, str): # String content - accumulate for this choice - str_key: Tuple[int, Optional[int]] = (choice_idx, None) + str_key: Tuple[int, int | None] = (choice_idx, None) if str_key not in combined_texts: combined_texts[str_key] = "" combined_texts[str_key] += content @@ -560,7 +703,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): for content_idx, content_item in enumerate(content): text_str = content_item.get("text") if text_str: - list_key: Tuple[int, Optional[int]] = ( + list_key: Tuple[int, int | None] = ( choice_idx, content_idx, ) @@ -607,7 +750,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): texts_to_check: List[str], images_to_check: List[str], tool_calls_to_check: List[Dict[str, Any]], - text_task_mappings: List[Tuple[int, Optional[int]]], + text_task_mappings: List[Tuple[int, int | None]], tool_call_task_mappings: List[Tuple[int, int]], ) -> None: """ @@ -619,7 +762,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): # Determine content source and tool calls based on choice type content = None - tool_calls: Optional[List[Any]] = None + tool_calls: List[Any] | None = None if isinstance(choice, litellm.Choices): content = choice.message.content tool_calls = choice.message.tool_calls @@ -662,7 +805,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): tool_calls_to_check.append(tool_call_dict) tool_call_task_mappings.append((choice_idx, int(tool_call_idx))) - def _convert_tool_call_to_dict(self, tool_call: Union[Dict[str, Any], Any]) -> Optional[Dict[str, Any]]: + def _convert_tool_call_to_dict(self, tool_call: Union[Dict[str, Any], Any]) -> Dict[str, Any] | None: """ Convert a tool call object to dictionary format. @@ -691,7 +834,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): self, response: "ModelResponse", responses: List[str], - task_mappings: List[Tuple[int, Optional[int]]], + task_mappings: List[Tuple[int, int | None]], ) -> None: """ Apply guardrail text responses back to output response. @@ -701,7 +844,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): for task_idx, guardrail_response in enumerate(responses): mapping = task_mappings[task_idx] choice_idx = cast(int, mapping[0]) - content_idx_optional = cast(Optional[int], mapping[1]) + content_idx_optional = cast(int | None, mapping[1]) choice = cast(Choices, response.choices[choice_idx]) @@ -755,7 +898,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): self, responses: List["ModelResponseStream"], guardrailed_texts: List[str], - task_mappings: List[Tuple[int, Optional[int]]], + task_mappings: List[Tuple[int, int | None]], ) -> None: """ Apply guardrail responses back to output streaming responses. @@ -771,16 +914,16 @@ class OpenAIChatCompletionsHandler(BaseTranslation): Override this method to customize how responses are applied to streaming responses. """ # Build a mapping of what guardrailed text to use for each (choice_idx, content_idx) - guardrail_map: Dict[Tuple[int, Optional[int]], str] = {} + guardrail_map: Dict[Tuple[int, int | None], str] = {} for task_idx, guardrail_response in enumerate(guardrailed_texts): mapping = task_mappings[task_idx] choice_idx = cast(int, mapping[0]) - content_idx_optional = cast(Optional[int], mapping[1]) + content_idx_optional = cast(int | None, mapping[1]) guardrail_map[(choice_idx, content_idx_optional)] = guardrail_response # Track which choices we've already set the guardrailed text for # Key: (choice_idx, content_idx), Value: boolean (True if already set) - already_set: Dict[Tuple[int, Optional[int]], bool] = {} + already_set: Dict[Tuple[int, int | None], bool] = {} # Iterate through all responses and update content for response_idx, response in enumerate(responses): @@ -797,7 +940,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): if isinstance(content, str): # String content - str_key: Tuple[int, Optional[int]] = (choice_idx_in_response, None) + str_key: Tuple[int, int | None] = (choice_idx_in_response, None) if str_key in guardrail_map: if str_key not in already_set: # First chunk - set the complete guardrailed text @@ -817,7 +960,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): # List content - handle each content item for content_idx, content_item in enumerate(content): if "text" in content_item: - list_key: Tuple[int, Optional[int]] = ( + list_key: Tuple[int, int | None] = ( choice_idx_in_response, content_idx, ) diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index f9ed8cea9b5..8c4bb1aa0c5 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -998,6 +998,8 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): response_modalities.append("IMAGE") elif modality == "audio": response_modalities.append("AUDIO") + elif modality == "video": + response_modalities.append("VIDEO") else: response_modalities.append("MODALITY_UNSPECIFIED") return response_modalities diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index c99347e189c..77ca423b866 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -19612,6 +19612,39 @@ }, "web_search_billing_unit": "per_query" }, + "gemini/gemini-omni-flash-preview": { + "input_cost_per_audio_token": 1.5e-06, + "input_cost_per_token": 1.5e-06, + "litellm_provider": "gemini", + "max_input_tokens": 1048576, + "max_output_tokens": 65535, + "max_tokens": 65535, + "mode": "chat", + "output_cost_per_reasoning_token": 9e-06, + "output_cost_per_token": 9e-06, + "output_cost_per_video_token": 1.75e-05, + "rpm": 2000, + "source": "https://ai.google.dev/gemini-api/docs/pricing", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text", + "video" + ], + "supports_audio_input": true, + "supports_reasoning": true, + "supports_system_messages": true, + "supports_video_input": true, + "supports_vision": true, + "tpm": 800000 + }, "gemini/gemini-3.1-pro-preview": { "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, @@ -19776,6 +19809,37 @@ }, "web_search_billing_unit": "per_query" }, + "gemini-omni-flash-preview": { + "input_cost_per_audio_token": 1.5e-06, + "input_cost_per_token": 1.5e-06, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 1048576, + "max_output_tokens": 65535, + "max_tokens": 65535, + "mode": "chat", + "output_cost_per_reasoning_token": 9e-06, + "output_cost_per_token": 9e-06, + "output_cost_per_video_token": 1.75e-05, + "source": "https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/gemini/omni-flash-preview", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text", + "video" + ], + "supports_audio_input": true, + "supports_reasoning": true, + "supports_system_messages": true, + "supports_video_input": true, + "supports_vision": true + }, "gemini-3.5-flash": { "cache_read_input_token_cost": 1.5e-07, "input_cost_per_audio_token": 1e-06, diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 744d8182715..b402212fb2e 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -1191,13 +1191,15 @@ async def _user_api_key_auth_builder( return await handle_oauth2_proxy_request(request=request) if general_settings.get("enable_jwt_auth", False) is True: - from litellm.proxy.proxy_server import premium_user - - if premium_user is not True: - raise ValueError(f"JWT Auth is an enterprise only feature. {CommonProxyErrors.not_premium_user.value}") is_jwt = jwt_handler.is_jwt(token=api_key) verbose_proxy_logger.debug("is_jwt: %s", is_jwt) if is_jwt: + from litellm.proxy.proxy_server import premium_user + + if premium_user is not True: + raise ValueError( + f"JWT Auth is an enterprise only feature. {CommonProxyErrors.not_premium_user.value}" + ) # Try JWT-to-Virtual-Key mapping first to avoid # unnecessary DB queries in auth_builder do_standard_jwt_auth = True diff --git a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/__init__.py index 63ead52baa6..a3ac1b94004 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/__init__.py @@ -38,6 +38,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" default_on=litellm_params.default_on, streaming_end_of_stream_only=_get_config_value(litellm_params, optional_params, "streaming_end_of_stream_only"), streaming_sampling_rate=_get_config_value(litellm_params, optional_params, "streaming_sampling_rate"), + streaming_transform_mode=_get_config_value(litellm_params, optional_params, "streaming_transform_mode"), ) litellm.logging_callback_manager.add_litellm_callback(_generic_guardrail_api_callback) diff --git a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py index e29d6e56353..6b1c9a8b2b3 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py +++ b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py @@ -58,7 +58,7 @@ _HEADER_PRESENT_PLACEHOLDER = "[present]" def _header_value_allowed( header_name: str, - extra_allowlist: Optional[Set[str]] = None, + extra_allowlist: Set[str] | None = None, ) -> bool: """Return True if this header's value may be forwarded (allowlist, including globs and extra_headers).""" lower = header_name.lower() @@ -74,8 +74,8 @@ def _header_value_allowed( def _sanitize_inbound_headers( headers: Any, - extra_allowlist: Optional[Set[str]] = None, -) -> Optional[Dict[str, str]]: + extra_allowlist: Set[str] | None = None, +) -> Dict[str, str] | None: """ Sanitize inbound headers before passing them to a 3rd party guardrail service. @@ -105,8 +105,8 @@ def _sanitize_inbound_headers( def _extract_inbound_headers( request_data: dict, logging_obj: Optional["LiteLLMLoggingObj"], - extra_allowlist: Optional[Set[str]] = None, -) -> Optional[Dict[str, str]]: + extra_allowlist: Set[str] | None = None, +) -> Dict[str, str] | None: """ Extract inbound headers from available request context. @@ -172,15 +172,16 @@ class GenericGuardrailAPI(CustomGuardrail): def __init__( self, - headers: Optional[Dict[str, Any]] = None, - api_base: Optional[str] = None, - api_key: Optional[str] = None, - additional_provider_specific_params: Optional[Dict[str, Any]] = None, + headers: Dict[str, Any] | None = None, + api_base: str | None = None, + api_key: str | None = None, + additional_provider_specific_params: Dict[str, Any] | None = None, unreachable_fallback: Literal["fail_closed", "fail_open"] = "fail_closed", - fail_on_error: Optional[bool] = True, - extra_headers: Optional[list] = None, - streaming_end_of_stream_only: Optional[bool] = None, - streaming_sampling_rate: Optional[int] = None, + fail_on_error: bool | None = True, + extra_headers: list | None = None, + streaming_end_of_stream_only: bool | None = None, + streaming_sampling_rate: int | None = None, + streaming_transform_mode: Literal["block_only", "incremental_diff"] | None = None, **kwargs, ): self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) @@ -221,6 +222,13 @@ class GenericGuardrailAPI(CustomGuardrail): raise ValueError(f"streaming_sampling_rate must be >= 1 (got {streaming_sampling_rate})") self.streaming_sampling_rate: int = 5 if streaming_sampling_rate is None else streaming_sampling_rate + # Read by UnifiedLLMGuardrails.async_post_call_streaming_iterator_hook. + # "block_only" (default) drops text rewrites on the streaming path; + # "incremental_diff" emits them as synthetic deltas. + self.streaming_transform_mode: Literal["block_only", "incremental_diff"] = ( + "block_only" if streaming_transform_mode is None else streaming_transform_mode + ) + # Set supported event hooks kwargs.setdefault("supported_event_hooks", list(self.get_supported_event_hooks())) @@ -280,7 +288,7 @@ class GenericGuardrailAPI(CustomGuardrail): input_type: Literal["request", "response"], logging_obj: Optional["LiteLLMLoggingObj"], error: Exception, - http_status_code: Optional[int] = None, + http_status_code: int | None = None, ) -> GenericGuardrailAPIInputs: status_suffix = f" http_status_code={http_status_code}" if http_status_code else "" verbose_proxy_logger.critical( @@ -326,6 +334,8 @@ class GenericGuardrailAPI(CustomGuardrail): return_inputs["tools"] = guardrail_response.tools elif tools: return_inputs["tools"] = tools + if guardrail_response.stream_holdback_chars is not None: + return_inputs["stream_holdback_chars"] = guardrail_response.stream_holdback_chars return return_inputs def _handle_guardrail_request_error( @@ -479,7 +489,7 @@ class GenericGuardrailAPI(CustomGuardrail): return self._handle_guardrail_request_error(e, inputs, input_type, logging_obj, is_unreachable=False) @staticmethod - def get_config_model() -> Optional[type["GuardrailConfigModel"]]: + def get_config_model() -> type["GuardrailConfigModel"] | None: from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import ( GenericGuardrailAPIConfigModel, ) diff --git a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py index e0b387e92c0..8e3abfbf159 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -8,7 +8,7 @@ Unified Guardrail, leveraging LiteLLM's /applyGuardrail endpoint import copy import json -from typing import TYPE_CHECKING, Any, AsyncGenerator, List, Optional, Union +from typing import TYPE_CHECKING, Any, AsyncGenerator, List, Union from fastapi import HTTPException @@ -21,7 +21,13 @@ from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_fo from litellm.llms import load_guardrail_translation_mappings from litellm.proxy._types import UserAPIKeyAuth from litellm.types.guardrails import GuardrailEventHooks -from litellm.types.utils import CallTypes, CallTypesLiteral +from litellm.types.utils import ( + CallTypes, + CallTypesLiteral, + Delta, + ModelResponseStream, + StreamingChoices, +) if TYPE_CHECKING: # Imported lazily at runtime (inside the streaming hook) to avoid a @@ -34,7 +40,12 @@ A2A_CALL_TYPES = (CallTypes.asend_message, CallTypes.send_message) GUARDRAIL_NAME = "unified_llm_guardrails" -def _get_a2a_request_id(responses_so_far: List[Any], request_data: dict) -> Optional[str]: +class _StreamTerminated(Exception): + """Internal signal that the incremental transform stream has already emitted + its terminal chunks (block message or in-stream error) and must stop.""" + + +def _get_a2a_request_id(responses_so_far: List[Any], request_data: dict) -> str | None: """Get JSON-RPC request id from first A2A chunk or request body for in-stream error reporting.""" for item in responses_so_far: if isinstance(item, dict) and "id" in item: @@ -216,7 +227,7 @@ class UnifiedLLMGuardrails(CustomLogger): verbose_proxy_logger.debug("async_post_call_success_hook response: %s", response) - call_type: Optional[CallTypesLiteral] = None + call_type: CallTypesLiteral | None = None if user_api_key_dict.request_route is not None: call_types = get_call_types_for_route(user_api_key_dict.request_route) if call_types is not None and len(call_types) > 0: # type: ignore @@ -292,6 +303,498 @@ class UnifiedLLMGuardrails(CustomLogger): for chunk in block_chunks: yield chunk + @staticmethod + def _resolve_transform_call_type( + user_api_key_dict: UserAPIKeyAuth, + mappings: dict, + ) -> str | None: + """Resolve the call type for the incremental_diff path, or None if the + route is unresolvable / unsupported. + + Incremental transformation needs a route we can resolve before the first + chunk and a handler that supports the streaming text-diff protocol (v1: + the OpenAI chat completions handler only). Returning None makes the caller + fall back to block_only. + """ + from litellm.llms.openai.chat.guardrail_translation.handler import ( + OpenAIChatCompletionsHandler, + ) + + if user_api_key_dict.request_route is None: + return None + call_types = get_call_types_for_route(user_api_key_dict.request_route) + if not call_types: + return None + call_type = call_types[0].value + try: + mapped = CallTypes(call_type) + except ValueError: + return None + handler_cls = mappings.get(mapped) + if handler_cls is None or not issubclass(handler_cls, OpenAIChatCompletionsHandler): + return None + return call_type + + async def _emit_streaming_http_error( + self, + exc: HTTPException, + call_type: str | None, + responses_so_far: list[Any], + request_data: dict, + ) -> AsyncGenerator[Any, None]: + """Surface a mid-stream HTTPException. For A2A (NDJSON) call types the + response has already started, so emit an in-stream JSON-RPC error chunk; + otherwise re-raise so the proxy can report it. + """ + if call_type is not None and CallTypes(call_type) in A2A_CALL_TYPES: + request_id = _get_a2a_request_id(responses_so_far, request_data) + detail = exc.detail if isinstance(exc.detail, dict) else {"message": str(exc.detail)} + error_chunk = ( + json.dumps( + { + "jsonrpc": "2.0", + "id": request_id, + "error": { + "code": -32603, + "message": detail.get("error", detail.get("message", str(exc.detail))), + "data": {k: v for k, v in detail.items() if k not in ("error", "message")}, + }, + } + ) + + "\n" + ) + yield error_chunk + return + raise exc + + def _build_transform_chunk( + self, + *, + reference_chunk: Any, + mutated_text_per_choice: dict[int, str], + emitted_text_per_choice: dict[int, str], + holdback_per_choice: dict[int, int], + finish_reason_per_choice: dict[int, str | None], + is_final: bool, + ) -> ModelResponseStream | None: + """Build the synthetic chunk carrying the newly-guardrailed deltas. + + For each choice, the new delta is the mutated accumulated text past what + has already been emitted, minus a trailing holdback (forced to 0 on the + final flush). ``emitted_text_per_choice`` holds the exact bytes already + sent per choice and is extended in place. Returns None when there is no + text to emit (e.g. a tool-call-only turn) or nothing new and this is not + the final chunk. + + Raises HTTPException(400, stream_transform_underflow) when the guardrail's + transform is not a forward extension of what has already been streamed + (shorter than, or rewrites, the already-sent prefix), since emitted bytes + cannot be retracted. This makes the framework fail closed rather than + silently leave un-transformed text on the wire; a guardrail that needs to + rewrite recent output must withhold it first via ``stream_holdback_chars``. + """ + if not mutated_text_per_choice: + # Fix #4 — on the final flush a deferred finish_reason (from a mixed + # content+tool_calls chunk whose passthrough suppressed it) still + # needs to reach the client, even if the guardrail returned no text + # to emit. Build a terminator chunk carrying finish_reason per choice. + if is_final and finish_reason_per_choice: + terminator_choices: list[StreamingChoices] = [] + for choice_idx, finish_reason in finish_reason_per_choice.items(): + if finish_reason is None: + continue + terminator_choices.append( + StreamingChoices( + index=choice_idx, + delta=Delta(content="", role=None, tool_calls=None), + finish_reason=finish_reason, + ) + ) + if terminator_choices: + return ModelResponseStream( + id=getattr(reference_chunk, "id", None), + created=getattr(reference_chunk, "created", None), + model=getattr(reference_chunk, "model", None), + choices=terminator_choices, + ) + return None + + deltas: dict[int, str] = {} + for choice_idx, text in mutated_text_per_choice.items(): + already = emitted_text_per_choice.get(choice_idx, "") + if not text.startswith(already): + raise HTTPException( + status_code=400, + detail={ + "error": "stream_transform_underflow", + "message": ( + f"Guardrail streaming transform for choice {choice_idx} is not a forward " + f"extension of the {len(already)} chars already streamed to the client " + "(it is shorter than, or rewrites, the emitted prefix); emitted bytes " + "cannot be retracted. Withhold recent output via stream_holdback_chars " + "before rewriting it." + ), + }, + ) + holdback = 0 if is_final else max(0, holdback_per_choice.get(choice_idx, 0)) + end = max(len(already), len(text) - holdback) + deltas[choice_idx] = text[len(already) : end] + + # Iterate the mutated choices (not just those in reference_chunk) so a + # choice with pending text is never dropped for n > 1. finish_reason is + # taken per choice from the accumulated map (a choice can finish in an + # earlier chunk than the stream's last one); tool_calls are dropped since + # v1 does not transform streamed tool calls (they pass through raw). + synthetic_choices: list[StreamingChoices] = [] + for choice_idx in mutated_text_per_choice: + delta_text = deltas.get(choice_idx, "") + finish_reason = finish_reason_per_choice.get(choice_idx) if is_final else None + # Skip a choice with nothing to say: no new content and no + # finish_reason to deliver. This avoids emitting an empty delta for an + # already-finished choice (e.g. one that terminated via a passed-through + # tool-call chunk, which already carried its own finish_reason). + if not delta_text and finish_reason is None: + continue + # role="assistant" on this choice's first emitted delta only. + role = "assistant" if not emitted_text_per_choice.get(choice_idx) else None + synthetic_choices.append( + StreamingChoices( + index=choice_idx, + delta=Delta(content=delta_text, role=role, tool_calls=None), + finish_reason=finish_reason, + ) + ) + + if not synthetic_choices: + return None + + for choice_idx in mutated_text_per_choice: + emitted_text_per_choice[choice_idx] = emitted_text_per_choice.get(choice_idx, "") + deltas.get( + choice_idx, "" + ) + + return ModelResponseStream( + id=getattr(reference_chunk, "id", None), + created=getattr(reference_chunk, "created", None), + model=getattr(reference_chunk, "model", None), + choices=synthetic_choices, + ) + + async def _emit_transform_round( + self, + *, + endpoint_translation: Any, + guardrail_to_apply: CustomGuardrail, + request_data: dict, + user_api_key_dict: UserAPIKeyAuth, + call_type: str, + reference_chunk: Any, + responses_so_far: list[Any], + responses_yielded: list[Any], + emitted_text_per_choice: dict[int, str], + finish_reason_per_choice: dict[int, str | None], + is_final: bool, + ) -> AsyncGenerator[Any, None]: + """Run one guardrail processing round and emit the resulting diff chunk. + + Raises ``_StreamTerminated`` (after emitting the terminal block message or + in-stream error) when the guardrail blocks or an underflow occurs. + """ + from litellm.integrations.custom_guardrail import ModifyResponseException + from litellm.llms.base_llm.guardrail_translation.base_translation import ( + StreamTransformSink, + ) + + sink = StreamTransformSink() + try: + await endpoint_translation.process_output_streaming_response( + responses_so_far=responses_so_far, + guardrail_to_apply=guardrail_to_apply, + litellm_logging_obj=request_data.get("litellm_logging_obj"), + user_api_key_dict=user_api_key_dict, + request_data=request_data, + stream_transform_sink=sink, + ) + synthetic = self._build_transform_chunk( + reference_chunk=reference_chunk, + mutated_text_per_choice=sink.mutated_text_per_choice, + emitted_text_per_choice=emitted_text_per_choice, + holdback_per_choice=sink.holdback_per_choice, + finish_reason_per_choice=finish_reason_per_choice, + is_final=is_final, + ) + except ModifyResponseException as e: + if e.original_response is None: + e.original_response = responses_so_far + async for block_chunk in self._handle_streaming_block( + e, + endpoint_translation, + stream_started=bool(responses_yielded), + responses_so_far=responses_yielded, + ): + yield block_chunk + raise _StreamTerminated() + except HTTPException as e: + async for error_item in self._emit_streaming_http_error(e, call_type, responses_so_far, request_data): + yield error_item + raise _StreamTerminated() + + if synthetic is not None: + responses_yielded.append(synthetic) + yield synthetic + + async def _run_incremental_transform_stream( + self, + *, + guardrail_to_apply: CustomGuardrail, + response: Any, + request_data: dict, + user_api_key_dict: UserAPIKeyAuth, + call_type: str, + sampling_rate: int, + end_of_stream_only: bool, + mappings: dict, + ) -> AsyncGenerator[Any, None]: + """Emit guardrail text transformations as new deltas on the stream. + + Raw chunks are withheld and accumulated; on each sampled processing round + (and once at end of stream) the guardrailed accumulated text is diffed + against what has already been emitted and the new portion is sent as a + synthetic chunk. A BLOCK terminates the stream via the shared block + handler; an underflow surfaces as an HTTPException. + """ + endpoint_translation = mappings[CallTypes(call_type)]() + responses_so_far: list[Any] = [] + responses_yielded: list[Any] = [] + emitted_text_per_choice: dict[int, str] = {} + finish_reason_per_choice: dict[int, str | None] = {} + chunk_counter = 0 + last_chunk: Any | None = None + + def _round(reference_chunk: Any, is_final: bool) -> AsyncGenerator[Any, None]: + return self._emit_transform_round( + endpoint_translation=endpoint_translation, + guardrail_to_apply=guardrail_to_apply, + request_data=request_data, + user_api_key_dict=user_api_key_dict, + call_type=call_type, + reference_chunk=reference_chunk, + responses_so_far=responses_so_far, + responses_yielded=responses_yielded, + emitted_text_per_choice=emitted_text_per_choice, + finish_reason_per_choice=finish_reason_per_choice, + is_final=is_final, + ) + + saw_tool_calls = False + saw_text_content = False + + try: + async for item in response: + # v1 transforms only text. A chunk carrying tool_calls is passed + # through raw so function-calling turns are not dropped, but ONLY + # its tool-call fields are forwarded: content is stripped so any + # response text (in the same delta, or in another choice of an n>1 + # chunk) can never bypass the transform. The original chunk is kept + # in responses_so_far so its text is still accumulated + redacted + + # emitted as synthetic deltas, and so the guardrail inspects the + # assembled tool calls at end of stream (see the block inspection + # below), matching block_only. finish_reason rides on the raw + # tool-only chunk, so it is not recorded for the text flush. + if self._chunk_has_tool_calls(item): + saw_tool_calls = True + responses_so_far.append(item) + last_chunk = item + # Fix #3 — flush accumulated text BEFORE the tool-call + # passthrough. Without this, a stream of text chunks that + # hasn't yet hit a sampled round can be trailed by a + # tool-call chunk carrying finish_reason="tool_calls"; an + # SSE-compliant client stops reading at that finish_reason + # and drops the end-of-stream text flush that would follow. + if saw_text_content: + async for out in _round(item, is_final=False): + yield out + # Fix #1 — pass finish_reason_per_choice into the + # passthrough so a mixed content+tool_call chunk defers its + # finish_reason to the final text terminator (see the + # _tool_call_passthrough_chunk docstring). + tool_only = self._tool_call_passthrough_chunk( + item, finish_reason_per_choice=finish_reason_per_choice + ) + responses_yielded.append(tool_only) + yield tool_only + continue + + chunk_counter += 1 + responses_so_far.append(item) + last_chunk = item + self._record_finish_reasons(item, finish_reason_per_choice) + if self._chunk_carries_text(item): + saw_text_content = True + # Skip the sampled round for a terminal chunk: the end-of-stream + # flush below processes it once with holdback forced to 0, so a + # sampled round here would guardrail the same content twice. + if ( + not end_of_stream_only + and not self._chunk_has_finish_reason(item) + and chunk_counter % sampling_rate == 0 + ): + async for out in _round(item, is_final=False): + yield out + + # v1 does not transform streamed tool calls, but they must still go + # through the guardrail's block decision. Run the block_only inspection + # over the full assembled response so tool calls cannot bypass it. + # + # Pass a deep copy of responses_so_far — the block path routes through + # ``_process_streaming_block_only`` which mutates ``delta.content`` + # in-place on the chunk objects it receives. For an n>1 chunk carrying + # text on one choice and tool_calls (with finish_reason) on another, + # ``has_stream_ended`` reads ``choices[0]`` alone and can miss the + # terminal signal, letting the block path rewrite the raw accumulator. + # The subsequent final ``_round`` would then re-read the already-mutated + # text, producing double-application for a non-idempotent guardrail or a + # ``stream_transform_underflow`` 400 from mismatched prefixes. A shallow + # list copy wouldn't help — the mutation is on the chunk objects + # themselves — so we deepcopy. + if saw_tool_calls: + async for out in self._inspect_full_response_for_block( + endpoint_translation=endpoint_translation, + guardrail_to_apply=guardrail_to_apply, + request_data=request_data, + user_api_key_dict=user_api_key_dict, + responses_so_far=copy.deepcopy(responses_so_far), + responses_yielded=responses_yielded, + ): + yield out + + if last_chunk is not None: + async for out in _round(last_chunk, is_final=True): + yield out + except _StreamTerminated: + return + + async def _inspect_full_response_for_block( + self, + *, + endpoint_translation: Any, + guardrail_to_apply: CustomGuardrail, + request_data: dict, + user_api_key_dict: UserAPIKeyAuth, + responses_so_far: list[Any], + responses_yielded: list[Any], + ) -> AsyncGenerator[Any, None]: + """Run the block-only guardrail inspection over the full assembled + response (text + tool calls) so nothing bypasses the block decision. + + The guardrail's returned transforms are discarded here (v1 does not + transform tool calls); only its block decision matters. A block is + surfaced the same way as elsewhere: ModifyResponseException terminates the + stream via the shared block handler; a GenericGuardrailAPI block raises and + propagates, matching block_only. + """ + from litellm.integrations.custom_guardrail import ModifyResponseException + + try: + await endpoint_translation.process_output_streaming_response( + responses_so_far=responses_so_far, + guardrail_to_apply=guardrail_to_apply, + litellm_logging_obj=request_data.get("litellm_logging_obj"), + user_api_key_dict=user_api_key_dict, + request_data=request_data, + stream_transform_sink=None, + ) + except ModifyResponseException as e: + if e.original_response is None: + e.original_response = responses_so_far + async for block_chunk in self._handle_streaming_block( + e, + endpoint_translation, + stream_started=bool(responses_yielded), + responses_so_far=responses_yielded, + ): + yield block_chunk + raise _StreamTerminated() + + @staticmethod + def _chunk_has_tool_calls(item: Any) -> bool: + for choice in getattr(item, "choices", None) or []: + delta = getattr(choice, "delta", None) + if getattr(delta, "tool_calls", None): + return True + return False + + @staticmethod + def _chunk_carries_text(item: Any) -> bool: + """True if any choice in this chunk has non-empty string ``delta.content``.""" + for choice in getattr(item, "choices", None) or []: + delta = getattr(choice, "delta", None) + content = getattr(delta, "content", None) + if isinstance(content, str) and content != "": + return True + return False + + @staticmethod + def _tool_call_passthrough_chunk( + item: Any, + finish_reason_per_choice: "dict[int, str | None] | None" = None, + ) -> ModelResponseStream: + """Copy of a chunk carrying tool calls with all text content stripped. + + Only tool_calls, role and finish_reason are forwarded; content is set to + None so response text can never be delivered raw (it flows through the + transform instead). Applies per choice so an n>1 chunk mixing a text + choice and a tool-call choice does not leak the text choice. + + For a choice that carries BOTH text content AND tool_calls, ``finish_reason`` + is suppressed on the passthrough and recorded on + ``finish_reason_per_choice`` (when provided) so the final synthetic text + chunk delivers it. Emitting the passthrough's ``finish_reason`` before the + text flush would let a spec-compliant SSE client stop reading at + ``finish_reason`` and silently drop the guardrailed text, defeating the + redaction purpose. + """ + synthetic_choices: list[StreamingChoices] = [] + for choice in getattr(item, "choices", None) or []: + delta = getattr(choice, "delta", None) + idx = getattr(choice, "index", 0) or 0 + original_finish = getattr(choice, "finish_reason", None) + has_text = isinstance(getattr(delta, "content", None), str) and getattr(delta, "content", "") != "" + if has_text and original_finish is not None and finish_reason_per_choice is not None: + finish_reason_per_choice[idx] = original_finish + passthrough_finish: str | None = None + else: + passthrough_finish = original_finish + synthetic_choices.append( + StreamingChoices( + index=idx, + delta=Delta( + content=None, + role=getattr(delta, "role", None), + tool_calls=getattr(delta, "tool_calls", None), + ), + finish_reason=passthrough_finish, + ) + ) + return ModelResponseStream( + id=getattr(item, "id", None), + created=getattr(item, "created", None), + model=getattr(item, "model", None), + choices=synthetic_choices, + ) + + @staticmethod + def _record_finish_reasons(item: Any, finish_reason_per_choice: dict[int, str | None]) -> None: + for choice in getattr(item, "choices", None) or []: + finish_reason = getattr(choice, "finish_reason", None) + if finish_reason is not None: + finish_reason_per_choice[getattr(choice, "index", 0) or 0] = finish_reason + + @staticmethod + def _chunk_has_finish_reason(item: Any) -> bool: + choices = getattr(item, "choices", None) or [] + return any(getattr(choice, "finish_reason", None) is not None for choice in choices) + async def async_post_call_streaming_iterator_hook( self, user_api_key_dict: UserAPIKeyAuth, @@ -334,6 +837,10 @@ class UnifiedLLMGuardrails(CustomLogger): sampling_rate = _streaming_flag("streaming_sampling_rate", 5) # Only apply the guardrail at end of stream (not per chunk). end_of_stream_only = _streaming_flag("streaming_end_of_stream_only", False) + # "block_only" (default) drops guardrail text rewrites on the streaming + # path; "incremental_diff" emits them as synthetic deltas (see + # _run_incremental_transform_stream). + streaming_transform_mode = _streaming_flag("streaming_transform_mode", "block_only") # Withhold every chunk until end-of-stream moderation passes, then # release the original chunks (clean) or only the block message # (blocked) -- moderating the whole response *before* any content @@ -380,6 +887,35 @@ class UnifiedLLMGuardrails(CustomLogger): if endpoint_guardrail_translation_mappings is None: endpoint_guardrail_translation_mappings = load_guardrail_translation_mappings() + # Streaming text transformation (incremental_diff) diverges enough from the + # block_only path that it runs as its own iterator. It requires a route we + # can resolve up front to an OpenAI-chat handler (the only supported v1 + # surface); anything else falls back to the block_only behavior below. + if streaming_transform_mode == "incremental_diff": + transform_call_type = self._resolve_transform_call_type( + user_api_key_dict=user_api_key_dict, + mappings=endpoint_guardrail_translation_mappings, + ) + if transform_call_type is not None: + async for transformed_item in self._run_incremental_transform_stream( + guardrail_to_apply=guardrail_to_apply, + response=response, + request_data=request_data, + user_api_key_dict=user_api_key_dict, + call_type=transform_call_type, + sampling_rate=sampling_rate, + end_of_stream_only=end_of_stream_only, + mappings=endpoint_guardrail_translation_mappings, + ): + yield transformed_item + return + verbose_proxy_logger.warning( + "UnifiedLLMGuardrails: streaming_transform_mode=incremental_diff is only supported " + "for the OpenAI chat completions streaming path with a resolvable request route; " + "falling back to block_only for %s", + getattr(guardrail_to_apply, "guardrail_name", None), + ) + # Infer call type from first chunk call_type = None chunk_counter = 0 diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py b/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py index d0ac8bb8998..9e198e13902 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py @@ -84,6 +84,24 @@ class GenericGuardrailAPIOptionalParams(BaseModel): ), ) + streaming_transform_mode: Optional[Literal["block_only", "incremental_diff"]] = Field( + default=None, + description=( + "Controls whether text modifications returned by the guardrail (action=" + "GUARDRAIL_INTERVENED with modified texts) reach the client on the streaming " + "path. 'block_only' (default) preserves the historical behavior: the raw " + "upstream chunks are streamed and only a BLOCK terminates the stream; text " + "rewrites are dropped. 'incremental_diff' withholds the raw chunks and instead " + "emits the guardrailed text as new deltas computed by diffing the mutated " + "accumulated text against what has already been sent, enabling PII masking, " + "pseudonym reversal, redaction and similar rewrites over HTTP. Only supported " + "for the OpenAI chat completions streaming path (string delta.content) and " + "ignored when streaming_end_of_stream_only is True except for a single " + "post-stream synthetic chunk. Defaults to 'block_only' in " + "GenericGuardrailAPI.__init__ when None." + ), + ) + class GenericGuardrailAPIConfigModel( GuardrailConfigModel[GenericGuardrailAPIOptionalParams], @@ -126,6 +144,20 @@ class GenericGuardrailAPIRequest(BaseModel): model: Optional[str] = None # the model being used for the LLM call +def coerce_stream_holdback_value(value: Any) -> int: + """Coerce a single ``stream_holdback_chars`` entry to a non-negative int. + + A guardrail returning a null, non-numeric, or negative holdback element must + not abort the streaming round, so malformed values degrade to 0 (no holdback) + rather than raising. Shared by response parsing (``from_dict``) and the + handler that applies holdback to in-process guardrail return values. + """ + try: + return max(0, int(value)) + except (TypeError, ValueError): + return 0 + + class GenericGuardrailAPIResponse: """Response model for the Generic Guardrail API""" @@ -134,6 +166,7 @@ class GenericGuardrailAPIResponse: tools: Optional[List[GuardrailToolParam]] action: str blocked_reason: Optional[str] + stream_holdback_chars: Optional[List[int]] def __init__( self, @@ -142,19 +175,29 @@ class GenericGuardrailAPIResponse: blocked_reason: Optional[str] = None, images: Optional[List[str]] = None, tools: Optional[List[GuardrailToolParam]] = None, + stream_holdback_chars: Optional[List[int]] = None, ): self.action = action self.blocked_reason = blocked_reason self.texts = texts self.images = images self.tools = tools + # Number of trailing chars, indexed the same as ``texts``, that the + # framework must withhold from streaming emission until the next + # processing round (word-boundary safety for text transformations). + self.stream_holdback_chars = stream_holdback_chars @classmethod def from_dict(cls, data: dict) -> "GenericGuardrailAPIResponse": + raw_holdback = data.get("stream_holdback_chars") + stream_holdback_chars = ( + [coerce_stream_holdback_value(value) for value in raw_holdback] if isinstance(raw_holdback, list) else None + ) return cls( action=data.get("action", "NONE"), blocked_reason=data.get("blocked_reason"), texts=data.get("texts"), images=data.get("images"), tools=data.get("tools"), + stream_holdback_chars=stream_holdback_chars, ) diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 0fafd4a27fe..5e4d069b8a9 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -209,6 +209,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): input_cost_per_query: Optional[float] # only for rerank models input_cost_per_image: Optional[float] # only for vertex ai models input_cost_per_image_token: Optional[float] # for gpt-image-1 and similar models + input_cost_per_video_token: Optional[float] # for gemini omni models with video input input_cost_per_audio_per_second: Optional[float] # only for vertex ai models input_cost_per_video_per_second: Optional[float] # only for vertex ai models input_cost_per_second: Optional[float] # for OpenAI Speech models @@ -234,6 +235,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): output_cost_per_character_above_128k_tokens: Optional[float] # only for vertex ai models output_cost_per_image: Optional[float] output_cost_per_image_token: Optional[float] + output_cost_per_video_token: Optional[float] # for gemini omni models with video output output_vector_size: Optional[int] output_cost_per_reasoning_token: Optional[float] output_cost_per_video_per_second: Optional[float] # only for vertex ai models @@ -3052,6 +3054,7 @@ class CustomPricingLiteLLMParams(BaseModel): output_cost_per_character_above_128k_tokens: Optional[float] = None output_cost_per_image: Optional[float] = None output_cost_per_image_token: Optional[float] = None + output_cost_per_video_token: Optional[float] = None output_cost_per_reasoning_token: Optional[float] = None output_cost_per_video_per_second: Optional[float] = None output_cost_per_audio_per_second: Optional[float] = None @@ -3061,6 +3064,7 @@ class CustomPricingLiteLLMParams(BaseModel): cache_read_input_token_cost_above_272k_tokens: Optional[float] = None cache_read_input_token_cost_above_512k_tokens: Optional[float] = None input_cost_per_image_token: Optional[float] = None + input_cost_per_video_token: Optional[float] = None input_cost_per_token_above_272k_tokens: Optional[float] = None input_cost_per_token_above_512k_tokens: Optional[float] = None output_cost_per_token_above_272k_tokens: Optional[float] = None @@ -3780,3 +3784,6 @@ class GenericGuardrailAPIInputs(TypedDict, total=False): AllMessageValues ] # structured messages sent to the LLM - indicates if text is from system or user model: Optional[str] # the model being used for the LLM call + stream_holdback_chars: List[ + int + ] # trailing chars to withhold from streaming emission per text (word-boundary safety) diff --git a/litellm/utils.py b/litellm/utils.py index 18b89ee0d13..0636d3683b7 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -5437,6 +5437,7 @@ def _get_model_info_helper( input_cost_per_second=_model_info.get("input_cost_per_second", None), input_cost_per_audio_token=_model_info.get("input_cost_per_audio_token", None), input_cost_per_image_token=_model_info.get("input_cost_per_image_token", None), + input_cost_per_video_token=_model_info.get("input_cost_per_video_token", None), input_cost_per_image=_model_info.get("input_cost_per_image", None), input_cost_per_audio_per_second=_model_info.get("input_cost_per_audio_per_second", None), input_cost_per_video_per_second=_model_info.get("input_cost_per_video_per_second", None), @@ -5480,6 +5481,7 @@ def _get_model_info_helper( output_cost_per_video_per_second=_model_info.get("output_cost_per_video_per_second", None), output_cost_per_image=_model_info.get("output_cost_per_image", None), output_cost_per_image_token=_model_info.get("output_cost_per_image_token", None), + output_cost_per_video_token=_model_info.get("output_cost_per_video_token", None), output_vector_size=_model_info.get("output_vector_size", None), citation_cost_per_token=_model_info.get("citation_cost_per_token", None), tiered_pricing=_model_info.get("tiered_pricing", None), diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 8f2de0d0398..df6e8c992ef 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -19690,6 +19690,39 @@ }, "web_search_billing_unit": "per_query" }, + "gemini/gemini-omni-flash-preview": { + "input_cost_per_audio_token": 1.5e-06, + "input_cost_per_token": 1.5e-06, + "litellm_provider": "gemini", + "max_input_tokens": 1048576, + "max_output_tokens": 65535, + "max_tokens": 65535, + "mode": "chat", + "output_cost_per_reasoning_token": 9e-06, + "output_cost_per_token": 9e-06, + "output_cost_per_video_token": 1.75e-05, + "rpm": 2000, + "source": "https://ai.google.dev/gemini-api/docs/pricing", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text", + "video" + ], + "supports_audio_input": true, + "supports_reasoning": true, + "supports_system_messages": true, + "supports_video_input": true, + "supports_vision": true, + "tpm": 800000 + }, "gemini/gemini-3.1-pro-preview": { "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, @@ -19854,6 +19887,37 @@ }, "web_search_billing_unit": "per_query" }, + "gemini-omni-flash-preview": { + "input_cost_per_audio_token": 1.5e-06, + "input_cost_per_token": 1.5e-06, + "litellm_provider": "vertex_ai-language-models", + "max_input_tokens": 1048576, + "max_output_tokens": 65535, + "max_tokens": 65535, + "mode": "chat", + "output_cost_per_reasoning_token": 9e-06, + "output_cost_per_token": 9e-06, + "output_cost_per_video_token": 1.75e-05, + "source": "https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/gemini/omni-flash-preview", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text", + "video" + ], + "supports_audio_input": true, + "supports_reasoning": true, + "supports_system_messages": true, + "supports_video_input": true, + "supports_vision": true + }, "gemini-3.5-flash": { "cache_read_input_token_cost": 1.5e-07, "input_cost_per_audio_token": 1e-06, diff --git a/tests/e2e/CLAUDE.md b/tests/e2e/CLAUDE.md index 6b1b9950656..15e35c30ef7 100644 --- a/tests/e2e/CLAUDE.md +++ b/tests/e2e/CLAUDE.md @@ -77,13 +77,13 @@ llm..... endpoint : chat_completions | messages | responses | embeddings | batches | files | rerank | images_generations | audio_speech | audio_transcriptions | moderations | realtime - route : openai | azure_openai | anthropic | bedrock_converse | vertex | azure_foundry - | cohere | together_ai + route : openai | azure_openai | anthropic | bedrock_converse | bedrock_invoke | vertex + | azure_foundry | cohere | together_ai (vocab varies per endpoint; messages is anthropic-format only) capability : basic | tool_use | prompt_cache_5m | vision | thinking | structured_output - | service_tier + | service_tier | mid_conversation_system streaming : stream | nonstream (omit where n/a) - assertion : works | cost_logged + assertion : works | cost_logged | cache_hit label (not in id): model = haiku-4.5 | sonnet-4.6 | opus-4.7 | gpt-* e.g. llm.chat_completions.bedrock_converse.tool_use.stream.works llm.messages.anthropic.prompt_cache_1h.nonstream.cache_hit diff --git a/tests/e2e/coverage_registry/llm_conversational.yaml b/tests/e2e/coverage_registry/llm_conversational.yaml index 7360aacb916..be8a291c6fc 100644 --- a/tests/e2e/coverage_registry/llm_conversational.yaml +++ b/tests/e2e/coverage_registry/llm_conversational.yaml @@ -39,6 +39,8 @@ - {id: llm.messages.anthropic.vision.nonstream.works, module: llm, tier: P0, subject_endpoint: messages, route: anthropic, capability: vision, streaming: nonstream, assertions: [works], source: "model_prices json", rationale: "Vision via Messages API"} - {id: llm.messages.anthropic.prompt_cache_5m.nonstream.works, module: llm, tier: P0, subject_endpoint: messages, route: anthropic, capability: prompt_cache_5m, streaming: nonstream, assertions: [works], source: "model_prices json", rationale: "Prompt caching via Messages API"} - {id: llm.messages.anthropic.thinking.nonstream.works, module: llm, tier: P1, subject_endpoint: messages, route: anthropic, capability: thinking, streaming: nonstream, assertions: [works], source: "model_prices json", rationale: "Extended thinking via Messages API"} +- {id: llm.messages.bedrock_invoke.mid_conversation_system.nonstream.cache_hit, module: llm, tier: P0, subject_endpoint: messages, route: bedrock_invoke, capability: mid_conversation_system, streaming: nonstream, assertions: [works, cache_hit], source: "llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py", rationale: "Flagged Claude 4.8+/5 must keep mid-conversation system reminders in messages; hoisting mutates the system prefix and collapses the prompt cache (#32578/#32831/#32882)", fail_before_fix: proven} +- {id: llm.messages.bedrock_invoke.mid_conversation_system.nonstream.works, module: llm, tier: P0, subject_endpoint: messages, route: bedrock_invoke, capability: mid_conversation_system, streaming: nonstream, assertions: [works], source: "llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py", rationale: "Claude <= 4.7 rejects role system inside messages; unflagged models must hoist reminders into top-level system or every Claude Code session 400s (#32831)", fail_before_fix: proven} - {id: llm.responses.openai.basic.nonstream.works, module: llm, tier: P0, subject_endpoint: responses, route: openai, capability: basic, streaming: nonstream, assertions: [works], source: "response_api_endpoints/endpoints.py:26", rationale: "Core endpoint; OpenAI Responses native"} - {id: llm.responses.openai.basic.stream.works, module: llm, tier: P0, subject_endpoint: responses, route: openai, capability: basic, streaming: stream, assertions: [works], source: "response_api_endpoints/endpoints.py:26", rationale: "Streaming via /v1/responses"} - {id: llm.responses.openai.basic.nonstream.cost_logged, module: llm, tier: P0, subject_endpoint: responses, route: openai, capability: basic, streaming: nonstream, assertions: [works, cost_logged], source: "response_api_endpoints/endpoints.py:26", rationale: "Cost logged on responses"} diff --git a/tests/e2e/coverage_registry/schema.py b/tests/e2e/coverage_registry/schema.py index 21be0acb18f..7482088d93f 100644 --- a/tests/e2e/coverage_registry/schema.py +++ b/tests/e2e/coverage_registry/schema.py @@ -45,6 +45,7 @@ LlmRoute = Literal[ "azure_foundry", "azure_openai", "bedrock_converse", + "bedrock_invoke", "cohere", "openai", "together_ai", @@ -53,6 +54,7 @@ LlmRoute = Literal[ LlmCapability = Literal[ "basic", + "mid_conversation_system", "prompt_cache_5m", "service_tier", "structured_output", diff --git a/tests/e2e/llm_translation/endpoints_client.py b/tests/e2e/llm_translation/endpoints_client.py index 0ab87472748..c3816928564 100644 --- a/tests/e2e/llm_translation/endpoints_client.py +++ b/tests/e2e/llm_translation/endpoints_client.py @@ -30,6 +30,28 @@ class MessagesRequest(BaseModel): messages: list[ChatMessage] +class CacheControl(BaseModel): + type: str = "ephemeral" + + +class TextBlock(BaseModel): + type: str = "text" + text: str + cache_control: CacheControl | None = None + + +class RichMessage(BaseModel): + role: str + content: list[TextBlock] + + +class RichMessagesRequest(BaseModel): + model: str + max_tokens: int = 64 + system: list[TextBlock] + messages: list[RichMessage] + + class EmbeddingsRequest(BaseModel): model: str input: str @@ -83,11 +105,19 @@ class AnthropicContentBlock(BaseModel): text: str | None = None +class MessagesUsage(BaseModel): + input_tokens: int = 0 + output_tokens: int = 0 + cache_creation_input_tokens: int = 0 + cache_read_input_tokens: int = 0 + + class MessagesResult(BaseModel): id: str | None = None role: str | None = None model: str | None = None content: list[AnthropicContentBlock] = [] + usage: MessagesUsage = MessagesUsage() @property def text(self) -> str: diff --git a/tests/e2e/llm_translation/test_cache_control.py b/tests/e2e/llm_translation/test_cache_control.py index 51ae6ccbc1d..ba701f4869c 100644 --- a/tests/e2e/llm_translation/test_cache_control.py +++ b/tests/e2e/llm_translation/test_cache_control.py @@ -15,7 +15,7 @@ service_tier lives in test_provider_features_e2e.py. The provider-native cache_control request shape is not expressible with the shared ``ChatBody`` (whose content is a plain string), so the cacheable body is -modelled locally with typed content blocks. +built from the typed content blocks shared in ``endpoints_client.py``. """ from __future__ import annotations @@ -27,6 +27,7 @@ from pydantic import BaseModel from e2e_config import unique_marker from e2e_http import Result, unwrap +from endpoints_client import CacheControl, RichMessage, TextBlock from lifecycle import ResourceManager from models import ChatResponse, LiteLLMParamsBody, Usage from passthrough_client import PassthroughClient @@ -38,21 +39,6 @@ BEDROCK_MODEL = "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0" VERTEX_MODEL = "vertex_ai/gemini-2.5-flash" -class CacheControl(BaseModel): - type: str = "ephemeral" - - -class TextBlock(BaseModel): - type: str = "text" - text: str - cache_control: CacheControl | None = None - - -class RichMessage(BaseModel): - role: str - content: list[TextBlock] - - class CacheChatBody(BaseModel): model: str messages: list[RichMessage] diff --git a/tests/e2e/llm_translation/test_messages_mid_conversation_system_e2e.py b/tests/e2e/llm_translation/test_messages_mid_conversation_system_e2e.py new file mode 100644 index 00000000000..1025d603fca --- /dev/null +++ b/tests/e2e/llm_translation/test_messages_mid_conversation_system_e2e.py @@ -0,0 +1,220 @@ +"""Live e2e: mid-conversation ``role: "system"`` handling on the Bedrock Invoke +/v1/messages path is model-aware (PRs #32578, #32831, #32882). + +Models flagged ``supports_mid_conversation_system`` in the cost map (Claude 4.8+ +and the 5 family) must keep a mid-conversation system reminder in place inside +``messages`` so the top-level ``system`` prefix stays byte-identical and the +prompt cache written on turn one is read back in full on turn two. Models +without the flag (Claude 4.7 and older) reject the role inside ``messages`` +outright, so the proxy must hoist the reminder into the top-level ``system`` +field and the call must still return a completion instead of a provider 400. + +The conversation shape mirrors what Claude Code sends mid-session: a cached +system prompt, a user turn carrying its own ``cache_control`` breakpoint, a +``role: "system"`` reminder, an assistant turn, and a fresh user turn. The +message-turn breakpoint is what makes the cache assertion able to fail: a cache +entry whose prefix spans ``system`` plus message turns is invalidated when the +reminder is hoisted (the ``system`` field mutates and a turn disappears from +``messages``), while an entry ending at the system block itself would survive +the hoist and mask the regression. +""" + +from __future__ import annotations + +import time + +import pytest +from pydantic import BaseModel + +from e2e_config import unique_marker +from e2e_http import Result, unwrap +from endpoints_client import ( + CacheControl, + EndpointsClient, + MessagesResult, + RichMessage, + RichMessagesRequest, + TextBlock, +) +from lifecycle import ResourceManager +from models import LiteLLMParamsBody + +pytestmark = pytest.mark.e2e + +FLAGGED_INVOKE_MODEL = "bedrock/invoke/us.anthropic.claude-sonnet-5" +UNFLAGGED_INVOKE_MODEL = "bedrock/invoke/us.anthropic.claude-haiku-4-5-20251001-v1:0" +AWS_REGION = "us-east-1" +CACHE_PRIMING_DEADLINE_SECONDS = 60.0 +CACHE_PRIMING_INTERVAL_SECONDS = 3.0 + + +def _cacheable_system_block(marker: str) -> TextBlock: + """A system prompt comfortably above Sonnet's 1024-token minimum cacheable + size, unique per run so no other run's cache entry can satisfy the read.""" + text = " ".join( + f"Reference paragraph {index} for run {marker}." for index in range(300) + ) + return TextBlock(text=text, cache_control=CacheControl()) + + +def _user_turn(text: str, *, cached: bool = False) -> RichMessage: + block = TextBlock(text=text, cache_control=CacheControl() if cached else None) + return RichMessage(role="user", content=[block]) + + +def _system_reminder_turn() -> RichMessage: + return RichMessage( + role="system", + content=[ + TextBlock( + text="Answer with exactly one word." + ) + ], + ) + + +def _post_messages( + client: EndpointsClient, key: str, body: RichMessagesRequest +) -> Result[MessagesResult]: + return client.gateway.transport.post( + "/v1/messages", + headers=client.gateway.transport.bearer(key), + json=body, + response_type=MessagesResult, + ) + + +def _register_invoke_deployment( + client: EndpointsClient, resources: ResourceManager, bedrock_model: str +) -> str: + model = f"e2e-midsys-{unique_marker()}" + model_id = client.create_model( + model, LiteLLMParamsBody(model=bedrock_model, aws_region_name=AWS_REGION) + ) + resources.defer(lambda: client.delete_model(model_id)) + return model + + +def _first_turn_user_text(marker: str) -> str: + """A first user turn heavy enough (hundreds of tokens) that losing its cache + entry is unambiguous in the usage numbers, unique per attempt so priming + retries never depend on the proxy's response cache behavior.""" + notes = " ".join(f"Session note {index} for attempt {marker}." for index in range(100)) + return f"Reply with one word.\n{notes}" + + +class PrimedCache(BaseModel): + first_user_text: str + prefix_read_tokens: int + first_turn_creation_tokens: int + + @property + def full_prefix_tokens(self) -> int: + return self.prefix_read_tokens + self.first_turn_creation_tokens + + +def _prime_prompt_cache( + client: EndpointsClient, key: str, model: str, system_block: TextBlock +) -> PrimedCache: + """Send first-turn calls (fresh cache-marked user turn each attempt, + identical system prefix) until one both reads the system prefix back from + cache and writes its own user-turn chunk, proving the cache is live in both + directions. Only the pre-reminder turn is ever retried here, so retries can + never warm a mutated-prefix cache entry and mask the regression the second + turn asserts on.""" + deadline = time.monotonic() + CACHE_PRIMING_DEADLINE_SECONDS + while True: + user_text = _first_turn_user_text(unique_marker()) + body = RichMessagesRequest( + model=model, + system=[system_block], + messages=[_user_turn(user_text, cached=True)], + ) + usage = unwrap(_post_messages(client, key, body)).usage + if usage.cache_read_input_tokens > 0 and usage.cache_creation_input_tokens > 0: + return PrimedCache( + first_user_text=user_text, + prefix_read_tokens=usage.cache_read_input_tokens, + first_turn_creation_tokens=usage.cache_creation_input_tokens, + ) + if time.monotonic() >= deadline: + pytest.fail( + f"{model}: prompt cache never became readable within " + f"{CACHE_PRIMING_DEADLINE_SECONDS}s (last usage: {usage})" + ) + time.sleep(CACHE_PRIMING_INTERVAL_SECONDS) + + +class TestBedrockInvokeMidConversationSystem: + @pytest.mark.covers( + "llm.messages.bedrock_invoke.mid_conversation_system.nonstream.cache_hit", + exercised_on=[], + ) + def test_flagged_model_keeps_prompt_cache_across_system_reminder( + self, endpoints_client: EndpointsClient, resources: ResourceManager + ) -> None: + model = _register_invoke_deployment( + endpoints_client, resources, FLAGGED_INVOKE_MODEL + ) + key = resources.key(models=[model]) + system_block = _cacheable_system_block(unique_marker()) + + primed = _prime_prompt_cache(endpoints_client, key, model, system_block) + + reminder_turn_body = RichMessagesRequest( + model=model, + system=[system_block], + messages=[ + _user_turn(primed.first_user_text, cached=True), + _system_reminder_turn(), + RichMessage(role="assistant", content=[TextBlock(text="OK.")]), + _user_turn("Reply with one word again.", cached=True), + ], + ) + second = unwrap(_post_messages(endpoints_client, key, reminder_turn_body)) + + assert second.text.strip(), ( + f"{model}: reminder turn returned no completion text" + ) + assert second.usage.cache_read_input_tokens >= primed.full_prefix_tokens, ( + f"{model}: turn with a mid-conversation system reminder read " + f"{second.usage.cache_read_input_tokens} cached tokens, expected at " + f"least the {primed.full_prefix_tokens} cached on turn one " + f"({primed.prefix_read_tokens} system prefix + " + f"{primed.first_turn_creation_tokens} first user turn); the reminder " + f"was hoisted into the top-level system field, which mutates the " + f"cached prefix and re-bills the conversation at cache-write pricing" + ) + + @pytest.mark.covers( + "llm.messages.bedrock_invoke.mid_conversation_system.nonstream.works", + exercised_on=[], + ) + def test_unflagged_model_hoists_system_reminder_and_succeeds( + self, endpoints_client: EndpointsClient, resources: ResourceManager + ) -> None: + model = _register_invoke_deployment( + endpoints_client, resources, UNFLAGGED_INVOKE_MODEL + ) + key = resources.key(models=[model]) + + body = RichMessagesRequest( + model=model, + system=[TextBlock(text="You are terse.")], + messages=[ + _user_turn(f"Say hi. Run {unique_marker()}."), + _system_reminder_turn(), + RichMessage(role="assistant", content=[TextBlock(text="Hi.")]), + _user_turn("Say bye."), + ], + ) + completion = unwrap(_post_messages(endpoints_client, key, body)) + + assert completion.role == "assistant", ( + f"{model}: unexpected role {completion.role!r}" + ) + assert completion.text.strip(), ( + f"{model}: conversation with a mid-conversation system reminder " + f"returned no text; the reminder was forwarded in place to a model " + f"that rejects role 'system' inside messages instead of being hoisted" + ) diff --git a/tests/test_litellm/integrations/test_s3_v2.py b/tests/test_litellm/integrations/test_s3_v2.py index 246b378c982..f0a33f2ebfc 100644 --- a/tests/test_litellm/integrations/test_s3_v2.py +++ b/tests/test_litellm/integrations/test_s3_v2.py @@ -1125,6 +1125,47 @@ async def test_combined_prefix_reflects_in_s3_object_key(): assert "myteam/apikey/" in key, f"Expected both prefixes in key: {key}" +def test_s3_object_key_sanitizes_slashes_in_file_name(): + """Response ids containing slashes (e.g. bedrock batch job ARNs) must not + create nested S3 folders; only path/prefix/date slashes are separators.""" + from litellm.integrations.s3 import get_s3_object_key + + start_time = datetime(2026, 2, 11, 0, 35, 18, 391582) + file_name = "time-00-35-18-391582_arn:aws:bedrock:us-east-1:123456789012:model-invocation-job/gl18r6skk9yy" + + key = get_s3_object_key( + s3_path="LiteLLMAPPLogs", + prefix="myteam/", + start_time=start_time, + s3_file_name=file_name, + ) + + assert key == ( + "LiteLLMAPPLogs/myteam/2026-02-11/" + "time-00-35-18-391582_arn:aws:bedrock:us-east-1:123456789012:model-invocation-job_gl18r6skk9yy.json" + ) + + +def test_create_s3_batch_logging_element_flat_key_for_arn_response_id(): + """End-to-end through the s3_v2 element builder: an ARN response id must + yield a flat file directly under the date segment.""" + logger = S3Logger(s3_use_team_prefix=False, s3_use_key_prefix=False) + payload = StandardLoggingPayload( + id="arn:aws:bedrock:us-east-1:123456789012:model-invocation-job/gl18r6skk9yy", + metadata={}, + messages=[], + ) + + start_time = datetime(2026, 2, 11, 0, 35, 18, 391582) + result = logger.create_s3_batch_logging_element(start_time, payload) + + assert result is not None + date_segment = "2026-02-11/" + file_segment = result.s3_object_key.split(date_segment, 1)[1] + assert "/" not in file_segment, f"Expected flat file under date segment, got: {result.s3_object_key}" + assert file_segment.endswith("model-invocation-job_gl18r6skk9yy.json") + + # -------------------------------------------------------------- # params_source / s3_callback_params_override (audit-log decoupling) # -------------------------------------------------------------- diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index ca23e61352e..b156faf3ea6 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -270,6 +270,105 @@ def test_image_tokens_fallback_to_base_cost(): assert round(completion_cost, 12) == round(expected_completion_cost, 12) +def test_video_output_tokens_gemini_omni_flash_preview(): + """Video output tokens are billed at output_cost_per_video_token, not the text rate and not zero.""" + model = "gemini-omni-flash-preview" + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + + text_tokens = 100 + video_tokens = 46336 + usage = Usage( + completion_tokens=text_tokens + video_tokens, + prompt_tokens=20, + total_tokens=20 + text_tokens + video_tokens, + completion_tokens_details=CompletionTokensDetailsWrapper( + text_tokens=text_tokens, + video_tokens=video_tokens, + ), + prompt_tokens_details=PromptTokensDetailsWrapper(text_tokens=20), + ) + model_cost_map = litellm.model_cost[f"gemini/{model}"] + assert model_cost_map["input_cost_per_token"] == 1.5e-06 + assert model_cost_map["output_cost_per_token"] == 9e-06 + assert model_cost_map["output_cost_per_video_token"] == 1.75e-05 + + prompt_cost, completion_cost = generic_cost_per_token( + model=model, + usage=usage, + custom_llm_provider="gemini", + ) + + assert round(prompt_cost, 10) == round( + model_cost_map["input_cost_per_token"] * usage.prompt_tokens, + 10, + ) + assert round(completion_cost, 10) == round( + (model_cost_map["output_cost_per_token"] * text_tokens) + + (model_cost_map["output_cost_per_video_token"] * video_tokens), + 10, + ) + + +def test_video_input_tokens_gemini_omni_flash_preview(): + """Video input tokens are billed at the standard input rate instead of being dropped.""" + model = "gemini-omni-flash-preview" + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + + usage = Usage( + completion_tokens=10, + prompt_tokens=10050, + total_tokens=10060, + completion_tokens_details=CompletionTokensDetailsWrapper(text_tokens=10), + prompt_tokens_details=PromptTokensDetailsWrapper(text_tokens=50, video_tokens=10000), + ) + model_cost_map = litellm.model_cost[f"gemini/{model}"] + + prompt_cost, _ = generic_cost_per_token( + model=model, + usage=usage, + custom_llm_provider="gemini", + ) + + assert round(prompt_cost, 10) == round( + model_cost_map["input_cost_per_token"] * usage.prompt_tokens, + 10, + ) + + +def test_video_tokens_fallback_to_base_cost(): + """Video output tokens fall back to the base output rate when output_cost_per_video_token is not set.""" + from unittest.mock import patch + + mock_model_info = { + "input_cost_per_token": 1e-6, + "output_cost_per_token": 2e-6, + } + + usage = Usage( + completion_tokens=1720, + prompt_tokens=14, + total_tokens=1734, + completion_tokens_details=CompletionTokensDetailsWrapper( + text_tokens=600, + video_tokens=1120, + ), + prompt_tokens_details=PromptTokensDetailsWrapper(text_tokens=14), + ) + + with patch( + "litellm.litellm_core_utils.llm_cost_calc.utils.get_model_info", + return_value=mock_model_info, + ): + prompt_cost, completion_cost = generic_cost_per_token( + model="test-model", usage=usage, custom_llm_provider="gemini" + ) + + assert round(prompt_cost, 12) == round(14 * 1e-6, 12) + assert round(completion_cost, 12) == round((600 + 1120) * 2e-6, 12) + + def test_generic_cost_per_token_above_200k_tokens(): # gemini-2.5-pro-exp-03-25 was removed; gemini-2.5-pro has same above-200k pricing model = "gemini-2.5-pro" @@ -1086,6 +1185,7 @@ def test_cache_writing_cost_with_zero_creation_tokens_and_ephemeral_details(): "text_tokens": 0, "audio_tokens": 0, "image_tokens": 0, + "video_tokens": 0, "character_count": 0, "image_count": 0, "video_length_seconds": 0.0, diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 51f13affa3f..40f9f4e7910 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -880,6 +880,12 @@ def test_vertex_ai_usage_metadata_with_image_tokens_in_prompt(): ) +def test_map_response_modalities_video(): + """The video modality maps to VIDEO instead of MODALITY_UNSPECIFIED, which Gemini rejects.""" + v = VertexGeminiConfig() + assert v.map_response_modalities(["text", "video"]) == ["TEXT", "VIDEO"] + + def test_vertex_ai_usage_metadata_accumulates_duplicate_modalities(): """Ensure _calculate_usage accumulates repeated modality entries.""" v = VertexGeminiConfig() diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index 90f46152837..a0248963cf1 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -3992,6 +3992,80 @@ async def test_non_admin_cli_session_token_reaches_production_auth_path(monkeypa assert result.is_session_token is True +@pytest.mark.asyncio +async def test_cli_session_token_authenticates_when_jwt_auth_enabled_without_license(monkeypatch): + """A lite login token is an encrypted (non-JWT) session blob. With + enable_jwt_auth on and no enterprise license (premium_user False), the JWT + premium gate used to fire for every request before the token was decoded, so + the CLI token 401'd with 'JWT Auth is an enterprise only feature' and was + never decrypted. The gate must apply only to actual JWTs; a non-JWT session + token has to keep authenticating on its own path regardless of license.""" + monkeypatch.delenv("EXPERIMENTAL_UI_LOGIN", raising=False) + cli_token = _mint_cli_session_token(monkeypatch) + + jwt_handler = MagicMock() + jwt_handler.is_jwt = JWTHandler.is_jwt + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth() + + mock_request = MagicMock() + mock_request.url.path = "/v1/messages" + mock_request.method = "POST" + mock_request.headers = {"authorization": f"Bearer {cli_token}"} + mock_request.query_params = {} + + with ( + patch("litellm.proxy.proxy_server.general_settings", {"enable_jwt_auth": True}), + patch("litellm.proxy.proxy_server.premium_user", False), + patch("litellm.proxy.proxy_server.jwt_handler", jwt_handler), + patch("litellm.proxy.proxy_server.master_key", "sk-master"), + patch("litellm.proxy.proxy_server.prisma_client", None), + ): + result = await user_api_key_auth( + request=mock_request, + api_key=f"Bearer {cli_token}", + ) + + assert result.user_id == "cli-admin" + assert result.team_id == "cli-team" + assert result.token is not None and result.token.startswith("cli-session-") + + +@pytest.mark.asyncio +async def test_real_jwt_still_requires_license_when_jwt_auth_enabled(monkeypatch): + """Guard for the reorder above: the enterprise gate must still reject an + actual JWT when there is no license. Moving the premium check inside the + is_jwt branch must not open JWT auth to non-premium deployments.""" + monkeypatch.delenv("EXPERIMENTAL_UI_LOGIN", raising=False) + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-salt-cli-test") + + jwt_token = "eyJhbGciOiJSUzI1NiJ9.eyJzdWIiOiJ1c2VyMSJ9.sig" + jwt_handler = MagicMock() + jwt_handler.is_jwt = JWTHandler.is_jwt + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth() + + mock_request = MagicMock() + mock_request.url.path = "/v1/messages" + mock_request.method = "POST" + mock_request.headers = {"authorization": f"Bearer {jwt_token}"} + mock_request.query_params = {} + + with ( + patch("litellm.proxy.proxy_server.general_settings", {"enable_jwt_auth": True}), + patch("litellm.proxy.proxy_server.premium_user", False), + patch("litellm.proxy.proxy_server.jwt_handler", jwt_handler), + patch("litellm.proxy.proxy_server.master_key", "sk-master"), + patch("litellm.proxy.proxy_server.prisma_client", None), + ): + with pytest.raises(Exception) as exc_info: + await user_api_key_auth( + request=mock_request, + api_key=f"Bearer {jwt_token}", + ) + + message = str(getattr(exc_info.value, "message", exc_info.value)) + assert "enterprise only feature" in message + + @pytest.mark.asyncio async def test_auth_path_caches_team_object_under_canonical_team_id_key(): """Regression for LIT-4000: the auth builder must cache the team object under diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py index 791fdd4077c..5be0d43c250 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py @@ -1150,6 +1150,23 @@ class TestGenericGuardrailAPIStreamingConfig: assert GenericGuardrailAPI.get_config_model() is GenericGuardrailAPIConfigModel + def test_streaming_transform_mode_defaults_block_only(self): + guardrail = GenericGuardrailAPI( + api_base="https://api.test.guardrail.com", + guardrail_name="test-generic-guardrail", + event_hook="post_call", + ) + assert guardrail.streaming_transform_mode == "block_only" + + def test_streaming_transform_mode_override(self): + guardrail = GenericGuardrailAPI( + api_base="https://api.test.guardrail.com", + guardrail_name="test-generic-guardrail", + event_hook="post_call", + streaming_transform_mode="incremental_diff", + ) + assert guardrail.streaming_transform_mode == "incremental_diff" + def test_initialize_guardrail_forwards_streaming_flags(self): from litellm.proxy.guardrails.guardrail_hooks.generic_guardrail_api import ( initialize_guardrail, @@ -1306,6 +1323,86 @@ class TestGenericGuardrailAPIStreamingConfig: assert guardrail.streaming_sampling_rate == 2 +class TestGenericGuardrailAPIResponseParsing: + """GenericGuardrailAPIResponse.from_dict handling of the streaming holdback field.""" + + def test_from_dict_parses_stream_holdback_chars(self): + from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import ( + GenericGuardrailAPIResponse, + ) + + response = GenericGuardrailAPIResponse.from_dict( + { + "action": "GUARDRAIL_INTERVENED", + "texts": ["Alice went to Berlin"], + "stream_holdback_chars": [5], + } + ) + + assert response.action == "GUARDRAIL_INTERVENED" + assert response.texts == ["Alice went to Berlin"] + assert response.stream_holdback_chars == [5] + + def test_from_dict_coerces_holdback_values_to_int(self): + from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import ( + GenericGuardrailAPIResponse, + ) + + response = GenericGuardrailAPIResponse.from_dict( + {"action": "GUARDRAIL_INTERVENED", "texts": ["x", "y"], "stream_holdback_chars": ["3", 0]} + ) + + assert response.stream_holdback_chars == [3, 0] + + def test_from_dict_holdback_absent_is_none(self): + from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import ( + GenericGuardrailAPIResponse, + ) + + response = GenericGuardrailAPIResponse.from_dict({"action": "NONE", "texts": ["hi"]}) + + assert response.stream_holdback_chars is None + + def test_from_dict_malformed_holdback_degrades_to_zero(self): + """A null/non-numeric/negative holdback element must not raise; it degrades + to 0 (no holdback) so a bad guardrail response can't abort the stream.""" + from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import ( + GenericGuardrailAPIResponse, + ) + + response = GenericGuardrailAPIResponse.from_dict( + { + "action": "GUARDRAIL_INTERVENED", + "texts": ["a", "b", "c", "d"], + "stream_holdback_chars": ["3", None, "bad", -2], + } + ) + + assert response.stream_holdback_chars == [3, 0, 0, 0] + + @pytest.mark.asyncio + async def test_apply_guardrail_flows_holdback_back_to_inputs(self, generic_guardrail): + """A GUARDRAIL_INTERVENED response with stream_holdback_chars is surfaced on + the returned inputs so the streaming framework can apply it.""" + mock_response = MagicMock() + mock_response.json.return_value = { + "action": "GUARDRAIL_INTERVENED", + "texts": ["Alice went to Berlin"], + "stream_holdback_chars": [5], + } + mock_response.raise_for_status = MagicMock() + + with patch.object(generic_guardrail.async_handler, "post", return_value=mock_response): + result = await generic_guardrail.apply_guardrail( + inputs={"texts": ["Zorg went to Xanadu"]}, + request_data={}, + input_type="response", + ) + + assert result["texts"] == ["Alice went to Berlin"] + assert result["stream_holdback_chars"] == [5] + + class TestGenericGuardrailAPIStreamingViaUnified: """Streaming output checks routed through UnifiedLLMGuardrails.""" diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py index 6e027fa4941..e84e9b74201 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py @@ -648,3 +648,845 @@ class TestUnifiedLLMGuardrails: # Response returned with pages intact assert result.pages[0].markdown == "Some text" + + +class _StreamingTextGuardrail(CustomGuardrail): + """Guardrail whose apply_guardrail rewrites (uppercases) response text. + + Optionally schedules a per-response-call ``stream_holdback_chars`` (indexed + like ``texts``) and can force the mutated text shorter than the input to + exercise the streaming underflow guard. + """ + + def __init__(self, *, holdback_schedule=None, shrink_to=None, shrink_after=0, sampling_rate=1): + super().__init__(guardrail_name="streaming-text-guardrail") + self.streaming_transform_mode = "incremental_diff" + self.streaming_sampling_rate = sampling_rate + self.streaming_end_of_stream_only = False + self.guardrail_config = {} + self._holdback_schedule = list(holdback_schedule or []) + self._shrink_to = shrink_to + self._shrink_after = shrink_after + self.response_calls = 0 + self.received_texts = [] + self.received_tool_calls = [] + + def should_run_guardrail(self, data, event_type): # type: ignore[override] + return True + + async def apply_guardrail(self, inputs, request_data, input_type, **kwargs): + texts = inputs.get("texts", []) + if input_type != "response": + return {"texts": [t.upper() for t in texts]} + + if inputs.get("tool_calls"): + self.received_tool_calls.append(inputs.get("tool_calls")) + self.received_texts.append(list(texts)) + idx = self.response_calls + self.response_calls += 1 + if self._shrink_to is not None and idx >= self._shrink_after: + transformed = [self._shrink_to for _ in texts] + else: + transformed = [t.upper() for t in texts] + result = {"texts": transformed} + if idx < len(self._holdback_schedule): + result["stream_holdback_chars"] = [self._holdback_schedule[idx]] * len(texts) + return result + + +def _stream_chunk(content, finish_reason=None, index=0): + return ModelResponseStream( + choices=[ + StreamingChoices( + index=index, + delta=Delta(content=content, role="assistant"), + finish_reason=finish_reason, + ) + ], + ) + + +async def _drive_stream(handler, guardrail, chunks, request_route="/v1/chat/completions"): + async def _mock_stream(): + for chunk in chunks: + yield chunk + + user_api_key_dict = UserAPIKeyAuth(api_key="test-key", request_route=request_route) + request_data = {"guardrail_to_apply": guardrail, "model": "gpt-4"} + out = [] + async for item in handler.async_post_call_streaming_iterator_hook( + user_api_key_dict=user_api_key_dict, + response=_mock_stream(), + request_data=request_data, + ): + out.append(item) + return out + + +def _delta_text(item): + if not getattr(item, "choices", None): + return "" + return item.choices[0].delta.content or "" + + +class TestStreamingTransform: + """Streaming text-transformation (incremental_diff) path on the OpenAI chat + completions streaming surface.""" + + @pytest.fixture(autouse=True) + def _use_openai_handler_mapping(self): + unified_module.endpoint_guardrail_translation_mappings = { + CallTypes.acompletion: OpenAIChatCompletionsHandler, + } + yield + unified_module.endpoint_guardrail_translation_mappings = None + + @pytest.mark.asyncio + async def test_block_only_drops_text_rewrites(self): + """Default block_only: the guardrail's uppercasing never reaches the + client; the original lowercase chunks are streamed verbatim.""" + guardrail = _StreamingTextGuardrail() + guardrail.streaming_transform_mode = "block_only" + + chunks = [ + _stream_chunk("hello "), + _stream_chunk("world"), + _stream_chunk("", finish_reason="stop"), + ] + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks) + streamed = "".join(_delta_text(i) for i in out) + + assert streamed == "hello world" + assert streamed != streamed.upper() + + @pytest.mark.asyncio + async def test_incremental_diff_emits_uppercased_deltas(self): + """incremental_diff: the client receives uppercased deltas whose + concatenation equals uppercase(full).""" + guardrail = _StreamingTextGuardrail() + + full = "hello world this is streaming" + words = ["hello ", "world ", "this ", "is ", "streaming"] + chunks = [_stream_chunk(w) for w in words] + chunks.append(_stream_chunk("", finish_reason="stop")) + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks) + streamed = "".join(_delta_text(i) for i in out) + + assert streamed == full.upper() + # No raw lowercase content leaked onto the wire. + assert "hello" not in streamed + + @pytest.mark.asyncio + async def test_incremental_diff_holdback_boundary(self): + """Holdback=5 on the first sample withholds the trailing chars until the + next round; the final concatenation matches with no loss or duplication.""" + guardrail = _StreamingTextGuardrail(holdback_schedule=[5, 0]) + + # No finish_reason: every sample uses the combined-text branch so the + # scheduled holdback is applied on the first round. + chunks = [_stream_chunk("abcdef"), _stream_chunk("ghij")] + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks) + deltas = [_delta_text(i) for i in out] + streamed = "".join(deltas) + + # First sample: "ABCDEF" with holdback 5 -> only "A" is emitted. + assert deltas[0] == "A" + assert streamed == "ABCDEFGHIJ" + + @pytest.mark.asyncio + async def test_incremental_diff_underflow_raises(self): + """A transform shorter than what was already streamed cannot retract + bytes: it raises HTTPException(stream_transform_underflow).""" + # First sample emits "ABCDEF" (6 chars); second sample shrinks to 3. + guardrail = _StreamingTextGuardrail(shrink_to="ABC", shrink_after=1) + + chunks = [_stream_chunk("abcdef"), _stream_chunk("ghij")] + + with pytest.raises(unified_module.HTTPException) as exc_info: + await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks) + + assert exc_info.value.status_code == 400 + assert exc_info.value.detail["error"] == "stream_transform_underflow" + + @pytest.mark.asyncio + async def test_incremental_diff_final_chunk_preserves_finish_reason(self): + """The final synthetic chunk carries the finish_reason of the last raw + chunk.""" + guardrail = _StreamingTextGuardrail() + + chunks = [ + _stream_chunk("hello "), + _stream_chunk("world"), + _stream_chunk("", finish_reason="stop"), + ] + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks) + + assert out, "expected at least one synthetic chunk" + assert out[-1].choices[0].finish_reason == "stop" + assert "".join(_delta_text(i) for i in out) == "HELLO WORLD" + + @pytest.mark.asyncio + async def test_end_of_stream_only_emits_single_final_chunk(self): + """incremental_diff + streaming_end_of_stream_only: a single post-stream + synthetic chunk carries the whole guardrailed text and the finish_reason.""" + guardrail = _StreamingTextGuardrail() + guardrail.streaming_end_of_stream_only = True + + chunks = [ + _stream_chunk("hello "), + _stream_chunk("world"), + _stream_chunk("", finish_reason="stop"), + ] + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks) + + non_empty = [i for i in out if _delta_text(i)] + assert len(non_empty) == 1 + assert _delta_text(non_empty[0]) == "HELLO WORLD" + assert out[-1].choices[0].finish_reason == "stop" + + @pytest.mark.asyncio + async def test_unsupported_route_falls_back_to_block_only(self): + """A route that does not resolve to the OpenAI chat handler falls back to + block_only rather than transforming.""" + guardrail = _StreamingTextGuardrail() + + chunks = [_stream_chunk("hello "), _stream_chunk("world", finish_reason="stop")] + + # request_route=None => no resolvable call type => block_only fallback. + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks, request_route=None) + streamed = "".join(_delta_text(i) for i in out) + + assert streamed == "hello world" + + @pytest.mark.asyncio + async def test_emit_streaming_http_error_a2a_yields_jsonrpc_chunk(self): + """The shared streaming error helper emits an in-stream JSON-RPC error for + A2A call types instead of raising.""" + import json + + handler = UnifiedLLMGuardrails() + exc = unified_module.HTTPException( + status_code=400, + detail={"error": "stream_transform_underflow", "message": "boom"}, + ) + + emitted = [] + async for item in handler._emit_streaming_http_error( + exc, + call_type=CallTypes.asend_message.value, + responses_so_far=[{"id": "req-1"}], + request_data={}, + ): + emitted.append(item) + + assert len(emitted) == 1 + payload = json.loads(emitted[0]) + assert payload["error"]["message"] == "stream_transform_underflow" + assert payload["id"] == "req-1" + + def test_final_chunk_preserves_per_choice_finish_reason(self): + """The final flush must carry each choice's own finish_reason, not + choices[0]'s, for n > 1 (e.g. "stop" vs "length").""" + reference_chunk = ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(content="a"))]) + + synthetic = UnifiedLLMGuardrails()._build_transform_chunk( + reference_chunk=reference_chunk, + mutated_text_per_choice={0: "A", 1: "B"}, + emitted_text_per_choice={}, + holdback_per_choice={}, + finish_reason_per_choice={0: "stop", 1: "length"}, + is_final=True, + ) + + by_index = {c.index: c for c in synthetic.choices} + assert by_index[0].finish_reason == "stop" + assert by_index[1].finish_reason == "length" + assert by_index[0].delta.content == "A" + assert by_index[1].delta.content == "B" + + def test_synthetic_chunk_drops_raw_tool_calls(self): + """v1 does not transform streamed tool calls; the synthetic chunk must not + pass raw upstream tool_calls through (they would bypass the guardrail).""" + reference_chunk = ModelResponseStream( + choices=[ + StreamingChoices( + index=0, + delta=Delta( + content="hi", + tool_calls=[ + { + "index": 0, + "id": "call_1", + "type": "function", + "function": {"name": "leak", "arguments": '{"ssn": "123-45-6789"}'}, + } + ], + ), + finish_reason=None, + ), + ], + ) + + synthetic = UnifiedLLMGuardrails()._build_transform_chunk( + reference_chunk=reference_chunk, + mutated_text_per_choice={0: "HI"}, + emitted_text_per_choice={}, + holdback_per_choice={}, + finish_reason_per_choice={}, + is_final=False, + ) + + assert synthetic.choices[0].delta.tool_calls is None + assert synthetic.choices[0].delta.content == "HI" + + def test_rewriting_already_emitted_prefix_raises(self): + """If a later transform rewrites bytes already streamed (not a forward + extension), the framework fails closed rather than leaking the original.""" + reference_chunk = ModelResponseStream( + choices=[StreamingChoices(index=0, delta=Delta(content="x"), finish_reason=None)], + ) + + with pytest.raises(unified_module.HTTPException) as exc_info: + UnifiedLLMGuardrails()._build_transform_chunk( + reference_chunk=reference_chunk, + # already streamed "My SSN is 123"; the guardrail now wants to + # redact those already-sent chars -> not a forward extension. + mutated_text_per_choice={0: "My SSN is [REDACTED]"}, + emitted_text_per_choice={0: "My SSN is 123"}, + holdback_per_choice={}, + finish_reason_per_choice={}, + is_final=False, + ) + + assert exc_info.value.status_code == 400 + assert exc_info.value.detail["error"] == "stream_transform_underflow" + + @pytest.mark.asyncio + async def test_tool_call_chunks_pass_through_and_not_dropped(self): + """A tool-call chunk is passed through raw under incremental_diff (v1 does + not transform tool calls) rather than being withheld and dropped, and no + bogus empty-choices chunk is emitted for a tool-call-only turn.""" + guardrail = _StreamingTextGuardrail() + + tool_chunk = ModelResponseStream( + choices=[ + StreamingChoices( + index=0, + delta=Delta( + content=None, + tool_calls=[ + { + "index": 0, + "id": "call_1", + "type": "function", + "function": {"name": "get_weather", "arguments": "{}"}, + } + ], + ), + finish_reason="tool_calls", + ) + ], + ) + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, [tool_chunk]) + + assert len(out) == 1 + assert out[0].choices[0].delta.tool_calls + assert out[0].choices[0].finish_reason == "tool_calls" + # The tool call is delivered raw but still inspected by the guardrail at + # end of stream (it can block), matching block_only. + assert guardrail.received_tool_calls + + @pytest.mark.asyncio + async def test_per_choice_finish_reason_when_choices_finish_in_different_chunks(self): + """n>1: a choice finishing before the stream's last chunk keeps its own + finish_reason (it must not be lost because it is not on last_chunk).""" + guardrail = _StreamingTextGuardrail() + guardrail.streaming_end_of_stream_only = True # only the flush emits + + chunks = [ + _stream_chunk("aa", index=0), + _stream_chunk("bb", index=1), + _stream_chunk("", finish_reason="stop", index=0), + _stream_chunk("", finish_reason="length", index=1), + ] + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks) + + by_index = {} + for item in out: + for choice in item.choices: + if choice.finish_reason is not None: + by_index[choice.index] = choice.finish_reason + assert by_index == {0: "stop", 1: "length"} + + @pytest.mark.asyncio + async def test_short_guardrail_texts_withheld_not_leaked(self): + """If the guardrail returns fewer texts than sent (contract violation), + the unmatched choice is withheld (fail closed), not emitted raw.""" + + class _DropsSecondChoice(_StreamingTextGuardrail): + async def apply_guardrail(self, inputs, request_data, input_type, **kwargs): + texts = inputs.get("texts", []) + if input_type != "response": + return {"texts": [t.upper() for t in texts]} + # Return only the first choice's transformed text. + return {"texts": [texts[0].upper()] if texts else []} + + guardrail = _DropsSecondChoice() + guardrail.streaming_end_of_stream_only = True + + chunks = [ + _stream_chunk("secret-a", index=0), + _stream_chunk("secret-b", index=1), + _stream_chunk("", finish_reason="stop", index=0), + _stream_chunk("", finish_reason="stop", index=1), + ] + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks) + + streamed = "".join(_delta_text(i) for i in out) + assert "SECRET-A" in streamed + # Choice 1 had no guardrailed text returned: withheld, never leaked raw. + assert "secret-b" not in streamed + assert "SECRET-B" not in streamed + + @pytest.mark.asyncio + async def test_no_spurious_chunk_after_text_then_tool_call_finish(self): + """When a choice streams text and then finishes via a tool-call chunk, the + raw tool-call chunk carries the finish_reason and no spurious empty chunk + for that choice is emitted afterwards (protocol: no delta after finish).""" + guardrail = _StreamingTextGuardrail() + + tool_chunk = ModelResponseStream( + choices=[ + StreamingChoices( + index=0, + delta=Delta( + content=None, + tool_calls=[ + { + "index": 0, + "id": "call_1", + "type": "function", + "function": {"name": "get_weather", "arguments": "{}"}, + } + ], + ), + finish_reason="tool_calls", + ) + ], + ) + chunks = [_stream_chunk("let me check "), tool_chunk] + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks) + + # Exactly: one synthetic text delta, then the raw tool-call chunk. No + # trailing empty chunk for choice 0 after it already finished. + assert len(out) == 2 + assert _delta_text(out[0]) == "LET ME CHECK " + assert out[1].choices[0].delta.tool_calls + assert out[1].choices[0].finish_reason == "tool_calls" + + @pytest.mark.asyncio + async def test_tool_call_blocking_guardrail_is_enforced(self): + """A guardrail that blocks on tool calls must terminate the incremental_diff + stream: tool calls go through the block decision, not bypass it.""" + from litellm.exceptions import GuardrailRaisedException + + class _ToolCallBlocker(_StreamingTextGuardrail): + async def apply_guardrail(self, inputs, request_data, input_type, **kwargs): + if input_type == "response" and inputs.get("tool_calls"): + raise GuardrailRaisedException( + guardrail_name="tc-block", + message="blocked tool call", + should_wrap_with_default_message=False, + ) + return await super().apply_guardrail(inputs, request_data, input_type, **kwargs) + + tool_chunk = ModelResponseStream( + choices=[ + StreamingChoices( + index=0, + delta=Delta( + content=None, + tool_calls=[ + { + "index": 0, + "id": "call_1", + "type": "function", + "function": {"name": "exfiltrate", "arguments": "{}"}, + } + ], + ), + finish_reason="tool_calls", + ) + ], + ) + + with pytest.raises(GuardrailRaisedException): + await _drive_stream(UnifiedLLMGuardrails(), _ToolCallBlocker(), [tool_chunk]) + + @pytest.mark.asyncio + async def test_mixed_content_and_tool_call_chunk_does_not_leak_text(self): + """A chunk carrying BOTH delta.content and a tool call must not be yielded + raw: the text has to go through the transform, only tool-call fields pass + through raw (content stripped).""" + guardrail = _StreamingTextGuardrail() + + mixed = ModelResponseStream( + choices=[ + StreamingChoices( + index=0, + delta=Delta( + content="secret", + role="assistant", + tool_calls=[ + { + "index": 0, + "id": "call_1", + "type": "function", + "function": {"name": "f", "arguments": "{}"}, + } + ], + ), + finish_reason="tool_calls", + ) + ], + ) + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, [mixed]) + + streamed = "".join(_delta_text(i) for i in out) + # Raw text never reaches the client; only the transformed text does. + assert "secret" not in streamed + assert "SECRET" in streamed + # The tool call is delivered, but its chunk carries no text. + tool_chunks = [i for i in out if i.choices[0].delta.tool_calls] + assert tool_chunks + assert all(not (c.choices[0].delta.content or "") for c in tool_chunks) + + @pytest.mark.asyncio + async def test_n_gt_1_text_and_tool_call_in_same_chunk_no_text_leak(self): + """n>1 chunk where one choice streams text and another a tool call: the + text choice must be transformed, not emitted raw alongside the tool call.""" + guardrail = _StreamingTextGuardrail() + + chunk = ModelResponseStream( + choices=[ + StreamingChoices(index=0, delta=Delta(content="secret", role="assistant"), finish_reason=None), + StreamingChoices( + index=1, + delta=Delta( + content=None, + role="assistant", + tool_calls=[ + { + "index": 0, + "id": "call_1", + "type": "function", + "function": {"name": "f", "arguments": "{}"}, + } + ], + ), + finish_reason="tool_calls", + ), + ], + ) + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, [chunk]) + + # choice 0's text is transformed, never delivered raw on the tool chunk. + for item in out: + for c in item.choices: + if c.delta.tool_calls: + assert not (c.delta.content or "") + all_text = "".join(c.delta.content or "" for i in out for c in i.choices) + assert "secret" not in all_text + assert "SECRET" in all_text + + @pytest.mark.asyncio + async def test_mixed_chunk_finish_reason_arrives_after_transformed_text(self): + """Fix #1: when a single chunk carries both delta.content and tool_calls + with finish_reason set, the passthrough must NOT emit finish_reason + before the transformed text — SSE clients that stop reading at + finish_reason would silently drop the guardrailed text. finish_reason + must ride on a terminator after the transformed text.""" + guardrail = _StreamingTextGuardrail() + mixed = ModelResponseStream( + choices=[ + StreamingChoices( + index=0, + delta=Delta( + content="secret", + role="assistant", + tool_calls=[ + { + "index": 0, + "id": "call_1", + "type": "function", + "function": {"name": "f", "arguments": "{}"}, + } + ], + ), + finish_reason="tool_calls", + ) + ], + ) + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, [mixed]) + + passthrough_idx = next(i for i, item in enumerate(out) if item.choices and item.choices[0].delta.tool_calls) + # Passthrough MUST NOT carry finish_reason for a mixed chunk (deferred). + assert out[passthrough_idx].choices[0].finish_reason is None, ( + f"passthrough of mixed chunk leaked finish_reason: {out[passthrough_idx].choices[0].finish_reason}" + ) + # finish_reason arrives via a terminator that comes AFTER the passthrough, + # so an SSE client reading top-down sees the transformed text before it + # sees the terminator. + finish_carriers = [ + i for i, item in enumerate(out) if item.choices and item.choices[0].finish_reason == "tool_calls" + ] + assert finish_carriers, "finish_reason=tool_calls never delivered" + assert min(finish_carriers) > passthrough_idx + # And the redacted text ("SECRET") reached the wire on some non-tool + # chunk (i.e. the text terminator). + transformed = "".join( + item.choices[0].delta.content or "" + for item in out + if item.choices and not item.choices[0].delta.tool_calls + ) + assert "SECRET" in transformed + assert "secret" not in transformed + + @pytest.mark.asyncio + async def test_text_flush_precedes_tool_call_passthrough(self): + """Fix #3: text chunks followed by a pure tool-call chunk carrying + finish_reason="tool_calls" must emit transformed text BEFORE the + passthrough, otherwise SSE-compliant clients stop reading at + finish_reason and drop the transformed text.""" + # sampling_rate 5: no mid-stream round would fire on 2 text chunks + # without the pre-tool-call flush. + guardrail = _StreamingTextGuardrail(sampling_rate=5) + chunks = [ + _stream_chunk("hello "), + _stream_chunk("world"), + ModelResponseStream( + choices=[ + StreamingChoices( + index=0, + delta=Delta( + content=None, + role="assistant", + tool_calls=[ + { + "index": 0, + "id": "call_1", + "type": "function", + "function": {"name": "f", "arguments": "{}"}, + } + ], + ), + finish_reason="tool_calls", + ) + ], + ), + ] + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks) + + text_indices = [ + i + for i, item in enumerate(out) + if item.choices and (item.choices[0].delta.content or "") and not item.choices[0].delta.tool_calls + ] + tool_indices = [i for i, item in enumerate(out) if item.choices and item.choices[0].delta.tool_calls] + assert text_indices, "transformed text was never emitted" + assert tool_indices, "tool-call passthrough missing" + assert max(text_indices) < min(tool_indices) + transformed = "".join(out[i].choices[0].delta.content or "" for i in text_indices) + assert "HELLO WORLD" in transformed + + @pytest.mark.asyncio + async def test_final_finish_reason_flushed_when_guardrail_suppresses_text(self): + """Fix #4: when the guardrail returns texts=[] (full suppression) and a + mixed content+tool_call chunk had deferred its finish_reason to the + text flush, the final flush must still emit a terminator chunk carrying + finish_reason. Otherwise the SSE stream ends without finish_reason.""" + + class _SuppressAll(_StreamingTextGuardrail): + async def apply_guardrail(self, inputs, request_data, input_type, **kwargs): + self.received_texts.append(list(inputs.get("texts") or [])) + return {"texts": []} + + mixed = ModelResponseStream( + choices=[ + StreamingChoices( + index=0, + delta=Delta( + content="secret", + role="assistant", + tool_calls=[ + { + "index": 0, + "id": "call_1", + "type": "function", + "function": {"name": "f", "arguments": "{}"}, + } + ], + ), + finish_reason="tool_calls", + ) + ], + ) + + out = await _drive_stream(UnifiedLLMGuardrails(), _SuppressAll(), [mixed]) + finishes = [c.finish_reason for item in out for c in item.choices if c.index == 0] + assert "tool_calls" in finishes + + @pytest.mark.asyncio + async def test_transform_sends_texts_sorted_by_choice_index(self): + """Fix #2: for n>1 streams where choice 1 emits before choice 0, the + transform must send texts to the guardrail in ascending choice-index + order so its returned texts realign to the correct choice indices.""" + + class _RecordingGuardrail(_StreamingTextGuardrail): + async def apply_guardrail(self, inputs, request_data, input_type, **kwargs): + self.received_texts.append(list(inputs.get("texts") or [])) + return {"texts": list(inputs.get("texts") or [])} + + guardrail = _RecordingGuardrail() + chunks = [ + _stream_chunk("beta", index=1), + _stream_chunk("alpha", index=0), + _stream_chunk("", index=0, finish_reason="stop"), + ] + await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks) + last = guardrail.received_texts[-1] + # Sorted ascending: alpha (index 0) before beta (index 1). + assert last[0].startswith("alpha") + assert last[1].startswith("beta") + + @pytest.mark.asyncio + async def test_non_idempotent_guardrail_not_double_applied_on_tool_call_streams(self): + """A non-idempotent guardrail must not see its own output as input on a + subsequent round. Regression guard for the bug where + ``_inspect_full_response_for_block`` shared ``responses_so_far`` with a + block-only path that mutated ``delta.content`` in place — the next + ``_round(is_final=True)`` would then re-read the already-guardrailed text + and re-apply the transform. + + Uses an n>1 chunk with text on choice 0 and tool_calls (with + finish_reason) on choice 1 — the exact shape that trips the + ``has_stream_ended=False → block path mutates anyway`` failure mode. + The test guardrail replaces the literal 'John' with '[REDACTED]', which is + non-idempotent: '[REDACTED]' does not contain 'John' so a second pass + produces the same output, BUT the raw accumulator would concat as + '[REDACTED]' + partial-raw, tripping stream_transform_underflow or + producing double-output. Assert the guardrail was called with the raw + text each time, not with any already-guardrailed prefix.""" + + class _RedactJohn(_StreamingTextGuardrail): + async def apply_guardrail(self, inputs, request_data, input_type, **kwargs): + texts = list(inputs.get("texts") or []) + self.received_texts.append(texts) + return {"texts": [t.replace("John", "[REDACTED]") for t in texts]} + + guardrail = _RedactJohn(sampling_rate=5) + + chunks = [ + _stream_chunk("John "), + _stream_chunk("went home."), + ModelResponseStream( + choices=[ + StreamingChoices(index=0, delta=Delta(content=None, role="assistant", tool_calls=None), finish_reason=None), + StreamingChoices( + index=1, + delta=Delta( + content=None, + role="assistant", + tool_calls=[ + { + "index": 0, + "id": "call_1", + "type": "function", + "function": {"name": "f", "arguments": "{}"}, + } + ], + ), + finish_reason="tool_calls", + ), + ], + ), + ] + + await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks) + + # Every text the guardrail saw as input must be raw ("John " not "[REDACTED]"). + # If the block path mutated the shared accumulator, the final _round would + # send "[REDACTED] went home." instead of "John went home.". + for received in guardrail.received_texts: + for text in received: + assert "[REDACTED]" not in text, ( + f"guardrail was re-invoked with its own already-redacted output " + f"— shared accumulator mutation regression: {received!r}" + ) + """The guardrail must receive the raw accumulated output each round, not a + transformed-prefix + raw-suffix mix (responses_so_far stays untouched).""" + guardrail = _StreamingTextGuardrail() + + chunks = [_stream_chunk("aa "), _stream_chunk("bb "), _stream_chunk("cc", finish_reason="stop")] + + await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks) + + # Every recorded input is the cumulative RAW (lowercase) text; if the + # accumulator were corrupted by write-back, later rounds would contain + # uppercased prefixes like "AA bb ". + for received in guardrail.received_texts: + assert received[0] == received[0].lower() + assert guardrail.received_texts[-1] == ["aa bb cc"] + + def test_accumulate_keys_by_choice_index_not_position(self): + """Single-choice chunks carrying a non-zero .index (n>1 streaming) must be + keyed by index, not enumerate position (which would collapse to 0).""" + handler = OpenAIChatCompletionsHandler() + chunks = [ + _stream_chunk("hello", index=1), + _stream_chunk(" world", index=1), + ] + + accumulated = handler._accumulate_string_content_by_choice_index(chunks) + + assert accumulated == {1: "hello world"} + + @pytest.mark.asyncio + async def test_terminal_chunk_not_guardrailed_twice(self): + """A terminal (finish_reason) chunk that is also a sampling boundary must + be processed once by the end-of-stream flush, not by a sampled round too.""" + guardrail = _StreamingTextGuardrail() # sampling_rate=1 + + chunks = [_stream_chunk("aa "), _stream_chunk("bb", finish_reason="stop")] + + await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks) + + # Round 1 (chunk 1) + end-of-stream flush = 2 calls. Without the terminal + # skip, chunk 2 would be guardrailed by a sampled round AND the flush (3). + assert guardrail.response_calls == 2 + + @pytest.mark.asyncio + async def test_malformed_holdback_from_in_process_guardrail_degrades(self): + """An in-process guardrail (bypassing from_dict) returning a null holdback + must degrade to 0 in the handler, not raise and abort the stream.""" + guardrail = _StreamingTextGuardrail(holdback_schedule=[None]) + + chunks = [_stream_chunk("abc"), _stream_chunk("def", finish_reason="stop")] + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks) + + # None holdback treated as 0: full text emitted, no crash. + assert "".join(_delta_text(i) for i in out) == "ABCDEF" diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 1a5c419754b..6d515ecdc73 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -615,6 +615,8 @@ def validate_model_cost_values(model_data, exceptions=None): "input_cost_per_audio_token", "output_cost_per_audio_token", "output_cost_per_image_token", + "input_cost_per_video_token", + "output_cost_per_video_token", "input_cost_per_audio_per_second", "input_cost_per_video_per_second", "input_cost_per_token_above_128k_tokens", @@ -732,6 +734,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "input_cost_per_image": {"type": "number"}, "input_cost_per_image_above_128k_tokens": {"type": "number"}, "input_cost_per_image_token": {"type": "number"}, + "input_cost_per_video_token": {"type": "number"}, "input_cost_per_token_above_200k_tokens": {"type": "number"}, "input_cost_per_token_above_256k_tokens": {"type": "number"}, "input_cost_per_token_above_272k_tokens": {"type": "number"}, @@ -807,6 +810,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "output_cost_per_character_above_128k_tokens": {"type": "number"}, "output_cost_per_image": {"type": "number"}, "output_cost_per_image_token": {"type": "number"}, + "output_cost_per_video_token": {"type": "number"}, "output_cost_per_pixel": {"type": "number"}, "output_cost_per_second": {"type": "number"}, "output_cost_per_second_1080p": {"type": "number"}, diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index fa2287bdaef..67139d2c932 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -301,17 +301,6 @@ "count": 1 } }, - "src/app/(dashboard)/guardrails/_components/edit_guardrail_form.tsx": { - "no-restricted-imports": { - "count": 1 - }, - "no-restricted-syntax": { - "count": 1 - }, - "react-hooks/set-state-in-effect": { - "count": 1 - } - }, "src/app/(dashboard)/guardrails/_components/guardrail_info.tsx": { "max-params": { "count": 1 @@ -339,14 +328,6 @@ "count": 1 } }, - "src/app/(dashboard)/guardrails/_components/guardrail_table.tsx": { - "no-nested-ternary": { - "count": 1 - }, - "no-restricted-imports": { - "count": 1 - } - }, "src/app/(dashboard)/guardrails/_components/tool_permission/ToolPermissionRulesEditor.tsx": { "no-restricted-imports": { "count": 1 diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.test.tsx index adaf7e4c63f..2ad332ae9ed 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.test.tsx @@ -105,6 +105,6 @@ describe("GuardrailsPanel", () => { expect(screen.getByText("Guardrails")).toBeInTheDocument(); // Activate the Guardrails tab so its content (including the Add button) is rendered fireEvent.click(screen.getByText("Guardrails")); - expect(screen.getByText("+ Add New Guardrail")).toBeInTheDocument(); + expect(screen.getByText("Add New Guardrail")).toBeInTheDocument(); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.tsx index af4c85d0b9f..33c6933634c 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.tsx @@ -1,7 +1,15 @@ import React, { useState, useEffect } from "react"; -import { Button, Dropdown, Tabs } from "antd"; -import { DownOutlined, PlusOutlined, CodeOutlined } from "@ant-design/icons"; +import { Tabs } from "antd"; +import { ChevronDown, Code, Plus } from "lucide-react"; import { getGuardrailsList, deleteGuardrailCall } from "@/components/networking"; +import { buttonVariants } from "@/components/ui/button"; +import { + DropdownMenu, + DropdownMenuContent, + DropdownMenuItem, + DropdownMenuTrigger, +} from "@/components/ui/dropdown-menu"; +import { cn } from "@/lib/cva.config"; import AddGuardrailForm from "./add_guardrail_form"; import GuardrailTable from "./guardrail_table"; import { isAdminRole } from "@/utils/roles"; @@ -133,30 +141,26 @@ const GuardrailsPanel: React.FC = ({ accessToken, userRole children: ( <>
- , - label: "Add Provider Guardrail", - onClick: handleAddGuardrail, - }, - { - key: "custom_code", - icon: , - label: "Create Custom Code Guardrail", - onClick: handleAddCustomCodeGuardrail, - }, - ], - }} - trigger={["click"]} - disabled={!accessToken} - > - - + + + + Add New Guardrail + + + + + + Add Provider Guardrail + + + + Create Custom Code Guardrail + + +
{selectedGuardrailId ? ( @@ -171,9 +175,6 @@ const GuardrailsPanel: React.FC = ({ accessToken, userRole guardrailsList={guardrailsList} isLoading={isLoading} onDeleteClick={handleDeleteClick} - accessToken={accessToken} - onGuardrailUpdated={fetchGuardrails} - isAdmin={isAdmin} onGuardrailClick={(id) => setSelectedGuardrailId(id)} /> )} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/edit_guardrail_form.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/edit_guardrail_form.tsx deleted file mode 100644 index a916984f8d7..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/edit_guardrail_form.tsx +++ /dev/null @@ -1,491 +0,0 @@ -import React, { useState, useEffect } from "react"; -import { Form, Typography, Select, Input, Switch, Modal } from "antd"; -import { Button, TextInput } from "@tremor/react"; -import { - guardrail_provider_map, - guardrailLogoMap, - getGuardrailProviders, - getSupportedModesForProvider, - toModeArray, - type SkipSystemMessageChoice, - type SkipToolMessageChoice, -} from "./guardrail_info_helpers"; -import { resolveLogoSrc } from "@/lib/assetPaths"; -import { getGuardrailUISettings, getGlobalLitellmHeaderName } from "@/components/networking"; -import PiiConfiguration from "./pii_configuration"; -import NotificationsManager from "@/components/molecules/notifications_manager"; - -const { Title, Text } = Typography; -const { Option } = Select; - -interface EditGuardrailFormProps { - visible: boolean; - onClose: () => void; - accessToken: string | null; - onSuccess: () => void; - guardrailId: string; - /** Full stored params merged into PUT so optional fields (e.g. content filter) are preserved. */ - fullLitellmParams?: Record | null; - initialValues: { - guardrail_name: string; - provider: string; - mode: string; - default_on: boolean; - pii_entities_config?: { [key: string]: string }; - skip_system_message_choice?: SkipSystemMessageChoice; - skip_tool_message_choice?: SkipToolMessageChoice; - [key: string]: unknown; - }; -} - -interface GuardrailSettings { - supported_entities: string[]; - supported_actions: string[]; - supported_modes: string[]; - supported_modes_by_provider?: Record; - pii_entity_categories: Array<{ - category: string; - entities: string[]; - }>; -} - -const EditGuardrailForm: React.FC = ({ - visible, - onClose, - accessToken, - onSuccess, - guardrailId, - fullLitellmParams, - initialValues, -}) => { - const [form] = Form.useForm(); - const [loading, setLoading] = useState(false); - const [selectedProvider, setSelectedProvider] = useState(initialValues?.provider || null); - const [guardrailSettings, setGuardrailSettings] = useState(null); - const [selectedEntities, setSelectedEntities] = useState([]); - const [selectedActions, setSelectedActions] = useState<{ [key: string]: string }>({}); - - // Fetch guardrail settings when the component mounts - useEffect(() => { - const fetchGuardrailSettings = async () => { - try { - if (!accessToken) return; - - const data = await getGuardrailUISettings(accessToken); - setGuardrailSettings(data); - } catch (error) { - console.error("Error fetching guardrail settings:", error); - NotificationsManager.fromBackend("Failed to load guardrail settings"); - } - }; - - fetchGuardrailSettings(); - }, [accessToken]); - - // Initialize selected entities and actions from initialValues - useEffect(() => { - if (initialValues?.pii_entities_config && Object.keys(initialValues.pii_entities_config).length > 0) { - const entities = Object.keys(initialValues.pii_entities_config); - setSelectedEntities(entities); - setSelectedActions(initialValues.pii_entities_config); - } - }, [initialValues]); - - const handleProviderChange = (value: string) => { - setSelectedProvider(value); - // Reset form fields that are provider-specific - form.setFieldsValue({ - config: undefined, - }); - - // Reset PII selections when changing provider - setSelectedEntities([]); - setSelectedActions({}); - }; - - const handleEntitySelect = (entity: string) => { - setSelectedEntities((prev) => { - if (prev.includes(entity)) { - return prev.filter((e) => e !== entity); - } else { - return [...prev, entity]; - } - }); - }; - - const handleActionSelect = (entity: string, action: string) => { - setSelectedActions((prev) => ({ - ...prev, - [entity]: action, - })); - }; - - const handleSubmit = async () => { - try { - setLoading(true); - const values = await form.validateFields(); - - // Get the guardrail provider value from the map - const guardrailProvider = guardrail_provider_map[values.provider]; - - const litellm_params: Record = - fullLitellmParams && typeof fullLitellmParams === "object" ? { ...fullLitellmParams } : {}; - - litellm_params.guardrail = guardrailProvider; - litellm_params.mode = values.mode; - litellm_params.default_on = values.default_on; - - const skipChoice = values.skip_system_message_choice as SkipSystemMessageChoice | undefined; - if (skipChoice === "yes") { - litellm_params.skip_system_message_in_guardrail = true; - } else if (skipChoice === "no") { - litellm_params.skip_system_message_in_guardrail = false; - } else { - delete litellm_params.skip_system_message_in_guardrail; - } - - const skipToolChoice = values.skip_tool_message_choice as SkipToolMessageChoice | undefined; - if (skipToolChoice === "yes") { - litellm_params.skip_tool_message_in_guardrail = true; - } else if (skipToolChoice === "no") { - litellm_params.skip_tool_message_in_guardrail = false; - } else { - delete litellm_params.skip_tool_message_in_guardrail; - } - - let guardrail_info: Record = {}; - - // For Presidio PII, add the entity and action configurations - if (values.provider === "PresidioPII" && selectedEntities.length > 0) { - const piiEntitiesConfig: { [key: string]: string } = {}; - selectedEntities.forEach((entity) => { - piiEntitiesConfig[entity] = selectedActions[entity] || "MASK"; // Default to MASK if no action selected - }); - - litellm_params.pii_entities_config = piiEntitiesConfig; - } - // Add config values to the guardrail_info if provided - else if (values.config) { - try { - const configObj = JSON.parse(values.config); - // For some guardrails, the config values need to be in litellm_params - // Especially for providers like Bedrock that need guardrailIdentifier and guardrailVersion - if (values.provider === "Bedrock" && configObj) { - if (configObj.guardrail_id) { - litellm_params.guardrailIdentifier = configObj.guardrail_id; - } - if (configObj.guardrail_version) { - litellm_params.guardrailVersion = configObj.guardrail_version; - } - } else { - // For other providers, add the config to guardrail_info - guardrail_info = configObj; - } - } catch (error) { - NotificationsManager.fromBackend("Invalid JSON in configuration"); - setLoading(false); - return; - } - } - - const guardrailData: { - guardrail_id: string; - guardrail: { - guardrail_name: string; - litellm_params: Record; - guardrail_info: Record; - }; - } = { - guardrail_id: guardrailId, - guardrail: { - guardrail_name: values.guardrail_name, - litellm_params, - guardrail_info, - }, - }; - - if (!accessToken) { - throw new Error("No access token available"); - } - - // Call the update endpoint - const url = `/guardrails/${guardrailId}`; - const response = await fetch(url, { - method: "PUT", - headers: { - [getGlobalLitellmHeaderName()]: `Bearer ${accessToken}`, - "Content-Type": "application/json", - }, - body: JSON.stringify(guardrailData), - }); - - if (!response.ok) { - const errorData = await response.text(); - throw new Error(errorData || "Failed to update guardrail"); - } - - NotificationsManager.success("Guardrail updated successfully"); - - // Reset and close - onSuccess(); - onClose(); - } catch (error) { - console.error("Failed to update guardrail:", error); - NotificationsManager.fromBackend( - "Failed to update guardrail: " + (error instanceof Error ? error.message : String(error)), - ); - } finally { - setLoading(false); - } - }; - - const renderPiiConfiguration = () => { - if (!guardrailSettings || !selectedProvider || selectedProvider !== "PresidioPII") return null; - - return ( - - ); - }; - - const renderProviderSpecificFields = () => { - if (!selectedProvider) return null; - - // For Presidio, we use the new PII configuration UI - if (selectedProvider === "PresidioPII") { - return renderPiiConfiguration(); - } - - switch (selectedProvider) { - case "Aporia": - return ( - - - - ); - case "AimSecurity": - return ( - - - - ); - case "Bedrock": - return ( - - - - ); - case "CatoNetworks": - return ( - - - - ); - case "GuardrailsAI": - return ( - - - - ); - case "LakeraAI": - return ( - - - - ); - case "PromptInjection": - return ( - - - - ); - default: - return ( - - - - ); - } - }; - - return ( - -
- - - - - - - - - - - - - - - - - - - - - - - - - {renderProviderSpecificFields()} - -
- - -
-
-
- ); -}; - -export default EditGuardrailForm; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrailTableColumns.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrailTableColumns.tsx new file mode 100644 index 00000000000..9ceb6ba244b --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrailTableColumns.tsx @@ -0,0 +1,176 @@ +"use client"; + +import { ColumnDef } from "@tanstack/react-table"; +import { MoreHorizontal, Trash2 } from "lucide-react"; + +import { DataTableSortHeader } from "@/components/shared/DataTable"; +import { DateCell, IdentityCell, StatusBadge } from "@/components/shared/table_cells"; +import { Guardrail, GuardrailDefinitionLocation } from "@/components/guardrails/types"; +import { buttonVariants } from "@/components/ui/button"; +import { + DropdownMenu, + DropdownMenuContent, + DropdownMenuItem, + DropdownMenuTrigger, +} from "@/components/ui/dropdown-menu"; +import { cn } from "@/lib/cva.config"; + +import { getGuardrailLogoAndName } from "./guardrail_info_helpers"; + +const CONFIG_DELETE_HINT = "Config guardrails are defined in the config file and cannot be deleted from the dashboard."; + +function GuardrailProviderCell({ provider }: { provider: string }) { + const { logo, displayName } = getGuardrailLogoAndName(provider); + return ( +
+ {logo ? ( + { + (event.currentTarget as HTMLImageElement).style.display = "none"; + }} + /> + ) : null} + {displayName} +
+ ); +} + +interface GuardrailRowActionsProps { + guardrail: Guardrail; + onDeleteClick: (guardrailId: string, guardrailName: string) => void; +} + +function GuardrailRowActions({ guardrail, onDeleteClick }: GuardrailRowActionsProps) { + const isConfigGuardrail = guardrail.guardrail_definition_location === GuardrailDefinitionLocation.CONFIG; + + return ( + + + + + + onDeleteClick(guardrail.guardrail_id, guardrail.guardrail_name || "Unnamed Guardrail")} + > + + Delete + + + + ); +} + +interface GuardrailTableColumnsDeps { + onGuardrailClick: (guardrailId: string) => void; + onDeleteClick: (guardrailId: string, guardrailName: string) => void; +} + +export const getGuardrailTableColumns = ({ + onGuardrailClick, + onDeleteClick, +}: GuardrailTableColumnsDeps): ColumnDef[] => [ + { + id: "guardrail_id", + accessorKey: "guardrail_id", + meta: { title: "Guardrail ID" }, + header: ({ column }) => , + size: 200, + enableSorting: true, + cell: ({ row }) => ( + onGuardrailClick(row.original.guardrail_id)} + /> + ), + }, + { + id: "guardrail_name", + accessorKey: "guardrail_name", + meta: { title: "Name" }, + header: ({ column }) => , + size: 200, + enableSorting: true, + cell: ({ row }) => { + const name = row.original.guardrail_name; + return ( + + {name || "-"} + + ); + }, + }, + { + id: "provider", + meta: { title: "Provider" }, + header: "Provider", + size: 180, + enableSorting: false, + cell: ({ row }) => , + }, + { + id: "mode", + meta: { title: "Mode" }, + header: "Mode", + size: 130, + enableSorting: false, + cell: ({ row }) => ( + {row.original.litellm_params.mode} + ), + }, + { + id: "default_on", + meta: { title: "Default On" }, + header: "Default On", + size: 120, + enableSorting: false, + cell: ({ row }) => { + const isDefaultOn = !!row.original.litellm_params?.default_on; + return ( + + ); + }, + }, + { + id: "created_at", + accessorKey: "created_at", + meta: { title: "Created At" }, + header: ({ column }) => , + size: 150, + enableSorting: true, + cell: ({ row }) => , + }, + { + id: "updated_at", + accessorKey: "updated_at", + meta: { title: "Updated At" }, + header: ({ column }) => , + size: 150, + enableSorting: true, + cell: ({ row }) => , + }, + { + id: "actions", + meta: { className: "text-right", headerClassName: "text-right" }, + header: () => Actions, + size: 64, + enableSorting: false, + enableHiding: false, + cell: ({ row }) => ( +
+ +
+ ), + }, +]; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_table.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_table.test.tsx index 70b5998ad95..4f556e74c16 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_table.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_table.test.tsx @@ -1,56 +1,61 @@ +import { render, screen } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { describe, it, expect, vi } from "vitest"; + import GuardrailTable from "./guardrail_table"; -import { render } from "@testing-library/react"; -import { describe, it, expect } from "vitest"; -import { GuardrailDefinitionLocation } from "@/components/guardrails/types"; +import { Guardrail, GuardrailDefinitionLocation } from "@/components/guardrails/types"; + +const baseProps = { + isLoading: false, + onDeleteClick: vi.fn(), + onGuardrailClick: vi.fn(), +}; + +const makeGuardrail = (overrides: Partial = {}): Guardrail => ({ + guardrail_id: "gr-1", + guardrail_name: "PII Redaction", + litellm_params: { guardrail: "presidio", mode: "pre_call", default_on: true }, + guardrail_info: null, + created_at: "2021-01-01", + updated_at: "2021-01-02", + guardrail_definition_location: GuardrailDefinitionLocation.DB, + ...overrides, +}); + describe("GuardrailTable", () => { - it("should render", () => { - const { getByText } = render( - {}} - accessToken={null} - onGuardrailUpdated={() => {}} - onGuardrailClick={() => {}} - />, - ); - expect(getByText("Guardrail ID")).toBeInTheDocument(); - expect(getByText("Name")).toBeInTheDocument(); - expect(getByText("Provider")).toBeInTheDocument(); - expect(getByText("Mode")).toBeInTheDocument(); - expect(getByText("Default On")).toBeInTheDocument(); - expect(getByText("Created At")).toBeInTheDocument(); - expect(getByText("Updated At")).toBeInTheDocument(); + it("renders every column header", () => { + render(); + for (const header of ["Guardrail ID", "Name", "Provider", "Mode", "Default On", "Created At", "Updated At"]) { + expect(screen.getByText(header)).toBeInTheDocument(); + } }); - it("should not allow deletion of config guardrails", () => { - const { getByTestId } = render( - {}} - accessToken={null} - onGuardrailUpdated={() => {}} - onGuardrailClick={() => {}} - />, - ); + it("deletes a DB guardrail through the actions menu", async () => { + const user = userEvent.setup(); + const onDeleteClick = vi.fn(); + const guardrail = makeGuardrail({ guardrail_id: "gr-9", guardrail_name: "Toxicity Filter" }); + render(); - const deleteGuardrailButton = getByTestId("config-delete-icon"); - expect(deleteGuardrailButton).toBeInTheDocument(); - expect(deleteGuardrailButton).toHaveClass("cursor-not-allowed text-gray-400"); - expect(deleteGuardrailButton).toHaveAttribute( - "title", - "Config guardrail cannot be deleted on the dashboard. Please delete it from the config file.", - ); + await user.click(screen.getByTestId("guardrail-actions-gr-9")); + await user.click(await screen.findByTestId("guardrail-action-delete")); + + expect(onDeleteClick).toHaveBeenCalledWith("gr-9", "Toxicity Filter"); + }); + + it("disables deletion for config guardrails so they cannot be removed from the dashboard", async () => { + const user = userEvent.setup(); + const onDeleteClick = vi.fn(); + const guardrail = makeGuardrail({ + guardrail_id: "cfg-1", + guardrail_name: "Config Guardrail", + guardrail_definition_location: GuardrailDefinitionLocation.CONFIG, + }); + render(); + + await user.click(screen.getByTestId("guardrail-actions-cfg-1")); + const deleteItem = await screen.findByTestId("guardrail-action-delete"); + + expect(deleteItem).toHaveAttribute("data-disabled"); + expect(onDeleteClick).not.toHaveBeenCalled(); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_table.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_table.tsx index 4ea4505f506..e6a14b2b2f4 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_table.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_table.tsx @@ -1,284 +1,61 @@ -import React, { useState } from "react"; -import { Table, TableBody, TableCell, TableHead, TableHeaderCell, TableRow, Icon } from "@tremor/react"; -import { TrashIcon, SwitchVerticalIcon, ChevronUpIcon, ChevronDownIcon } from "@heroicons/react/outline"; -import { Tooltip } from "antd"; -import { DateCell, IdCell, StatusBadge } from "@/components/shared/table_cells"; -import { - ColumnDef, - flexRender, - getCoreRowModel, - getSortedRowModel, - SortingState, - useReactTable, -} from "@tanstack/react-table"; -import { - getGuardrailLogoAndName, - guardrail_provider_map, - skipSystemMessageToChoice, - skipToolMessageToChoice, -} from "./guardrail_info_helpers"; -import EditGuardrailForm from "./edit_guardrail_form"; -import { Guardrail, GuardrailDefinitionLocation } from "@/components/guardrails/types"; +"use client"; + +import { SortingState } from "@tanstack/react-table"; +import { Inbox } from "lucide-react"; +import React, { useMemo, useState } from "react"; + +import { DataTable } from "@/components/shared/DataTable"; +import { Guardrail } from "@/components/guardrails/types"; + +import { getGuardrailTableColumns } from "./guardrailTableColumns"; interface GuardrailTableProps { guardrailsList: Guardrail[]; isLoading: boolean; onDeleteClick: (guardrailId: string, guardrailName: string) => void; - accessToken: string | null; - onGuardrailUpdated: () => void; - isAdmin?: boolean; onGuardrailClick: (id: string) => void; } +const DEFAULT_SORTING: SortingState = [{ id: "created_at", desc: true }]; + +function EmptyState() { + return ( +
+
+ +
+
No guardrails yet
+
Add a guardrail to start filtering requests and responses.
+
+ ); +} + const GuardrailTable: React.FC = ({ guardrailsList, isLoading, onDeleteClick, - accessToken, - onGuardrailUpdated, - isAdmin = false, onGuardrailClick, }) => { - const [sorting, setSorting] = useState([{ id: "created_at", desc: true }]); - const [editModalVisible, setEditModalVisible] = useState(false); - const [selectedGuardrail, setSelectedGuardrail] = useState(null); + const [sorting, setSorting] = useState(DEFAULT_SORTING); - const handleEditClick = (guardrail: Guardrail) => { - setSelectedGuardrail(guardrail); - setEditModalVisible(true); - }; - - const handleEditSuccess = () => { - setEditModalVisible(false); - setSelectedGuardrail(null); - onGuardrailUpdated(); - }; - - const columns: ColumnDef[] = [ - { - header: "Guardrail ID", - accessorKey: "guardrail_id", - cell: (info: any) => , - }, - { - header: "Name", - accessorKey: "guardrail_name", - cell: ({ row }) => { - const guardrail = row.original; - return ( - - {guardrail.guardrail_name || "-"} - - ); - }, - }, - { - header: "Provider", - accessorKey: "litellm_params.guardrail", - cell: ({ row }) => { - const guardrail = row.original; - const { logo, displayName } = getGuardrailLogoAndName(guardrail.litellm_params.guardrail); - return ( -
- {logo && ( - {`${displayName} { - // Hide broken image - (e.target as HTMLImageElement).style.display = "none"; - }} - /> - )} - {displayName} -
- ); - }, - }, - { - header: "Mode", - accessorKey: "litellm_params.mode", - cell: ({ row }) => { - const guardrail = row.original; - return {guardrail.litellm_params.mode}; - }, - }, - { - header: "Default On", - accessorKey: "litellm_params.default_on", - cell: ({ row }) => { - const isDefaultOn = !!row.original.litellm_params?.default_on; - return ( - - ); - }, - }, - { - header: "Created At", - accessorKey: "created_at", - cell: ({ row }) => , - }, - { - header: "Updated At", - accessorKey: "updated_at", - cell: ({ row }) => , - }, - { - id: "actions", - header: "Actions", - cell: ({ row }) => { - const guardrail = row.original; - const isConfigGuardrail = guardrail.guardrail_definition_location === GuardrailDefinitionLocation.CONFIG; - return ( -
- {isConfigGuardrail ? ( - - - - ) : ( - - - guardrail.guardrail_id && - onDeleteClick(guardrail.guardrail_id, guardrail.guardrail_name || "Unnamed Guardrail") - } - className="cursor-pointer hover:text-red-500" - /> - - )} -
- ); - }, - }, - ]; - - const table = useReactTable({ - data: guardrailsList, - columns, - state: { - sorting, - }, - onSortingChange: setSorting, - getCoreRowModel: getCoreRowModel(), - getSortedRowModel: getSortedRowModel(), - enableSorting: true, - }); + const columns = useMemo( + () => getGuardrailTableColumns({ onGuardrailClick, onDeleteClick }), + [onGuardrailClick, onDeleteClick], + ); return ( -
-
- - - {table.getHeaderGroups().map((headerGroup) => ( - - {headerGroup.headers.map((header) => ( - -
-
- {header.isPlaceholder ? null : flexRender(header.column.columnDef.header, header.getContext())} -
- {header.id !== "actions" && ( -
- {header.column.getIsSorted() ? ( - { - asc: , - desc: , - }[header.column.getIsSorted() as string] - ) : ( - - )} -
- )} -
-
- ))} -
- ))} -
- - {isLoading ? ( - - -
-

Loading...

-
-
-
- ) : guardrailsList.length > 0 ? ( - table.getRowModel().rows.map((row) => ( - - {row.getVisibleCells().map((cell) => ( - - {flexRender(cell.column.columnDef.cell, cell.getContext())} - - ))} - - )) - ) : ( - - -
-

No guardrails found

-
-
-
- )} -
-
-
- - {/* Edit Modal */} - {selectedGuardrail && ( - setEditModalVisible(false)} - accessToken={accessToken} - onSuccess={handleEditSuccess} - guardrailId={selectedGuardrail.guardrail_id || ""} - fullLitellmParams={selectedGuardrail.litellm_params} - initialValues={{ - guardrail_name: selectedGuardrail.guardrail_name || "", - provider: - Object.keys(guardrail_provider_map).find( - (key) => guardrail_provider_map[key] === selectedGuardrail?.litellm_params.guardrail, - ) || "", - mode: selectedGuardrail.litellm_params.mode, - default_on: selectedGuardrail.litellm_params.default_on, - pii_entities_config: selectedGuardrail.litellm_params.pii_entities_config, - skip_system_message_choice: skipSystemMessageToChoice( - selectedGuardrail.litellm_params?.skip_system_message_in_guardrail, - ), - skip_tool_message_choice: skipToolMessageToChoice( - selectedGuardrail.litellm_params?.skip_tool_message_in_guardrail, - ), - ...selectedGuardrail.guardrail_info, - }} - /> - )} -
+ guardrail.guardrail_id || String(index)} + sortingMode="client" + sorting={sorting} + onSortingChange={setSorting} + isLoading={isLoading} + loadingMessage="Loading guardrails…" + noDataMessage={} + size="compact" + /> ); }; diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx index 0d3468d4ada..60547415908 100644 --- a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx @@ -358,6 +358,23 @@ describe("DataTable loading", () => { expect(names()).toEqual(["Charlie", "Alice", "Bob"]); }); + it("gives compact skeleton rows the same height as loaded rows so loading does not shrink the table", () => { + const { rerender } = render( + , + ); + const skeletonRow = screen.getAllByTestId("skeleton-row").at(0); + const loadedRowHeight = "h-8"; + expect(skeletonRow?.className).toContain(loadedRowHeight); + + rerender(); + expect(document.querySelector("[data-row-id]")?.className).toContain(loadedRowHeight); + }); + + it("does not force the compact height on default-size skeleton rows", () => { + render(); + expect(screen.getAllByTestId("skeleton-row").at(0)?.className).not.toContain("h-8"); + }); + it("varies skeleton shape and width per column instead of one fixed bar", () => { const columns: ColumnDef[] = [ { accessorKey: "name", header: "Name", meta: { skeleton: "twoLine" }, cell: () => null }, diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx index d55831b87ca..fa0e672026e 100644 --- a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx +++ b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx @@ -389,7 +389,11 @@ function SkeletonRows({ return ( {rowKeys.map((rowKey) => ( - + {cells.map((column, columnKey) => ( diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 59c0532ccbe..547e84a3cd2 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -25706,6 +25706,8 @@ export interface components { input_cost_per_video_per_second_above_15s_interval?: number | null; /** Input Cost Per Video Per Second Above 8S Interval */ input_cost_per_video_per_second_above_8s_interval?: number | null; + /** Input Cost Per Video Token */ + input_cost_per_video_token?: number | null; /** Itpm */ itpm?: number | null; /** Litellm Credential Name */ @@ -25787,6 +25789,8 @@ export interface components { output_cost_per_token_priority?: number | null; /** Output Cost Per Video Per Second */ output_cost_per_video_per_second?: number | null; + /** Output Cost Per Video Token */ + output_cost_per_video_token?: number | null; /** Output Vector Size */ output_vector_size?: number | null; /** Quality Router Config */ @@ -33536,6 +33540,8 @@ export interface components { input_cost_per_video_per_second_above_15s_interval?: number | null; /** Input Cost Per Video Per Second Above 8S Interval */ input_cost_per_video_per_second_above_8s_interval?: number | null; + /** Input Cost Per Video Token */ + input_cost_per_video_token?: number | null; /** Itpm */ itpm?: number | null; /** Litellm Credential Name */ @@ -33617,6 +33623,8 @@ export interface components { output_cost_per_token_priority?: number | null; /** Output Cost Per Video Per Second */ output_cost_per_video_per_second?: number | null; + /** Output Cost Per Video Token */ + output_cost_per_video_token?: number | null; /** Output Vector Size */ output_vector_size?: number | null; /** Quality Router Config */