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 88

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

Browse source 
# Conflicts:
#	tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py
This commit is contained in:
mateo-berri 2026-09-08 12:45:27 -07:00
commit 6bfcfdbc50
10 changed files with 446 additions and 44 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

20
litellm/integrations/custom_guardrail.py
View file

@ -601,6 +601,12 @@ class CustomGuardrail(CustomLogger):
event_hook: GuardrailEventHooks | list[GuardrailEventHooks] | Mode | None,
supported_event_hooks: list[GuardrailEventHooks],
) -> None:
allowed_hooks: Final = frozenset(supported_event_hooks) | (
frozenset((GuardrailEventHooks.logging_only,))
if self.uses_apply_guardrail_interface() and not self.use_native_lifecycle_hooks
else frozenset()
)
def _validate_event_hook_list_is_in_supported_event_hooks(
event_hook: list[GuardrailEventHooks] | list[str],
supported_event_hooks: list[GuardrailEventHooks],
@ -608,7 +614,7 @@ class CustomGuardrail(CustomLogger):
for hook in event_hook:
if isinstance(hook, str):
hook = GuardrailEventHooks(hook)
if hook not in supported_event_hooks:
if hook not in allowed_hooks:
raise ValueError(f"Event hook {hook} is not in the supported event hooks {supported_event_hooks}")
if event_hook is None:
@ -629,7 +635,7 @@ class CustomGuardrail(CustomLogger):
default_list = event_hook.default if isinstance(event_hook.default, list) else [event_hook.default]
_validate_event_hook_list_is_in_supported_event_hooks(default_list, supported_event_hooks)
elif isinstance(event_hook, GuardrailEventHooks):
if event_hook not in supported_event_hooks:
if event_hook not in allowed_hooks:
raise ValueError(f"Event hook {event_hook} is not in the supported event hooks {supported_event_hooks}")
@staticmethod
@ -773,7 +779,7 @@ class CustomGuardrail(CustomLogger):
def uses_apply_guardrail_interface(self) -> bool:
return type(self).apply_guardrail is not CustomGuardrail.apply_guardrail
def _deployment_pre_call_target(self) -> "CustomLogger":
def _deployment_hook_target(self) -> "CustomLogger":
if not self.uses_apply_guardrail_interface() or self.use_native_lifecycle_hooks:
return self
try:
@ -802,7 +808,7 @@ class CustomGuardrail(CustomLogger):
# CHECK IF GUARDRAIL REJECTS THE REQUEST
if call_type == CallTypes.completion or call_type == CallTypes.acompletion:
target: Final = self._deployment_pre_call_target()
target: Final = self._deployment_hook_target()
if target is not self:
kwargs["guardrail_to_apply"] = self
result: Final = await target.async_pre_call_hook(
@ -845,7 +851,9 @@ class CustomGuardrail(CustomLogger):
return None
# CHECK IF GUARDRAIL REJECTS THE REQUEST
result: Final = await self.async_post_call_success_hook(
target: Final = self._deployment_hook_target()
hook_request_data: Final = {**request_data, "guardrail_to_apply": self} if target is not self else request_data
result: Final = await target.async_post_call_success_hook(
user_api_key_dict=UserAPIKeyAuth(
user_id=request_data.get("user_api_key_user_id"),
team_id=request_data.get("user_api_key_team_id"),
@ -853,7 +861,7 @@ class CustomGuardrail(CustomLogger):
api_key=request_data.get("user_api_key_hash"),
request_route=request_data.get("user_api_key_request_route"),
),
data=request_data,
data=hook_request_data,
response=response,
)

2
litellm/llms/bedrock/embed/embedding.py
View file

@ -51,7 +51,7 @@ def _sign_get_request(
except ImportError:
raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.")
request: Final = AWSRequest(method="GET", url=url, data=None, headers=dict(headers))
request: Final = AWSRequest(method="GET", url=url, data=None, headers=headers)
SigV4Auth(credentials, "bedrock", aws_region_name).add_auth(request)
return request.prepare()

12
litellm/llms/custom_httpx/llm_http_handler.py
View file

@ -1742,6 +1742,12 @@ class BaseLLMHTTPHandler:
except Exception as e:
raise self._handle_error(e=e, provider_config=provider_config)
logging_obj.post_call(
api_key=api_key,
original_response=response.text,
additional_args={"complete_input_dict": data},
)
return self._transform_ocr_response(
provider_config=provider_config,
model=model,
@ -1805,6 +1811,12 @@ class BaseLLMHTTPHandler:
except Exception as e:
raise self._handle_error(e=e, provider_config=provider_config)
logging_obj.post_call(
api_key=api_key,
original_response=response.text,
additional_args={"complete_input_dict": data},
)
# Use async response transform for async operations
return await provider_config.async_transform_ocr_response(
model=model,

2
litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py
View file

@ -1192,7 +1192,7 @@ def _sign_aws_json_post(
except ImportError:
raise ImportError(f"Missing boto3 to call {service_name}. Run 'pip install boto3'.")
aws_request: Final = AWSRequest(method="POST", url=url, data=body, headers=dict(headers))
aws_request: Final = AWSRequest(method="POST", url=url, data=body, headers=headers)
SigV4Auth(get_credentials(), service_name, aws_region_name).add_auth(aws_request)
return aws_request.prepare()

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

@ -6,7 +6,7 @@ pass/fail actions (allow, block, next, modify_response) and data forwarding.
"""
import time
from collections.abc import Sequence
from collections.abc import Mapping, Sequence
from typing import Any, Final, Literal
import litellm
@ -16,7 +16,11 @@ from litellm.integrations.custom_guardrail import (
ModifyResponseException,
)
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.core_helpers import independent_snapshot
from litellm.litellm_core_utils.core_helpers import (
get_metadata_variable_name_from_kwargs,
get_or_create_metadata_bucket,
independent_snapshot,
)
from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import (
UnifiedLLMGuardrails,
)
@ -25,6 +29,7 @@ from litellm.types.proxy.policy_engine.pipeline_types import (
PipelineStep,
PipelineStepResult,
)
from litellm.types.utils import StandardLoggingGuardrailInformation
try:
from fastapi.exceptions import HTTPException
@ -118,6 +123,7 @@ class PipelineExecutor:
return _allow_result(step_results=step_results, working_data=working_data, request_data=data)
if action == "block":
_carry_working_guardrail_information(working_data=working_data, request_data=data)
return PipelineExecutionResult(
terminal_action="block",
step_results=step_results,
@ -126,6 +132,7 @@ class PipelineExecutor:
)
if action == "modify_response":
_carry_working_guardrail_information(working_data=working_data, request_data=data)
return PipelineExecutionResult(
terminal_action="modify_response",
step_results=step_results,
@ -168,34 +175,33 @@ class PipelineExecutor:
verbose_proxy_logger.warning("Pipeline: guardrail '%s' not found in callbacks", step.guardrail)
return ("error", None, f"Guardrail '{step.guardrail}' not found", None)
# Inject guardrail name into metadata so should_run_guardrail() allows it
if "metadata" not in data:
data["metadata"] = {}
data["metadata"]["guardrails"] = [step.guardrail]
# A scan_raw_request step evaluates the pristine pre-pipeline
# snapshot instead of `data` (which earlier pass_data steps in
# this same pipeline may have already rewritten), same reason
# the normal sequential/parallel guardrail loops do this.
scans_raw_request: Final = callback.scan_raw_request
hook_input: Final[dict] = ( # mutable-ok: same request-payload shape as data
independent_snapshot(raw_request_snapshot)
if scans_raw_request and raw_request_snapshot is not None
else data
)
if hook_input is not data:
hook_input.setdefault("metadata", {})["guardrails"] = [step.guardrail]
snapshot_entries_before: Final = len(_recorded_guardrail_information(hook_input))
# Use unified_guardrail path if callback implements apply_guardrail
target: CustomLogger = callback
use_unified: Final = "apply_guardrail" in type(callback).__dict__ and not callback.use_native_lifecycle_hooks
if use_unified:
hook_input["guardrail_to_apply"] = callback
target = UnifiedLLMGuardrails()
try:
# Inject guardrail name into metadata so should_run_guardrail() allows it
if "metadata" not in data:
data["metadata"] = {}
data["metadata"]["guardrails"] = [step.guardrail]
# A scan_raw_request step evaluates the pristine pre-pipeline
# snapshot instead of `data` (which earlier pass_data steps in
# this same pipeline may have already rewritten), same reason
# the normal sequential/parallel guardrail loops do this.
scans_raw_request: Final = callback.scan_raw_request
hook_input: Final[dict] = ( # mutable-ok: same request-payload shape as data
independent_snapshot(raw_request_snapshot)
if scans_raw_request and raw_request_snapshot is not None
else data
)
if hook_input is not data:
hook_input.setdefault("metadata", {})["guardrails"] = [step.guardrail]
# Use unified_guardrail path if callback implements apply_guardrail
target: CustomLogger = callback
use_unified: Final = (
"apply_guardrail" in type(callback).__dict__ and not callback.use_native_lifecycle_hooks
)
if use_unified:
hook_input["guardrail_to_apply"] = callback
target = UnifiedLLMGuardrails()
if mode == "pre_call":
response = await target.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
@ -233,6 +239,12 @@ class PipelineExecutor:
else:
verbose_proxy_logger.error("Pipeline: unexpected error from guardrail '%s': %s", step.guardrail, e)
return ("error", None, str(e), e)
finally:
if hook_input is not data:
_append_guardrail_information(
request_data=data,
entries=_recorded_guardrail_information(hook_input)[snapshot_entries_before:],
)
@staticmethod
def find_guardrail_callback(guardrail_name: str) -> CustomGuardrail | None:
@ -283,6 +295,40 @@ def _restore_request_guardrails(
return {**working_data, "metadata": stripped} # mutable-ok: request dict
_GUARDRAIL_INFORMATION_KEY: Final = "standard_logging_guardrail_information"
def _recorded_guardrail_information(source: Mapping[str, object]) -> list[StandardLoggingGuardrailInformation]:
bucket: Final = source.get(get_metadata_variable_name_from_kwargs(source))
recorded: Final = bucket.get(_GUARDRAIL_INFORMATION_KEY) if isinstance(bucket, dict) else None
return recorded if isinstance(recorded, list) else []
def _append_guardrail_information(
request_data: dict[str, object], # mutable-ok: same request-payload shape as execute_steps' data
entries: Sequence[StandardLoggingGuardrailInformation],
) -> None:
if not entries:
return
_, request_bucket = get_or_create_metadata_bucket(request_data)
existing: Final = request_bucket.get(_GUARDRAIL_INFORMATION_KEY)
if isinstance(existing, list):
existing.extend(entries)
return
request_bucket[_GUARDRAIL_INFORMATION_KEY] = list(entries)
def _carry_working_guardrail_information(
working_data: Mapping[str, object],
request_data: dict[str, object], # mutable-ok: same request-payload shape as execute_steps' data
) -> None:
recorded: Final = _recorded_guardrail_information(working_data)
existing: Final = _recorded_guardrail_information(request_data)
if recorded is existing:
return
_append_guardrail_information(request_data=request_data, entries=[e for e in recorded if e not in existing])
def _pipeline_action_for_outcome(step: PipelineStep, outcome: str) -> str:
"""
Map pipeline step outcome to the configured action.

2
tests/test_litellm/enterprise/enterprise_callbacks/test_secret_detection.py
View file

@ -262,7 +262,7 @@ def test_proxied_traffic_stays_on_native_hooks():
never sees ``data["prompt"]``."""
guardrail = _guardrail()
assert guardrail.uses_apply_guardrail_interface() is True
assert guardrail._deployment_pre_call_target() is guardrail
assert guardrail._deployment_hook_target() is guardrail
@pytest.mark.asyncio

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

@ -1,5 +1,5 @@
import asyncio
from typing import TYPE_CHECKING, Literal, Optional
from typing import TYPE_CHECKING, ClassVar, Final, Literal, Optional
from unittest.mock import AsyncMock
import pytest
@ -10,6 +10,7 @@ from litellm.integrations.custom_guardrail import (
log_guardrail_information,
)
from litellm.proxy._types import CallTypes, UserAPIKeyAuth
from litellm.types.guardrails import GuardrailEventHooks, Mode
from litellm.types.utils import GenericGuardrailAPIInputs, GuardrailTracingDetail
if TYPE_CHECKING:
@ -2378,11 +2379,108 @@ def _logged_call(messages: list | str) -> tuple[dict, object]:
return kwargs, response
class _NativeApplyGuardrail(_InheritedApplyGuardrail):
use_native_lifecycle_hooks: ClassVar[bool] = True
@pytest.mark.parametrize("guardrail_type", (CustomGuardrail, _NativeApplyGuardrail, _InheritedApplyGuardrail))
@pytest.mark.parametrize(
"event_hook",
(
GuardrailEventHooks.logging_only,
"logging_only",
[GuardrailEventHooks.pre_call, GuardrailEventHooks.logging_only],
["pre_call", "logging_only"],
Mode(tags={"audit": "logging_only"}, default="pre_call"),
Mode(tags={"audit": ["pre_call", "logging_only"]}),
Mode(tags={"enforce": "pre_call"}, default="logging_only"),
Mode(tags={}, default=["pre_call", "logging_only"]),
),
)
def test_logging_only_requires_framework_support_or_explicit_declaration(
guardrail_type: type[CustomGuardrail],
event_hook: GuardrailEventHooks | str | list[GuardrailEventHooks] | list[str] | Mode,
) -> None:
supported: Final = [GuardrailEventHooks.pre_call]
if guardrail_type is _InheritedApplyGuardrail:
guardrail: Final = guardrail_type(event_hook=event_hook, supported_event_hooks=supported)
assert guardrail.event_hook == event_hook
assert supported == [GuardrailEventHooks.pre_call]
else:
with pytest.raises(ValueError, match=r"logging_only.*not in the supported event hooks"):
guardrail_type(event_hook=event_hook, supported_event_hooks=supported)
explicitly_supported: Final = guardrail_type(
event_hook=event_hook,
supported_event_hooks=[GuardrailEventHooks.pre_call, GuardrailEventHooks.logging_only],
)
assert explicitly_supported.event_hook == event_hook
@pytest.mark.parametrize(
"event_hook",
(
GuardrailEventHooks.post_call,
"post_call",
[GuardrailEventHooks.logging_only, GuardrailEventHooks.post_call],
["logging_only", "post_call"],
Mode(tags={"enforce": "post_call"}, default="logging_only"),
Mode(tags={"enforce": ["logging_only", "post_call"]}),
Mode(tags={"audit": "logging_only"}, default="post_call"),
Mode(tags={}, default=["logging_only", "post_call"]),
),
)
def test_framework_logging_only_does_not_allow_other_unsupported_modes(
event_hook: GuardrailEventHooks | str | list[GuardrailEventHooks] | list[str] | Mode,
) -> None:
with pytest.raises(ValueError, match=r"post_call.*not in the supported event hooks"):
_InheritedApplyGuardrail(event_hook=event_hook, supported_event_hooks=[GuardrailEventHooks.pre_call])
class TestLoggingOnlyApplyGuardrail:
"""LIT-4876 regression: a guardrail in mode logging_only that implements only
apply_guardrail must still run against the logged request and response and
record guardrail_information, instead of inheriting the CustomLogger no-op."""
@pytest.mark.parametrize(
"event_hook",
(
GuardrailEventHooks.logging_only,
"logging_only",
[GuardrailEventHooks.pre_call, GuardrailEventHooks.logging_only],
["pre_call", "logging_only"],
Mode(tags={"audit": "logging_only"}, default="pre_call"),
Mode(tags={"audit": ["pre_call", "logging_only"]}),
Mode(tags={"enforce": "pre_call"}, default="logging_only"),
Mode(tags={}, default=["pre_call", "logging_only"]),
),
)
@pytest.mark.asyncio
async def test_content_filter_accepts_logging_only_and_records_detection(
self, event_hook: GuardrailEventHooks | str | list[GuardrailEventHooks] | list[str] | Mode
) -> None:
from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import (
ContentFilterGuardrail,
)
from litellm.types.guardrails import BlockedWord, ContentFilterAction, GuardrailEventHooks
guardrail: Final = ContentFilterGuardrail(
guardrail_name="content-review",
event_hook=event_hook,
default_on=True,
blocked_words=[BlockedWord(keyword="hello", action=ContentFilterAction.BLOCK)],
)
kwargs, response = _logged_call([{"role": "user", "content": "hello there"}])
out_kwargs, out_response = await guardrail.async_logging_hook(kwargs, response, CallTypes.acompletion.value)
assert out_response is response
assert out_kwargs["messages"] == kwargs["messages"]
assert (
out_kwargs["standard_logging_object"]["guardrail_information"][0]["guardrail_status"]
== "guardrail_intervened"
)
@pytest.mark.asyncio
async def test_runs_apply_guardrail_observe_only_and_records_verdict(self):
guardrail = _ApplyOnlyObserver()
@ -2610,3 +2708,36 @@ class TestCustomGuardrailPostCallSuccessDeploymentHook:
)
assert result is replacement
@pytest.mark.asyncio
async def test_apply_guardrail_interface_modifies_deployment_response(self):
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.utils import ModelResponse
class ReplacingGuardrail(CustomGuardrail):
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict[str, object],
input_type: Literal["request", "response"],
logging_obj: Optional["LiteLLMLoggingObj"] = None,
) -> GenericGuardrailAPIInputs:
assert input_type == "response"
return {**inputs, "texts": ["filtered response"]}
guardrail = ReplacingGuardrail(
guardrail_name="test-guardrail",
event_hook=GuardrailEventHooks.post_call,
)
response = ModelResponse(choices=[{"message": {"role": "assistant", "content": "original response"}}])
request_data = {"guardrails": ["test-guardrail"]}
result = await guardrail.async_post_call_success_deployment_hook(
request_data=request_data,
response=response,
call_type=CallTypes.acompletion,
)
assert result is response
assert response.choices[0].message.content == "filtered response"
assert request_data == {"guardrails": ["test-guardrail"]}

64
tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py
View file

@ -32,6 +32,7 @@ from litellm.llms.azure.videos.transformation import AzureVideoConfig
from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import (
AmazonAnthropicClaudeMessagesConfig,
)
from litellm.llms.mistral.ocr.transformation import MistralOCRConfig
from litellm.llms.openai.videos.transformation import OpenAIVideoConfig
from litellm.types.llms.openai import ResponsesAPIResponse
from litellm.types.router import GenericLiteLLMParams
@ -41,6 +42,69 @@ from tests.test_litellm.llms.bedrock.event_loop_probe import EventLoopProbe
_ACTIVE_KEY = "_code_interpreter_interception_active"
_SANDBOX_KEY = "_code_interpreter_interception_sandbox_key"
OCR_RESPONSE = {
"pages": [{"index": 0, "markdown": "OCR output", "images": []}],
"model": "mistral-ocr-latest",
"usage_info": {"pages_processed": 1},
}
def _ocr_sync_client() -> HTTPHandler:
client = HTTPHandler()
client.client = httpx.Client(transport=httpx.MockTransport(lambda _request: httpx.Response(200, json=OCR_RESPONSE)))
return client
def _ocr_async_client() -> AsyncHTTPHandler:
client = AsyncHTTPHandler()
client.client = httpx.AsyncClient(
transport=httpx.MockTransport(lambda _request: httpx.Response(200, json=OCR_RESPONSE))
)
return client
def test_ocr_calls_post_call_with_raw_provider_response():
logging_obj = Mock()
response = BaseLLMHTTPHandler().ocr(
model="mistral-ocr-latest",
document={"type": "document_url", "document_url": "https://example.com/document.pdf"},
optional_params={},
timeout=5,
logging_obj=logging_obj,
api_key="test-key",
api_base="https://api.mistral.ai/v1/ocr",
custom_llm_provider="mistral",
client=_ocr_sync_client(),
provider_config=MistralOCRConfig(),
)
assert response.pages[0].markdown == "OCR output"
logging_obj.post_call.assert_called_once()
assert json.loads(logging_obj.post_call.call_args.kwargs["original_response"]) == OCR_RESPONSE
@pytest.mark.asyncio
async def test_async_ocr_calls_post_call_with_raw_provider_response():
logging_obj = Mock()
response = await BaseLLMHTTPHandler().async_ocr(
model="mistral-ocr-latest",
document={"type": "document_url", "document_url": "https://example.com/document.pdf"},
optional_params={},
timeout=5,
logging_obj=logging_obj,
api_key="test-key",
api_base="https://api.mistral.ai/v1/ocr",
custom_llm_provider="mistral",
client=_ocr_async_client(),
provider_config=MistralOCRConfig(),
)
assert response.pages[0].markdown == "OCR output"
logging_obj.post_call.assert_called_once()
assert json.loads(logging_obj.post_call.call_args.kwargs["original_response"]) == OCR_RESPONSE
def test_prepare_fake_stream_request():
# Initialize the BaseLLMHTTPHandler

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

@ -4,20 +4,26 @@ Tests for the pipeline executor.
Uses mock guardrails to validate pipeline execution without external services.
"""
import copy
from typing import Literal
from unittest.mock import MagicMock
import pytest
import litellm
from litellm.caching.dual_cache import DualCache
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.guardrails.guardrail_hooks.custom_code.custom_code_guardrail import (
CustomCodeGuardrail,
)
from litellm.proxy.policy_engine.pipeline_executor import PipelineExecutor
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.proxy.policy_engine.pipeline_types import (
GuardrailPipeline,
PipelineStep,
)
from litellm.types.utils import CallTypesLiteral
try:
from fastapi.exceptions import HTTPException
@ -158,11 +164,146 @@ class ContentCheckGuardrail(CustomGuardrail):
return None
class RecordingGuardrail(CustomGuardrail):
def __init__(self, guardrail_name: str, scan_raw_request: bool = False, block: bool = True):
super().__init__(
guardrail_name=guardrail_name,
event_hook="pre_call",
default_on=True,
scan_raw_request=scan_raw_request,
)
self.block = block
def should_run_guardrail(self, data: dict[str, object], event_type: GuardrailEventHooks) -> bool:
return True
async def async_pre_call_hook(
self,
user_api_key_dict: UserAPIKeyAuth,
cache: DualCache,
data: dict[str, object],
call_type: CallTypesLiteral,
) -> dict[str, object]:
self.add_standard_logging_guardrail_information_to_request_data(
guardrail_json_response={"detected": ["aws_access_key"]},
request_data=data,
guardrail_status="guardrail_intervened" if self.block else "success",
)
if self.block:
raise HTTPException(status_code=400, detail="Content policy violation")
return copy.deepcopy(data)
# ─────────────────────────────────────────────────────────────────────────────
# Tests
# ─────────────────────────────────────────────────────────────────────────────
@pytest.mark.skipif(HTTPException is None, reason="fastapi not installed")
@pytest.mark.asyncio
@pytest.mark.parametrize("scan_raw_request", [False, True])
@pytest.mark.parametrize("on_fail", ["block", "modify_response"])
async def test_terminal_block_carries_guardrail_information_to_request(
monkeypatch: pytest.MonkeyPatch, scan_raw_request: bool, on_fail: Literal["block", "modify_response"]
):
"""
Spend logging and the Guardrails Monitor read standard_logging_guardrail_information
off the caller's request dict. A blocking step records it on the executor's
working copy (or the raw-request snapshot), so the terminal result must carry it
back onto the request or the block is never counted.
"""
guard = RecordingGuardrail(guardrail_name="credentials-api-keys", scan_raw_request=scan_raw_request)
monkeypatch.setattr(litellm, "callbacks", [guard])
data = {
"messages": [{"role": "user", "content": "key AKIAIOSFODNN7EXAMPLE"}],
"metadata": {"user_api_key_hash": "abc"},
}
result = await PipelineExecutor.execute_steps(
steps=[PipelineStep(guardrail="credentials-api-keys", on_fail=on_fail, on_pass="next")],
mode="pre_call",
data=data,
user_api_key_dict=MagicMock(),
call_type="completion",
policy_name="baseline-pii-protection",
raw_request_snapshot={"messages": data["messages"], "metadata": {"user_api_key_hash": "abc"}},
)
assert result.terminal_action == on_fail
recorded = data["metadata"]["standard_logging_guardrail_information"]
assert [entry["guardrail_name"] for entry in recorded] == ["credentials-api-keys"]
assert recorded[0]["guardrail_status"] == "guardrail_intervened"
assert data["metadata"]["user_api_key_hash"] == "abc"
assert "guardrails" not in data["metadata"]
@pytest.mark.skipif(HTTPException is None, reason="fastapi not installed")
@pytest.mark.asyncio
async def test_terminal_block_merges_guardrail_information_without_duplicates(monkeypatch: pytest.MonkeyPatch):
"""A pass_data step that returns a rewritten copy of the request, and a scan_raw_request step
that evaluates a deep copy taken before the pipeline ran, both leave earlier entries in two
dicts at once. Those must be carried back once while every step's own entry is kept."""
first = RecordingGuardrail(guardrail_name="pii-scan", block=False)
second = RecordingGuardrail(guardrail_name="credentials-api-keys", scan_raw_request=True)
monkeypatch.setattr(litellm, "callbacks", [first, second])
earlier = {"guardrail_name": "earlier-guard", "guardrail_status": "success"}
data = {"messages": [{"role": "user", "content": "hi"}], "metadata": {}}
data["metadata"]["standard_logging_guardrail_information"] = [earlier]
result = await PipelineExecutor.execute_steps(
steps=[
PipelineStep(guardrail="pii-scan", on_fail="block", on_pass="next", pass_data=True),
PipelineStep(guardrail="credentials-api-keys", on_fail="block", on_pass="next"),
],
mode="pre_call",
data=data,
user_api_key_dict=MagicMock(),
call_type="completion",
policy_name="baseline-pii-protection",
raw_request_snapshot={
"messages": data["messages"],
"metadata": {"standard_logging_guardrail_information": [dict(earlier)]},
},
)
assert result.terminal_action == "block"
recorded = data["metadata"]["standard_logging_guardrail_information"]
assert [entry["guardrail_name"] for entry in recorded] == ["earlier-guard", "pii-scan", "credentials-api-keys"]
@pytest.mark.skipif(HTTPException is None, reason="fastapi not installed")
@pytest.mark.asyncio
async def test_repeated_scan_raw_request_step_is_counted_once_per_evaluation(monkeypatch: pytest.MonkeyPatch):
"""Running the same raw-scan guardrail twice yields two identical entries; both must reach the caller,
while the entries the raw snapshot already held before the pipeline ran are not copied again."""
guard = RecordingGuardrail(guardrail_name="credentials-raw", scan_raw_request=True, block=False)
monkeypatch.setattr(litellm, "callbacks", [guard])
earlier = {"guardrail_name": "earlier-guard", "guardrail_status": "success"}
data = {"messages": [{"role": "user", "content": "hi"}], "metadata": {}}
data["metadata"]["standard_logging_guardrail_information"] = [earlier]
result = await PipelineExecutor.execute_steps(
steps=[
PipelineStep(guardrail="credentials-raw", on_fail="block", on_pass="next"),
PipelineStep(guardrail="credentials-raw", on_fail="block", on_pass="next"),
],
mode="pre_call",
data=data,
user_api_key_dict=MagicMock(),
call_type="completion",
policy_name="raw-scan-policy",
raw_request_snapshot={
"messages": data["messages"],
"metadata": {"standard_logging_guardrail_information": [dict(earlier)]},
},
)
assert result.terminal_action == "allow"
assert result.modified_data is not None
recorded = result.modified_data["metadata"]["standard_logging_guardrail_information"]
assert [entry["guardrail_name"] for entry in recorded] == ["earlier-guard", "credentials-raw", "credentials-raw"]
@pytest.mark.skipif(HTTPException is None, reason="fastapi not installed")
@pytest.mark.asyncio
async def test_escalation_step1_fails_step2_blocks(monkeypatch):

10
tests/test_litellm/proxy/test_proxy_logging_hook_detection.py
View file

@ -637,14 +637,14 @@ def test_callback_capabilities_excludes_opted_out_guardrail_from_iterator_overri
assert [cb for cb, _ in caps.iterator_overrides if cb is opted_out] == []
def test_deployment_pre_call_target_stays_native_when_opted_out():
def test_deployment_hook_target_stays_native_when_opted_out():
"""Model-level guardrails resolve their target here rather than through ProxyLogging."""
assert _KeepsNativeHooks()._deployment_pre_call_target() is not None
assert _KeepsNativeHooks()._deployment_hook_target() is not None
opted_out = _KeepsNativeHooks()
assert opted_out._deployment_pre_call_target() is opted_out
assert _AppliesGuardrail()._deployment_pre_call_target() is not None
assert opted_out._deployment_hook_target() is opted_out
assert _AppliesGuardrail()._deployment_hook_target() is not None
routed = _AppliesGuardrail()
assert routed._deployment_pre_call_target() is not routed
assert routed._deployment_hook_target() is not routed
@pytest.mark.asyncio