mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-12 08:42:18 +00:00
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:
parent
1a5b56c9a2
commit
880f28451f
15 changed files with 422 additions and 21 deletions
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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 {})
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)]}}
|
||||
)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue