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

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

Browse source
This commit is contained in:
mateo-berri 2026-09-09 15:52:48 -07:00
commit 5fdbb2a1c8
42 changed files with 3224 additions and 284 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

194
litellm/llms/anthropic/chat/guardrail_translation/handler.py
View file

@ -13,10 +13,11 @@ Pattern Overview:
"""
import json
from collections.abc import Iterator, Mapping, Sequence
from collections.abc import Mapping, MutableSequence, Sequence
from copy import deepcopy
from dataclasses import dataclass
from itertools import chain, repeat
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Protocol, cast, overload, runtime_checkable
from typing_extensions import ReadOnly, TypedDict, assert_never
@ -41,6 +42,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import (
merge_guardrailed_scoped_messages,
merge_returned_tools_into_request_tools,
scoped_structured_message_indices,
stream_item_field,
stream_item_fingerprint,
)
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import (
@ -153,6 +155,46 @@ class ExtractedInput:
EMPTY_EXTRACTED_INPUT: Final = ExtractedInput(scanned=(), images=())
@dataclass(frozen=True, slots=True)
class _ToolCallShape:
name: str | None
arguments: str
@dataclass(frozen=True, slots=True)
class _SSEFieldRewrite:
"""One field of one nested section of a buffered SSE event, rewritten."""
section: str
field: str
value: object
class _SSEEventRewriter(Protocol):
def __call__(self, event: Mapping[str, object]) -> _SSEFieldRewrite | None: ...
def _rewritten_event(event: Mapping[str, object], rewrite_event: _SSEEventRewriter) -> Mapping[str, object]:
rewrite: Final = rewrite_event(event)
section: Final = None if rewrite is None else event.get(rewrite.section)
if rewrite is None or not isinstance(section, Mapping):
return event
return {**event, rewrite.section: {**section, rewrite.field: rewrite.value}} # mutable-ok: json.dumps needs a dict
def _tool_call_shapes(tool_calls: Sequence[object]) -> tuple[_ToolCallShape, ...]:
"""The guardrail-visible shape of each tool call, whether the guardrail handed
back the ``ChatCompletionMessageToolCall`` objects it was given or plain dicts."""
functions: Final = tuple(stream_item_field(tool_call, "function") for tool_call in tool_calls)
return tuple(
_ToolCallShape(
name=name if isinstance(name := stream_item_field(function, "name"), str) else None,
arguments=arguments if isinstance(arguments := stream_item_field(function, "arguments"), str) else "",
)
for function in functions
)
class _AnthropicSSEDelta(TypedDict, total=False):
type: ReadOnly[str]
text: ReadOnly[str]
@ -170,7 +212,7 @@ class AnthropicMessagesHandler(BaseTranslation):
them through guardrail rewrites; downstream provider handling is out of scope.
"""
delivers_ended_stream_text_rewrites = True
delivers_ended_stream_rewrites = True
assembles_streamed_response = True
def __init__(self):
@ -1056,6 +1098,7 @@ class AnthropicMessagesHandler(BaseTranslation):
first_choice.message.tool_calls,
)
string_so_far = first_choice.message.content
pre_guardrail_tool_calls: Final = _tool_call_shapes(tool_calls_list or ())
guardrail_inputs: Final = GenericGuardrailAPIInputs()
if string_so_far:
guardrail_inputs["texts"] = [string_so_far]
@ -1090,6 +1133,19 @@ class AnthropicMessagesHandler(BaseTranslation):
and guardrailed_texts[0] != string_so_far
):
self._write_ended_stream_text_rewrite(responses_so_far, guardrailed_texts[0])
if deliver_ended_stream_rewrites:
returned_tool_calls: Final = _guardrailed_inputs.get("tool_calls")
self._write_ended_stream_tool_call_rewrites(
responses_so_far,
pre_guardrail_tool_calls=pre_guardrail_tool_calls,
post_guardrail_tool_calls=_tool_call_shapes(
returned_tool_calls
if isinstance(returned_tool_calls, list)
and len(returned_tool_calls) == len(pre_guardrail_tool_calls)
else tool_calls_list or ()
),
guardrail_name=guardrail_to_apply.guardrail_name or "unknown",
)
else:
verbose_proxy_logger.debug("Skipping output guardrail - model response has no choices")
return responses_so_far
@ -1212,44 +1268,124 @@ class AnthropicMessagesHandler(BaseTranslation):
@staticmethod
def _write_ended_stream_text_rewrite(
responses_so_far: list[Any], # mutable-ok: rewrites the caller's buffered chunks in place
responses_so_far: MutableSequence[object], # mutable-ok: rewrites the caller's buffered chunks in place
rewritten_text: str,
) -> None:
"""Deliver an ended-stream guardrail text rewrite by rewriting the
buffered chunks in place: the first ``text_delta`` carries the full
rewritten text and every later one is blanked, leaving the surrounding
message and content-block framing untouched. Handles both chunk formats
this stream carries (parsed event dicts and raw SSE bytes)."""
message and content-block framing untouched."""
replacements: Final = chain((rewritten_text,), repeat(""))
for idx, item in enumerate(responses_so_far):
if isinstance(item, dict):
delta = item.get("delta")
if item.get("type") == "content_block_delta" and isinstance(delta, dict):
if delta.get("type") == "text_delta":
delta["text"] = next(replacements)
elif isinstance(item, (bytes, bytearray)):
responses_so_far[idx] = ( # rebind-ok: delivers the rewrite into the caller's buffer
AnthropicMessagesHandler._rewrite_sse_text_deltas(bytes(item), replacements)
)
def rewrite_text_delta(event: Mapping[str, object]) -> _SSEFieldRewrite | None:
delta: Final = event.get("delta")
if event.get("type") != "content_block_delta" or not isinstance(delta, Mapping):
return None
if delta.get("type") != "text_delta":
return None
return _SSEFieldRewrite("delta", "text", next(replacements))
AnthropicMessagesHandler._rewrite_ended_stream_events(responses_so_far, rewrite_text_delta)
@classmethod
def _write_ended_stream_tool_call_rewrites(
cls,
responses_so_far: MutableSequence[object], # mutable-ok: rewrites the caller's buffered chunks in place
*,
pre_guardrail_tool_calls: tuple[_ToolCallShape, ...],
post_guardrail_tool_calls: tuple[_ToolCallShape, ...],
guardrail_name: str,
) -> None:
"""Deliver ended-stream guardrail tool-call rewrites by rewriting the
buffered chunks in place: the rebuilt response lists tool calls in the
order of the stream's ``tool_use`` blocks, so the nth rewritten call lands
on the nth block, its first ``input_json_delta`` carrying the full rewritten
arguments, every later one blanked, and ``content_block_start`` carrying the
rewritten name. Blocks that do not line up with the rebuilt tool calls make
the rewrite undeliverable, so the pipeline executor discards it and releases
the original chunks."""
if post_guardrail_tool_calls == pre_guardrail_tool_calls:
return
block_indices: Final = tuple(
index
for item in responses_so_far
for event in cls._iter_sse_events(item)
if event.get("type") == "content_block_start"
and isinstance(block := event.get("content_block"), Mapping)
and block.get("type") == "tool_use"
and isinstance(index := event.get("index"), int)
)
if len(block_indices) != len(post_guardrail_tool_calls):
from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
raise UndeliverableStreamRewrite(guardrail_name)
rewrites_by_block: Final = MappingProxyType(
{
index: after
for index, before, after in zip(block_indices, pre_guardrail_tool_calls, post_guardrail_tool_calls)
if after != before
}
)
argument_replacements: Final = MappingProxyType(
{index: chain((rewrite.arguments,), repeat("")) for index, rewrite in rewrites_by_block.items()}
)
def rewrite_tool_use(event: Mapping[str, object]) -> _SSEFieldRewrite | None:
index: Final = event.get("index")
if not isinstance(index, int) or index not in rewrites_by_block:
return None
match event.get("type"):
case "content_block_start":
name: Final = rewrites_by_block[index].name
if name is None:
return None
return _SSEFieldRewrite("content_block", "name", name)
case "content_block_delta":
delta: Final = event.get("delta")
if not isinstance(delta, Mapping) or delta.get("type") != "input_json_delta":
return None
return _SSEFieldRewrite("delta", "partial_json", next(argument_replacements[index]))
case _:
return None
cls._rewrite_ended_stream_events(responses_so_far, rewrite_tool_use)
@staticmethod
def _rewrite_sse_text_deltas(sse_bytes: bytes, replacements: "Iterator[str]") -> bytes:
"""Rewrite every ``text_delta`` data line in one SSE chunk with the next
replacement text, leaving all other events and framing byte-identical."""
def _rewrite_ended_stream_events(
responses_so_far: MutableSequence[object], # mutable-ok: rewrites the caller's buffered chunks in place
rewrite_event: _SSEEventRewriter,
) -> None:
"""Replace every buffered event ``rewrite_event`` returns a rewrite for, in
both chunk formats this stream carries (parsed event dicts and raw SSE
bytes), leaving every other event and the framing untouched."""
rewritten_items: Final = tuple(
AnthropicMessagesHandler._rewrite_buffered_item(item, rewrite_event) for item in responses_so_far
)
responses_so_far[:] = rewritten_items # rebind-ok: delivers the rewrites into the caller's buffer
@staticmethod
def _rewrite_buffered_item(item: object, rewrite_event: _SSEEventRewriter) -> object:
if isinstance(item, dict):
return _rewritten_event(_as_str_mapping(item), rewrite_event)
if isinstance(item, (bytes, bytearray)):
return AnthropicMessagesHandler._rewrite_sse_events(bytes(item), rewrite_event)
return item
@staticmethod
def _rewrite_sse_events(sse_bytes: bytes, rewrite_event: _SSEEventRewriter) -> bytes:
"""Rewrite the data lines of one SSE chunk that ``rewrite_event`` rewrites,
leaving all other events and framing byte-identical."""
try:
decoded: Final = sse_bytes.decode("utf-8")
except UnicodeDecodeError:
return sse_bytes
return "\n\n".join(
AnthropicMessagesHandler._rewrite_sse_block(block, replacements) for block in decoded.split("\n\n")
"\n".join(AnthropicMessagesHandler._rewrite_sse_line(line, rewrite_event) for line in block.split("\n"))
for block in decoded.split("\n\n")
).encode("utf-8")
@staticmethod
def _rewrite_sse_block(block: str, replacements: "Iterator[str]") -> str:
return "\n".join(AnthropicMessagesHandler._rewrite_sse_line(line, replacements) for line in block.split("\n"))
@staticmethod
def _rewrite_sse_line(line: str, replacements: "Iterator[str]") -> str:
def _rewrite_sse_line(line: str, rewrite_event: _SSEEventRewriter) -> str:
if not line.startswith("data:"):
return line
try:
@ -1258,14 +1394,10 @@ class AnthropicMessagesHandler(BaseTranslation):
)
except json.JSONDecodeError:
return line
if not isinstance(data, dict) or data.get("type") != "content_block_delta":
if not isinstance(data, dict):
return line
delta: Final = data.get("delta")
if not isinstance(delta, dict) or delta.get("type") != "text_delta":
return line
return "data: " + json.dumps(
{**data, "delta": {**delta, "text": next(replacements)}} # mutable-ok: json.dumps needs plain dicts
)
rewritten: Final = _rewritten_event(_as_str_mapping(data), rewrite_event)
return line if rewritten is data else "data: " + json.dumps(rewritten)
def get_streaming_scan_key(self, responses_so_far: Sequence[object]) -> StreamingScanKey | None:
stream_ended: Final = self._check_streaming_has_ended(responses_so_far)

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

@ -52,13 +52,14 @@ class StreamingScanKey:
class BaseTranslation(ABC):
delivers_ended_stream_text_rewrites: ClassVar[bool] = False
delivers_ended_stream_rewrites: ClassVar[bool] = False
"""Whether ``process_output_streaming_response`` accepts
``deliver_ended_stream_rewrites=True`` and, on an ended (fully buffered)
stream, writes guardrail text rewrites back across ``responses_so_far`` so
a buffered pipeline can release rewritten chunks. Tool-call rewrites, and
text rewrites on every other translation, are undeliverable: the pipeline
executor discards them and releases the original chunks."""
stream, writes guardrail text and tool-call rewrites back across
``responses_so_far`` so a buffered pipeline can release rewritten chunks,
raising ``UndeliverableStreamRewrite`` for a shape it cannot place. Rewrites
on every other translation are undeliverable: the pipeline executor
discards them and releases the original chunks."""
assembles_streamed_response: ClassVar[bool] = False
"""Whether ``process_output_streaming_response`` stores the assembled response of an
@ -189,9 +190,9 @@ class BaseTranslation(ABC):
transformations (see ``StreamTransformSink``); base handlers ignore it.
``deliver_ended_stream_rewrites`` is passed True only when the caller
holds the whole buffered stream and the subclass declares
``delivers_ended_stream_text_rewrites``: the handler then writes
guardrail text rewrites back across ``responses_so_far`` instead of
discarding them.
``delivers_ended_stream_rewrites``: the handler then writes
guardrail text and tool-call rewrites back across ``responses_so_far``
instead of discarding them.
"""
return responses_so_far

9
litellm/llms/hosted_vllm/image_edit/__init__.py Normal file
View file

@ -0,0 +1,9 @@
from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig
from .transformation import HostedVLLMImageEditConfig
__all__ = ("HostedVLLMImageEditConfig",)
def get_hosted_vllm_image_edit_config(model: str) -> BaseImageEditConfig:
return HostedVLLMImageEditConfig()

43
litellm/llms/hosted_vllm/image_edit/transformation.py Normal file
View file

@ -0,0 +1,43 @@
from typing import Final
from litellm.llms.openai.image_edit.transformation import OpenAIImageEditConfig
from litellm.secret_managers.main import get_secret_str
PARAMS_VLLM_OMNI_DOES_NOT_ACCEPT: Final = frozenset({"mask", "quality", "input_fidelity"})
class HostedVLLMImageEditConfig(OpenAIImageEditConfig):
def get_supported_openai_params(self, model: str) -> list: # mutable-ok: BaseImageEditConfig contract
return [ # mutable-ok: BaseImageEditConfig returns list
param
for param in super().get_supported_openai_params(model)
if param not in PARAMS_VLLM_OMNI_DOES_NOT_ACCEPT
]
def validate_environment(
self,
headers: dict, # mutable-ok: BaseImageEditConfig contract
model: str,
api_key: str | None = None,
litellm_params: dict | None = None, # mutable-ok: BaseImageEditConfig contract
api_base: str | None = None,
) -> dict: # mutable-ok: BaseImageEditConfig contract
resolved_key: Final = api_key or get_secret_str("HOSTED_VLLM_API_KEY") or "fake-api-key"
return {**headers, "Authorization": f"Bearer {resolved_key}"} # mutable-ok: httpx headers are a dict
def get_complete_url(
self,
model: str,
api_base: str | None,
litellm_params: dict, # mutable-ok: BaseImageEditConfig contract
) -> str:
resolved_api_base: Final = api_base or get_secret_str("HOSTED_VLLM_API_BASE")
if resolved_api_base is None:
raise ValueError(
"api_base not set for Hosted VLLM images edits API. "
"Set via api_base parameter or HOSTED_VLLM_API_BASE environment variable"
)
trimmed: Final = resolved_api_base.rstrip("/")
if trimmed.endswith("/v1"):
return f"{trimmed}/images/edits"
return f"{trimmed}/v1/images/edits"

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

