scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_google_interactions_cost

Browse source
This commit is contained in:
mateo-berri 2026-07-14 18:45:01 -07:00
commit 887fc0a73c
37 changed files with 2752 additions and 942 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
litellm/integrations/s3.py
View file

@ -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

35
litellm/litellm_core_utils/llm_cost_calc/utils.py
View file

@ -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).

38
litellm/llms/base_llm/guardrail_translation/base_translation.py
View file

@ -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.

207
litellm/llms/openai/chat/guardrail_translation/handler.py
View file

@ -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,
)

2
litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py
View file

@ -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

64
litellm/model_prices_and_context_window_backup.json
View file

@ -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,

10
litellm/proxy/auth/user_api_key_auth.py
View file

@ -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

1
litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/__init__.py
View file

@ -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)

40
litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py
View file

@ -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,
)

544
litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py
View file

@ -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

43
litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py
View file

@ -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,
)

7
litellm/types/utils.py
View file

@ -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)

2
litellm/utils.py
View file

@ -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),

64
model_prices_and_context_window.json
View file

@ -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,

8
tests/e2e/CLAUDE.md
View file

@ -77,13 +77,13 @@ llm.<endpoint>.<route>.<capability>.<streaming>.<assertion>
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

2
tests/e2e/coverage_registry/llm_conversational.yaml
View file

@ -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"}

2
tests/e2e/coverage_registry/schema.py
View file

@ -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",

30
tests/e2e/llm_translation/endpoints_client.py
View file

@ -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:

18
tests/e2e/llm_translation/test_cache_control.py
View file

@ -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]

220
tests/e2e/llm_translation/test_messages_mid_conversation_system_e2e.py Normal file
View file

@ -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="<system-reminder>Answer with exactly one word.</system-reminder>"
)
],
)
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"
)

41
tests/test_litellm/integrations/test_s3_v2.py
View file

@ -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)
# --------------------------------------------------------------

100
tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py
View file

@ -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,

6
tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py
View file

@ -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()

74
tests/test_litellm/proxy/auth/test_user_api_key_auth.py
View file

@ -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

97
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py
View file

