mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-01 07:03:02 +00:00
Merge a247d9e553 into afec145edc
This commit is contained in:
commit
857fdb079a
19 changed files with 1263 additions and 64 deletions
|
|
@ -170,6 +170,7 @@ HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS: Final = get_env_int_in_range(
|
|||
# Streams exceeding this duration are terminated with a Timeout error.
|
||||
# None (default) = no limit. Set env var to a number of seconds to enable globally.
|
||||
LITELLM_MAX_STREAMING_DURATION_SECONDS: Final = get_env_float_or_none("LITELLM_MAX_STREAMING_DURATION_SECONDS")
|
||||
FAILURE_LOG_MASK_TIMEOUT_SECONDS: Final = 1.0
|
||||
MID_STREAM_FAILURE_LOG_WAIT_SECONDS: Final = 1.0
|
||||
|
||||
# Maximum number of base64 characters to keep in logging payloads.
|
||||
|
|
|
|||
|
|
@ -298,6 +298,7 @@ class CustomGuardrail(CustomLogger):
|
|||
use_native_lifecycle_hooks: ClassVar[bool] = False
|
||||
|
||||
records_own_guardrail_information: ClassVar[bool] = False
|
||||
masks_logged_request: ClassVar[bool] = False
|
||||
logging_only_scope: LoggingOnlyScope | None
|
||||
logging_only_continue_on_input_failure: bool
|
||||
|
||||
|
|
|
|||
|
|
@ -504,7 +504,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
|
|||
"""For masking logged request/response. Return a modified version of the request/result."""
|
||||
return kwargs, result
|
||||
|
||||
def logging_hook(self, kwargs: dict, result: object, call_type: str) -> tuple[dict, object]:
|
||||
def logging_hook(self, kwargs: dict, result: object, call_type: str) -> tuple[dict[str, object], object]:
|
||||
"""For masking logged request/response. Return a modified version of the request/result."""
|
||||
return kwargs, result
|
||||
|
||||
|
|
|
|||
|
|
@ -2,9 +2,11 @@
|
|||
## Common Utility file for Logging handler
|
||||
# Logging function -> log the exact model details + what's being sent | Non-Blocking
|
||||
import asyncio
|
||||
import concurrent.futures
|
||||
import copy
|
||||
import datetime
|
||||
import functools
|
||||
import itertools
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
|
|
@ -43,6 +45,7 @@ from litellm.constants import (
|
|||
DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT,
|
||||
DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT,
|
||||
EMPTY_MAPPING,
|
||||
FAILURE_LOG_MASK_TIMEOUT_SECONDS,
|
||||
INTERNAL_CALLBACK_NAMES,
|
||||
KUBERNETES_POD_ROUTING_KEY,
|
||||
LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL,
|
||||
|
|
@ -70,6 +73,7 @@ from litellm.integrations.deepeval.deepeval import DeepEvalLogger
|
|||
from litellm.integrations.langtrace import langtrace_trace_endpoint
|
||||
from litellm.integrations.mlflow import MlflowLogger
|
||||
from litellm.integrations.sqs import SQSLogger
|
||||
from litellm.litellm_core_utils.asyncify import is_event_loop_running
|
||||
from litellm.litellm_core_utils.classifier_logging import (
|
||||
classifier_audit_fields,
|
||||
classifier_input_snapshot,
|
||||
|
|
@ -129,6 +133,7 @@ from litellm.litellm_core_utils.redact_messages import (
|
|||
redact_message_input_output_from_custom_logger,
|
||||
redact_message_input_output_from_logging,
|
||||
redacted_model_call_details_for_custom_logger,
|
||||
redacted_request_entries,
|
||||
redacted_response_obj_for_custom_logger,
|
||||
should_redact_message_logging,
|
||||
)
|
||||
|
|
@ -136,7 +141,7 @@ from litellm.litellm_core_utils.served_output_texts import (
|
|||
SERVED_OUTPUT_TEXTS_KEY,
|
||||
overlay_served_output_texts,
|
||||
)
|
||||
from litellm.litellm_core_utils.thread_pool_executor import executor
|
||||
from litellm.litellm_core_utils.thread_pool_executor import executor, failure_log_mask_executor
|
||||
from litellm.llms.base_llm.ocr.transformation import OCRResponse
|
||||
from litellm.llms.base_llm.search.transformation import SearchResponse
|
||||
from litellm.responses.utils import ResponseAPILoggingUtils
|
||||
|
|
@ -144,6 +149,7 @@ from litellm.types.agents import LiteLLMSendMessageResponse
|
|||
from litellm.types.containers.main import ContainerObject
|
||||
from litellm.types.decisions import DecisionsResponse, OpenAIDecisionResponse
|
||||
from litellm.types.extraction import ExtractionResponse
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
from litellm.types.integrations.s3_v2 import S3PartitionGranularity
|
||||
from litellm.types.interactions import (
|
||||
InteractionsAPIResponse,
|
||||
|
|
@ -572,6 +578,62 @@ _PER_CALL_LITELLM_PARAM_DICTS: Final = frozenset({"metadata", "litellm_metadata"
|
|||
_LITELLM_PARAMS: Final = TypeAdapter(dict[str, object])
|
||||
_SERVED_ANTHROPIC_BODY: Final = TypeAdapter(dict[str, object])
|
||||
_REDACTED_RESULT: Final = TypeAdapter(object)
|
||||
_ORPHANED_FAILURE_LOG_TASKS: Final[set["asyncio.Future[None]"]] = set() # mutable-ok: keeps orphaned tasks alive
|
||||
|
||||
|
||||
_ABSENT: Final = object()
|
||||
_LOGGED_PROMPT_PATHS: Final = (
|
||||
("messages",),
|
||||
("input",),
|
||||
("optional_params", "messages"),
|
||||
("litellm_params", "proxy_server_request", "body", "messages"),
|
||||
)
|
||||
|
||||
|
||||
def _failure_log_hook_kwargs(details: Mapping[str, object]) -> dict[str, object]: # mutable-ok: hooks mask in place
|
||||
return {**details, "messages": copy.deepcopy(details.get("messages"))}
|
||||
|
||||
|
||||
def _logged_field(container: object, key: str) -> object:
|
||||
if not isinstance(container, Mapping):
|
||||
return _ABSENT
|
||||
return cast(Mapping[str, object], container).get(key, _ABSENT) # cast-ok: logged request dicts are str-keyed JSON
|
||||
|
||||
|
||||
def _with_field(inner: object, step: tuple[object, str]) -> object:
|
||||
container, key = step
|
||||
return {**cast(Mapping[str, object], container), key: inner} # cast-ok: containers above a found copy are Mappings
|
||||
|
||||
|
||||
def _with_prompt_copy_replaced(
|
||||
details: Mapping[str, object], path: Sequence[str], raw: object, replacement: object
|
||||
) -> Mapping[str, object]:
|
||||
containers: Final = tuple(itertools.accumulate(path, _logged_field, initial=details))
|
||||
if containers[-1] != raw:
|
||||
return details
|
||||
rebuilt: Final = functools.reduce(_with_field, reversed(tuple(zip(containers, path))), replacement)
|
||||
return cast(Mapping[str, object], rebuilt) # cast-ok: the outermost rebuilt container is details
|
||||
|
||||
|
||||
def _with_prompt_copies_replaced(
|
||||
details: Mapping[str, object], raw: object, replacement: object
|
||||
) -> dict[str, object]: # mutable-ok: becomes the Logging object's model_call_details
|
||||
if raw is None:
|
||||
return dict(details)
|
||||
return dict(
|
||||
functools.reduce(
|
||||
lambda fields, path: _with_prompt_copy_replaced(fields, path, raw, replacement),
|
||||
_LOGGED_PROMPT_PATHS,
|
||||
details,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _keep_alive_until_done(task: "asyncio.Future[None]") -> None:
|
||||
_ORPHANED_FAILURE_LOG_TASKS.add(task)
|
||||
task.add_done_callback(_ORPHANED_FAILURE_LOG_TASKS.discard)
|
||||
|
||||
|
||||
_GUARDRAIL_INFORMATION_KEY: Final = "standard_logging_guardrail_information"
|
||||
|
||||
|
||||
|
|
@ -626,6 +688,14 @@ def _timestamp_seconds(moment: object) -> float | None:
|
|||
return None
|
||||
|
||||
|
||||
def _masks_synchronously(guardrail: CustomGuardrail) -> bool:
|
||||
return next(cls for cls in type(guardrail).__mro__ if "logging_hook" in vars(cls)) is not CustomLogger
|
||||
|
||||
|
||||
def _runs_on_logging_only(guardrail: CustomGuardrail, data: Mapping[str, object]) -> bool:
|
||||
return guardrail.should_run_guardrail(data=data, event_type=GuardrailEventHooks.logging_only) is True
|
||||
|
||||
|
||||
def same_attempt_marker(first: object, second: object) -> bool:
|
||||
return first is second or first == second
|
||||
|
||||
|
|
@ -654,6 +724,9 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
baseline_cache_context: "BaselineCacheContext | None" = None
|
||||
baseline_observation: "CapturedBaselineObservation | None" = None
|
||||
_on_failure_payload_built: "Callable[[], None] | None" = None
|
||||
_failure_log_masked: bool = False
|
||||
_unmasked_messages: object = None
|
||||
_failure_logged_messages: object = None
|
||||
|
||||
@property
|
||||
def _defer_async_logging(self) -> bool:
|
||||
|
|
@ -2702,7 +2775,9 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
sync_sdk: Final = self.is_sync_litellm_request(litellm_params)
|
||||
passthrough: Final = self.call_type == CallTypes.pass_through.value
|
||||
if sync_sdk and not prefer_async_handlers and not passthrough:
|
||||
self.failure_handler(exception, traceback_exception)
|
||||
run_failure_handler: Final = functools.partial(self.failure_handler, exception, traceback_exception)
|
||||
await self.async_mask_failure_log_request(on_cancelled=run_failure_handler)
|
||||
run_failure_handler()
|
||||
return
|
||||
|
||||
try:
|
||||
|
|
@ -3502,28 +3577,20 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
hook_callbacks: Final = tuple(dict.fromkeys((*callbacks, *delivery_callbacks)))
|
||||
for callback in hook_callbacks:
|
||||
if isinstance(callback, CustomGuardrail):
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
|
||||
if (
|
||||
callback.should_run_guardrail(
|
||||
data=call_details,
|
||||
event_type=GuardrailEventHooks.logging_only,
|
||||
)
|
||||
is not True
|
||||
):
|
||||
if not _runs_on_logging_only(callback, call_details):
|
||||
continue
|
||||
|
||||
self.shadow_eval_request_snapshot = None
|
||||
call_details, result = callback.logging_hook( # pyright: ignore[reportUnknownVariableType] # logging_hook is typed with a bare dict
|
||||
kwargs=call_details,
|
||||
result=result,
|
||||
call_type=self.call_type,
|
||||
call_type=self._logging_hook_call_type,
|
||||
)
|
||||
elif isinstance(callback, CustomLogger):
|
||||
call_details, result = callback.logging_hook( # pyright: ignore[reportUnknownVariableType] # logging_hook is typed with a bare dict
|
||||
kwargs=call_details,
|
||||
result=result,
|
||||
call_type=self.call_type,
|
||||
call_type=self._logging_hook_call_type,
|
||||
)
|
||||
if not for_async_request:
|
||||
self.model_call_details = call_details
|
||||
|
|
@ -4094,22 +4161,14 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
for callback in hook_callbacks:
|
||||
try:
|
||||
if isinstance(callback, CustomGuardrail):
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
|
||||
if (
|
||||
callback.should_run_guardrail(
|
||||
data=self.model_call_details,
|
||||
event_type=GuardrailEventHooks.logging_only,
|
||||
)
|
||||
is not True
|
||||
):
|
||||
if not _runs_on_logging_only(callback, self.model_call_details):
|
||||
continue
|
||||
|
||||
self.shadow_eval_request_snapshot = None
|
||||
self.model_call_details, result = await callback.async_logging_hook(
|
||||
kwargs=self.model_call_details,
|
||||
result=result,
|
||||
call_type=self.call_type,
|
||||
call_type=self._logging_hook_call_type,
|
||||
)
|
||||
elif isinstance(callback, CustomLogger):
|
||||
result = redact_message_input_output_from_custom_logger(
|
||||
|
|
@ -4118,7 +4177,7 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
self.model_call_details, result = await callback.async_logging_hook(
|
||||
kwargs=self.model_call_details,
|
||||
result=result,
|
||||
call_type=self.call_type,
|
||||
call_type=self._logging_hook_call_type,
|
||||
)
|
||||
except Exception: # noqa: BLE001 # one failing hook must not skip later callbacks (slot release)
|
||||
verbose_logger.error(
|
||||
|
|
@ -4295,6 +4354,180 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
except Exception as e:
|
||||
verbose_logger.debug("Error in _handle_callback_failure: %s", e)
|
||||
|
||||
@property
|
||||
def _logging_hook_call_type(self) -> str:
|
||||
return cast(str, self.call_type) # cast-ok: Logging.__init__ takes call_type untyped and callers pass names
|
||||
|
||||
def _async_failure_callbacks(self) -> Sequence[str | Callable[..., object] | CustomLogger]:
|
||||
return self.get_combined_callback_list(
|
||||
dynamic_success_callbacks=self.dynamic_async_failure_callbacks,
|
||||
global_callbacks=litellm._async_failure_callback,
|
||||
)
|
||||
|
||||
def _failure_log_maskers(self) -> tuple[CustomGuardrail, ...]:
|
||||
callbacks: Final = (
|
||||
*self.get_combined_callback_list(self.dynamic_failure_callbacks, litellm.failure_callback),
|
||||
*self._async_failure_callbacks(),
|
||||
)
|
||||
return tuple(
|
||||
callback
|
||||
for callback in dict.fromkeys(callbacks)
|
||||
if isinstance(callback, CustomGuardrail) and callback.masks_logged_request
|
||||
)
|
||||
|
||||
@property
|
||||
def messages_sent_to_model(self) -> object:
|
||||
return self._unmasked_messages if self._failure_log_masked else self.model_call_details.get("messages")
|
||||
|
||||
def failure_logged_proxy_server_request(self, proxy_server_request: Mapping[str, object]) -> Mapping[str, object]:
|
||||
if not self._failure_log_masked:
|
||||
return proxy_server_request
|
||||
return _with_prompt_copy_replaced(
|
||||
proxy_server_request, ("body", "messages"), self._unmasked_messages, self._failure_logged_messages
|
||||
)
|
||||
|
||||
def _claim_failure_log_mask(self) -> bool:
|
||||
if self._failure_log_masked:
|
||||
return False
|
||||
self._failure_log_masked = True
|
||||
self._unmasked_messages = self.model_call_details.get("messages")
|
||||
self._failure_logged_messages = self._unmasked_messages
|
||||
return True
|
||||
|
||||
def _log_failure_with_messages(self, details: Mapping[str, object], logged_messages: object) -> None:
|
||||
self.model_call_details = _with_prompt_copies_replaced(details, self._unmasked_messages, logged_messages)
|
||||
self._failure_logged_messages = logged_messages
|
||||
|
||||
def _redact_failure_log_request(self) -> None:
|
||||
redacted: Final = {**self.model_call_details, **redacted_request_entries(self.model_call_details)}
|
||||
self._log_failure_with_messages(redacted, redacted["messages"])
|
||||
|
||||
@property
|
||||
def masks_failure_log(self) -> bool:
|
||||
return any(
|
||||
_runs_on_logging_only(guardrail, self.model_call_details) for guardrail in self._failure_log_maskers()
|
||||
)
|
||||
|
||||
def redact_failure_log_request(self) -> None:
|
||||
if self.masks_failure_log and self._claim_failure_log_mask():
|
||||
self._redact_failure_log_request()
|
||||
|
||||
def _report_failure_log_mask_error(self, guardrail: CustomGuardrail) -> None:
|
||||
verbose_logger.error(
|
||||
"LiteLLM.LoggingError: [Non-Blocking] Exception occurred masking the failure log with %s, so its request "
|
||||
"is redacted %s",
|
||||
guardrail.guardrail_name,
|
||||
traceback.format_exc(),
|
||||
)
|
||||
self._handle_callback_failure(callback=guardrail)
|
||||
|
||||
def _redact_failure_log_request_after_timeout(self) -> None:
|
||||
verbose_logger.error(
|
||||
"LiteLLM.LoggingError: [Non-Blocking] Masking the failure log took over %ss, so its request is redacted",
|
||||
FAILURE_LOG_MASK_TIMEOUT_SECONDS,
|
||||
)
|
||||
self._redact_failure_log_request()
|
||||
|
||||
async def _mask_failure_log_request_with(self, maskers: tuple[CustomGuardrail, ...]) -> None:
|
||||
try:
|
||||
await asyncio.wait_for(self._run_failure_log_maskers(maskers), timeout=FAILURE_LOG_MASK_TIMEOUT_SECONDS)
|
||||
except asyncio.TimeoutError:
|
||||
self._redact_failure_log_request_after_timeout()
|
||||
|
||||
async def _run_failure_log_maskers(self, maskers: tuple[CustomGuardrail, ...]) -> None:
|
||||
masked: Final = await self._masked_failure_log_details(self.model_call_details, maskers)
|
||||
if masked is None:
|
||||
self._redact_failure_log_request()
|
||||
return
|
||||
self._log_failure_with_messages(self.model_call_details, masked.get("messages"))
|
||||
|
||||
async def _masked_failure_log_details(
|
||||
self, details: Mapping[str, object] | None, maskers: tuple[CustomGuardrail, ...]
|
||||
) -> Mapping[str, object] | None:
|
||||
if details is None or not maskers:
|
||||
return details
|
||||
return await self._masked_failure_log_details(
|
||||
await self._async_failure_log_mask_step(details, maskers[0]), maskers[1:]
|
||||
)
|
||||
|
||||
async def _async_failure_log_mask_step(
|
||||
self, details: Mapping[str, object], guardrail: CustomGuardrail
|
||||
) -> Mapping[str, object] | None:
|
||||
if not _runs_on_logging_only(guardrail, details):
|
||||
return details
|
||||
try:
|
||||
hooked_details: Final = _METADATA_MAPPING.validate_python(
|
||||
(
|
||||
await guardrail.async_logging_hook(
|
||||
kwargs=_failure_log_hook_kwargs(details), result=None, call_type=self._logging_hook_call_type
|
||||
)
|
||||
)[0]
|
||||
)
|
||||
except Exception: # noqa: BLE001 # a failing mask redacts the log instead of stopping failure logging
|
||||
self._report_failure_log_mask_error(guardrail)
|
||||
return None
|
||||
return hooked_details
|
||||
|
||||
async def async_mask_failure_log_request(
|
||||
self,
|
||||
on_cancelled: Callable[[], None] | None = None,
|
||||
on_cancelled_async: Callable[[], Awaitable[None]] | None = None,
|
||||
) -> None:
|
||||
if not self._claim_failure_log_mask():
|
||||
return
|
||||
maskers: Final = self._failure_log_maskers()
|
||||
if not maskers:
|
||||
return
|
||||
mask: Final = asyncio.create_task(self._mask_failure_log_request_with(maskers))
|
||||
try:
|
||||
await asyncio.shield(mask)
|
||||
except asyncio.CancelledError:
|
||||
_keep_alive_until_done(mask)
|
||||
if on_cancelled is not None:
|
||||
mask.add_done_callback(lambda _: on_cancelled())
|
||||
if on_cancelled_async is not None:
|
||||
mask.add_done_callback(lambda _: _keep_alive_until_done(asyncio.ensure_future(on_cancelled_async())))
|
||||
raise
|
||||
|
||||
def mask_failure_log_request(self, mask_executor: concurrent.futures.Executor = failure_log_mask_executor) -> None:
|
||||
if is_event_loop_running() or not self._claim_failure_log_mask():
|
||||
return
|
||||
maskers: Final = self._failure_log_maskers()
|
||||
if not maskers:
|
||||
return
|
||||
masking: Final = mask_executor.submit(
|
||||
functools.reduce, self._sync_failure_log_mask_step, maskers, self.model_call_details
|
||||
)
|
||||
try:
|
||||
masked: Final = masking.result(timeout=FAILURE_LOG_MASK_TIMEOUT_SECONDS)
|
||||
except concurrent.futures.TimeoutError:
|
||||
_ = masking.cancel()
|
||||
self._redact_failure_log_request_after_timeout()
|
||||
return
|
||||
except concurrent.futures.CancelledError:
|
||||
self._redact_failure_log_request()
|
||||
return
|
||||
if masked is None:
|
||||
self._redact_failure_log_request()
|
||||
return
|
||||
self._log_failure_with_messages(self.model_call_details, masked.get("messages"))
|
||||
|
||||
def _sync_failure_log_mask_step(
|
||||
self, details: Mapping[str, object] | None, guardrail: CustomGuardrail
|
||||
) -> Mapping[str, object] | None:
|
||||
if details is None or not _runs_on_logging_only(guardrail, details):
|
||||
return details
|
||||
if not _masks_synchronously(guardrail):
|
||||
return None
|
||||
try:
|
||||
hooked: Final = guardrail.logging_hook(
|
||||
kwargs=_failure_log_hook_kwargs(details), result=None, call_type=self._logging_hook_call_type
|
||||
)
|
||||
except Exception: # noqa: BLE001 # a failing mask redacts the log instead of stopping failure logging
|
||||
self._report_failure_log_mask_error(guardrail)
|
||||
return None
|
||||
return hooked[0]
|
||||
|
||||
def _failure_handler_helper_fn(self, exception, traceback_exception, start_time=None, end_time=None):
|
||||
if start_time is None:
|
||||
start_time = self.start_time
|
||||
|
|
@ -4420,6 +4653,7 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
is_sync_request: Final = self.is_sync_litellm_request(litellm_params)
|
||||
|
||||
try:
|
||||
self.mask_failure_log_request()
|
||||
start_time, end_time = self._failure_handler_helper_fn(
|
||||
exception=exception,
|
||||
traceback_exception=traceback_exception,
|
||||
|
|
@ -4619,6 +4853,8 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
await self.special_failure_handlers(exception=exception)
|
||||
if not self.should_run_logging(event_type="async_failure"): # prevent double logging
|
||||
return
|
||||
self.mark_logging_complete(event_type="async_failure")
|
||||
await self.async_mask_failure_log_request()
|
||||
start_time, end_time = self._failure_handler_helper_fn(
|
||||
exception=exception,
|
||||
traceback_exception=traceback_exception,
|
||||
|
|
@ -4627,14 +4863,10 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
)
|
||||
if self._on_failure_payload_built is not None:
|
||||
self._on_failure_payload_built()
|
||||
callbacks: Final = self.get_combined_callback_list(
|
||||
dynamic_success_callbacks=self.dynamic_async_failure_callbacks,
|
||||
global_callbacks=litellm._async_failure_callback,
|
||||
)
|
||||
callbacks: Final = self._async_failure_callbacks()
|
||||
|
||||
result: Final = None # result sent to all loggers, init this to None incase it's not created
|
||||
|
||||
self.mark_logging_complete(event_type="async_failure")
|
||||
for callback in callbacks:
|
||||
try:
|
||||
litellm_params = self.model_call_details.get("litellm_params", {})
|
||||
|
|
|
|||
|
|
@ -67,11 +67,11 @@ def redacted_model_call_details_for_custom_logger(model_call_details: dict, cust
|
|||
return model_call_details
|
||||
return {
|
||||
**{key: value for key, value in model_call_details.items() if key != SERVED_OUTPUT_TEXTS_KEY},
|
||||
**_redacted_entries_for_opted_out_logger(model_call_details),
|
||||
**redacted_request_entries(model_call_details),
|
||||
}
|
||||
|
||||
|
||||
def _redacted_entries_for_opted_out_logger(model_call_details: Mapping[str, object]) -> Mapping[str, object]:
|
||||
def redacted_request_entries(model_call_details: Mapping[str, object]) -> Mapping[str, object]:
|
||||
optional_params: Final = _as_mapping(model_call_details.get("optional_params"))
|
||||
litellm_params: Final = _as_mapping(model_call_details.get("litellm_params"))
|
||||
standard_logging_object: Final = _as_mapping(model_call_details.get("standard_logging_object"))
|
||||
|
|
|
|||
|
|
@ -23,6 +23,7 @@ from typing_extensions import NotRequired, TypedDict
|
|||
import litellm
|
||||
from litellm import verbose_logger
|
||||
from litellm._uuid import uuid
|
||||
from litellm.constants import FAILURE_LOG_MASK_TIMEOUT_SECONDS
|
||||
from litellm.litellm_core_utils.asyncify import asyncify
|
||||
from litellm.litellm_core_utils.core_helpers import (
|
||||
UPSTREAM_USAGE_COST,
|
||||
|
|
@ -2509,7 +2510,8 @@ class CustomStreamWrapper:
|
|||
"""Let this stream's failure log read its own deployment before a fallback reuses the logging object."""
|
||||
if self._failure_logging is None:
|
||||
return
|
||||
await asyncio.wait((self._failure_logging,), timeout=timeout)
|
||||
mask_wait: Final = FAILURE_LOG_MASK_TIMEOUT_SECONDS if self.logging_obj.masks_failure_log else 0.0
|
||||
await asyncio.wait((self._failure_logging,), timeout=None if timeout is None else timeout + mask_wait)
|
||||
|
||||
def bill_failure_after_earlier_rounds(self, earlier: "BilledRoundsSpend") -> None:
|
||||
self._earlier_rounds_spend = earlier
|
||||
|
|
|
|||
|
|
@ -80,3 +80,7 @@ executor: Final = BoundedLoggingThreadPoolExecutor(
|
|||
max_workers=MAX_THREADS,
|
||||
max_pending_tasks=LOGGING_EXECUTOR_MAX_PENDING_TASKS,
|
||||
)
|
||||
failure_log_mask_executor: Final = BoundedLoggingThreadPoolExecutor(
|
||||
max_workers=MAX_THREADS,
|
||||
max_pending_tasks=LOGGING_EXECUTOR_MAX_PENDING_TASKS,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -7,7 +7,7 @@
|
|||
|
||||
import json
|
||||
import os
|
||||
from typing import TYPE_CHECKING, Final, Literal, TypedDict
|
||||
from typing import TYPE_CHECKING, ClassVar, Final, Literal, TypedDict
|
||||
|
||||
from fastapi import HTTPException
|
||||
|
||||
|
|
@ -48,6 +48,8 @@ class GuardrailsAIResponsePreCall(TypedDict):
|
|||
|
||||
|
||||
class GuardrailsAI(CustomGuardrail):
|
||||
masks_logged_request: ClassVar[bool] = True
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
guard_name: str,
|
||||
|
|
|
|||
|
|
@ -17,7 +17,7 @@ from contextlib import asynccontextmanager
|
|||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from itertools import accumulate, count, pairwise
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, cast
|
||||
from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, Protocol, cast
|
||||
|
||||
import aiohttp
|
||||
from pydantic import ConfigDict, JsonValue, TypeAdapter, with_config
|
||||
|
|
@ -238,6 +238,7 @@ def _resolve_overlapping_spans(analyze_results: Sequence[PresidioAnalyzeResponse
|
|||
|
||||
|
||||
class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
|
||||
masks_logged_request: ClassVar[bool] = True
|
||||
user_api_key_cache = None
|
||||
ad_hoc_recognizers: list[str] | None = None
|
||||
|
||||
|
|
@ -389,7 +390,8 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail):
|
|||
async with self._session_lock:
|
||||
if self._http_session is None or self._http_session.closed:
|
||||
self._http_session = aiohttp.ClientSession()
|
||||
yield self._http_session
|
||||
session: Final = self._http_session
|
||||
yield session
|
||||
else:
|
||||
# Background thread/loop -> use loop-bound session cache
|
||||
# This avoids "attached to a different loop" or "no running event loop" errors
|
||||
|
|
|
|||
|
|
@ -1708,12 +1708,24 @@ def _failure_fields_to_lift(
|
|||
request_data (it is not serialisable), so the caller merges these fields
|
||||
onto request_data first: the request start and first-handoff instants for
|
||||
duration and preprocessing latency, the call type, recovered or estimated
|
||||
usage for token counts, and the standard logging object for deployment
|
||||
attribution on failed-request spend logs."""
|
||||
usage for token counts, the standard logging object for deployment
|
||||
attribution on failed-request spend logs, and the request body with the
|
||||
prompt the failure log masked."""
|
||||
_logging_obj: Final = request_data.get("litellm_logging_obj")
|
||||
if _logging_obj is None:
|
||||
return _EMPTY_LIFT
|
||||
_model_call_details: Final = getattr(_logging_obj, "model_call_details", {})
|
||||
_proxy_server_request: Final = request_data.get("proxy_server_request")
|
||||
_logged_proxy_server_request: Final = (
|
||||
_logging_obj.failure_logged_proxy_server_request(_as_str_mapping(_proxy_server_request))
|
||||
if isinstance(_logging_obj, Logging) and isinstance(_proxy_server_request, Mapping)
|
||||
else None
|
||||
)
|
||||
_sent_model_call_details: Final = (
|
||||
{**_model_call_details, "messages": _logging_obj.messages_sent_to_model}
|
||||
if isinstance(_logging_obj, Logging)
|
||||
else _model_call_details
|
||||
)
|
||||
_first_handoff: Final = _model_call_details.get("first_api_call_start_time")
|
||||
_litellm_params: Final = getattr(_logging_obj, "litellm_params", None)
|
||||
_base_model: Final = get_base_model_from_metadata(model_call_details=_model_call_details)
|
||||
|
|
@ -1721,7 +1733,7 @@ def _failure_fields_to_lift(
|
|||
_success_already_logged: Final = isinstance(_logging_obj, Logging) and _logging_obj.async_success_log_enqueued
|
||||
_billed: Final = provider_billed and not _success_already_logged
|
||||
_usage_to_lift: Final = _failure_usage_to_lift(
|
||||
model_call_details=_model_call_details,
|
||||
model_call_details=_sent_model_call_details,
|
||||
request_body=request_data,
|
||||
dispatched=_model_call_details.get("current_attempt_api_call_start_time", _first_handoff) is not None,
|
||||
provider_billed=_billed,
|
||||
|
|
@ -1730,7 +1742,7 @@ def _failure_fields_to_lift(
|
|||
dispatched_call_type=dispatched_call_type,
|
||||
)
|
||||
_prompt_estimate: Final = (
|
||||
_estimated_prompt_usage(_model_call_details, request_data, dispatched_call_type)
|
||||
_estimated_prompt_usage(_sent_model_call_details, request_data, dispatched_call_type)
|
||||
if _billed and _first_handoff is not None and not _model_call_details.get("failed_attempts_usage")
|
||||
else None
|
||||
)
|
||||
|
|
@ -1747,6 +1759,7 @@ def _failure_fields_to_lift(
|
|||
("response_cost", None if _usage_to_lift is None else (_usage_to_lift[1] or 0.0)),
|
||||
("standard_logging_object", _model_call_details.get("standard_logging_object")),
|
||||
("billed_failed_attempts", _earlier_attempts or None),
|
||||
("proxy_server_request", _logged_proxy_server_request),
|
||||
)
|
||||
return MappingProxyType({key: value for key, value in _entries if value is not None})
|
||||
|
||||
|
|
@ -4100,8 +4113,9 @@ class ProxyLogging:
|
|||
litellm_logging_obj: Logging | None = request_data.get("litellm_logging_obj", None)
|
||||
set_up_by_request: Final = litellm_logging_obj is not None
|
||||
request_fields: Final = cast(Mapping[str, object], request_data) # cast-ok: request bodies are str-keyed JSON
|
||||
authenticated: Final = error_type != ProxyErrorTypes.auth_error
|
||||
logged_metadata_keys: Final = _proxy_only_error_logged_metadata_keys(
|
||||
authenticated=error_type != ProxyErrorTypes.auth_error, past_pre_call=set_up_by_request
|
||||
authenticated=authenticated, past_pre_call=set_up_by_request
|
||||
)
|
||||
user_api_key_logged_metadata: Final = LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(
|
||||
user_api_key_dict=user_api_key_dict
|
||||
|
|
@ -4210,6 +4224,8 @@ class ProxyLogging:
|
|||
input=input,
|
||||
api_key="",
|
||||
)
|
||||
if not authenticated:
|
||||
litellm_logging_obj.redact_failure_log_request()
|
||||
|
||||
await self._dispatch_proxy_only_failure_handlers(
|
||||
litellm_logging_obj=litellm_logging_obj,
|
||||
|
|
|
|||
|
|
@ -10,6 +10,7 @@ import binascii
|
|||
import contextvars
|
||||
import copy
|
||||
import datetime
|
||||
import functools
|
||||
import hashlib
|
||||
import inspect
|
||||
import io
|
||||
|
|
@ -2535,12 +2536,16 @@ def client(original_function):
|
|||
end_time = _deployment_call_end_time if _deployment_call_end_time is not None else datetime.datetime.now() # noqa: DTZ005 # matches the naive datetimes this whole function already times start_time/end_time with
|
||||
if logging_obj and not _is_litellm_internal_call:
|
||||
unbind_budget_reservation_from_callbacks(logging_obj.litellm_params)
|
||||
try:
|
||||
logging_obj.failure_handler(
|
||||
e, traceback_exception, start_time, end_time, for_async_request=True
|
||||
) # DO NOT MAKE THREADED - router retry fallback relies on this!
|
||||
except Exception as e:
|
||||
raise e
|
||||
run_failure_handler: Final = functools.partial(
|
||||
logging_obj.failure_handler, e, traceback_exception, start_time, end_time, for_async_request=True
|
||||
)
|
||||
await logging_obj.async_mask_failure_log_request(
|
||||
on_cancelled=run_failure_handler,
|
||||
on_cancelled_async=functools.partial(
|
||||
logging_obj.async_failure_handler, e, traceback_exception, start_time, end_time
|
||||
),
|
||||
)
|
||||
run_failure_handler() # DO NOT MAKE THREADED - router retry fallback relies on this!
|
||||
try:
|
||||
await logging_obj.async_failure_handler(e, traceback_exception, start_time, end_time)
|
||||
except Exception as e:
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
import asyncio
|
||||
import contextlib
|
||||
import concurrent.futures
|
||||
import copy
|
||||
import datetime
|
||||
import importlib.abc
|
||||
|
|
@ -10,11 +11,11 @@ import sys
|
|||
import threading
|
||||
import time
|
||||
import traceback
|
||||
from collections.abc import AsyncIterator, Callable, Iterator, Mapping, Sequence
|
||||
from collections.abc import AsyncIterator, Callable, Coroutine, Iterator, Mapping, Sequence
|
||||
from datetime import datetime as datetime_standard_logging, datetime as datetime_unit_test, datetime as dt_object
|
||||
from importlib.machinery import ModuleSpec
|
||||
from types import MappingProxyType, ModuleType
|
||||
from typing import Final, Literal
|
||||
from typing import ClassVar, Final, Literal, ParamSpec, TypeVar, cast
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import httpx
|
||||
|
|
@ -31,6 +32,7 @@ from litellm._logging import session_id_var, trace_id_var, verbose_logger
|
|||
from litellm._service_logger import ServiceLogging
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.constants import (
|
||||
FAILURE_LOG_MASK_TIMEOUT_SECONDS,
|
||||
KUBERNETES_POD_ROUTING_KEY,
|
||||
LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL,
|
||||
LOGGING_WORKER_MAX_TIME_PER_COROUTINE,
|
||||
|
|
@ -38,6 +40,7 @@ from litellm.constants import (
|
|||
SENTRY_PII_DENYLIST,
|
||||
)
|
||||
from litellm.cost_calculator import ocr_batch_cost
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.litellm_core_utils import litellm_logging
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging
|
||||
|
|
@ -52,12 +55,14 @@ from litellm.litellm_core_utils.litellm_logging import (
|
|||
)
|
||||
from litellm.litellm_core_utils.llm_response_utils.response_metadata import update_response_metadata
|
||||
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
|
||||
from litellm.litellm_core_utils.thread_pool_executor import BoundedLoggingThreadPoolExecutor
|
||||
from litellm.llms.base_llm.ocr.transformation import OCRUsageInfo
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.hooks.cache_control_check import PROXY_CacheControlCheck
|
||||
from litellm.proxy.hooks.max_iterations_limiter import PROXY_MaxIterationsHandler
|
||||
from litellm.proxy.hooks.parallel_request_limiter_v3 import PROXY_MaxParallelRequestsHandler_v3
|
||||
from litellm.proxy.utils import InternalUsageCache
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
from litellm.types.llms.openai import ResponseAPIUsage, ResponseCompletedEvent, ResponsesAPIResponse
|
||||
from litellm.types.rerank import RerankResponse
|
||||
from litellm.types.utils import (
|
||||
|
|
@ -8775,6 +8780,531 @@ async def test_streaming_success_callbacks_survive_guardrail_logging_hook_failur
|
|||
releasing.async_log_success_event.assert_awaited_once()
|
||||
|
||||
|
||||
_RAW_PROMPT: Final = "My name is John Smith"
|
||||
_RAW_LOG: Final = [{"role": "user", "content": _RAW_PROMPT}]
|
||||
_MASKED_LOG: Final = [{"role": "user", "content": "My name is <PERSON>"}]
|
||||
_REDACTED_LOG: Final = [{"role": "user", "content": REDACTED_BY_LITELLM}]
|
||||
|
||||
|
||||
class _NameMaskingGuardrail(CustomGuardrail):
|
||||
masks_logged_request: ClassVar[bool] = True
|
||||
|
||||
def __init__(self, event_hook: GuardrailEventHooks = GuardrailEventHooks.logging_only) -> None:
|
||||
super().__init__(guardrail_name="name-masker", event_hook=event_hook)
|
||||
self.scans = 0
|
||||
|
||||
def _mask(self, kwargs: dict[str, object]) -> dict[str, object]:
|
||||
self.scans += 1
|
||||
for message in cast(list[dict[str, str]], kwargs["messages"]):
|
||||
message["content"] = message["content"].replace("John Smith", "<PERSON>")
|
||||
return kwargs
|
||||
|
||||
def logging_hook(
|
||||
self, kwargs: dict[str, object], result: object, call_type: str
|
||||
) -> tuple[dict[str, object], object]:
|
||||
return self._mask(kwargs), result
|
||||
|
||||
async def async_logging_hook(
|
||||
self, kwargs: dict[str, object], result: object, call_type: str
|
||||
) -> tuple[dict[str, object], object]:
|
||||
await asyncio.sleep(0)
|
||||
return self._mask(kwargs), result
|
||||
|
||||
|
||||
class _UnflaggedNameMaskingGuardrail(_NameMaskingGuardrail):
|
||||
masks_logged_request: ClassVar[bool] = False
|
||||
|
||||
|
||||
class _RaisingMaskingGuardrail(_NameMaskingGuardrail):
|
||||
def _mask(self, kwargs: dict[str, object]) -> dict[str, object]:
|
||||
raise RuntimeError("masking failed")
|
||||
|
||||
|
||||
class _BlockingNameMaskingGuardrail(_NameMaskingGuardrail):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.scanning: Final = asyncio.Event()
|
||||
self.release: Final = asyncio.Event()
|
||||
|
||||
async def async_logging_hook(
|
||||
self, kwargs: dict[str, object], result: object, call_type: str
|
||||
) -> tuple[dict[str, object], object]:
|
||||
self.scanning.set()
|
||||
await self.release.wait()
|
||||
return self._mask(kwargs), result
|
||||
|
||||
|
||||
class _HangingSyncMaskingGuardrail(_NameMaskingGuardrail):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.release: Final = threading.Event()
|
||||
|
||||
def logging_hook(
|
||||
self, kwargs: dict[str, object], result: object, call_type: str
|
||||
) -> tuple[dict[str, object], object]:
|
||||
_ = self.release.wait(timeout=30)
|
||||
return self._mask(kwargs), result
|
||||
|
||||
|
||||
class _AsyncOnlyNameMaskingGuardrail(CustomGuardrail):
|
||||
masks_logged_request: ClassVar[bool] = True
|
||||
|
||||
async def async_logging_hook(
|
||||
self, kwargs: dict[str, object], result: object, call_type: str
|
||||
) -> tuple[dict[str, object], object]:
|
||||
return await _NameMaskingGuardrail().async_logging_hook(kwargs, result, call_type)
|
||||
|
||||
|
||||
class _FailureLogRecorder(CustomLogger):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.logged: list[tuple[object, object]] = []
|
||||
|
||||
def _record(self, kwargs: dict[str, object]) -> None:
|
||||
payload: Final = cast(StandardLoggingPayload, kwargs["standard_logging_object"])
|
||||
self.logged.append((copy.deepcopy(kwargs["messages"]), copy.deepcopy(payload["messages"])))
|
||||
|
||||
def log_failure_event(
|
||||
self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object
|
||||
) -> None:
|
||||
self._record(kwargs)
|
||||
|
||||
async def async_log_failure_event(
|
||||
self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object
|
||||
) -> None:
|
||||
self._record(kwargs)
|
||||
|
||||
|
||||
def _register_failure_logging(
|
||||
monkeypatch: pytest.MonkeyPatch, guardrail: CustomGuardrail
|
||||
) -> tuple[_FailureLogRecorder, list[object]]:
|
||||
recorder: Final = _FailureLogRecorder()
|
||||
sync_logged: Final[list[object]] = []
|
||||
|
||||
def record_sync_failure(
|
||||
kwargs: dict[str, object], completion_response: object, start_time: object, end_time: object
|
||||
) -> None:
|
||||
sync_logged.append(copy.deepcopy(kwargs["messages"]))
|
||||
|
||||
for name in ("input_callback", "success_callback", "_async_success_callback", "_async_failure_callback"):
|
||||
monkeypatch.setattr(litellm, name, [])
|
||||
monkeypatch.setattr(litellm, "callbacks", [guardrail, recorder])
|
||||
monkeypatch.setattr(litellm, "failure_callback", [record_sync_failure])
|
||||
return recorder, sync_logged
|
||||
|
||||
|
||||
def _bad_key_error() -> litellm.AuthenticationError:
|
||||
return litellm.AuthenticationError(message="bad key", llm_provider="openai", model="gpt-4o-mini")
|
||||
|
||||
|
||||
async def _until(condition: Callable[[], bool]) -> None:
|
||||
while not condition():
|
||||
await asyncio.sleep(0)
|
||||
|
||||
|
||||
class _SkippableClockLoop(asyncio.SelectorEventLoop):
|
||||
skipped_seconds: float = 0.0
|
||||
|
||||
def time(self) -> float:
|
||||
return super().time() + self.skipped_seconds
|
||||
|
||||
def skip(self, seconds: float) -> None:
|
||||
self.skipped_seconds += seconds
|
||||
|
||||
|
||||
def _run_on_skippable_clock(scenario: Callable[[_SkippableClockLoop], Coroutine[object, object, None]]) -> None:
|
||||
loop: Final = _SkippableClockLoop()
|
||||
try:
|
||||
loop.run_until_complete(scenario(loop))
|
||||
finally:
|
||||
loop.run_until_complete(loop.shutdown_default_executor())
|
||||
loop.close()
|
||||
|
||||
|
||||
_P = ParamSpec("_P")
|
||||
_T = TypeVar("_T")
|
||||
|
||||
|
||||
class _OutlastedTimeBoundFuture(concurrent.futures.Future[_T]):
|
||||
def __init__(self, work: concurrent.futures.Future[_T], on_wait: Callable[[float | None], None]) -> None:
|
||||
super().__init__()
|
||||
self._work: Final = work
|
||||
self._on_wait: Final = on_wait
|
||||
|
||||
def result(self, timeout: float | None = None) -> _T:
|
||||
self._on_wait(timeout)
|
||||
raise concurrent.futures.TimeoutError
|
||||
|
||||
def cancel(self) -> bool:
|
||||
return self._work.cancel()
|
||||
|
||||
|
||||
class _TimeBoundOutlastingExecutor(concurrent.futures.Executor):
|
||||
def __init__(self, workers: concurrent.futures.Executor) -> None:
|
||||
self._workers: Final = workers
|
||||
self.waits: tuple[float | None, ...] = ()
|
||||
|
||||
def _record_wait(self, timeout: float | None) -> None:
|
||||
self.waits = (*self.waits, timeout)
|
||||
|
||||
def submit(self, fn: Callable[_P, _T], /, *args: _P.args, **kwargs: _P.kwargs) -> concurrent.futures.Future[_T]:
|
||||
return _OutlastedTimeBoundFuture(self._workers.submit(fn, *args, **kwargs), self._record_wait)
|
||||
|
||||
def shutdown(self, wait: bool = True, *, cancel_futures: bool = False) -> None:
|
||||
self._workers.shutdown(wait=wait, cancel_futures=cancel_futures)
|
||||
|
||||
|
||||
async def _fail_chat_call(call: Literal["completion", "acompletion"], messages: list[dict[str, str]]) -> None:
|
||||
error: Final = _bad_key_error()
|
||||
request: Final = (
|
||||
litellm.acompletion(model="gpt-4o-mini", messages=messages, mock_response=error)
|
||||
if call == "acompletion"
|
||||
else asyncio.to_thread(litellm.completion, model="gpt-4o-mini", messages=messages, mock_response=error)
|
||||
)
|
||||
with pytest.raises(litellm.AuthenticationError):
|
||||
await request
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("call", ["completion", "acompletion"])
|
||||
async def test_failure_log_carries_the_prompt_a_logging_only_guardrail_masked(
|
||||
monkeypatch: pytest.MonkeyPatch, call: Literal["completion", "acompletion"]
|
||||
) -> None:
|
||||
guardrail: Final = _NameMaskingGuardrail()
|
||||
recorder, sync_logged = _register_failure_logging(monkeypatch, guardrail)
|
||||
messages: Final = copy.deepcopy(_RAW_LOG)
|
||||
|
||||
await _fail_chat_call(call, messages)
|
||||
|
||||
assert recorder.logged == [(_MASKED_LOG, _MASKED_LOG)]
|
||||
assert sync_logged == [_MASKED_LOG]
|
||||
assert messages == _RAW_LOG
|
||||
assert guardrail.scans == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"make_guardrail",
|
||||
[_UnflaggedNameMaskingGuardrail, lambda: _NameMaskingGuardrail(event_hook=GuardrailEventHooks.pre_call)],
|
||||
ids=["logging_only_without_masking_opt_in", "masking_opt_in_in_pre_call_mode"],
|
||||
)
|
||||
@pytest.mark.parametrize("call", ["completion", "acompletion"])
|
||||
async def test_failure_log_keeps_the_prompt_when_no_logging_only_guardrail_masks_it(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
make_guardrail: Callable[[], _NameMaskingGuardrail],
|
||||
call: Literal["completion", "acompletion"],
|
||||
) -> None:
|
||||
guardrail: Final = make_guardrail()
|
||||
recorder, sync_logged = _register_failure_logging(monkeypatch, guardrail)
|
||||
|
||||
await _fail_chat_call(call, copy.deepcopy(_RAW_LOG))
|
||||
|
||||
assert recorder.logged == [(_RAW_LOG, _RAW_LOG)]
|
||||
assert sync_logged == [_RAW_LOG]
|
||||
assert guardrail.scans == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("call", ["completion", "acompletion"])
|
||||
async def test_failure_is_logged_redacted_when_the_masking_guardrail_raises(
|
||||
monkeypatch: pytest.MonkeyPatch, call: Literal["completion", "acompletion"]
|
||||
) -> None:
|
||||
recorder, sync_logged = _register_failure_logging(monkeypatch, _RaisingMaskingGuardrail())
|
||||
|
||||
await _fail_chat_call(call, copy.deepcopy(_RAW_LOG))
|
||||
|
||||
assert recorder.logged == [(_REDACTED_LOG, _REDACTED_LOG)]
|
||||
assert sync_logged == [_REDACTED_LOG]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_failure_log_still_goes_out_masked_when_the_caller_is_cancelled_mid_mask(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
guardrail: Final = _BlockingNameMaskingGuardrail()
|
||||
recorder, sync_logged = _register_failure_logging(monkeypatch, guardrail)
|
||||
call: Final = asyncio.create_task(
|
||||
litellm.acompletion(model="gpt-4o-mini", messages=copy.deepcopy(_RAW_LOG), mock_response=_bad_key_error())
|
||||
)
|
||||
await asyncio.wait_for(guardrail.scanning.wait(), timeout=10)
|
||||
|
||||
call.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await call
|
||||
guardrail.release.set()
|
||||
await asyncio.wait_for(_until(lambda: bool(sync_logged) and bool(recorder.logged)), timeout=10)
|
||||
|
||||
assert sync_logged == [_MASKED_LOG]
|
||||
assert recorder.logged == [(_MASKED_LOG, _MASKED_LOG)]
|
||||
|
||||
|
||||
def _sync_failure_logging(call_id: str, guardrail: CustomGuardrail, recorder: _FailureLogRecorder) -> LitellmLogging:
|
||||
logging_obj: Final = LitellmLogging(
|
||||
model="gpt-4o-mini",
|
||||
messages=copy.deepcopy(_RAW_LOG),
|
||||
stream=False,
|
||||
call_type="completion",
|
||||
start_time=datetime.datetime(2026, 1, 1),
|
||||
litellm_call_id=call_id,
|
||||
function_id=call_id,
|
||||
dynamic_failure_callbacks=[guardrail, recorder],
|
||||
)
|
||||
logging_obj.update_environment_variables(litellm_params={"metadata": {}}, optional_params={})
|
||||
return logging_obj
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("no_global_callbacks")
|
||||
def test_sync_failure_log_is_redacted_when_the_mask_outlasts_its_time_bound() -> None:
|
||||
guardrail: Final = _HangingSyncMaskingGuardrail()
|
||||
recorder: Final = _FailureLogRecorder()
|
||||
logging_obj: Final = _sync_failure_logging("hanging-mask-failure", guardrail, recorder)
|
||||
mask_executor: Final = _TimeBoundOutlastingExecutor(
|
||||
BoundedLoggingThreadPoolExecutor(max_workers=1, max_pending_tasks=1)
|
||||
)
|
||||
|
||||
try:
|
||||
logging_obj.mask_failure_log_request(mask_executor=mask_executor)
|
||||
finally:
|
||||
guardrail.release.set()
|
||||
mask_executor.shutdown(wait=True)
|
||||
logging_obj.failure_handler(litellm.ServiceUnavailableError("503", "openai", "gpt-4o-mini"), "tb")
|
||||
|
||||
assert mask_executor.waits == (FAILURE_LOG_MASK_TIMEOUT_SECONDS,)
|
||||
assert recorder.logged == [(_REDACTED_LOG, _REDACTED_LOG)]
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("no_global_callbacks")
|
||||
def test_sync_failure_is_logged_redacted_when_the_logging_executor_drops_the_mask() -> None:
|
||||
recorder: Final = _FailureLogRecorder()
|
||||
logging_obj: Final = _sync_failure_logging("full-executor-failure", _NameMaskingGuardrail(), recorder)
|
||||
|
||||
logging_obj.mask_failure_log_request(
|
||||
mask_executor=BoundedLoggingThreadPoolExecutor(max_workers=1, max_pending_tasks=0)
|
||||
)
|
||||
logging_obj.failure_handler(litellm.ServiceUnavailableError("503", "openai", "gpt-4o-mini"), "tb")
|
||||
|
||||
assert recorder.logged == [(_REDACTED_LOG, _REDACTED_LOG)]
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("no_global_callbacks")
|
||||
def test_sync_failure_is_logged_redacted_when_the_masking_guardrail_has_only_an_async_hook() -> None:
|
||||
recorder: Final = _FailureLogRecorder()
|
||||
logging_obj: Final = _sync_failure_logging(
|
||||
"async-only-mask-failure",
|
||||
_AsyncOnlyNameMaskingGuardrail(
|
||||
guardrail_name="async-only-name-masker", event_hook=GuardrailEventHooks.logging_only
|
||||
),
|
||||
recorder,
|
||||
)
|
||||
|
||||
logging_obj.mask_failure_log_request()
|
||||
logging_obj.failure_handler(litellm.ServiceUnavailableError("503", "openai", "gpt-4o-mini"), "tb")
|
||||
|
||||
assert recorder.logged == [(_REDACTED_LOG, _REDACTED_LOG)]
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("no_global_callbacks")
|
||||
def test_sync_failure_mask_still_queued_at_its_time_bound_never_runs() -> None:
|
||||
guardrail: Final = _NameMaskingGuardrail()
|
||||
recorder: Final = _FailureLogRecorder()
|
||||
logging_obj: Final = _sync_failure_logging("queued-mask-failure", guardrail, recorder)
|
||||
mask_executor: Final = _TimeBoundOutlastingExecutor(
|
||||
BoundedLoggingThreadPoolExecutor(max_workers=1, max_pending_tasks=2)
|
||||
)
|
||||
busy_worker_released: Final = threading.Event()
|
||||
_ = mask_executor.submit(busy_worker_released.wait, 30)
|
||||
|
||||
try:
|
||||
logging_obj.mask_failure_log_request(mask_executor=mask_executor)
|
||||
finally:
|
||||
busy_worker_released.set()
|
||||
mask_executor.shutdown(wait=True)
|
||||
logging_obj.failure_handler(litellm.ServiceUnavailableError("503", "openai", "gpt-4o-mini"), "tb")
|
||||
|
||||
assert guardrail.scans == 0
|
||||
assert recorder.logged == [(_REDACTED_LOG, _REDACTED_LOG)]
|
||||
|
||||
|
||||
class _ProxyRequestRecorder(_FailureLogRecorder):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.request_copies: list[tuple[object, object]] = []
|
||||
|
||||
def _record(self, kwargs: dict[str, object]) -> None:
|
||||
super()._record(kwargs)
|
||||
litellm_params: Final = cast(dict[str, object], kwargs["litellm_params"])
|
||||
self.request_copies.append(
|
||||
copy.deepcopy((kwargs["optional_params"], litellm_params["proxy_server_request"]))
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("no_global_callbacks")
|
||||
@pytest.mark.parametrize("sync", [True, False])
|
||||
@pytest.mark.parametrize(
|
||||
("guardrail_type", "logged_messages"),
|
||||
[(_NameMaskingGuardrail, _MASKED_LOG), (_RaisingMaskingGuardrail, _REDACTED_LOG)],
|
||||
)
|
||||
async def test_failure_log_masks_the_prompt_copies_a_proxy_rejection_logs(
|
||||
sync: bool, guardrail_type: type[_NameMaskingGuardrail], logged_messages: list[dict[str, str]]
|
||||
) -> None:
|
||||
messages: Final = copy.deepcopy(_RAW_LOG)
|
||||
guardrail: Final = guardrail_type()
|
||||
recorder: Final = _ProxyRequestRecorder()
|
||||
logging_obj: Final = LitellmLogging(
|
||||
model="gpt-4o-mini",
|
||||
messages=messages,
|
||||
stream=False,
|
||||
call_type="completion" if sync else "acompletion",
|
||||
start_time=datetime.datetime(2026, 1, 1),
|
||||
litellm_call_id="proxy-rejection-copies",
|
||||
function_id="proxy-rejection-copies",
|
||||
dynamic_failure_callbacks=[guardrail, recorder],
|
||||
dynamic_async_failure_callbacks=[guardrail, recorder],
|
||||
)
|
||||
logging_obj.update_environment_variables(
|
||||
litellm_params={"metadata": {}, "proxy_server_request": {"url": "/chat/completions", "body": {"messages": messages}}},
|
||||
optional_params={"messages": messages},
|
||||
)
|
||||
logging_obj.model_call_details["input"] = messages
|
||||
error: Final = litellm.RateLimitError("429", "openai", "gpt-4o-mini")
|
||||
|
||||
if sync:
|
||||
await asyncio.to_thread(logging_obj.failure_handler, error, "tb")
|
||||
else:
|
||||
await logging_obj.async_failure_handler(error, "tb")
|
||||
|
||||
assert len(recorder.request_copies) == 1
|
||||
optional_params, proxy_server_request = recorder.request_copies[0]
|
||||
assert cast(dict[str, object], optional_params)["messages"] == logged_messages
|
||||
assert _RAW_PROMPT not in str(proxy_server_request)
|
||||
assert "/chat/completions" in str(proxy_server_request)
|
||||
assert messages == _RAW_LOG
|
||||
|
||||
|
||||
async def _fail_once_the_mask_outlasts_its_time_bound(
|
||||
clock: _SkippableClockLoop, guardrail: _BlockingNameMaskingGuardrail, recorder: _FailureLogRecorder
|
||||
) -> None:
|
||||
call: Final = asyncio.create_task(
|
||||
litellm.acompletion(model="gpt-4o-mini", messages=copy.deepcopy(_RAW_LOG), mock_response=_bad_key_error())
|
||||
)
|
||||
await asyncio.wait_for(guardrail.scanning.wait(), timeout=10)
|
||||
clock.skip(FAILURE_LOG_MASK_TIMEOUT_SECONDS)
|
||||
with pytest.raises(litellm.AuthenticationError):
|
||||
await asyncio.wait_for(call, timeout=10)
|
||||
await asyncio.wait_for(_until(lambda: bool(recorder.logged)), timeout=10)
|
||||
|
||||
|
||||
def test_failure_log_is_redacted_when_the_mask_outlasts_its_time_bound(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
guardrail: Final = _BlockingNameMaskingGuardrail()
|
||||
recorder, sync_logged = _register_failure_logging(monkeypatch, guardrail)
|
||||
|
||||
_run_on_skippable_clock(lambda clock: _fail_once_the_mask_outlasts_its_time_bound(clock, guardrail, recorder))
|
||||
|
||||
assert sync_logged == [_REDACTED_LOG]
|
||||
assert recorder.logged == [(_REDACTED_LOG, _REDACTED_LOG)]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("no_global_callbacks")
|
||||
async def test_concurrent_failure_notifications_log_one_masked_failure() -> None:
|
||||
recorder: Final = _FailureLogRecorder()
|
||||
logging_obj: Final = LitellmLogging(
|
||||
model="gpt-4o-mini",
|
||||
messages=copy.deepcopy(_RAW_LOG),
|
||||
stream=False,
|
||||
call_type="acompletion",
|
||||
start_time=datetime.datetime(2026, 1, 1),
|
||||
litellm_call_id="concurrent-masked-failure",
|
||||
function_id="concurrent-masked-failure",
|
||||
dynamic_async_failure_callbacks=[_NameMaskingGuardrail(), recorder],
|
||||
)
|
||||
logging_obj.update_environment_variables(litellm_params={"acompletion": True, "metadata": {}}, optional_params={})
|
||||
error: Final = litellm.ServiceUnavailableError("503", "openai", "gpt-4o-mini")
|
||||
|
||||
await asyncio.gather(*(logging_obj.async_failure_handler(error, "tb") for _ in range(3)))
|
||||
|
||||
assert recorder.logged == [(_MASKED_LOG, _MASKED_LOG)]
|
||||
|
||||
|
||||
class _DeploymentSuccessRecorder(CustomLogger):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.logged: list[tuple[str, str, str]] = []
|
||||
|
||||
async def async_log_success_event(
|
||||
self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object
|
||||
) -> None:
|
||||
payload: Final = cast(StandardLoggingPayload, kwargs["standard_logging_object"])
|
||||
self.logged.append((payload["model_id"] or "", payload["model_group"] or "", payload["api_base"]))
|
||||
|
||||
|
||||
def _deployment_params(model_id: str, model_group: str) -> dict[str, object]:
|
||||
return {
|
||||
"acompletion": True,
|
||||
"api_base": f"https://{model_group}.example.com",
|
||||
"metadata": {"model_group": model_group, "model_info": {"id": model_id}},
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("no_global_callbacks")
|
||||
async def test_fallback_started_while_the_failure_log_is_masked_logs_its_own_deployment() -> None:
|
||||
guardrail: Final = _BlockingNameMaskingGuardrail()
|
||||
failures: Final = _FailureLogRecorder()
|
||||
successes: Final = _DeploymentSuccessRecorder()
|
||||
logging_obj: Final = LitellmLogging(
|
||||
model="gpt-4o-mini",
|
||||
messages=copy.deepcopy(_RAW_LOG),
|
||||
stream=False,
|
||||
call_type="acompletion",
|
||||
start_time=datetime.datetime(2026, 1, 1),
|
||||
litellm_call_id="fallback-during-mask",
|
||||
function_id="fallback-during-mask",
|
||||
dynamic_async_failure_callbacks=[guardrail, failures],
|
||||
dynamic_async_success_callbacks=[successes],
|
||||
)
|
||||
logging_obj.update_environment_variables(litellm_params=_deployment_params("rejected", "primary"), optional_params={})
|
||||
rejected_attempt_log: Final = asyncio.create_task(
|
||||
logging_obj.async_failure_handler(litellm.RateLimitError("429", "openai", "gpt-4o-mini"), "tb")
|
||||
)
|
||||
await asyncio.wait_for(guardrail.scanning.wait(), timeout=10)
|
||||
|
||||
logging_obj.update_environment_variables(litellm_params=_deployment_params("served", "backup"), optional_params={})
|
||||
guardrail.release.set()
|
||||
await asyncio.wait_for(rejected_attempt_log, timeout=10)
|
||||
await logging_obj.async_success_handler(
|
||||
result=ModelResponse(
|
||||
model="gpt-4o-mini",
|
||||
choices=[{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}],
|
||||
)
|
||||
)
|
||||
|
||||
assert failures.logged == [(_MASKED_LOG, _MASKED_LOG)]
|
||||
assert successes.logged == [("served", "backup", "https://backup.example.com")]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("no_global_callbacks")
|
||||
async def test_sync_sdk_failure_dispatched_from_async_code_logs_the_masked_prompt() -> None:
|
||||
recorder: Final = _FailureLogRecorder()
|
||||
logging_obj: Final = LitellmLogging(
|
||||
model="gpt-4o-mini",
|
||||
messages=copy.deepcopy(_RAW_LOG),
|
||||
stream=True,
|
||||
call_type="completion",
|
||||
start_time=datetime.datetime(2026, 1, 1),
|
||||
litellm_call_id="sync-sdk-masked-failure",
|
||||
function_id="sync-sdk-masked-failure",
|
||||
dynamic_failure_callbacks=[_NameMaskingGuardrail(), recorder],
|
||||
)
|
||||
logging_obj.update_environment_variables(litellm_params={"metadata": {}}, optional_params={})
|
||||
|
||||
await logging_obj.dispatch_failure_handlers(litellm.ServiceUnavailableError("503", "openai", "gpt-4o-mini"), "tb")
|
||||
|
||||
assert recorder.logged == [(_MASKED_LOG, _MASKED_LOG)]
|
||||
|
||||
|
||||
def _async_request_logging_obj(call_id: str, callbacks: list[CustomLogger], stream: bool = False) -> LitellmLogging:
|
||||
logging_obj: Final = LitellmLogging(
|
||||
model="gpt-4o-mini",
|
||||
|
|
|
|||
|
|
@ -4,15 +4,16 @@ import itertools
|
|||
import json
|
||||
import os
|
||||
import queue
|
||||
import selectors
|
||||
import subprocess
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
from collections.abc import Mapping, Sequence
|
||||
from collections.abc import Coroutine, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from types import CodeType, FrameType, MappingProxyType, TracebackType
|
||||
from typing import AsyncIterator, Callable, Final, Iterator, Literal, NoReturn, Optional, cast
|
||||
from typing import AsyncIterator, Callable, ClassVar, Final, Iterator, Literal, NoReturn, Optional, cast
|
||||
from unittest.mock import AsyncMock, MagicMock, Mock, patch
|
||||
|
||||
import httpx
|
||||
|
|
@ -23,7 +24,9 @@ from pydantic import BaseModel, TypeAdapter
|
|||
import litellm
|
||||
from litellm import verbose_logger
|
||||
from litellm._logging import session_id_var, trace_id_var
|
||||
from litellm.constants import FAILURE_LOG_MASK_TIMEOUT_SECONDS
|
||||
from litellm.exceptions import MidStreamFallbackError
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.litellm_core_utils.get_litellm_params import get_litellm_params
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
|
|
@ -35,6 +38,7 @@ from litellm.litellm_core_utils.streaming_handler import (
|
|||
from litellm.llms.custom_httpx.http_handler import HTTPHandler
|
||||
from litellm.llms.openai.chat.gpt_transformation import OpenAIChatCompletionStreamingHandler
|
||||
from litellm.responses.mcp.chat_completions_handler import _add_mcp_metadata_to_response
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
from litellm.types.utils import (
|
||||
ChoiceLogprobs,
|
||||
CompletionTokensDetailsWrapper,
|
||||
|
|
@ -3478,7 +3482,9 @@ def _stuck_stream_wrapper(
|
|||
provider_stream: Iterator[ModelResponseStream] | AsyncIterator[ModelResponseStream],
|
||||
call_type: Literal["completion", "acompletion"],
|
||||
recorder: _StuckStreamFailureRecorder,
|
||||
masker: CustomGuardrail | None = None,
|
||||
) -> CustomStreamWrapper:
|
||||
failure_callbacks: Final[list[CustomLogger]] = [recorder] if masker is None else [masker, recorder]
|
||||
logging_obj: Final = Logging(
|
||||
model="test-model",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
|
|
@ -3487,8 +3493,8 @@ def _stuck_stream_wrapper(
|
|||
start_time=datetime.datetime(2026, 1, 1),
|
||||
litellm_call_id=f"stuck-stream-{call_type}",
|
||||
function_id=f"stuck-stream-{call_type}",
|
||||
dynamic_failure_callbacks=[recorder],
|
||||
dynamic_async_failure_callbacks=[recorder],
|
||||
dynamic_failure_callbacks=failure_callbacks,
|
||||
dynamic_async_failure_callbacks=failure_callbacks,
|
||||
)
|
||||
logging_obj.update_environment_variables(
|
||||
model="test-model",
|
||||
|
|
@ -3503,6 +3509,83 @@ def _stuck_stream_wrapper(
|
|||
)
|
||||
|
||||
|
||||
class _BlockedStuckStreamFailureRecorder(_StuckStreamFailureRecorder):
|
||||
def __init__(self, release: asyncio.Event) -> None:
|
||||
super().__init__()
|
||||
self.release = release
|
||||
|
||||
async def async_log_failure_event(
|
||||
self, kwargs: object, response_obj: object, start_time: object, end_time: object
|
||||
) -> None:
|
||||
await self.release.wait()
|
||||
await super().async_log_failure_event(kwargs, response_obj, start_time, end_time)
|
||||
|
||||
|
||||
class _FailureLogMasker(CustomGuardrail):
|
||||
masks_logged_request: ClassVar[bool] = True
|
||||
|
||||
async def async_logging_hook(
|
||||
self, kwargs: dict[str, object], result: object, call_type: str
|
||||
) -> tuple[dict[str, object], object]:
|
||||
return kwargs, result
|
||||
|
||||
|
||||
class _VirtualClockSelector(selectors.DefaultSelector):
|
||||
now: float = 0.0
|
||||
|
||||
def select(self, timeout: float | None = None) -> list[tuple[selectors.SelectorKey, int]]:
|
||||
if timeout is None:
|
||||
return super().select(None)
|
||||
ready: Final = super().select(0)
|
||||
if not ready:
|
||||
self.now += timeout
|
||||
return ready
|
||||
|
||||
|
||||
class _VirtualClockLoop(asyncio.SelectorEventLoop):
|
||||
def __init__(self) -> None:
|
||||
self.clock: Final = _VirtualClockSelector()
|
||||
super().__init__(self.clock)
|
||||
|
||||
def time(self) -> float:
|
||||
return self.clock.now
|
||||
|
||||
|
||||
def _on_virtual_clock(scenario: Coroutine[object, object, float]) -> float:
|
||||
loop: Final = _VirtualClockLoop()
|
||||
try:
|
||||
return loop.run_until_complete(scenario)
|
||||
finally:
|
||||
loop.run_until_complete(loop.shutdown_default_executor())
|
||||
loop.close()
|
||||
|
||||
|
||||
async def _seconds_waited_for_a_blocked_failure_log(masker: CustomGuardrail | None) -> float:
|
||||
release: Final = asyncio.Event()
|
||||
wrapper: Final = _stuck_stream_wrapper(
|
||||
_StuckProvider().async_stream(), "acompletion", _BlockedStuckStreamFailureRecorder(release), masker
|
||||
)
|
||||
with pytest.raises(Exception, match="repeating the same chunk"):
|
||||
async for _ in wrapper:
|
||||
pass
|
||||
clock: Final = asyncio.get_running_loop()
|
||||
started: Final = clock.time()
|
||||
await wrapper.wait_for_failure_logging(timeout=0.1)
|
||||
waited: Final = clock.time() - started
|
||||
release.set()
|
||||
await wrapper.wait_for_failure_logging(timeout=None)
|
||||
return waited
|
||||
|
||||
|
||||
def test_mid_stream_failure_wait_extends_by_the_mask_time_only_when_the_failure_log_is_masked() -> None:
|
||||
assert _on_virtual_clock(_seconds_waited_for_a_blocked_failure_log(None)) == pytest.approx(0.1)
|
||||
assert _on_virtual_clock(
|
||||
_seconds_waited_for_a_blocked_failure_log(
|
||||
_FailureLogMasker(guardrail_name="failure-log-masker", event_hook=GuardrailEventHooks.logging_only)
|
||||
)
|
||||
) == pytest.approx(0.1 + FAILURE_LOG_MASK_TIMEOUT_SECONDS)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_stream_stuck_on_one_chunk_closes_the_provider_stream_and_logs_one_failure():
|
||||
"""https://github.com/BerriAI/litellm/issues/19428 and https://github.com/BerriAI/litellm/issues/13786"""
|
||||
|
|
|
|||
|
|
@ -646,6 +646,7 @@ async def test_allm_passthrough_route_429_streaming_raises():
|
|||
mock_logging_obj.update_environment_variables = MagicMock()
|
||||
mock_logging_obj.async_flush_passthrough_collected_chunks = AsyncMock()
|
||||
mock_logging_obj.async_failure_handler = AsyncMock()
|
||||
mock_logging_obj.async_mask_failure_log_request = AsyncMock()
|
||||
|
||||
with (
|
||||
patch(
|
||||
|
|
|
|||
|
|
@ -227,3 +227,43 @@ async def test_guardrails_ai_process_input():
|
|||
|
||||
# Should keep original content when no output fields are present
|
||||
assert result["messages"][0]["content"] == "Original message"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_guardrails_ai_logging_only_masks_the_failure_log(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.proxy.guardrails.guardrail_hooks.guardrails_ai.guardrails_ai import GuardrailsAIResponse
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
|
||||
logged: list[object] = []
|
||||
|
||||
class _FailureLogRecorder(CustomLogger):
|
||||
async def async_log_failure_event(
|
||||
self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object
|
||||
) -> None:
|
||||
logged.append(kwargs["messages"])
|
||||
|
||||
guardrail = GuardrailsAI(
|
||||
guardrail_name="pii-guard",
|
||||
guard_name="pii-guard",
|
||||
api_base="http://test.example.com",
|
||||
event_hook=GuardrailEventHooks.logging_only,
|
||||
)
|
||||
for name in ("input_callback", "success_callback", "failure_callback", "_async_success_callback"):
|
||||
monkeypatch.setattr(litellm, name, [])
|
||||
monkeypatch.setattr(litellm, "_async_failure_callback", [])
|
||||
monkeypatch.setattr(litellm, "callbacks", [guardrail, _FailureLogRecorder()])
|
||||
|
||||
with patch.object(
|
||||
guardrail,
|
||||
"make_guardrails_ai_api_request",
|
||||
AsyncMock(return_value=GuardrailsAIResponse(validatedOutput="My name is <PERSON>")),
|
||||
):
|
||||
with pytest.raises(litellm.AuthenticationError):
|
||||
await litellm.acompletion(
|
||||
model="gpt-4o-mini",
|
||||
messages=[{"role": "user", "content": "My name is John Smith"}],
|
||||
mock_response=litellm.AuthenticationError(message="bad key", llm_provider="openai", model="gpt-4o-mini"),
|
||||
)
|
||||
|
||||
assert logged == [[{"role": "user", "content": "My name is <PERSON>"}]]
|
||||
|
|
|
|||
|
|
@ -11,7 +11,7 @@ import re
|
|||
from collections.abc import Iterable, Sequence
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import Final, Literal
|
||||
from typing import Final, Literal, cast
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import aiohttp
|
||||
|
|
@ -26,6 +26,7 @@ import litellm
|
|||
from litellm import mock_completion
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.guardrails.guardrail_hooks.presidio import (
|
||||
OPTIONAL_PresidioPIIMasking,
|
||||
|
|
@ -33,7 +34,7 @@ from litellm.proxy.guardrails.guardrail_hooks.presidio import (
|
|||
)
|
||||
from litellm.types.guardrails import LitellmParams, PiiAction, PiiEntityType
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.presidio import PresidioAnalyzeRequest, PresidioAnalyzeResponseItem
|
||||
from litellm.types.utils import Choices, Delta, Message, ModelResponse, StreamingChoices
|
||||
from litellm.types.utils import Choices, Delta, Message, ModelResponse, StandardLoggingPayload, StreamingChoices
|
||||
|
||||
|
||||
def _make_mock_session_iterator(json_response, status=200, content_type="application/json", text_response=""):
|
||||
|
|
@ -4425,6 +4426,45 @@ async def test_pii_masking_replays_a_byte_identical_prefix_across_turns(mock_use
|
|||
assert later[3]["content"] == "Now compare against <PERSON> too."
|
||||
|
||||
|
||||
class _FailureLogRecorder(CustomLogger):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.logged: list[tuple[object, object]] = []
|
||||
|
||||
async def async_log_failure_event(
|
||||
self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object
|
||||
) -> None:
|
||||
payload: Final = cast(StandardLoggingPayload, kwargs["standard_logging_object"])
|
||||
self.logged.append((copy.deepcopy(kwargs["messages"]), copy.deepcopy(payload["messages"])))
|
||||
|
||||
|
||||
async def test_logging_only_presidio_masks_the_prompt_in_the_failure_log(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""A failed call is logged with the prompt a logging_only Presidio guardrail masked, the same as a
|
||||
successful one. The analyzer and anonymizer are the in-process fake handed over through api_base."""
|
||||
recorder: Final = _FailureLogRecorder()
|
||||
for name in ("input_callback", "success_callback", "failure_callback", "_async_success_callback"):
|
||||
monkeypatch.setattr(litellm, name, [])
|
||||
monkeypatch.setattr(litellm, "_async_failure_callback", [])
|
||||
async with TestServer(_fake_presidio_app()) as server:
|
||||
guardrail: Final = OPTIONAL_PresidioPIIMasking(
|
||||
presidio_analyzer_api_base=str(server.make_url("/")),
|
||||
presidio_anonymizer_api_base=str(server.make_url("/")),
|
||||
pii_entities_config={PiiEntityType.PERSON: PiiAction.MASK},
|
||||
logging_only=True,
|
||||
)
|
||||
monkeypatch.setattr(litellm, "callbacks", [guardrail, recorder])
|
||||
with pytest.raises(litellm.AuthenticationError):
|
||||
await litellm.acompletion(
|
||||
model="gpt-4o-mini",
|
||||
messages=[{"role": "user", "content": "My name is John Smith"}],
|
||||
mock_response=litellm.AuthenticationError(message="bad key", llm_provider="openai", model="gpt-4o-mini"),
|
||||
)
|
||||
await guardrail._close_http_session()
|
||||
|
||||
masked: Final = [{"role": "user", "content": "My name is <PERSON>"}]
|
||||
assert recorder.logged == [(masked, masked)]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("surface", ["mcp_arguments", "mcp_result", "llm_output"])
|
||||
@pytest.mark.parametrize("action", [PiiAction.MASK, PiiAction.BLOCK])
|
||||
|
|
@ -4857,11 +4897,17 @@ def _card_and_email_spans(text: str) -> tuple[PresidioAnalyzeResponseItem, ...]:
|
|||
)
|
||||
|
||||
|
||||
_PENDING_PRESIDIO_REPLIES: Final[set[asyncio.Task[None]]] = set() # mutable-ok: keeps pending replies alive
|
||||
|
||||
|
||||
class _InMemoryPresidioTransport(asyncio.Transport):
|
||||
def __init__(self, protocol: ResponseHandler, analyzed: asyncio.Queue[PresidioAnalyzeRequest]) -> None:
|
||||
def __init__(
|
||||
self, protocol: ResponseHandler, analyzed: asyncio.Queue[PresidioAnalyzeRequest], replies_wait_for: int
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self._protocol: Final = protocol
|
||||
self._analyzed: Final = analyzed
|
||||
self._replies_wait_for: Final = replies_wait_for
|
||||
self._received = b""
|
||||
self._closing = False
|
||||
|
||||
|
|
@ -4881,7 +4927,12 @@ class _InMemoryPresidioTransport(asyncio.Transport):
|
|||
+ f"Content-Length: {len(reply)}\r\n\r\n".encode()
|
||||
+ reply
|
||||
)
|
||||
asyncio.get_running_loop().call_soon(self._protocol.data_received, response)
|
||||
_PENDING_PRESIDIO_REPLIES.add(asyncio.get_running_loop().create_task(self._reply_once_analyzed(response)))
|
||||
|
||||
async def _reply_once_analyzed(self, response: bytes) -> None:
|
||||
while self._analyzed.qsize() < self._replies_wait_for:
|
||||
await asyncio.sleep(0.01)
|
||||
self._protocol.data_received(response)
|
||||
|
||||
def _reply(self, path: str, body: bytes) -> tuple[PresidioAnalyzeResponseItem, ...] | _PresidioAnonymizeReply:
|
||||
if path == "/analyze":
|
||||
|
|
@ -4922,15 +4973,16 @@ class _InMemoryPresidioTransport(asyncio.Transport):
|
|||
|
||||
|
||||
class _InMemoryPresidioConnector(aiohttp.BaseConnector):
|
||||
def __init__(self, analyzed: asyncio.Queue[PresidioAnalyzeRequest]) -> None:
|
||||
def __init__(self, analyzed: asyncio.Queue[PresidioAnalyzeRequest], replies_wait_for: int) -> None:
|
||||
super().__init__()
|
||||
self._analyzed: Final = analyzed
|
||||
self._replies_wait_for: Final = replies_wait_for
|
||||
|
||||
async def _create_connection( # pyright: ignore[reportImplicitOverride] # aiohttp's connector extension point
|
||||
self, req: aiohttp.ClientRequest, traces: Sequence[object], timeout: aiohttp.ClientTimeout
|
||||
) -> ResponseHandler:
|
||||
protocol: Final = ResponseHandler(asyncio.get_running_loop())
|
||||
protocol.connection_made(_InMemoryPresidioTransport(protocol, self._analyzed))
|
||||
protocol.connection_made(_InMemoryPresidioTransport(protocol, self._analyzed, self._replies_wait_for))
|
||||
return protocol
|
||||
|
||||
|
||||
|
|
@ -4939,17 +4991,37 @@ def _drain(analyzed: asyncio.Queue[PresidioAnalyzeRequest]) -> tuple[PresidioAna
|
|||
|
||||
|
||||
def _guardrail_with_in_memory_presidio(
|
||||
analyzed: asyncio.Queue[PresidioAnalyzeRequest],
|
||||
analyzed: asyncio.Queue[PresidioAnalyzeRequest], replies_wait_for: int = 0
|
||||
) -> OPTIONAL_PresidioPIIMasking:
|
||||
guardrail: Final = OPTIONAL_PresidioPIIMasking(
|
||||
pii_entities_config=_BLOCK_CARD_MASK_EMAIL,
|
||||
presidio_analyzer_api_base="http://presidio-analyzer.test/",
|
||||
presidio_anonymizer_api_base="http://presidio-anonymizer.test/",
|
||||
)
|
||||
guardrail._http_session = aiohttp.ClientSession(connector=_InMemoryPresidioConnector(analyzed))
|
||||
guardrail._http_session = aiohttp.ClientSession(connector=_InMemoryPresidioConnector(analyzed, replies_wait_for))
|
||||
return guardrail
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_presidio_scans_every_message_of_a_request_at_once() -> None:
|
||||
texts: Final = (f"My email is {_EMAIL}", f"Write to {_EMAIL} today")
|
||||
analyzed: Final[asyncio.Queue[PresidioAnalyzeRequest]] = asyncio.Queue()
|
||||
guardrail: Final = _guardrail_with_in_memory_presidio(analyzed, replies_wait_for=len(texts))
|
||||
try:
|
||||
_ = await asyncio.wait_for(
|
||||
guardrail.async_logging_hook(
|
||||
kwargs={"messages": [{"role": "user", "content": text} for text in texts]},
|
||||
result=None,
|
||||
call_type="acompletion",
|
||||
),
|
||||
timeout=5,
|
||||
)
|
||||
finally:
|
||||
await guardrail._close_http_session()
|
||||
|
||||
assert sorted(request.get("text") or "" for request in _drain(analyzed)) == sorted(texts)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_check_pii_raises_blocked_entity_for_card() -> None:
|
||||
text: Final = f"My credit card number is {_CARD_NUMBER} and my email is {_EMAIL}"
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ import json
|
|||
import smtplib
|
||||
import traceback
|
||||
from collections.abc import Mapping
|
||||
from typing import Final
|
||||
from typing import ClassVar, Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
|
@ -3379,6 +3379,36 @@ def _dispatched_failure_logging_obj(model: str, litellm_params: dict) -> object:
|
|||
return logging_obj
|
||||
|
||||
|
||||
class _LongPlaceholderMaskingGuardrail(CustomGuardrail):
|
||||
masks_logged_request: ClassVar[bool] = True
|
||||
|
||||
async def async_logging_hook(
|
||||
self, kwargs: dict[str, object], result: object, call_type: str
|
||||
) -> tuple[dict[str, object], object]:
|
||||
return {**kwargs, "messages": [{"role": "user", "content": "<PERSON> " * 40}]}, result
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_billed_failure_is_estimated_on_the_prompt_sent_not_the_masked_failure_log(monkeypatch):
|
||||
from litellm.litellm_core_utils.get_litellm_params import get_litellm_params
|
||||
|
||||
def billed_prompt_tokens(logging_obj: object) -> int:
|
||||
lifted: Final = _failure_fields_to_lift({"litellm_logging_obj": logging_obj}, provider_billed=True)
|
||||
return lifted["combined_usage_object"].prompt_tokens
|
||||
|
||||
unmasked: Final = _dispatched_failure_logging_obj(FAILURE_USAGE_MODEL, get_litellm_params(metadata={}))
|
||||
masked: Final = _dispatched_failure_logging_obj(FAILURE_USAGE_MODEL, get_litellm_params(metadata={}))
|
||||
monkeypatch.setattr(
|
||||
litellm,
|
||||
"_async_failure_callback",
|
||||
[_LongPlaceholderMaskingGuardrail(guardrail_name="masker", event_hook=GuardrailEventHooks.logging_only)],
|
||||
)
|
||||
await masked.async_mask_failure_log_request()
|
||||
|
||||
assert masked.model_call_details["messages"] != ONE_USER_MESSAGE
|
||||
assert billed_prompt_tokens(masked) == billed_prompt_tokens(unmasked)
|
||||
|
||||
|
||||
def test_the_deployment_rate_is_found_on_the_logging_objects_real_litellm_params(monkeypatch):
|
||||
"""``get_litellm_params()`` always carries a ``litellm_metadata`` key (None on chat
|
||||
completions) next to the router-stamped ``metadata.model_info.id``; the failure row must
|
||||
|
|
|
|||
|
|
@ -8,18 +8,20 @@ import json
|
|||
from collections.abc import Callable
|
||||
from datetime import datetime
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
from typing import ClassVar, Final
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
from fastapi import HTTPException
|
||||
|
||||
import litellm
|
||||
from litellm.constants import PROXY_REJECTED_BEFORE_ROUTING_KEY
|
||||
from litellm.constants import PROXY_REJECTED_BEFORE_ROUTING_KEY, REDACTED_BY_LITELLM
|
||||
from litellm.exceptions import GuardrailRaisedException
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.proxy._types import ProxyErrorTypes, UserAPIKeyAuth
|
||||
from litellm.proxy.utils import REJECTED_BEFORE_REQUEST_SETUP, REJECTED_BEFORE_REQUEST_SETUP_KEY, ProxyLogging
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
from litellm.types.utils import CachingDetails
|
||||
|
||||
|
||||
|
|
@ -500,6 +502,181 @@ async def test_proxy_only_failure_log_keeps_caller_tags_only_after_auth(
|
|||
assert recorded[0]["standard_logging_object"]["request_tags"] == logged_tags
|
||||
|
||||
|
||||
class _PromptMaskingGuardrail(CustomGuardrail):
|
||||
masks_logged_request: ClassVar[bool] = True
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__(guardrail_name="prompt-masker", event_hook=GuardrailEventHooks.logging_only)
|
||||
self.scans = 0
|
||||
|
||||
async def async_logging_hook(
|
||||
self, kwargs: dict[str, object], result: object, call_type: str
|
||||
) -> tuple[dict[str, object], object]:
|
||||
self.scans += 1
|
||||
return {**kwargs, "messages": [{"role": "user", "content": "My name is <PERSON>"}]}, result
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("error_type", "masker_scans", "logged_content"),
|
||||
[(ProxyErrorTypes.auth_error, 0, REDACTED_BY_LITELLM), (None, 1, "My name is <PERSON>")],
|
||||
)
|
||||
async def test_proxy_only_failure_log_masks_after_auth_and_redacts_before_it(
|
||||
proxy_logging,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
error_type: ProxyErrorTypes | None,
|
||||
masker_scans: int,
|
||||
logged_content: str,
|
||||
) -> None:
|
||||
from litellm.proxy import proxy_server
|
||||
|
||||
recorded: Final[list[dict[str, object]]] = []
|
||||
|
||||
class _RecordingLogger(CustomLogger):
|
||||
async def async_log_failure_event(
|
||||
self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object
|
||||
) -> None:
|
||||
recorded.append(kwargs)
|
||||
|
||||
masker: Final = _PromptMaskingGuardrail()
|
||||
monkeypatch.setattr(proxy_server, "llm_router", None)
|
||||
monkeypatch.setattr(litellm, "callbacks", [masker, _RecordingLogger()])
|
||||
proxy_logging.alert_types = []
|
||||
|
||||
await proxy_logging.post_call_failure_hook(
|
||||
request_data={"model": "internal-model", "messages": [{"role": "user", "content": "My name is John Smith"}]},
|
||||
original_exception=HTTPException(status_code=401, detail="Invalid proxy server token passed"),
|
||||
user_api_key_dict=UserAPIKeyAuth(request_route="/chat/completions"),
|
||||
error_type=error_type,
|
||||
route="/chat/completions",
|
||||
)
|
||||
|
||||
logged_messages: Final = [{"role": "user", "content": logged_content}]
|
||||
assert masker.scans == masker_scans
|
||||
assert [(kwargs["messages"], kwargs["standard_logging_object"]["messages"]) for kwargs in recorded] == [
|
||||
(logged_messages, logged_messages)
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("error_type", "logged_content"),
|
||||
[(ProxyErrorTypes.auth_error, REDACTED_BY_LITELLM), (None, "My name is <PERSON>")],
|
||||
)
|
||||
async def test_failure_hooks_get_the_request_body_with_the_prompt_the_failure_log_masked(
|
||||
proxy_logging, monkeypatch: pytest.MonkeyPatch, error_type: ProxyErrorTypes | None, logged_content: str
|
||||
) -> None:
|
||||
from litellm.proxy import proxy_server
|
||||
|
||||
received_bodies: Final[list[object]] = []
|
||||
|
||||
class _SpendLogLikeHook(CustomLogger):
|
||||
async def async_post_call_failure_hook(
|
||||
self,
|
||||
request_data: dict[str, object],
|
||||
original_exception: Exception,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
traceback_str: str | None = None,
|
||||
) -> None:
|
||||
received_bodies.append(request_data["proxy_server_request"])
|
||||
|
||||
monkeypatch.setattr(proxy_server, "llm_router", None)
|
||||
monkeypatch.setattr(litellm, "callbacks", [_PromptMaskingGuardrail(), _SpendLogLikeHook()])
|
||||
proxy_logging.alert_types = []
|
||||
messages: Final = [{"role": "user", "content": "My name is John Smith"}]
|
||||
|
||||
await proxy_logging.post_call_failure_hook(
|
||||
request_data={
|
||||
"model": "internal-model",
|
||||
"messages": messages,
|
||||
"proxy_server_request": {
|
||||
"url": "http://localhost/chat/completions",
|
||||
"body": {"model": "internal-model", "messages": list(messages)},
|
||||
},
|
||||
},
|
||||
original_exception=HTTPException(status_code=429, detail="Rate limit exceeded"),
|
||||
user_api_key_dict=UserAPIKeyAuth(request_route="/chat/completions"),
|
||||
error_type=error_type,
|
||||
route="/chat/completions",
|
||||
)
|
||||
|
||||
assert received_bodies == [
|
||||
{
|
||||
"url": "http://localhost/chat/completions",
|
||||
"body": {"model": "internal-model", "messages": [{"role": "user", "content": logged_content}]},
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("maskers", "logged_content"),
|
||||
[((), "My name is John Smith"), ((_PromptMaskingGuardrail(),), "My name is <PERSON>")],
|
||||
ids=["without_a_masker", "with_a_masker"],
|
||||
)
|
||||
async def test_failure_hooks_get_the_request_body_the_failure_log_carried_after_the_router_retried_the_call(
|
||||
proxy_logging: ProxyLogging, monkeypatch: pytest.MonkeyPatch, maskers: tuple[CustomGuardrail, ...], logged_content: str
|
||||
) -> None:
|
||||
from litellm.proxy import proxy_server
|
||||
|
||||
received_bodies: Final[list[object]] = []
|
||||
|
||||
class _SpendLogLikeHook(CustomLogger):
|
||||
async def async_post_call_failure_hook(
|
||||
self,
|
||||
request_data: dict[str, object],
|
||||
original_exception: Exception,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
traceback_str: str | None = None,
|
||||
) -> None:
|
||||
received_bodies.append(request_data["proxy_server_request"])
|
||||
|
||||
monkeypatch.setattr(proxy_server, "llm_router", None)
|
||||
monkeypatch.setattr(litellm, "callbacks", [*maskers, _SpendLogLikeHook()])
|
||||
proxy_logging.alert_types = []
|
||||
messages: Final = [{"role": "user", "content": "My name is John Smith"}]
|
||||
router: Final = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "internal-model",
|
||||
"litellm_params": {
|
||||
"model": "openai/gpt-4.1",
|
||||
"api_key": "sk-test",
|
||||
"mock_response": "litellm.RateLimitError",
|
||||
},
|
||||
}
|
||||
],
|
||||
num_retries=1,
|
||||
retry_after=0,
|
||||
)
|
||||
logging_obj, _ = litellm.utils.function_setup(
|
||||
original_function="acompletion",
|
||||
rules_obj=litellm.utils.Rules(),
|
||||
start_time=datetime(2026, 1, 1),
|
||||
model="internal-model",
|
||||
messages=messages,
|
||||
litellm_call_id="retried-call",
|
||||
)
|
||||
with pytest.raises(litellm.RateLimitError):
|
||||
await router.acompletion(model="internal-model", messages=messages, litellm_logging_obj=logging_obj)
|
||||
|
||||
await proxy_logging.post_call_failure_hook(
|
||||
request_data={
|
||||
"model": "internal-model",
|
||||
"messages": messages,
|
||||
"litellm_logging_obj": logging_obj,
|
||||
"proxy_server_request": {"body": {"model": "internal-model", "messages": messages}},
|
||||
},
|
||||
original_exception=HTTPException(status_code=429, detail="Rate limit exceeded"),
|
||||
user_api_key_dict=UserAPIKeyAuth(request_route="/chat/completions"),
|
||||
route="/chat/completions",
|
||||
)
|
||||
|
||||
assert received_bodies == [
|
||||
{"body": {"model": "internal-model", "messages": [{"role": "user", "content": logged_content}]}}
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_call_failure_hook_attribution_does_not_count_against_the_deployment(
|
||||
proxy_logging, make_user_api_key_auth, monkeypatch
|
||||
|
|
|
|||
|
|
@ -52,6 +52,7 @@ def _make_logging_obj(
|
|||
logging_obj.get_chat_completion_prompt.return_value = prompt_return
|
||||
logging_obj.async_get_chat_completion_prompt = AsyncMock(return_value=prompt_return)
|
||||
logging_obj.async_failure_handler = AsyncMock()
|
||||
logging_obj.async_mask_failure_log_request = AsyncMock()
|
||||
logging_obj.model_call_details = {}
|
||||
return logging_obj
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue