diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 639003c9a6a..551e08f870e 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -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) diff --git a/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py b/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py index 3a65a8928bf..8cca405cd97 100644 --- a/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py +++ b/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py @@ -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): diff --git a/litellm/llms/pass_through/guardrail_translation/handler.py b/litellm/llms/pass_through/guardrail_translation/handler.py index ac91b7bd260..712e51f7afc 100644 --- a/litellm/llms/pass_through/guardrail_translation/handler.py +++ b/litellm/llms/pass_through/guardrail_translation/handler.py @@ -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 {}) ) diff --git a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py index 228dc417142..f93204f1bc9 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py @@ -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 diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index f377756b75a..7696d9fc344 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -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( diff --git a/litellm/proxy/pass_through_endpoints/success_handler.py b/litellm/proxy/pass_through_endpoints/success_handler.py index eb0ece24c0e..335b4889037 100644 --- a/litellm/proxy/pass_through_endpoints/success_handler.py +++ b/litellm/proxy/pass_through_endpoints/success_handler.py @@ -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) diff --git a/litellm/proxy/spend_tracking/cost_finalization.py b/litellm/proxy/spend_tracking/cost_finalization.py index c42579714e8..2db4b05f1e6 100644 --- a/litellm/proxy/spend_tracking/cost_finalization.py +++ b/litellm/proxy/spend_tracking/cost_finalization.py @@ -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 diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index eb36d1ae361..d4a2b6299e2 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -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 diff --git a/tests/integration/authorization/test_passthrough_auth_guardrails.py b/tests/integration/authorization/test_passthrough_auth_guardrails.py index c8b91077d91..690b018d8f3 100644 --- a/tests/integration/authorization/test_passthrough_auth_guardrails.py +++ b/tests/integration/authorization/test_passthrough_auth_guardrails.py @@ -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] diff --git a/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py b/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py index aec1672b335..02bc6e0d938 100644 --- a/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py +++ b/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py @@ -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 diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai.py index afba8f1b39d..bacb576d6c9 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai.py @@ -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] diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index f9a3f04e9e7..1fd7ee4d70c 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -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( diff --git a/tests/unit/proxy/pass_through_endpoints/test_passthrough_post_call_guardrails.py b/tests/unit/proxy/pass_through_endpoints/test_passthrough_post_call_guardrails.py index 0a750335fd6..667625cdb7d 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_passthrough_post_call_guardrails.py +++ b/tests/unit/proxy/pass_through_endpoints/test_passthrough_post_call_guardrails.py @@ -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, diff --git a/tests/unit/proxy/spend_tracking/test_cost_finalization.py b/tests/unit/proxy/spend_tracking/test_cost_finalization.py index c9515c1232e..3dffc02e7b3 100644 --- a/tests/unit/proxy/spend_tracking/test_cost_finalization.py +++ b/tests/unit/proxy/spend_tracking/test_cost_finalization.py @@ -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) diff --git a/tests/unit/proxy/utils/proxy_logging/test_pre_call_hook.py b/tests/unit/proxy/utils/proxy_logging/test_pre_call_hook.py index 26ee9f04d01..77542f1c3f4 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_pre_call_hook.py +++ b/tests/unit/proxy/utils/proxy_logging/test_pre_call_hook.py @@ -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)]}} + )