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

Merge pull request #30890 from BerriAI/litellm_backport_1_87_x_0620

Browse source 
chore(release): backport #29311, #29444, #29447, #29598, #30480, #30543, #30542, #30573 to stable/1.87.x and cut 1.87.4
This commit is contained in:
yuneng-jiang 2026-06-20 14:44:03 -07:00 • committed by GitHub
commit faa2f13a05
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
40 changed files with 2729 additions and 231 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
litellm/constants.py
View file

@ -190,6 +190,10 @@ DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET_GEMINI_2_5_FLASH_LITE = int(
# Override with LITELLM_MAX_CALLBACKS env var for large deployments (e.g., many teams with guardrails)
MAX_CALLBACKS = get_env_int("LITELLM_MAX_CALLBACKS", 100)
# Metadata key recording which pre_call guardrails the proxy loop already ran,
# so the deployment-level hook does not re-run them for the same request
PRE_CALL_EXECUTED_GUARDRAILS_KEY = "_pre_call_executed_guardrails"
# Generic fallback for unknown models
DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET = int(
os.getenv("DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET", 128)

155
litellm/integrations/anthropic_cache_control_hook.py
View file

@ -27,6 +27,11 @@ else:
LiteLLMLoggingObj = Any
# Anthropic (and Bedrock Claude) reject requests with more than 4 cache_control
# breakpoints: "A maximum of 4 blocks with cache_control may be provided."
MAX_CACHE_CONTROL_BLOCKS = 4
class AnthropicCacheControlHook(CustomPromptManagement):
def get_chat_completion_prompt(
self,
@ -61,16 +66,30 @@ class AnthropicCacheControlHook(CustomPromptManagement):
processed_messages = copy.deepcopy(messages)
# Separate message-level and non-message-level injection points
remaining_points = []
message_points: List[CacheControlMessageInjectionPoint] = []
remaining_points: List[CacheControlInjectionPoint] = []
for point in injection_points:
if point.get("location") == "message":
point = cast(CacheControlMessageInjectionPoint, point)
processed_messages = self._process_message_injection(
point=point, messages=processed_messages
)
message_points.append(cast(CacheControlMessageInjectionPoint, point))
else:
remaining_points.append(point)
# Non-message points (currently Bedrock tool_config) are handled in the
# provider transform, where each tool_config point appends at most one
# cachePoint to the tools. That block also counts toward Anthropic's
# limit, so reserve a slot for it here to leave room.
reserved_blocks = (
1
if any(p.get("location") == "tool_config" for p in remaining_points)
else 0
)
processed_messages = self._apply_message_injections(
points=message_points,
messages=processed_messages,
max_blocks=MAX_CACHE_CONTROL_BLOCKS - reserved_blocks,
)
# Pass through non-message injection points for provider-specific handling
if remaining_points:
non_default_params["cache_control_injection_points"] = remaining_points
@ -78,14 +97,71 @@ class AnthropicCacheControlHook(CustomPromptManagement):
return model, processed_messages, non_default_params
@staticmethod
def _process_message_injection(
point: CacheControlMessageInjectionPoint, messages: List[AllMessageValues]
def _apply_message_injections(
points: List[CacheControlMessageInjectionPoint],
messages: List[AllMessageValues],
max_blocks: int,
) -> List[AllMessageValues]:
"""Process message-level cache control injection."""
control: ChatCompletionCachedContent = point.get(
"control", None
) or ChatCompletionCachedContent(type="ephemeral")
"""Apply message-level cache control injection points in order.
Anthropic allows at most ``MAX_CACHE_CONTROL_BLOCKS`` cache_control
breakpoints per request. Client-supplied breakpoints count toward that
limit, so we never inject onto a message that already carries
cache_control (preserving the client's TTL) and we stop injecting once
``max_blocks`` is reached. Injection points are honored in config order,
so earlier points win when slots are scarce.
"""
used_blocks = sum(
AnthropicCacheControlHook._count_cache_control_blocks(msg)
for msg in messages
)
limit_reached = False
for point in points:
if used_blocks >= max_blocks:
limit_reached = True
break
control: ChatCompletionCachedContent = point.get(
"control", None
) or ChatCompletionCachedContent(type="ephemeral")
for target_index in AnthropicCacheControlHook._resolve_target_indices(
point=point, messages=messages
):
if used_blocks >= max_blocks:
limit_reached = True
break
if AnthropicCacheControlHook._message_has_cache_control(
messages[target_index]
):
# Client already marked this message; don't overwrite it.
continue
messages[target_index] = (
AnthropicCacheControlHook._safe_insert_cache_control_in_message(
messages[target_index], control
)
)
used_blocks += 1
if limit_reached:
break
if limit_reached:
verbose_logger.warning(
f"AnthropicCacheControlHook: Reached the Anthropic limit of "
f"{MAX_CACHE_CONTROL_BLOCKS} cache_control blocks. Skipping further injection."
)
return messages
@staticmethod
def _resolve_target_indices(
point: CacheControlMessageInjectionPoint, messages: List[AllMessageValues]
) -> List[int]:
"""Resolve which message indices an injection point targets."""
_targetted_index: Optional[Union[int, str]] = point.get("index", None)
targetted_index: Optional[int] = None
if isinstance(_targetted_index, str):
@ -96,36 +172,49 @@ class AnthropicCacheControlHook(CustomPromptManagement):
else:
targetted_index = _targetted_index
targetted_role = point.get("role", None)
# Case 1: Target by specific index
if targetted_index is not None:
original_index = targetted_index
# Handle negative indices (convert to positive)
if targetted_index < 0:
targetted_index += len(messages)
if 0 <= targetted_index < len(messages):
messages[targetted_index] = (
AnthropicCacheControlHook._safe_insert_cache_control_in_message(
messages[targetted_index], control
)
)
else:
verbose_logger.warning(
f"AnthropicCacheControlHook: Provided index {original_index} is out of bounds for message list of length {len(messages)}. "
f"Targeted index was {targetted_index}. Skipping cache control injection for this point."
)
return [targetted_index]
verbose_logger.warning(
f"AnthropicCacheControlHook: Provided index {original_index} is out of bounds for message list of length {len(messages)}. "
f"Targeted index was {targetted_index}. Skipping cache control injection for this point."
)
return []
# Case 2: Target by role
elif targetted_role is not None:
for msg in messages:
if msg.get("role") == targetted_role:
msg = (
AnthropicCacheControlHook._safe_insert_cache_control_in_message(
message=msg, control=control
)
)
return messages
targetted_role = point.get("role", None)
if targetted_role is not None:
return [
idx
for idx, msg in enumerate(messages)
if msg.get("role") == targetted_role
]
return []
@staticmethod
def _count_cache_control_blocks(message: AllMessageValues) -> int:
"""Count cache_control breakpoints on a message (message + content level)."""
count = 0
if message.get("cache_control") is not None:
count += 1
content = message.get("content")
if isinstance(content, list):
for block in content:
if isinstance(block, dict) and block.get("cache_control") is not None:
count += 1
return count
@staticmethod
def _message_has_cache_control(message: AllMessageValues) -> bool:
"""Return True if the message already carries any cache_control."""
return AnthropicCacheControlHook._count_cache_control_blocks(message) > 0
@staticmethod
def _safe_insert_cache_control_in_message(

54
litellm/integrations/custom_guardrail.py
View file

@ -1,3 +1,4 @@
import secrets
from datetime import datetime
from typing import (
TYPE_CHECKING,
@ -43,12 +44,19 @@ if TYPE_CHECKING:
dc = DualCache()
from litellm.constants import PRE_CALL_EXECUTED_GUARDRAILS_KEY
from litellm.exceptions import (
BlockedPiiEntityError,
GuardrailRaisedException,
ModifyResponseException,
)
# Per-process secret tagging each recorded marker. The deployment hook only
# honors markers carrying this token, so a caller cannot forge the metadata
# field to suppress a guardrail on the direct-SDK path that never reaches the
# proxy's metadata sanitizer.
_PRE_CALL_EXECUTED_TOKEN = secrets.token_hex(16)
class CustomGuardrail(CustomLogger):
# If True, during_call runs async_moderation_hook instead of the unified apply_guardrail path.
@ -325,6 +333,49 @@ class CustomGuardrail(CustomLogger):
return False
def _pre_call_marker(self) -> Optional[str]:
name = self.guardrail_name
if not name:
return None
return f"{_PRE_CALL_EXECUTED_TOKEN}:{name}"
def mark_pre_call_hook_ran(self, data: Dict[str, Any]) -> None:
"""
Record that this guardrail's ``async_pre_call_hook`` already ran for this
request, so the deployment-level hook does not run it a second time.
The proxy runs pre-call guardrails in ``ProxyLogging.pre_call_hook``. The
router later spreads a deployment's model-level ``guardrails`` into the
top-level request kwargs, which would otherwise re-trigger the same hook
from ``async_pre_call_deployment_hook``.
"""
marker = self._pre_call_marker()
if marker is None:
return
for meta_key in ("metadata", "litellm_metadata"):
meta = data.get(meta_key)
if isinstance(meta, dict):
executed = meta.get(PRE_CALL_EXECUTED_GUARDRAILS_KEY)
if isinstance(executed, list):
if marker not in executed:
executed.append(marker)
else:
meta[PRE_CALL_EXECUTED_GUARDRAILS_KEY] = [marker]
return
data["metadata"] = {PRE_CALL_EXECUTED_GUARDRAILS_KEY: [marker]}
def _pre_call_hook_already_ran(self, data: Dict[str, Any]) -> bool:
marker = self._pre_call_marker()
if marker is None:
return False
for meta_key in ("metadata", "litellm_metadata"):
meta = data.get(meta_key)
if isinstance(meta, dict):
executed = meta.get(PRE_CALL_EXECUTED_GUARDRAILS_KEY)
if isinstance(executed, list) and marker in executed:
return True
return False
async def async_pre_call_deployment_hook(
self, kwargs: Dict[str, Any], call_type: Optional[CallTypes]
) -> Optional[dict]:
@ -335,6 +386,9 @@ class CustomGuardrail(CustomLogger):
if litellm_guardrails is None or not isinstance(litellm_guardrails, list):
return kwargs
if self._pre_call_hook_already_ran(kwargs):
return kwargs
if (
self.should_run_guardrail(
data=kwargs, event_type=GuardrailEventHooks.pre_call

97
litellm/integrations/datadog/datadog.py
View file

@ -41,6 +41,7 @@ from litellm.integrations.datadog.datadog_handler import (
)
from litellm.litellm_core_utils.dd_tracing import tracer
from litellm.llms.custom_httpx.http_handler import (
MaskedHTTPStatusError,
_get_httpx_client,
get_async_httpx_client,
httpxSpecialProvider,
@ -68,6 +69,22 @@ DD_LOGGED_SUCCESS_SERVICE_TYPES = [
]
def _resolve_dd_batch_size() -> int:
raw = os.getenv("DD_BATCH_SIZE")
if raw is None:
return DD_MAX_BATCH_SIZE
try:
value = int(raw)
except ValueError:
verbose_logger.warning(
"Datadog: ignoring invalid DD_BATCH_SIZE=%r, using %s",
raw,
DD_MAX_BATCH_SIZE,
)
return DD_MAX_BATCH_SIZE
return max(1, min(value, DD_MAX_BATCH_SIZE))
class DataDogLogger(
CustomBatchLogger,
AdditionalLoggingUtils,
@ -128,7 +145,9 @@ class DataDogLogger(
asyncio.create_task(self.periodic_flush())
self.flush_lock = asyncio.Lock()
super().__init__(
**kwargs, flush_lock=self.flush_lock, batch_size=DD_MAX_BATCH_SIZE
**kwargs,
flush_lock=self.flush_lock,
batch_size=_resolve_dd_batch_size(),
)
except Exception as e:
verbose_logger.exception(
@ -339,28 +358,14 @@ class DataDogLogger(
"[DATADOG MOCK] Mock mode enabled - API calls will be intercepted"
)
response = await self.async_send_compressed_data(batch_to_send)
if response.status_code == 413:
verbose_logger.exception(DD_ERRORS.DATADOG_413_ERROR.value)
self.log_queue = batch_to_send + self.log_queue
return
response.raise_for_status()
if response.status_code != 202:
raise Exception(
f"Response from datadog API status_code: {response.status_code}, text: {response.text}"
)
undelivered = await self._send_with_413_split(batch_to_send)
if undelivered:
self.log_queue = undelivered + self.log_queue
if self.is_mock_mode:
verbose_logger.debug(
f"[DATADOG MOCK] Batch of {len(batch_to_send)} events successfully mocked"
)
else:
verbose_logger.debug(
"Datadog: Response from datadog API status_code: %s, text: %s",
response.status_code,
response.text,
)
except Exception as e:
self.log_queue = batch_to_send + self.log_queue
@ -368,6 +373,62 @@ class DataDogLogger(
f"Datadog Error sending batch API - {str(e)}\n{traceback.format_exc()}"
)
async def _send_with_413_split(self, batch: List) -> List:
"""
Send a batch, halving any sub-batch that 413s (payload too large) and retrying the
halves, since Datadog enforces a 5MB uncompressed limit per request.
A 413 surfaces as a raised MaskedHTTPStatusError (httpx raise_for_status), not a
returned response, so both paths are handled. A lone event that still 413s is
dropped to avoid wedging the queue on an undeliverable payload. Returns the events
that could not be delivered because of a non-413 (transient) error, so the caller
re-queues only those and never the events already accepted by Datadog.
"""
pending: List[List] = [batch]
while pending:
chunk = pending.pop()
if not chunk:
continue
try:
response = await self.async_send_compressed_data(chunk)
except Exception as e:
if isinstance(e, MaskedHTTPStatusError) and e.status_code == 413:
response = e.response
else:
verbose_logger.exception(
f"Datadog Error sending batch API - {str(e)}"
)
return self._undelivered(chunk, pending)
if response.status_code == 413:
if len(chunk) == 1:
verbose_logger.error(DD_ERRORS.DATADOG_413_ERROR.value)
continue
mid = len(chunk) // 2
pending.append(chunk[mid:])
pending.append(chunk[:mid])
continue
if response.status_code != 202:
verbose_logger.error(
"Datadog: unexpected response status_code=%s, text=%s",
response.status_code,
response.text,
)
return self._undelivered(chunk, pending)
verbose_logger.debug(
"Datadog: delivered %s events, status_code=%s, text=%s",
len(chunk),
response.status_code,
response.text,
)
return []
@staticmethod
def _undelivered(chunk: List, pending: List[List]) -> List:
return chunk + [event for remaining in reversed(pending) for event in remaining]
async def flush_queue(self):
if self.flush_lock is None:
return

106
litellm/litellm_core_utils/litellm_logging.py
View file

@ -1607,6 +1607,90 @@ class Logging(LiteLLMLoggingBaseClass):
) -> Optional[float]:
return self._response_cost_calculator(result=result, cache_hit=cache_hit)
@staticmethod
def _is_sync_litellm_request(litellm_params: dict) -> bool:
"""True for sync SDK entrypoints (``completion``), false for async (``acompletion``, etc.)."""
return (
litellm_params.get(CallTypes.acompletion.value, False) is not True
and litellm_params.get(CallTypes.aresponses.value, False) is not True
and litellm_params.get(CallTypes.aembedding.value, False) is not True
and litellm_params.get(CallTypes.aimage_generation.value, False) is not True
and litellm_params.get(CallTypes.atranscription.value, False) is not True
)
def _is_assembled_stream_success(self, result=None) -> bool:
"""Final assembled stream export (not a per-chunk success call).
Per-chunk callers pass a ``ModelResponseStream`` (or ``None``); the
final assembled response is any other non-``None`` value (typically a
``ModelResponse``). Treating a chunk as the assembled response would
prematurely set the ``has_dispatched_final_stream_success`` dedup
guard and silently suppress the real final stream log.
"""
if self.stream is not True:
return False
if result is not None and not isinstance(result, ModelResponseStream):
return True
return (
"async_complete_streaming_response" in self.model_call_details
or self.model_call_details.get("complete_streaming_response") is not None
)
async def dispatch_success_handlers(
self,
result=None,
start_time=None,
end_time=None,
cache_hit=None,
prefer_async_handlers: bool = False,
**kwargs,
) -> None:
"""Route success logging to async and/or sync handlers for this request.
``prefer_async_handlers`` only bypasses the sync-SDK-only shortcut (e.g.
``async for`` on a stream from ``completion()``). Legacy string callbacks
still run via ``executor.submit(success_handler)`` when configured.
"""
from litellm.litellm_core_utils.thread_pool_executor import executor
if self._is_assembled_stream_success(result):
if self.model_call_details.get("has_dispatched_final_stream_success"):
return
self.model_call_details["has_dispatched_final_stream_success"] = True
litellm_params = self.model_call_details.get("litellm_params", {}) or {}
sync_sdk = self._is_sync_litellm_request(litellm_params)
passthrough = self.call_type == CallTypes.pass_through.value
if sync_sdk and not prefer_async_handlers and not passthrough:
self.success_handler(
result,
start_time=start_time,
end_time=end_time,
cache_hit=cache_hit,
**kwargs,
)
return
await self.async_success_handler(
result,
start_time=start_time,
end_time=end_time,
cache_hit=cache_hit,
**kwargs,
)
if not self._should_run_sync_callbacks_for_async_calls():
return
executor.submit(
self.success_handler,
result,
start_time=start_time,
end_time=end_time,
cache_hit=cache_hit,
**kwargs,
)
def should_run_logging(
self,
event_type: Literal[
@ -2029,13 +2113,7 @@ class Logging(LiteLLMLoggingBaseClass):
standard_logging_object=kwargs.get("standard_logging_object", None),
)
litellm_params = self.model_call_details.get("litellm_params", {})
is_sync_request = (
litellm_params.get(CallTypes.acompletion.value, False) is not True
and litellm_params.get(CallTypes.aresponses.value, False) is not True
and litellm_params.get(CallTypes.aembedding.value, False) is not True
and litellm_params.get(CallTypes.aimage_generation.value, False) is not True
and litellm_params.get(CallTypes.atranscription.value, False) is not True
)
is_sync_request = self._is_sync_litellm_request(litellm_params)
try:
## BUILD COMPLETE STREAMED RESPONSE
complete_streaming_response: Optional[
@ -2491,9 +2569,11 @@ class Logging(LiteLLMLoggingBaseClass):
print_verbose(
"Logging Details LiteLLM-Async Success Call, cache_hit={}".format(cache_hit)
)
if not self.should_run_logging(
if not self._is_assembled_stream_success(
result
) and not self.should_run_logging(
event_type="async_success"
): # prevent double logging
): # prevent double logging (non-streaming)
return
## CALCULATE COST FOR BATCH JOBS
@ -2943,13 +3023,7 @@ class Logging(LiteLLMLoggingBaseClass):
): # prevent double logging
return
litellm_params = self.model_call_details.get("litellm_params", {})
is_sync_request = (
litellm_params.get(CallTypes.acompletion.value, False) is not True
and litellm_params.get(CallTypes.aresponses.value, False) is not True
and litellm_params.get(CallTypes.aembedding.value, False) is not True
and litellm_params.get(CallTypes.aimage_generation.value, False) is not True
and litellm_params.get(CallTypes.atranscription.value, False) is not True
)
is_sync_request = self._is_sync_litellm_request(litellm_params)
try:
start_time, end_time = self._failure_handler_helper_fn(

16
litellm/litellm_core_utils/logging_callback_manager.py
View file

@ -394,6 +394,22 @@ class LoggingCallbackManager:
+ litellm._async_failure_callback
)
def remove_callback_from_all_lists(self, obj, require_self=False) -> None:
"""
Remove a callback object from every callback list it may have been
promoted into, so a re-initialized callback leaves no stale instance behind.
"""
for callback_list in (
litellm.callbacks,
litellm.success_callback,
litellm.failure_callback,
litellm._async_success_callback,
litellm._async_failure_callback,
):
self.remove_callback_from_list_by_object(
callback_list, obj, require_self=require_self
)
def get_active_additional_logging_utils_from_custom_logger(
self,
) -> Set[AdditionalLoggingUtils]:

20
litellm/litellm_core_utils/streaming_handler.py
View file

@ -1808,8 +1808,10 @@ class CustomStreamWrapper:
processed_chunk, None, None, cache_hit
)
)
## SYNC LOGGING
self.logging_obj.success_handler(processed_chunk, None, None, cache_hit)
## SYNC LOGGING — only for sync SDK entrypoints; async proxy paths export via async_success_handler
litellm_params = self.logging_obj.model_call_details.get("litellm_params", {})
if self.logging_obj._is_sync_litellm_request(litellm_params):
self.logging_obj.success_handler(processed_chunk, None, None, cache_hit)
def finish_reason_handler(self):
model_response = self.model_response_creator()
@ -2206,23 +2208,19 @@ class CustomStreamWrapper:
cache_hit,
)
else:
# prefer_async_handlers routes CustomLogger to async_success_handler
# when consumers use ``async for`` on sync-SDK streams. Legacy string
# callbacks still run via executor.submit inside dispatch_success_handlers.
asyncio.create_task(
self.logging_obj.async_success_handler(
self.logging_obj.dispatch_success_handlers(
complete_streaming_response,
cache_hit=cache_hit,
start_time=None,
end_time=None,
prefer_async_handlers=True,
)
)
executor.submit(
self.logging_obj.success_handler,
complete_streaming_response,
cache_hit=cache_hit,
start_time=None,
end_time=None,
)
raise StopAsyncIteration # Re-raise StopIteration
else:
self.sent_last_chunk = True

1
litellm/proxy/_types.py
View file

@ -3683,6 +3683,7 @@ class ProxyException(Exception):
provider_specific_fields: Optional[dict] = None,
):
self.message = str(message)
super().__init__(self.message)
self.type = type
self.param = param
self.openai_code = openai_code or code

56
litellm/proxy/common_request_processing.py
View file

@ -1290,7 +1290,7 @@ class ProxyBaseLLMRequestProcessing:
# (ProxyLogging._fire_deferred_stream_logging) fires the
# closure after the full streaming pipeline finishes.
# The closure runs non-apply_guardrail hooks on the
# assembled response, then fires both logging handlers.
# assembled response, then fires success logging.
# Only for CustomStreamWrapper — raw async generators from
# passthrough routes bypass CSW and would orphan the closure.
from litellm.litellm_core_utils.streaming_handler import (
@ -1411,33 +1411,18 @@ class ProxyBaseLLMRequestProcessing:
logging_obj._on_deferred_stream_complete = None # type: ignore[union-attr]
try:
asyncio.create_task(
logging_obj.async_success_handler(
logging_obj.dispatch_success_handlers(
response,
cache_hit=None,
start_time=None,
end_time=None,
prefer_async_handlers=True,
)
)
except Exception as e:
verbose_proxy_logger.exception(
"Error in orphaned streaming async logging: %s", e
)
try:
from litellm.litellm_core_utils.thread_pool_executor import (
executor as _exc,
)
_exc.submit(
logging_obj.success_handler,
response,
cache_hit=None,
start_time=None,
end_time=None,
)
except Exception as e:
verbose_proxy_logger.exception(
"Error in orphaned streaming sync logging: %s", e
)
# Always return the client-requested model name (not provider-prefixed internal identifiers)
# for OpenAI-compatible responses.
@ -1639,7 +1624,7 @@ class ProxyBaseLLMRequestProcessing:
) -> None:
"""
Run non-streaming post-call guardrail hooks on an assembled streaming
response, then fire both async and sync logging handlers.
response, then fire success logging via ``dispatch_success_handlers``.
Called by ProxyLogging._fire_deferred_stream_logging after the full
streaming pipeline (including unified_guardrail end-of-stream blocks)
@ -1655,8 +1640,6 @@ class ProxyBaseLLMRequestProcessing:
Extracted as a static method so tests can call the production
implementation directly rather than reimplementing the closure.
"""
from litellm.litellm_core_utils.thread_pool_executor import executor
_response = assembled_response
try:
from litellm.proxy.proxy_server import llm_router as _global_llm_router
@ -1715,31 +1698,23 @@ class ProxyBaseLLMRequestProcessing:
)
finally:
try:
# Proxy streaming always runs in async context and proxy spend
# logging is async-only; force async dispatch so DB/spend
# callbacks fire regardless of the call-type heuristic in
# _is_sync_litellm_request (which only recognizes a subset of
# async markers stored in litellm_params).
asyncio.create_task(
captured_logging_obj.async_success_handler(
captured_logging_obj.dispatch_success_handlers(
_response,
cache_hit=cache_hit,
start_time=None,
end_time=None,
prefer_async_handlers=True,
)
)
except Exception as e:
verbose_proxy_logger.exception(
"Error in deferred streaming async logging: %s",
e,
)
try:
executor.submit(
captured_logging_obj.success_handler,
_response,
cache_hit=cache_hit,
start_time=None,
end_time=None,
)
except Exception as e:
verbose_proxy_logger.exception(
"Error in deferred streaming sync logging: %s",
"Error in deferred streaming success logging: %s",
e,
)
@ -1824,6 +1799,13 @@ class ProxyBaseLLMRequestProcessing:
except Exception:
pass
if isinstance(e, ProxyException):
e.headers = {
**e.headers,
**{k: v if isinstance(v, str) else str(v) for k, v in headers.items()},
}
raise e
if isinstance(e, HTTPException):
raw_detail = getattr(e, "detail", str(e))
message, structured_fields = _serialize_http_exception_detail(raw_detail)

2
litellm/proxy/common_utils/callback_utils.py
View file

@ -4,6 +4,7 @@ from typing import TYPE_CHECKING, Any, Callable, Dict, Iterable, List, Literal,
import litellm
from litellm import get_secret
from litellm._logging import verbose_proxy_logger
from litellm.constants import PRE_CALL_EXECUTED_GUARDRAILS_KEY
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker
from litellm.proxy._types import CommonProxyErrors, LiteLLMPromptInjectionParams
@ -482,6 +483,7 @@ LITELLM_PROXY_INTERNAL_METADATA_KEYS = frozenset(
"guardrail_config",
"_guardrail_pipelines",
"_pipeline_managed_guardrails",
PRE_CALL_EXECUTED_GUARDRAILS_KEY,
"disable_global_guardrails",
"disable_global_guardrail",
"opted_out_global_guardrails",

34
litellm/proxy/guardrails/guardrail_hooks/aim/aim.py
View file

@ -9,7 +9,6 @@ import json
import os
from typing import TYPE_CHECKING, Any, AsyncGenerator, Optional, Type, Union
from fastapi import HTTPException
from pydantic import BaseModel
from websockets.asyncio.client import ClientConnection, connect
@ -21,7 +20,7 @@ from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
httpxSpecialProvider,
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy._types import ProxyException, UserAPIKeyAuth
from litellm.proxy.guardrails._content_utils import (
apply_redacted_messages_back,
build_inspection_messages,
@ -129,6 +128,16 @@ class AimGuardrail(CustomGuardrail):
verbose_proxy_logger.error(f"Aim: {action_type} action")
return data
@staticmethod
def _rejection(message: str, *, openai_code: str | None = None) -> ProxyException:
return ProxyException(
message=message,
type="invalid_request_error",
param=None,
code=400,
openai_code=openai_code,
)
def _handle_block_action(self, analysis_result: Any, required_action: Any) -> None:
detection_message = required_action.get("detection_message", None)
verbose_proxy_logger.info(
@ -136,7 +145,7 @@ class AimGuardrail(CustomGuardrail):
policies=list(analysis_result["policy_drill_down"].keys()),
),
)
raise HTTPException(status_code=400, detail=detection_message)
raise self._rejection(detection_message, openai_code="content_policy_violation")
def _anonymize_request(self, res: Any, data: dict) -> dict:
verbose_proxy_logger.info("Aim: anonymize action")
@ -148,14 +157,11 @@ class AimGuardrail(CustomGuardrail):
# parts from a multimodal request — degrade to block so the
# multimodal payload is never silently rewritten.
if has_non_string_content(data):
raise HTTPException(
status_code=400,
detail=(
"Aim: anonymize action requested for multimodal input "
"but mask-in-place would drop non-text parts. Send the "
"request with plain string content to use anonymize, "
"or rely on block-mode policies."
),
raise self._rejection(
"Aim: anonymize action requested for multimodal input "
"but mask-in-place would drop non-text parts. Send the "
"request with plain string content to use anonymize, "
"or rely on block-mode policies."
)
redacted_messages = [
{
@ -287,9 +293,9 @@ class AimGuardrail(CustomGuardrail):
if aim_output_guardrail_result and aim_output_guardrail_result.get(
"detection_message"
):
raise HTTPException(
status_code=400,
detail=aim_output_guardrail_result.get("detection_message"),
raise self._rejection(
aim_output_guardrail_result.get("detection_message"),
openai_code="content_policy_violation",
)
if aim_output_guardrail_result and aim_output_guardrail_result.get(
"redacted_output"

64
litellm/proxy/guardrails/guardrail_registry.py
View file

@ -5,6 +5,8 @@ import os
from datetime import datetime, timezone
from typing import Any, Dict, List, Literal, Optional, Set, Type, cast
from pydantic import ValidationError
import litellm
from litellm import Router
from litellm._logging import verbose_proxy_logger
@ -598,21 +600,25 @@ class InMemoryGuardrailHandler:
def delete_in_memory_guardrail(self, guardrail_id: str) -> None:
"""
Delete a guardrail in memory and remove from litellm callbacks.
The callback is purged from every callback list, not just
litellm.callbacks: request handling promotes guardrail callbacks into the
success/failure/async lists, so removing it from only litellm.callbacks
leaves the old instance stranded in those lists on every re-initialization.
"""
# Remove from in-memory storage
self.IN_MEMORY_GUARDRAILS.pop(guardrail_id, None)
self._sources.pop(guardrail_id, None)
# Remove the callback from litellm.callbacks
custom_guardrail_callback = self.guardrail_id_to_custom_guardrail.pop(
guardrail_id, None
)
if custom_guardrail_callback:
litellm.logging_callback_manager.remove_callback_from_list_by_object(
callback_list=litellm.callbacks,
obj=custom_guardrail_callback,
require_self=False,
)
if custom_guardrail_callback is None:
return
litellm.logging_callback_manager.remove_callback_from_all_lists(
custom_guardrail_callback
)
def list_in_memory_guardrails(self) -> List[Guardrail]:
"""
@ -654,6 +660,34 @@ class InMemoryGuardrailHandler:
self.delete_in_memory_guardrail(guardrail_id)
return stale_ids
@staticmethod
def _normalize_litellm_params_for_comparison(
params: Optional[Any],
) -> Optional[Dict[str, Any]]:
"""
Render litellm_params to a canonical dict so an in-memory LitellmParams and
the raw dict loaded from the DB compare equal when they describe the same
config. The in-memory side is a LitellmParams whose model_dump() carries
every field default and coerces enums, while the DB side is the raw stored
dict holding only the keys originally provided. Comparing those two shapes
directly never matches, so each DB poll would re-initialize the guardrail
forever; normalizing both through LitellmParams keeps the diff meaningful.
"""
if params is None:
return None
if isinstance(params, LitellmParams):
return params.model_dump()
if isinstance(params, dict):
try:
return LitellmParams(**params).model_dump()
except ValidationError as e:
verbose_proxy_logger.warning(
f"Could not normalize guardrail litellm_params for comparison; "
f"treating the guardrail as changed. Error: {e}"
)
return params
return params
def _has_guardrail_params_changed(
self, guardrail_id: str, new_guardrail: Guardrail
) -> bool:
@ -670,19 +704,11 @@ class InMemoryGuardrailHandler:
return True
# Compare litellm_params
existing_params = existing.get("litellm_params")
new_params = new_guardrail.get("litellm_params")
# Convert to dicts for comparison
existing_dict = (
existing_params.model_dump()
if isinstance(existing_params, LitellmParams)
else existing_params
existing_dict = self._normalize_litellm_params_for_comparison(
existing.get("litellm_params")
)
new_dict = (
new_params.model_dump()
if isinstance(new_params, LitellmParams)
else new_params
new_dict = self._normalize_litellm_params_for_comparison(
new_guardrail.get("litellm_params")
)
# Compare and identify specific differences

2
litellm/proxy/litellm_pre_call_utils.py
View file

@ -13,6 +13,7 @@ from starlette.datastructures import Headers
import litellm
from litellm._logging import verbose_logger, verbose_proxy_logger
from litellm._service_logger import ServiceLogging
from litellm.constants import PRE_CALL_EXECUTED_GUARDRAILS_KEY
from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
from litellm.litellm_core_utils.url_utils import is_url_destination_allowed_by_host
@ -161,6 +162,7 @@ _UNTRUSTED_METADATA_CONTROL_FIELDS = (
"secret_fields",
"_guardrail_pipelines",
"_pipeline_managed_guardrails",
PRE_CALL_EXECUTED_GUARDRAILS_KEY,
)
_UNTRUSTED_REQUEST_HEADER_CONTROL_FIELDS = frozenset(

7
litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py
View file

@ -150,6 +150,13 @@ class AnthropicPassthroughLoggingHandler:
handles streaming and non-streaming responses
"""
# Only record complete_streaming_response for actual streaming responses.
# perform_redaction scrubs this field only when stream is True, so setting
# it on a non-streaming response would bypass message redaction.
if logging_obj.model_call_details.get("stream") is True:
logging_obj.model_call_details["complete_streaming_response"] = (
litellm_model_response
)
try:
# Get custom_llm_provider from logging object if available (e.g., azure_ai for Azure Anthropic)
custom_llm_provider = logging_obj.model_call_details.get(

6
litellm/proxy/pass_through_endpoints/pass_through_endpoints.py
View file

@ -871,6 +871,9 @@ async def pass_through_request( # noqa: PLR0915
)
if stream:
logging_obj.stream = True
logging_obj.model_call_details["stream"] = True
if is_multipart:
response = (
await HttpPassThroughEndpointHelpers.make_multipart_http_request(
@ -931,6 +934,9 @@ async def pass_through_request( # noqa: PLR0915
verbose_proxy_logger.debug("response.headers= %s", response.headers)
if _is_streaming_response(response) is True:
logging_obj.stream = True
logging_obj.model_call_details["stream"] = True
try:
response.raise_for_status()
except httpx.HTTPStatusError as e:

22
litellm/proxy/pass_through_endpoints/streaming_handler.py
View file

@ -7,7 +7,6 @@ import httpx
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.litellm_core_utils.thread_pool_executor import executor
from litellm.proxy._types import PassThroughEndpointLoggingResultValues
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType
@ -145,25 +144,16 @@ class PassThroughStreamingHandler:
end_time=end_time,
model=model,
)
await litellm_logging_obj.async_success_handler(
# Always reached from an async context (anthropic_messages,
# google_genai, and proxy pass-through stream tasks). prefer_async_handlers
# keeps async-only loggers running even when call_type isn't pass_through
# and litellm_params lacks an async flag (e.g. aanthropic_messages).
await litellm_logging_obj.dispatch_success_handlers(
result=standard_logging_response_object,
start_time=start_time,
end_time=end_time,
cache_hit=False,
**kwargs,
)
if (
litellm_logging_obj._should_run_sync_callbacks_for_async_calls()
is False
):
return
executor.submit(
litellm_logging_obj.success_handler,
result=standard_logging_response_object,
end_time=end_time,
cache_hit=False,
start_time=start_time,
prefer_async_handlers=True,
**kwargs,
)
except Exception as e:

24
litellm/proxy/pass_through_endpoints/success_handler.py
View file

@ -11,7 +11,6 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import (
PassthroughStandardLoggingPayload,
)
from litellm.types.utils import StandardPassThroughResponseObject
from litellm.utils import executor as thread_pool_executor
from .llm_provider_handlers.anthropic_passthrough_logging_handler import (
AnthropicPassthroughLoggingHandler,
@ -94,19 +93,15 @@ class PassThroughEndpointLogging:
cache_hit: bool,
**kwargs,
):
"""Helper function to handle both sync and async logging operations"""
# Submit to thread pool for sync logging
thread_pool_executor.submit(
logging_obj.success_handler,
standard_logging_response_object,
start_time,
end_time,
cache_hit,
**kwargs,
)
# Handle async logging
await logging_obj.async_success_handler(
"""Log pass-through success via the shared async dispatch path."""
# Always reached from pass_through_async_success_handler, which runs in
# an async context. call_type is "pass_through_endpoint" here, so the
# passthrough guard in dispatch_success_handlers already forces the
# async handler to run; pass prefer_async_handlers explicitly to match
# the streaming sibling (_route_streaming_logging_to_handler) and keep
# async-only loggers (e.g. the proxy spend logger) firing regardless of
# how the call-type classification evolves.
await logging_obj.dispatch_success_handlers(
result=(
json.dumps(result)
if isinstance(result, dict)
@ -115,6 +110,7 @@ class PassThroughEndpointLogging:
start_time=start_time,
end_time=end_time,
cache_hit=False,
prefer_async_handlers=True,
**kwargs,
)

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

@ -171,6 +171,10 @@ class PipelineExecutor:
data=data,
call_type=call_type, # type: ignore
)
if isinstance(callback, CustomGuardrail):
callback.mark_pre_call_hook_ran(data)
if isinstance(response, dict):
callback.mark_pre_call_hook_ran(response)
elif mode == "post_call":
response = await target.async_post_call_success_hook(
user_api_key_dict=user_api_key_dict,

7
litellm/proxy/utils.py
View file

@ -1126,6 +1126,8 @@ class ProxyLogging:
response=response, data=data, call_type=call_type
)
callback.mark_pre_call_hook_ran(data)
except Exception as e:
status = "error"
error_type = type(e).__name__
@ -1965,7 +1967,7 @@ class ProxyLogging:
litellm_call_id=request_data.get("litellm_call_id", ""), status="fail"
)
if AlertType.llm_exceptions in self.alert_types and not isinstance(
original_exception, HTTPException
original_exception, (HTTPException, ProxyException)
):
"""
Just alert on LLM API exceptions. Do not alert on user errors
@ -2069,6 +2071,7 @@ class ProxyLogging:
e.g should only return True for:
- Authentication Errors from user_api_key_auth
- HTTP HTTPException (rate limit errors)
- ProxyException (guardrail blocks, budget / rate-limit errors)
"""
#########################################################
@ -2085,7 +2088,7 @@ class ProxyLogging:
):
return False
return isinstance(original_exception, HTTPException) or (
return isinstance(original_exception, (HTTPException, ProxyException)) or (
error_type == ProxyErrorTypes.auth_error
)

1
litellm/types/utils.py
View file

@ -3171,6 +3171,7 @@ all_litellm_params = (
"allowed_openai_params",
"litellm_session_id",
"use_litellm_proxy",
"use_chat_completions_api",
"prompt_label",
"shared_session",
"search_tool_name",

4
pyproject.toml
View file

@ -1,6 +1,6 @@
[project]
name = "litellm"
version = "1.87.3"
version = "1.87.4"
description = "Library to easily interface with LLM API providers"
readme = "README.md"
requires-python = ">=3.10, <3.14"
@ -257,7 +257,7 @@ source-exclude = [
profile = "black"
[tool.commitizen]
version = "1.87.3"
version = "1.87.4"
version_files = [
"pyproject.toml:^version",
]

23
tests/litellm_utils_tests/test_logging_callback_manager.py
View file

@ -192,6 +192,29 @@ def test_remove_callback_from_list_by_object():
assert len(litellm._async_failure_callback) == 0
def test_remove_callback_from_all_lists():
manager = LoggingCallbackManager()
manager._reset_all_callbacks()
class TestLogger(CustomLogger):
pass
obj = TestLogger()
manager.add_litellm_callback(obj)
manager.add_litellm_success_callback(obj)
manager.add_litellm_failure_callback(obj)
manager.add_litellm_async_success_callback(obj)
manager.add_litellm_async_failure_callback(obj)
manager.remove_callback_from_all_lists(obj)
assert obj not in litellm.callbacks
assert obj not in litellm.success_callback
assert obj not in litellm.failure_callback
assert obj not in litellm._async_success_callback
assert obj not in litellm._async_failure_callback
def test_reset_callbacks(callback_manager):
# Add various callbacks
callback_manager.add_litellm_callback("test")

135
tests/local_testing/test_aim_guardrails.py
View file

@ -6,10 +6,10 @@ import sys
from unittest.mock import AsyncMock, patch, call
import pytest
from fastapi.exceptions import HTTPException
from httpx import Request, Response
from litellm import DualCache
from litellm.proxy._types import ProxyException
from litellm.proxy.guardrails.guardrail_hooks.aim.aim import (
AimGuardrail,
AimGuardrailMissingSecrets,
@ -101,7 +101,7 @@ async def test_block_callback(mode: str):
],
}
with pytest.raises(HTTPException, match="Jailbreak detected"):
with pytest.raises(ProxyException, match="Jailbreak detected") as exc_info:
with patch(
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
return_value=Response(
@ -135,6 +135,137 @@ async def test_block_callback(mode: str):
call_type="completion",
)
exc = exc_info.value
assert exc.code == "400"
assert exc.type == "invalid_request_error"
assert exc.param is None
assert exc.openai_code == "content_policy_violation"
@pytest.mark.asyncio
async def test_output_block_raises_proxy_exception():
"""An output-side block is a content-policy violation, like the input block:
it must surface a conformant ProxyException, not a bare HTTPException whose
type/param serialize as the literal string "None". Regression for LIT-3751."""
init_guardrails_v2(
all_guardrails=[
{
"guardrail_name": "gibberish-guard",
"litellm_params": {
"guardrail": "aim",
"mode": "post_call",
"api_key": "hs-aim-key",
},
},
],
config_file_path="",
)
aim_guardrails = [
callback for callback in litellm.callbacks if isinstance(callback, AimGuardrail)
]
assert len(aim_guardrails) == 1
aim_guardrail = aim_guardrails[0]
block_on_output = Response(
json={
"analysis_result": {"policy_drill_down": {"PII": {}}},
"required_action": {
"action_type": "block_action",
"detection_message": "Output blocked: leaked secret",
"policy_name": "blocking policy",
},
},
status_code=200,
request=Request(method="POST", url="http://aim"),
)
response = ModelResponse(
choices=[
{
"finish_reason": "stop",
"index": 0,
"message": {"content": "here is the secret", "role": "assistant"},
}
]
)
with patch(
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
return_value=block_on_output,
):
with pytest.raises(ProxyException, match="Output blocked") as exc_info:
await aim_guardrail.async_post_call_success_hook(
data={"messages": [{"role": "user", "content": "tell me a secret"}]},
response=response,
user_api_key_dict=UserAPIKeyAuth(),
)
exc = exc_info.value
assert exc.code == "400"
assert exc.type == "invalid_request_error"
assert exc.param is None
assert exc.openai_code == "content_policy_violation"
@pytest.mark.asyncio
async def test_anonymize_multimodal_rejection_raises_proxy_exception():
"""Anonymize on multimodal input degrades to a 400 because mask-in-place would
drop non-text parts. That is a usage error, not a content-policy violation, so
it must raise a conformant ProxyException WITHOUT the content_policy_violation
code. Regression for LIT-3751."""
init_guardrails_v2(
all_guardrails=[
{
"guardrail_name": "gibberish-guard",
"litellm_params": {
"guardrail": "aim",
"mode": "pre_call",
"api_key": "hs-aim-key",
},
},
],
config_file_path="",
)
aim_guardrails = [
callback for callback in litellm.callbacks if isinstance(callback, AimGuardrail)
]
assert len(aim_guardrails) == 1
aim_guardrail = aim_guardrails[0]
data = {
"messages": [
{
"role": "user",
"content": [
{"type": "text", "text": "Hi my name is Brian"},
{
"type": "image_url",
"image_url": {"url": "data:image/png;base64,iVBORw0KGgo="},
},
],
},
],
}
with patch(
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
return_value=response_with_detections,
):
with pytest.raises(
ProxyException, match="anonymize action requested for multimodal"
) as exc_info:
await aim_guardrail.async_pre_call_hook(
data=data,
cache=DualCache(),
user_api_key_dict=UserAPIKeyAuth(),
call_type="completion",
)
exc = exc_info.value
assert exc.code == "400"
assert exc.type == "invalid_request_error"
assert exc.param is None
assert exc.openai_code != "content_policy_violation"
@pytest.mark.asyncio
@pytest.mark.parametrize("mode", ["pre_call", "during_call"])

1
tests/pass_through_unit_tests/test_unit_test_anthropic_pass_through.py
View file

@ -318,6 +318,7 @@ def test_handle_logging_anthropic_collected_chunks(all_chunks):
from litellm.types.utils import ModelResponse
litellm_logging_obj = Mock()
litellm_logging_obj.model_call_details = {}
pass_through_logging_obj = Mock()
sent_args = {

117
tests/pass_through_unit_tests/test_unit_test_streaming.py
View file

@ -97,6 +97,123 @@ async def test_chunk_processor_yields_raw_bytes(endpoint_type, url_route):
), "Collected chunks do not match raw chunks"
@pytest.mark.asyncio
async def test_route_streaming_logging_runs_async_handler_for_sdk_passthrough():
"""
SDK pass-through streaming (anthropic_messages, google generate_content) must run
the async success handler so async-only loggers record the assembled stream.
Regression for duplicate-trace dedupe: dispatch_success_handlers treated these as
sync SDK requests because call_type is not ``pass_through_endpoint`` and
litellm_params carries no ``acompletion`` flag, so only the sync success_handler
ran and CustomLogger.async_log_success_event never fired.
"""
import time
from litellm.types.utils import CallTypes
logging_obj = LiteLLMLoggingObj(
model="claude-sonnet-4-5",
messages=[{"role": "user", "content": "hi"}],
stream=True,
call_type=CallTypes.anthropic_messages.value,
start_time=time.time(),
litellm_call_id="test-id",
function_id="fn",
)
logging_obj.model_call_details["litellm_params"] = {"anthropic_messages": True}
with (
patch.object(
PassThroughStreamingHandler,
"_build_passthrough_logging_result",
return_value=({"id": "slp"}, {}),
),
patch.object(
logging_obj, "async_success_handler", new_callable=AsyncMock
) as mock_async,
patch.object(
logging_obj, "success_handler", new_callable=MagicMock
) as mock_sync,
patch.object(
logging_obj,
"_should_run_sync_callbacks_for_async_calls",
return_value=False,
),
):
await PassThroughStreamingHandler._route_streaming_logging_to_handler(
litellm_logging_obj=logging_obj,
passthrough_success_handler_obj=MagicMock(),
url_route="/v1/messages",
request_body={},
endpoint_type=EndpointType.ANTHROPIC,
start_time=datetime.now(),
raw_bytes=[],
end_time=datetime.now(),
)
mock_async.assert_awaited_once()
mock_sync.assert_not_called()
@pytest.mark.asyncio
async def test_handle_logging_runs_async_handler_for_passthrough():
"""
Non-streaming pass-through logging (_handle_logging) must always run the
async success handler so async-only loggers (e.g. the proxy spend logger)
record the request.
_handle_logging is only ever reached from pass_through_async_success_handler
(an async context), so it forces async dispatch via prefer_async_handlers.
This pins that contract independent of the call-type classification: even a
call_type that _is_sync_litellm_request would classify as sync (here
"completion" with no async marker in litellm_params) must still reach
async_success_handler. Without prefer_async_handlers=True the sync-only
branch would return early and async_log_success_event would never fire.
"""
import time
from litellm.types.utils import CallTypes
logging_obj = LiteLLMLoggingObj(
model="claude-sonnet-4-5",
messages=[{"role": "user", "content": "hi"}],
stream=False,
call_type=CallTypes.completion.value,
start_time=time.time(),
litellm_call_id="test-id",
function_id="fn",
)
logging_obj.model_call_details["litellm_params"] = {}
handler = PassThroughEndpointLogging()
with (
patch.object(
logging_obj, "async_success_handler", new_callable=AsyncMock
) as mock_async,
patch.object(
logging_obj, "success_handler", new_callable=MagicMock
) as mock_sync,
patch.object(
logging_obj,
"_should_run_sync_callbacks_for_async_calls",
return_value=False,
),
):
await handler._handle_logging(
logging_obj=logging_obj,
standard_logging_response_object={"id": "slp"},
result="",
start_time=datetime.now(),
end_time=datetime.now(),
cache_hit=False,
)
mock_async.assert_awaited_once()
mock_sync.assert_not_called()
def test_convert_raw_bytes_to_str_lines():
"""
Test that the _convert_raw_bytes_to_str_lines method correctly converts raw bytes to a list of strings

19
tests/proxy_unit_tests/test_proxy_reject_logging.py
View file

@ -95,6 +95,21 @@ router = Router(
)
def _register_proxy_test_logger(callback_logger: testLogger) -> None:
"""
Register the test logger on global callback lists.
``function_setup`` dedupes by object identity; each parametrized case
constructs a new ``testLogger`` and must replace the global lists, not
only ``litellm.callbacks``.
"""
litellm.callbacks = [callback_logger]
litellm.success_callback = [callback_logger]
litellm.failure_callback = [callback_logger]
litellm._async_success_callback = [callback_logger]
litellm._async_failure_callback = [callback_logger]
@pytest.mark.parametrize(
"route, body",
[
@ -115,7 +130,7 @@ router = Router(
"/v1/embeddings",
{
"input": "The food was delicious and the waiter...",
"model": "text-embedding-ada-002",
"model": "fake-model",
"encoding_format": "float",
},
),
@ -133,7 +148,7 @@ async def test_chat_completion_request_with_redaction(route, body):
setattr(proxy_server, "llm_router", router)
_test_logger = testLogger()
litellm.callbacks = [_test_logger]
_register_proxy_test_logger(_test_logger)
litellm.set_verbose = True
# Prepare the query string

183
tests/test_litellm/integrations/datadog/test_datadog_logger_batching.py
View file

@ -1,10 +1,49 @@
from unittest.mock import AsyncMock, Mock, patch
import httpx
import pytest
from httpx import Request, Response
from litellm.integrations.datadog.datadog import DataDogLogger
from litellm.types.integrations.datadog import DatadogPayload
from litellm.llms.custom_httpx.http_handler import MaskedHTTPStatusError
from litellm.types.integrations.datadog import DD_MAX_BATCH_SIZE, DatadogPayload
def _payloads(n):
return [
DatadogPayload(
ddsource="litellm",
ddtags="env:test",
hostname="host",
message=f'{{"event": {i}}}',
service="svc",
status="info",
)
for i in range(n)
]
def _raised_413():
request = Request("POST", "https://example.com")
response = Response(413, request=request, text="Payload Too Large")
return MaskedHTTPStatusError(
httpx.HTTPStatusError("413", request=request, response=response)
)
def _make_send(max_ok, delivered, *, raise_413=True):
"""Datadog double: 413 batches larger than max_ok, 202 (recording delivery) otherwise."""
async def _send(data):
request = Request("POST", "https://example.com")
if len(data) > max_ok:
if raise_413:
raise _raised_413()
return Response(413, request=request, text="Payload Too Large")
delivered.extend(event["message"] for event in data)
return Response(202, request=request, text="Accepted")
return _send
@pytest.fixture
@ -75,40 +114,152 @@ async def test_failure_hook_threshold_flush_uses_flush_queue(datadog_env):
@pytest.mark.asyncio
async def test_async_send_batch_requeues_events_on_413(datadog_env):
async def test_413_splits_oversized_batch_and_delivers_every_event(datadog_env):
"""A raised 413 (the real httpx path) halves the batch until each piece is accepted."""
with patch("asyncio.create_task"):
logger = DataDogLogger()
logger.log_queue = [
DatadogPayload(
ddsource="litellm",
ddtags="env:test",
hostname="host",
message=f'{{"event": {i}}}',
service="svc",
status="info",
logger.log_queue = _payloads(4)
delivered: list = []
logger.async_send_compressed_data = AsyncMock(side_effect=_make_send(1, delivered))
await logger.async_send_batch()
assert sorted(delivered) == [f'{{"event": {i}}}' for i in range(4)]
assert logger.log_queue == []
@pytest.mark.asyncio
async def test_413_does_not_requeue_oversized_batch(datadog_env):
"""Regression for the infinite 413 loop: an undeliverable batch must not be re-queued."""
with patch("asyncio.create_task"):
logger = DataDogLogger()
logger.log_queue = _payloads(4)
logger.async_send_compressed_data = AsyncMock(side_effect=_make_send(0, []))
await logger.async_send_batch()
await logger.async_send_batch()
assert logger.log_queue == []
@pytest.mark.asyncio
async def test_413_drops_single_oversized_event(datadog_env):
with patch("asyncio.create_task"):
logger = DataDogLogger()
logger.log_queue = _payloads(1)
send = AsyncMock(side_effect=_make_send(0, []))
logger.async_send_compressed_data = send
await logger.async_send_batch()
assert send.await_count == 1
assert logger.log_queue == []
@pytest.mark.asyncio
async def test_413_returned_response_also_splits(datadog_env):
"""Defensive path: a 413 returned (not raised) is handled the same way."""
with patch("asyncio.create_task"):
logger = DataDogLogger()
logger.log_queue = _payloads(4)
delivered: list = []
logger.async_send_compressed_data = AsyncMock(
side_effect=_make_send(1, delivered, raise_413=False)
)
await logger.async_send_batch()
assert sorted(delivered) == [f'{{"event": {i}}}' for i in range(4)]
assert logger.log_queue == []
@pytest.mark.asyncio
async def test_partial_delivery_then_transient_error_requeues_only_undelivered(
datadog_env,
):
"""A transient error after a partial split delivery must not duplicate delivered events."""
with patch("asyncio.create_task"):
logger = DataDogLogger()
logger.log_queue = _payloads(4)
delivered: list = []
async def _send(data):
messages = [event["message"] for event in data]
if len(data) > 2:
raise _raised_413()
if messages == ['{"event": 2}', '{"event": 3}']:
raise RuntimeError("transient network error")
delivered.extend(messages)
return Response(
202, request=Request("POST", "https://example.com"), text="Accepted"
)
for i in range(2)
logger.async_send_compressed_data = AsyncMock(side_effect=_send)
await logger.async_send_batch()
assert delivered == ['{"event": 0}', '{"event": 1}']
assert [event["message"] for event in logger.log_queue] == [
'{"event": 2}',
'{"event": 3}',
]
@pytest.mark.asyncio
async def test_unexpected_non_202_status_requeues(datadog_env):
"""A non-413, non-202 response is treated as undelivered and re-queued."""
with patch("asyncio.create_task"):
logger = DataDogLogger()
logger.log_queue = _payloads(2)
logger.async_send_compressed_data = AsyncMock(
return_value=Response(
413,
request=Request("POST", "https://example.com"),
text="Payload Too Large",
200, request=Request("POST", "https://example.com"), text="OK"
)
)
await logger.async_send_batch()
assert logger.async_send_compressed_data.await_count == 1
assert len(logger.log_queue) == 2
assert [event["message"] for event in logger.log_queue] == [
'{"event": 0}',
'{"event": 1}',
]
@pytest.mark.parametrize(
"value, expected",
[
("50", 50),
("1", 1),
("0", 1),
("-5", 1),
(str(DD_MAX_BATCH_SIZE + 100), DD_MAX_BATCH_SIZE),
("not_an_int", DD_MAX_BATCH_SIZE),
],
)
def test_dd_batch_size_env_resolution(monkeypatch, value, expected):
monkeypatch.setenv("DD_API_KEY", "test_api_key")
monkeypatch.setenv("DD_SITE", "test.datadoghq.com")
monkeypatch.setenv("DD_BATCH_SIZE", value)
with patch("asyncio.create_task"):
logger = DataDogLogger()
assert logger.batch_size == expected
def test_dd_batch_size_defaults_to_max(monkeypatch):
monkeypatch.setenv("DD_API_KEY", "test_api_key")
monkeypatch.setenv("DD_SITE", "test.datadoghq.com")
monkeypatch.delenv("DD_BATCH_SIZE", raising=False)
with patch("asyncio.create_task"):
logger = DataDogLogger()
assert logger.batch_size == DD_MAX_BATCH_SIZE
@pytest.mark.asyncio
async def test_async_send_batch_handles_empty_queue(datadog_env):
with patch("asyncio.create_task"):

354
tests/test_litellm/integrations/test_anthropic_cache_control_hook.py
View file

@ -1087,3 +1087,357 @@ async def test_anthropic_cache_control_hook_string_negative_index():
f"Expected cachePoint in last message content, got: {last_message_content}. "
"String index '-1' was not parsed correctly (str.isdigit() returns False for negative strings)."
)
def _count_cache_control(messages: List[AllMessageValues]) -> int:
"""Count cache_control breakpoints across messages (message + content level)."""
count = 0
for message in messages:
if message.get("cache_control") is not None:
count += 1
content = message.get("content")
if isinstance(content, list):
for block in content:
if isinstance(block, dict) and block.get("cache_control") is not None:
count += 1
return count
def _build_injection_points():
return [
{
"location": "message",
"role": "system",
"control": {"type": "ephemeral", "ttl": "1h"},
},
{
"location": "message",
"index": -1,
"control": {"type": "ephemeral", "ttl": "5m"},
},
]
def test_cache_control_hook_caps_at_four_blocks_with_client_cache_control():
"""Regression for LIT-3667 / Anthropic 'A maximum of 4 blocks ... Found 5'.
A Hermes-style request already carries 4 client cache_control breakpoints on
its system messages. With both auto-inject points configured the hook must
NOT add a 5th breakpoint, and must NOT overwrite the client's existing
breakpoints (TTL must be preserved).
"""
hook = AnthropicCacheControlHook()
messages: List[AllMessageValues] = [
{
"role": "system",
"content": [
{
"type": "text",
"text": f"System block {i}",
"cache_control": {"type": "ephemeral", "ttl": "1h"},
}
],
}
for i in range(4)
]
messages.append({"role": "user", "content": "hello"})
_, processed, _ = hook.get_chat_completion_prompt(
model="bedrock/us.anthropic.claude-opus-4-6-v1:0",
messages=messages,
non_default_params={
"cache_control_injection_points": _build_injection_points()
},
prompt_id=None,
prompt_variables=None,
dynamic_callback_params={},
)
assert (
_count_cache_control(processed) == 4
), "Hook must cap cache_control at Anthropic's limit of 4 blocks"
# Client TTL on system blocks must be preserved (not overwritten by config).
for i in range(4):
assert processed[i]["content"][-1]["cache_control"] == {
"type": "ephemeral",
"ttl": "1h",
}
# The last (user) message must not receive a 5th breakpoint.
user_message = processed[-1]
assert user_message.get("cache_control") is None
user_content = user_message.get("content")
if isinstance(user_content, list):
assert all(
block.get("cache_control") is None
for block in user_content
if isinstance(block, dict)
)
def test_cache_control_hook_caps_at_four_blocks_without_client_cache_control():
"""Four plain system messages + role:system + index:-1 must stay at 4 blocks.
role:system fills all four slots, so the index:-1 point is skipped.
"""
hook = AnthropicCacheControlHook()
messages: List[AllMessageValues] = [
{"role": "system", "content": f"System {i}"} for i in range(4)
]
messages.append({"role": "user", "content": "hello"})
_, processed, _ = hook.get_chat_completion_prompt(
model="bedrock/us.anthropic.claude-opus-4-6-v1:0",
messages=messages,
non_default_params={
"cache_control_injection_points": _build_injection_points()
},
prompt_id=None,
prompt_variables=None,
dynamic_callback_params={},
)
assert _count_cache_control(processed) == 4
# All four system messages cached; user message skipped (limit reached).
assert all(processed[i].get("cache_control") is not None for i in range(4))
assert processed[-1].get("cache_control") is None
def test_cache_control_hook_does_not_overwrite_existing_cache_control():
"""If a targeted message already has client cache_control, do not inject."""
hook = AnthropicCacheControlHook()
messages: List[AllMessageValues] = [
{
"role": "system",
"content": [
{
"type": "text",
"text": "Cached by client",
"cache_control": {"type": "ephemeral", "ttl": "1h"},
}
],
},
{"role": "user", "content": "hello"},
]
_, processed, _ = hook.get_chat_completion_prompt(
model="bedrock/us.anthropic.claude-opus-4-6-v1:0",
messages=messages,
# Target the already-cached system message with a different TTL.
non_default_params={
"cache_control_injection_points": [
{
"location": "message",
"index": 0,
"control": {"type": "ephemeral", "ttl": "5m"},
}
]
},
prompt_id=None,
prompt_variables=None,
dynamic_callback_params={},
)
# Client's 1h TTL must be preserved, not replaced by the config's 5m.
assert processed[0]["content"][-1]["cache_control"] == {
"type": "ephemeral",
"ttl": "1h",
}
assert _count_cache_control(processed) == 1
@pytest.mark.asyncio
async def test_cache_control_hook_bedrock_payload_caps_cachepoints_at_four():
"""End-to-end: outgoing Bedrock payload must not exceed 4 cachePoint blocks.
Reproduces the customer report where 4 client cache_control system blocks
plus auto-inject produced 5 cachePoint blocks and Bedrock returned 400.
"""
with patch.dict(
os.environ,
{
"AWS_ACCESS_KEY_ID": "fake_access_key_id",
"AWS_SECRET_ACCESS_KEY": "fake_secret_access_key",
"AWS_REGION_NAME": "us-east-1",
},
):
litellm.callbacks = [AnthropicCacheControlHook()]
mock_response = MagicMock()
mock_response.json.return_value = {
"output": {"message": {"role": "assistant", "content": "ok"}},
"stopReason": "end_turn",
"usage": {"inputTokens": 100, "outputTokens": 4, "totalTokens": 104},
}
mock_response.status_code = 200
client = AsyncHTTPHandler()
with patch.object(client, "post", return_value=mock_response) as mock_post:
messages = [
{
"role": "system",
"content": [
{
"type": "text",
"text": f"System block {i}",
"cache_control": {"type": "ephemeral", "ttl": "1h"},
}
],
}
for i in range(4)
]
messages.append({"role": "user", "content": "hello"})
await litellm.acompletion(
model="bedrock/us.anthropic.claude-opus-4-6-v1:0",
messages=messages,
max_tokens=32,
cache_control_injection_points=_build_injection_points(),
client=client,
)
request_body = json.loads(mock_post.call_args.kwargs["data"])
cache_points = sum(
1
for block in request_body.get("system", [])
if isinstance(block, dict) and "cachePoint" in block
)
for msg in request_body.get("messages", []):
content = msg.get("content", [])
if isinstance(content, list):
cache_points += sum(
1
for block in content
if isinstance(block, dict) and "cachePoint" in block
)
assert cache_points <= 4, (
f"Bedrock payload exceeded Anthropic's 4 cache_control block limit: "
f"found {cache_points} cachePoint blocks"
)
def test_cache_control_hook_reserves_slot_for_tool_config_point():
"""A tool_config injection point consumes one of the 4 slots downstream.
With role:system targeting 4 system messages plus a tool_config point, the
hook must inject at most 3 message-level blocks so the tool_config cachePoint
appended by the Bedrock transform keeps the total at 4, not 5.
"""
hook = AnthropicCacheControlHook()
messages: List[AllMessageValues] = [
{"role": "system", "content": f"System {i}"} for i in range(4)
]
messages.append({"role": "user", "content": "hello"})
_, processed, non_default_params = hook.get_chat_completion_prompt(
model="bedrock/us.anthropic.claude-opus-4-6-v1:0",
messages=messages,
non_default_params={
"cache_control_injection_points": [
{
"location": "message",
"role": "system",
"control": {"type": "ephemeral", "ttl": "1h"},
},
{"location": "tool_config"},
]
},
prompt_id=None,
prompt_variables=None,
dynamic_callback_params={},
)
assert _count_cache_control(processed) == 3
# The tool_config point is passed through for the provider transform.
assert non_default_params["cache_control_injection_points"] == [
{"location": "tool_config"}
]
@pytest.mark.asyncio
async def test_cache_control_hook_bedrock_payload_caps_with_tool_config_point():
"""End-to-end: message + tool_config injection must not exceed 4 cachePoints."""
with patch.dict(
os.environ,
{
"AWS_ACCESS_KEY_ID": "fake_access_key_id",
"AWS_SECRET_ACCESS_KEY": "fake_secret_access_key",
"AWS_REGION_NAME": "us-east-1",
},
):
litellm.callbacks = [AnthropicCacheControlHook()]
mock_response = MagicMock()
mock_response.json.return_value = {
"output": {"message": {"role": "assistant", "content": "ok"}},
"stopReason": "end_turn",
"usage": {"inputTokens": 100, "outputTokens": 4, "totalTokens": 104},
}
mock_response.status_code = 200
client = AsyncHTTPHandler()
with patch.object(client, "post", return_value=mock_response) as mock_post:
messages = [
{"role": "system", "content": f"System block {i}"} for i in range(4)
]
messages.append({"role": "user", "content": "What is the weather?"})
await litellm.acompletion(
model="bedrock/us.anthropic.claude-opus-4-6-v1:0",
messages=messages,
max_tokens=32,
tools=[
{
"type": "function",
"function": {
"name": "get_weather",
"description": "Get weather for a location",
"parameters": {
"type": "object",
"properties": {"location": {"type": "string"}},
"required": ["location"],
},
},
}
],
cache_control_injection_points=[
{
"location": "message",
"role": "system",
"control": {"type": "ephemeral", "ttl": "1h"},
},
{"location": "tool_config"},
],
client=client,
)
request_body = json.loads(mock_post.call_args.kwargs["data"])
cache_points = sum(
1
for block in request_body.get("system", [])
if isinstance(block, dict) and "cachePoint" in block
)
for msg in request_body.get("messages", []):
content = msg.get("content", [])
if isinstance(content, list):
cache_points += sum(
1
for block in content
if isinstance(block, dict) and "cachePoint" in block
)
for tool in request_body.get("toolConfig", {}).get("tools", []):
if isinstance(tool, dict) and "cachePoint" in tool:
cache_points += 1
assert cache_points <= 4, (
f"Bedrock payload exceeded Anthropic's 4 cache_control block limit "
f"when mixing message and tool_config injection: found {cache_points}"
)

106
tests/test_litellm/integrations/test_custom_guardrail.py
View file

@ -84,6 +84,112 @@ class TestCustomGuardrailDeploymentHook:
assert result["messages"] == mock_result["messages"]
assert result["messages"] != original_messages
@pytest.mark.asyncio
async def test_deployment_hook_skips_when_pre_call_already_ran(self):
"""The deployment hook must not re-run async_pre_call_hook once the proxy
pre-call loop has already run it for this request."""
class CountingGuardrail(CustomGuardrail):
def __init__(self):
super().__init__(guardrail_name="g1", default_on=True)
self.pre_call_count = 0
async def async_pre_call_hook(
self, user_api_key_dict, cache, data, call_type
):
self.pre_call_count += 1
return data
guardrail = CountingGuardrail()
kwargs = {
"messages": [{"role": "user", "content": "hi"}],
"model": "gpt-3.5-turbo",
"guardrails": ["g1"],
"metadata": {},
}
guardrail.mark_pre_call_hook_ran(kwargs)
await guardrail.async_pre_call_deployment_hook(
kwargs=kwargs, call_type=CallTypes.completion
)
assert guardrail.pre_call_count == 0
@pytest.mark.asyncio
async def test_deployment_hook_runs_when_not_marked(self):
"""Without the proxy marker (direct-SDK usage) the deployment hook is the
only execution path and must still run the guardrail exactly once."""
class CountingGuardrail(CustomGuardrail):
def __init__(self):
super().__init__(guardrail_name="g1", default_on=True)
self.pre_call_count = 0
async def async_pre_call_hook(
self, user_api_key_dict, cache, data, call_type
):
self.pre_call_count += 1
return data
guardrail = CountingGuardrail()
kwargs = {
"messages": [{"role": "user", "content": "hi"}],
"model": "gpt-3.5-turbo",
"guardrails": ["g1"],
"metadata": {},
}
await guardrail.async_pre_call_deployment_hook(
kwargs=kwargs, call_type=CallTypes.completion
)
assert guardrail.pre_call_count == 1
def test_mark_pre_call_hook_ran_uses_litellm_metadata(self):
"""The marker is recorded in litellm_metadata when that is the metadata
bucket in use, and is then visible to the skip check."""
from litellm.constants import PRE_CALL_EXECUTED_GUARDRAILS_KEY
guardrail = CustomGuardrail(guardrail_name="g1")
kwargs = {"litellm_metadata": {}}
guardrail.mark_pre_call_hook_ran(kwargs)
assert kwargs["litellm_metadata"][PRE_CALL_EXECUTED_GUARDRAILS_KEY]
assert guardrail._pre_call_hook_already_ran(kwargs) is True
@pytest.mark.asyncio
async def test_deployment_hook_ignores_forged_caller_marker(self):
"""A direct-SDK caller controls request metadata but cannot know the
per-process token, so a hand-crafted marker must not suppress a
requested guardrail in async_pre_call_deployment_hook."""
from litellm.constants import PRE_CALL_EXECUTED_GUARDRAILS_KEY
class CountingGuardrail(CustomGuardrail):
def __init__(self):
super().__init__(guardrail_name="g1", default_on=True)
self.pre_call_count = 0
async def async_pre_call_hook(
self, user_api_key_dict, cache, data, call_type
):
self.pre_call_count += 1
return data
guardrail = CountingGuardrail()
kwargs = {
"messages": [{"role": "user", "content": "hi"}],
"model": "gpt-3.5-turbo",
"guardrails": ["g1"],
"metadata": {PRE_CALL_EXECUTED_GUARDRAILS_KEY: ["g1"]},
}
await guardrail.async_pre_call_deployment_hook(
kwargs=kwargs, call_type=CallTypes.completion
)
assert guardrail.pre_call_count == 1
class TestCustomGuardrailShouldRunGuardrail:

212
tests/test_litellm/litellm_core_utils/test_litellm_logging.py
View file

@ -1,6 +1,7 @@
import os
import sys
from unittest.mock import MagicMock, patch
import asyncio
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
@ -786,6 +787,211 @@ def test_success_handler_runs_sync_callbacks_for_sync_requests(logging_obj, call
dummy_logger.log_stream_event.assert_not_called()
def test_is_sync_litellm_request():
assert LitellmLogging._is_sync_litellm_request({}) is True
assert LitellmLogging._is_sync_litellm_request({"acompletion": True}) is False
@pytest.mark.asyncio
async def test_dispatch_success_handlers_invokes_callbacks_once_for_final_stream(
logging_obj,
):
"""Second final-stream dispatch must not re-export (CSW + deferred guardrail paths)."""
import litellm
from litellm.integrations.custom_logger import CustomLogger
class MockCallback(CustomLogger):
pass
mock_callback = MockCallback()
original_async_callbacks = list(litellm._async_success_callback or [])
litellm._async_success_callback = [mock_callback]
result = ModelResponse(
id="resp-dedupe",
model="gpt-4o-mini",
choices=[
{
"message": {"role": "assistant", "content": "hi"},
"finish_reason": "stop",
"index": 0,
}
],
usage={"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
)
try:
logging_obj.stream = True
logging_obj.model_call_details["litellm_params"] = {"acompletion": True}
with (
patch.object(
mock_callback, "async_log_success_event", new_callable=AsyncMock
) as mock_async_log,
patch.object(mock_callback, "log_success_event") as mock_sync_log,
patch.object(
logging_obj,
"_success_handler_helper_fn",
return_value=(time.time(), time.time(), result),
),
patch.object(
logging_obj,
"_get_assembled_streaming_response",
return_value=result,
),
patch.object(
logging_obj,
"_should_run_sync_callbacks_for_async_calls",
return_value=True,
),
):
await logging_obj.dispatch_success_handlers(result=result)
await logging_obj.dispatch_success_handlers(result=result)
mock_async_log.assert_awaited_once()
mock_sync_log.assert_not_called()
finally:
litellm._async_success_callback = original_async_callbacks
@pytest.mark.asyncio
async def test_dispatch_success_handlers_sync_path_invokes_callback_once_for_final_stream(
logging_obj,
):
"""Sync dispatch path must also dedupe when dispatch is called twice."""
import litellm
from litellm.integrations.custom_logger import CustomLogger
class MockCallback(CustomLogger):
pass
mock_callback = MockCallback()
original_success_callbacks = list(litellm.success_callback or [])
litellm.success_callback = [mock_callback]
result = ModelResponse(
id="resp-sync-dedupe",
model="gpt-4o-mini",
choices=[
{
"message": {"role": "assistant", "content": "hi"},
"finish_reason": "stop",
"index": 0,
}
],
usage={"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
)
try:
logging_obj.stream = True
logging_obj.model_call_details["litellm_params"] = {}
with (
patch.object(mock_callback, "log_success_event") as mock_sync_log,
patch.object(
mock_callback, "async_log_success_event", new_callable=AsyncMock
) as mock_async_log,
patch.object(
logging_obj,
"_success_handler_helper_fn",
return_value=(time.time(), time.time(), result),
),
patch.object(
logging_obj,
"_get_assembled_streaming_response",
return_value=result,
),
):
await logging_obj.dispatch_success_handlers(result=result)
await logging_obj.dispatch_success_handlers(result=result)
mock_sync_log.assert_called_once()
mock_async_log.assert_not_awaited()
finally:
litellm.success_callback = original_success_callbacks
@pytest.mark.asyncio
async def test_dispatch_prefer_async_handlers_runs_legacy_callbacks(
logging_obj,
):
"""``prefer_async_handlers`` must not skip executor.submit for string callbacks."""
result = ModelResponse(
id="resp-prefer-async",
model="gpt-4o-mini",
choices=[
{
"message": {"role": "assistant", "content": "hi"},
"finish_reason": "stop",
"index": 0,
}
],
)
logging_obj.stream = True
logging_obj.model_call_details["litellm_params"] = {}
with (
patch.object(
logging_obj, "async_success_handler", new_callable=AsyncMock
) as mock_async,
patch.object(
logging_obj, "success_handler", new_callable=MagicMock
) as mock_sync,
patch.object(
logging_obj,
"_should_run_sync_callbacks_for_async_calls",
return_value=True,
),
patch(
"litellm.litellm_core_utils.litellm_logging.executor.submit"
) as mock_submit,
):
await logging_obj.dispatch_success_handlers(
result=result,
prefer_async_handlers=True,
)
mock_async.assert_awaited_once()
mock_sync.assert_not_called()
mock_submit.assert_called_once()
@pytest.mark.asyncio
async def test_dispatch_success_handlers_invokes_async_callback_for_pass_through(
logging_obj,
):
"""Pass-through must use async_success_handler (CustomLogger skips sync success_handler)."""
import litellm
from litellm.integrations.custom_logger import CustomLogger
from litellm.types.utils import CallTypes
class MockCallback(CustomLogger):
pass
mock_callback = MockCallback()
original_async_callbacks = list(litellm._async_success_callback or [])
litellm._async_success_callback = [mock_callback]
logging_obj.call_type = CallTypes.pass_through.value
logging_obj.stream = False
logging_obj.model_call_details["litellm_params"] = {}
try:
with (
patch.object(
mock_callback, "async_log_success_event", new_callable=AsyncMock
) as mock_async_log,
patch.object(mock_callback, "log_success_event") as mock_sync_log,
):
await logging_obj.dispatch_success_handlers(result={"id": "pt-1"})
mock_async_log.assert_awaited_once()
mock_sync_log.assert_not_called()
finally:
litellm._async_success_callback = original_async_callbacks
def test_success_handler_skips_guardrail_logging_hook_when_disabled(logging_obj):
"""Ensure CustomGuardrail logging_hook is skipped when should_run_guardrail is False."""
import datetime
@ -1351,7 +1557,7 @@ async def test_e2e_generate_cold_storage_object_key_with_custom_logger_s3_path()
Test that _generate_cold_storage_object_key uses s3_path from custom logger instance.
"""
from datetime import datetime, timezone
from unittest.mock import MagicMock, patch
from unittest.mock import AsyncMock, MagicMock, patch
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
@ -1404,7 +1610,7 @@ async def test_e2e_generate_cold_storage_object_key_with_logger_no_s3_path():
Test that _generate_cold_storage_object_key falls back to empty s3_path when logger has no s3_path.
"""
from datetime import datetime, timezone
from unittest.mock import MagicMock, patch
from unittest.mock import AsyncMock, MagicMock, patch
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup

16
tests/test_litellm/litellm_core_utils/test_streaming_handler.py
View file

@ -569,8 +569,6 @@ async def test_streaming_with_usage_and_logging(sync_mode: bool):
== final_usage_block
)
print(mock_log_success_event.call_args.kwargs.keys())
def test_streaming_handler_with_stop_chunk(
initialized_custom_stream_wrapper: CustomStreamWrapper,
@ -2036,23 +2034,19 @@ async def test_azure_streaming_role_preserved_with_include_usage(sync_mode: bool
chunks.append(chunk)
# The prompt_filter chunk should be forwarded with choices=[]
assert len(chunks[0].choices) == 0, (
f"Expected prompt_filter chunk with choices=[], got {len(chunks[0].choices)} choices"
)
assert (
len(chunks[0].choices) == 0
), f"Expected prompt_filter chunk with choices=[], got {len(chunks[0].choices)} choices"
# At least one chunk must have role='assistant' in its delta
has_role = any(
len(c.choices) > 0
and getattr(c.choices[0].delta, "role", None) == "assistant"
len(c.choices) > 0 and getattr(c.choices[0].delta, "role", None) == "assistant"
for c in chunks
)
assert has_role, (
"No chunk contained role='assistant' in delta (issue #24221). "
"Chunk deltas: "
+ str([
c.choices[0].delta if c.choices else "no choices"
for c in chunks
])
+ str([c.choices[0].delta if c.choices else "no choices" for c in chunks])
)

74
tests/test_litellm/llms/openai/test_use_chat_completions_api_no_leak.py Normal file
View file

@ -0,0 +1,74 @@
"""
Regression test for issue #28146.
`use_chat_completions_api` is a LiteLLM-internal control flag (it forces the
/responses -> /chat/completions bridge). When set as a model-level param in the
proxy config, it must never be forwarded to the upstream provider's request
body. OpenAI/Anthropic reject unknown body params with HTTP 400.
"""
import os
import sys
from unittest.mock import MagicMock
sys.path.insert(0, os.path.abspath("../../../.."))
import litellm
from litellm.types.utils import all_litellm_params
from litellm.utils import get_non_default_completion_params
def test_use_chat_completions_api_is_a_known_litellm_param():
assert "use_chat_completions_api" in all_litellm_params
def test_use_chat_completions_api_not_forwarded_as_provider_param():
forwarded = get_non_default_completion_params(
{"use_chat_completions_api": True, "temperature": 0.5}
)
assert "use_chat_completions_api" not in forwarded
def test_completion_does_not_leak_flag_into_provider_request_body():
mock_response = MagicMock()
mock_response.model_dump.return_value = {
"id": "chatcmpl-1",
"object": "chat.completion",
"created": 1234567890,
"model": "gpt-4o-mini",
"choices": [
{
"index": 0,
"message": {"role": "assistant", "content": "hi"},
"finish_reason": "stop",
}
],
"usage": {
"prompt_tokens": 1,
"completion_tokens": 1,
"total_tokens": 2,
},
}
mock_raw_response = MagicMock()
mock_raw_response.headers = {}
mock_raw_response.parse.return_value = mock_response
mock_client = MagicMock()
mock_client.chat.completions.with_raw_response.create.return_value = (
mock_raw_response
)
litellm.completion(
model="openai/gpt-4o-mini",
messages=[{"role": "user", "content": "hi"}],
use_chat_completions_api=True,
api_key="sk-test",
client=mock_client,
)
create_kwargs = (
mock_client.chat.completions.with_raw_response.create.call_args.kwargs
)
assert "use_chat_completions_api" not in create_kwargs
assert "use_chat_completions_api" not in (create_kwargs.get("extra_body") or {})

103
tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py
View file

@ -18,7 +18,7 @@ import asyncio
import os
import sys
from typing import Any
from unittest.mock import MagicMock, patch
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
@ -38,6 +38,24 @@ from litellm.types.guardrails import GuardrailEventHooks
# ---------------------------------------------------------------------------
def _attach_mock_success_dispatch(mock_logging_obj, async_success_fn):
"""Match production entrypoint: ``_run_deferred_stream_guardrails`` uses dispatch."""
async def dispatch_success_handlers(
result=None, start_time=None, end_time=None, cache_hit=None, **kwargs
):
await async_success_fn(
result,
start_time=start_time,
end_time=end_time,
cache_hit=cache_hit,
**kwargs,
)
mock_logging_obj.dispatch_success_handlers = dispatch_success_handlers
mock_logging_obj.async_success_handler = async_success_fn
class PostCallGuardrail(CustomGuardrail):
"""A post-call guardrail."""
@ -454,7 +472,7 @@ class TestDeferredStreamingClosure:
async def track_async_success(*args, **kwargs):
pass
mock_logging_obj.async_success_handler = track_async_success
_attach_mock_success_dispatch(mock_logging_obj, track_async_success)
tracking_guardrail = TrackingGuardrail()
tracking_logger = TrackingLogger()
@ -511,7 +529,7 @@ class TestDeferredStreamingClosure:
nonlocal logged_response
logged_response = args[0] if args else None
mock_logging_obj.async_success_handler = track_async_success
_attach_mock_success_dispatch(mock_logging_obj, track_async_success)
class ModifyingGuardrail(CustomGuardrail):
def __init__(self):
@ -573,7 +591,7 @@ class TestDeferredStreamingClosure:
nonlocal logging_called
logging_called = True
mock_logging_obj.async_success_handler = track_async_success
_attach_mock_success_dispatch(mock_logging_obj, track_async_success)
guardrail = BlockingGuardrail()
@ -621,7 +639,7 @@ class TestDeferredStreamingClosure:
async def track_async_success(*args, **kwargs):
pass
mock_logging_obj.async_success_handler = track_async_success
_attach_mock_success_dispatch(mock_logging_obj, track_async_success)
guardrail = TransientErrorGuardrail()
@ -656,7 +674,7 @@ class TestDeferredStreamingClosure:
nonlocal logged_response
logged_response = args[0] if args else None
mock_logging_obj.async_success_handler = track_async_success
_attach_mock_success_dispatch(mock_logging_obj, track_async_success)
class TestGuardrail(CustomGuardrail):
def __init__(self):
@ -739,7 +757,7 @@ class TestDeferredStreamingClosure:
async def track_async_success(*args, **kwargs):
pass
mock_logging_obj.async_success_handler = track_async_success
_attach_mock_success_dispatch(mock_logging_obj, track_async_success)
guardrail = ApplyGuardrailType()
@ -792,7 +810,7 @@ class TestDeferredStreamingClosure:
async def track_async_success(*args, **kwargs):
pass
mock_logging_obj.async_success_handler = track_async_success
_attach_mock_success_dispatch(mock_logging_obj, track_async_success)
guardrail = IteratorHookGuardrail()
@ -847,7 +865,7 @@ class TestDeferredStreamingClosure:
async def track_async_success(*args, **kwargs):
pass
mock_logging_obj.async_success_handler = track_async_success
_attach_mock_success_dispatch(mock_logging_obj, track_async_success)
guardrail = InspectingGuardrail()
@ -914,7 +932,7 @@ class TestDeferredStreamingClosure:
async def track_async_success(*args, **kwargs):
pass
mock_logging_obj.async_success_handler = track_async_success
_attach_mock_success_dispatch(mock_logging_obj, track_async_success)
guardrail_a = TaggedGuardrail("guardrail-a")
guardrail_b = TaggedGuardrail("guardrail-b")
@ -962,7 +980,7 @@ class TestDeferredStreamingClosure:
nonlocal logging_called
logging_called = True
mock_logging_obj.async_success_handler = track_async_success
_attach_mock_success_dispatch(mock_logging_obj, track_async_success)
def exploding_merge(data, llm_router):
raise RuntimeError("Simulated init failure")
@ -986,6 +1004,67 @@ class TestDeferredStreamingClosure:
logging_called is True
), "Logging must fire even when guardrail initialization raises"
@pytest.mark.asyncio
async def test_deferred_logging_forces_async_for_sync_classified_call_type(self):
"""
Regression: proxy deferred streaming logging must reach the async success
handler (which runs the async-only DB/spend logger) even when the call
type is classified as a sync SDK request by _is_sync_litellm_request.
Without prefer_async_handlers=True, an async proxy stream whose
litellm_params lacks a recognized async marker would enter the sync
branch of dispatch_success_handlers and silently skip spend tracking.
Uses the real dispatch_success_handlers via the production
_run_deferred_stream_guardrails entrypoint.
"""
import time
from litellm.litellm_core_utils.litellm_logging import (
Logging as LiteLLMLoggingObj,
)
logging_obj = LiteLLMLoggingObj(
model="gpt-4o-mini",
messages=[{"role": "user", "content": "hi"}],
stream=True,
call_type="completion", # not pass_through_endpoint
start_time=time.time(),
litellm_call_id="test-id",
function_id="fn",
)
# litellm_params with no recognized async marker -> classified sync.
logging_obj.model_call_details["litellm_params"] = {}
assert LiteLLMLoggingObj._is_sync_litellm_request({}) is True
with (
patch.object(
logging_obj, "async_success_handler", new_callable=AsyncMock
) as mock_async,
patch.object(
logging_obj, "success_handler", new_callable=MagicMock
) as mock_sync,
patch.object(
logging_obj,
"_should_run_sync_callbacks_for_async_calls",
return_value=False,
),
patch("litellm.callbacks", [PostCallGuardrail()]),
):
await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails(
captured_data={"model": "gpt-4o-mini", "metadata": {}},
captured_user_api_key_dict=UserAPIKeyAuth(api_key="test"),
captured_logging_obj=logging_obj,
assembled_response=MagicMock(),
cache_hit=False,
)
await asyncio.sleep(0)
await asyncio.sleep(0)
mock_async.assert_awaited_once()
mock_sync.assert_not_called()
# ---------------------------------------------------------------------------
# 7. _fire_deferred_stream_logging
@ -1054,7 +1133,7 @@ class TestFireDeferredStreamLogging:
nonlocal logged_response
logged_response = args[0] if args else None
mock_logging_obj.async_success_handler = track_async_success
_attach_mock_success_dispatch(mock_logging_obj, track_async_success)
class InfoWritingGuardrail(CustomGuardrail):
def __init__(self):

169
tests/test_litellm/proxy/guardrails/test_guardrail_registry.py
View file

@ -180,3 +180,172 @@ def test_sync_guardrail_from_db_marks_source_db_when_unchanged():
handler.sync_guardrail_from_db(g)
assert handler.get_source("collide") == "db"
def _db_litellm_params() -> dict:
"""
Shape produced by GuardrailRegistry.get_all_guardrails_from_db: litellm_params
is a raw dict (not a LitellmParams), holding only the keys originally stored,
a non-schema extra key, and plain-string enum values.
"""
return {
"guardrail": "litellm_content_filter",
"mode": "pre_call",
"default_on": True,
"version": 2,
"blocked_words": [{"keyword": "secret", "action": "BLOCK"}],
}
def test_unchanged_db_params_do_not_register_as_changed():
"""
A DB poll returns litellm_params as a raw dict while the in-memory copy is a
LitellmParams whose model_dump() fills every field default and coerces enums.
The two shapes must compare equal when the config is identical; otherwise
every poll cycle re-initializes the guardrail indefinitely.
"""
handler = InMemoryGuardrailHandler()
raw = _db_litellm_params()
gid = "11111111-1111-1111-1111-111111111111"
handler.IN_MEMORY_GUARDRAILS[gid] = Guardrail(
guardrail_id=gid,
guardrail_name="cf",
litellm_params=LitellmParams(**raw),
)
new = Guardrail(guardrail_id=gid, guardrail_name="cf", litellm_params=dict(raw))
assert handler._has_guardrail_params_changed(gid, new) is False
def test_changed_db_params_register_as_changed():
"""Normalizing both sides must still surface a genuine config change."""
handler = InMemoryGuardrailHandler()
raw = _db_litellm_params()
gid = "22222222-2222-2222-2222-222222222222"
handler.IN_MEMORY_GUARDRAILS[gid] = Guardrail(
guardrail_id=gid,
guardrail_name="cf",
litellm_params=LitellmParams(**raw),
)
changed = {**raw, "blocked_words": [{"keyword": "different", "action": "BLOCK"}]}
new = Guardrail(guardrail_id=gid, guardrail_name="cf", litellm_params=changed)
assert handler._has_guardrail_params_changed(gid, new) is True
def test_unnormalizable_db_params_register_as_changed_without_raising():
"""
A DB row whose litellm_params fail LitellmParams validation must not crash the
poll loop. The comparison falls back to treating the guardrail as changed so it
re-initializes (and surfaces the bad row in logs) rather than propagating the
validation error up through the polling cycle.
"""
handler = InMemoryGuardrailHandler()
raw = _db_litellm_params()
gid = "55555555-5555-5555-5555-555555555555"
handler.IN_MEMORY_GUARDRAILS[gid] = Guardrail(
guardrail_id=gid,
guardrail_name="cf",
litellm_params=LitellmParams(**raw),
)
malformed = {**raw, "default_on": "not-a-bool-xyz"}
new = Guardrail(guardrail_id=gid, guardrail_name="cf", litellm_params=malformed)
assert handler._has_guardrail_params_changed(gid, new) is True
def _all_callback_lists():
import litellm
return [
litellm.callbacks,
litellm.success_callback,
litellm.failure_callback,
litellm._async_success_callback,
litellm._async_failure_callback,
]
def test_delete_in_memory_guardrail_removes_callback_from_all_lists():
"""
Request handling promotes guardrail callbacks from litellm.callbacks into the
success/failure/async lists. delete_in_memory_guardrail must purge the callback
from every list, otherwise a re-initialized guardrail leaves its old instance
stranded in those lists and instances accumulate.
"""
handler = InMemoryGuardrailHandler()
callback = CustomGuardrail(
guardrail_name="cf-delete",
default_on=True,
event_hook=GuardrailEventHooks.pre_call,
)
gid = "33333333-3333-3333-3333-333333333333"
handler.IN_MEMORY_GUARDRAILS[gid] = _make_guardrail(gid, "cf-delete")
handler._sources[gid] = "db"
handler.guardrail_id_to_custom_guardrail[gid] = callback
lists = _all_callback_lists()
snapshots = [list(cb_list) for cb_list in lists]
try:
for cb_list in lists:
cb_list.append(callback)
handler.delete_in_memory_guardrail(gid)
for cb_list in lists:
assert callback not in cb_list
finally:
for cb_list, snapshot in zip(lists, snapshots):
cb_list[:] = snapshot
def test_repeated_db_sync_does_not_accumulate_runner_instances():
"""
End-to-end regression for the OOM: across repeated DB polls (with the config
genuinely changing each cycle to force re-initialization), exactly one live
guardrail instance must exist across all callback lists. On the unfixed code
the stale instance lingers in the success/failure lists and the distinct count
climbs above one.
"""
import litellm
handler = InMemoryGuardrailHandler()
gid = "44444444-4444-4444-4444-444444444444"
name = "cf-accum"
def db_guardrail(word: str) -> Guardrail:
params = {
**_db_litellm_params(),
"blocked_words": [{"keyword": word, "action": "BLOCK"}],
}
return Guardrail(guardrail_id=gid, guardrail_name=name, litellm_params=params)
def promote_into_request_lists() -> None:
manager = litellm.logging_callback_manager
for callback in list(litellm.callbacks):
manager.add_litellm_success_callback(callback)
manager.add_litellm_failure_callback(callback)
manager.add_litellm_async_success_callback(callback)
manager.add_litellm_async_failure_callback(callback)
def distinct_runner_instances() -> int:
seen = set()
for callback in litellm.logging_callback_manager._get_all_callbacks():
if (
isinstance(callback, CustomGuardrail)
and getattr(callback, "guardrail_name", None) == name
):
seen.add(id(callback))
return len(seen)
lists = _all_callback_lists()
snapshots = [list(cb_list) for cb_list in lists]
try:
for cycle in range(5):
handler.sync_guardrail_from_db(db_guardrail(f"word-{cycle}"))
promote_into_request_lists()
assert distinct_runner_instances() == 1
finally:
for cb_list, snapshot in zip(lists, snapshots):
cb_list[:] = snapshot

332
tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py
View file

@ -1374,3 +1374,335 @@ class TestPureTextFastPathParity:
AnthropicPassthroughLoggingHandler._collapse_pure_text_chunks(all_chunks)
is None
)
class TestStreamFalseDeduplication:
"""
Regression tests for the duplicate-callback bug where a streaming pass-through
request had stream=False hardcoded on its Logging object.
Before the fix:
- logging_obj.stream was always False for pass-through requests
- _is_assembled_stream_success() checked `self.stream is not True` and returned
False immediately, so has_dispatched_final_stream_success was never set
- Any second dispatch_success_handlers call went through unchecked
After the fix:
- pass_through_endpoints.py sets logging_obj.stream = True after detecting stream
- _create_anthropic_response_logging_payload sets complete_streaming_response on
model_call_details so callbacks see the correct assembled response state
- _is_assembled_stream_success returns True, dedup guard fires on first dispatch
"""
@staticmethod
def _sse(event, data):
return f"event: {event}\ndata: {json.dumps(data)}\n\n".encode()
@staticmethod
def _make_logging_obj(stream: bool = False) -> LiteLLMLoggingObj:
logging_obj = LiteLLMLoggingObj(
model="claude-3-5-sonnet-20241022",
messages=[{"role": "user", "content": "hello"}],
stream=stream,
call_type="pass_through_endpoint",
start_time=datetime.now(),
litellm_call_id="test-call-id",
function_id="1245",
)
return logging_obj
@staticmethod
def _build_chunks():
frames = [
TestStreamFalseDeduplication._sse(
"message_start",
{
"type": "message_start",
"message": {
"id": "msg_abc",
"type": "message",
"role": "assistant",
"model": "claude-3-5-sonnet-20241022",
"content": [],
"stop_reason": None,
"stop_sequence": None,
"usage": {"input_tokens": 10, "output_tokens": 0},
},
},
),
TestStreamFalseDeduplication._sse(
"content_block_start",
{
"type": "content_block_start",
"index": 0,
"content_block": {"type": "text", "text": ""},
},
),
TestStreamFalseDeduplication._sse(
"content_block_delta",
{
"type": "content_block_delta",
"index": 0,
"delta": {"type": "text_delta", "text": "Hello"},
},
),
TestStreamFalseDeduplication._sse(
"content_block_stop", {"type": "content_block_stop", "index": 0}
),
TestStreamFalseDeduplication._sse(
"message_delta",
{
"type": "message_delta",
"delta": {"stop_reason": "end_turn", "stop_sequence": None},
"usage": {"output_tokens": 5},
},
),
TestStreamFalseDeduplication._sse("message_stop", {"type": "message_stop"}),
]
from litellm.proxy.pass_through_endpoints.streaming_handler import (
PassThroughStreamingHandler,
)
return PassThroughStreamingHandler._convert_raw_bytes_to_str_lines(frames)
def test_complete_streaming_response_set_on_model_call_details(self):
"""
After the fix, _create_anthropic_response_logging_payload must set
complete_streaming_response on logging_obj.model_call_details so that
callbacks like _PROXY_track_cost_callback see the assembled response
instead of None.
Before the fix: model_call_details had no complete_streaming_response key.
The log showed: "kwargs stream: True + complete streaming response: None"
"""
from litellm.types.passthrough_endpoints.pass_through_endpoints import (
EndpointType,
)
# pass_through_request sets the stream flag before the streaming handler
# reconstructs the response; mirror that here.
logging_obj = self._make_logging_obj(stream=True)
logging_obj.model_call_details["stream"] = True
all_chunks = list(self._build_chunks())
result = AnthropicPassthroughLoggingHandler._handle_logging_anthropic_collected_chunks(
litellm_logging_obj=logging_obj,
passthrough_success_handler_obj=MagicMock(),
url_route="/anthropic/v1/messages",
request_body={"model": "claude-3-5-sonnet-20241022", "stream": True},
endpoint_type=EndpointType.ANTHROPIC,
start_time=datetime.now(),
all_chunks=all_chunks,
end_time=datetime.now(),
)
# The assembled response must be stored on model_call_details so callbacks
# can identify this as a completed streaming call, not an in-progress one.
assert (
logging_obj.model_call_details.get("complete_streaming_response")
is not None
), "complete_streaming_response must be set on model_call_details after assembly"
# The returned result must match what was stored
assert result["result"] is logging_obj.model_call_details.get(
"complete_streaming_response"
)
def test_dedup_guard_fires_when_stream_true_on_logging_obj(self):
"""
When logging_obj.stream is True (set by pass_through_endpoints.py after
detecting a streaming request), dispatch_success_handlers must set
has_dispatched_final_stream_success=True on the first call so that any
second call is a no-op.
This is the _is_assembled_stream_success gate: with stream=False it
always returned False and the guard was permanently disabled.
"""
from litellm.types.passthrough_endpoints.pass_through_endpoints import (
EndpointType,
)
from litellm.types.utils import ModelResponse
# Simulate what pass_through_endpoints.py now does after stream detection
logging_obj = self._make_logging_obj(stream=False)
logging_obj.stream = True # fix applied
logging_obj.model_call_details["stream"] = True
# Simulate what _create_anthropic_response_logging_payload now does
mock_response = ModelResponse(model="claude-3-5-sonnet-20241022")
logging_obj.model_call_details["complete_streaming_response"] = mock_response
assert logging_obj._is_assembled_stream_success(result=mock_response) is True
# First dispatch sets the flag
assert not logging_obj.model_call_details.get(
"has_dispatched_final_stream_success"
)
logging_obj.model_call_details["has_dispatched_final_stream_success"] = True
# Second dispatch would be blocked — simulate the guard check
would_skip = bool(
logging_obj._is_assembled_stream_success(result=mock_response)
and logging_obj.model_call_details.get(
"has_dispatched_final_stream_success"
)
)
assert would_skip is True, (
"Dedup guard must block a second dispatch_success_handlers call for the "
"same assembled streaming response"
)
def test_sse_fallback_path_sets_stream_true_for_dedup(self):
"""
When a nominally non-streaming request receives an SSE response
(_is_streaming_response returns True), the fallback branch in
pass_through_endpoints.py must set logging_obj.stream = True so the
dedup guard activates.
Before the fix the fallback path never set stream=True, so
_is_assembled_stream_success always returned False and duplicate
callback dispatches were never blocked.
"""
from litellm.types.utils import ModelResponse
# logging_obj starts with stream=False, as created before the request
logging_obj = self._make_logging_obj(stream=False)
assert logging_obj._is_assembled_stream_success(result=MagicMock()) is False
# Simulate what the SSE fallback branch in pass_through_endpoints.py now does
logging_obj.stream = True
logging_obj.model_call_details["stream"] = True
mock_response = ModelResponse(model="claude-3-5-sonnet-20241022")
logging_obj.model_call_details["complete_streaming_response"] = mock_response
# With stream=True the dedup guard must be active
assert logging_obj._is_assembled_stream_success(result=mock_response) is True
logging_obj.model_call_details["has_dispatched_final_stream_success"] = True
would_skip = bool(
logging_obj._is_assembled_stream_success(result=mock_response)
and logging_obj.model_call_details.get(
"has_dispatched_final_stream_success"
)
)
assert would_skip is True
def test_stream_false_logging_obj_bypasses_dedup_guard(self):
"""
Demonstrates the pre-fix state: with stream=False on the logging object,
_is_assembled_stream_success always returns False regardless of whether
complete_streaming_response is set. This means the dedup guard can never
fire, so duplicate dispatches go through unchecked.
This test documents the old broken behavior so the fix is clearly justified.
"""
from litellm.types.utils import ModelResponse
logging_obj = self._make_logging_obj(stream=False)
mock_response = ModelResponse(model="claude-3-5-sonnet-20241022")
logging_obj.model_call_details["complete_streaming_response"] = mock_response
# With stream=False, _is_assembled_stream_success returns False even though
# complete_streaming_response is present — the guard is permanently disabled.
assert logging_obj._is_assembled_stream_success(result=mock_response) is False
class TestNonStreamingResponseRedaction:
"""
Regression tests ensuring _create_anthropic_response_logging_payload only sets
complete_streaming_response for streaming responses. perform_redaction scrubs
that field exclusively when model_call_details["stream"] is True, so storing it
on a non-streaming response would deliver the unredacted response to logging
callbacks when message logging is disabled.
"""
@staticmethod
def _make_logging_obj(stream: bool) -> LiteLLMLoggingObj:
logging_obj = LiteLLMLoggingObj(
model="claude-3-5-sonnet-20241022",
messages=[{"role": "user", "content": "hello"}],
stream=stream,
call_type="pass_through_endpoint",
start_time=datetime.now(),
litellm_call_id="test-call-id",
function_id="1245",
)
# pass_through_request mirrors the stream flag onto model_call_details,
# which is the key perform_redaction inspects.
logging_obj.model_call_details["stream"] = stream
return logging_obj
def test_non_streaming_does_not_set_complete_streaming_response(self):
from litellm.types.utils import ModelResponse
logging_obj = self._make_logging_obj(stream=False)
response = ModelResponse(model="claude-3-5-sonnet-20241022")
AnthropicPassthroughLoggingHandler._create_anthropic_response_logging_payload(
litellm_model_response=response,
model="claude-3-5-sonnet-20241022",
kwargs={},
start_time=datetime.now(),
end_time=datetime.now(),
logging_obj=logging_obj,
)
assert (
"complete_streaming_response" not in logging_obj.model_call_details
), "non-streaming responses must not populate complete_streaming_response"
def test_streaming_sets_complete_streaming_response(self):
from litellm.types.utils import ModelResponse
logging_obj = self._make_logging_obj(stream=True)
response = ModelResponse(model="claude-3-5-sonnet-20241022")
AnthropicPassthroughLoggingHandler._create_anthropic_response_logging_payload(
litellm_model_response=response,
model="claude-3-5-sonnet-20241022",
kwargs={},
start_time=datetime.now(),
end_time=datetime.now(),
logging_obj=logging_obj,
)
assert (
logging_obj.model_call_details.get("complete_streaming_response")
is response
)
def test_non_streaming_response_is_redacted_when_message_logging_off(self):
from litellm.litellm_core_utils.redact_messages import (
redact_message_input_output_from_logging,
)
from litellm.types.utils import Choices, Message, ModelResponse
logging_obj = self._make_logging_obj(stream=False)
response = ModelResponse(
model="claude-3-5-sonnet-20241022",
choices=[Choices(message=Message(role="assistant", content="secret"))],
)
AnthropicPassthroughLoggingHandler._create_anthropic_response_logging_payload(
litellm_model_response=response,
model="claude-3-5-sonnet-20241022",
kwargs={},
start_time=datetime.now(),
end_time=datetime.now(),
logging_obj=logging_obj,
)
logging_obj.model_call_details["litellm_params"] = {
"metadata": {"headers": {"x-litellm-enable-message-redaction": True}}
}
redacted = redact_message_input_output_from_logging(
model_call_details=logging_obj.model_call_details,
result=response,
)
leaked = logging_obj.model_call_details.get("complete_streaming_response")
assert leaked is None
assert redacted.choices[0].message.content == "redacted-by-litellm"

125
tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py
View file

@ -989,6 +989,131 @@ async def test_pass_through_request_contains_proxy_server_request_in_kwargs():
assert metadata["user_api_key_user_id"] == "test-user-id"
@pytest.mark.asyncio
async def test_pass_through_request_streaming_marks_logging_obj_as_stream():
"""
Regression: a streaming pass-through request must flag its logging object as
streaming (logging_obj.stream and model_call_details["stream"]) before the
response is dispatched, so cost/success callbacks treat it as a stream and the
streaming dedup guard fires instead of double-logging.
"""
with patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging:
with patch(
"litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_async_httpx_client"
) as mock_get_client:
with patch(
"litellm.proxy.pass_through_endpoints.pass_through_endpoints.PassThroughStreamingHandler.chunk_processor"
) as mock_chunk_processor:
mock_proxy_logging.pre_call_hook = AsyncMock(
return_value={"model": "claude-3", "stream": True}
)
mock_proxy_logging.post_call_failure_hook = AsyncMock()
upstream_response = MagicMock()
upstream_response.status_code = 200
upstream_response.headers = {}
upstream_response.raise_for_status = MagicMock()
async_client = MagicMock()
async_client.build_request = MagicMock(return_value=MagicMock())
async_client.send = AsyncMock(return_value=upstream_response)
mock_get_client.return_value = MagicMock(client=async_client)
async def _empty_chunks(*args, **kwargs):
return
yield # pragma: no cover
mock_chunk_processor.return_value = _empty_chunks()
mock_request = MagicMock(spec=Request)
mock_request.method = "POST"
mock_request.url = "http://test-proxy.com/v1/messages"
mock_request.body = AsyncMock(
return_value=b'{"model": "claude-3", "stream": true}'
)
mock_request.headers = Headers({})
mock_request.query_params = QueryParams({})
await pass_through_request(
request=mock_request,
target="http://target-api.com/v1/messages",
custom_headers={},
user_api_key_dict=MagicMock(),
stream=True,
)
async_client.send.assert_awaited_once()
assert async_client.send.call_args.kwargs["stream"] is True
mock_chunk_processor.assert_called_once()
logging_obj = mock_chunk_processor.call_args.kwargs[
"litellm_logging_obj"
]
assert logging_obj.stream is True
assert logging_obj.model_call_details["stream"] is True
@pytest.mark.asyncio
async def test_pass_through_request_sse_response_marks_logging_obj_as_stream():
"""
Regression: a request that is not flagged as streaming up front but whose
upstream response comes back as an SSE stream (content-type text/event-stream)
must still flag its logging object as streaming before dispatch. Otherwise the
cost/success callbacks treat the assembled stream as a non-stream and the dedup
guard never fires, double-logging the request.
"""
with patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging:
with patch(
"litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_async_httpx_client"
) as mock_get_client:
with patch(
"litellm.proxy.pass_through_endpoints.pass_through_endpoints.PassThroughStreamingHandler.chunk_processor"
) as mock_chunk_processor:
mock_proxy_logging.pre_call_hook = AsyncMock(
return_value={"model": "claude-3"}
)
mock_proxy_logging.post_call_failure_hook = AsyncMock()
upstream_response = MagicMock()
upstream_response.status_code = 200
upstream_response.headers = {"content-type": "text/event-stream"}
upstream_response.raise_for_status = MagicMock()
async_client = MagicMock()
async_client.request = AsyncMock(return_value=upstream_response)
mock_get_client.return_value = MagicMock(client=async_client)
async def _empty_chunks(*args, **kwargs):
return
yield # pragma: no cover
mock_chunk_processor.return_value = _empty_chunks()
mock_request = MagicMock(spec=Request)
mock_request.method = "POST"
mock_request.url = "http://test-proxy.com/v1/messages"
mock_request.body = AsyncMock(return_value=b'{"model": "claude-3"}')
mock_request.headers = Headers({})
mock_request.query_params = QueryParams({})
await pass_through_request(
request=mock_request,
target="http://target-api.com/v1/messages",
custom_headers={},
user_api_key_dict=MagicMock(),
stream=False,
)
async_client.request.assert_awaited_once()
mock_chunk_processor.assert_called_once()
logging_obj = mock_chunk_processor.call_args.kwargs[
"litellm_logging_obj"
]
assert logging_obj.stream is True
assert logging_obj.model_call_details["stream"] is True
@pytest.mark.asyncio
async def test_create_pass_through_endpoint():
"""

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

@ -2244,6 +2244,41 @@ class TestHandleLLMApiExceptionDictDetail:
assert proxy_exc.message == "Content blocked by guardrail"
assert proxy_exc.provider_specific_fields is None
async def test_already_normalized_proxy_exception_is_honored(self):
"""A ProxyException raised mid-request (e.g. a guardrail block) is already
the OpenAI wire format. The funnel must re-raise it untouched instead of
re-deriving the status from a (nonexistent) status_code attribute and
defaulting to 500. Regression for LIT-3751."""
from litellm.proxy._types import ProxyException
exc = ProxyException(
message='"Leroy Jenkins" detected as name',
type="invalid_request_error",
param=None,
code=400,
openai_code="content_policy_violation",
)
proxy_exc = await self._invoke(exc)
assert proxy_exc is exc
assert proxy_exc.code == "400"
assert proxy_exc.type == "invalid_request_error"
assert proxy_exc.param is None
assert proxy_exc.openai_code == "content_policy_violation"
assert proxy_exc.message == '"Leroy Jenkins" detected as name'
# The body the OpenAI-SDK client actually receives. The HTTP status line
# comes from int(exc.code) == 400; the wire ``code`` stays the status
# string. ``openai_code`` ("content_policy_violation") is intentionally
# NOT serialized here - to_dict() emits only ``code`` - so this asserts
# the real contract rather than the write-only attribute.
assert int(proxy_exc.code) == 400
assert proxy_exc.to_dict() == {
"message": '"Leroy Jenkins" detected as name',
"type": "invalid_request_error",
"param": None,
"code": "400",
}
class TestAsyncStreamingDataGeneratorFastPath:
"""Fast/slow path branching in async_streaming_data_generator."""

152
tests/test_litellm/proxy/test_model_level_guardrails.py
View file

@ -19,7 +19,6 @@ from litellm.proxy.utils import (
_merge_guardrails_with_existing,
)
# ---------------------------------------------------------------------------
# Unit tests for _check_and_merge_model_level_guardrails
# ---------------------------------------------------------------------------
@ -159,6 +158,157 @@ class TestCheckAndMergeModelLevelGuardrails:
assert "existing" in result["metadata"]["guardrails"]
# ---------------------------------------------------------------------------
# Regression test: pre_call hook must run exactly once with model-level guardrails
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_pre_call_hook_runs_once_with_model_level_guardrails():
"""
A guardrail attached at the model level (litellm_params.guardrails) is
spread into the top-level request kwargs by the router. The proxy pre-call
loop (async_pre_call_hook) and the deployment-level hook
(async_pre_call_deployment_hook) must together invoke async_pre_call_hook
exactly once, not twice.
"""
from litellm.caching.caching import DualCache
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.proxy._types import CallTypes, UserAPIKeyAuth
from litellm.proxy.utils import ProxyLogging
from litellm.types.guardrails import GuardrailEventHooks
class CountingGuardrail(CustomGuardrail):
def __init__(self):
super().__init__(
guardrail_name="counting-guardrail",
event_hook=GuardrailEventHooks.pre_call,
default_on=True,
)
self.pre_call_count = 0
async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type):
self.pre_call_count += 1
return data
guardrail = CountingGuardrail()
with patch("litellm.callbacks", [guardrail]):
ProxyLogging._callback_capabilities_cache.clear()
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
user_api_key_dict = UserAPIKeyAuth(api_key="test-key")
data = {
"model": "gpt-4",
"messages": [{"role": "user", "content": "hello"}],
"metadata": {},
}
# Path A: proxy pre-call loop runs the guardrail and records that it ran
data = await proxy_logging.pre_call_hook(
user_api_key_dict=user_api_key_dict,
data=data,
call_type="acompletion",
)
# Path B: the router spreads the deployment's model-level guardrails into
# the top-level kwargs, then litellm.acompletion fires the deployment hook
data["guardrails"] = ["counting-guardrail"]
await guardrail.async_pre_call_deployment_hook(data, CallTypes.acompletion)
assert guardrail.pre_call_count == 1
@pytest.mark.asyncio
async def test_pre_call_hook_runs_once_when_hook_returns_fresh_dict():
"""
async_pre_call_hook may return a brand-new request dict instead of mutating
or spreading the one it received. The exactly-once marker must live on the
data that flows downstream, so the deployment hook still skips the guardrail
even when the proxy loop swapped in a fresh dict that never carried it.
"""
from litellm.caching.caching import DualCache
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.proxy._types import CallTypes, UserAPIKeyAuth
from litellm.proxy.utils import ProxyLogging
from litellm.types.guardrails import GuardrailEventHooks
class FreshDictGuardrail(CustomGuardrail):
def __init__(self):
super().__init__(
guardrail_name="counting-guardrail",
event_hook=GuardrailEventHooks.pre_call,
default_on=True,
)
self.pre_call_count = 0
async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type):
self.pre_call_count += 1
return {"model": data["model"], "messages": data["messages"]}
guardrail = FreshDictGuardrail()
with patch("litellm.callbacks", [guardrail]):
ProxyLogging._callback_capabilities_cache.clear()
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
user_api_key_dict = UserAPIKeyAuth(api_key="test-key")
data = {
"model": "gpt-4",
"messages": [{"role": "user", "content": "hello"}],
"metadata": {},
}
data = await proxy_logging.pre_call_hook(
user_api_key_dict=user_api_key_dict,
data=data,
call_type="acompletion",
)
data["guardrails"] = ["counting-guardrail"]
await guardrail.async_pre_call_deployment_hook(data, CallTypes.acompletion)
assert guardrail.pre_call_count == 1
@pytest.mark.asyncio
async def test_deployment_hook_runs_pre_call_without_proxy_loop():
"""
Direct-SDK usage (litellm.acompletion(..., guardrails=[...]) without the
proxy) never runs the proxy pre-call loop, so the deployment hook is the
only place the guardrail executes and it must still run exactly once.
"""
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.proxy._types import CallTypes
from litellm.types.guardrails import GuardrailEventHooks
class CountingGuardrail(CustomGuardrail):
def __init__(self):
super().__init__(
guardrail_name="counting-guardrail",
event_hook=GuardrailEventHooks.pre_call,
default_on=True,
)
self.pre_call_count = 0
async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type):
self.pre_call_count += 1
return data
guardrail = CountingGuardrail()
data = {
"model": "gpt-4",
"messages": [{"role": "user", "content": "hello"}],
"guardrails": ["counting-guardrail"],
"metadata": {},
}
await guardrail.async_pre_call_deployment_hook(data, CallTypes.acompletion)
assert guardrail.pre_call_count == 1
# ---------------------------------------------------------------------------
# Integration test: post_call_success_hook with model-level guardrails
# ---------------------------------------------------------------------------

114
tests/test_litellm/proxy/test_proxy_utils.py
View file

@ -321,3 +321,117 @@ class TestPostCallFailureHookLiftsFirstApiCallStartTime:
await self._run(request_data)
assert "first_api_call_start_time" not in request_data
assert "litellm_logging_obj" not in request_data
class TestPostCallFailureHookLLMExceptionAlerting:
"""The llm_exceptions alert is for infra / LLM-API failures, not user
errors (https://github.com/BerriAI/litellm/issues/3395). Already-normalized
client errors must be excluded so a guardrail content-policy block never
pages on-call. ProxyException is such an error; before LIT-3751 only
HTTPException was excluded, so AIM blocks paged as if the LLM API failed."""
async def _alerted(self, exc) -> bool:
import asyncio
from unittest.mock import AsyncMock, patch
from litellm.proxy._types import AlertType, UserAPIKeyAuth
proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache())
proxy_logging_obj.alert_types = [AlertType.llm_exceptions]
alerting_handler = AsyncMock()
with (
patch.object(proxy_logging_obj, "update_request_status", new=AsyncMock()),
patch.object(proxy_logging_obj, "alerting_handler", new=alerting_handler),
):
await proxy_logging_obj.post_call_failure_hook(
request_data={},
original_exception=exc,
user_api_key_dict=UserAPIKeyAuth(),
)
await asyncio.sleep(0) # let the fire-and-forget alert task run
return alerting_handler.called
@pytest.mark.asyncio
async def test_proxy_exception_does_not_alert(self):
from litellm.proxy._types import ProxyException
exc = ProxyException(
message="content blocked",
type="invalid_request_error",
param=None,
code=400,
openai_code="content_policy_violation",
)
assert await self._alerted(exc) is False
@pytest.mark.asyncio
async def test_http_exception_does_not_alert(self):
assert (
await self._alerted(HTTPException(status_code=400, detail="blocked"))
is False
)
@pytest.mark.asyncio
async def test_genuine_llm_api_error_still_alerts(self):
assert await self._alerted(Exception("upstream 503")) is True
class TestPostCallFailureHookProxyExceptionLogging:
"""A guardrail block raises a ProxyException; on an LLM route it must still
drive proxy-only failure logging (_handle_logging_proxy_only_error) so the
blocked request is recorded, exactly as the old HTTPException did. Before
LIT-3751 the classifier only matched HTTPException, so switching AIM to
ProxyException silently dropped the rejected prompt from failure logs."""
async def _logged(self, exc, *, request_route) -> bool:
from unittest.mock import AsyncMock, patch
from litellm.proxy._types import UserAPIKeyAuth
proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache())
proxy_logging_obj.alert_types = []
handle_mock = AsyncMock()
with (
patch.object(proxy_logging_obj, "update_request_status", new=AsyncMock()),
patch.object(
proxy_logging_obj,
"_handle_logging_proxy_only_error",
new=handle_mock,
),
):
await proxy_logging_obj.post_call_failure_hook(
request_data={},
original_exception=exc,
user_api_key_dict=UserAPIKeyAuth(
api_key="sk-test", request_route=request_route
),
)
return handle_mock.await_count > 0
def _block(self):
from litellm.proxy._types import ProxyException
return ProxyException(
message="content blocked",
type="invalid_request_error",
param=None,
code=400,
openai_code="content_policy_violation",
)
@pytest.mark.asyncio
async def test_proxy_exception_on_llm_route_is_logged(self):
assert (
await self._logged(self._block(), request_route="/v1/chat/completions")
is True
)
@pytest.mark.asyncio
async def test_generic_exception_on_llm_route_is_not_logged(self):
# A raw provider/unknown exception is logged by the LLM call path, not here.
assert (
await self._logged(
Exception("upstream 503"), request_route="/v1/chat/completions"
)
is False
)

4
uv.lock generated
View file

@ -9,7 +9,7 @@ resolution-markers = [
]
[options]
exclude-newer = "2026-06-10T23:28:01.952719Z"
exclude-newer = "2026-06-17T19:22:29.366621Z"
exclude-newer-span = "P3D"
[manifest]
@ -3273,7 +3273,7 @@ wheels = [
[[package]]
name = "litellm"
version = "1.87.3"
version = "1.87.4"
source = { editable = "." }
dependencies = [
{ name = "aiohttp" },