@ -49,6 +49,8 @@ from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import
coerce_stream_holdback_value,
)
from litellm.types.utils import (
ChatCompletionDeltaToolCall,
ChatCompletionMessageToolCall,
Choices,
GenericGuardrailAPIInputs,
ModelResponse,
@ -78,7 +80,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
Methods can be overridden to customize behavior for different message formats.
"""
delivers_ended_stream_text_rewrites = True
delivers_ended_stream_rewrites = True
assembles_streamed_response = True
def get_structured_messages(self, data: dict) -> list[AllMessageValues] | None:
@ -611,13 +613,14 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
deliver_ended_stream_rewrites: bool,
) -> None:
"""Ended-stream path: rebuild the full response, run the non-streaming
output guardrail against it, and (when opted in) write any text rewrite
back across the buffered chunks."""
output guardrail against it, and (when opted in) write any text or
tool-call rewrite back across the buffered chunks."""
model_response: Final = cast(
ModelResponse,
stream_chunk_builder(chunks=responses_so_far, logging_obj=litellm_logging_obj),
)
pre_guardrail_texts: Final = self._string_choice_contents(model_response)
pre_guardrail_tool_calls: Final = self._function_tool_call_shapes(model_response)
await self.process_output_response(
response=model_response,
guardrail_to_apply=guardrail_to_apply,
@ -625,13 +628,21 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
user_api_key_dict=user_api_key_dict,
request_data=request_data,
)
if deliver_ended_stream_rewrites:
await self._write_ended_stream_text_rewrites(
responses_so_far=responses_so_far,
guardrailed_response=model_response,
pre_guardrail_texts=pre_guardrail_texts,
guardrail_name=guardrail_to_apply.guardrail_name or "unknown",
)
if not deliver_ended_stream_rewrites:
return
guardrail_name: Final = guardrail_to_apply.guardrail_name or "unknown"
await self._write_ended_stream_text_rewrites(
responses_so_far=responses_so_far,
guardrailed_response=model_response,
pre_guardrail_texts=pre_guardrail_texts,
guardrail_name=guardrail_name,
)
self._write_ended_stream_tool_call_rewrites(
responses_so_far=responses_so_far,
guardrailed_response=model_response,
pre_guardrail_tool_calls=pre_guardrail_tool_calls,
guardrail_name=guardrail_name,
)
def build_stream_error_items(
self,
@ -1044,6 +1055,71 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
task_mappings=[(target_choice_index, None) for _ in changed], # mutable-ok: callee takes lists
)
@staticmethod
def _function_tool_call_shapes(response: "ModelResponse") -> tuple[tuple[str | None, str], ...]:
return tuple(
(tool_call.function.name, tool_call.function.arguments)
for choice in response.choices
for tool_call in choice.message.tool_calls or ()
if isinstance(tool_call, ChatCompletionMessageToolCall)
)
@staticmethod
def _function_tool_call_fragments(
responses_so_far: Sequence["ModelResponseStream"],
) -> tuple[tuple[ChatCompletionDeltaToolCall, ...], ...]:
"""Group the stream's function tool-call fragments by their tool-call index, in
the index order ``stream_chunk_builder`` lists the rebuilt tool calls, keeping
only the indices the builder keeps (an id and a name somewhere in the stream)."""
fragments: Final = tuple(
tool_call
for response in responses_so_far
for choice in response.choices
for tool_call in choice.delta.tool_calls or ()
if isinstance(tool_call, ChatCompletionDeltaToolCall)
)
identified: Final = frozenset(fragment.index for fragment in fragments if fragment.id)
named: Final = frozenset(fragment.index for fragment in fragments if fragment.function.name)
return tuple(
tuple(fragment for fragment in fragments if fragment.index == index) for index in sorted(identified & named)
)
def _write_ended_stream_tool_call_rewrites(
self,
responses_so_far: list["ModelResponseStream"], # mutable-ok: rewrites the caller's buffered chunks in place
guardrailed_response: "ModelResponse",
pre_guardrail_tool_calls: tuple[tuple[str | None, str], ...],
guardrail_name: str,
) -> None:
"""Write ended-stream guardrail tool-call rewrites back across the buffered
chunks: the rewritten name and full arguments land in the tool call's first
fragment and the arguments of its later fragments are blanked, mirroring the
text write-back. A rewrite on a stream carrying more than one distinct choice
index, or whose fragments do not line up with the rebuilt tool calls, is
reported as undeliverable, so the pipeline executor discards it and releases
the original chunks."""
post_guardrail_tool_calls: Final = self._function_tool_call_shapes(guardrailed_response)
if post_guardrail_tool_calls == pre_guardrail_tool_calls:
return
stream_choice_indices: Final = frozenset(
choice.index for response in responses_so_far for choice in response.choices
)
fragments_by_tool_call: Final = self._function_tool_call_fragments(responses_so_far)
if len(stream_choice_indices) != 1 or len(fragments_by_tool_call) != len(post_guardrail_tool_calls):
from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
raise UndeliverableStreamRewrite(guardrail_name)
for before, (name, arguments), fragments in zip(
pre_guardrail_tool_calls, post_guardrail_tool_calls, fragments_by_tool_call
):
if (name, arguments) == before:
continue
head, *tail = fragments
head.function.name = name
head.function.arguments = arguments
for fragment in tail:
fragment.function.arguments = ""
async def _apply_guardrail_responses_to_output_streaming(
self,
responses: list["ModelResponseStream"],

155
litellm/llms/openai/responses/guardrail_translation/handler.py
View file

@ -101,6 +101,18 @@ if TYPE_CHECKING:
from litellm.types.llms.openai import ResponseInputParam
class _ToolCallShape(NamedTuple):
name: str | None
arguments: str
def _tool_call_shapes(tool_calls: Sequence[ChatCompletionToolCallChunk]) -> tuple[_ToolCallShape, ...]:
return tuple(
_ToolCallShape(name=tool_call["function"].get("name"), arguments=tool_call["function"].get("arguments", ""))
for tool_call in tool_calls
)
class ResponseOutputEnvelope(TypedDict, total=False):
"""Dict form of a Responses API response, as far as guardrail write-back reads it."""
@ -128,6 +140,10 @@ _TERMINAL_ENVELOPE_EVENT_TYPES: Final = frozenset(
)
_FUNCTION_CALL_ARGUMENT_EVENT_TYPES: Final = frozenset(
{"response.function_call_arguments.delta", "response.function_call_arguments.done"}
)
_OUTPUT_ITEM_EVENT_TYPES: Final = frozenset({"response.output_item.added", "response.output_item.done"})
_PATCHABLE_ITEM_FIELDS: Final[Mapping[str, str]] = MappingProxyType(
{"function_call_output": "output", "message": "content"}
)
@ -340,7 +356,7 @@ class OpenAIResponsesHandler(BaseTranslation):
Methods can be overridden to customize behavior for different message formats.
"""
delivers_ended_stream_text_rewrites = True
delivers_ended_stream_rewrites = True
assembles_streamed_response = True
def get_structured_messages(self, data: dict) -> list[AllMessageValues] | None:
@ -755,6 +771,7 @@ class OpenAIResponsesHandler(BaseTranslation):
if response_model:
inputs["model"] = response_model
pre_guardrail_tool_calls: Final = _tool_call_shapes(tool_calls_to_check)
guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail(
inputs=inputs,
request_data=request_data,
@ -763,6 +780,12 @@ class OpenAIResponsesHandler(BaseTranslation):
)
guardrailed_texts: Final = guardrailed_inputs.get("texts", [])
returned_tool_calls: Final = guardrailed_inputs.get("tool_calls")
post_guardrail_tool_calls: Final = _tool_call_shapes(
returned_tool_calls
if isinstance(returned_tool_calls, list) and len(returned_tool_calls) == len(tool_calls_to_check)
else tool_calls_to_check
)
# Write guardrailed texts back into the output items in-place.
# final_chunk is a reference into responses_so_far so this
@ -785,6 +808,13 @@ class OpenAIResponsesHandler(BaseTranslation):
stream_events=responses_so_far[:-1],
rewrites_by_position=rewrites_by_position,
)
self._deliver_ended_stream_tool_call_rewrites(
responses_so_far=responses_so_far,
outputs=outputs,
pre_guardrail_tool_calls=pre_guardrail_tool_calls,
post_guardrail_tool_calls=post_guardrail_tool_calls,
guardrail_name=guardrail_to_apply.guardrail_name or "unknown",
)
return responses_so_far
# ------------------------------------------------------------------ #
@ -895,6 +925,129 @@ class OpenAIResponsesHandler(BaseTranslation):
continue
OpenAIResponsesHandler._write_event_field(content[content_idx], "text", rewritten)
def _deliver_ended_stream_tool_call_rewrites(
self,
responses_so_far: Sequence[object],
outputs: Sequence[object],
pre_guardrail_tool_calls: tuple[_ToolCallShape, ...],
post_guardrail_tool_calls: tuple[_ToolCallShape, ...],
guardrail_name: str,
) -> None:
"""Write ended-stream guardrail tool-call rewrites into the completed
envelope's ``function_call`` items and sync the earlier stream events,
keyed by ``call_id``. The guardrail sees the envelope's function calls
in output order, which is how a rewritten call finds its ``call_id``;
the stream events find their call through the ``call_id`` on
``output_item`` events and the ``item_id`` on argument events, since an
event's ``output_index`` need not match the envelope's (the chat bridge
numbers tool calls from 1 while the envelope lists them after the
message). A rewrite whose calls do not line up with the envelope, or
whose events cannot be found, is reported as undeliverable, so the
pipeline executor discards it and releases the original events."""
if post_guardrail_tool_calls == pre_guardrail_tool_calls:
return
function_call_items: Final = tuple(
output_item for output_item in outputs if stream_item_field(output_item, "type") == "function_call"
)
call_ids: Final = tuple(
call_id
for output_item in function_call_items
if isinstance(call_id := stream_item_field(output_item, "call_id"), str) and call_id
)
stream_events: Final = responses_so_far[:-1]
call_id_by_item_id: Final = self._function_call_ids_by_item_id(stream_events)
event_call_ids: Final = tuple(
self._function_call_event_call_id(event, call_id_by_item_id) for event in stream_events
)
rewrites_by_call_id: Final = MappingProxyType(
{
call_id: after
for call_id, before, after in zip(call_ids, pre_guardrail_tool_calls, post_guardrail_tool_calls)
if after != before
}
)
unresolved_argument_event: Final = any(
call_id is None and stream_item_field(event, "type") in _FUNCTION_CALL_ARGUMENT_EVENT_TYPES
for event, call_id in zip(stream_events, event_call_ids)
)
if (
len(call_ids) != len(function_call_items)
or len(frozenset(call_ids)) != len(call_ids)
or len(call_ids) != len(post_guardrail_tool_calls)
or unresolved_argument_event
or not rewrites_by_call_id.keys() <= frozenset(event_call_ids)
):
from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
raise UndeliverableStreamRewrite(guardrail_name)
for output_item, rewrite in (
(output_item, rewrites_by_call_id[call_id])
for output_item, call_id in zip(function_call_items, call_ids)
if call_id in rewrites_by_call_id
):
self._write_function_call_item(output_item, rewrite.name, rewrite.arguments)
delta_replacements: Final = MappingProxyType(
{call_id: chain((rewrite.arguments,), repeat("")) for call_id, rewrite in rewrites_by_call_id.items()}
)
for event, call_id in zip(stream_events, event_call_ids):
if call_id not in rewrites_by_call_id:
continue
match stream_item_field(event, "type"):
case "response.function_call_arguments.delta":
self._write_event_field(event, "delta", next(delta_replacements[call_id]))
case "response.function_call_arguments.done":
self._write_event_field(event, "arguments", rewrites_by_call_id[call_id].arguments)
case "response.output_item.added":
self._write_function_call_item(
stream_item_field(event, "item"), rewrites_by_call_id[call_id].name, None
)
case "response.output_item.done":
self._write_function_call_item(
stream_item_field(event, "item"),
rewrites_by_call_id[call_id].name,
rewrites_by_call_id[call_id].arguments,
)
case _:
pass
@staticmethod
def _function_call_ids_by_item_id(stream_events: Sequence[object]) -> Mapping[str, str]:
items: Final = tuple(
stream_item_field(event, "item")
for event in stream_events
if stream_item_field(event, "type") in _OUTPUT_ITEM_EVENT_TYPES
)
return MappingProxyType(
{
item_id: call_id
for item in items
if stream_item_field(item, "type") == "function_call"
and isinstance(item_id := stream_item_field(item, "id"), str)
and isinstance(call_id := stream_item_field(item, "call_id"), str)
}
)
@staticmethod
def _function_call_event_call_id(event: object, call_id_by_item_id: Mapping[str, str]) -> str | None:
event_type: Final = stream_item_field(event, "type")
if event_type in _FUNCTION_CALL_ARGUMENT_EVENT_TYPES:
item_id: Final = stream_item_field(event, "item_id")
return call_id_by_item_id.get(item_id) if isinstance(item_id, str) else None
if event_type not in _OUTPUT_ITEM_EVENT_TYPES:
return None
item: Final = stream_item_field(event, "item")
call_id: Final = stream_item_field(item, "call_id")
return call_id if stream_item_field(item, "type") == "function_call" and isinstance(call_id, str) else None
@staticmethod
def _write_function_call_item(item: object, name: str | None, arguments: str | None) -> None:
if item is None:
return
if name is not None:
OpenAIResponsesHandler._write_event_field(item, "name", name)
if arguments is not None:
OpenAIResponsesHandler._write_event_field(item, "arguments", arguments)
def _check_streaming_has_ended(self, responses_so_far: Sequence[object]) -> bool:
"""
Check if the streaming has ended.

12
litellm/proxy/_types.py
View file

@ -2836,6 +2836,18 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
"UI username/password login. Default is False."
),
)
disable_env_credential_login: bool | None = Field(
None,
description=(
"If True, disables signing in to the Admin UI with the environment credentials: "
"UI_USERNAME/UI_PASSWORD, or the master key when UI_PASSWORD is unset (that fallback "
"means env-credential login is always live by default). Database users with passwords "
"are unaffected. LOCKOUT RISK: create at least one proxy admin user with a password "
"before enabling, or nobody can sign in to the UI. A locked-out admin can still "
"administer the proxy over the API with the master key, and can unset this setting "
"and restart the proxy to restore env-credential login. Default is False."
),
)
disable_budget_reservation: bool | None = Field(
None,
description=(

39
litellm/proxy/auth/login_utils.py
View file

@ -85,6 +85,29 @@ def get_ui_credentials(master_key: str | None) -> tuple[str, str]:
return ui_username, ui_password
def _matches_env_credentials(username: str, password: str, master_key: str | None) -> bool:
ui_username, ui_password = get_ui_credentials(master_key)
return secrets.compare_digest(username.encode("utf-8"), ui_username.encode("utf-8")) and secrets.compare_digest(
password.encode("utf-8"), ui_password.encode("utf-8")
)
def is_env_credential_login_enabled(general_settings: Mapping[str, object]) -> bool:
"""Whether a login with UI_USERNAME/UI_PASSWORD (or the master-key fallback) can succeed.
Two settings can turn it off: `disable_env_credential_login` unconditionally, and
`disable_password_login_when_sso_enabled` as a side effect, since its gate rejects
every username/password login before the env comparison runs. Feeds both the
`authenticate_user` gate and the Admin UI warning banner, so the banner never nags
about a login path that is already unreachable.
"""
if general_settings.get("disable_env_credential_login") is True:
return False
if general_settings.get("disable_password_login_when_sso_enabled") is True and is_sso_provider_fully_configured():
return False
return True
class LoginResult:
"""Result object containing authentication data from login."""
@ -129,7 +152,8 @@ async def authenticate_user(
master_key: Master key for the proxy (required)
prisma_client: Prisma database client (optional)
general_settings: Proxy general_settings, checked for
`disable_password_login_when_sso_enabled`
`disable_password_login_when_sso_enabled` and
`disable_env_credential_login`
Returns:
LoginResult: Object containing authentication data
@ -170,8 +194,6 @@ async def authenticate_user(
code=500,
)
ui_username, ui_password = get_ui_credentials(master_key)
# Check if we can find the `username` in the db. On the UI, users can enter username=their email
_user_row: LiteLLM_UserTable | None = None
user_role: (
@ -197,8 +219,8 @@ async def authenticate_user(
- Login with UI_USERNAME and UI_PASSWORD
- Login with Invite Link `user_email` and `password` combination
"""
if secrets.compare_digest(username.encode("utf-8"), ui_username.encode("utf-8")) and secrets.compare_digest(
password.encode("utf-8"), ui_password.encode("utf-8")
if general_settings.get("disable_env_credential_login") is not True and _matches_env_credentials(
username, password, master_key
):
# Non SSO -> If user is using UI_USERNAME and UI_PASSWORD they are Proxy admin
user_role = LitellmUserRoles.PROXY_ADMIN
@ -340,8 +362,13 @@ async def authenticate_user(
code=401,
)
else:
env_credentials_hint: Final = (
"\nCheck 'UI_USERNAME', 'UI_PASSWORD' in .env file"
if is_env_credential_login_enabled(general_settings)
else ""
)
raise ProxyException(
message="Invalid credentials used to access UI.\nCheck 'UI_USERNAME', 'UI_PASSWORD' in .env file",
message=f"Invalid credentials used to access UI.{env_credentials_hint}",
type=ProxyErrorTypes.auth_error,
param="invalid_credentials",
code=401,

8
litellm/proxy/common_request_processing.py
View file

@ -190,6 +190,7 @@ from litellm.proxy.litellm_pre_call_utils import (
refresh_proxy_server_request_body_snapshot,
reject_url_valued_destination,
)
from litellm.proxy.policy_engine.response_retrieval import attach_post_call_pipelines_to_retrieval
from litellm.types.utils import (
ModelResponse,
ModelResponseStream,
@ -1849,7 +1850,6 @@ class ProxyBaseLLMRequestProcessing:
data=self.data,
user_api_key_dict=user_api_key_dict,
)
# Calculate request queue time after add_litellm_data_to_request
# which sets arrival_time in proxy_server_request. Ends at start_time
# (not a freshly captured time.time() here) so this window is exactly
@ -1997,6 +1997,12 @@ class ProxyBaseLLMRequestProcessing:
data=self.data,
call_type=route_type,
)
if route_type == "aget_responses":
attach_post_call_pipelines_to_retrieval(
data=self.data,
user_api_key_dict=user_api_key_dict,
llm_router=llm_router,
)
# Refresh AFTER pre_call_hook: guardrails (e.g. Presidio PII masking) may
# have mutated `self.data` in place, and the audit-trail snapshot taken in

10
litellm/proxy/health_endpoints/_health_endpoints.py
View file

@ -1582,6 +1582,13 @@ async def _show_no_redis_warning() -> bool:
return await count_live_proxy_workers(prisma_client) != 1
def _show_env_credential_login_warning() -> bool:
from litellm.proxy.auth.login_utils import is_env_credential_login_enabled
from litellm.proxy.proxy_server import general_settings
return is_env_credential_login_enabled(general_settings)
async def _get_health_readiness_details(
response: Response | None = None,
) -> dict[str, Any]:
@ -1623,6 +1630,7 @@ async def _get_health_readiness_details(
log_level_name: Final = logging.getLevelName(verbose_logger.getEffectiveLevel())
is_detailed_debug: Final = verbose_logger.isEnabledFor(logging.DEBUG)
show_no_redis_warning: Final = await _show_no_redis_warning()
show_env_credential_login_warning: Final = _show_env_credential_login_warning()
# check DB
if prisma_client is not None: # if db passed in, check if it's connected
@ -1650,6 +1658,7 @@ async def _get_health_readiness_details(
"log_level": log_level_name,
"is_detailed_debug": is_detailed_debug,
"show_no_redis_warning": show_no_redis_warning,
"show_env_credential_login_warning": show_env_credential_login_warning,
}
else:
return {
@ -1662,6 +1671,7 @@ async def _get_health_readiness_details(
"log_level": log_level_name,
"is_detailed_debug": is_detailed_debug,
"show_no_redis_warning": show_no_redis_warning,
"show_env_credential_login_warning": show_env_credential_login_warning,
}
except Exception as e:
raise HTTPException(status_code=503, detail=f"Service Unhealthy ({e})")

18
litellm/proxy/hooks/dynamic_rate_limiter_v3.py
View file

@ -32,6 +32,10 @@ from litellm.proxy.hooks.rate_limiter_utils import (
resolve_llm_provider_for_rate_limit,
)
from litellm.proxy.utils import InternalUsageCache
from litellm.router_utils.add_retry_fallback_headers import (
ensure_response_additional_headers,
response_has_hidden_params,
)
from litellm.types.router import ModelGroupInfo
from litellm.types.utils import CallTypesLiteral
@ -659,22 +663,12 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
data=data, user_api_key_dict=user_api_key_dict, response=response
)
# Add additional priority-specific headers
if isinstance(response, ModelResponse):
if response_has_hidden_params(response):
priority: Final = self._get_priority_from_user_api_key_dict(user_api_key_dict=user_api_key_dict)
# Get existing additional headers
additional_headers: Final = getattr(response, "_hidden_params", {}).get("additional_headers", {}) or {}
# Add priority information
additional_headers: Final = ensure_response_additional_headers(response)
additional_headers["x-litellm-priority"] = priority or "default"
additional_headers["x-litellm-rate-limiter-version"] = "v3"
# Update response
if not hasattr(response, "_hidden_params"):
response._hidden_params = {}
response._hidden_params["additional_headers"] = additional_headers
return response
except Exception as e:

33
litellm/proxy/hooks/parallel_request_limiter_v3.py
View file

@ -52,6 +52,10 @@ from litellm.proxy.hooks.batch_enqueued_tokens import (
canonical_provider_batch_id,
)
from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit
from litellm.router_utils.add_retry_fallback_headers import (
ensure_response_additional_headers,
response_has_hidden_params,
)
from litellm.types.caching import RedisPipelineIncrementOperation
from litellm.types.llms.openai import BaseLiteLLMOpenAIResponseObject, ResponseAPIUsage
from litellm.types.utils import (
@ -4677,34 +4681,17 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
Post-call hook to update rate limit headers in the response.
"""
try:
from pydantic import BaseModel
stash: Final = get_request_stash()
litellm_proxy_rate_limit_response: Final = stash.rate_limit_response if stash is not None else None
if litellm_proxy_rate_limit_response is not None:
# Update response headers
if hasattr(response, "_hidden_params"):
_hidden_params = getattr(response, "_hidden_params")
else:
_hidden_params = None
if _hidden_params is not None and (
isinstance(_hidden_params, BaseModel) or isinstance(_hidden_params, dict)
):
if isinstance(_hidden_params, BaseModel):
_hidden_params = _hidden_params.model_dump()
_additional_headers: Final = self._merge_ratelimit_statuses_into_additional_headers(
additional_headers=_hidden_params.get("additional_headers", {}) or {},
if litellm_proxy_rate_limit_response is not None and response_has_hidden_params(response):
additional_headers: Final = ensure_response_additional_headers(response)
additional_headers.update(
self._merge_ratelimit_statuses_into_additional_headers(
additional_headers={},
statuses=litellm_proxy_rate_limit_response["statuses"],
)
setattr(
response,
"_hidden_params",
{**_hidden_params, "additional_headers": _additional_headers},
)
)
except Exception as e:
verbose_proxy_logger.exception("Error in rate limit post-call hook: %s", e)

10
litellm/proxy/litellm_pre_call_utils.py
View file

@ -773,6 +773,16 @@ def apply_missing_session_id_policy(
return
if policy == "omit":
metadata[SESSION_ID_OMITTED_METADATA_KEY] = True
requester_metadata: Final = data.get("metadata")
requester_session_id: Final = (
requester_metadata.get("session_id") if isinstance(requester_metadata, dict) else None
)
if (
(body_session_id := data.get("litellm_session_id"))
and not metadata.get("session_id")
and not requester_session_id
):
metadata["session_id"] = body_session_id
return
if data.get("litellm_session_id") or metadata.get("session_id"):
return

46
litellm/proxy/policy_engine/pipeline_executor.py
View file

@ -82,6 +82,10 @@ def _rewrote(sent: tuple[object, ...] | None, returned: tuple[object, ...] | Non
return sent is not None and returned is not None and returned != sent
def _changed_count(sent: tuple[object, ...] | None, returned: tuple[object, ...] | None) -> bool:
return sent is not None and returned is not None and len(returned) != len(sent)
_GuardrailMethodT = TypeVar("_GuardrailMethodT", bound=Callable[..., object])
@ -93,10 +97,11 @@ def _logged_by_inner_guardrail(method: _GuardrailMethodT) -> _GuardrailMethodT:
class _StreamRewriteObserver(CustomGuardrail):
"""Stand-in handed to the endpoint translation in place of a streaming pipeline step's
guardrail. It records whether the guardrail returned different output than it was given,
which for guardrails like Bedrock's ANONYMIZED action is only known at runtime. Text
rewrites are deliverable on translations that write them back across the buffered chunks
(``delivers_ended_stream_text_rewrites``); tool-call rewrites and text rewrites on any
other translation are discarded by the executor, which releases the original chunks.
which for guardrails like Bedrock's ANONYMIZED action is only known at runtime. Text and
tool-call rewrites are deliverable on translations that write them back across the
buffered chunks (``delivers_ended_stream_rewrites``); rewrites on any other translation,
and a rewrite that drops or adds a tool call on any translation, are discarded by the
executor, which releases the original chunks.
The inner guardrail's ``apply_guardrail`` already records the guardrail information
and span, so the observer's stays out of ``log_guardrail_information``."""
@ -105,6 +110,7 @@ class _StreamRewriteObserver(CustomGuardrail):
self.inner: Final = inner
self.rewrote_texts = False
self.rewrote_tool_calls = False
self.changed_tool_call_count = False
def structured_messages_cover_full_request(self) -> bool:
return self.inner.structured_messages_cover_full_request()
@ -122,9 +128,11 @@ class _StreamRewriteObserver(CustomGuardrail):
outputs: Final = await self.inner.apply_guardrail(
inputs=inputs, request_data=request_data, input_type=input_type, logging_obj=logging_obj
)
returned_tool_shapes: Final = _tool_call_shapes(outputs.get("tool_calls"))
self.rewrote_texts = self.rewrote_texts or _rewrote(sent_texts, _text_snapshot(outputs.get("texts")))
self.rewrote_tool_calls = self.rewrote_tool_calls or _rewrote(
sent_tool_shapes, _tool_call_shapes(outputs.get("tool_calls"))
self.rewrote_tool_calls = self.rewrote_tool_calls or _rewrote(sent_tool_shapes, returned_tool_shapes)
self.changed_tool_call_count = self.changed_tool_call_count or _changed_count(
sent_tool_shapes, returned_tool_shapes
)
return outputs
@ -388,23 +396,23 @@ class PipelineExecutor:
litellm_logging_obj: "LiteLLMLoggingObj | None",
) -> None:
"""Run one streaming post_call step through the endpoint translation, delivering
text rewrites on translations that support ended-stream write-back. A guardrail
without the unified interface runs its legacy post-call hook against the assembled
response through ``_LegacyHookStreamAdapter``. A rewrite that cannot reach the client
yet (a tool-call rewrite, a text rewrite on a translation without write-back, or one
the translation or adapter refused with ``UndeliverableStreamRewrite``) is discarded:
the buffered chunks go back to the originals and the step passes, so the client gets
the stream the merge base sent, and the guardrail stays out of the applied-guardrails
header since its output never reached the client. The response an earlier step's translation stored under
``request_data["response"]`` is dropped first, so this step's hook sees the stream as
the steps before it left it."""
text and tool-call rewrites on translations that support ended-stream write-back. A
guardrail without the unified interface runs its legacy post-call hook against the
assembled response through ``_LegacyHookStreamAdapter``. A rewrite that cannot reach the
client yet (one on a translation without write-back, one that drops or adds a tool call,
or one the translation or adapter refused with ``UndeliverableStreamRewrite``) is
discarded: the buffered chunks go back to the originals and the step passes, so the
client gets the stream the merge base sent, and the guardrail stays out of the
applied-guardrails header since its output never reached the client. The response an
earlier step's translation stored under ``request_data["response"]`` is dropped first,
so this step's hook sees the stream as the steps before it left it."""
scanner: Final = (
callback
if PipelineExecutor.supports_unified_execution(callback)
else _LegacyHookStreamAdapter(callback, endpoint_translation, user_api_key_dict)
)
observer: Final = _StreamRewriteObserver(scanner)
deliver_rewrites: Final = type(endpoint_translation).delivers_ended_stream_text_rewrites
deliver_rewrites: Final = type(endpoint_translation).delivers_ended_stream_rewrites
originals: Final = copy.deepcopy(streaming_chunks)
hook_input.pop("response", None) # rebind-ok: an earlier step's stored response goes so this step's is stored
try:
@ -428,7 +436,9 @@ class PipelineExecutor:
except UndeliverableStreamRewrite:
_release_original_chunks(step.guardrail, streaming_chunks, originals)
return
if observer.rewrote_tool_calls or (observer.rewrote_texts and not deliver_rewrites):
if observer.changed_tool_call_count or (
not deliver_rewrites and (observer.rewrote_texts or observer.rewrote_tool_calls)
):
_release_original_chunks(step.guardrail, streaming_chunks, originals)
return
if not callback.records_own_guardrail_information:

152
litellm/proxy/policy_engine/response_retrieval.py Normal file
View file

@ -0,0 +1,152 @@
from collections.abc import Mapping
from dataclasses import dataclass
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, Literal, TypeAlias
from pydantic import TypeAdapter
from litellm._logging import verbose_proxy_logger
from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket
from litellm.proxy.common_utils.callback_utils import (
add_guardrail_to_applied_guardrails_header,
add_policy_sources_to_metadata,
add_policy_to_applied_policies_header,
)
from litellm.proxy.common_utils.http_parsing_utils import get_tags_from_request_body
from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry
from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher
from litellm.proxy.policy_engine.policy_registry import get_policy_registry
from litellm.proxy.policy_engine.policy_resolver import PolicyResolver
from litellm.responses.utils import ResponsesAPIRequestUtils
from litellm.router_utils.common_utils import resolve_model_group_alias
from litellm.types.proxy.policy_engine import PolicyMatchContext
from litellm.types.proxy.policy_engine.pipeline_types import GuardrailPipeline
if TYPE_CHECKING:
from litellm.proxy._types import UserAPIKeyAuth
from litellm.router import Router
PolicyPipelines: TypeAlias = tuple[tuple[str, GuardrailPipeline], ...]
_POLICY_PIPELINES_ADAPTER: Final = TypeAdapter(PolicyPipelines)
@dataclass(frozen=True, slots=True)
class UngovernedRetrieval:
reason: Literal["no router", "response id names no deployment", "deployment no longer in the router"]
def _model_group_for_response_id(response_id: object, llm_router: "Router | None") -> str | UngovernedRetrieval:
if llm_router is None:
return UngovernedRetrieval("no router")
model_id: Final = (
ResponsesAPIRequestUtils.get_model_id_from_response_id(response_id) if isinstance(response_id, str) else None
)
if model_id is None:
return UngovernedRetrieval("response id names no deployment")
deployment: Final = llm_router.get_deployment(model_id)
if deployment is None:
return UngovernedRetrieval("deployment no longer in the router")
hidden_by: Final = _submit_model_hidden_by(deployment.model_name, llm_router.model_group_alias)
if hidden_by is not None:
verbose_proxy_logger.warning(
"Policy engine: background response %s re-matches policies on retrieval as model group %s (%s), "
"so a policy attached to the model name it was submitted as does not run on it",
response_id,
deployment.model_name,
hidden_by,
)
return deployment.model_name
def _submit_model_hidden_by(model_group: str, model_group_alias: Mapping[str, object]) -> str | None:
if "*" in model_group:
return "a wildcard deployment"
aliases: Final = tuple(
alias for alias in model_group_alias if resolve_model_group_alias(model_group_alias, alias) == model_group
)
if not aliases:
return None
return f"the target of model_group_alias {', '.join(aliases)}"
def _retrieval_context(
data: Mapping[str, object], user_api_key_dict: "UserAPIKeyAuth", model_group: str
) -> PolicyMatchContext:
team_alias: Final = user_api_key_dict.team_alias
key_alias: Final = user_api_key_dict.key_alias
return PolicyMatchContext(
team_alias=team_alias if isinstance(team_alias, str) else None,
key_alias=key_alias if isinstance(key_alias, str) else None,
model=model_group,
tags=get_tags_from_request_body(data) or None,
)
def _post_call_pipelines_for_context(context: PolicyMatchContext) -> tuple[PolicyPipelines, Mapping[str, str]]:
matches: Final = get_attachment_registry().get_attached_policies_with_reasons(context)
if not matches:
return (), MappingProxyType({})
applied_policy_names: Final = PolicyMatcher.get_policies_with_matching_conditions(
policy_names=[match["policy_name"] for match in matches], # mutable-ok: the matcher takes a list
context=context,
)
post_call_pipelines: Final = tuple(
(policy_name, pipeline)
for policy_name, pipeline in PolicyResolver.resolve_pipelines_for_context(
context=context, policy_names=applied_policy_names
)
if pipeline.mode == "post_call"
)
return post_call_pipelines, MappingProxyType({match["policy_name"]: match["matched_via"] for match in matches})
def attach_post_call_pipelines_to_retrieval(
data: dict[str, object], # mutable-ok: request-state dict the policy engine hooks all write in place
user_api_key_dict: "UserAPIKeyAuth",
llm_router: "Router | None",
) -> None:
if not get_policy_registry().is_initialized():
return
model_group: Final = _model_group_for_response_id(data.get("response_id"), llm_router)
if isinstance(model_group, UngovernedRetrieval):
verbose_proxy_logger.warning(
"Policy engine: background response %s is retrieved without its post_call policy pipelines (%s)",
data.get("response_id"),
model_group.reason,
)
return
context: Final = _retrieval_context(data, user_api_key_dict, model_group)
post_call_pipelines, policy_sources = _post_call_pipelines_for_context(context)
_, bucket = get_or_create_metadata_bucket(data)
already_attached: Final = _POLICY_PIPELINES_ADAPTER.validate_python(bucket.get("_guardrail_pipelines") or ())
attached_policy_names: Final = frozenset(policy_name for policy_name, _pipeline in already_attached)
added: Final = tuple(
(policy_name, pipeline)
for policy_name, pipeline in post_call_pipelines
if policy_name not in attached_policy_names
)
if not added:
return
pipelines: Final = (*already_attached, *added)
bucket["_guardrail_pipelines"] = pipelines
bucket["_pipeline_managed_guardrails"] = frozenset(
step.guardrail for _policy_name, pipeline in pipelines for step in pipeline.steps
)
for policy_name, _pipeline in added:
add_policy_to_applied_policies_header(request_data=data, policy_name=policy_name)
for _policy_name, pipeline in added:
for step in pipeline.steps:
add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=step.guardrail)
add_policy_sources_to_metadata(
request_data=data,
policy_sources={ # mutable-ok: add_policy_sources_to_metadata takes a dict
policy_name: policy_sources[policy_name] for policy_name, _pipeline in added
},
)
verbose_proxy_logger.debug(
"Policy engine: attached post_call pipelines to the retrieval of background response %s (model group %s): %s",
data.get("response_id"),
model_group,
", ".join(policy_name for policy_name, _pipeline in added),
)

202
litellm/proxy/utils.py
View file

@ -94,6 +94,7 @@ from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting
from litellm.integrations.SlackAlerting.utils import _add_langfuse_trace_id_to_alert
from litellm.litellm_core_utils.core_helpers import (
coerce_token_limit,
get_or_create_metadata_bucket,
independent_snapshot,
is_expected_client_error,
)
@ -157,6 +158,8 @@ from litellm.proxy.hooks.sensitive_data_routing import (
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup, add_guardrails_from_auth_metadata
from litellm.proxy.management_helpers.key_settings_audit import with_settings_updated_at
from litellm.proxy.policy_engine.pipeline_executor import PipelineExecutor
from litellm.proxy.policy_engine.policy_registry import get_policy_registry
from litellm.proxy.policy_engine.policy_resolver import PolicyResolver
from litellm.repositories.budget_repository import BudgetRepository
from litellm.repositories.config_repository import ConfigRepository
from litellm.repositories.table_repositories import (
@ -172,6 +175,7 @@ from litellm.repositories.verification_token_repository import (
)
from litellm.secret_managers.main import str_to_bool
from litellm.types.integrations.slack_alerting import DEFAULT_ALERT_TYPES
from litellm.types.llms.openai import ResponsesAPIResponse
from litellm.types.mcp import (
MCPDuringCallResponseObject,
MCPPreCallRequestObject,
@ -538,17 +542,127 @@ def _post_call_pipelines(data: Mapping[str, object]) -> tuple[tuple[str, "Guardr
)
def _warn_background_skips_post_call_pipelines(data: Mapping[str, object]) -> None:
if data.get("background") is not True:
return
policy_names: Final = tuple(policy_name for policy_name, _pipeline in _post_call_pipelines(data))
if not policy_names:
return
verbose_proxy_logger.warning(
"Policies with post_call guardrail pipelines do not run on background responses yet; "
"the response is released ungoverned by them: %s",
", ".join(policy_names),
_PENDING_BACKGROUND_RESPONSE_STATUSES: Final = frozenset(("queued", "in_progress"))
def _is_pending_background_response(response: LLMResponseTypes) -> bool:
return isinstance(response, ResponsesAPIResponse) and response.status in _PENDING_BACKGROUND_RESPONSE_STATUSES
def _guardrails_outside_pipeline(policy_name: str, pipeline: "GuardrailPipeline") -> frozenset[str]:
resolved: Final = PolicyResolver.resolve_policy_guardrails(
policy_name=policy_name, policies=get_policy_registry().get_all_policies()
)
return frozenset(resolved.guardrails) - frozenset(step.guardrail for step in pipeline.steps)
def _guardrails_run_standalone_pre_call(data: Mapping[str, object]) -> frozenset[str]:
return frozenset(
callback.guardrail_name
for callback in litellm.callbacks
if isinstance(callback, CustomGuardrail)
and callback.guardrail_name is not None
and callback.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_call)
)
def _without_names(
bucket: dict[str, object], # mutable-ok: the applied_* header slots live in the request-state dict hooks write
slot: str,
names: frozenset[str],
) -> None:
claimed: Final = bucket.get(slot)
if not isinstance(claimed, list):
return
remaining: Final = [ # mutable-ok: the slot stays a list, the shape every applied_* header writer appends to
name for name in claimed if name not in names
]
if remaining:
bucket[slot] = remaining # rebind-ok: the slot lives in the shared request-state dict, rewritten in place
else:
bucket.pop(slot)
def _withdraw_deferred_claims(
data: dict[str, object], # mutable-ok: same request-payload shape as post_call_success_hook's data
deferred: Sequence[tuple[str, "GuardrailPipeline"]],
) -> None:
outside_by_policy: Final = MappingProxyType(
{policy_name: _guardrails_outside_pipeline(policy_name, pipeline) for policy_name, pipeline in deferred}
)
running_elsewhere: Final = pipeline_managed_guardrail_names(data, "pre_call").union(
_guardrails_run_standalone_pre_call(data), *outside_by_policy.values()
)
withdrawn_policies: Final = frozenset(name for name, outside in outside_by_policy.items() if not outside)
withdrawn_guardrails: Final = _pipeline_step_guardrail_names(deferred) - running_elsewhere
_, bucket = get_or_create_metadata_bucket(data)
_without_names(bucket, "applied_policies", withdrawn_policies)
_without_names(bucket, "applied_guardrails", withdrawn_guardrails)
sources: Final = bucket.get("policy_sources")
if not isinstance(sources, dict):
return
remaining_sources: Final = { # mutable-ok: policy_sources stays a dict, the shape its writer updates in place
name: reason for name, reason in sources.items() if name not in withdrawn_policies
}
if remaining_sources:
bucket["policy_sources"] = remaining_sources
else:
bucket.pop("policy_sources")
def _defer_post_call_pipelines(
data: dict[str, object], # mutable-ok: same request-payload shape as post_call_success_hook's data
response: ResponsesAPIResponse,
) -> None:
deferred: Final = _post_call_pipelines(data)
if not deferred:
return
verbose_proxy_logger.debug(
"Post_call guardrail pipelines wait for background response %s (status=%s) to be retrieved complete: %s",
response.id,
response.status,
", ".join(policy_name for policy_name, _pipeline in deferred),
)
tag_matched: Final = _tag_matched_deferrals(data, deferred)
if tag_matched:
verbose_proxy_logger.warning(
"Policy engine: background response %s matched post_call policies through a request tag at submit; "
"retrieval re-matches only the key, team, and model scopes, so a tag carried in the request body "
"does not govern the completed response: %s",
response.id,
", ".join(tag_matched),
)
body_selected: Final = _body_selected_deferrals(data, deferred)
if body_selected:
verbose_proxy_logger.warning(
"Policy engine: background response %s matched post_call policies through the request body's policies "
"list at submit; retrieval carries no request body, so those policies do not govern the completed "
"response: %s",
response.id,
", ".join(body_selected),
)
_withdraw_deferred_claims(data, deferred)
def _tag_matched_deferrals(
data: Mapping[str, object], deferred: Sequence[tuple[str, "GuardrailPipeline"]]
) -> tuple[str, ...]:
sources: Final = _policy_state_metadata(data).get("policy_sources")
if not isinstance(sources, dict):
return ()
return tuple(
policy_name
for policy_name, _pipeline in deferred
if policy_name in sources and "tag:" in str(sources[policy_name])
)
def _body_selected_deferrals(
data: Mapping[str, object], deferred: Sequence[tuple[str, "GuardrailPipeline"]]
) -> tuple[str, ...]:
sources: Final = _policy_state_metadata(data).get("policy_sources")
attributed: Final = frozenset(sources) if isinstance(sources, dict) else frozenset()
return tuple(policy_name for policy_name, _pipeline in deferred if policy_name not in attributed)
def _pipeline_unsupported_streaming_guardrails(
@ -585,17 +699,11 @@ def _streaming_pipeline_translation(user_api_key_dict: UserAPIKeyAuth) -> "BaseT
return None if resolved is None else resolved[1]
def _route_supports_streaming_pipelines(
user_api_key_dict: UserAPIKeyAuth, translation: "BaseTranslation | None"
) -> bool:
return not user_api_key_dict.request_route or translation is not None
def stream_gated_guardrail_names(
request_data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth
) -> frozenset[str]:
translation: Final = _streaming_pipeline_translation(user_api_key_dict)
if not _route_supports_streaming_pipelines(user_api_key_dict, translation):
if translation is None:
return frozenset()
return _pipeline_step_guardrail_names(
tuple(
@ -625,7 +733,7 @@ def _streamable_post_call_pipelines(
if not post_call_pipelines:
return ()
translation: Final = _streaming_pipeline_translation(user_api_key_dict)
if not _route_supports_streaming_pipelines(user_api_key_dict, translation):
if translation is None:
verbose_proxy_logger.warning(
"Policies with post_call guardrail pipelines cannot scan streaming responses on route %s yet "
"(no endpoint guardrail translation); the stream skips the pipelines and their guardrails run "
@ -2017,8 +2125,6 @@ class ProxyLogging:
)
try:
_warn_background_skips_post_call_pipelines(data)
# Execute guardrail pipelines before the normal callback loop
data, _ = await self._maybe_execute_pipelines( # rebind-ok: pipeline edits feed the callback loop below
data=data,
@ -2988,6 +3094,24 @@ class ProxyLogging:
daemon=True,
).start()
async def _run_post_call_pipelines(
self,
data: dict[str, object], # mutable-ok: same request-payload shape as post_call_success_hook's data
user_api_key_dict: UserAPIKeyAuth,
response: LLMResponseTypes,
) -> LLMResponseTypes | None:
if _is_pending_background_response(response):
_defer_post_call_pipelines(data, response)
return None
_, pipeline_response = await self._maybe_execute_pipelines(
data=data,
user_api_key_dict=user_api_key_dict,
call_type=getattr(data.get("litellm_logging_obj"), "call_type", None) or "acompletion",
event_hook="post_call",
response=response,
)
return pipeline_response
async def post_call_success_hook(
self,
data: dict,
@ -3007,11 +3131,9 @@ class ProxyLogging:
from litellm.proxy.proxy_server import llm_router
from litellm.types.guardrails import GuardrailEventHooks
_, pipeline_response = await self._maybe_execute_pipelines(
pipeline_response: Final = await self._run_post_call_pipelines(
data=data,
user_api_key_dict=user_api_key_dict,
call_type=getattr(data.get("litellm_logging_obj"), "call_type", None) or "acompletion",
event_hook="post_call",
response=response,
)
if pipeline_response is not None:
@ -3468,12 +3590,16 @@ class ProxyLogging:
),
)
if post_call_pipelines:
pipeline_translation: Final = (
resolve_endpoint_translation(user_api_key_dict, None) if post_call_pipelines else None
)
if pipeline_translation is not None:
current_response = self._pipeline_gated_stream(
response=current_response,
user_api_key_dict=user_api_key_dict,
request_data=request_data,
pipelines=post_call_pipelines,
translation=pipeline_translation,
)
try:
@ -3497,6 +3623,7 @@ class ProxyLogging:
user_api_key_dict: UserAPIKeyAuth,
request_data: dict, # mutable-ok: same request-payload shape the hooks mutate
pipelines: "tuple[tuple[str, GuardrailPipeline], ...]",
translation: "tuple[str, BaseTranslation]",
) -> "AsyncGenerator[Any, None]":
"""
Execute post_call policy pipelines against a streamed response.
@ -3506,14 +3633,13 @@ class ProxyLogging:
assembled output through the endpoint guardrail translation, the same
machinery flat post_call guardrails use at end of stream. An allow
releases the buffered chunks: verbatim when no guardrail rewrote the
output, rewritten in place when one rewrote text and the translation
delivers ended-stream rewrites (later steps then re-scan the rewritten
chunks, so rewrites chain). A rewrite the translation cannot deliver
yet (a tool-call rewrite, or a text rewrite on a route without
write-back) is discarded by the executor and the original chunks are
released, as is a buffered shape no translation resolves; a block or
modify_response terminates with the translation's block chunks or the
raised error.
output, rewritten in place when one rewrote text or a tool call and the
translation delivers ended-stream rewrites (later steps then re-scan the
rewritten chunks, so rewrites chain). A rewrite the translation cannot
deliver yet (one on a route without write-back, or a shape the route
refuses) is discarded by the executor and the original chunks are
released; a block or modify_response terminates with the translation's
block chunks or the raised error.
"""
buffered: Final[list[object]] = [] # mutable-ok: accumulates the stream before the pipeline verdict
async for item in response:
@ -3521,17 +3647,7 @@ class ProxyLogging:
if not buffered:
return
resolved: Final = resolve_endpoint_translation(user_api_key_dict, buffered[0])
if resolved is None:
verbose_proxy_logger.warning(
"Policies with post_call guardrail pipelines cannot scan this streaming response shape yet; "
"the stream is released ungoverned by them: %s",
", ".join(policy_name for policy_name, _pipeline in pipelines),
)
for buffered_item in buffered:
yield buffered_item
return
call_type, endpoint_translation = resolved
call_type, endpoint_translation = translation
for policy_name, pipeline in pipelines:
result: PipelineExecutionResult = await PipelineExecutor.execute_steps(

6
litellm/router_utils/add_retry_fallback_headers.py
View file

@ -50,6 +50,12 @@ def prepare_response_for_header_attachment(response: object) -> object | None:
return response
def response_has_hidden_params(response: object) -> bool:
if isinstance(response, dict):
return "_hidden_params" in response
return hasattr(response, "_hidden_params")
def ensure_response_additional_headers(response: object) -> dict[str, object]:
hidden_params: Final = get_hidden_params_dict(response, create=isinstance(response, dict))
_write_hidden_params(response, hidden_params)

4
litellm/utils.py
View file

@ -9241,6 +9241,10 @@ class ProviderConfigManager:
from litellm.llms.openai.image_edit import get_openai_image_edit_config
return get_openai_image_edit_config(model=model)
elif LlmProviders.HOSTED_VLLM == provider:
from litellm.llms.hosted_vllm.image_edit import get_hosted_vllm_image_edit_config
return get_hosted_vllm_image_edit_config(model=model)
elif LlmProviders.AZURE == provider:
from litellm.llms.azure.image_edit.transformation import (
AzureImageEditConfig,

114
tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py
View file

@ -315,6 +315,120 @@ class TestAnthropicMessagesHandlerStreamingOutputProcessing:
assert "event: message_start" in raw and "event: message_stop" in raw
assert '"stop_reason": "end_turn"' in raw
@staticmethod
def _ended_tool_use_sse_chunks() -> list:
events = [
("message_start", {"type": "message_start", "message": {"id": "msg_1", "type": "message", "role": "assistant", "model": "claude-sonnet-4-5", "content": [], "stop_reason": None, "usage": {"input_tokens": 1, "output_tokens": 0}}}),
("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "tool_use", "id": "toolu_1", "name": "lookup_fruit", "input": {}}}),
("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": ""}}),
("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": '{"fruit": "persim'}}),
("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": 'mon"}'}}),
("content_block_stop", {"type": "content_block_stop", "index": 0}),
("message_delta", {"type": "message_delta", "delta": {"stop_reason": "tool_use", "stop_sequence": None}, "usage": {"output_tokens": 2}}),
("message_stop", {"type": "message_stop"}),
]
return [f"event: {name}\ndata: {json.dumps(payload)}\n\n".encode() for name, payload in events]
@staticmethod
def _argument_masking_guardrail() -> CustomGuardrail:
class MaskArguments(CustomGuardrail):
async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
for tool_call in inputs.get("tool_calls", []):
tool_call.function.arguments = '{"fruit": "[MASKED]"}'
return inputs
return MaskArguments(guardrail_name="test")
@staticmethod
def _partial_jsons(chunks: list) -> list:
return [
json.loads(line[len("data:") :].strip())["delta"]["partial_json"]
for chunk in chunks
for line in chunk.decode().split("\n")
if line.startswith("data:") and json.loads(line[len("data:") :].strip()).get("type") == "content_block_delta"
]
@pytest.mark.asyncio
async def test_deliver_ended_stream_rewrites_writes_tool_use_input_back_into_sse_chunks(self):
handler = AnthropicMessagesHandler()
chunks = self._ended_tool_use_sse_chunks()
result = await handler.process_output_streaming_response(
responses_so_far=chunks,
guardrail_to_apply=self._argument_masking_guardrail(),
litellm_logging_obj=MagicMock(),
deliver_ended_stream_rewrites=True,
)
assert result is chunks
assert self._partial_jsons(chunks) == ['{"fruit": "[MASKED]"}', "", ""]
raw = b"".join(chunks).decode()
assert '"name": "lookup_fruit"' in raw and '"id": "toolu_1"' in raw
assert '"stop_reason": "tool_use"' in raw
assert "persim" not in raw
@pytest.mark.asyncio
async def test_deliver_ended_stream_rewrites_writes_tool_use_name_back_into_sse_chunks(self):
class RenameTool(CustomGuardrail):
async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
for tool_call in inputs.get("tool_calls", []):
tool_call.function.name = "lookup_fruit_reviewed"
return inputs
handler = AnthropicMessagesHandler()
chunks = self._ended_tool_use_sse_chunks()
await handler.process_output_streaming_response(
responses_so_far=chunks,
guardrail_to_apply=RenameTool(guardrail_name="test"),
litellm_logging_obj=MagicMock(),
deliver_ended_stream_rewrites=True,
)
raw = b"".join(chunks).decode()
assert '"name": "lookup_fruit_reviewed"' in raw and '"id": "toolu_1"' in raw
assert '"name": "lookup_fruit"' not in raw
assert json.loads("".join(self._partial_jsons(chunks))) == {"fruit": "persimmon"}
@pytest.mark.asyncio
async def test_ended_stream_tool_use_rewrite_leaves_chunks_untouched_by_default(self):
handler = AnthropicMessagesHandler()
chunks = self._ended_tool_use_sse_chunks()
original = [bytes(chunk) for chunk in chunks]
await handler.process_output_streaming_response(
responses_so_far=chunks,
guardrail_to_apply=self._argument_masking_guardrail(),
litellm_logging_obj=MagicMock(),
)
assert chunks == original
@pytest.mark.asyncio
async def test_deliver_ended_stream_tool_use_rewrite_with_server_tool_use_block_fails_closed(self):
from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
handler = AnthropicMessagesHandler()
server_tool_use = [
("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "server_tool_use", "id": "srvtoolu_1", "name": "web_search", "input": {}}}),
("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": '{"query": "fruit"}'}}),
("content_block_stop", {"type": "content_block_stop", "index": 0}),
]
tool_use = self._ended_tool_use_sse_chunks()
chunks = (
tool_use[:1]
+ [f"event: {name}\ndata: {json.dumps(payload)}\n\n".encode() for name, payload in server_tool_use]
+ [chunk.replace(b'"index": 0', b'"index": 1') for chunk in tool_use[1:]]
)
with pytest.raises(UndeliverableStreamRewrite):
await handler.process_output_streaming_response(
responses_so_far=chunks,
guardrail_to_apply=self._argument_masking_guardrail(),
litellm_logging_obj=MagicMock(),
deliver_ended_stream_rewrites=True,
)
@pytest.mark.asyncio
async def test_ended_stream_rewrite_leaves_chunks_untouched_by_default(self):
handler = AnthropicMessagesHandler()

155
tests/test_litellm/llms/hosted_vllm/image_edit/test_hosted_vllm_image_edit_transformation.py Normal file
View file

@ -0,0 +1,155 @@
import httpx
import pytest
import litellm
from litellm.llms.custom_httpx.http_handler import HTTPHandler
from litellm.llms.hosted_vllm.image_edit import get_hosted_vllm_image_edit_config
from litellm.llms.hosted_vllm.image_edit.transformation import HostedVLLMImageEditConfig
from litellm.types.utils import LlmProviders
from litellm.utils import ProviderConfigManager
PNG_BYTES = b"\x89PNG\r\n\x1a\nfakepng"
MODEL = "Qwen/Qwen-Image-Edit-2511"
@pytest.fixture(autouse=True)
def _clear_hosted_vllm_env(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.delenv("HOSTED_VLLM_API_KEY", raising=False)
monkeypatch.delenv("HOSTED_VLLM_API_BASE", raising=False)
def test_provider_config_registration():
config = ProviderConfigManager.get_provider_image_edit_config(
model=f"hosted_vllm/{MODEL}",
provider=LlmProviders.HOSTED_VLLM,
)
assert isinstance(config, HostedVLLMImageEditConfig)
assert isinstance(get_hosted_vllm_image_edit_config(MODEL), HostedVLLMImageEditConfig)
@pytest.mark.parametrize(
"api_base",
["http://localhost:8091", "http://localhost:8091/", "http://localhost:8091/v1", "http://localhost:8091/v1/"],
)
def test_get_complete_url_appends_images_edits(api_base: str):
config = HostedVLLMImageEditConfig()
assert (
config.get_complete_url(model=MODEL, api_base=api_base, litellm_params={})
== "http://localhost:8091/v1/images/edits"
)
def test_get_complete_url_falls_back_to_env(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setenv("HOSTED_VLLM_API_BASE", "http://vllm-omni:8000/v1")
config = HostedVLLMImageEditConfig()
assert (
config.get_complete_url(model=MODEL, api_base=None, litellm_params={})
== "http://vllm-omni:8000/v1/images/edits"
)
def test_get_complete_url_requires_api_base():
config = HostedVLLMImageEditConfig()
with pytest.raises(ValueError, match="api_base not set"):
config.get_complete_url(model=MODEL, api_base=None, litellm_params={})
def test_validate_environment_defaults_to_fake_api_key():
headers = HostedVLLMImageEditConfig().validate_environment(headers={}, model=MODEL)
assert headers == {"Authorization": "Bearer fake-api-key"}
def test_validate_environment_uses_provided_api_key_and_keeps_headers():
headers = HostedVLLMImageEditConfig().validate_environment(
headers={"X-Test": "1"},
model=MODEL,
api_key="my-custom-key",
)
assert headers == {"X-Test": "1", "Authorization": "Bearer my-custom-key"}
def test_validate_environment_falls_back_to_env_api_key(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setenv("HOSTED_VLLM_API_KEY", "env-key")
headers = HostedVLLMImageEditConfig().validate_environment(headers={}, model=MODEL)
assert headers["Authorization"] == "Bearer env-key"
def test_image_edit_posts_multipart_to_vllm_omni():
captured: list[httpx.Request] = []
def handler(request: httpx.Request) -> httpx.Response:
captured.append(request)
return httpx.Response(200, json={"created": 1712697600, "data": [{"b64_json": "aW1n"}]})
response = litellm.image_edit(
model=f"hosted_vllm/{MODEL}",
image=PNG_BYTES,
prompt="add a hat",
api_base="http://localhost:8091",
api_key="test-key",
client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(handler))),
seed=42,
)
assert response.data
assert len(captured) == 1
request = captured[0]
assert str(request.url) == "http://localhost:8091/v1/images/edits"
assert request.headers["authorization"] == "Bearer test-key"
assert request.headers["content-type"].startswith("multipart/form-data")
assert b'name="image[]"' in request.content
assert PNG_BYTES in request.content
assert f'name="model"\r\n\r\n{MODEL}'.encode() in request.content
assert b'name="prompt"\r\n\r\nadd a hat' in request.content
assert b'name="seed"\r\n\r\n42' in request.content
@pytest.mark.parametrize("param", ["mask", "quality", "input_fidelity"])
def test_params_vllm_omni_ignores_are_not_advertised(param: str):
supported = HostedVLLMImageEditConfig().get_supported_openai_params(MODEL)
assert param not in supported
assert {"image", "prompt", "n", "size", "response_format", "background", "user"} <= set(supported)
def test_image_edit_rejects_quality_unless_dropped():
captured: list[httpx.Request] = []
def handler(request: httpx.Request) -> httpx.Response:
captured.append(request)
return httpx.Response(200, json={"created": 1712697600, "data": [{"b64_json": "aW1n"}]})
client = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(handler)))
with pytest.raises(litellm.UnsupportedParamsError, match="quality"):
litellm.image_edit(
model=f"hosted_vllm/{MODEL}",
image=PNG_BYTES,
prompt="add a hat",
api_base="http://localhost:8091",
client=client,
quality="low",
)
assert captured == []
litellm.image_edit(
model=f"hosted_vllm/{MODEL}",
image=PNG_BYTES,
prompt="add a hat",
api_base="http://localhost:8091",
client=client,
quality="low",
drop_params=True,
)
assert len(captured) == 1
assert b'name="quality"' not in captured[0].content
assert b'name="prompt"\r\n\r\nadd a hat' in captured[0].content

152
tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py
View file

@ -1113,6 +1113,102 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput:
assert chunks[1].choices[0].delta.content in (None, "")
assert chunks[1].choices[0].finish_reason == "stop"
@staticmethod
def _ended_tool_call_stream_chunks() -> list:
from litellm.types.utils import (
ChatCompletionDeltaToolCall,
Delta,
Function,
ModelResponseStream,
StreamingChoices,
)
def chunk(tool_call: ChatCompletionDeltaToolCall | None, finish_reason: Optional[str] = None):
return ModelResponseStream(
id="chatcmpl-123",
created=1234567890,
model="gpt-4",
object="chat.completion.chunk",
choices=[
StreamingChoices(
index=0,
delta=Delta(tool_calls=[tool_call] if tool_call else None),
finish_reason=finish_reason,
)
],
)
def fragment(arguments: str, name: Optional[str] = None, call_id: Optional[str] = None):
return ChatCompletionDeltaToolCall(
id=call_id, index=0, type="function", function=Function(name=name, arguments=arguments)
)
return [
chunk(fragment("", name="lookup_fruit", call_id="call_1")),
chunk(fragment('{"fruit":')),
chunk(fragment(' "persimmon"}')),
chunk(None, finish_reason="tool_calls"),
]
@pytest.mark.asyncio
async def test_deliver_ended_stream_rewrites_writes_tool_call_arguments_back_into_chunks(self):
handler = OpenAIChatCompletionsHandler()
guardrail = MockGuardrail(guardrail_name="test")
chunks = self._ended_tool_call_stream_chunks()
result = await handler.process_output_streaming_response(
responses_so_far=chunks,
guardrail_to_apply=guardrail,
litellm_logging_obj=None,
deliver_ended_stream_rewrites=True,
)
assert result is chunks
fragments = [chunk.choices[0].delta.tool_calls for chunk in chunks[:3]]
assert [fragment[0].function.arguments for fragment in fragments] == ['{"fruit": "PERSIMMON"}', "", ""]
assert fragments[0][0].function.name == "lookup_fruit"
assert fragments[0][0].id == "call_1"
assert chunks[3].choices[0].delta.tool_calls is None
assert chunks[3].choices[0].finish_reason == "tool_calls"
@pytest.mark.asyncio
async def test_deliver_ended_stream_rewrites_writes_tool_call_name_back_into_chunks(self):
class RenameTool(CustomGuardrail):
async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
for tool_call in inputs.get("tool_calls", []):
tool_call["function"]["name"] = "lookup_fruit_reviewed"
return inputs
handler = OpenAIChatCompletionsHandler()
chunks = self._ended_tool_call_stream_chunks()
await handler.process_output_streaming_response(
responses_so_far=chunks,
guardrail_to_apply=RenameTool(guardrail_name="test"),
litellm_logging_obj=None,
deliver_ended_stream_rewrites=True,
)
fragments = [chunk.choices[0].delta.tool_calls[0] for chunk in chunks[:3]]
assert [fragment.function.name for fragment in fragments] == ["lookup_fruit_reviewed", None, None]
assert json.loads("".join(fragment.function.arguments for fragment in fragments)) == {"fruit": "persimmon"}
assert fragments[0].id == "call_1"
@pytest.mark.asyncio
async def test_ended_stream_tool_call_rewrite_leaves_chunks_untouched_by_default(self):
handler = OpenAIChatCompletionsHandler()
guardrail = MockGuardrail(guardrail_name="test")
chunks = self._ended_tool_call_stream_chunks()
await handler.process_output_streaming_response(
responses_so_far=chunks,
guardrail_to_apply=guardrail,
litellm_logging_obj=None,
)
fragments = [chunk.choices[0].delta.tool_calls for chunk in chunks[:3]]
assert [fragment[0].function.arguments for fragment in fragments] == ["", '{"fruit":', ' "persimmon"}']
@pytest.mark.asyncio
async def test_ended_stream_rewrite_leaves_chunks_untouched_by_default(self):
handler = OpenAIChatCompletionsHandler()
@ -1179,6 +1275,62 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput:
deliver_ended_stream_rewrites=True,
)
@staticmethod
def _two_choice_tool_call_stream_chunks() -> list:
from litellm.types.utils import (
ChatCompletionDeltaToolCall,
Delta,
Function,
ModelResponseStream,
StreamingChoices,
)
def chunk(
choice_index: int, tool_call: ChatCompletionDeltaToolCall | None, finish_reason: Optional[str] = None
) -> ModelResponseStream:
return ModelResponseStream(
id="chatcmpl-123",
created=1234567890,
model="gpt-4",
object="chat.completion.chunk",
choices=[
StreamingChoices(
index=choice_index,
delta=Delta(tool_calls=[tool_call] if tool_call else None),
finish_reason=finish_reason,
)
],
)
def fragment(arguments: str, name: Optional[str] = None, call_id: Optional[str] = None):
return ChatCompletionDeltaToolCall(
id=call_id, index=0, type="function", function=Function(name=name, arguments=arguments)
)
return [
chunk(0, fragment("", name="lookup_fruit", call_id="call_1")),
chunk(1, fragment("", name="lookup_fruit", call_id="call_2")),
chunk(0, fragment('{"fruit": "persimmon"}')),
chunk(1, fragment('{"fruit": "durian"}')),
chunk(0, None, finish_reason="tool_calls"),
chunk(1, None, finish_reason="tool_calls"),
]
@pytest.mark.asyncio
async def test_deliver_ended_stream_tool_call_rewrite_on_multi_choice_stream_fails_closed(self):
from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
handler = OpenAIChatCompletionsHandler()
chunks = self._two_choice_tool_call_stream_chunks()
with pytest.raises(UndeliverableStreamRewrite):
await handler.process_output_streaming_response(
responses_so_far=chunks,
guardrail_to_apply=MockGuardrail(guardrail_name="test"),
litellm_logging_obj=None,
deliver_ended_stream_rewrites=True,
)
@pytest.mark.asyncio
async def test_deliver_ended_stream_clean_multi_choice_stream_released_untouched(self):
handler = OpenAIChatCompletionsHandler()

225
tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py
View file

@ -17,6 +17,7 @@ from fastapi import HTTPException
from openai.types.responses import ResponseFunctionToolCall
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms import get_guardrail_translation_mapping
from litellm.llms.openai.responses.guardrail_translation.handler import (
OpenAIResponsesHandler,
@ -1195,6 +1196,230 @@ class TestOpenAIResponsesHandlerStreamingOutputProcessing:
assert events[4]["item"]["content"][0]["text"] == "hello [MASKED]"
assert events[5]["response"]["output"][0]["content"][0]["text"] == "hello [MASKED]"
@staticmethod
def _ended_function_call_stream_events() -> List[dict]:
def item(arguments: str, status: str) -> dict:
return {
"type": "function_call",
"id": "fc_123",
"call_id": "call_123",
"name": "lookup_fruit",
"arguments": arguments,
"status": status,
}
return [
{"type": "response.output_item.added", "output_index": 0, "item": item("", "in_progress")},
{"type": "response.function_call_arguments.delta", "item_id": "fc_123", "output_index": 0, "delta": '{"fruit":'},
{"type": "response.function_call_arguments.delta", "item_id": "fc_123", "output_index": 0, "delta": ' "persimmon"}'},
{
"type": "response.function_call_arguments.done",
"item_id": "fc_123",
"output_index": 0,
"arguments": '{"fruit": "persimmon"}',
},
{"type": "response.output_item.done", "output_index": 0, "item": item('{"fruit": "persimmon"}', "completed")},
{
"type": "response.completed",
"response": {
"id": "resp_123",
"created_at": 1,
"model": "gpt-4o",
"output": [item('{"fruit": "persimmon"}', "completed")],
"status": "completed",
},
},
]
@staticmethod
def _argument_masking_guardrail() -> CustomGuardrail:
class MaskArguments(CustomGuardrail):
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict,
input_type: Literal["request", "response"],
logging_obj: LiteLLMLoggingObj | None = None,
) -> GenericGuardrailAPIInputs:
tool_calls = [
{**tool_call, "function": {**tool_call["function"], "arguments": '{"fruit": "[MASKED]"}'}}
for tool_call in inputs.get("tool_calls", [])
]
return {**inputs, "tool_calls": tool_calls}
return MaskArguments(guardrail_name="test-mask-arguments")
@pytest.mark.asyncio
async def test_deliver_ended_stream_rewrites_syncs_function_call_events(self):
handler = OpenAIResponsesHandler()
events = self._ended_function_call_stream_events()
result = await handler.process_output_streaming_response(
responses_so_far=events,
guardrail_to_apply=self._argument_masking_guardrail(),
litellm_logging_obj=None,
deliver_ended_stream_rewrites=True,
)
assert result is events
assert events[0]["item"]["arguments"] == ""
assert events[1]["delta"] == '{"fruit": "[MASKED]"}'
assert events[2]["delta"] == ""
assert events[3]["arguments"] == '{"fruit": "[MASKED]"}'
assert events[4]["item"]["arguments"] == '{"fruit": "[MASKED]"}'
assert events[5]["response"]["output"][0]["arguments"] == '{"fruit": "[MASKED]"}'
assert events[5]["response"]["output"][0]["name"] == "lookup_fruit"
@pytest.mark.asyncio
async def test_deliver_ended_stream_rewrites_syncs_typed_function_call_events(self):
from litellm.types.llms.openai import (
FunctionCallArgumentsDeltaEvent,
FunctionCallArgumentsDoneEvent,
OutputItemAddedEvent,
OutputItemDoneEvent,
ResponseCompletedEvent,
ResponsesAPIResponse,
)
handler = OpenAIResponsesHandler()
typed_events: List[Any] = [
model.model_validate(event)
for model, event in zip(
(
OutputItemAddedEvent,
FunctionCallArgumentsDeltaEvent,
FunctionCallArgumentsDeltaEvent,
FunctionCallArgumentsDoneEvent,
OutputItemDoneEvent,
ResponseCompletedEvent,
),
self._ended_function_call_stream_events(),
)
]
completed_event = typed_events[5]
assert isinstance(completed_event, ResponseCompletedEvent)
assert isinstance(completed_event.response, ResponsesAPIResponse)
assert isinstance(completed_event.response.output[0], ResponseFunctionToolCall)
await handler.process_output_streaming_response(
responses_so_far=typed_events,
guardrail_to_apply=self._argument_masking_guardrail(),
litellm_logging_obj=None,
deliver_ended_stream_rewrites=True,
)
assert typed_events[1].delta == '{"fruit": "[MASKED]"}'
assert typed_events[2].delta == ""
assert typed_events[3].arguments == '{"fruit": "[MASKED]"}'
assert typed_events[4].item.arguments == '{"fruit": "[MASKED]"}'
assert completed_event.response.output[0].arguments == '{"fruit": "[MASKED]"}'
assert completed_event.response.output[0].name == "lookup_fruit"
@staticmethod
def _bridged_function_call_stream_events() -> List[dict]:
reasoning = {"type": "reasoning", "id": "rs_1", "summary": []}
text = {"type": "output_text", "text": "Looking that up", "annotations": []}
message = {"type": "message", "id": "msg_1", "role": "assistant", "status": "completed", "content": [text]}
def function_call(arguments: str, status: str) -> dict:
return {
"type": "function_call",
"id": "fc_1",
"call_id": "call_1",
"name": "lookup_fruit",
"arguments": arguments,
"status": status,
}
return [
{"type": "response.output_item.added", "output_index": 0, "item": dict(reasoning)},
{"type": "response.output_item.done", "output_index": 0, "item": dict(reasoning)},
{"type": "response.output_item.added", "output_index": 0, "item": {**message, "status": "in_progress", "content": []}},
{"type": "response.output_text.delta", "item_id": "msg_1", "output_index": 0, "content_index": 0, "delta": "Looking that up"},
{"type": "response.output_item.done", "output_index": 0, "item": {**message, "content": [dict(text)]}},
{"type": "response.output_item.added", "output_index": 1, "item": function_call("", "in_progress")},
{"type": "response.function_call_arguments.delta", "item_id": "fc_1", "output_index": 1, "delta": '{"fruit":'},
{"type": "response.function_call_arguments.delta", "item_id": "fc_1", "output_index": 1, "delta": ' "persimmon"}'},
{
"type": "response.function_call_arguments.done",
"item_id": "fc_1",
"output_index": 1,
"arguments": '{"fruit": "persimmon"}',
},
{"type": "response.output_item.done", "output_index": 1, "item": function_call('{"fruit": "persimmon"}', "completed")},
{
"type": "response.completed",
"response": {
"id": "resp_1",
"model": "claude-haiku-4-5",
"output": [
dict(reasoning),
{**message, "content": [dict(text)]},
function_call('{"fruit": "persimmon"}', "completed"),
],
},
},
]
@pytest.mark.asyncio
async def test_deliver_ended_stream_rewrites_keys_bridged_function_call_events_by_call_id(self):
handler = OpenAIResponsesHandler()
events = self._bridged_function_call_stream_events()
await handler.process_output_streaming_response(
responses_so_far=events,
guardrail_to_apply=self._argument_masking_guardrail(),
litellm_logging_obj=None,
deliver_ended_stream_rewrites=True,
)
assert events[6]["delta"] == '{"fruit": "[MASKED]"}'
assert events[7]["delta"] == ""
assert events[8]["arguments"] == '{"fruit": "[MASKED]"}'
assert events[5]["item"]["name"] == "lookup_fruit"
assert events[9]["item"]["arguments"] == '{"fruit": "[MASKED]"}'
assert events[10]["response"]["output"][2]["arguments"] == '{"fruit": "[MASKED]"}'
assert events[3]["delta"] == "Looking that up"
assert events[4]["item"]["content"][0]["text"] == "Looking that up"
assert events[10]["response"]["output"][1]["content"][0]["text"] == "Looking that up"
assert events[1]["item"] == {"type": "reasoning", "id": "rs_1", "summary": []}
@pytest.mark.asyncio
@pytest.mark.parametrize("mismatch", ["orphan_call_id", "duplicate_call_id"])
async def test_deliver_ended_stream_function_call_rewrite_without_matching_events_fails_closed(self, mismatch):
from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite
handler = OpenAIResponsesHandler()
events = self._ended_function_call_stream_events()
envelope_item = events[5]["response"]["output"][0]
if mismatch == "orphan_call_id":
events[5]["response"]["output"] = [{**envelope_item, "call_id": "call_999"}]
else:
events[5]["response"]["output"] = [dict(envelope_item), dict(envelope_item)]
with pytest.raises(UndeliverableStreamRewrite):
await handler.process_output_streaming_response(
responses_so_far=events,
guardrail_to_apply=self._argument_masking_guardrail(),
litellm_logging_obj=None,
deliver_ended_stream_rewrites=True,
)
@pytest.mark.asyncio
async def test_ended_stream_function_call_rewrite_leaves_events_untouched_by_default(self):
handler = OpenAIResponsesHandler()
events = self._ended_function_call_stream_events()
await handler.process_output_streaming_response(
responses_so_far=events,
guardrail_to_apply=self._argument_masking_guardrail(),
litellm_logging_obj=None,
)
assert events[1]["delta"] == '{"fruit":'
assert events[3]["arguments"] == '{"fruit": "persimmon"}'
assert events[5]["response"]["output"][0]["arguments"] == '{"fruit": "persimmon"}'
@pytest.mark.asyncio
@pytest.mark.parametrize("terminal_type", ["response.incomplete", "response.failed"])
async def test_deliver_ended_stream_rewrites_syncs_non_completed_terminals(self, terminal_type):

157
tests/test_litellm/proxy/auth/test_login_utils.py
View file

@ -23,6 +23,7 @@ from litellm.proxy.auth.login_utils import (
LoginResult,
authenticate_user,
get_ui_credentials,
is_env_credential_login_enabled,
)
@ -185,6 +186,7 @@ async def test_authenticate_user_invalid_credentials():
assert exc_info.value.type == ProxyErrorTypes.auth_error
assert exc_info.value.code == "401"
assert "Invalid credentials" in exc_info.value.message
assert "UI_USERNAME" in exc_info.value.message
@pytest.mark.asyncio
@ -799,3 +801,158 @@ class TestDisablePasswordLoginWhenSSOEnabled:
assert isinstance(result, LoginResult)
assert result.user_id == LITELLM_PROXY_ADMIN_NAME
class TestDisableEnvCredentialLogin:
"""`disable_env_credential_login` must reject a login with the env
credentials (UI_USERNAME/UI_PASSWORD, or the master-key fallback when
UI_PASSWORD is unset) while leaving database-user password logins
untouched, so admins with real accounts keep a way in."""
@pytest.mark.asyncio
async def test_rejects_correct_env_credentials_when_disabled(self):
master_key = "sk-1234"
ui_username = "admin"
ui_password = "env-only-password"
mock_prisma_client = MagicMock()
mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None)
with patch.dict(os.environ, {"UI_USERNAME": ui_username, "UI_PASSWORD": ui_password}):
with pytest.raises(ProxyException) as exc_info:
await authenticate_user(
username=ui_username,
password=ui_password,
master_key=master_key,
prisma_client=mock_prisma_client,
general_settings={"disable_env_credential_login": True},
)
assert exc_info.value.type == ProxyErrorTypes.auth_error
assert exc_info.value.code == "401"
assert "UI_USERNAME" not in exc_info.value.message
assert "UI_PASSWORD" not in exc_info.value.message
@pytest.mark.asyncio
async def test_rejects_master_key_fallback_when_disabled(self):
"""With UI_PASSWORD unset, the master key IS the env password, so the
setting must reject it too or it protects nothing by default."""
master_key = "sk-1234"
mock_prisma_client = MagicMock()
mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None)
with patch.dict(os.environ, {"UI_USERNAME": "admin"}, clear=True):
with pytest.raises(ProxyException) as exc_info:
await authenticate_user(
username="admin",
password=master_key,
master_key=master_key,
prisma_client=mock_prisma_client,
general_settings={"disable_env_credential_login": True},
)
assert exc_info.value.code == "401"
@pytest.mark.asyncio
async def test_db_user_login_still_works_when_disabled(self):
master_key = "sk-1234"
user_email = "admin@example.com"
password = "Str0ng!Passw0rd"
mock_user = LiteLLM_UserTable(
user_id="db-admin-1",
user_email=user_email,
password=hash_token(token=password),
user_role=LitellmUserRoles.PROXY_ADMIN,
)
mock_prisma_client = MagicMock()
mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=mock_user)
with patch.dict(
os.environ,
{
"UI_USERNAME": "admin",
"UI_PASSWORD": "env-password",
"DATABASE_URL": "postgresql://test:test@localhost/test",
},
clear=True,
):
with ExitStack() as stack:
stack.enter_context(
patch( # test-quality-ok: internal orchestration, no HTTP boundary; matches pre-existing tests
"litellm.proxy.auth.login_utils.generate_key_helper_fn",
new_callable=AsyncMock,
return_value={"token": "db-user-token"},
)
)
result = await authenticate_user(
username=user_email,
password=password,
master_key=master_key,
prisma_client=mock_prisma_client,
general_settings={"disable_env_credential_login": True},
)
assert isinstance(result, LoginResult)
assert result.user_id == "db-admin-1"
assert result.user_role == LitellmUserRoles.PROXY_ADMIN
@pytest.mark.asyncio
async def test_env_login_still_works_when_setting_absent(self):
"""Env-credential login is the bootstrap path on a fresh install and
must stay on by default."""
master_key = "sk-1234"
ui_username = "admin"
mock_prisma_client = MagicMock()
mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None)
with patch.dict(
os.environ,
{
"UI_USERNAME": ui_username,
"UI_PASSWORD": master_key,
"DATABASE_URL": "postgresql://test:test@localhost/test",
},
clear=True,
):
with ExitStack() as stack:
_patch_successful_admin_login_deps(stack)
result = await authenticate_user(
username=ui_username,
password=master_key,
master_key=master_key,
prisma_client=mock_prisma_client,
general_settings={},
)
assert isinstance(result, LoginResult)
assert result.user_id == LITELLM_PROXY_ADMIN_NAME
class TestIsEnvCredentialLoginEnabled:
"""Drives the Admin UI warning banner: it must be True exactly when a
login with the env credentials could actually succeed."""
def test_enabled_by_default(self):
assert is_env_credential_login_enabled({}) is True
def test_disabled_by_dedicated_setting(self):
assert is_env_credential_login_enabled({"disable_env_credential_login": True}) is False
def test_explicit_false_keeps_it_enabled(self):
assert is_env_credential_login_enabled({"disable_env_credential_login": False}) is True
def test_disabled_when_sso_gate_blocks_all_password_logins(self):
"""`disable_password_login_when_sso_enabled` with SSO configured
rejects every username/password login before the env comparison runs,
so the banner must not nag about an already-unreachable path."""
with ExitStack() as stack:
_patch_sso_configured(stack, configured=True)
assert is_env_credential_login_enabled({"disable_password_login_when_sso_enabled": True}) is False
def test_enabled_when_sso_gate_is_set_but_sso_not_configured(self):
with ExitStack() as stack:
_patch_sso_configured(stack, configured=False)
assert is_env_credential_login_enabled({"disable_password_login_when_sso_enabled": True}) is True

27
tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py
View file

@ -1301,6 +1301,33 @@ def test_health_readiness_details_returns_diagnostic_fields(monkeypatch):
assert "cache" in response_data
@pytest.mark.parametrize(
"general_settings, expected_warning",
[
({}, True),
({"disable_env_credential_login": True}, False),
],
)
def test_health_readiness_details_reports_env_credential_login_warning(monkeypatch, general_settings, expected_warning):
"""
The Admin UI banner is driven by this flag: it must be True while
env-credential login is possible and False once
`disable_env_credential_login` turns that login path off.
"""
app = FastAPI()
app.include_router(_health_endpoints_module.router)
app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)
client = TestClient(app)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", general_settings)
response = client.get("/health/readiness/details")
assert response.status_code == 200, response.text
assert response.json()["show_env_credential_login_warning"] is expected_warning
def test_health_readiness_allows_explicit_legacy_public_details(monkeypatch):
"""
Operators can explicitly preserve the legacy public readiness payload.

57
tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter_v3.py
View file

@ -1861,3 +1861,60 @@ async def test_tpm_only_model_enforces_priority_and_model_capacity(monkeypatch):
)
assert capacity_blocked.value.status_code == 429
assert "Model capacity reached" in capacity_blocked.value.detail["error"]
@pytest.mark.asyncio
async def test_post_call_success_hook_attaches_priority_headers_to_dict_response():
from litellm.proxy.hooks.parallel_request_limiter_v3 import (
RateLimitResponse,
RateLimitStatus,
get_or_create_request_stash,
)
handler = DynamicRateLimitHandler(internal_usage_cache=DualCache())
get_or_create_request_stash().rate_limit_response = RateLimitResponse(
overall_code="OK",
statuses=[
RateLimitStatus(
code="OK",
current_limit=75,
limit_remaining=74,
rate_limit_type="requests",
descriptor_key="priority_model",
)
],
)
response = {
"id": "msg_123",
"type": "message",
"role": "assistant",
"content": [],
"_hidden_params": {"additional_headers": {"x-litellm-attempted-retries": 0}},
}
await handler.async_post_call_success_hook(
data={"model": "anthropic-haiku"},
user_api_key_dict=UserAPIKeyAuth(metadata={"priority": "premium"}),
response=response,
)
additional_headers = response["_hidden_params"]["additional_headers"]
assert additional_headers["x-litellm-attempted-retries"] == 0
assert additional_headers["x-ratelimit-priority_model-limit-requests"] == 75
assert additional_headers["x-ratelimit-priority_model-remaining-requests"] == 74
assert additional_headers["x-litellm-priority"] == "premium"
assert additional_headers["x-litellm-rate-limiter-version"] == "v3"
@pytest.mark.asyncio
async def test_post_call_success_hook_leaves_raw_provider_dict_untouched():
handler = DynamicRateLimitHandler(internal_usage_cache=DualCache())
response = {"id": "msg_123", "type": "message", "role": "assistant", "content": []}
await handler.async_post_call_success_hook(
data={"model": "anthropic-haiku"},
user_api_key_dict=UserAPIKeyAuth(metadata={"priority": "premium"}),
response=response,
)
assert response == {"id": "msg_123", "type": "message", "role": "assistant", "content": []}

65
tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py
View file

@ -6171,3 +6171,68 @@ async def test_success_hook_leaves_stash_untouched_for_non_batch_responses():
data={}, user_api_key_dict=user, response=ModelResponse(usage=Usage(total_tokens=5))
)
assert get_request_stash().batch_enqueued_reservation == reservation
@pytest.mark.asyncio
async def test_post_call_success_hook_attaches_ratelimit_headers_to_dict_response():
from litellm.proxy.hooks.parallel_request_limiter_v3 import RateLimitResponse, RateLimitStatus
handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache()))
get_or_create_request_stash().rate_limit_response = RateLimitResponse(
overall_code="OK",
statuses=[
RateLimitStatus(
code="OK",
current_limit=100,
limit_remaining=99,
rate_limit_type="requests",
descriptor_key="model_saturation_check",
)
],
)
response = {
"id": "msg_123",
"type": "message",
"role": "assistant",
"content": [],
"_hidden_params": {"additional_headers": {"x-litellm-attempted-retries": 0}},
}
await handler.async_post_call_success_hook(
data={"model": "anthropic-haiku"},
user_api_key_dict=UserAPIKeyAuth(api_key=hash_token("sk-dict-response")),
response=response,
)
additional_headers = response["_hidden_params"]["additional_headers"]
assert additional_headers["x-litellm-attempted-retries"] == 0
assert additional_headers["x-ratelimit-model_saturation_check-limit-requests"] == 100
assert additional_headers["x-ratelimit-model_saturation_check-remaining-requests"] == 99
@pytest.mark.asyncio
async def test_post_call_success_hook_leaves_raw_provider_dict_untouched():
from litellm.proxy.hooks.parallel_request_limiter_v3 import RateLimitResponse, RateLimitStatus
handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache()))
get_or_create_request_stash().rate_limit_response = RateLimitResponse(
overall_code="OK",
statuses=[
RateLimitStatus(
code="OK",
current_limit=100,
limit_remaining=99,
rate_limit_type="requests",
descriptor_key="model_saturation_check",
)
],
)
response = {"id": "msg_123", "type": "message", "role": "assistant", "content": []}
await handler.async_post_call_success_hook(
data={"model": "anthropic-haiku"},
user_api_key_dict=UserAPIKeyAuth(api_key=hash_token("sk-raw-dict")),
response=response,
)
assert response == {"id": "msg_123", "type": "message", "role": "assistant", "content": []}

47
tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py
View file

@ -1063,7 +1063,7 @@ class _TextReturningGuardrail(CustomGuardrail):
class _TextTranslation:
delivers_ended_stream_text_rewrites = False
delivers_ended_stream_rewrites = False
def __init__(self):
self.seen_guardrail_names = []
@ -1085,7 +1085,7 @@ class _WritingTranslation:
"""Writes the guardrail's text (and tool-call) outputs back into the buffered chunks the way the
chat/Responses/Messages handlers do on an ended stream."""
delivers_ended_stream_text_rewrites = True
delivers_ended_stream_rewrites = True
async def process_output_streaming_response(
self,
@ -1104,12 +1104,13 @@ class _WritingTranslation:
logging_obj=litellm_logging_obj,
)
responses_so_far[0]["text"] = outputs["texts"][0]
responses_so_far[0]["tool_call"] = outputs["tool_calls"][0]
if len(outputs["tool_calls"]) == 1:
responses_so_far[0]["tool_call"] = outputs["tool_calls"][0]
return responses_so_far
class _RefusingTranslation:
delivers_ended_stream_text_rewrites = True
delivers_ended_stream_rewrites = True
async def process_output_streaming_response(
self,
@ -1228,13 +1229,47 @@ async def test_streaming_step_delivers_text_rewrite_through_writing_translation(
@pytest.mark.asyncio
async def test_streaming_step_discards_tool_call_rewrite_and_restores_written_text(monkeypatch, caplog):
async def test_streaming_step_delivers_tool_call_rewrite_through_writing_translation(monkeypatch, caplog):
monkeypatch.setattr(litellm, "callbacks", [_TextAndToolCallRewritingGuardrail(rewrite_tool_call=True)])
chunks = [_chunk()]
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
result = await _run_streaming_step(_WritingTranslation(), chunks)
assert result.terminal_action == "allow"
assert chunks[0]["text"] == "hello [MASKED]"
assert chunks[0]["tool_call"]["function"]["arguments"] == '{"ssn": "[MASKED]"}'
assert not any("discarded" in record.getMessage() for record in caplog.records)
class _ToolCallDroppingGuardrail(CustomGuardrail):
def __init__(self):
super().__init__(guardrail_name="masker", event_hook="post_call", default_on=True)
async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
return {**inputs, "texts": ["hello [MASKED]"], "tool_calls": []}
@pytest.mark.asyncio
async def test_streaming_step_discards_whole_rewrite_when_guardrail_drops_a_tool_call(monkeypatch, caplog):
monkeypatch.setattr(litellm, "callbacks", [_ToolCallDroppingGuardrail()])
chunks = [_chunk()]
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
result = await _run_streaming_step(_WritingTranslation(), chunks)
_assert_passed_with_discard_warning(result, caplog)
assert chunks == [_chunk()]
@pytest.mark.asyncio
async def test_streaming_step_discards_tool_call_rewrite_when_translation_lacks_write_back(monkeypatch, caplog):
monkeypatch.setattr(litellm, "callbacks", [_TextAndToolCallRewritingGuardrail(rewrite_tool_call=True)])
chunks = [_chunk()]
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
result = await _run_streaming_step(_TextTranslation(), chunks)
_assert_passed_with_discard_warning(result, caplog)
assert chunks == [_chunk()]
@ -1326,7 +1361,7 @@ class _LegacyScanningTranslation:
chat, Responses, and Messages handlers, hands hooks a route-native shape, and re-extracts one
text per entry of a replacement's "texts"."""
delivers_ended_stream_text_rewrites = True
delivers_ended_stream_rewrites = True
def post_call_hook_response(self, response):
return {"native": True, "text": response["text"], "tool_calls": response["tool_calls"]}

262
tests/test_litellm/proxy/policy_engine/test_response_retrieval.py Normal file
View file

@ -0,0 +1,262 @@
import logging
from collections.abc import Iterator, Mapping
import pytest
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry
from litellm.proxy.policy_engine.policy_registry import get_policy_registry
from litellm.proxy.policy_engine.response_retrieval import attach_post_call_pipelines_to_retrieval
from litellm.responses.utils import ResponsesAPIRequestUtils
from litellm.types.router import Deployment, LiteLLM_Params
GOVERNED_MODEL_GROUP = "gpt-5.4-mini"
GOVERNED_MODEL_ID = "deployment-governed"
UNGOVERNED_MODEL_GROUP = "gpt-4.1-mini"
UNGOVERNED_MODEL_ID = "deployment-ungoverned"
WILDCARD_MODEL_GROUP = "openai/*"
WILDCARD_MODEL_ID = "deployment-wildcard"
class FakeRouter:
def __init__(self, deployments: dict[str, Deployment], model_group_alias: dict[str, object] | None = None):
self._deployments = deployments
self.model_group_alias = model_group_alias or {}
def get_deployment(self, model_id: str) -> Deployment | None:
return self._deployments.get(model_id)
def _deployment(model_group: str, model_id: str) -> Deployment:
return Deployment(
model_name=model_group,
litellm_params=LiteLLM_Params(model=f"openai/{model_group}"),
model_info={"id": model_id},
)
def _router(model_group_alias: dict[str, object] | None = None) -> FakeRouter:
return FakeRouter(
{
GOVERNED_MODEL_ID: _deployment(GOVERNED_MODEL_GROUP, GOVERNED_MODEL_ID),
UNGOVERNED_MODEL_ID: _deployment(UNGOVERNED_MODEL_GROUP, UNGOVERNED_MODEL_ID),
WILDCARD_MODEL_ID: _deployment(WILDCARD_MODEL_GROUP, WILDCARD_MODEL_ID),
},
model_group_alias,
)
def _encoded_response_id(model_id: str) -> str:
return ResponsesAPIRequestUtils._build_responses_api_response_id(
custom_llm_provider="openai", model_id=model_id, response_id="resp_upstream"
)
def _pipeline_policy(guardrail: str, mode: str = "post_call") -> dict[str, object]:
return {
"guardrails": {"add": [guardrail]},
"pipeline": {"mode": mode, "steps": [{"guardrail": guardrail, "on_pass": "allow", "on_fail": "block"}]},
}
@pytest.fixture
def policy_engine() -> Iterator[None]:
policy_registry = get_policy_registry()
attachment_registry = get_attachment_registry()
policy_registry.load_policies(
{
"response-governance": _pipeline_policy("output-word-filter"),
"input-governance": _pipeline_policy("input-word-filter", mode="pre_call"),
"team-governance": _pipeline_policy("team-word-filter"),
"tag-governance": _pipeline_policy("tag-word-filter"),
}
)
attachment_registry.load_attachments(
[
{"policy": "response-governance", "models": [GOVERNED_MODEL_GROUP]},
{"policy": "input-governance", "models": [GOVERNED_MODEL_GROUP]},
{"policy": "team-governance", "teams": ["governed-team"]},
{"policy": "tag-governance", "tags": ["governed"]},
]
)
yield
policy_registry.clear()
attachment_registry.clear()
def _retrieval_data(model_id: str) -> dict[str, object]:
return {"response_id": _encoded_response_id(model_id), "litellm_metadata": {}}
def _attached_pipelines(data: Mapping[str, object]) -> tuple[tuple[str, str], ...]:
bucket = data["litellm_metadata"]
assert isinstance(bucket, dict)
return tuple(
(policy_name, ",".join(step.guardrail for step in pipeline.steps))
for policy_name, pipeline in bucket["_guardrail_pipelines"]
)
def test_attaches_model_scoped_post_call_pipeline_to_retrieval(policy_engine: None) -> None:
data = _retrieval_data(GOVERNED_MODEL_ID)
attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=_router())
assert _attached_pipelines(data) == (("response-governance", "output-word-filter"),)
assert data["litellm_metadata"]["_pipeline_managed_guardrails"] == frozenset({"output-word-filter"})
assert data["litellm_metadata"]["applied_policies"] == ["response-governance"]
assert data["litellm_metadata"]["applied_guardrails"] == ["output-word-filter"]
assert data["litellm_metadata"]["policy_sources"] == {"response-governance": "model:gpt-5.4-mini"}
assert "model" not in data
assert "guardrails" not in data["litellm_metadata"]
def test_key_and_team_context_also_governs_retrieval(policy_engine: None) -> None:
data = _retrieval_data(UNGOVERNED_MODEL_ID)
attach_post_call_pipelines_to_retrieval(
data=data, user_api_key_dict=UserAPIKeyAuth(team_alias="governed-team"), llm_router=_router()
)
assert _attached_pipelines(data) == (("team-governance", "team-word-filter"),)
def test_tag_attached_policy_is_not_re_matched_when_the_retrieval_carries_no_tag(policy_engine: None) -> None:
data = _retrieval_data(GOVERNED_MODEL_ID)
attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=_router())
assert _attached_pipelines(data) == (("response-governance", "output-word-filter"),)
def test_tag_attached_policy_governs_a_retrieval_whose_metadata_carries_the_tag(policy_engine: None) -> None:
data: dict[str, object] = {
"response_id": _encoded_response_id(UNGOVERNED_MODEL_ID),
"litellm_metadata": {"tags": ["governed"]},
}
attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=_router())
assert _attached_pipelines(data) == (("tag-governance", "tag-word-filter"),)
assert data["litellm_metadata"]["policy_sources"] == {"tag-governance": "tag:governed"}
def test_retrieval_of_an_ungoverned_model_attaches_nothing(policy_engine: None) -> None:
data = _retrieval_data(UNGOVERNED_MODEL_ID)
attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=_router())
assert data == _retrieval_data(UNGOVERNED_MODEL_ID)
def test_already_attached_policy_is_not_attached_twice(policy_engine: None) -> None:
data = _retrieval_data(GOVERNED_MODEL_ID)
router = _router()
attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=router)
attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=router)
assert _attached_pipelines(data) == (("response-governance", "output-word-filter"),)
assert data["litellm_metadata"]["applied_policies"] == ["response-governance"]
def _hidden_submit_model_warnings(caplog: pytest.LogCaptureFixture) -> list[str]:
return [
record.getMessage()
for record in caplog.records
if record.levelno == logging.WARNING and "the model name it was submitted as" in record.getMessage()
]
def test_wildcard_deployment_attaches_nothing_for_the_submitted_model_and_warns(
policy_engine: None, caplog: pytest.LogCaptureFixture
) -> None:
data = _retrieval_data(WILDCARD_MODEL_ID)
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=_router())
assert data == _retrieval_data(WILDCARD_MODEL_ID)
assert [
"as model group openai/* (a wildcard deployment)" in message
for message in _hidden_submit_model_warnings(caplog)
] == [True]
def test_aliased_model_group_still_attaches_its_own_policies_and_warns(
policy_engine: None, caplog: pytest.LogCaptureFixture
) -> None:
data = _retrieval_data(GOVERNED_MODEL_ID)
router = _router({"gpt-mini": GOVERNED_MODEL_GROUP, "gpt-hidden": {"model": GOVERNED_MODEL_GROUP, "hidden": True}})
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=router)
assert _attached_pipelines(data) == (("response-governance", "output-word-filter"),)
assert [
"(the target of model_group_alias gpt-mini, gpt-hidden)" in message
for message in _hidden_submit_model_warnings(caplog)
] == [True]
def test_plain_model_group_retrieval_does_not_warn_about_the_submitted_model(
policy_engine: None, caplog: pytest.LogCaptureFixture
) -> None:
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
attach_post_call_pipelines_to_retrieval(
data=_retrieval_data(GOVERNED_MODEL_ID),
user_api_key_dict=UserAPIKeyAuth(),
llm_router=_router({"other-alias": UNGOVERNED_MODEL_GROUP}),
)
assert _hidden_submit_model_warnings(caplog) == []
def _ungoverned_retrieval_warnings(caplog: pytest.LogCaptureFixture) -> list[str]:
return [
record.getMessage()
for record in caplog.records
if record.levelno == logging.WARNING
and "retrieved without its post_call policy pipelines" in record.getMessage()
]
@pytest.mark.parametrize(
("response_id", "reason"),
[
("resp_plain_upstream_id", "response id names no deployment"),
(_encoded_response_id("deployment-missing-from-router"), "deployment no longer in the router"),
(None, "response id names no deployment"),
],
)
def test_unresolvable_response_id_attaches_nothing_and_warns(
policy_engine: None, caplog: pytest.LogCaptureFixture, response_id: str, reason: str
) -> None:
data = {"response_id": response_id, "litellm_metadata": {}}
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=_router())
assert data == {"response_id": response_id, "litellm_metadata": {}}
assert [message.endswith(f"({reason})") for message in _ungoverned_retrieval_warnings(caplog)] == [True]
def test_without_a_router_attaches_nothing_and_warns(policy_engine: None, caplog: pytest.LogCaptureFixture) -> None:
data = _retrieval_data(GOVERNED_MODEL_ID)
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=None)
assert data == _retrieval_data(GOVERNED_MODEL_ID)
assert [message.endswith("(no router)") for message in _ungoverned_retrieval_warnings(caplog)] == [True]
def test_without_policy_engine_attaches_nothing_quietly(caplog: pytest.LogCaptureFixture) -> None:
get_policy_registry().clear()
data = _retrieval_data(GOVERNED_MODEL_ID)
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=None)
assert data == _retrieval_data(GOVERNED_MODEL_ID)
assert _ungoverned_retrieval_warnings(caplog) == []

