fix(proxy): keep main's outcome where pass-through guardrails reached too far

- Drop a top-level guardrails or guardrail_config the merge can't combine,
  which returned 500 where main stripped the field.
- Count only a pre_call policy pipeline when withholding the Anthropic user_id.
- Read /bedrock invoke replies whole for generic guardrails, as on main.
- Keep the provider's body for logging when a hook returns a non-dict.
- Settle a no-cost blocked response at 0 so its failure row keeps its usage.
- Leave custom endpoints uncharged on a failed post-call hook, as on main.
- Keep a legacy Lakera string list's empty entries.
This commit is contained in:
yucheng 2026-11-03 04:06:39 -08:00
parent 1a5b56c9a2
commit 880f28451f
15 changed files with 422 additions and 21 deletions

View file

@ -2566,7 +2566,10 @@ class Logging(LiteLLMLoggingBaseClass):
**kwargs,
)
def prepare_response_cost(self, response: object, now: datetime.datetime) -> PreparedCost | None:
def prepare_response_cost(
self, response: object, now: datetime.datetime, settle_unknown_cost_at_zero: bool = False
) -> PreparedCost | None:
"""The response's cost, ready to finalize, or ``None`` when it has no cost result, unless that settles at 0"""
if self.prepared_cost is not None:
return self.prepared_cost
recorded_cache_hit: Final[object] = self.model_call_details.get("cache_hit")
@ -2592,8 +2595,10 @@ class Logging(LiteLLMLoggingBaseClass):
publish_logging_payload=False,
)
payload: Final = prepared[3]
raw_cost: Final[object] = self.model_call_details.get("response_cost")
if payload is None or not isinstance(raw_cost, (int, float)):
recorded_cost: Final[object] = self.model_call_details.get("response_cost")
unknown_cost: Final = 0.0 if settle_unknown_cost_at_zero else None
raw_cost: Final = recorded_cost if isinstance(recorded_cost, (int, float)) else unknown_cost
if payload is None or raw_cost is None:
return None
recorded_usage: Final[object] = self.model_call_details.get("combined_usage_object")
response_usage: Final[object] = getattr(prepared[2], "usage", None)

View file

@ -1416,6 +1416,7 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation):
litellm_logging_obj=litellm_logging_obj,
user_api_key_dict=user_api_key_dict,
request_data=request_data,
read_whole_response=True,
)
if not isinstance(response, dict):

View file

@ -233,6 +233,7 @@ class PassThroughEndpointHandler(BaseTranslation):
litellm_logging_obj: Optional["LiteLLMLoggingObj"] = None,
user_api_key_dict: Optional["UserAPIKeyAuth"] = None,
request_data: dict | None = None,
read_whole_response: bool = False,
) -> object:
"""
Process output response by applying guardrails to targeted fields.
@ -242,6 +243,7 @@ class PassThroughEndpointHandler(BaseTranslation):
guardrail_to_apply: The guardrail instance to apply
litellm_logging_obj: Optional logging object
user_api_key_dict: User API key metadata to pass to guardrails
read_whole_response: Without targeted fields, read the response as JSON, ids and usage included
"""
response_fields: Final = as_str_mapping(response)
if not isinstance(response, dict):
@ -261,7 +263,7 @@ class PassThroughEndpointHandler(BaseTranslation):
# Extract text to check
text_to_check: Final = (
self._extract_text_for_guardrail(response, field_expressions)
if field_expressions
if field_expressions or read_whole_response
else _response_text(response_fields or {})
)

View file

@ -235,7 +235,8 @@ class lakeraAI_Moderation(CustomGuardrail):
}
)
elif "input" in data and isinstance(data["input"], list):
text = "\n".join(filter(None, map(_input_item_text, data["input"])))
item_texts: Final = tuple(map(_input_item_text, _CONTENT_PARTS.validate_python(data["input"])))
text = "\n".join(item_texts) if any(item_texts) else ""
if not text:
verbose_proxy_logger.debug("Skipping lakera prompt injection, input has no text")
return

View file