@ -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."""

842
tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py
View file

@ -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"

4
tests/test_litellm/test_utils.py
View file

@ -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"},

19
ui/litellm-dashboard/eslint-suppressions.json
View file

@ -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

2
ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.test.tsx
View file

@ -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();
});
});

59
ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.tsx
View file

@ -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<GuardrailsPanelProps> = ({ accessToken, userRole
children: (
<>
<div className="flex justify-between items-center mb-4">
<Dropdown
menu={{
items: [
{
key: "provider",
icon: <PlusOutlined />,
label: "Add Provider Guardrail",
onClick: handleAddGuardrail,
},
{
key: "custom_code",
icon: <CodeOutlined />,
label: "Create Custom Code Guardrail",
onClick: handleAddCustomCodeGuardrail,
},
],
}}
trigger={["click"]}
disabled={!accessToken}
>
<Button disabled={!accessToken}>
+ Add New Guardrail <DownOutlined className="ml-2" />
</Button>
</Dropdown>
<DropdownMenu>
<DropdownMenuTrigger
disabled={!accessToken}
className={cn(buttonVariants({ variant: "default" }))}
>
<Plus />
Add New Guardrail
<ChevronDown />
</DropdownMenuTrigger>
<DropdownMenuContent align="start" className="w-56">
<DropdownMenuItem onClick={handleAddGuardrail}>
<Plus />
Add Provider Guardrail
</DropdownMenuItem>
<DropdownMenuItem onClick={handleAddCustomCodeGuardrail}>
<Code />
Create Custom Code Guardrail
</DropdownMenuItem>
</DropdownMenuContent>
</DropdownMenu>
</div>
{selectedGuardrailId ? (
@ -171,9 +175,6 @@ const GuardrailsPanel: React.FC<GuardrailsPanelProps> = ({ accessToken, userRole
guardrailsList={guardrailsList}
isLoading={isLoading}
onDeleteClick={handleDeleteClick}
accessToken={accessToken}
onGuardrailUpdated={fetchGuardrails}
isAdmin={isAdmin}
onGuardrailClick={(id) => setSelectedGuardrailId(id)}
/>
)}

491
ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/edit_guardrail_form.tsx
View file

@ -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<string, unknown> | 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<string, string[]>;
pii_entity_categories: Array<{
category: string;
entities: string[];
}>;
}
const EditGuardrailForm: React.FC<EditGuardrailFormProps> = ({
visible,
onClose,
accessToken,
onSuccess,
guardrailId,
fullLitellmParams,
initialValues,
}) => {
const [form] = Form.useForm();
const [loading, setLoading] = useState(false);
const [selectedProvider, setSelectedProvider] = useState<string | null>(initialValues?.provider || null);
const [guardrailSettings, setGuardrailSettings] = useState<GuardrailSettings | null>(null);
const [selectedEntities, setSelectedEntities] = useState<string[]>([]);
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<string, unknown> =
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<string, unknown> = {};
// 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<string, unknown>;
guardrail_info: Record<string, unknown>;
};
} = {
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 (
<PiiConfiguration
entities={guardrailSettings.supported_entities}
actions={guardrailSettings.supported_actions}
selectedEntities={selectedEntities}
selectedActions={selectedActions}
onEntitySelect={handleEntitySelect}
onActionSelect={handleActionSelect}
entityCategories={guardrailSettings.pii_entity_categories}
/>
);
};
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 (
<Form.Item label="Aporia Configuration" name="config" tooltip="JSON configuration for Aporia">
<Input.TextArea
rows={4}
placeholder={`{
"api_key": "your_aporia_api_key",
"project_name": "your_project_name"
}`}
/>
</Form.Item>
);
case "AimSecurity":
return (
<Form.Item label="Aim Security Configuration" name="config" tooltip="JSON configuration for Aim Security">
<Input.TextArea
rows={4}
placeholder={`{
"api_key": "your_aim_api_key"
}`}
/>
</Form.Item>
);
case "Bedrock":
return (
<Form.Item
label="Amazon Bedrock Configuration"
name="config"
tooltip="JSON configuration for Amazon Bedrock guardrails"
>
<Input.TextArea
rows={4}
placeholder={`{
"guardrail_id": "your_guardrail_id",
"guardrail_version": "your_guardrail_version"
}`}
/>
</Form.Item>
);
case "CatoNetworks":
return (
<Form.Item label="Cato Networks Configuration" name="config" tooltip="JSON configuration for Cato Networks">
<Input.TextArea
rows={4}
placeholder={`{
"api_key": "your_cato_api_key"
}`}
/>
</Form.Item>
);
case "GuardrailsAI":
return (
<Form.Item label="Guardrails.ai Configuration" name="config" tooltip="JSON configuration for Guardrails.ai">
<Input.TextArea
rows={4}
placeholder={`{
"api_key": "your_guardrails_api_key",
"guardrail_id": "your_guardrail_id"
}`}
/>
</Form.Item>
);
case "LakeraAI":
return (
<Form.Item label="Lakera AI Configuration" name="config" tooltip="JSON configuration for Lakera AI">
<Input.TextArea
rows={4}
placeholder={`{
"api_key": "your_lakera_api_key"
}`}
/>
</Form.Item>
);
case "PromptInjection":
return (
<Form.Item
label="Prompt Injection Configuration"
name="config"
tooltip="JSON configuration for prompt injection detection"
>
<Input.TextArea
rows={4}
placeholder={`{
"threshold": 0.8
}`}
/>
</Form.Item>
);
default:
return (
<Form.Item label="Custom Configuration" name="config" tooltip="JSON configuration for your custom guardrail">
<Input.TextArea
rows={4}
placeholder={`{
"key1": "value1",
"key2": "value2"
}`}
/>
</Form.Item>
);
}
};
return (
<Modal title="Edit Guardrail" open={visible} onCancel={onClose} footer={null} width={700}>
<Form form={form} layout="vertical" initialValues={initialValues}>
<Form.Item
name="guardrail_name"
label="Guardrail Name"
rules={[{ required: true, message: "Please enter a guardrail name" }]}
>
<TextInput placeholder="Enter a name for this guardrail" />
</Form.Item>
<Form.Item
name="provider"
label="Guardrail Provider"
rules={[{ required: true, message: "Please select a provider" }]}
>
<Select
placeholder="Select a guardrail provider"
onChange={handleProviderChange}
disabled={true} // Disable changing provider in edit mode
optionLabelProp="label"
>
{Object.entries(getGuardrailProviders()).map(([key, value]) => (
<Option key={key} value={key} label={value}>
<div style={{ display: "flex", alignItems: "center" }}>
{guardrailLogoMap[value] && (
<img
src={resolveLogoSrc(guardrailLogoMap[value])}
alt=""
style={{
height: "20px",
width: "20px",
marginRight: "8px",
objectFit: "contain",
}}
onError={(e) => {
// Hide broken image icon if image fails to load
e.currentTarget.style.display = "none";
}}
/>
)}
<span>{value}</span>
</div>
</Option>
))}
</Select>
</Form.Item>
<Form.Item
name="mode"
label="Mode"
tooltip="How the guardrail should be applied"
rules={[{ required: true, message: "Please select a mode" }]}
>
<Select>
{(() => {
const modes = getSupportedModesForProvider(guardrailSettings, selectedProvider) ?? [
"pre_call",
"post_call",
];
const currentModes = toModeArray(initialValues?.mode);
const unsupportedCurrent = currentModes.filter((m) => !modes.includes(m));
return [...unsupportedCurrent, ...modes].map((mode) => (
<Option key={mode} value={mode}>
{unsupportedCurrent.includes(mode)
? `${mode} (not supported by ${selectedProvider}, pick another)`
: mode}
</Option>
));
})()}
</Select>
</Form.Item>
<Form.Item
name="default_on"
label="Always On"
tooltip="If enabled, this guardrail will be applied to all requests by default"
valuePropName="checked"
>
<Switch />
</Form.Item>
<Form.Item
name="skip_system_message_choice"
label="Skip system messages in guardrail"
tooltip="Unified guardrails only: whether role: system content is omitted from guardrail input (LLM still receives full messages). Use global default follows litellm_settings.skip_system_message_in_guardrail."
>
<Select>
<Option value="inherit">Use global default</Option>
<Option value="yes">Yes — exclude from guardrail scan</Option>
<Option value="no">No — always include in scan</Option>
</Select>
</Form.Item>
<Form.Item
name="skip_tool_message_choice"
label="Skip tool messages in guardrail"
tooltip="Unified guardrails only: whether role: tool content is omitted from guardrail input (LLM still receives full messages). Use global default follows litellm_settings.skip_tool_message_in_guardrail."
>
<Select>
<Option value="inherit">Use global default</Option>
<Option value="yes">Yes — exclude from guardrail scan</Option>
<Option value="no">No — always include in scan</Option>
</Select>
</Form.Item>
{renderProviderSpecificFields()}
<div className="flex justify-end space-x-2 mt-4">
<Button variant="secondary" onClick={onClose}>
Cancel
</Button>
<Button onClick={handleSubmit} loading={loading}>
Update Guardrail
</Button>
</div>
</Form>
</Modal>
);
};
export default EditGuardrailForm;

176
ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrailTableColumns.tsx Normal file
View file

@ -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 (
<div className="flex items-center gap-2">
{logo ? (
<img
src={logo}
alt=""
className="size-4 shrink-0"
onError={(event) => {
(event.currentTarget as HTMLImageElement).style.display = "none";
}}
/>
) : null}
<span className="truncate text-sm">{displayName}</span>
</div>
);
}
interface GuardrailRowActionsProps {
guardrail: Guardrail;
onDeleteClick: (guardrailId: string, guardrailName: string) => void;
}
function GuardrailRowActions({ guardrail, onDeleteClick }: GuardrailRowActionsProps) {
const isConfigGuardrail = guardrail.guardrail_definition_location === GuardrailDefinitionLocation.CONFIG;
return (
<DropdownMenu>
<DropdownMenuTrigger
aria-label="Open guardrail actions"
data-testid={`guardrail-actions-${guardrail.guardrail_id}`}
className={cn(buttonVariants({ variant: "ghost", size: "icon-sm" }), "text-muted-foreground")}
>
<MoreHorizontal className="size-4" />
</DropdownMenuTrigger>
<DropdownMenuContent align="end" className="w-52">
<DropdownMenuItem
variant="destructive"
disabled={isConfigGuardrail}
data-testid="guardrail-action-delete"
title={isConfigGuardrail ? CONFIG_DELETE_HINT : undefined}
onClick={() => onDeleteClick(guardrail.guardrail_id, guardrail.guardrail_name || "Unnamed Guardrail")}
>
<Trash2 />
Delete
</DropdownMenuItem>
</DropdownMenuContent>
</DropdownMenu>
);
}
interface GuardrailTableColumnsDeps {
onGuardrailClick: (guardrailId: string) => void;
onDeleteClick: (guardrailId: string, guardrailName: string) => void;
}
export const getGuardrailTableColumns = ({
onGuardrailClick,
onDeleteClick,
}: GuardrailTableColumnsDeps): ColumnDef<Guardrail>[] => [
{
id: "guardrail_id",
accessorKey: "guardrail_id",
meta: { title: "Guardrail ID" },
header: ({ column }) => <DataTableSortHeader column={column} title="Guardrail ID" />,
size: 200,
enableSorting: true,
cell: ({ row }) => (
<IdentityCell
title={row.original.guardrail_id}
titleClassName="font-mono text-xs font-normal"
onClick={() => onGuardrailClick(row.original.guardrail_id)}
/>
),
},
{
id: "guardrail_name",
accessorKey: "guardrail_name",
meta: { title: "Name" },
header: ({ column }) => <DataTableSortHeader column={column} title="Name" />,
size: 200,
enableSorting: true,
cell: ({ row }) => {
const name = row.original.guardrail_name;
return (
<span className="block truncate text-sm font-medium" title={name ?? undefined}>
{name || "-"}
</span>
);
},
},
{
id: "provider",
meta: { title: "Provider" },
header: "Provider",
size: 180,
enableSorting: false,
cell: ({ row }) => <GuardrailProviderCell provider={row.original.litellm_params.guardrail} />,
},
{
id: "mode",
meta: { title: "Mode" },
header: "Mode",
size: 130,
enableSorting: false,
cell: ({ row }) => (
<span className="font-mono text-xs text-muted-foreground">{row.original.litellm_params.mode}</span>
),
},
{
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 (
<StatusBadge tone={isDefaultOn ? "success" : "neutral"} label={isDefaultOn ? "Default On" : "Default Off"} />
);
},
},
{
id: "created_at",
accessorKey: "created_at",
meta: { title: "Created At" },
header: ({ column }) => <DataTableSortHeader column={column} title="Created At" />,
size: 150,
enableSorting: true,
cell: ({ row }) => <DateCell value={row.original.created_at} />,
},
{
id: "updated_at",
accessorKey: "updated_at",
meta: { title: "Updated At" },
header: ({ column }) => <DataTableSortHeader column={column} title="Updated At" />,
size: 150,
enableSorting: true,
cell: ({ row }) => <DateCell value={row.original.updated_at} />,
},
{
id: "actions",
meta: { className: "text-right", headerClassName: "text-right" },
header: () => <span className="sr-only">Actions</span>,
size: 64,
enableSorting: false,
enableHiding: false,
cell: ({ row }) => (
<div className="flex justify-end">
<GuardrailRowActions guardrail={row.original} onDeleteClick={onDeleteClick} />
</div>
),
},
];

103
ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_table.test.tsx
View file

@ -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 => ({
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(
<GuardrailTable
guardrailsList={[]}
isLoading={false}
onDeleteClick={() => {}}
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(<GuardrailTable guardrailsList={[]} {...baseProps} />);
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(
<GuardrailTable
guardrailsList={[
{
guardrail_id: "1",
guardrail_name: "Guardrail 1",
litellm_params: { guardrail: "presidio", mode: "pre_call", default_on: true },
guardrail_info: null,
created_at: "2021-01-01",
updated_at: "2021-01-01",
guardrail_definition_location: GuardrailDefinitionLocation.CONFIG,
},
]}
isLoading={false}
onDeleteClick={() => {}}
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(<GuardrailTable guardrailsList={[guardrail]} {...baseProps} onDeleteClick={onDeleteClick} />);
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(<GuardrailTable guardrailsList={[guardrail]} {...baseProps} onDeleteClick={onDeleteClick} />);
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();
});
});

305
ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_table.tsx
View file

@ -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 (
<div className="flex flex-col items-center gap-1 py-6">
<div className="mb-1 flex size-10 items-center justify-center rounded-lg bg-muted">
<Inbox className="size-5 text-muted-foreground" />
</div>
<div className="text-sm font-medium text-foreground">No guardrails yet</div>
<div className="text-sm text-muted-foreground">Add a guardrail to start filtering requests and responses.</div>
</div>
);
}
const GuardrailTable: React.FC<GuardrailTableProps> = ({
guardrailsList,
isLoading,
onDeleteClick,
accessToken,
onGuardrailUpdated,
isAdmin = false,
onGuardrailClick,
}) => {
const [sorting, setSorting] = useState<SortingState>([{ id: "created_at", desc: true }]);
const [editModalVisible, setEditModalVisible] = useState(false);
const [selectedGuardrail, setSelectedGuardrail] = useState<Guardrail | null>(null);
const [sorting, setSorting] = useState<SortingState>(DEFAULT_SORTING);
const handleEditClick = (guardrail: Guardrail) => {
setSelectedGuardrail(guardrail);
setEditModalVisible(true);
};
const handleEditSuccess = () => {
setEditModalVisible(false);
setSelectedGuardrail(null);
onGuardrailUpdated();
};
const columns: ColumnDef<Guardrail>[] = [
{
header: "Guardrail ID",
accessorKey: "guardrail_id",
cell: (info: any) => <IdCell value={info.getValue()} onClick={onGuardrailClick} />,
},
{
header: "Name",
accessorKey: "guardrail_name",
cell: ({ row }) => {
const guardrail = row.original;
return (
<Tooltip title={guardrail.guardrail_name}>
<span className="text-xs font-medium">{guardrail.guardrail_name || "-"}</span>
</Tooltip>
);
},
},
{
header: "Provider",
accessorKey: "litellm_params.guardrail",
cell: ({ row }) => {
const guardrail = row.original;
const { logo, displayName } = getGuardrailLogoAndName(guardrail.litellm_params.guardrail);
return (
<div className="flex items-center space-x-2">
{logo && (
<img
src={logo}
alt={`${displayName} logo`}
className="w-4 h-4"
onError={(e) => {
// Hide broken image
(e.target as HTMLImageElement).style.display = "none";
}}
/>
)}
<span className="text-xs">{displayName}</span>
</div>
);
},
},
{
header: "Mode",
accessorKey: "litellm_params.mode",
cell: ({ row }) => {
const guardrail = row.original;
return <span className="text-xs">{guardrail.litellm_params.mode}</span>;
},
},
{
header: "Default On",
accessorKey: "litellm_params.default_on",
cell: ({ row }) => {
const isDefaultOn = !!row.original.litellm_params?.default_on;
return (
<StatusBadge tone={isDefaultOn ? "success" : "neutral"} label={isDefaultOn ? "Default On" : "Default Off"} />
);
},
},
{
header: "Created At",
accessorKey: "created_at",
cell: ({ row }) => <DateCell value={row.original.created_at} />,
},
{
header: "Updated At",
accessorKey: "updated_at",
cell: ({ row }) => <DateCell value={row.original.updated_at} />,
},
{
id: "actions",
header: "Actions",
cell: ({ row }) => {
const guardrail = row.original;
const isConfigGuardrail = guardrail.guardrail_definition_location === GuardrailDefinitionLocation.CONFIG;
return (
<div className="flex space-x-2">
{isConfigGuardrail ? (
<Tooltip title="Config guardrail cannot be deleted on the dashboard. Please delete it from the config file.">
<Icon
data-testid="config-delete-icon"
icon={TrashIcon}
size="sm"
className="cursor-not-allowed text-gray-400"
title="Config guardrail cannot be deleted on the dashboard. Please delete it from the config file."
aria-label="Delete guardrail (config)"
/>
</Tooltip>
) : (
<Tooltip title="Delete guardrail">
<Icon
icon={TrashIcon}
size="sm"
onClick={() =>
guardrail.guardrail_id &&
onDeleteClick(guardrail.guardrail_id, guardrail.guardrail_name || "Unnamed Guardrail")
}
className="cursor-pointer hover:text-red-500"
/>
</Tooltip>
)}
</div>
);
},
},
];
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 (
<div className="rounded-lg custom-border relative">
<div className="overflow-x-auto">
<Table className="[&_td]:py-0.5 [&_th]:py-1">
<TableHead>
{table.getHeaderGroups().map((headerGroup) => (
<TableRow key={headerGroup.id}>
{headerGroup.headers.map((header) => (
<TableHeaderCell
key={header.id}
className={`py-1 h-8 ${
header.id === "actions" ? "sticky right-0 bg-white shadow-[-4px_0_8px_-6px_rgba(0,0,0,0.1)]" : ""
}`}
onClick={header.column.getToggleSortingHandler()}
>
<div className="flex items-center justify-between gap-2">
<div className="flex items-center">
{header.isPlaceholder ? null : flexRender(header.column.columnDef.header, header.getContext())}
</div>
{header.id !== "actions" && (
<div className="w-4">
{header.column.getIsSorted() ? (
{
asc: <ChevronUpIcon className="h-4 w-4 text-blue-500" />,
desc: <ChevronDownIcon className="h-4 w-4 text-blue-500" />,
}[header.column.getIsSorted() as string]
) : (
<SwitchVerticalIcon className="h-4 w-4 text-gray-400" />
)}
</div>
)}
</div>
</TableHeaderCell>
))}
</TableRow>
))}
</TableHead>
<TableBody>
{isLoading ? (
<TableRow>
<TableCell colSpan={columns.length} className="h-8 text-center">
<div className="text-center text-gray-500">
<p>Loading...</p>
</div>
</TableCell>
</TableRow>
) : guardrailsList.length > 0 ? (
table.getRowModel().rows.map((row) => (
<TableRow key={row.id} className="h-8">
{row.getVisibleCells().map((cell) => (
<TableCell
key={cell.id}
className={`py-0.5 max-h-8 overflow-hidden text-ellipsis whitespace-nowrap ${
cell.column.id === "actions"
? "sticky right-0 bg-white shadow-[-4px_0_8px_-6px_rgba(0,0,0,0.1)]"
: ""
}`}
>
{flexRender(cell.column.columnDef.cell, cell.getContext())}
</TableCell>
))}
</TableRow>
))
) : (
<TableRow>
<TableCell colSpan={columns.length} className="h-8 text-center">
<div className="text-center text-gray-500">
<p>No guardrails found</p>
</div>
</TableCell>
</TableRow>
)}
</TableBody>
</Table>
</div>
{/* Edit Modal */}
{selectedGuardrail && (
<EditGuardrailForm
visible={editModalVisible}
onClose={() => 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,
}}
/>
)}
</div>
<DataTable
data={guardrailsList}
columns={columns}
getRowId={(guardrail, index) => guardrail.guardrail_id || String(index)}
sortingMode="client"
sorting={sorting}
onSortingChange={setSorting}
isLoading={isLoading}
loadingMessage="Loading guardrails…"
noDataMessage={<EmptyState />}
size="compact"
/>
);
};

17
ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx
View file

@ -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(
<DataTable data={CHARLIE_ALICE_BOB} columns={nameCellColumns} size="compact" isLoading />,
);
const skeletonRow = screen.getAllByTestId("skeleton-row").at(0);
const loadedRowHeight = "h-8";
expect(skeletonRow?.className).toContain(loadedRowHeight);
rerender(<DataTable data={CHARLIE_ALICE_BOB} columns={nameCellColumns} size="compact" />);
expect(document.querySelector("[data-row-id]")?.className).toContain(loadedRowHeight);
});
it("does not force the compact height on default-size skeleton rows", () => {
render(<DataTable data={CHARLIE_ALICE_BOB} columns={nameCellColumns} isLoading />);
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<Person, unknown>[] = [
{ accessorKey: "name", header: "Name", meta: { skeleton: "twoLine" }, cell: () => null },

6
ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx
View file

@ -389,7 +389,11 @@ function SkeletonRows<TData>({
return (
<Fragment>
{rowKeys.map((rowKey) => (
<TableRow key={`skeleton-${rowKey}`} className="hover:bg-transparent" data-testid="skeleton-row">
<TableRow
key={`skeleton-${rowKey}`}
className={cn("hover:bg-transparent", size === "compact" ? "h-8" : "")}
data-testid="skeleton-row"
>
{cells.map((column, columnKey) => (
<TableCell key={column?.id ?? columnKey} className={size === "compact" ? "px-2 py-1" : ""}>
<SkeletonCell column={column} index={columnKey} />

8
ui/litellm-dashboard/src/lib/http/schema.d.ts generated vendored
View file

@ -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 */