115
tests/test_litellm/proxy/test_common_request_processing.py
View file

@ -3,7 +3,7 @@ import copy
import datetime
import json
from types import MappingProxyType, SimpleNamespace
from typing import AsyncGenerator, Callable, Final, Optional
from typing import AsyncGenerator, Callable, Final, Iterator, Optional
from unittest.mock import AsyncMock, MagicMock, patch
import httpx
@ -8307,3 +8307,116 @@ async def test_handle_llm_api_exception_forwards_provider_headers_on_http_status
assert exc_info.value.headers is not None
assert exc_info.value.headers["llm_provider-x-amzn-requestid"] == "req-passthrough-500"
class TestBackgroundResponseRetrievalGovernance:
"""LIT-7175: retrieving a background Response attaches the model's post_call policy pipelines."""
GOVERNED_MODEL_GROUP = "gpt-5.4-mini"
GOVERNED_MODEL_ID = "deployment-governed"
@pytest.fixture
def policy_engine(self) -> Iterator[None]:
from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry
from litellm.proxy.policy_engine.policy_registry import get_policy_registry
get_policy_registry().load_policies(
{
"response-governance": {
"guardrails": {"add": ["output-word-filter"]},
"pipeline": {
"mode": "post_call",
"steps": [{"guardrail": "output-word-filter", "on_pass": "allow", "on_fail": "block"}],
},
}
}
)
get_attachment_registry().load_attachments(
[{"policy": "response-governance", "models": [self.GOVERNED_MODEL_GROUP]}]
)
yield
get_policy_registry().clear()
get_attachment_registry().clear()
def _router(self) -> MagicMock:
from litellm.types.router import Deployment, LiteLLM_Params
router = MagicMock()
router.get_deployment.side_effect = lambda model_id: (
Deployment(
model_name=self.GOVERNED_MODEL_GROUP,
litellm_params=LiteLLM_Params(model=f"openai/{self.GOVERNED_MODEL_GROUP}"),
model_info={"id": model_id},
)
if model_id == self.GOVERNED_MODEL_ID
else None
)
return router
async def _pre_call(self, route_type: str, monkeypatch: pytest.MonkeyPatch) -> dict[str, object]:
from litellm.responses.utils import ResponsesAPIRequestUtils
client_facing_response_id = "resp_opaque-client-facing-id"
encoded_response_id = ResponsesAPIRequestUtils._build_responses_api_response_id(
custom_llm_provider="openai", model_id=self.GOVERNED_MODEL_ID, response_id="resp_upstream"
)
processing_obj = ProxyBaseLLMRequestProcessing(
data={"response_id": client_facing_response_id, "litellm_metadata": {}}
)
mock_request = MagicMock(spec=Request)
mock_request.headers = {}
async def passthrough_add_litellm_data_to_request(
data: dict[str, object], **kwargs: object
) -> dict[str, object]:
return data
async def decrypting_pre_call_hook(
user_api_key_dict: ProxyUserAPIKeyAuth, data: dict[str, object], call_type: str
) -> dict[str, object]:
if data.get("response_id") == client_facing_response_id:
data["response_id"] = encoded_response_id
return data
monkeypatch.setattr(
litellm.proxy.common_request_processing,
"add_litellm_data_to_request",
passthrough_add_litellm_data_to_request,
)
proxy_logging_obj = MagicMock(spec=ProxyLogging)
proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=decrypting_pre_call_hook)
proxy_config = MagicMock(spec=ProxyConfig)
proxy_config._get_hierarchical_router_settings = AsyncMock(return_value=None)
returned_data, _ = await processing_obj.common_processing_pre_call_logic(
request=mock_request,
general_settings={},
user_api_key_dict=ProxyUserAPIKeyAuth(),
proxy_logging_obj=proxy_logging_obj,
proxy_config=proxy_config,
route_type=route_type,
llm_router=self._router(),
)
return returned_data
@pytest.mark.asyncio
async def test_retrieving_a_background_response_attaches_its_model_post_call_pipeline(
self, policy_engine: None, monkeypatch: pytest.MonkeyPatch
) -> None:
data = await self._pre_call("aget_responses", monkeypatch)
assert data["response_id"].startswith("resp_bGl0ZWxsbTpjdXN0b21f")
pipelines = data["litellm_metadata"]["_guardrail_pipelines"]
assert [(policy_name, [step.guardrail for step in pipeline.steps]) for policy_name, pipeline in pipelines] == [
("response-governance", ["output-word-filter"])
]
assert data["litellm_metadata"]["applied_policies"] == ["response-governance"]
assert data["model"] is None
@pytest.mark.asyncio
async def test_submitting_a_response_does_not_attach_pipelines_from_its_response_id(
self, policy_engine: None, monkeypatch: pytest.MonkeyPatch
) -> None:
data = await self._pre_call("aresponses", monkeypatch)
assert "_guardrail_pipelines" not in data["litellm_metadata"]
assert "applied_policies" not in data["litellm_metadata"]