@ -886,6 +886,17 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils):
_update_stream_param_based_on_request_body = update_stream_param_based_on_request_body
def _drop_unmergeable_guardrail_fields(
parsed_body: dict[str, object], # mutable-ok: strips the request body in place before the guardrail merge
) -> None:
"""Drop a ``guardrails`` that isn't a list or an object and a ``guardrail_config`` that isn't an object, which the
merge can't combine with the metadata's own. The provider receives neither field."""
if "guardrails" in parsed_body and not isinstance(parsed_body["guardrails"], list | dict):
parsed_body.pop("guardrails")
if "guardrail_config" in parsed_body and not isinstance(parsed_body["guardrail_config"], dict):
parsed_body.pop("guardrail_config")
async def _inherit_auth_guardrails(
parsed_body: dict[str, object] | None, # mutable-ok: merges the guardrails into the live request body
user_api_key_dict: UserAPIKeyAuth,
@ -901,6 +912,7 @@ async def _inherit_auth_guardrails(
"""
if configured_guardrails or not inherit or not parsed_body:
return configured_guardrails
_drop_unmergeable_guardrail_fields(parsed_body)
if not isinstance(parsed_body.get("metadata"), dict):
parsed_body["metadata"] = {} # rebind-ok: the merge below writes into the live request body's metadata
client_named: Final = (as_str_mapping(parsed_body["metadata"]) or {}).get("guardrails")
@ -1012,12 +1024,16 @@ def _provider_charge_settlement(
async def _run_post_call_guardrails_keeping_the_charge(
run_guardrails: Callable[[], Awaitable[_GuardedResponse]],
settle_provider_charge: Callable[[], Awaitable[None]],
settles_the_charge: bool = True,
) -> _GuardedResponse:
"""A guardrail that blocks the response, or fails, does so after the provider answered and billed the call, so the
charge is settled before the failure path writes its row, as on the native routes."""
charge is settled before the failure path writes its row, as on the native routes. A custom endpoint runs every
hook here, not only guardrails, and doesn't settle, so its failure row stays uncharged."""
try:
return await run_guardrails()
except Exception:
if not settles_the_charge:
raise
try:
await settle_provider_charge()
except Exception: # noqa: BLE001 - the guardrail's outcome still reaches the client
@ -1649,7 +1665,7 @@ async def pass_through_request(
)
hooked_body: Final = _json_object(_parsed_body)
client_user_id: Final = _anthropic_user_id_to_forward(
hooked_body, str(url), proxy_logging_obj.request_has_pre_call_guardrails
hooked_body, str(url), proxy_logging_obj.request_is_scanned_before_the_provider
)
kwargs = HttpPassThroughEndpointHelpers.init_kwargs_for_pass_through_endpoint( # rebind-ok: pre-existing rebinding on a rename-only line
user_api_key_dict=user_api_key_dict,
@ -2127,7 +2143,7 @@ async def pass_through_request(
inherit_auth_guardrails,
)
post_call_guardrail_data = hook_data
response_body = await _run_post_call_guardrails_keeping_the_charge(
guarded_body: Final = await _run_post_call_guardrails_keeping_the_charge(
partial(
proxy_logging_obj.post_call_success_hook,
data=hook_data,
@ -2145,14 +2161,16 @@ async def pass_through_request(
custom_llm_provider,
cast(dict[str, object], kwargs), # cast-ok: the logging fields init_kwargs built above
),
settles_the_charge=inherit_auth_guardrails,
)
if isinstance(response_body, dict):
if isinstance(guarded_body, dict):
response_body = guarded_body
content = json.dumps(response_body).encode("utf-8")
_content_modified = True
else:
verbose_proxy_logger.debug(
"pass_through_endpoint: post_call_success_hook returned %s, expected dict — using original response",
type(response_body).__name__,
type(guarded_body).__name__,
)
elif response_body is None:
verbose_proxy_logger.debug(

View file

@ -2,7 +2,7 @@ import asyncio
import json
from collections.abc import Mapping, Sequence
from contextlib import suppress
from datetime import datetime
from datetime import datetime, timezone
from types import MappingProxyType
from typing import Any, Final
from urllib.parse import urlparse
@ -85,9 +85,13 @@ async def finalize_cost_without_success_logs(
standard_logging_response_object: object,
**_success_log_fields: object, # kwargs-ok: accepts every field a log dispatch receives; none are needed here
) -> None:
"""Settles the provider's charge on the logging object for a response that won't be logged as a success"""
from litellm.proxy.spend_tracking.cost_finalization import finalize_proxy_response_cost
"""Settles the provider's charge on the logging object for a response that won't be logged as a success. A response
with no cost result settles at 0, the spend its success row would carry, so the failure row keeps its usage."""
from litellm.proxy.spend_tracking.cost_finalization import finalize_proxy_response_cost, prepare_response_cost
prepare_response_cost(
logging_obj, standard_logging_response_object, datetime.now(timezone.utc), settle_unknown_cost_at_zero=True
)
await finalize_proxy_response_cost(logging_obj, standard_logging_response_object)

View file

@ -11,9 +11,11 @@ from litellm.proxy.hooks.search_free_tier import SharedCounter, allocate_search_
from litellm.types.utils import CostBreakdown, StandardLoggingPayload
def prepare_response_cost(logging_obj: Logging, response: object, now: datetime) -> PreparedCost | None:
def prepare_response_cost(
logging_obj: Logging, response: object, now: datetime, settle_unknown_cost_at_zero: bool = False
) -> PreparedCost | None:
try:
return logging_obj.prepare_response_cost(response, now)
return logging_obj.prepare_response_cost(response, now, settle_unknown_cost_at_zero)
except Exception: # noqa: BLE001 # preparation must not prevent serving an already completed response
verbose_proxy_logger.exception("Unable to prepare completed response cost for %s", logging_obj.litellm_call_id)
return None

View file

@ -155,6 +155,7 @@ from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting
from litellm.integrations.SlackAlerting.utils import add_langfuse_trace_id_to_alert
from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_for_route
from litellm.litellm_core_utils.core_helpers import (
as_object_list,
coerce_token_limit,
get_or_create_metadata_bucket,
independent_snapshot,
@ -719,6 +720,20 @@ def pipeline_managed_guardrail_names(
)
def _runs_after_the_provider(pipeline_entry: object) -> bool:
if not isinstance(pipeline_entry, tuple):
return False
pipeline: Final = cast("tuple[object, ...]", pipeline_entry)[-1] # cast-ok: the policy engine's (name, pipeline)
return getattr(pipeline, "mode", None) == "post_call"
def _holds_pipeline_before_the_provider(pipelines: object) -> bool:
if not pipelines:
return False
entries: Final = as_object_list(pipelines)
return entries is None or not all(map(_runs_after_the_provider, entries))
def _ran_or_runs_outside_pipelines(
callback: CustomGuardrail | None,
data: Mapping[str, object],
@ -3108,8 +3123,18 @@ class ProxyLogging:
request_data: Mapping[str, object],
event_type: GuardrailEventHooks = GuardrailEventHooks.pre_call,
) -> bool:
if _policy_state_metadata(request_data).get("_guardrail_pipelines"):
return True
return bool(_policy_state_metadata(request_data).get("_guardrail_pipelines")) or (
self._request_has_guardrail_callback(request_data, event_type)
)
def request_is_scanned_before_the_provider(self, request_data: Mapping[str, object]) -> bool:
"""``request_has_pre_call_guardrails``, except that a resolved policy pipeline running post_call doesn't count"""
return _holds_pipeline_before_the_provider(
_policy_state_metadata(request_data).get("_guardrail_pipelines")
) or self._request_has_guardrail_callback(request_data, GuardrailEventHooks.pre_call)
@staticmethod
def _request_has_guardrail_callback(request_data: Mapping[str, object], event_type: GuardrailEventHooks) -> bool:
caps: Final = ProxyLogging.callback_capabilities()
if caps.has_content_enforcer:
return True

View file

@ -44,6 +44,11 @@ OPENAI_ROUTE: Final = "/openai/v1/chat/completions"
ANTHROPIC_ROUTE: Final = "/anthropic/v1/messages"
GEMINI_ROUTE: Final = "/gemini/v1beta/models/gemini-2.5-flash:generateContent"
CUSTOM_ROUTE: Final = "/passthrough-auth-custom"
CUSTOM_GUARDED_ROUTE: Final = "/passthrough-auth-custom-guarded"
UNPRICED_MODEL: Final = "passthrough-auth-unpriced-model"
MONITOR_POST_CALL: Final = "passthrough-auth-monitor-post-call"
MONITORED_PROMPT: Final = "reply-with-the-monitored-id"
MONITORED_REPLY_ID: Final = "msg_passthrough_auth_monitored"
PHONE_POST_CALL: Final = "passthrough-auth-phone-post-call"
PIPELINE_BLOCKS: Final = "passthrough-pipeline-blocks-the-response"
PIPELINE_ALLOWS: Final = "passthrough-pipeline-allows-the-response"
@ -94,6 +99,18 @@ def _openai_reply(_request: Request) -> Reply:
return Reply(body=_OPENAI_BODY)
def _lakera_pii_only(_request: Request) -> Reply:
return Reply(
body=json.dumps(
{
"flagged": True,
"breakdown": [{"detector_type": "pii/email", "detected": True, "message_id": 0}],
"payload": [],
}
).encode()
)
def _lakera_clean(_request: Request) -> Reply:
return Reply(body=json.dumps({"flagged": False, "breakdown": [], "payload": []}).encode())
@ -125,8 +142,20 @@ _UPSTREAM_BODIES: Final = MappingProxyType(
)
def _anthropic_reply(_request: Request) -> Reply:
return Reply(body=_ANTHROPIC_BODY)
_MONITORED_ANTHROPIC_BODY: Final = json.dumps(
_JSON.validate_json(message_json(MONITORED_REPLY_ID, "claude-sonnet-5-5", "Hi"))
).encode()
_UNPRICED_ANTHROPIC_BODY: Final = json.dumps(
_JSON.validate_json(message_json("msg_passthrough_auth_unpriced", UNPRICED_MODEL, "Hi"))
).encode()
def _anthropic_reply(request: Request) -> Reply:
if UNPRICED_MODEL.encode() in request.body:
return Reply(body=_UNPRICED_ANTHROPIC_BODY)
return Reply(body=_MONITORED_ANTHROPIC_BODY if MONITORED_PROMPT.encode() in request.body else _ANTHROPIC_BODY)
def _gemini_reply(_request: Request) -> Reply:
@ -153,6 +182,7 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I
key_post_call_vendor: Final = stack.enter_context(wire_server(_lakera_flagged))
generic_post_call_vendor: Final = stack.enter_context(wire_server(_verdict))
blocking_post_call_vendor: Final = stack.enter_context(wire_server(_block_every_reply))
monitor_post_call_vendor: Final = stack.enter_context(wire_server(_lakera_pii_only))
model_post_call_vendor: Final = stack.enter_context(wire_server(_verdict))
upstreams: Final = MappingProxyType(
{
@ -180,7 +210,14 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I
"path": CUSTOM_ROUTE,
"target": f"{upstreams[OPENAI_ROUTE].url}/v1/chat/completions",
"auth": True,
}
},
{
"path": CUSTOM_GUARDED_ROUTE,
"target": f"{upstreams[OPENAI_ROUTE].url}/v1/chat/completions",
"auth": True,
"cost_per_request": 0.25,
"guardrails": {BLOCKING_POST_CALL: None},
},
],
},
"policies": {
@ -252,6 +289,17 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I
"api_key": "synthetic-guardrail-key",
},
},
{
"guardrail_name": MONITOR_POST_CALL,
"litellm_params": {
"guardrail": "lakera_v2",
"mode": "post_call",
"default_on": False,
"on_flagged": "monitor",
"api_base": monitor_post_call_vendor.url,
"api_key": "synthetic-guardrail-key",
},
},
{
"guardrail_name": GENERIC_POST_CALL,
"litellm_params": {
@ -689,3 +737,100 @@ def test_response_a_post_call_guardrail_blocks_still_charges_the_provider_call(
assert [row["status"] for row in rows] == ["failure"], rows
assert float(str(rows[0]["spend"])) > 0, rows
assert int(str(rows[0]["total_tokens"])) > 0, rows
def _row_fields(gateway: Gateway, key: str) -> tuple[Mapping[str, object], ...]:
return tuple(
eventually(
lambda: gateway.rows(
'SELECT status, spend, total_tokens, api_base, request_id FROM "LiteLLM_SpendLogs" WHERE api_key = %s',
(key_hash(key),),
),
bool,
)
)
@pytest.mark.parametrize("route", [OPENAI_ROUTE, ANTHROPIC_ROUTE, GEMINI_ROUTE])
@pytest.mark.parametrize("value", [None, False, 7, "not-a-list"])
def test_a_body_guardrails_that_is_not_a_list_is_dropped_for_a_guarded_key(
rig: AuthGuardrailRig, route: str, value: JsonValue
) -> None:
gateway: Final = rig.owned.gateway
with gateway.scenario() as scenario:
key: Final = scenario.key(metadata={"guardrails": [GUARDRAIL]})
_drain(rig)
response: Final = _send(gateway, route, {**_route_body(route, "Hello"), "guardrails": value}, key)
assert response.status_code == 200, response.text
assert len(rig.vendor.drain()) == 1
assert [b'"guardrails"' in request.body for request in rig.upstreams[route].drain()] == [False]
def test_a_blocked_response_with_no_cost_result_keeps_its_usage_on_the_failure_row(rig: AuthGuardrailRig) -> None:
gateway: Final = rig.owned.gateway
with gateway.scenario() as scenario:
key: Final = scenario.key(metadata={"guardrails": [BLOCKING_POST_CALL]})
_drain(rig)
response: Final = _send(
gateway, ANTHROPIC_ROUTE, {**_route_body(ANTHROPIC_ROUTE, "Hello"), "model": UNPRICED_MODEL}, key
)
assert response.status_code == 400, response.text
assert len(rig.upstreams[ANTHROPIC_ROUTE].drain()) == 1
rows: Final = _row_fields(gateway, key)
assert [row["status"] for row in rows] == ["failure"], rows
assert float(str(rows[0]["spend"])) == 0, rows
assert int(str(rows[0]["total_tokens"])) > 0, rows
assert rows[0]["api_base"], rows
def test_custom_endpoint_writes_no_charge_for_a_response_its_post_call_guardrail_blocks(rig: AuthGuardrailRig) -> None:
gateway: Final = rig.owned.gateway
with gateway.scenario() as scenario:
key: Final = scenario.key(
team_id=scenario.team(), metadata={"allowed_passthrough_routes": [CUSTOM_GUARDED_ROUTE]}
)
_drain(rig)
response: Final = gateway.request("POST", CUSTOM_GUARDED_ROUTE, _route_body(OPENAI_ROUTE, "Hello"), key=key)
assert response.status_code == 400, response.text
assert len(rig.upstreams[OPENAI_ROUTE].drain()) == 1
assert [float(str(row["spend"])) for row in _row_fields(gateway, key)] == [0], response.text
def test_a_monitored_lakera_flag_keeps_the_provider_id_on_the_row(rig: AuthGuardrailRig) -> None:
gateway: Final = rig.owned.gateway
with gateway.scenario() as scenario:
key: Final = scenario.key(metadata={"guardrails": [MONITOR_POST_CALL]})
_drain(rig)
response: Final = _send(gateway, ANTHROPIC_ROUTE, _route_body(ANTHROPIC_ROUTE, MONITORED_PROMPT), key)
assert response.status_code == 200, response.text
assert response.content == _MONITORED_ANTHROPIC_BODY
assert [row["request_id"] for row in _row_fields(gateway, key)] == [MONITORED_REPLY_ID]
@pytest.mark.parametrize(
("policy", "forwarded"), [(PIPELINE_ALLOWS, True), (PIPELINE_BEFORE_THE_CALL, False)], ids=["post-call", "pre-call"]
)
def test_only_a_pipeline_that_runs_before_the_provider_withholds_the_client_user_id(
rig: AuthGuardrailRig, policy: str, forwarded: bool
) -> None:
gateway: Final = rig.owned.gateway
with gateway.scenario() as scenario:
key: Final = scenario.key()
_drain(rig)
body: Final = {
**_route_body(ANTHROPIC_ROUTE, "Hello"),
"metadata": {"tags": [policy], "user_id": "client-user-1"},
}
response: Final = _send(gateway, ANTHROPIC_ROUTE, body, key)
assert response.status_code == 200, response.text
assert [b"client-user-1" in request.body for request in rig.upstreams[ANTHROPIC_ROUTE].drain()] == [forwarded]

View file

@ -832,6 +832,28 @@ class TestBedrockPassthroughGuardrailHandlerOutput:
sent_texts = guardrail.apply_guardrail.call_args.kwargs["inputs"]["texts"]
assert "blocked model output" in sent_texts[0]
@pytest.mark.asyncio
async def test_non_converse_response_is_read_whole_with_its_envelope(self):
response = {
"id": "msg_marker",
"type": "message",
"role": "assistant",
"stop_reason": "end_turn",
"stop_sequence": "stop-marker",
"content": [{"type": "text", "text": "model output"}],
"usage": {"input_tokens": 5551234},
}
guardrail = _make_guardrail({"texts": []})
await BedrockPassthroughGuardrailHandler().process_output_response(
response=response,
guardrail_to_apply=guardrail,
request_data={"endpoint": "model/anthropic.claude-3-sonnet/invoke"},
)
sent_text = guardrail.apply_guardrail.call_args.kwargs["inputs"]["texts"][0]
assert all(part in sent_text for part in ("msg_marker", "stop-marker", "5551234", "model output"))
def _build_event_stream_frame(event_type: str, payload: dict) -> bytes:
import json

View file

@ -55,3 +55,11 @@ async def test_a_pass_through_body_with_no_messages_or_input_sends_nothing() ->
@pytest.mark.asyncio
async def test_a_native_body_with_no_messages_or_input_still_reaches_lakera() -> None:
assert await _sent_to_lakera({"prompt": "hello"}, call_type="text_completion") == ("",)
@pytest.mark.asyncio
@pytest.mark.parametrize(("strings", "text"), [(["a", "", "b"], "a\n\nb"), (["a", ""], "a\n")])
async def test_empty_strings_of_a_list_input_keep_their_place(strings: list[object], text: str) -> None:
sent: Final = await _sent_to_lakera({"input": strings})
assert [json.loads(body)["input"] for body in sent] == [text]

View file

@ -56,6 +56,8 @@ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
HttpPassThroughEndpointHelpers,
InitPassThroughEndpointHelpers,
SafeRouteAdder,
_anthropic_user_id_to_forward,
_inherit_auth_guardrails,
WebsocketFrameModelAllowed,
WebsocketFrameModelGate,
WebsocketFrameModelRefused,
@ -13328,6 +13330,70 @@ class TestInheritAuthGuardrails:
def _chat_body(**extra: object) -> dict[str, object]:
return {"model": "gpt-4o", "messages": [{"role": "user", "content": "hi"}], **extra}
@pytest.mark.asyncio
@pytest.mark.parametrize("value", [None, False, True, 7, "name"])
@pytest.mark.parametrize(
("key_metadata", "body_metadata"),
[({"guardrails": ["key-guard"]}, None), ({}, {"guardrails": ["body-guard"]}), ({}, None)],
ids=["key-guardrails", "metadata-guardrails", "no-guardrails"],
)
async def test_a_top_level_guardrails_the_merge_cant_add_to_is_dropped(self, value, key_metadata, body_metadata):
body: Final[dict[str, object]] = {
**self._chat_body(guardrails=value),
**({"metadata": body_metadata} if body_metadata is not None else {}),
}
with patch("litellm.proxy.proxy_server.premium_user", True):
await _inherit_auth_guardrails(body, UserAPIKeyAuth(api_key="sk-test", metadata=key_metadata), None, True)
assert "guardrails" not in body
assert (body["metadata"].get("guardrails") or []) == [
*key_metadata.get("guardrails", ()),
*(body_metadata or {}).get("guardrails", ()),
]
@pytest.mark.asyncio
@pytest.mark.parametrize("value", [None, False, 7, "config", ["config"]])
async def test_a_guardrail_config_the_merge_cant_add_to_is_dropped(self, value):
body: Final[dict[str, object]] = self._chat_body(
guardrail_config=value, metadata={"guardrail_config": {"kept": True}}
)
await _inherit_auth_guardrails(body, UserAPIKeyAuth(api_key="sk-test"), None, True)
assert "guardrail_config" not in body
assert body["metadata"]["guardrail_config"] == {"kept": True}
@pytest.mark.asyncio
@pytest.mark.parametrize(
("body_guardrails", "merged"),
[(["body-guard"], ["key-guard", "body-guard"]), ({"body-guard": {}}, ["key-guard", "body-guard"])],
ids=["list", "object"],
)
async def test_a_top_level_guardrails_list_or_object_still_merges(self, body_guardrails, merged):
body: Final[dict[str, object]] = self._chat_body(guardrails=body_guardrails)
with patch("litellm.proxy.proxy_server.premium_user", True):
await _inherit_auth_guardrails(
body, UserAPIKeyAuth(api_key="sk-test", metadata={"guardrails": ["key-guard"]}), None, True
)
assert body["metadata"]["guardrails"] == merged
@pytest.mark.parametrize(("mode", "forwarded"), [("post_call", "user-1"), ("pre_call", None)])
def test_only_a_pipeline_that_runs_before_the_provider_withholds_the_user_id(self, mode, forwarded):
pipeline: Final = GuardrailPipeline(mode=mode, steps=[PipelineStep(guardrail="guard")])
body: Final = {"model": "claude", "metadata": {"user_id": "user-1", "_guardrail_pipelines": [("p", pipeline)]}}
with patch.object(litellm, "callbacks", []):
user_id: Final = _anthropic_user_id_to_forward(
body,
"https://api.anthropic.com/v1/messages",
ProxyLogging(user_api_key_cache=DualCache()).request_is_scanned_before_the_provider,
)
assert user_id == forwarded
@pytest.mark.asyncio
async def test_key_team_and_project_guardrails_reach_pre_and_post_call(self):
auth: Final = UserAPIKeyAuth(

View file

@ -228,6 +228,54 @@ class TestPassthroughPostCallGuardrails:
assert body["error"]["guardrail_name"] == "rubrik"
assert body["error"]["model"] == "gemini-2.0-flash"
@patch(_COLLECT, return_value=["rubrik"])
async def test_configured_endpoint_writes_no_charge_for_a_failed_post_call_hook(self, mock_collect):
mock_proxy_logging = MagicMock()
mock_proxy_logging.pre_call_hook = AsyncMock(return_value={})
mock_proxy_logging.post_call_success_hook = AsyncMock(side_effect=HTTPException(status_code=418, detail="no"))
mock_proxy_logging.post_call_failure_hook = AsyncMock()
with (
_common_patches(mock_proxy_logging, _make_httpx_response(_GEMINI_RESPONSE)),
patch(f"{_PT_MOD}._blocked_response_cost_logging") as settlement,
pytest.raises(ProxyException),
):
await pass_through_request(
request=_make_mock_request(),
target="https://example.com/v1/generateContent",
custom_headers={"Content-Type": "application/json"},
user_api_key_dict=_make_user_api_key_dict(),
stream=False,
)
settlement.pass_through_async_success_handler.assert_not_called()
@patch(_COLLECT, return_value=["rubrik"])
async def test_a_hook_result_that_is_not_a_dict_leaves_the_body_the_logs_read(self, mock_collect):
mock_proxy_logging = MagicMock()
mock_proxy_logging.pre_call_hook = AsyncMock(return_value={})
mock_proxy_logging.post_call_success_hook = AsyncMock(return_value=MagicMock(id="rebuilt"))
mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value={})
success_handler = AsyncMock()
with (
_common_patches(mock_proxy_logging, _make_httpx_response(_GEMINI_RESPONSE)),
patch(
f"{_PT_MOD}.pass_through_endpoint_logging",
MagicMock(pass_through_async_success_handler=success_handler),
),
):
result = await pass_through_request(
request=_make_mock_request(),
target="https://example.com/v1/generateContent",
custom_headers={"Content-Type": "application/json"},
user_api_key_dict=_make_user_api_key_dict(),
stream=False,
)
assert success_handler.await_args.kwargs["response_body"] == _GEMINI_RESPONSE
assert json.loads(result.body) == _GEMINI_RESPONSE
@patch(_COLLECT, return_value=["rubrik"])
async def test_deny_forwards_guardrail_logging_info_to_failure_hook(
self,

View file

@ -1,6 +1,7 @@
import asyncio
from datetime import datetime, timezone
from typing import Final
from unittest.mock import patch
import pytest
@ -633,3 +634,19 @@ async def test_guardrail_json_rejection_retains_completed_provider_usage(
assert (payload["prompt_tokens"], payload["completion_tokens"]) == (10, 5)
assert payload["response_cost"] == 0.5
ProxyLogging._callback_capabilities_cache.clear()
def test_a_response_with_no_cost_result_settles_at_zero_only_when_asked(logging_obj: Logging) -> None:
from litellm.proxy.spend_tracking.cost_finalization import prepare_response_cost
unpriced: Final = response()
del unpriced._hidden_params["response_cost"]
with patch.object(logging_obj, "response_cost_calculator", return_value=None):
assert prepare_response_cost(logging_obj, unpriced, NOW) is None
prepared: Final = prepare_response_cost(logging_obj, unpriced, NOW, settle_unknown_cost_at_zero=True)
assert prepared is not None
assert prepared.amount == 0.0
assert prepared.payload["response_cost"] == 0.0
assert (prepared.usage.prompt_tokens, prepared.usage.completion_tokens) == (10, 5)

View file

@ -2,7 +2,7 @@
from __future__ import annotations
from typing import Any, Dict
from typing import Any, Dict, Final
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
@ -16,6 +16,7 @@ from litellm.integrations.custom_logger import CustomLogger
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.utils import ProxyLogging
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.proxy.policy_engine.pipeline_types import GuardrailPipeline, PipelineStep
from litellm.types.utils import CallTypesLiteral
@ -1014,3 +1015,39 @@ def test_request_has_pre_call_guardrails_reads_pipelines_from_both_metadata_buck
monkeypatch.setattr(litellm, "callbacks", [])
assert proxy_logging.request_has_pre_call_guardrails(request_data) is expected
@pytest.mark.parametrize(
("pipelines", "scanned"),
[
([("policy", GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="g")]))], False),
([("policy", GuardrailPipeline(mode="pre_call", steps=[PipelineStep(guardrail="g")]))], True),
(
[
("after", GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="g")])),
("before", GuardrailPipeline(mode="pre_call", steps=[PipelineStep(guardrail="g")])),
],
True,
),
(["p1"], True),
(5, True),
([], False),
],
ids=["post-call", "pre-call", "both", "unresolved-entry", "non-list", "empty"],
)
def test_only_a_pipeline_running_before_the_provider_scans_the_request(proxy_logging, monkeypatch, pipelines, scanned):
monkeypatch.setattr(litellm, "callbacks", [])
request_data: Final = {"metadata": {"_guardrail_pipelines": pipelines}}
assert proxy_logging.request_is_scanned_before_the_provider(request_data) is scanned
assert proxy_logging.request_has_pre_call_guardrails(request_data) is bool(pipelines)
def test_a_pre_call_guardrail_scans_the_request_beside_a_post_call_pipeline(proxy_logging, monkeypatch):
guardrail: Final = CustomGuardrail(guardrail_name="g", default_on=True, event_hook=GuardrailEventHooks.pre_call)
monkeypatch.setattr(litellm, "callbacks", [guardrail])
pipeline: Final = GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="other")])
assert proxy_logging.request_is_scanned_before_the_provider(
{"metadata": {"_guardrail_pipelines": [("p", pipeline)]}}
)