101
tests/test_litellm/proxy/test_litellm_pre_call_utils.py
View file

@ -7678,6 +7678,107 @@ async def test_missing_session_id_omit_keeps_client_supplied_session_id():
assert _spend_log_session_id(updated) == "client-session-1"
@pytest.mark.asyncio
@pytest.mark.parametrize(
"client_body",
[
{"model": "gpt-4o", "messages": [], "litellm_session_id": "cust-sess-1"},
{"model": "gpt-4o", "messages": [], "litellm_session_id": "cust-sess-1", "metadata": {"trace_id": "trace-1"}},
],
)
async def test_missing_session_id_omit_keeps_body_litellm_session_id(
monkeypatch: pytest.MonkeyPatch, client_body: dict[str, object]
):
from litellm.litellm_core_utils.get_litellm_params import get_litellm_params
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
monkeypatch.setattr(litellm, "request_correlation_in_logs", True)
updated = await add_litellm_data_to_request(
data=client_body,
request=_request_for("/v1/chat/completions"),
user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"),
proxy_config=MagicMock(),
general_settings={"missing_session_id": "omit"},
)
callback_session_id = StandardLoggingPayloadSetup.get_standard_logging_payload_session_id(
logging_obj=SimpleNamespace(litellm_session_id=""),
litellm_params=get_litellm_params(litellm_session_id="cust-sess-1", metadata=updated["metadata"]),
)
assert callback_session_id == "cust-sess-1"
assert updated["metadata"]["session_id"] == "cust-sess-1"
assert _spend_log_session_id(updated) == "cust-sess-1"
@pytest.mark.asyncio
async def test_missing_session_id_omit_body_litellm_session_id_does_not_override_metadata_session_id():
updated = await add_litellm_data_to_request(
data={
"model": "gpt-4o",
"messages": [],
"litellm_session_id": "cust-sess-1",
"metadata": {"session_id": "meta-sess-1"},
},
request=_request_for("/v1/chat/completions"),
user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"),
proxy_config=MagicMock(),
general_settings={"missing_session_id": "omit"},
)
assert updated["metadata"]["session_id"] == "meta-sess-1"
assert _spend_log_session_id(updated) == "meta-sess-1"
@pytest.mark.asyncio
@pytest.mark.parametrize("path", ["/v1/responses", "/v1/messages"])
async def test_missing_session_id_omit_keeps_metadata_session_id_on_litellm_metadata_routes(path: str):
updated = await add_litellm_data_to_request(
data={
"model": "gpt-4o",
"input": "hi",
"litellm_session_id": "cust-sess-1",
"metadata": {"session_id": "meta-sess-1"},
},
request=_request_for(path),
user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"),
proxy_config=MagicMock(),
general_settings={"missing_session_id": "omit"},
)
assert updated["litellm_metadata"]["session_id"] == "meta-sess-1"
assert _spend_log_session_id(updated, "litellm_metadata") == "meta-sess-1"
@pytest.mark.asyncio
@pytest.mark.parametrize("path", ["/v1/responses", "/v1/messages"])
async def test_missing_session_id_omit_keeps_body_litellm_session_id_on_litellm_metadata_routes(path: str):
updated = await add_litellm_data_to_request(
data={"model": "gpt-4o", "input": "hi", "litellm_session_id": "cust-sess-1"},
request=_request_for(path),
user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"),
proxy_config=MagicMock(),
general_settings={"missing_session_id": "omit"},
)
assert updated["litellm_metadata"]["session_id"] == "cust-sess-1"
assert _spend_log_session_id(updated, "litellm_metadata") == "cust-sess-1"
@pytest.mark.asyncio
async def test_missing_session_id_omit_ignores_empty_body_litellm_session_id():
updated = await add_litellm_data_to_request(
data={"model": "gpt-4o", "messages": [], "litellm_session_id": ""},
request=_request_for("/v1/chat/completions"),
user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"),
proxy_config=MagicMock(),
general_settings={"missing_session_id": "omit"},
)
assert "session_id" not in updated["metadata"]
assert _spend_log_session_id(updated) is None
@pytest.mark.asyncio
async def test_missing_session_id_generate_reuses_traceparent_trace_id():
"""A W3C traceparent already decides SpendLogs.session_id, so the callback session id must reuse it."""

691
tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py
View file

@ -13,6 +13,7 @@ import asyncio
import json
from copy import deepcopy
import logging
from collections.abc import Iterator
from typing import Any, Callable, Dict, List
from unittest.mock import AsyncMock, MagicMock, patch
@ -26,11 +27,13 @@ from litellm.integrations.custom_guardrail import (
ModifyResponseException,
)
from litellm.integrations.prometheus import PrometheusLogger
from litellm.llms.base_llm.guardrail_translation.utils import stream_item_field
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.common_utils.callback_utils import add_guardrail_to_applied_guardrails_header
from litellm.proxy.utils import ProxyLogging, _streamable_post_call_pipelines, stream_gated_guardrail_names
from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ContentFilterGuardrail
from litellm.types.guardrails import BlockedWord, ContentFilterAction, GuardrailEventHooks
from litellm.types.llms.openai import ResponsesAPIResponse
from litellm.types.proxy.policy_engine.pipeline_types import (
GuardrailPipeline,
PipelineStep,
@ -155,9 +158,7 @@ async def test_execute_guardrail_hook_unknown_hook_type_raises(proxy_logging, ma
@pytest.mark.asyncio
async def test_execute_guardrail_with_load_balancing_routes_through_router(
proxy_logging, make_user_api_key_auth
):
async def test_execute_guardrail_with_load_balancing_routes_through_router(proxy_logging, make_user_api_key_auth):
cb = _make_guardrail()
router = MagicMock()
router.get_available_guardrail = MagicMock(return_value={"callback": cb})
@ -173,9 +174,7 @@ async def test_execute_guardrail_with_load_balancing_routes_through_router(
@pytest.mark.asyncio
async def test_execute_guardrail_with_load_balancing_router_none_raises(
proxy_logging, make_user_api_key_auth
):
async def test_execute_guardrail_with_load_balancing_router_none_raises(proxy_logging, make_user_api_key_auth):
with patch("litellm.proxy.proxy_server.llm_router", None):
with pytest.raises(ValueError, match="Router not initialized"):
await proxy_logging._execute_guardrail_with_load_balancing(
@ -188,9 +187,7 @@ async def test_execute_guardrail_with_load_balancing_router_none_raises(
@pytest.mark.asyncio
async def test_execute_guardrail_with_load_balancing_no_callback_raises(
proxy_logging, make_user_api_key_auth
):
async def test_execute_guardrail_with_load_balancing_no_callback_raises(proxy_logging, make_user_api_key_auth):
router = MagicMock()
router.get_available_guardrail = MagicMock(return_value={"callback": None})
with patch("litellm.proxy.proxy_server.llm_router", router):
@ -210,9 +207,7 @@ async def test_execute_guardrail_with_load_balancing_no_callback_raises(
@pytest.mark.asyncio
async def test_process_guardrail_callback_skipped_when_should_run_false(
proxy_logging, make_user_api_key_auth
):
async def test_process_guardrail_callback_skipped_when_should_run_false(proxy_logging, make_user_api_key_auth):
cb = _make_guardrail()
cb.should_run_guardrail = MagicMock(return_value=False)
out = await proxy_logging._process_guardrail_callback(
@ -226,9 +221,7 @@ async def test_process_guardrail_callback_skipped_when_should_run_false(
@pytest.mark.asyncio
async def test_process_guardrail_callback_returns_data_on_success(
proxy_logging, make_user_api_key_auth, monkeypatch
):
async def test_process_guardrail_callback_returns_data_on_success(proxy_logging, make_user_api_key_auth, monkeypatch):
cb = _make_guardrail()
cb.should_run_guardrail = MagicMock(return_value=True)
proxy_logging._should_use_guardrail_load_balancing = MagicMock(return_value=False)
@ -343,14 +336,14 @@ async def test_maybe_execute_pipelines_no_pipelines_returns_data(proxy_logging,
@pytest.mark.asyncio
async def test_maybe_execute_pipelines_skips_pipelines_with_other_mode(proxy_logging, make_user_api_key_auth, monkeypatch):
async def test_maybe_execute_pipelines_skips_pipelines_with_other_mode(
proxy_logging, make_user_api_key_auth, monkeypatch
):
pipeline = MagicMock()
pipeline.mode = "post_call" # not pre_call
data = {"metadata": {"_guardrail_pipelines": [("p1", pipeline)]}, "model": "m", "messages": []}
executed = MagicMock()
monkeypatch.setattr(
"litellm.proxy.policy_engine.pipeline_executor.PipelineExecutor.execute_steps", executed
)
monkeypatch.setattr("litellm.proxy.policy_engine.pipeline_executor.PipelineExecutor.execute_steps", executed)
out, replacement = await proxy_logging._maybe_execute_pipelines(
data=data,
user_api_key_dict=make_user_api_key_auth(),
@ -538,9 +531,7 @@ def test_handle_pipeline_result_block_enriches_with_guardrail_name_and_mode():
litellm.callbacks = [cb]
try:
with pytest.raises(HTTPException) as info:
ProxyLogging._handle_pipeline_result(
result=result, data={"model": "m"}, policy_name="p"
)
ProxyLogging._handle_pipeline_result(result=result, data={"model": "m"}, policy_name="p")
finally:
litellm.callbacks = saved
@ -652,9 +643,7 @@ async def test_run_guardrail_with_metrics_records_error_and_enriches(monkeypatch
monkeypatch.setattr(litellm, "callbacks", [prom])
with pytest.raises(HTTPException):
await ProxyLogging._run_guardrail_with_metrics(
callback=cb, coro=task(), hook_type="post_call"
)
await ProxyLogging._run_guardrail_with_metrics(callback=cb, coro=task(), hook_type="post_call")
assert detail["guardrail_name"] == "presidio"
recorded = prom._record_guardrail_metrics.call_args.kwargs
@ -682,9 +671,7 @@ def _moderation_guardrail() -> MagicMock:
@pytest.mark.asyncio
async def test_during_call_hook_records_latency_metric(
proxy_logging, make_user_api_key_auth, monkeypatch
):
async def test_during_call_hook_records_latency_metric(proxy_logging, make_user_api_key_auth, monkeypatch):
cb = _moderation_guardrail()
prom = _prometheus_callback()
monkeypatch.setattr(litellm, "callbacks", [prom, cb])
@ -703,9 +690,7 @@ async def test_during_call_hook_records_latency_metric(
@pytest.mark.asyncio
async def test_post_call_success_hook_records_latency_metric(
proxy_logging, make_user_api_key_auth, monkeypatch
):
async def test_post_call_success_hook_records_latency_metric(proxy_logging, make_user_api_key_auth, monkeypatch):
cb = _moderation_guardrail()
prom = _prometheus_callback()
monkeypatch.setattr(litellm, "callbacks", [prom, cb])
@ -733,9 +718,7 @@ async def test_post_call_success_hook_records_latency_metric(
async def test_process_prompt_template_no_op_when_no_prompt_spec(proxy_logging, monkeypatch):
from litellm.proxy.prompts import prompt_registry
monkeypatch.setattr(
prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "resolve_prompt_spec", lambda *a, **kw: None
)
monkeypatch.setattr(prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "resolve_prompt_spec", lambda *a, **kw: None)
data: Dict[str, Any] = {"messages": [{"role": "user"}], "model": "m", "temperature": 0.1}
await proxy_logging._process_prompt_template(
data=data,
@ -760,9 +743,7 @@ async def test_process_prompt_template_applies_when_spec_resolves(proxy_logging,
"get_prompt_callback_for_prompt",
lambda *a, **kw: custom_logger,
)
monkeypatch.setattr(
prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "resolve_prompt_spec", lambda *a, **kw: prompt_spec
)
monkeypatch.setattr(prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "resolve_prompt_spec", lambda *a, **kw: prompt_spec)
logging_obj = MagicMock()
logging_obj.async_get_chat_completion_prompt = AsyncMock(
@ -810,9 +791,7 @@ async def test_process_prompt_template_async_get_prompt_error_raises(proxy_loggi
"get_prompt_callback_for_prompt",
lambda *a, **kw: custom_logger,
)
monkeypatch.setattr(
prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "resolve_prompt_spec", lambda *a, **kw: prompt_spec
)
monkeypatch.setattr(prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "resolve_prompt_spec", lambda *a, **kw: prompt_spec)
logging_obj = MagicMock()
logging_obj.async_get_chat_completion_prompt = AsyncMock(side_effect=RuntimeError("bad prompt"))
with pytest.raises(RuntimeError):
@ -913,9 +892,7 @@ async def test_process_prompt_template_aresponses_swaps_model_and_merges_input(p
"get_prompt_callback_for_prompt",
lambda *a, **kw: custom_logger,
)
monkeypatch.setattr(
prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "resolve_prompt_spec", lambda *a, **kw: prompt_spec
)
monkeypatch.setattr(prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "resolve_prompt_spec", lambda *a, **kw: prompt_spec)
logging_obj = MagicMock()
logging_obj.async_get_chat_completion_prompt = AsyncMock(
@ -966,6 +943,7 @@ def _post_call_pipeline_data(
"metadata": {
"_guardrail_pipelines": [("response-governance", pipeline)],
"_pipeline_managed_guardrails": {guardrail},
"policy_sources": {"response-governance": "model:m"},
},
**extra,
}
@ -1124,9 +1102,7 @@ async def test_pre_call_hook_still_runs_guardrail_managed_only_by_post_call_pipe
},
}
await proxy_logging.pre_call_hook(
user_api_key_dict=make_user_api_key_auth(), data=data, call_type="completion"
)
await proxy_logging.pre_call_hook(user_api_key_dict=make_user_api_key_auth(), data=data, call_type="completion")
assert seen["count"] == 1
@ -1289,11 +1265,7 @@ async def test_post_call_pipeline_block_keeps_guardrail_metadata_writes(
monkeypatch.setattr(
litellm,
"callbacks",
[
BlockingWriterGuardrail(
guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False
)
],
[BlockingWriterGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False)],
)
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
data = _post_call_pipeline_data()
@ -1379,9 +1351,7 @@ async def test_pre_call_pipeline_managed_parallel_guardrail_runs_exactly_once(
},
}
await proxy_logging.pre_call_hook(
user_api_key_dict=make_user_api_key_auth(), data=data, call_type="completion"
)
await proxy_logging.pre_call_hook(user_api_key_dict=make_user_api_key_auth(), data=data, call_type="completion")
assert seen["count"] == 1
@ -1420,10 +1390,295 @@ async def test_streaming_request_whose_pipeline_guardrail_is_missing_streams_ver
assert any("response-governance" in message and "gr-post" in message for message in _warnings(caplog))
def _background_response(status: str, text: str = "") -> ResponsesAPIResponse:
output = (
[{"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": text}]}] if text else []
)
return ResponsesAPIResponse(id="resp_bg", created_at=0, output=output, status=status)
def _output_blocking_callbacks(seen: dict[str, object]) -> list[CustomGuardrail]:
class OutputBlockingGuardrail(CustomGuardrail):
async def async_post_call_success_hook(self, data, user_api_key_dict, response):
seen["response"] = response
raise HTTPException(status_code=400, detail={"error": "output blocked"})
return [
OutputBlockingGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False)
]
@pytest.mark.asyncio
async def test_pre_call_hook_accepts_background_request_with_post_call_pipeline(
proxy_logging, make_user_api_key_auth, monkeypatch, caplog
@pytest.mark.parametrize("pending_status", ["queued", "in_progress"])
async def test_post_call_success_hook_waits_for_pending_background_response_before_running_pipeline(
proxy_logging: ProxyLogging,
make_user_api_key_auth: Callable[..., UserAPIKeyAuth],
monkeypatch: pytest.MonkeyPatch,
caplog: pytest.LogCaptureFixture,
pending_status: str,
) -> None:
seen: dict[str, object] = {}
monkeypatch.setattr(litellm, "callbacks", _output_blocking_callbacks(seen))
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
data = _post_call_pipeline_data(background=True)
response = _background_response(pending_status)
with caplog.at_level(logging.DEBUG, logger="LiteLLM Proxy"):
out = await proxy_logging.post_call_success_hook(
data=data, response=response, user_api_key_dict=make_user_api_key_auth()
)
assert out is response
assert "response" not in seen
assert not _warnings(caplog)
assert any(
"response-governance" in record.getMessage() and pending_status in record.getMessage()
for record in caplog.records
)
@pytest.mark.asyncio
@pytest.mark.parametrize("final_status", ["completed", "incomplete"])
async def test_post_call_success_hook_runs_pipeline_on_retrieved_background_response(
proxy_logging: ProxyLogging,
make_user_api_key_auth: Callable[..., UserAPIKeyAuth],
monkeypatch: pytest.MonkeyPatch,
final_status: str,
) -> None:
seen: dict[str, object] = {}
monkeypatch.setattr(litellm, "callbacks", _output_blocking_callbacks(seen))
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
data = _post_call_pipeline_data()
response = _background_response(final_status, text="kumquat")
with pytest.raises(HTTPException) as info:
await proxy_logging.post_call_success_hook(
data=data, response=response, user_api_key_dict=make_user_api_key_auth()
)
assert info.value.detail["error"] == "output blocked"
assert seen["response"] is response
def _output_passing_callbacks() -> list[CustomGuardrail]:
class OutputPassingGuardrail(CustomGuardrail):
async def async_post_call_success_hook(self, data, user_api_key_dict, response):
return response
return [
OutputPassingGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False)
]
def _claimed_post_call_pipeline_data(
*policy_names: str, extra_guardrails: dict[str, list[str]] | None = None, policy_source: str | None = "model:m"
):
from litellm.proxy.policy_engine.policy_registry import get_policy_registry
step = {"guardrail": "gr-post", "on_pass": "allow", "on_fail": "block"}
get_policy_registry().load_policies(
{
policy_name: {
"guardrails": {"add": ["gr-post", *(extra_guardrails or {}).get(policy_name, [])]},
"pipeline": {"mode": "post_call", "steps": [step]},
}
for policy_name in policy_names
}
)
pipeline = GuardrailPipeline(mode="post_call", steps=[PipelineStep(**step)])
return {
"model": "m",
"messages": [{"role": "user", "content": "hi"}],
"metadata": {
"_guardrail_pipelines": [(policy_name, pipeline) for policy_name in policy_names],
"_pipeline_managed_guardrails": {"gr-post"},
"applied_policies": list(policy_names),
"applied_guardrails": ["gr-post", *(g for gs in (extra_guardrails or {}).values() for g in gs)],
"policy_sources": {policy_name: policy_source for policy_name in policy_names if policy_source is not None},
},
}
@pytest.fixture
def clear_policy_registry() -> Iterator[None]:
from litellm.proxy.policy_engine.policy_registry import get_policy_registry
yield
get_policy_registry().clear()
@pytest.mark.asyncio
async def test_pending_background_response_withdraws_the_deferred_policy_claims(
proxy_logging: ProxyLogging,
make_user_api_key_auth: Callable[..., UserAPIKeyAuth],
monkeypatch: pytest.MonkeyPatch,
clear_policy_registry: None,
) -> None:
monkeypatch.setattr(litellm, "callbacks", _output_blocking_callbacks({}))
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
data = _claimed_post_call_pipeline_data("response-governance")
out = await proxy_logging.post_call_success_hook(
data=data, response=_background_response("queued"), user_api_key_dict=make_user_api_key_auth()
)
assert out.status == "queued"
assert "applied_policies" not in data["metadata"]
assert "policy_sources" not in data["metadata"]
assert "applied_guardrails" not in data["metadata"]
@pytest.mark.asyncio
async def test_pending_background_response_warns_when_the_deferred_policy_was_matched_through_a_tag(
proxy_logging: ProxyLogging,
make_user_api_key_auth: Callable[..., UserAPIKeyAuth],
monkeypatch: pytest.MonkeyPatch,
clear_policy_registry: None,
caplog: pytest.LogCaptureFixture,
) -> None:
monkeypatch.setattr(litellm, "callbacks", _output_blocking_callbacks({}))
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
data = _claimed_post_call_pipeline_data("response-governance", policy_source="tag:governed+model:m")
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
out = await proxy_logging.post_call_success_hook(
data=data, response=_background_response("queued"), user_api_key_dict=make_user_api_key_auth()
)
assert out.status == "queued"
assert "policy_sources" not in data["metadata"]
assert [message for message in _warnings(caplog) if "through a request tag" in message] == [
"Policy engine: background response resp_bg matched post_call policies through a request tag at submit; "
"retrieval re-matches only the key, team, and model scopes, so a tag carried in the request body "
"does not govern the completed response: response-governance"
]
@pytest.mark.asyncio
async def test_pending_background_response_warns_when_the_deferred_policy_came_from_the_request_body(
proxy_logging: ProxyLogging,
make_user_api_key_auth: Callable[..., UserAPIKeyAuth],
monkeypatch: pytest.MonkeyPatch,
clear_policy_registry: None,
caplog: pytest.LogCaptureFixture,
) -> None:
monkeypatch.setattr(litellm, "callbacks", _output_blocking_callbacks({}))
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
data = _claimed_post_call_pipeline_data("body-governance", policy_source=None)
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
out = await proxy_logging.post_call_success_hook(
data=data, response=_background_response("queued"), user_api_key_dict=make_user_api_key_auth()
)
assert out.status == "queued"
assert "policy_sources" not in data["metadata"]
assert _warnings(caplog) == [
"Policy engine: background response resp_bg matched post_call policies through the request body's policies "
"list at submit; retrieval carries no request body, so those policies do not govern the completed "
"response: body-governance"
]
@pytest.mark.asyncio
async def test_pending_background_response_matched_through_its_model_does_not_warn(
proxy_logging: ProxyLogging,
make_user_api_key_auth: Callable[..., UserAPIKeyAuth],
monkeypatch: pytest.MonkeyPatch,
clear_policy_registry: None,
caplog: pytest.LogCaptureFixture,
) -> None:
monkeypatch.setattr(litellm, "callbacks", _output_blocking_callbacks({}))
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
await proxy_logging.post_call_success_hook(
data=_claimed_post_call_pipeline_data("response-governance"),
response=_background_response("queued"),
user_api_key_dict=make_user_api_key_auth(),
)
assert _warnings(caplog) == []
@pytest.mark.asyncio
async def test_pending_background_response_keeps_the_claim_of_a_policy_that_runs_outside_its_pipeline(
proxy_logging: ProxyLogging,
make_user_api_key_auth: Callable[..., UserAPIKeyAuth],
monkeypatch: pytest.MonkeyPatch,
clear_policy_registry: None,
) -> None:
monkeypatch.setattr(litellm, "callbacks", _output_blocking_callbacks({}))
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
data = _claimed_post_call_pipeline_data(
"input-and-output-governance",
"response-governance",
extra_guardrails={"input-and-output-governance": ["gr-pre"]},
)
await proxy_logging.post_call_success_hook(
data=data, response=_background_response("in_progress"), user_api_key_dict=make_user_api_key_auth()
)
assert data["metadata"]["applied_policies"] == ["input-and-output-governance"]
assert data["metadata"]["applied_guardrails"] == ["gr-pre"]
assert data["metadata"]["policy_sources"] == {"input-and-output-governance": "model:m"}
@pytest.mark.asyncio
async def test_pending_background_response_keeps_the_claim_of_a_default_on_guardrail_that_ran_pre_call(
proxy_logging: ProxyLogging,
make_user_api_key_auth: Callable[..., UserAPIKeyAuth],
monkeypatch: pytest.MonkeyPatch,
clear_policy_registry: None,
) -> None:
class DualStageGuardrail(CustomGuardrail):
async def async_post_call_success_hook(self, data, user_api_key_dict, response):
return response
monkeypatch.setattr(
litellm,
"callbacks",
[DualStageGuardrail(guardrail_name="gr-post", event_hook=["pre_call", "post_call"], default_on=True)],
)
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
data = _claimed_post_call_pipeline_data("response-governance")
await proxy_logging.post_call_success_hook(
data=data, response=_background_response("queued"), user_api_key_dict=make_user_api_key_auth()
)
assert "applied_policies" not in data["metadata"]
assert data["metadata"]["applied_guardrails"] == ["gr-post"]
@pytest.mark.asyncio
async def test_retrieved_background_response_keeps_the_policy_claim_once_its_pipeline_ran(
proxy_logging: ProxyLogging,
make_user_api_key_auth: Callable[..., UserAPIKeyAuth],
monkeypatch: pytest.MonkeyPatch,
clear_policy_registry: None,
) -> None:
monkeypatch.setattr(litellm, "callbacks", _output_passing_callbacks())
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
data = _claimed_post_call_pipeline_data("response-governance")
await proxy_logging.post_call_success_hook(
data=data, response=_background_response("completed", text="fine"), user_api_key_dict=make_user_api_key_auth()
)
assert data["metadata"]["applied_policies"] == ["response-governance"]
assert data["metadata"]["policy_sources"] == {"response-governance": "model:m"}
assert data["metadata"]["applied_guardrails"] == ["gr-post"]
@pytest.mark.asyncio
async def test_pre_call_hook_stays_quiet_on_background_request_with_post_call_pipeline(
proxy_logging: ProxyLogging,
make_user_api_key_auth: Callable[..., UserAPIKeyAuth],
monkeypatch: pytest.MonkeyPatch,
caplog: pytest.LogCaptureFixture,
) -> None:
monkeypatch.setattr(litellm, "callbacks", [])
data = _post_call_pipeline_data(background=True)
@ -1437,36 +1692,7 @@ async def test_pre_call_hook_accepts_background_request_with_post_call_pipeline(
assert out is not None
assert out.get("background") is True
assert any("response-governance" in message and "background" in message for message in _warnings(caplog))
@pytest.mark.asyncio
async def test_pre_call_hook_stays_quiet_on_background_request_without_post_call_pipeline(
proxy_logging, make_user_api_key_auth, monkeypatch, caplog
):
seen: Dict[str, Any] = {}
monkeypatch.setattr(litellm, "callbacks", [_unified_stream_guardrail(seen)])
pre_call = GuardrailPipeline(mode="pre_call", steps=[PipelineStep(guardrail="gr-post", on_fail="block")])
data = {
"model": "m",
"messages": [{"role": "user", "content": "hi"}],
"background": True,
"metadata": {
"_guardrail_pipelines": [("request-governance", pre_call)],
"_pipeline_managed_guardrails": {"gr-post"},
},
}
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
out = await proxy_logging.pre_call_hook(
user_api_key_dict=make_user_api_key_auth(),
data=data,
call_type="aresponses",
guardrails_only=True,
)
assert out is not None
assert not any("background" in message for message in _warnings(caplog))
assert not _warnings(caplog)
# ---------------------------------------------------------------------------
@ -1577,7 +1803,9 @@ def test_streamable_post_call_pipelines_keeps_hook_guardrails_and_drops_iterator
steps=[PipelineStep(guardrail="gr-post", on_fail="next"), PipelineStep(guardrail="gr-iterator", on_fail="block")],
)
pre_call = GuardrailPipeline(mode="pre_call", steps=[PipelineStep(guardrail="gr-iterator", on_fail="block")])
data = {"metadata": {"_guardrail_pipelines": [("governed", governed), ("ungoverned", ungoverned), ("req", pre_call)]}}
data = {
"metadata": {"_guardrail_pipelines": [("governed", governed), ("ungoverned", ungoverned), ("req", pre_call)]}
}
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
streamable = _streamable_post_call_pipelines(data, make_user_api_key_auth(request_route="/v1/chat/completions"))
@ -1589,7 +1817,7 @@ def test_streamable_post_call_pipelines_keeps_hook_guardrails_and_drops_iterator
@pytest.mark.parametrize(
"request_route",
[None, "/v1/completions", "/v1beta/models/gemini-2.5-flash:streamGenerateContent", "/a2a/agent"],
["/v1/completions", "/v1beta/models/gemini-2.5-flash:streamGenerateContent", "/a2a/agent"],
)
def test_streamable_post_call_pipelines_keeps_legacy_hooks_off_routes_that_assemble_no_response(
make_user_api_key_auth, monkeypatch, caplog, request_route
@ -1993,7 +2221,9 @@ def _rewriting_stream_guardrail(transform: Callable[[Dict[str, Any]], Dict[str,
async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
return {**inputs, **transform(inputs)}
return RewritingStreamGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False)
return RewritingStreamGuardrail(
guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False
)
def _tool_call_stream_chunks() -> List[Any]:
@ -2017,7 +2247,7 @@ def _echoed_tool_call_dicts(arguments: str) -> List[Dict[str, Any]]:
@pytest.mark.asyncio
@pytest.mark.parametrize("on_fail, on_error", [("block", None), ("next", "next")])
async def test_streaming_iterator_hook_pipeline_releases_originals_on_runtime_tool_call_rewrite(
async def test_streaming_iterator_hook_pipeline_delivers_runtime_tool_call_rewrite(
proxy_logging, make_user_api_key_auth, monkeypatch, on_fail, on_error, caplog
):
transform = lambda inputs: {"tool_calls": _echoed_tool_call_dicts('{"ssn": "[MASKED]"}')} # noqa: E731
@ -2036,9 +2266,12 @@ async def test_streaming_iterator_hook_pipeline_releases_originals_on_runtime_to
delivered.append(item)
assert len(delivered) == 2
assert delivered[0].choices[0].delta.tool_calls[0].function.arguments == '{"ssn": "123"}'
delivered_tool_call = delivered[0].choices[0].delta.tool_calls[0]
assert delivered_tool_call.function.arguments == '{"ssn": "[MASKED]"}'
assert delivered_tool_call.function.name == "lookup"
assert delivered_tool_call.id == "call_1"
assert delivered[1].choices[0].finish_reason == "tool_calls"
assert any("'gr-post'" in message and "discarded" in message for message in _warnings(caplog))
assert not any("discarded" in message for message in _warnings(caplog))
@pytest.mark.asyncio
@ -2146,37 +2379,97 @@ async def test_streaming_iterator_hook_pipeline_releases_stream_echoed_in_anothe
@pytest.mark.asyncio
async def test_streaming_iterator_hook_pipeline_releases_originals_on_unresolvable_response_shape(
async def test_streaming_iterator_hook_skips_pipeline_and_warns_without_request_route(
proxy_logging, make_user_api_key_auth, monkeypatch, caplog
):
seen: Dict[str, Any] = {}
monkeypatch.setattr(litellm, "callbacks", [_unified_stream_guardrail(seen)])
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
data = _post_call_pipeline_data(stream=True)
chunks = [object(), object()]
delivered: List[Any] = []
chunks = _stream_chunks()
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
async for item in proxy_logging.async_post_call_streaming_iterator_hook(
user_api_key_dict=make_user_api_key_auth(),
response=_async_chunk_iter(chunks),
request_data=data,
):
delivered.append(item)
delivered = [
item
async for item in proxy_logging.async_post_call_streaming_iterator_hook(
user_api_key_dict=make_user_api_key_auth(),
response=_async_chunk_iter(chunks),
request_data=data,
)
]
assert [item is chunk for item, chunk in zip(delivered, chunks)] == [True, True]
assert len(delivered) == 2
assert seen.get("count") is None
assert any("response-governance" in message and "shape" in message for message in _warnings(caplog))
assert any("response-governance" in message and "route None" in message for message in _warnings(caplog))
@pytest.mark.asyncio
async def test_per_chunk_streaming_hook_runs_pipeline_managed_guardrail_without_request_route(
proxy_logging, make_user_api_key_auth, monkeypatch
):
seen: Dict[str, Any] = {}
class UnifiedRecordingGuardrail(CustomGuardrail):
async def async_post_call_streaming_hook(self, user_api_key_dict, response):
seen[self.guardrail_name] = seen.get(self.guardrail_name, 0) + 1
return None
async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
return inputs
monkeypatch.setattr(
litellm,
"callbacks",
[UnifiedRecordingGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=True)],
)
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
data = _post_call_pipeline_data(stream=True)
result = await proxy_logging.async_post_call_streaming_hook(
data=data,
response=_stream_chunks()[0],
user_api_key_dict=make_user_api_key_auth(),
)
assert result is not None
assert seen["gr-post"] == 1
def _anthropic_sse_chunks() -> List[bytes]:
events = [
("message_start", {"type": "message_start", "message": {"id": "msg_1", "type": "message", "role": "assistant", "model": "m", "content": [], "stop_reason": None, "usage": {"input_tokens": 1, "output_tokens": 0}}}),
("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}),
("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "hello world"}}),
(
"message_start",
{
"type": "message_start",
"message": {
"id": "msg_1",
"type": "message",
"role": "assistant",
"model": "m",
"content": [],
"stop_reason": None,
"usage": {"input_tokens": 1, "output_tokens": 0},
},
},
),
(
"content_block_start",
{"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}},
),
(
"content_block_delta",
{"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "hello world"}},
),
("content_block_stop", {"type": "content_block_stop", "index": 0}),
("message_delta", {"type": "message_delta", "delta": {"stop_reason": "end_turn", "stop_sequence": None}, "usage": {"output_tokens": 2}}),
(
"message_delta",
{
"type": "message_delta",
"delta": {"stop_reason": "end_turn", "stop_sequence": None},
"usage": {"output_tokens": 2},
},
),
("message_stop", {"type": "message_stop"}),
]
return [f"event: {name}\ndata: {json.dumps(payload)}\n\n".encode() for name, payload in events]
@ -2243,7 +2536,13 @@ async def test_streaming_iterator_hook_pipeline_delivers_text_rewrite_on_anthrop
assert "hello [MASKED]" in raw
assert "hello world" not in raw
assert raw.count("event: content_block_delta") == 1
for expected_event in ("message_start", "content_block_start", "content_block_stop", "message_delta", "message_stop"):
for expected_event in (
"message_start",
"content_block_start",
"content_block_stop",
"message_delta",
"message_stop",
):
assert f"event: {expected_event}" in raw
@ -2316,9 +2615,7 @@ async def test_per_chunk_streaming_hook_skips_pipeline_managed_guardrail(
managed = UnifiedRecordingGuardrail(
guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=True
)
free = RecordingGuardrail(
guardrail_name="gr-free", event_hook=GuardrailEventHooks.post_call, default_on=True
)
free = RecordingGuardrail(guardrail_name="gr-free", event_hook=GuardrailEventHooks.post_call, default_on=True)
monkeypatch.setattr(litellm, "callbacks", [managed, free])
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
data = _post_call_pipeline_data(stream=True)
@ -2363,3 +2660,177 @@ async def test_per_chunk_streaming_hook_runs_guardrail_whose_pipeline_cannot_str
assert result is not None
assert seen["count"] == 1
assert seen["response"] == "hello "
def _mask_tool_call_arguments(inputs: Dict[str, Any]) -> Dict[str, Any]:
return {
"tool_calls": [
{
"id": stream_item_field(tool_call, "id"),
"type": "function",
"function": {
"name": stream_item_field(stream_item_field(tool_call, "function"), "name"),
"arguments": '{"fruit": "[MASKED]"}',
},
}
for tool_call in inputs.get("tool_calls", [])
]
}
def _anthropic_tool_use_sse_chunks() -> List[bytes]:
events = [
("message_start", {"type": "message_start", "message": {"id": "msg_1", "type": "message", "role": "assistant", "model": "m", "content": [], "stop_reason": None, "usage": {"input_tokens": 1, "output_tokens": 0}}}),
("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "tool_use", "id": "toolu_1", "name": "lookup_fruit", "input": {}}}),
("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": '{"fruit": "persim'}}),
("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": 'mon"}'}}),
("content_block_stop", {"type": "content_block_stop", "index": 0}),
("message_delta", {"type": "message_delta", "delta": {"stop_reason": "tool_use", "stop_sequence": None}, "usage": {"output_tokens": 2}}),
("message_stop", {"type": "message_stop"}),
]
return [f"event: {name}\ndata: {json.dumps(payload)}\n\n".encode() for name, payload in events]
@pytest.mark.asyncio
async def test_streaming_iterator_hook_pipeline_delivers_tool_use_rewrite_on_anthropic_sse(
proxy_logging, make_user_api_key_auth, monkeypatch
):
monkeypatch.setattr(litellm, "callbacks", [_rewriting_stream_guardrail(_mask_tool_call_arguments)])
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
data = _post_call_pipeline_data(stream=True)
delivered = [
item
async for item in proxy_logging.async_post_call_streaming_iterator_hook(
user_api_key_dict=make_user_api_key_auth(request_route="/v1/messages"),
response=_async_chunk_iter(_anthropic_tool_use_sse_chunks()),
request_data=data,
)
]
raw = b"".join(delivered).decode()
assert '{\\"fruit\\": \\"[MASKED]\\"}' in raw
assert "persim" not in raw
assert '"name": "lookup_fruit"' in raw and '"id": "toolu_1"' in raw
assert '"stop_reason": "tool_use"' in raw
assert raw.count("event: content_block_delta") == 2
def _responses_function_call_events() -> List[Dict[str, Any]]:
def item(arguments: str, status: str) -> Dict[str, Any]:
return {
"type": "function_call",
"id": "fc_1",
"call_id": "call_1",
"name": "lookup_fruit",
"arguments": arguments,
"status": status,
}
return [
{"type": "response.output_item.added", "output_index": 0, "item": item("", "in_progress")},
{"type": "response.function_call_arguments.delta", "item_id": "fc_1", "output_index": 0, "delta": '{"fruit":'},
{"type": "response.function_call_arguments.delta", "item_id": "fc_1", "output_index": 0, "delta": ' "persimmon"}'},
{"type": "response.function_call_arguments.done", "item_id": "fc_1", "output_index": 0, "arguments": '{"fruit": "persimmon"}'},
{"type": "response.output_item.done", "output_index": 0, "item": item('{"fruit": "persimmon"}', "completed")},
{
"type": "response.completed",
"response": {"id": "resp_1", "created_at": 1, "model": "m", "output": [item('{"fruit": "persimmon"}', "completed")], "status": "completed"},
},
]
@pytest.mark.asyncio
async def test_streaming_iterator_hook_pipeline_delivers_function_call_rewrite_on_responses_events(
proxy_logging, make_user_api_key_auth, monkeypatch
):
monkeypatch.setattr(litellm, "callbacks", [_rewriting_stream_guardrail(_mask_tool_call_arguments)])
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
data = _post_call_pipeline_data(stream=True)
delivered = [
item
async for item in proxy_logging.async_post_call_streaming_iterator_hook(
user_api_key_dict=make_user_api_key_auth(request_route="/v1/responses"),
response=_async_chunk_iter(_responses_function_call_events()),
request_data=data,
)
]
assert [event["type"] for event in delivered] == [event["type"] for event in _responses_function_call_events()]
assert [event["delta"] for event in delivered if event["type"] == "response.function_call_arguments.delta"] == ['{"fruit": "[MASKED]"}', ""]
assert delivered[3]["arguments"] == '{"fruit": "[MASKED]"}'
assert delivered[4]["item"]["arguments"] == '{"fruit": "[MASKED]"}'
assert delivered[5]["response"]["output"][0]["arguments"] == '{"fruit": "[MASKED]"}'
assert "persimmon" not in json.dumps(delivered)
def _drop_tool_calls(inputs: Dict[str, Any]) -> Dict[str, Any]:
return {"tool_calls": []}
@pytest.mark.asyncio
async def test_streaming_iterator_hook_pipeline_discards_dropped_tool_call_on_chat_chunks(
proxy_logging, make_user_api_key_auth, monkeypatch, caplog
):
monkeypatch.setattr(litellm, "callbacks", [_rewriting_stream_guardrail(_drop_tool_calls)])
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
data = _post_call_pipeline_data(stream=True)
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
delivered = [
item
async for item in proxy_logging.async_post_call_streaming_iterator_hook(
user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"),
response=_async_chunk_iter(_tool_call_stream_chunks()),
request_data=data,
)
]
assert delivered[0].choices[0].delta.tool_calls[0].function.arguments == '{"ssn": "123"}'
assert delivered[1].choices[0].finish_reason == "tool_calls"
assert any("'gr-post'" in message and "discarded" in message for message in _warnings(caplog))
@pytest.mark.asyncio
async def test_streaming_iterator_hook_pipeline_discards_dropped_tool_call_on_anthropic_sse(
proxy_logging, make_user_api_key_auth, monkeypatch, caplog
):
monkeypatch.setattr(litellm, "callbacks", [_rewriting_stream_guardrail(_drop_tool_calls)])
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
data = _post_call_pipeline_data(stream=True)
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
delivered = [
item
async for item in proxy_logging.async_post_call_streaming_iterator_hook(
user_api_key_dict=make_user_api_key_auth(request_route="/v1/messages"),
response=_async_chunk_iter(_anthropic_tool_use_sse_chunks()),
request_data=data,
)
]
assert delivered == _anthropic_tool_use_sse_chunks()
assert any("'gr-post'" in message and "discarded" in message for message in _warnings(caplog))
@pytest.mark.asyncio
async def test_streaming_iterator_hook_pipeline_discards_dropped_tool_call_on_responses_events(
proxy_logging, make_user_api_key_auth, monkeypatch, caplog
):
monkeypatch.setattr(litellm, "callbacks", [_rewriting_stream_guardrail(_drop_tool_calls)])
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
data = _post_call_pipeline_data(stream=True)
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
delivered = [
item
async for item in proxy_logging.async_post_call_streaming_iterator_hook(
user_api_key_dict=make_user_api_key_auth(request_route="/v1/responses"),
response=_async_chunk_iter(_responses_function_call_events()),
request_data=data,
)
]
assert delivered == _responses_function_call_events()
assert any("'gr-post'" in message and "discarded" in message for message in _warnings(caplog))

1
ui/litellm-dashboard/src/app/(dashboard)/hooks/healthReadiness/useHealthReadinessDetails.ts
View file

@ -14,6 +14,7 @@ export interface HealthReadinessDetailsResponse {
log_level?: string;
is_detailed_debug?: boolean;
show_no_redis_warning?: boolean;
show_env_credential_login_warning?: boolean;
}
const fetchHealthReadinessDetails = async (accessToken: string): Promise<HealthReadinessDetailsResponse> => {

4
ui/litellm-dashboard/src/app/(dashboard)/layout.test.tsx
View file

@ -29,6 +29,10 @@ vi.mock("@/components/NoRedisWarningBanner", () => ({
NoRedisWarningBanner: () => null,
}));
vi.mock("@/components/EnvCredentialLoginWarningBanner", () => ({
EnvCredentialLoginWarningBanner: () => null,
}));
vi.mock("@/components/LicenseExpiryBanner", () => ({
LicenseExpiryBanner: () => null,
}));

3
ui/litellm-dashboard/src/app/(dashboard)/layout.tsx
View file

@ -10,6 +10,7 @@ import SidebarProvider from "@/app/(dashboard)/components/SidebarProvider";
import { useRouter, useSearchParams } from "next/navigation";
import { DebugWarningBanner } from "@/components/DebugWarningBanner";
import { NoRedisWarningBanner } from "@/components/NoRedisWarningBanner";
import { EnvCredentialLoginWarningBanner } from "@/components/EnvCredentialLoginWarningBanner";
import { LicenseExpiryBanner } from "@/components/LicenseExpiryBanner";
import { UserBanner } from "@/components/UserBanner";
import { uiHref } from "@/utils/uiHref";
@ -113,6 +114,7 @@ function DashboardShell({ children }: { children: React.ReactNode }) {
<Navbar accessToken={accessToken} isPublicPage={false} />
<DebugWarningBanner accessToken={accessToken} />
<NoRedisWarningBanner accessToken={accessToken} />
<EnvCredentialLoginWarningBanner accessToken={accessToken} />
<LicenseExpiryBanner accessToken={accessToken} />
<UserBanner accessToken={accessToken} />
<main className="flex min-h-0 flex-1 overflow-hidden">
@ -132,6 +134,7 @@ function DashboardShell({ children }: { children: React.ReactNode }) {
<DashboardHeader />
<DebugWarningBanner accessToken={accessToken} />
<NoRedisWarningBanner accessToken={accessToken} />
<EnvCredentialLoginWarningBanner accessToken={accessToken} />
<LicenseExpiryBanner accessToken={accessToken} />
<UserBanner accessToken={accessToken} />
<main className="min-w-0 flex-1 overflow-y-auto">{children}</main>

53
ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.test.tsx
View file

@ -5,6 +5,7 @@ import { beforeEach, describe, expect, it, vi } from "vitest";
import ChatUI from "./ChatUI";
import * as fetchModelsModule from "@/components/llm_calls/fetch_models";
import { makeOpenAIChatCompletionRequest } from "@/components/llm_calls/chat_completion";
import { makeAnthropicMessagesRequest } from "../../llm_calls/anthropic_messages";
vi.mock("@/components/llm_calls/fetch_models", () => ({
fetchAvailableModels: vi.fn(),
@ -14,6 +15,10 @@ vi.mock("@/components/llm_calls/chat_completion", () => ({
makeOpenAIChatCompletionRequest: vi.fn().mockResolvedValue(undefined),
}));
vi.mock("../../llm_calls/anthropic_messages", () => ({
makeAnthropicMessagesRequest: vi.fn().mockResolvedValue(undefined),
}));
vi.mock("@/components/networking", () => ({
tagListCall: vi.fn().mockResolvedValue({}),
vectorStoreListCall: vi.fn().mockResolvedValue({ data: [] }),
@ -32,6 +37,8 @@ beforeEach(() => {
const CHAT_REQUEST_ARG_COUNT = 26;
const STREAMING_ENABLED_ARG_INDEX = 25;
const MESSAGES_REQUEST_ARG_COUNT = 19;
const MESSAGES_STREAMING_ENABLED_ARG_INDEX = 18;
async function openComboboxByPlaceholder(placeholder: string) {
const user = userEvent.setup();
@ -378,6 +385,52 @@ describe("ChatUI", () => {
expect(requestArgs[STREAMING_ENABLED_ARG_INDEX]).toBe(false);
});
it("should send the /v1/messages request non-streaming after Stream responses is unchecked", async () => {
const user = userEvent.setup();
render(
<ChatUI
accessToken="1234567890"
token="1234567890"
userRole="user"
userID="1234567890"
disabledPersonalKeyCreation={false}
/>,
);
await waitFor(() => {
expect(screen.getByText("Test Key")).toBeInTheDocument();
});
await selectComboboxOption("Select an endpoint", "/v1/messages");
await selectComboboxOption("Select a Model", "Model 1");
await user.click(await screen.findByTestId("model-settings-button"));
const streamingCheckbox = await screen.findByRole("checkbox", { name: /Stream responses/i });
expect(streamingCheckbox).toBeChecked();
await user.click(streamingCheckbox);
await waitFor(() => {
expect(screen.getByRole("checkbox", { name: /Stream responses/i })).not.toBeChecked();
});
const messageInput = screen.getByPlaceholderText("Type your message... (Shift+Enter for new line)");
await act(async () => {
fireEvent.change(messageInput, { target: { value: "hello" } });
});
await act(async () => {
fireEvent.keyDown(messageInput, { key: "Enter", code: "Enter" });
});
await waitFor(() => {
expect(makeAnthropicMessagesRequest).toHaveBeenCalledTimes(1);
});
const requestArgs = vi.mocked(makeAnthropicMessagesRequest).mock.calls[0];
expect(requestArgs).toHaveLength(MESSAGES_REQUEST_ARG_COUNT);
expect(requestArgs[MESSAGES_STREAMING_ENABLED_ARG_INDEX]).toBe(false);
});
it("should force streaming in simplified mode even when the playground setting is off", async () => {
sessionStorage.setItem("streamingEnabled", "false");

6
ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx
View file

@ -1025,6 +1025,7 @@ const ChatUI: React.FC<ChatUIProps> = ({
mcpServers,
mcpServerToolRestrictions,
mcpToolsets,
streamingEnabled,
);
} else if (endpointType === EndpointType.EMBEDDINGS) {
await makeOpenAIEmbeddingsRequest(
@ -1174,7 +1175,10 @@ const ChatUI: React.FC<ChatUIProps> = ({
return !model.mode || model.mode === "chat";
};
const supportsStreamingToggle = endpointType === EndpointType.CHAT || endpointType === EndpointType.RESPONSES;
const supportsStreamingToggle =
endpointType === EndpointType.CHAT ||
endpointType === EndpointType.RESPONSES ||
endpointType === EndpointType.ANTHROPIC_MESSAGES;
const modelsForEndpoint = useMemo(
() => filterModelsForEndpoint(modelInfo, endpointType as EndpointType),
[modelInfo, endpointType],

66
ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/anthropic_messages.test.tsx
View file

@ -7,13 +7,27 @@ vi.mock("@/components/networking", () => ({
}));
const mockMessagesStream = vi.fn();
const mockMessagesCreate = vi.fn();
vi.mock("@anthropic-ai/sdk", () => ({
default: vi.fn(function () {
return { messages: { stream: mockMessagesStream } };
return { messages: { stream: mockMessagesStream, create: mockMessagesCreate } };
}),
}));
const NON_STREAMING_ARGS = [
undefined, // traceId
undefined, // vector_store_ids
undefined, // guardrails
undefined, // policies
undefined, // selectedMCPServers
undefined, // customBaseUrl
undefined, // mcpServers
undefined, // mcpServerToolRestrictions
undefined, // mcpToolsets
false, // streamingEnabled
] as const;
describe("anthropic_messages prompt cache usage", () => {
const captureUsage = async (usage: Record<string, unknown>): Promise<TokenUsage> => {
async function* mockStream() {
@ -59,3 +73,53 @@ describe("anthropic_messages prompt cache usage", () => {
expect(usageData.promptTokens).toBe(5000);
});
});
describe("anthropic_messages non-streaming", () => {
afterEach(() => {
vi.clearAllMocks();
});
it("sends stream:false through messages.create and renders the full reply at once", async () => {
mockMessagesCreate.mockResolvedValue({
content: [
{ type: "thinking", thinking: "considering" },
{ type: "text", text: "OK" },
],
usage: { input_tokens: 12, output_tokens: 3, cache_read_input_tokens: 7 },
});
const updateTextUI = vi.fn();
const onReasoningContent = vi.fn();
const onUsageData = vi.fn();
await makeAnthropicMessagesRequest(
[{ role: "user", content: "Hello" }],
updateTextUI,
"claude-haiku-4-5",
"test-token",
undefined,
undefined,
onReasoningContent,
undefined,
onUsageData,
...NON_STREAMING_ARGS,
);
expect(mockMessagesStream).not.toHaveBeenCalled();
expect(mockMessagesCreate).toHaveBeenCalledTimes(1);
expect(mockMessagesCreate.mock.calls[0][0]).toMatchObject({ model: "claude-haiku-4-5", stream: false });
expect(updateTextUI).toHaveBeenCalledWith("assistant", "OK", "claude-haiku-4-5");
expect(onReasoningContent).toHaveBeenCalledWith("considering");
const expectedUsage: TokenUsage = { completionTokens: 3, promptTokens: 12, totalTokens: 15, cacheReadTokens: 7 };
expect(onUsageData).toHaveBeenCalledWith(expectedUsage);
});
it("keeps streaming as the default when the flag is omitted", async () => {
async function* emptyStream() {}
mockMessagesStream.mockReturnValue(emptyStream());
await makeAnthropicMessagesRequest([{ role: "user", content: "Hello" }], vi.fn(), "claude-haiku-4-5", "test-token");
expect(mockMessagesCreate).not.toHaveBeenCalled();
expect(mockMessagesStream.mock.calls[0][0]).toMatchObject({ stream: true });
});
});

33
ui/litellm-dashboard/src/app/(dashboard)/playground/llm_calls/anthropic_messages.tsx
View file

@ -7,6 +7,13 @@ import { getProxyBaseUrl } from "@/components/networking";
import { toast } from "@/lib/toast";
import { extractPromptCacheTokens } from "@/utils/promptCacheUsage";
const toTokenUsage = (usage: Anthropic.Usage): TokenUsage => ({
completionTokens: usage.output_tokens,
promptTokens: usage.input_tokens,
totalTokens: usage.input_tokens + usage.output_tokens,
...extractPromptCacheTokens(usage),
});
export async function makeAnthropicMessagesRequest(
messages: MessageType[],
updateTextUI: (role: string, delta: string, model?: string) => void,
@ -26,6 +33,7 @@ export async function makeAnthropicMessagesRequest(
mcpServers?: MCPServer[],
mcpServerToolRestrictions?: Record<string, string[]>,
mcpToolsets?: MCPToolset[],
streamingEnabled: boolean = true,
) {
if (!accessToken) {
throw new Error("Virtual Key is required");
@ -58,7 +66,7 @@ export async function makeAnthropicMessagesRequest(
const requestBody: any = {
model: selectedModel,
messages: messages.map((m) => ({ role: m.role, content: m.content })),
stream: true,
stream: streamingEnabled,
max_tokens: 1024,
// @ts-ignore - litellm specific parameter
litellm_trace_id: traceId,
@ -74,6 +82,20 @@ export async function makeAnthropicMessagesRequest(
if (vector_store_ids) requestBody.vector_store_ids = vector_store_ids;
if (guardrails) requestBody.guardrails = guardrails;
if (policies) requestBody.policies = policies;
if (!streamingEnabled) {
const message: Anthropic.Message = await client.messages.create({ ...requestBody, stream: false }, { signal });
for (const block of message.content) {
if (block.type === "text") {
updateTextUI("assistant", block.text, selectedModel);
} else if (block.type === "thinking" && onReasoningContent) {
onReasoningContent(block.thinking);
}
}
onUsageData?.(toTokenUsage(message.usage));
return;
}
// Use the streaming helper method for cleaner async iteration
// @ts-ignore - The SDK types might not include all litellm-specific parameters
const stream = client.messages.stream(requestBody, { signal });
@ -105,14 +127,7 @@ export async function makeAnthropicMessagesRequest(
// Process usage data from message_delta events
if (messageStreamEvent.type === "message_delta" && (messageStreamEvent as any).usage && onUsageData) {
const usage = (messageStreamEvent as any).usage;
const usageData: TokenUsage = {
completionTokens: usage.output_tokens,
promptTokens: usage.input_tokens,
totalTokens: usage.input_tokens + usage.output_tokens,
...extractPromptCacheTokens(usage),
};
onUsageData(usageData);
onUsageData(toTokenUsage((messageStreamEvent as any).usage));
}
}
} catch (error) {

76
ui/litellm-dashboard/src/components/EnvCredentialLoginWarningBanner.test.tsx Normal file
View file

@ -0,0 +1,76 @@
import { renderWithProviders, screen } from "../../tests/test-utils";
import { vi } from "vitest";
import { EnvCredentialLoginWarningBanner } from "./EnvCredentialLoginWarningBanner";
import type { HealthReadinessDetailsResponse } from "@/app/(dashboard)/hooks/healthReadiness/useHealthReadinessDetails";
import type { UseQueryResult } from "@tanstack/react-query";
vi.mock("@/app/(dashboard)/hooks/healthReadiness/useHealthReadinessDetails", () => ({
useHealthReadinessDetails: vi.fn(),
}));
vi.mock("@/contexts/AuthContext", () => ({
useAuth: vi.fn(),
}));
import { useHealthReadinessDetails } from "@/app/(dashboard)/hooks/healthReadiness/useHealthReadinessDetails";
import { useAuth } from "@/contexts/AuthContext";
const mockDetails = (data: Partial<HealthReadinessDetailsResponse> | undefined) => {
vi.mocked(useHealthReadinessDetails).mockReturnValue({ data } as UseQueryResult<HealthReadinessDetailsResponse>);
};
const mockRole = (userRole: string) => {
vi.mocked(useAuth).mockReturnValue({ userRole } as ReturnType<typeof useAuth>);
};
describe("EnvCredentialLoginWarningBanner", () => {
it("should warn an admin when env-credential login is enabled", () => {
mockRole("Admin");
mockDetails({ status: "healthy", show_env_credential_login_warning: true });
renderWithProviders(<EnvCredentialLoginWarningBanner accessToken="token" />);
expect(screen.getByRole("alert")).toBeInTheDocument();
expect(screen.getByText("Environment-credential login is enabled")).toBeInTheDocument();
});
it("should tell the admin to create a regular admin account before disabling", () => {
mockRole("Admin");
mockDetails({ status: "healthy", show_env_credential_login_warning: true });
renderWithProviders(<EnvCredentialLoginWarningBanner accessToken="token" />);
expect(screen.getByText(/First create a regular admin account/i)).toBeInTheDocument();
expect(screen.getByText("general_settings.disable_env_credential_login: true")).toBeInTheDocument();
});
it("should warn an admin viewer too", () => {
mockRole("Admin Viewer");
mockDetails({ status: "healthy", show_env_credential_login_warning: true });
renderWithProviders(<EnvCredentialLoginWarningBanner accessToken="token" />);
expect(screen.getByRole("alert")).toBeInTheDocument();
});
it("should render nothing for a non-admin even when the proxy reports the warning", () => {
mockRole("Internal User");
mockDetails({ status: "healthy", show_env_credential_login_warning: true });
const { container } = renderWithProviders(<EnvCredentialLoginWarningBanner accessToken="token" />);
expect(container).toBeEmptyDOMElement();
});
it("should render nothing when env-credential login is disabled", () => {
mockRole("Admin");
mockDetails({ status: "healthy", show_env_credential_login_warning: false });
const { container } = renderWithProviders(<EnvCredentialLoginWarningBanner accessToken="token" />);
expect(container).toBeEmptyDOMElement();
});
it("should render nothing when readiness details are unavailable", () => {
mockRole("Admin");
mockDetails(undefined);
const { container } = renderWithProviders(<EnvCredentialLoginWarningBanner accessToken={null} />);
expect(container).toBeEmptyDOMElement();
});
it("should pass the access token to the readiness hook", () => {
mockRole("Admin");
mockDetails(undefined);
renderWithProviders(<EnvCredentialLoginWarningBanner accessToken="my-token" />);
expect(useHealthReadinessDetails).toHaveBeenCalledWith("my-token");
});
});

35
ui/litellm-dashboard/src/components/EnvCredentialLoginWarningBanner.tsx Normal file
View file

@ -0,0 +1,35 @@
"use client";
import React from "react";
import { TriangleAlert } from "lucide-react";
import { useHealthReadinessDetails } from "@/app/(dashboard)/hooks/healthReadiness/useHealthReadinessDetails";
import { useAuth } from "@/contexts/AuthContext";
import { isAdminRole } from "@/utils/roles";
export const EnvCredentialLoginWarningBanner: React.FC<{ accessToken: string | null }> = ({ accessToken }) => {
const { userRole } = useAuth();
const { data: healthData } = useHealthReadinessDetails(accessToken);
if (!isAdminRole(userRole) || !healthData?.show_env_credential_login_warning) {
return null;
}
return (
<div
role="alert"
className="flex items-start gap-3 border-b border-destructive/40 bg-destructive/10 px-4 py-3 text-sm text-destructive"
>
<TriangleAlert className="mt-0.5 size-5 shrink-0" aria-hidden="true" />
<div>
<p className="font-semibold">Environment-credential login is enabled</p>
<p>
Anyone with <code className="font-mono">UI_USERNAME</code>/<code className="font-mono">UI_PASSWORD</code> (or
the master key, when <code className="font-mono">UI_PASSWORD</code> is unset) can sign in as a proxy admin
with a shared static secret. First create a regular admin account with its own password, then set{" "}
<code className="font-mono">general_settings.disable_env_credential_login: true</code> to turn this login path
off.
</p>
</div>
</div>
);
};

4
ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx
View file

@ -1046,9 +1046,7 @@ describe("AddAutoRouterTab", () => {
await waitFor(() => expect(handleAddAutoRouterSubmit).toHaveBeenCalled());
expect(vi.mocked(handleAddAutoRouterSubmit).mock.calls.at(-1)?.[0]).toMatchObject({
complexity_router_config: {
tier_model_configs: {
REASONING: [{ model_name: "claude-opus-5", litellm_params: { reasoning_effort: "high" } }],
},
tier_model_configs: ANTHROPIC_PRESET.complexity_router_config.tier_model_configs,
},
});
});

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

@ -25870,6 +25870,11 @@ export interface components {
* @description If True, disables the optimistic per-request budget reservation introduced in v1.84.0. WARNING: This weakens hard budget enforcement. Without the reservation, a burst of concurrent requests from a single key can each pass the read-time spend check before any of them is charged, allowing a configured budget to be exceeded under high concurrency. Budgets are still evaluated on every request at read time, so an already-exhausted budget is still rejected. Enable only if your deployment is experiencing phantom BudgetExceededError responses caused by leaked reservations (see GitHub issue #27639). An INFO notice is logged once per worker at config load while this flag is active as a reminder that hard enforcement is relaxed.
*/
disable_budget_reservation?: boolean | null;
/**
* Disable Env Credential Login
* @description If True, disables signing in to the Admin UI with the environment credentials: UI_USERNAME/UI_PASSWORD, or the master key when UI_PASSWORD is unset (that fallback means env-credential login is always live by default). Database users with passwords are unaffected. LOCKOUT RISK: create at least one proxy admin user with a password before enabling, or nobody can sign in to the UI. A locked-out admin can still administer the proxy over the API with the master key, and can unset this setting and restart the proxy to restore env-credential login. Default is False.
*/
disable_env_credential_login?: boolean | null;
/**
* Disable Password Login When Sso Enabled
* @description If True and SSO is configured (MICROSOFT_CLIENT_ID, GOOGLE_CLIENT_ID, GENERIC_CLIENT_ID, or SAML_IDP_METADATA_URL/XML), disables username/password login on /login, /v2/login, and /v3/login so SSO is the only way to reach the Admin UI. An admin locked out of the UI can still administer the proxy over the API with the master key; unset this setting and restart the proxy to restore UI username/password login. Default is False.