From 1e59cd9e3561be8df396619fcf3e6f87a89a5e57 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 17 Jul 2026 14:14:54 -0400 Subject: [PATCH 01/60] fix(utils): honor string drop_params values from config and DB deployments --- litellm/litellm_core_utils/core_helpers.py | 12 +++++++ litellm/types/router.py | 7 ++++ litellm/utils.py | 9 +++-- .../litellm_core_utils/test_core_helpers.py | 22 ++++++++++++ tests/test_litellm/test_router.py | 34 +++++++++++++++++++ tests/test_litellm/test_utils.py | 28 +++++++++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 4 +++ 7 files changed, 113 insertions(+), 3 deletions(-) diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index 002a46771e3..838019264c9 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -36,6 +36,18 @@ def safe_divide_seconds(seconds: float, denominator: float, default: Optional[fl return float(seconds / denominator) +def normalize_drop_params(value: object) -> bool | None: + if isinstance(value, bool): + return value + if isinstance(value, str): + lowered = value.strip().lower() + if lowered == "true": + return True + if lowered == "false": + return False + return None + + def safe_divide( numerator: Union[int, float], denominator: Union[int, float], diff --git a/litellm/types/router.py b/litellm/types/router.py index 28e4a8272e8..77d1815d28d 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -23,6 +23,7 @@ from pydantic import BaseModel, ConfigDict, Field, field_validator, model_valida from typing_extensions import Protocol, Required, TypedDict, runtime_checkable from litellm._uuid import uuid +from litellm.litellm_core_utils.core_helpers import normalize_drop_params from .completion import CompletionRequest from .embedding import EmbeddingRequest @@ -233,6 +234,7 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams): None # timeout when making stream=True calls, if str, pass in as os.environ/ ) max_retries: Optional[int] = None + drop_params: Optional[bool] = None organization: Optional[str] = None # for openai orgs configurable_clientside_auth_params: CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS = None litellm_credential_name: Optional[str] = None @@ -311,6 +313,11 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams): return filtered return data + @field_validator("drop_params", mode="before") + @classmethod + def coerce_drop_params(cls, value: object) -> Optional[bool]: + return normalize_drop_params(value) + def __contains__(self, key): # Define custom behavior for the 'in' operator return hasattr(self, key) diff --git a/litellm/utils.py b/litellm/utils.py index e19d2b36a52..e10b63ceb37 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -60,6 +60,7 @@ from litellm._lazy_imports import ( _get_token_counter_new, ) from litellm._uuid import uuid +from litellm.litellm_core_utils.core_helpers import normalize_drop_params from litellm.litellm_core_utils.fallback_generalizations import ( match_capability_generalizations, ) @@ -2852,7 +2853,7 @@ def get_optional_params_transcription( passed_params.pop("OPENAI_TRANSCRIPTION_PARAMS") custom_llm_provider = passed_params.pop("custom_llm_provider") - drop_params = passed_params.pop("drop_params") + drop_params = normalize_drop_params(passed_params.pop("drop_params")) special_params = passed_params.pop("kwargs") for k, v in special_params.items(): passed_params[k] = v @@ -2960,7 +2961,7 @@ def get_optional_params_image_gen( model = passed_params.pop("model", None) custom_llm_provider = passed_params.pop("custom_llm_provider") provider_config = passed_params.pop("provider_config", None) - drop_params = passed_params.pop("drop_params", None) + drop_params = normalize_drop_params(passed_params.pop("drop_params", None)) additional_drop_params = passed_params.pop("additional_drop_params", None) special_params = passed_params.pop("kwargs") for k, v in special_params.items(): @@ -3084,7 +3085,7 @@ def get_optional_params_embeddings( custom_llm_provider = passed_params.pop("custom_llm_provider", None) special_params = passed_params.pop("kwargs") - drop_params = passed_params.pop("drop_params", None) + drop_params = normalize_drop_params(passed_params.pop("drop_params", None)) additional_drop_params = passed_params.pop("additional_drop_params", None) allowed_openai_params = passed_params.pop("allowed_openai_params", None) or [] # Remove function objects from passed_params to avoid JSON serialization errors @@ -3797,6 +3798,8 @@ def get_optional_params( ): passed_params = locals().copy() special_params = passed_params.pop("kwargs") + drop_params = normalize_drop_params(drop_params) + passed_params["drop_params"] = drop_params # Remove base_model from passed_params so it doesn't interfere with # non_default_params / _check_valid_arg — it's a routing hint, not an # OpenAI param. diff --git a/tests/test_litellm/litellm_core_utils/test_core_helpers.py b/tests/test_litellm/litellm_core_utils/test_core_helpers.py index b67ea91bb0b..a7f93e3c997 100644 --- a/tests/test_litellm/litellm_core_utils/test_core_helpers.py +++ b/tests/test_litellm/litellm_core_utils/test_core_helpers.py @@ -5,6 +5,7 @@ import pytest from litellm.litellm_core_utils.core_helpers import ( _FINISH_REASON_MAP, map_finish_reason, + normalize_drop_params, reconstruct_model_name, redact_nested_match_and_regex_keys, ) @@ -201,3 +202,24 @@ class TestRedactNestedMatchAndRegexKeys: def test_passes_through_none_and_str(self): assert redact_nested_match_and_regex_keys(None) is None assert redact_nested_match_and_regex_keys("plain") == "plain" + + +@pytest.mark.parametrize( + "value, expected", + [ + (True, True), + (False, False), + ("true", True), + ("True", True), + (" TRUE ", True), + ("false", False), + ("False", False), + (None, None), + ("yes", None), + ("", None), + (1, None), + (0, None), + ], +) +def test_normalize_drop_params(value, expected): + assert normalize_drop_params(value) is expected diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index c2c98c8869c..81927d7e959 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -5430,3 +5430,37 @@ class TestRouterRequestTimeoutPropagation: ) == 60 ) + + +@pytest.mark.asyncio +async def test_router_deployment_drop_params_string_true_is_honored(monkeypatch): + from litellm import Router + + monkeypatch.setattr(litellm, "drop_params", False) + router = Router( + model_list=[ + { + "model_name": "gpt-5-nano", + "litellm_params": { + "model": "openai/gpt-5-nano", + "api_key": "sk-fake", + "temperature": 1, + "reasoning_effort": "minimal", + "drop_params": "true", + "mock_response": "Hello, world!", + }, + } + ], + num_retries=0, + ) + + deployment = router.get_deployment_by_model_group_name(model_group_name="gpt-5-nano") + assert deployment is not None + assert deployment.litellm_params.drop_params is True + + response = await router.acompletion( + model="gpt-5-nano", + messages=[{"role": "user", "content": "hi"}], + temperature=0.1, + ) + assert response.choices[0].message.content == "Hello, world!" diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 073ff17991e..0fbf9c0db20 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -4814,3 +4814,31 @@ def test_is_prompt_caching_valid_prompt_explicit_min_token_count_overrides_model is_prompt_caching_valid_prompt(model="claude-opus-4-8", messages=PROMPT_CACHE_MESSAGES, min_token_count=8192) is False ) + + +class TestDropParamsStringCoercion: + @pytest.mark.parametrize("drop_params", ["true", "True", True]) + def test_truthy_drop_params_drops_unsupported_temperature(self, drop_params, monkeypatch): + from litellm.utils import get_optional_params + + monkeypatch.setattr(litellm, "drop_params", False) + result = get_optional_params( + model="gpt-5-nano", + custom_llm_provider="openai", + temperature=0.1, + drop_params=drop_params, + ) + assert "temperature" not in result + + @pytest.mark.parametrize("drop_params", ["false", False, None]) + def test_falsy_drop_params_still_raises(self, drop_params, monkeypatch): + from litellm.utils import get_optional_params + + monkeypatch.setattr(litellm, "drop_params", False) + with pytest.raises(litellm.UnsupportedParamsError): + get_optional_params( + model="gpt-5-nano", + custom_llm_provider="openai", + temperature=0.1, + drop_params=drop_params, + ) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 0d8f55164f9..62103e4f742 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -25667,6 +25667,8 @@ export interface components { default_api_key_rpm_limit?: number | null; /** Default Api Key Tpm Limit */ default_api_key_tpm_limit?: number | null; + /** Drop Params */ + drop_params?: boolean | null; /** Gcs Bucket Name */ gcs_bucket_name?: string | null; /** Input Cost Per Audio Per Second */ @@ -33503,6 +33505,8 @@ export interface components { default_api_key_rpm_limit?: number | null; /** Default Api Key Tpm Limit */ default_api_key_tpm_limit?: number | null; + /** Drop Params */ + drop_params?: boolean | null; /** Gcs Bucket Name */ gcs_bucket_name?: string | null; /** Input Cost Per Audio Per Second */ From dfcea2c1866630313ec3083794a922ff6971583a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 28 Aug 2026 17:16:00 -0700 Subject: [PATCH 02/60] fix(policy_engine): execute post_call guardrail pipelines on responses --- litellm/proxy/utils.py | 47 +++++- .../proxy_logging/test_guardrail_pipeline.py | 152 +++++++++++++++++- 2 files changed, 196 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index d880b529727..bd8ba6ae9f0 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -455,6 +455,30 @@ def _pipeline_managed_guardrail_names(data: Mapping[str, object]) -> frozenset[s ) +def _raise_for_streaming_post_call_pipelines(data: Mapping[str, object]) -> None: + if data.get("stream") is not True: + return + post_call_policies: Final = tuple( + policy_name for policy_name, pipeline in _policy_pipelines(data) if pipeline.mode == "post_call" + ) + if not post_call_policies: + return + raise HTTPException( + status_code=400, + detail={ + "error": { + "message": ( + "Policies with post_call guardrail pipelines cannot govern streaming responses yet: " + f"{', '.join(post_call_policies)}. Retry with stream=false, or move these policies' output " + "guardrails from pipeline steps to guardrails.add, which scans streamed output." + ), + "type": "guardrail_pipeline_error", + "policies": list(post_call_policies), + } + }, + ) + + def _prompt_block_text(block: object) -> str: if isinstance(block, str): return block @@ -1578,6 +1602,7 @@ class ProxyLogging: call_type: str, event_hook: str, raw_request_snapshot: dict | None = None, # mutable-ok: same request-payload shape as data + response: LLMResponseTypes | None = None, ) -> dict: """ Execute guardrail pipelines if any are configured for this request. @@ -1596,6 +1621,8 @@ class ProxyLogging: if not pipelines: return data + step_input: Final = {**data, "response": response} if response is not None else data + for policy_name, pipeline in pipelines: if pipeline.mode != event_hook: continue @@ -1603,7 +1630,7 @@ class ProxyLogging: result: PipelineExecutionResult = await PipelineExecutor.execute_steps( steps=pipeline.steps, mode=pipeline.mode, - data=data, + data=step_input, user_api_key_dict=user_api_key_dict, call_type=call_type, policy_name=policy_name, @@ -1614,6 +1641,7 @@ class ProxyLogging: result=result, data=data, policy_name=policy_name, + original_response=response, ) return data @@ -1623,14 +1651,18 @@ class ProxyLogging: result: PipelineExecutionResult, data: dict, policy_name: str, + original_response: LLMResponseTypes | None = None, ) -> dict: """ Handle a PipelineExecutionResult — allow, block, or modify_response. Returns data dict if allowed, raises on block/modify_response. + ``original_response`` is set on the post_call path, where allowed + modifications land on the response object in place, so the request + payload (already sent upstream) is left untouched. """ if result.terminal_action == "allow": - if result.modified_data is not None: + if result.modified_data is not None and original_response is None: data.update(result.modified_data) return data @@ -1671,6 +1703,7 @@ class ProxyLogging: request_data=data, guardrail_name=f"pipeline:{policy_name}", detection_info=None, + original_response=original_response, ) return data @@ -1786,6 +1819,8 @@ class ProxyLogging: ) try: + _raise_for_streaming_post_call_pipelines(data) + # Execute guardrail pipelines before the normal callback loop data = await self._maybe_execute_pipelines( data=data, @@ -2774,6 +2809,14 @@ class ProxyLogging: from litellm.proxy.proxy_server import llm_router from litellm.types.guardrails import GuardrailEventHooks + await self._maybe_execute_pipelines( + data=data, + user_api_key_dict=user_api_key_dict, + call_type=getattr(data.get("litellm_logging_obj"), "call_type", None) or "acompletion", + event_hook="post_call", + response=response, + ) + guardrail_callbacks: Final[list[CustomGuardrail]] = [] other_callbacks: Final[list[CustomLogger]] = [] try: diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py index e99e34d65d4..b7e86b68fe4 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -23,7 +23,7 @@ from litellm.integrations.custom_guardrail import ( ModifyResponseException, ) from litellm.integrations.prometheus import PrometheusLogger -from litellm.proxy.utils import ProxyLogging +from litellm.proxy.utils import ProxyLogging, _raise_for_streaming_post_call_pipelines from litellm.types.guardrails import GuardrailEventHooks from litellm.types.proxy.policy_engine.pipeline_types import ( GuardrailPipeline, @@ -865,3 +865,153 @@ async def test_process_prompt_template_aresponses_swaps_model_and_merges_input(p hook_kwargs = logging_obj.async_get_chat_completion_prompt.await_args.kwargs assert hook_kwargs["messages"] == [{"role": "user", "content": "Who are you?"}] assert hook_kwargs["prompt_spec"] is prompt_spec + + +# --------------------------------------------------------------------------- +# post_call pipeline execution (LIT-6410) +# --------------------------------------------------------------------------- + + +def _post_call_pipeline_data(guardrail: str = "gr-post", **extra: Any) -> Dict[str, Any]: + pipeline = GuardrailPipeline( + mode="post_call", + steps=[PipelineStep(guardrail=guardrail, on_pass="allow", on_fail="block")], + ) + return { + "model": "m", + "messages": [{"role": "user", "content": "hi"}], + "metadata": { + "_guardrail_pipelines": [("response-governance", pipeline)], + "_pipeline_managed_guardrails": {guardrail}, + }, + **extra, + } + + +@pytest.mark.asyncio +async def test_post_call_success_hook_runs_post_call_pipeline_and_reraises_block( + proxy_logging, make_user_api_key_auth, monkeypatch +): + seen: Dict[str, Any] = {} + + class OutputBlockingGuardrail(CustomGuardrail): + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + seen["response"] = response + raise HTTPException(status_code=400, detail={"error": "output blocked"}) + + monkeypatch.setattr( + litellm, + "callbacks", + [OutputBlockingGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False)], + ) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data() + response = litellm.ModelResponse() + + with pytest.raises(HTTPException) as info: + await proxy_logging.post_call_success_hook( + data=data, response=response, user_api_key_dict=make_user_api_key_auth() + ) + + assert info.value.detail["error"] == "output blocked" + assert seen["response"] is response + + +@pytest.mark.asyncio +async def test_post_call_pipeline_pass_runs_once_and_leaves_request_data_untouched( + proxy_logging, make_user_api_key_auth, monkeypatch +): + seen: Dict[str, Any] = {"count": 0} + + class RecordingGuardrail(CustomGuardrail): + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + seen["count"] += 1 + seen["response"] = response + return None + + monkeypatch.setattr( + litellm, + "callbacks", + [RecordingGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False)], + ) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data() + response = litellm.ModelResponse() + + out = await proxy_logging.post_call_success_hook( + data=data, response=response, user_api_key_dict=make_user_api_key_auth() + ) + + assert out is response + assert seen["response"] is response + assert seen["count"] == 1 + assert "response" not in data + assert "guardrails" not in data["metadata"] + + +def test_handle_pipeline_result_modify_response_carries_original_response(): + result = MagicMock() + result.terminal_action = "modify_response" + result.modify_response_message = "filtered" + response = litellm.ModelResponse() + + with pytest.raises(ModifyResponseException) as info: + ProxyLogging._handle_pipeline_result( + result=result, data={"model": "m"}, policy_name="p", original_response=response + ) + + assert info.value.original_response is response + + +def test_handle_pipeline_result_allow_discards_modifications_on_post_call(): + data = {"a": 1, "metadata": {"guardrails": ["other"]}} + result = MagicMock() + result.terminal_action = "allow" + result.modified_data = {"metadata": {"guardrails": ["gr-post"]}, "response": object()} + + out = ProxyLogging._handle_pipeline_result( + result=result, data=data, policy_name="p", original_response=litellm.ModelResponse() + ) + + assert out is data + assert data == {"a": 1, "metadata": {"guardrails": ["other"]}} + + +@pytest.mark.asyncio +async def test_pre_call_hook_rejects_streaming_request_with_post_call_pipeline( + proxy_logging, make_user_api_key_auth, monkeypatch +): + monkeypatch.setattr(litellm, "callbacks", []) + data = _post_call_pipeline_data(stream=True) + + with pytest.raises(HTTPException) as info: + await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data=data, + call_type="completion", + guardrails_only=True, + ) + + assert info.value.status_code == 400 + assert info.value.detail["error"]["policies"] == ["response-governance"] + assert "stream=false" in info.value.detail["error"]["message"] + + +def test_raise_for_streaming_post_call_pipelines_ignores_non_streaming_and_pre_call(): + post_call = GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="g", on_fail="block")]) + pre_call = GuardrailPipeline(mode="pre_call", steps=[PipelineStep(guardrail="g", on_fail="block")]) + + assert ( + _raise_for_streaming_post_call_pipelines( + {"stream": False, "metadata": {"_guardrail_pipelines": [("p", post_call)]}} + ) + is None + ) + assert _raise_for_streaming_post_call_pipelines({"metadata": {"_guardrail_pipelines": [("p", post_call)]}}) is None + assert ( + _raise_for_streaming_post_call_pipelines( + {"stream": True, "metadata": {"_guardrail_pipelines": [("p", pre_call)]}} + ) + is None + ) + assert _raise_for_streaming_post_call_pipelines({"stream": True}) is None From e6edd62f5d0f010d34c203d9df8192462a1622c0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 28 Aug 2026 17:40:48 -0700 Subject: [PATCH 03/60] fix(policy_engine): propagate post_call pipeline replacement responses to the client --- .../proxy/policy_engine/pipeline_executor.py | 19 +++-- litellm/proxy/utils.py | 32 +++++--- .../proxy_logging/test_guardrail_pipeline.py | 81 ++++++++++++++++++- .../utils/proxy_logging/test_pre_call_hook.py | 2 +- 4 files changed, 112 insertions(+), 22 deletions(-) diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index a5619821197..190784a2a60 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -108,8 +108,10 @@ class PipelineExecutor: action, ) - # Forward modified data to next step if pass_data is True - if step.pass_data and modified_data is not None: + # Forward modified data to the next step if pass_data is True; + # post_call response replacements always chain, matching the flat + # callback loop where each hook sees the previous hook's response + if modified_data is not None and (step.pass_data or mode == "post_call"): working_data = {**working_data, **modified_data} # Handle terminal actions @@ -227,11 +229,14 @@ class PipelineExecutor: # same contract as run_in_parallel/scan_raw_request elsewhere: any # data it returned is discarded, since applying it on top of the # raw snapshot would silently undo whatever an earlier step in - # this pipeline already did. - modified_data = None - if response is not None and isinstance(response, dict) and not scans_raw_request: - modified_data = response - return ("pass", modified_data, None, None) + # this pipeline already did. A post_call hook's non-None return is + # a replacement response (the flat callback-loop contract), carried + # under the same "response" key the step input uses. + if response is None or scans_raw_request: + return ("pass", None, None, None) + if mode == "post_call": + return ("pass", {"response": response}, None, None) + return ("pass", response if isinstance(response, dict) else None, None, None) except Exception as e: if CustomGuardrail._is_guardrail_intervention(e): diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index bd8ba6ae9f0..25c1068cc6d 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1603,7 +1603,7 @@ class ProxyLogging: event_hook: str, raw_request_snapshot: dict | None = None, # mutable-ok: same request-payload shape as data response: LLMResponseTypes | None = None, - ) -> dict: + ) -> tuple[dict, LLMResponseTypes | None]: """ Execute guardrail pipelines if any are configured for this request. @@ -1615,18 +1615,21 @@ class ProxyLogging: ``scan_raw_request`` evaluates the pristine request, not whatever an earlier ``pass_data`` step in the same pipeline already rewrote. - Returns the (possibly modified) data dict. + Returns the (possibly modified) data dict, plus the replacement + response when a post_call pipeline step returned one (None when the + response is unchanged), matching the flat callback-loop contract. """ pipelines: Final = _policy_pipelines(data) if not pipelines: - return data - - step_input: Final = {**data, "response": response} if response is not None else data + return data, None + current_response = response # rebind-ok: chains each pipeline's replacement response into the next for policy_name, pipeline in pipelines: if pipeline.mode != event_hook: continue + step_input: dict = {**data, "response": current_response} if current_response is not None else data + result: PipelineExecutionResult = await PipelineExecutor.execute_steps( steps=pipeline.steps, mode=pipeline.mode, @@ -1641,10 +1644,13 @@ class ProxyLogging: result=result, data=data, policy_name=policy_name, - original_response=response, + original_response=current_response, ) - return data + if current_response is not None and result.modified_data is not None: + current_response = result.modified_data.get("response", current_response) + + return data, current_response if current_response is not response else None @staticmethod def _handle_pipeline_result( @@ -1657,9 +1663,9 @@ class ProxyLogging: Handle a PipelineExecutionResult — allow, block, or modify_response. Returns data dict if allowed, raises on block/modify_response. - ``original_response`` is set on the post_call path, where allowed - modifications land on the response object in place, so the request - payload (already sent upstream) is left untouched. + ``original_response`` is set on the post_call path, where the request + payload (already sent upstream) must stay untouched; a replacement + response carried in ``modified_data`` is adopted by the caller. """ if result.terminal_action == "allow": if result.modified_data is not None and original_response is None: @@ -1822,7 +1828,7 @@ class ProxyLogging: _raise_for_streaming_post_call_pipelines(data) # Execute guardrail pipelines before the normal callback loop - data = await self._maybe_execute_pipelines( + data, _ = await self._maybe_execute_pipelines( data=data, user_api_key_dict=user_api_key_dict, call_type=call_type, @@ -2809,13 +2815,15 @@ class ProxyLogging: from litellm.proxy.proxy_server import llm_router from litellm.types.guardrails import GuardrailEventHooks - await self._maybe_execute_pipelines( + _, pipeline_response = await self._maybe_execute_pipelines( data=data, user_api_key_dict=user_api_key_dict, call_type=getattr(data.get("litellm_logging_obj"), "call_type", None) or "acompletion", event_hook="post_call", response=response, ) + if pipeline_response is not None: + response = pipeline_response # rebind-ok: adopt the pipeline's replacement response, same contract as the callback loops below guardrail_callbacks: Final[list[CustomGuardrail]] = [] other_callbacks: Final[list[CustomLogger]] = [] diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py index b7e86b68fe4..8bc71f6e178 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -326,13 +326,14 @@ def test_process_guardrail_metadata_invalid_data_raises(proxy_logging): @pytest.mark.asyncio async def test_maybe_execute_pipelines_no_pipelines_returns_data(proxy_logging, make_user_api_key_auth): data = {"messages": [{"role": "user"}], "model": "m", "temperature": 0.1} - out = await proxy_logging._maybe_execute_pipelines( + out, replacement = await proxy_logging._maybe_execute_pipelines( data=data, user_api_key_dict=make_user_api_key_auth(), call_type="completion", event_hook="pre_call", ) assert out == {"messages": [{"role": "user"}], "model": "m", "temperature": 0.1} + assert replacement is None @pytest.mark.asyncio @@ -344,7 +345,7 @@ async def test_maybe_execute_pipelines_skips_pipelines_with_other_mode(proxy_log monkeypatch.setattr( "litellm.proxy.policy_engine.pipeline_executor.PipelineExecutor.execute_steps", executed ) - out = await proxy_logging._maybe_execute_pipelines( + out, replacement = await proxy_logging._maybe_execute_pipelines( data=data, user_api_key_dict=make_user_api_key_auth(), call_type="completion", @@ -352,6 +353,7 @@ async def test_maybe_execute_pipelines_skips_pipelines_with_other_mode(proxy_log ) executed.assert_not_called() assert out is data + assert replacement is None @pytest.mark.parametrize( @@ -949,6 +951,81 @@ async def test_post_call_pipeline_pass_runs_once_and_leaves_request_data_untouch assert "guardrails" not in data["metadata"] +@pytest.mark.asyncio +async def test_post_call_pipeline_replacement_response_reaches_caller( + proxy_logging, make_user_api_key_auth, monkeypatch +): + masked = litellm.ModelResponse() + + class MaskingGuardrail(CustomGuardrail): + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + return masked + + monkeypatch.setattr( + litellm, + "callbacks", + [MaskingGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False)], + ) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data() + + out = await proxy_logging.post_call_success_hook( + data=data, response=litellm.ModelResponse(), user_api_key_dict=make_user_api_key_auth() + ) + + assert out is masked + assert "response" not in data + + +@pytest.mark.asyncio +async def test_post_call_pipeline_replacement_chains_to_next_step_without_pass_data( + proxy_logging, make_user_api_key_auth, monkeypatch +): + masked = litellm.ModelResponse() + seen: Dict[str, Any] = {} + + class MaskingGuardrail(CustomGuardrail): + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + return masked + + class RecordingGuardrail(CustomGuardrail): + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + seen["response"] = response + return None + + pipeline = GuardrailPipeline( + mode="post_call", + steps=[ + PipelineStep(guardrail="gr-mask", on_pass="next", on_fail="block"), + PipelineStep(guardrail="gr-audit", on_pass="allow", on_fail="block"), + ], + ) + monkeypatch.setattr( + litellm, + "callbacks", + [ + MaskingGuardrail(guardrail_name="gr-mask", event_hook=GuardrailEventHooks.post_call, default_on=False), + RecordingGuardrail(guardrail_name="gr-audit", event_hook=GuardrailEventHooks.post_call, default_on=False), + ], + ) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = { + "model": "m", + "messages": [{"role": "user", "content": "hi"}], + "metadata": { + "_guardrail_pipelines": [("response-governance", pipeline)], + "_pipeline_managed_guardrails": {"gr-mask", "gr-audit"}, + }, + } + + out = await proxy_logging.post_call_success_hook( + data=data, response=litellm.ModelResponse(), user_api_key_dict=make_user_api_key_auth() + ) + + assert out is masked + assert seen["response"] is masked + + def test_handle_pipeline_result_modify_response_carries_original_response(): result = MagicMock() result.terminal_action = "modify_response" diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_pre_call_hook.py b/tests/test_litellm/proxy/utils/proxy_logging/test_pre_call_hook.py index 0971ce09d79..06cf328a20c 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_pre_call_hook.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_pre_call_hook.py @@ -660,7 +660,7 @@ async def test_scan_raw_request_snapshot_taken_before_pipelines( for msg in data.get("messages", []): if "SECRET" in msg.get("content", ""): msg["content"] = msg["content"].replace("SECRET", "[REDACTED]") - return data + return data, None monkeypatch.setattr(ProxyLogging, "_maybe_execute_pipelines", fake_pipelines) monkeypatch.setattr(litellm, "callbacks", [_BlockOnSecretGuardrail(scan_raw_request=True)]) From aeac6a412c98c07cce64c5ddbd74d005f2e60e7e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 28 Aug 2026 17:59:15 -0700 Subject: [PATCH 04/60] fix(policy_engine): skip pipeline-managed guardrails in the response-path guardrail loop --- litellm/proxy/utils.py | 9 ++++++- .../proxy_logging/test_guardrail_pipeline.py | 26 +++++++++++++++++++ 2 files changed, 34 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 25c1068cc6d..75cc5c259a7 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -2825,6 +2825,7 @@ class ProxyLogging: if pipeline_response is not None: response = pipeline_response # rebind-ok: adopt the pipeline's replacement response, same contract as the callback loops below + pipeline_managed: Final = _pipeline_managed_guardrail_names(data) guardrail_callbacks: Final[list[CustomGuardrail]] = [] other_callbacks: Final[list[CustomLogger]] = [] try: @@ -2849,12 +2850,18 @@ class ProxyLogging: guardrail_data: Final = _check_and_merge_model_level_guardrails(data=data, llm_router=llm_router) parallel_guardrails: Final[tuple[CustomGuardrail, ...]] = tuple( - callback for callback in guardrail_callbacks if getattr(callback, "run_in_parallel", False) + callback + for callback in guardrail_callbacks + if getattr(callback, "run_in_parallel", False) + and not (callback.guardrail_name and callback.guardrail_name in pipeline_managed) ) for callback in guardrail_callbacks: # Main - V2 Guardrails implementation + if callback.guardrail_name and callback.guardrail_name in pipeline_managed: + continue + if getattr(callback, "run_in_parallel", False): continue diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py index 8bc71f6e178..c92f5fa6c55 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -951,6 +951,32 @@ async def test_post_call_pipeline_pass_runs_once_and_leaves_request_data_untouch assert "guardrails" not in data["metadata"] +@pytest.mark.asyncio +async def test_post_call_pipeline_managed_default_on_guardrail_runs_exactly_once( + proxy_logging, make_user_api_key_auth, monkeypatch +): + seen: Dict[str, Any] = {"count": 0} + + class CountingGuardrail(CustomGuardrail): + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + seen["count"] += 1 + return None + + monkeypatch.setattr( + litellm, + "callbacks", + [CountingGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=True)], + ) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data() + + await proxy_logging.post_call_success_hook( + data=data, response=litellm.ModelResponse(), user_api_key_dict=make_user_api_key_auth() + ) + + assert seen["count"] == 1 + + @pytest.mark.asyncio async def test_post_call_pipeline_replacement_response_reaches_caller( proxy_logging, make_user_api_key_auth, monkeypatch From 55569729b05d601c139e43b8faba447983e89f74 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 28 Aug 2026 18:16:35 -0700 Subject: [PATCH 05/60] fix(policy_engine): scope pipeline-managed guardrail skips to the pipeline's mode --- litellm/proxy/utils.py | 20 ++--- .../proxy_logging/test_guardrail_pipeline.py | 73 +++++++++++++++++++ 2 files changed, 84 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 75cc5c259a7..d3f2e1d7301 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -11,7 +11,7 @@ import sys import threading import time import traceback -from collections.abc import AsyncGenerator, Awaitable, Callable, Collection, Coroutine, Mapping, Sequence +from collections.abc import AsyncGenerator, Awaitable, Callable, Coroutine, Mapping, Sequence from dataclasses import dataclass, field from datetime import date, datetime, timedelta, timezone from email.mime.multipart import MIMEMultipart @@ -446,12 +446,14 @@ def _policy_pipelines(data: Mapping[str, object]) -> tuple[tuple[str, "Guardrail ) -def _pipeline_managed_guardrail_names(data: Mapping[str, object]) -> frozenset[str]: - managed: Final = _policy_state_metadata(data).get("_pipeline_managed_guardrails") - return ( - frozenset(cast("Collection[str]", managed)) # cast-ok: the policy engine wrote these guardrail names - if managed - else frozenset() +def _pipeline_managed_guardrail_names( + data: Mapping[str, object], mode: Literal["pre_call", "post_call"] +) -> frozenset[str]: + return frozenset( + step.guardrail + for _policy_name, pipeline in _policy_pipelines(data) + if pipeline.mode == mode + for step in pipeline.steps ) @@ -1837,7 +1839,7 @@ class ProxyLogging: ) # Get pipeline-managed guardrails to skip in normal loop - pipeline_managed: Final = _pipeline_managed_guardrail_names(data) + pipeline_managed: Final = _pipeline_managed_guardrail_names(data, "pre_call") caps: Final = ProxyLogging._callback_capabilities() # Skip the per-request callback walk entirely when nothing in @@ -2825,7 +2827,7 @@ class ProxyLogging: if pipeline_response is not None: response = pipeline_response # rebind-ok: adopt the pipeline's replacement response, same contract as the callback loops below - pipeline_managed: Final = _pipeline_managed_guardrail_names(data) + pipeline_managed: Final = _pipeline_managed_guardrail_names(data, "post_call") guardrail_callbacks: Final[list[CustomGuardrail]] = [] other_callbacks: Final[list[CustomLogger]] = [] try: diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py index c92f5fa6c55..4e1ccf71c5c 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -977,6 +977,79 @@ async def test_post_call_pipeline_managed_default_on_guardrail_runs_exactly_once assert seen["count"] == 1 +@pytest.mark.asyncio +async def test_post_call_hook_still_runs_guardrail_managed_only_by_pre_call_pipeline( + proxy_logging, make_user_api_key_auth, monkeypatch +): + seen: Dict[str, Any] = {"count": 0} + + class DualStageGuardrail(CustomGuardrail): + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + seen["count"] += 1 + return None + + pre_call_pipeline = GuardrailPipeline( + mode="pre_call", + steps=[PipelineStep(guardrail="gr-dual", on_pass="allow", on_fail="block")], + ) + monkeypatch.setattr( + litellm, + "callbacks", + [DualStageGuardrail(guardrail_name="gr-dual", event_hook=["pre_call", "post_call"], default_on=True)], + ) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = { + "model": "m", + "messages": [{"role": "user", "content": "hi"}], + "metadata": { + "_guardrail_pipelines": [("request-governance", pre_call_pipeline)], + "_pipeline_managed_guardrails": {"gr-dual"}, + }, + } + + await proxy_logging.post_call_success_hook( + data=data, response=litellm.ModelResponse(), user_api_key_dict=make_user_api_key_auth() + ) + + assert seen["count"] == 1 + + +@pytest.mark.asyncio +async def test_pre_call_hook_still_runs_guardrail_managed_only_by_post_call_pipeline( + proxy_logging, make_user_api_key_auth, monkeypatch +): + seen: Dict[str, Any] = {"count": 0} + + class DualStageGuardrail(CustomGuardrail): + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + seen["count"] += 1 + return data + + post_call_pipeline = GuardrailPipeline( + mode="post_call", + steps=[PipelineStep(guardrail="gr-dual", on_pass="allow", on_fail="block")], + ) + monkeypatch.setattr( + litellm, + "callbacks", + [DualStageGuardrail(guardrail_name="gr-dual", event_hook=["pre_call", "post_call"], default_on=True)], + ) + data = { + "model": "m", + "messages": [{"role": "user", "content": "hi"}], + "metadata": { + "_guardrail_pipelines": [("response-governance", post_call_pipeline)], + "_pipeline_managed_guardrails": {"gr-dual"}, + }, + } + + await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), data=data, call_type="completion" + ) + + assert seen["count"] == 1 + + @pytest.mark.asyncio async def test_post_call_pipeline_replacement_response_reaches_caller( proxy_logging, make_user_api_key_auth, monkeypatch From 996019cd23423c7b2a35dbd50f4b3b3571362cd3 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 00:14:23 -0700 Subject: [PATCH 06/60] fix(policy_engine): keep post_call pipeline guardrail logging and reject background bypass Post_call pipelines run step hooks against a copied request dict, so guardrail writes into the metadata bucket (applied_guardrails for the response header, standard_logging_guardrail_information for spend logs) were dropped when the guardrail was the first writer. Merge those writes back onto the request on the post_call allow path, keeping the request payload and the executor's per-step guardrails activation flag out of it. Background /v1/responses requests dodge the streaming 400: pre_call sees stream unset, then the polling task forces stream=true with pre-call logic skipped and the streaming branch returns before post_call_success_hook, silently bypassing post_call pipelines. Reject background=true at pre_call the same way as stream=true. Also pin the run_in_parallel pipeline-managed exclusion in both hook loops with regression tests. --- litellm/proxy/utils.py | 49 +++++- .../proxy_logging/test_guardrail_pipeline.py | 148 +++++++++++++++++- 2 files changed, 187 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index d3f2e1d7301..0f97473312c 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -457,8 +457,37 @@ def _pipeline_managed_guardrail_names( ) +def _merge_pipeline_metadata_bucket(data: dict, bucket_key: str, modified_bucket_value: object) -> None: + if not isinstance(modified_bucket_value, dict): + return + modified_bucket: Final = cast("dict[str, object]", modified_bucket_value) # cast-ok: metadata buckets are str-keyed + surviving_writes: Final = {key: value for key, value in modified_bucket.items() if key != "guardrails"} + existing_bucket: Final = data.get(bucket_key) + if isinstance(existing_bucket, dict): + cast("dict[str, object]", existing_bucket).update(surviving_writes) # cast-ok: metadata buckets are str-keyed + else: + data[bucket_key] = surviving_writes + + +def _merge_pipeline_metadata_writes(data: dict, modified_data: Mapping[str, object]) -> None: + """ + Copy metadata-bucket writes from a pipeline's working copy back onto the request. + + Post_call pipelines run step hooks against a copied request dict so the payload + already sent upstream stays untouched, but hooks record proxy-internal logging + state in the metadata buckets (``applied_guardrails`` for response headers, + ``standard_logging_guardrail_information`` for spend logs), and those writes + must reach the request dict the proxy keeps reading after the pipeline returns. + + The ``guardrails`` key is the executor's per-step activation flag for + ``should_run_guardrail``, not a hook write, so it stays in the working copy. + """ + for bucket_key in ("metadata", "litellm_metadata"): + _merge_pipeline_metadata_bucket(data, bucket_key, modified_data.get(bucket_key)) + + def _raise_for_streaming_post_call_pipelines(data: Mapping[str, object]) -> None: - if data.get("stream") is not True: + if data.get("stream") is not True and data.get("background") is not True: return post_call_policies: Final = tuple( policy_name for policy_name, pipeline in _policy_pipelines(data) if pipeline.mode == "post_call" @@ -470,9 +499,10 @@ def _raise_for_streaming_post_call_pipelines(data: Mapping[str, object]) -> None detail={ "error": { "message": ( - "Policies with post_call guardrail pipelines cannot govern streaming responses yet: " - f"{', '.join(post_call_policies)}. Retry with stream=false, or move these policies' output " - "guardrails from pipeline steps to guardrails.add, which scans streamed output." + "Policies with post_call guardrail pipelines cannot govern streaming or background " + f"responses yet: {', '.join(post_call_policies)}. Retry with stream=false and " + "background=false, or move these policies' output guardrails from pipeline steps to " + "guardrails.add, which scans streamed output." ), "type": "guardrail_pipeline_error", "policies": list(post_call_policies), @@ -1667,11 +1697,16 @@ class ProxyLogging: Returns data dict if allowed, raises on block/modify_response. ``original_response`` is set on the post_call path, where the request payload (already sent upstream) must stay untouched; a replacement - response carried in ``modified_data`` is adopted by the caller. + response carried in ``modified_data`` is adopted by the caller, and + metadata-bucket writes (applied guardrails, guardrail logging info) + are merged back so headers and spend logs still see them. """ if result.terminal_action == "allow": - if result.modified_data is not None and original_response is None: - data.update(result.modified_data) + if result.modified_data is not None: + if original_response is None: + data.update(result.modified_data) + else: + _merge_pipeline_metadata_writes(data, result.modified_data) return data if result.terminal_action == "block": diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py index 4e1ccf71c5c..34a25d75bc0 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -23,6 +23,7 @@ from litellm.integrations.custom_guardrail import ( ModifyResponseException, ) from litellm.integrations.prometheus import PrometheusLogger +from litellm.proxy.common_utils.callback_utils import add_guardrail_to_applied_guardrails_header from litellm.proxy.utils import ProxyLogging, _raise_for_streaming_post_call_pipelines from litellm.types.guardrails import GuardrailEventHooks from litellm.types.proxy.policy_engine.pipeline_types import ( @@ -1139,18 +1140,132 @@ def test_handle_pipeline_result_modify_response_carries_original_response(): assert info.value.original_response is response -def test_handle_pipeline_result_allow_discards_modifications_on_post_call(): +def test_handle_pipeline_result_allow_on_post_call_keeps_metadata_writes_only(): data = {"a": 1, "metadata": {"guardrails": ["other"]}} result = MagicMock() result.terminal_action = "allow" - result.modified_data = {"metadata": {"guardrails": ["gr-post"]}, "response": object()} + result.modified_data = { + "a": 2, + "metadata": {"guardrails": ["other"], "applied_guardrails": ["gr-post"]}, + "response": object(), + } out = ProxyLogging._handle_pipeline_result( result=result, data=data, policy_name="p", original_response=litellm.ModelResponse() ) assert out is data - assert data == {"a": 1, "metadata": {"guardrails": ["other"]}} + assert data["a"] == 1 + assert "response" not in data + assert data["metadata"] == {"guardrails": ["other"], "applied_guardrails": ["gr-post"]} + + +@pytest.mark.asyncio +async def test_post_call_pipeline_guardrail_metadata_writes_reach_request_data( + proxy_logging, make_user_api_key_auth, monkeypatch +): + class HeaderWritingGuardrail(CustomGuardrail): + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name="gr-post") + self.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response={"verdict": "pass"}, + request_data=data, + guardrail_status="success", + ) + return None + + monkeypatch.setattr( + litellm, + "callbacks", + [HeaderWritingGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False)], + ) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data() + + await proxy_logging.post_call_success_hook( + data=data, response=litellm.ModelResponse(), user_api_key_dict=make_user_api_key_auth() + ) + + assert data["metadata"]["applied_guardrails"] == ["gr-post"] + slg_entries = data["metadata"]["standard_logging_guardrail_information"] + assert len(slg_entries) == 1 + assert slg_entries[0]["guardrail_name"] == "gr-post" + + +@pytest.mark.asyncio +async def test_post_call_pipeline_managed_parallel_guardrail_runs_exactly_once( + proxy_logging, make_user_api_key_auth, monkeypatch +): + seen: Dict[str, Any] = {"count": 0} + + class CountingGuardrail(CustomGuardrail): + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + seen["count"] += 1 + return None + + monkeypatch.setattr( + litellm, + "callbacks", + [ + CountingGuardrail( + guardrail_name="gr-post", + event_hook=GuardrailEventHooks.post_call, + default_on=True, + run_in_parallel=True, + ) + ], + ) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data() + + await proxy_logging.post_call_success_hook( + data=data, response=litellm.ModelResponse(), user_api_key_dict=make_user_api_key_auth() + ) + + assert seen["count"] == 1 + + +@pytest.mark.asyncio +async def test_pre_call_pipeline_managed_parallel_guardrail_runs_exactly_once( + proxy_logging, make_user_api_key_auth, monkeypatch +): + seen: Dict[str, Any] = {"count": 0} + + class CountingGuardrail(CustomGuardrail): + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + seen["count"] += 1 + return data + + pre_call_pipeline = GuardrailPipeline( + mode="pre_call", + steps=[PipelineStep(guardrail="gr-pre", on_pass="allow", on_fail="block")], + ) + monkeypatch.setattr( + litellm, + "callbacks", + [ + CountingGuardrail( + guardrail_name="gr-pre", + event_hook=GuardrailEventHooks.pre_call, + default_on=True, + run_in_parallel=True, + ) + ], + ) + data = { + "model": "m", + "messages": [{"role": "user", "content": "hi"}], + "metadata": { + "_guardrail_pipelines": [("request-governance", pre_call_pipeline)], + "_pipeline_managed_guardrails": {"gr-pre"}, + }, + } + + await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), data=data, call_type="completion" + ) + + assert seen["count"] == 1 @pytest.mark.asyncio @@ -1173,6 +1288,26 @@ async def test_pre_call_hook_rejects_streaming_request_with_post_call_pipeline( assert "stream=false" in info.value.detail["error"]["message"] +@pytest.mark.asyncio +async def test_pre_call_hook_rejects_background_request_with_post_call_pipeline( + proxy_logging, make_user_api_key_auth, monkeypatch +): + monkeypatch.setattr(litellm, "callbacks", []) + data = _post_call_pipeline_data(background=True) + + with pytest.raises(HTTPException) as info: + await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data=data, + call_type="aresponses", + guardrails_only=True, + ) + + assert info.value.status_code == 400 + assert info.value.detail["error"]["policies"] == ["response-governance"] + assert "background=false" in info.value.detail["error"]["message"] + + def test_raise_for_streaming_post_call_pipelines_ignores_non_streaming_and_pre_call(): post_call = GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="g", on_fail="block")]) pre_call = GuardrailPipeline(mode="pre_call", steps=[PipelineStep(guardrail="g", on_fail="block")]) @@ -1183,6 +1318,12 @@ def test_raise_for_streaming_post_call_pipelines_ignores_non_streaming_and_pre_c ) is None ) + assert ( + _raise_for_streaming_post_call_pipelines( + {"background": False, "metadata": {"_guardrail_pipelines": [("p", post_call)]}} + ) + is None + ) assert _raise_for_streaming_post_call_pipelines({"metadata": {"_guardrail_pipelines": [("p", post_call)]}}) is None assert ( _raise_for_streaming_post_call_pipelines( @@ -1191,3 +1332,4 @@ def test_raise_for_streaming_post_call_pipelines_ignores_non_streaming_and_pre_c is None ) assert _raise_for_streaming_post_call_pipelines({"stream": True}) is None + assert _raise_for_streaming_post_call_pipelines({"background": True}) is None From c5bcf3a73594ce5fad662a47781d89dbe7955718 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 12:06:43 -0700 Subject: [PATCH 07/60] feat(policy_engine): execute post_call guardrail pipelines on streaming responses --- .../unified_guardrail/unified_guardrail.py | 46 ++++- .../proxy/policy_engine/pipeline_executor.py | 47 ++++- litellm/proxy/utils.py | 187 ++++++++++++++++-- .../test_unified_guardrail.py | 2 +- .../proxy_logging/test_guardrail_pipeline.py | 168 +++++++++++++++- 5 files changed, 417 insertions(+), 33 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py index e95e97bfe74..60b4444e1d4 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -19,7 +19,7 @@ from litellm.cost_calculator import _infer_call_type from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_for_route -from litellm.llms import load_guardrail_translation_mappings +from litellm.llms import get_guardrail_translation_mapping, load_guardrail_translation_mappings from litellm.proxy._types import UserAPIKeyAuth from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import ( @@ -62,6 +62,36 @@ def _as_endpoint_translation(translation: _EndpointTranslation) -> _EndpointTran return translation +def resolve_endpoint_translation( + user_api_key_dict: UserAPIKeyAuth, first_response_item: object | None +) -> "tuple[str, BaseTranslation] | None": + """ + Resolve the endpoint guardrail translation for a streamed response: the + request route wins, falling back to inferring the call type from the first + response chunk (the same resolution order the streaming iterator hook uses). + Returns None when the call type is unresolvable or has no translation. + """ + route_call_types: Final = ( + get_call_types_for_route(user_api_key_dict.request_route) if user_api_key_dict.request_route else None + ) + call_type: Final = ( + route_call_types[0].value + if route_call_types + else ( + _infer_call_type(call_type=None, completion_response=first_response_item) + if first_response_item is not None + else None + ) + ) + if call_type is None: + return None + try: + handler_cls: Final = get_guardrail_translation_mapping(CallTypes(call_type)) + except ValueError: + return None + return call_type, handler_cls() + + def _chunk_choices(item: object) -> Sequence[object]: choices: Final[Sequence[object]] = getattr(item, "choices", None) or [] return choices @@ -346,7 +376,7 @@ class UnifiedLLMGuardrails(CustomLogger): return response - async def _handle_streaming_block( + async def handle_streaming_block( self, exc: "ModifyResponseException", endpoint_translation: _EndpointTranslation, @@ -402,7 +432,7 @@ class UnifiedLLMGuardrails(CustomLogger): return None return call_type - async def _emit_streaming_http_error( + async def emit_streaming_http_error( self, exc: HTTPException, call_type: str | None, @@ -577,7 +607,7 @@ class UnifiedLLMGuardrails(CustomLogger): except ModifyResponseException as e: if e.original_response is None: e.original_response = responses_so_far - async for block_chunk in self._handle_streaming_block( + async for block_chunk in self.handle_streaming_block( e, endpoint_translation, stream_started=bool(responses_yielded), @@ -586,7 +616,7 @@ class UnifiedLLMGuardrails(CustomLogger): yield block_chunk raise _StreamTerminated() except HTTPException as e: - async for error_item in self._emit_streaming_http_error(e, call_type, responses_so_far, request_data): + async for error_item in self.emit_streaming_http_error(e, call_type, responses_so_far, request_data): yield error_item raise _StreamTerminated() @@ -758,7 +788,7 @@ class UnifiedLLMGuardrails(CustomLogger): except ModifyResponseException as e: if e.original_response is None: e.original_response = responses_so_far - async for block_chunk in self._handle_streaming_block( + async for block_chunk in self.handle_streaming_block( e, endpoint_translation, stream_started=bool(responses_yielded), @@ -1060,7 +1090,7 @@ class UnifiedLLMGuardrails(CustomLogger): # The current chunk was appended to responses_so_far but not # yet yielded, so exclude it: the continuation must reflect # only what the client has actually received. - async for block_chunk in self._handle_streaming_block( + async for block_chunk in self.handle_streaming_block( e, endpoint_translation, stream_started=chunks_yielded, @@ -1124,7 +1154,7 @@ class UnifiedLLMGuardrails(CustomLogger): # terminating SSE sequence with the block message rather than # propagating into a bare error blob that truncates the stream. # The withheld original chunks are never released. - async for block_chunk in self._handle_streaming_block( + async for block_chunk in self.handle_streaming_block( e, endpoint_translation, stream_started=bool(responses_yielded), diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index 190784a2a60..c422a7c0964 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -6,7 +6,7 @@ pass/fail actions (allow, block, next, modify_response) and data forwarding. """ import time -from typing import Any, Final, Literal +from typing import TYPE_CHECKING, Any, Final, Literal import litellm from litellm._logging import verbose_proxy_logger @@ -25,6 +25,11 @@ from litellm.types.proxy.policy_engine.pipeline_types import ( PipelineStepResult, ) +if TYPE_CHECKING: + from litellm.llms.base_llm.guardrail_translation.base_translation import ( + BaseTranslation, + ) + try: from fastapi.exceptions import HTTPException except ImportError: @@ -43,6 +48,8 @@ class PipelineExecutor: call_type: str, policy_name: str, raw_request_snapshot: dict | None = None, # mutable-ok: same request-payload shape as data + streaming_chunks: list[Any] | None = None, # mutable-ok: shared buffered-stream chunks, read per step + endpoint_translation: "BaseTranslation | None" = None, ) -> PipelineExecutionResult: """ Execute pipeline steps sequentially with conditional actions. @@ -59,6 +66,12 @@ class PipelineExecutor: step whose guardrail opted into ``scan_raw_request`` evaluates the original request instead of whatever an earlier ``pass_data`` step in this same pipeline already rewrote. + streaming_chunks: buffered chunks of a completed stream. When set + (with ``endpoint_translation``), post_call steps scan the + assembled streamed output through the endpoint translation + instead of calling ``async_post_call_success_hook``. + endpoint_translation: the guardrail translation for the streamed + endpoint, resolved by the caller. Returns: PipelineExecutionResult with terminal action and step results @@ -83,6 +96,8 @@ class PipelineExecutor: user_api_key_dict=user_api_key_dict, call_type=call_type, raw_request_snapshot=raw_request_snapshot, + streaming_chunks=streaming_chunks, + endpoint_translation=endpoint_translation, ) duration = time.perf_counter() - start_time @@ -154,6 +169,8 @@ class PipelineExecutor: user_api_key_dict: Any, call_type: str, raw_request_snapshot: dict | None = None, # mutable-ok: same request-payload shape as data + streaming_chunks: list[Any] | None = None, # mutable-ok: shared buffered-stream chunks, read per step + endpoint_translation: "BaseTranslation | None" = None, ) -> tuple[ Literal["pass", "fail", "error"], dict | None, @@ -198,10 +215,8 @@ class PipelineExecutor: # Use unified_guardrail path if callback implements apply_guardrail target: CustomLogger = callback - use_unified: Final = ( - "apply_guardrail" in type(callback).__dict__ and not callback.use_native_lifecycle_hooks - ) - if use_unified: + use_unified: Final = PipelineExecutor.supports_unified_execution(callback) + if use_unified and streaming_chunks is None: hook_input["guardrail_to_apply"] = callback target = UnifiedLLMGuardrails() @@ -216,6 +231,22 @@ class PipelineExecutor: callback.mark_pre_call_hook_ran(data) if isinstance(response, dict): callback.mark_pre_call_hook_ran(response) + elif mode == "post_call" and streaming_chunks is not None: + if not use_unified or endpoint_translation is None: + return ( + "error", + None, + f"Guardrail '{step.guardrail}' does not support streaming pipeline execution", + None, + ) + await endpoint_translation.process_output_streaming_response( + responses_so_far=streaming_chunks, + guardrail_to_apply=callback, + litellm_logging_obj=data.get("litellm_logging_obj"), + user_api_key_dict=user_api_key_dict, + request_data=hook_input, + ) + response = None elif mode == "post_call": response = await target.async_post_call_success_hook( user_api_key_dict=user_api_key_dict, @@ -246,6 +277,12 @@ class PipelineExecutor: verbose_proxy_logger.error("Pipeline: unexpected error from guardrail '%s': %s", step.guardrail, e) return ("error", None, str(e), e) + @staticmethod + def supports_unified_execution(callback: CustomGuardrail) -> bool: + """Whether this guardrail runs through the unified apply_guardrail path, + the interface streaming pipeline execution requires.""" + return "apply_guardrail" in type(callback).__dict__ and not callback.use_native_lifecycle_hooks + @staticmethod def find_guardrail_callback(guardrail_name: str) -> CustomGuardrail | None: """Look up an initialized guardrail callback by name from litellm.callbacks.""" diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 80e991640a9..5642eb5383e 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -19,7 +19,7 @@ from email.mime.text import MIMEText from types import MappingProxyType from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, Protocol, TypeVar, Union, cast, overload -from typing_extensions import ReadOnly, TypedDict +from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm import _custom_logger_compatible_callbacks_literal from litellm.constants import ( @@ -486,29 +486,80 @@ def _merge_pipeline_metadata_writes(data: dict, modified_data: Mapping[str, obje _merge_pipeline_metadata_bucket(data, bucket_key, modified_data.get(bucket_key)) +def _pipeline_step_supports_streaming(guardrail_name: str) -> bool: + callback: Final = PipelineExecutor.find_guardrail_callback(guardrail_name) + return callback is not None and PipelineExecutor.supports_unified_execution(callback) + + +class _PipelineErrorBody(TypedDict): + message: ReadOnly[str] + type: ReadOnly[str] + policies: ReadOnly[tuple[str, ...]] + guardrails: NotRequired[ReadOnly[tuple[str, ...]]] + + +class _PipelineErrorDetail(TypedDict): + error: ReadOnly[_PipelineErrorBody] + + def _raise_for_streaming_post_call_pipelines(data: Mapping[str, object]) -> None: - if data.get("stream") is not True and data.get("background") is not True: + """ + Reject up front the requests whose post_call pipelines could never run. + + Background responses skip the post_call hooks entirely, so a pipeline + governing one would silently never execute. Streaming responses execute + pipelines against the buffered stream through the endpoint guardrail + translations, which requires every step's guardrail to support the unified + apply_guardrail interface; steps that cannot (native-lifecycle guardrails, + or guardrails not registered at all) keep the 400 rather than letting + ungoverned output stream through. + """ + is_stream: Final = data.get("stream") is True + is_background: Final = data.get("background") is True + if not is_stream and not is_background: return - post_call_policies: Final = tuple( - policy_name for policy_name, pipeline in _policy_pipelines(data) if pipeline.mode == "post_call" + post_call_pipelines: Final = tuple( + (policy_name, pipeline) for policy_name, pipeline in _policy_pipelines(data) if pipeline.mode == "post_call" ) - if not post_call_policies: + if not post_call_pipelines: return - raise HTTPException( - status_code=400, - detail={ + post_call_policies: Final = tuple(policy_name for policy_name, _pipeline in post_call_pipelines) + if is_background: + background_detail: Final[_PipelineErrorDetail] = { "error": { "message": ( - "Policies with post_call guardrail pipelines cannot govern streaming or background " - f"responses yet: {', '.join(post_call_policies)}. Retry with stream=false and " - "background=false, or move these policies' output guardrails from pipeline steps to " - "guardrails.add, which scans streamed output." + "Policies with post_call guardrail pipelines cannot govern background " + f"responses: {', '.join(post_call_policies)}. Retry with background=false." ), "type": "guardrail_pipeline_error", - "policies": list(post_call_policies), + "policies": post_call_policies, } - }, + } + raise HTTPException(status_code=400, detail=background_detail) + unsupported_guardrails: Final = tuple( + dict.fromkeys( + step.guardrail + for _policy_name, pipeline in post_call_pipelines + for step in pipeline.steps + if not _pipeline_step_supports_streaming(step.guardrail) + ) ) + if not unsupported_guardrails: + return + unsupported_detail: Final[_PipelineErrorDetail] = { + "error": { + "message": ( + "Policies with post_call guardrail pipelines cannot govern streaming responses " + "because these pipeline guardrails do not support the unified apply_guardrail " + f"interface: {', '.join(unsupported_guardrails)}. Retry with stream=false, or move " + "them from pipeline steps to guardrails.add, which scans streamed output." + ), + "type": "guardrail_pipeline_error", + "policies": post_call_policies, + "guardrails": unsupported_guardrails, + } + } + raise HTTPException(status_code=400, detail=unsupported_detail) def _prompt_block_text(block: object) -> str: @@ -1689,7 +1740,7 @@ class ProxyLogging: result: PipelineExecutionResult, data: dict, policy_name: str, - original_response: LLMResponseTypes | None = None, + original_response: "LLMResponseTypes | Sequence[object] | None" = None, ) -> dict: """ Handle a PipelineExecutionResult — allow, block, or modify_response. @@ -1699,7 +1750,9 @@ class ProxyLogging: payload (already sent upstream) must stay untouched; a replacement response carried in ``modified_data`` is adopted by the caller, and metadata-bucket writes (applied guardrails, guardrail logging info) - are merged back so headers and spend logs still see them. + are merged back so headers and spend logs still see them. On the + streaming path it is the buffered chunk list, carried into + ``ModifyResponseException.original_response`` for usage reporting. """ if result.terminal_action == "allow": if result.modified_data is not None: @@ -3195,11 +3248,16 @@ class ProxyLogging: # dict lookups + llm_router.get_deployment() per callback per chunk. _cached_guardrail_data: dict | None = None _guardrail_data_computed = False + pipeline_managed: Final = ( + _pipeline_managed_guardrail_names(data, "post_call") if caps.has_guardrail else frozenset() + ) for callback in litellm.callbacks: try: _callback: CustomLogger | None = None if isinstance(callback, CustomGuardrail): + if callback.guardrail_name in pipeline_managed: + continue # Main - V2 Guardrails implementation from litellm.types.guardrails import GuardrailEventHooks @@ -3256,12 +3314,17 @@ class ProxyLogging: 1. /chat/completions """ caps: Final = ProxyLogging._callback_capabilities() + post_call_pipelines: Final = tuple( + (policy_name, pipeline) + for policy_name, pipeline in _policy_pipelines(request_data) + if pipeline.mode == "post_call" + ) # Fast path: no real overrides. Internal proxy CustomLogger callbacks # (e.g. _PROXY_MaxBudgetLimiter, ManagedFiles) inherit the default # ``async for chunk: yield chunk`` body, so wrapping the iterator # through each of them adds N pass-through trampolines per chunk for # zero behavior change. Skip the chain entirely and stream through. - if not caps.iterator_overrides: + if not caps.iterator_overrides and not post_call_pipelines: try: async for chunk in response: yield chunk @@ -3281,8 +3344,11 @@ class ProxyLogging: current_response = response stream_needs_translation: Final = ProxyLogging._stream_requires_guardrail_translation(user_api_key_dict) + pipeline_managed_names: Final = _pipeline_managed_guardrail_names(request_data, "post_call") for resolved_callback, kind in caps.iterator_overrides: if isinstance(resolved_callback, CustomGuardrail): + if resolved_callback.guardrail_name in pipeline_managed_names: + continue if ( resolved_callback.should_run_guardrail(data=request_data, event_type=GuardrailEventHooks.post_call) is not True @@ -3322,6 +3388,17 @@ class ProxyLogging: ), ) + # Policy pipelines run last, over the fully buffered stream, so a + # pipeline verdict covers whatever the flat guardrail chain above + # already let through. + if post_call_pipelines: + current_response = self._pipeline_gated_stream( + response=current_response, + user_api_key_dict=user_api_key_dict, + request_data=request_data, + pipelines=post_call_pipelines, + ) + try: async for chunk in current_response: yield chunk @@ -3337,6 +3414,82 @@ class ProxyLogging: # we reach this point the metadata is fully populated. ProxyLogging._fire_deferred_stream_logging(request_data) + async def _pipeline_gated_stream( + self, + response: "AsyncGenerator[object, None]", + user_api_key_dict: UserAPIKeyAuth, + request_data: dict, + pipelines: "tuple[tuple[str, GuardrailPipeline], ...]", + ) -> "AsyncGenerator[Any, None]": + """ + Execute post_call policy pipelines against a streamed response. + + Buffers the whole stream (nothing reaches the client until every + pipeline allows it), then runs each pipeline's steps against the + assembled output through the endpoint guardrail translation, the same + machinery flat post_call guardrails use at end of stream. An allow + releases the buffered chunks verbatim; a block or modify_response + terminates with the translation's block chunks or the raised error. + """ + from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( + resolve_endpoint_translation, + ) + + buffered: Final[list[object]] = [] # mutable-ok: accumulates the stream before the pipeline verdict + async for item in response: + buffered.append(item) + if not buffered: + return + + resolved: Final = resolve_endpoint_translation(user_api_key_dict, buffered[0]) + if resolved is None: + policy_names: Final = tuple(policy_name for policy_name, _pipeline in pipelines) + unresolvable_detail: Final[_PipelineErrorDetail] = { + "error": { + "message": ( + "Policy pipelines could not govern this streaming response shape; " + f"the response was withheld: {', '.join(policy_names)}." + ), + "type": "guardrail_pipeline_error", + "policies": policy_names, + } + } + raise HTTPException(status_code=500, detail=unresolvable_detail) + call_type, endpoint_translation = resolved + + for policy_name, pipeline in pipelines: + result: PipelineExecutionResult = await PipelineExecutor.execute_steps( + steps=pipeline.steps, + mode="post_call", + data=request_data, + user_api_key_dict=user_api_key_dict, + call_type=call_type, + policy_name=policy_name, + streaming_chunks=buffered, + endpoint_translation=endpoint_translation, + ) + try: + ProxyLogging._handle_pipeline_result( + result, data=request_data, policy_name=policy_name, original_response=buffered + ) + except ModifyResponseException as e: + if e.original_response is None: + e.original_response = buffered + async for block_chunk in unified_guardrail.handle_streaming_block( + e, endpoint_translation, stream_started=False, responses_so_far=() + ): + yield block_chunk + return + except HTTPException as e: + async for error_chunk in unified_guardrail.emit_streaming_http_error( + e, call_type, buffered, request_data + ): + yield error_chunk + return + + for buffered_item in buffered: + yield buffered_item + @staticmethod def _fire_deferred_stream_logging(request_data: dict) -> None: """ diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py index 8b9ecfbbeee..fe4acbf8277 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py @@ -1026,7 +1026,7 @@ class TestStreamingTransform: ) emitted = [] - async for item in handler._emit_streaming_http_error( + async for item in handler.emit_streaming_http_error( exc, call_type=CallTypes.asend_message.value, responses_so_far=[{"id": "req-1"}], diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py index 34a25d75bc0..2b36ae111dc 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -1284,7 +1284,8 @@ async def test_pre_call_hook_rejects_streaming_request_with_post_call_pipeline( ) assert info.value.status_code == 400 - assert info.value.detail["error"]["policies"] == ["response-governance"] + assert info.value.detail["error"]["policies"] == ("response-governance",) + assert info.value.detail["error"]["guardrails"] == ("gr-post",) assert "stream=false" in info.value.detail["error"]["message"] @@ -1304,7 +1305,7 @@ async def test_pre_call_hook_rejects_background_request_with_post_call_pipeline( ) assert info.value.status_code == 400 - assert info.value.detail["error"]["policies"] == ["response-governance"] + assert info.value.detail["error"]["policies"] == ("response-governance",) assert "background=false" in info.value.detail["error"]["message"] @@ -1333,3 +1334,166 @@ def test_raise_for_streaming_post_call_pipelines_ignores_non_streaming_and_pre_c ) assert _raise_for_streaming_post_call_pipelines({"stream": True}) is None assert _raise_for_streaming_post_call_pipelines({"background": True}) is None + + +# --------------------------------------------------------------------------- +# post_call pipelines on streaming responses +# --------------------------------------------------------------------------- + + +def _unified_stream_guardrail(seen: Dict[str, Any], block: bool = False) -> CustomGuardrail: + class UnifiedStreamGuardrail(CustomGuardrail): + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + seen["count"] = seen.get("count", 0) + 1 + seen["input_type"] = input_type + if block: + raise HTTPException(status_code=400, detail={"error": "output blocked"}) + return inputs + + return UnifiedStreamGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False) + + +def _stream_chunks() -> List[Any]: + return [ + litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": "hello "}, "finish_reason": None}]), + litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": "world"}, "finish_reason": "stop"}]), + ] + + +async def _async_chunk_iter(chunks: List[Any]): + for chunk in chunks: + yield chunk + + +@pytest.mark.asyncio +async def test_pre_call_hook_allows_streaming_when_pipeline_guardrail_supports_unified( + proxy_logging, make_user_api_key_auth, monkeypatch +): + seen: Dict[str, Any] = {} + monkeypatch.setattr(litellm, "callbacks", [_unified_stream_guardrail(seen)]) + data = _post_call_pipeline_data(stream=True) + + out = await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data=data, + call_type="completion", + guardrails_only=True, + ) + + assert out is not None + assert out.get("stream") is True + + +@pytest.mark.asyncio +@pytest.mark.parametrize("native_lifecycle", [False, True]) +async def test_pre_call_hook_rejects_streaming_when_pipeline_guardrail_lacks_unified_support( + proxy_logging, make_user_api_key_auth, monkeypatch, native_lifecycle +): + if native_lifecycle: + + class NativeOnlyGuardrail(CustomGuardrail): + use_native_lifecycle_hooks = True + + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + return inputs + + else: + + class NativeOnlyGuardrail(CustomGuardrail): + pass + + monkeypatch.setattr( + litellm, + "callbacks", + [NativeOnlyGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False)], + ) + data = _post_call_pipeline_data(stream=True) + + with pytest.raises(HTTPException) as info: + await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data=data, + call_type="completion", + guardrails_only=True, + ) + + assert info.value.status_code == 400 + assert info.value.detail["error"]["guardrails"] == ("gr-post",) + assert "apply_guardrail" in info.value.detail["error"]["message"] + + +@pytest.mark.asyncio +async def test_streaming_iterator_hook_pipeline_allow_releases_buffered_chunks( + proxy_logging, make_user_api_key_auth, monkeypatch +): + seen: Dict[str, Any] = {} + monkeypatch.setattr(litellm, "callbacks", [_unified_stream_guardrail(seen)]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data(stream=True) + chunks = _stream_chunks() + + delivered = [ + item + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"), + response=_async_chunk_iter(chunks), + request_data=data, + ) + ] + + assert [id(item) for item in delivered] == [id(chunk) for chunk in chunks] + assert seen["count"] == 1 + assert seen["input_type"] == "response" + + +@pytest.mark.asyncio +async def test_streaming_iterator_hook_pipeline_block_withholds_all_chunks( + proxy_logging, make_user_api_key_auth, monkeypatch +): + seen: Dict[str, Any] = {} + monkeypatch.setattr(litellm, "callbacks", [_unified_stream_guardrail(seen, block=True)]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data(stream=True) + delivered: List[Any] = [] + + async def _drain() -> None: + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"), + response=_async_chunk_iter(_stream_chunks()), + request_data=data, + ): + delivered.append(item) + + with pytest.raises(HTTPException) as info: + await _drain() + + assert delivered == [] + assert info.value.status_code == 400 + assert "output blocked" in str(info.value.detail) + + +@pytest.mark.asyncio +async def test_streaming_iterator_hook_pipeline_withholds_unresolvable_response_shape( + proxy_logging, make_user_api_key_auth, monkeypatch +): + seen: Dict[str, Any] = {} + monkeypatch.setattr(litellm, "callbacks", [_unified_stream_guardrail(seen)]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data(stream=True) + delivered: List[Any] = [] + + async def _drain() -> None: + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=make_user_api_key_auth(), + response=_async_chunk_iter([object(), object()]), + request_data=data, + ): + delivered.append(item) + + with pytest.raises(HTTPException) as info: + await _drain() + + assert delivered == [] + assert info.value.status_code == 500 + assert "withheld" in info.value.detail["error"]["message"] + assert seen.get("count") is None From fa5a10941e713968fbe7251a99d867ef0050dd69 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 13:28:11 -0700 Subject: [PATCH 08/60] fix(policy_engine): fail closed on streaming for content-rewriting pipeline steps and untranslatable routes --- .../unified_guardrail/unified_guardrail.py | 19 ++- litellm/proxy/utils.py | 111 ++++++++++++------ .../proxy_logging/test_guardrail_pipeline.py | 84 +++++++++++-- 3 files changed, 155 insertions(+), 59 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py index 60b4444e1d4..89527c05eb6 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -876,6 +876,14 @@ class UnifiedLLMGuardrails(CustomLogger): choices: Final = _chunk_choices(item) return any(getattr(choice, "finish_reason", None) is not None for choice in choices) + def resolve_streaming_flag(self, guardrail_to_apply: CustomGuardrail | None, name: str, default: object) -> object: + """Streaming flag resolution order (later wins): default < guardrail + attribute < guardrail_config dict < this callback's optional_params.""" + attribute_value: Final = default if guardrail_to_apply is None else getattr(guardrail_to_apply, name, default) + config: Final = None if guardrail_to_apply is None else getattr(guardrail_to_apply, "guardrail_config", None) + config_value: Final = config.get(name, attribute_value) if isinstance(config, dict) else attribute_value + return self.optional_params.get(name, config_value) + async def async_post_call_streaming_iterator_hook( self, user_api_key_dict: UserAPIKeyAuth, @@ -906,17 +914,8 @@ class UnifiedLLMGuardrails(CustomLogger): if guardrail_to_apply is None: guardrail_to_apply = request_data.pop("guardrail_to_apply", None) - # Get streaming configuration. Resolution order (later wins): default - # < guardrail attribute < guardrail_config dict < this callback's - # optional_params. def _streaming_flag(name: str, default: object) -> Any: - value = default - if guardrail_to_apply is not None: - value = getattr(guardrail_to_apply, name, value) - config: Final[Mapping[str, object]] = getattr(guardrail_to_apply, "guardrail_config", {}) - if isinstance(config, dict): - value = config.get(name, value) - return self.optional_params.get(name, value) + return self.resolve_streaming_flag(guardrail_to_apply, name, default) sampling_rate: Final[int] = _streaming_flag("streaming_sampling_rate", 5) # Only apply the guardrail at end of stream (not per chunk). diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 5642eb5383e..8b2c38d457f 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -139,6 +139,7 @@ from litellm.proxy.db.token_auth import ( ) from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( UnifiedLLMGuardrails, + resolve_endpoint_translation, ) from litellm.proxy.hooks import PROXY_HOOKS, get_proxy_hook from litellm.proxy.hooks.cache_control_check import _PROXY_CacheControlCheck @@ -486,11 +487,19 @@ def _merge_pipeline_metadata_writes(data: dict, modified_data: Mapping[str, obje _merge_pipeline_metadata_bucket(data, bucket_key, modified_data.get(bucket_key)) -def _pipeline_step_supports_streaming(guardrail_name: str) -> bool: +def _pipeline_step_supports_unified_streaming(guardrail_name: str) -> bool: callback: Final = PipelineExecutor.find_guardrail_callback(guardrail_name) return callback is not None and PipelineExecutor.supports_unified_execution(callback) +def _pipeline_step_rewrites_streamed_content(guardrail_name: str) -> bool: + callback: Final = PipelineExecutor.find_guardrail_callback(guardrail_name) + if callback is None: + return False + transform_mode: Final = unified_guardrail.resolve_streaming_flag(callback, "streaming_transform_mode", "block_only") + return callback.mask_response_content or transform_mode == "incremental_diff" + + class _PipelineErrorBody(TypedDict): message: ReadOnly[str] type: ReadOnly[str] @@ -502,17 +511,20 @@ class _PipelineErrorDetail(TypedDict): error: ReadOnly[_PipelineErrorBody] -def _raise_for_streaming_post_call_pipelines(data: Mapping[str, object]) -> None: +def _raise_for_streaming_post_call_pipelines(data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth) -> None: """ Reject up front the requests whose post_call pipelines could never run. Background responses skip the post_call hooks entirely, so a pipeline governing one would silently never execute. Streaming responses execute pipelines against the buffered stream through the endpoint guardrail - translations, which requires every step's guardrail to support the unified - apply_guardrail interface; steps that cannot (native-lifecycle guardrails, - or guardrails not registered at all) keep the 400 rather than letting - ungoverned output stream through. + translation of the request route, releasing the buffered chunks verbatim + on allow. That needs every step's guardrail to support the unified + apply_guardrail interface and to only allow or block (a step that rewrites + streamed content, via mask_response_content or + streaming_transform_mode=incremental_diff, would have its rewrite silently + dropped), and needs the route to have a translation at all; anything else + keeps the 400 rather than letting ungoverned output stream through. """ is_stream: Final = data.get("stream") is True is_background: Final = data.get("background") is True @@ -536,30 +548,61 @@ def _raise_for_streaming_post_call_pipelines(data: Mapping[str, object]) -> None } } raise HTTPException(status_code=400, detail=background_detail) - unsupported_guardrails: Final = tuple( - dict.fromkeys( - step.guardrail - for _policy_name, pipeline in post_call_pipelines - for step in pipeline.steps - if not _pipeline_step_supports_streaming(step.guardrail) - ) + step_guardrails: Final = tuple( + dict.fromkeys(step.guardrail for _policy_name, pipeline in post_call_pipelines for step in pipeline.steps) ) - if not unsupported_guardrails: + unsupported_guardrails: Final = tuple( + guardrail for guardrail in step_guardrails if not _pipeline_step_supports_unified_streaming(guardrail) + ) + if unsupported_guardrails: + unsupported_detail: Final[_PipelineErrorDetail] = { + "error": { + "message": ( + "Policies with post_call guardrail pipelines cannot govern streaming responses " + "because these pipeline guardrails do not support the unified apply_guardrail " + f"interface: {', '.join(unsupported_guardrails)}. Retry with stream=false, or drop " + "them from the pipeline steps so guardrails.add scans them on streamed output." + ), + "type": "guardrail_pipeline_error", + "policies": post_call_policies, + "guardrails": unsupported_guardrails, + } + } + raise HTTPException(status_code=400, detail=unsupported_detail) + rewriting_guardrails: Final = tuple( + guardrail for guardrail in step_guardrails if _pipeline_step_rewrites_streamed_content(guardrail) + ) + if rewriting_guardrails: + rewriting_detail: Final[_PipelineErrorDetail] = { + "error": { + "message": ( + "Policies with post_call guardrail pipelines cannot govern streaming responses " + "because these pipeline guardrails rewrite streamed content (mask_response_content " + "or streaming_transform_mode=incremental_diff), which pipeline steps would release " + f"unmodified: {', '.join(rewriting_guardrails)}. Retry with stream=false, or drop " + "them from the pipeline steps so guardrails.add applies them to streamed output." + ), + "type": "guardrail_pipeline_error", + "policies": post_call_policies, + "guardrails": rewriting_guardrails, + } + } + raise HTTPException(status_code=400, detail=rewriting_detail) + route: Final = user_api_key_dict.request_route + if not route or resolve_endpoint_translation(user_api_key_dict, None) is not None: return - unsupported_detail: Final[_PipelineErrorDetail] = { + route_detail: Final[_PipelineErrorDetail] = { "error": { "message": ( - "Policies with post_call guardrail pipelines cannot govern streaming responses " - "because these pipeline guardrails do not support the unified apply_guardrail " - f"interface: {', '.join(unsupported_guardrails)}. Retry with stream=false, or move " - "them from pipeline steps to guardrails.add, which scans streamed output." + "Policies with post_call guardrail pipelines cannot govern streaming responses on " + f"route {route} because it has no endpoint guardrail translation to scan the stream " + f"through: {', '.join(post_call_policies)}. Retry with stream=false." ), "type": "guardrail_pipeline_error", "policies": post_call_policies, - "guardrails": unsupported_guardrails, } } - raise HTTPException(status_code=400, detail=unsupported_detail) + raise HTTPException(status_code=400, detail=route_detail) def _prompt_block_text(block: object) -> str: @@ -1915,7 +1958,7 @@ class ProxyLogging: ) try: - _raise_for_streaming_post_call_pipelines(data) + _raise_for_streaming_post_call_pipelines(data, user_api_key_dict) # Execute guardrail pipelines before the normal callback loop data, _ = await self._maybe_execute_pipelines( @@ -3431,10 +3474,6 @@ class ProxyLogging: releases the buffered chunks verbatim; a block or modify_response terminates with the translation's block chunks or the raised error. """ - from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( - resolve_endpoint_translation, - ) - buffered: Final[list[object]] = [] # mutable-ok: accumulates the stream before the pipeline verdict async for item in response: buffered.append(item) @@ -3444,17 +3483,15 @@ class ProxyLogging: resolved: Final = resolve_endpoint_translation(user_api_key_dict, buffered[0]) if resolved is None: policy_names: Final = tuple(policy_name for policy_name, _pipeline in pipelines) - unresolvable_detail: Final[_PipelineErrorDetail] = { - "error": { - "message": ( - "Policy pipelines could not govern this streaming response shape; " - f"the response was withheld: {', '.join(policy_names)}." - ), - "type": "guardrail_pipeline_error", - "policies": policy_names, - } - } - raise HTTPException(status_code=500, detail=unresolvable_detail) + raise ProxyException( + message=( + "Policy pipelines could not govern this streaming response shape; " + f"the response was withheld: {', '.join(policy_names)}." + ), + type="guardrail_pipeline_error", + param=None, + code=500, + ) call_type, endpoint_translation = resolved for policy_name, pipeline in pipelines: diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py index 2b36ae111dc..6c90637158e 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -23,6 +23,7 @@ from litellm.integrations.custom_guardrail import ( ModifyResponseException, ) from litellm.integrations.prometheus import PrometheusLogger +from litellm.proxy._types import ProxyException from litellm.proxy.common_utils.callback_utils import add_guardrail_to_applied_guardrails_header from litellm.proxy.utils import ProxyLogging, _raise_for_streaming_post_call_pipelines from litellm.types.guardrails import GuardrailEventHooks @@ -1309,31 +1310,35 @@ async def test_pre_call_hook_rejects_background_request_with_post_call_pipeline( assert "background=false" in info.value.detail["error"]["message"] -def test_raise_for_streaming_post_call_pipelines_ignores_non_streaming_and_pre_call(): +def test_raise_for_streaming_post_call_pipelines_ignores_non_streaming_and_pre_call(make_user_api_key_auth): post_call = GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="g", on_fail="block")]) pre_call = GuardrailPipeline(mode="pre_call", steps=[PipelineStep(guardrail="g", on_fail="block")]) + auth = make_user_api_key_auth(request_route="/custom/stream") assert ( _raise_for_streaming_post_call_pipelines( - {"stream": False, "metadata": {"_guardrail_pipelines": [("p", post_call)]}} + {"stream": False, "metadata": {"_guardrail_pipelines": [("p", post_call)]}}, auth ) is None ) assert ( _raise_for_streaming_post_call_pipelines( - {"background": False, "metadata": {"_guardrail_pipelines": [("p", post_call)]}} + {"background": False, "metadata": {"_guardrail_pipelines": [("p", post_call)]}}, auth ) is None ) - assert _raise_for_streaming_post_call_pipelines({"metadata": {"_guardrail_pipelines": [("p", post_call)]}}) is None + assert ( + _raise_for_streaming_post_call_pipelines({"metadata": {"_guardrail_pipelines": [("p", post_call)]}}, auth) + is None + ) assert ( _raise_for_streaming_post_call_pipelines( - {"stream": True, "metadata": {"_guardrail_pipelines": [("p", pre_call)]}} + {"stream": True, "metadata": {"_guardrail_pipelines": [("p", pre_call)]}}, auth ) is None ) - assert _raise_for_streaming_post_call_pipelines({"stream": True}) is None - assert _raise_for_streaming_post_call_pipelines({"background": True}) is None + assert _raise_for_streaming_post_call_pipelines({"stream": True}, auth) is None + assert _raise_for_streaming_post_call_pipelines({"background": True}, auth) is None # --------------------------------------------------------------------------- @@ -1366,15 +1371,16 @@ async def _async_chunk_iter(chunks: List[Any]): @pytest.mark.asyncio +@pytest.mark.parametrize("request_route", [None, "/v1/chat/completions"]) async def test_pre_call_hook_allows_streaming_when_pipeline_guardrail_supports_unified( - proxy_logging, make_user_api_key_auth, monkeypatch + proxy_logging, make_user_api_key_auth, monkeypatch, request_route ): seen: Dict[str, Any] = {} monkeypatch.setattr(litellm, "callbacks", [_unified_stream_guardrail(seen)]) data = _post_call_pipeline_data(stream=True) out = await proxy_logging.pre_call_hook( - user_api_key_dict=make_user_api_key_auth(), + user_api_key_dict=make_user_api_key_auth(request_route=request_route), data=data, call_type="completion", guardrails_only=True, @@ -1422,6 +1428,60 @@ async def test_pre_call_hook_rejects_streaming_when_pipeline_guardrail_lacks_uni assert "apply_guardrail" in info.value.detail["error"]["message"] +@pytest.mark.asyncio +@pytest.mark.parametrize( + "rewrite_attribute, value", + [ + ("mask_response_content", True), + ("streaming_transform_mode", "incremental_diff"), + ("guardrail_config", {"streaming_transform_mode": "incremental_diff"}), + ], +) +async def test_pre_call_hook_rejects_streaming_when_pipeline_guardrail_rewrites_streamed_content( + proxy_logging, make_user_api_key_auth, monkeypatch, rewrite_attribute, value +): + seen: Dict[str, Any] = {} + guardrail = _unified_stream_guardrail(seen) + setattr(guardrail, rewrite_attribute, value) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + data = _post_call_pipeline_data(stream=True) + + with pytest.raises(HTTPException) as info: + await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"), + data=data, + call_type="completion", + guardrails_only=True, + ) + + assert info.value.status_code == 400 + assert info.value.detail["error"]["guardrails"] == ("gr-post",) + assert "rewrite streamed content" in info.value.detail["error"]["message"] + assert seen.get("count") is None + + +@pytest.mark.asyncio +async def test_pre_call_hook_rejects_streaming_when_route_has_no_guardrail_translation( + proxy_logging, make_user_api_key_auth, monkeypatch +): + seen: Dict[str, Any] = {} + monkeypatch.setattr(litellm, "callbacks", [_unified_stream_guardrail(seen)]) + data = _post_call_pipeline_data(stream=True) + + with pytest.raises(HTTPException) as info: + await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/custom/stream"), + data=data, + call_type="completion", + guardrails_only=True, + ) + + assert info.value.status_code == 400 + assert info.value.detail["error"]["policies"] == ("response-governance",) + assert "/custom/stream" in info.value.detail["error"]["message"] + assert seen.get("count") is None + + @pytest.mark.asyncio async def test_streaming_iterator_hook_pipeline_allow_releases_buffered_chunks( proxy_logging, make_user_api_key_auth, monkeypatch @@ -1490,10 +1550,10 @@ async def test_streaming_iterator_hook_pipeline_withholds_unresolvable_response_ ): delivered.append(item) - with pytest.raises(HTTPException) as info: + with pytest.raises(ProxyException) as info: await _drain() assert delivered == [] - assert info.value.status_code == 500 - assert "withheld" in info.value.detail["error"]["message"] + assert info.value.code == "500" + assert "withheld" in info.value.message assert seen.get("count") is None From 1bed9bae43a7c28e46d6b38d4b58d10d35d87c58 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 13:41:33 -0700 Subject: [PATCH 09/60] chore(policy_engine): drop control-flow comment flagged in review --- litellm/proxy/utils.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 8b2c38d457f..231d4f18aa0 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -3431,9 +3431,6 @@ class ProxyLogging: ), ) - # Policy pipelines run last, over the fully buffered stream, so a - # pipeline verdict covers whatever the flat guardrail chain above - # already let through. if post_call_pipelines: current_response = self._pipeline_gated_stream( response=current_response, From d51198fdeb3ccc7fcf70fbb116f94d8f360ef4d6 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 14:07:34 -0700 Subject: [PATCH 10/60] test(policy_engine): cover streaming pipeline gate branches Adds regression tests for the modify_response block on the Anthropic route, the gate with no iterator overrides, and the per-chunk hook skipping pipeline-managed guardrails. Corrects the gate docstring: an allow releases the chunks as the endpoint translation left them, not verbatim --- litellm/proxy/utils.py | 7 +- .../proxy_logging/test_guardrail_pipeline.py | 110 ++++++++++++++++++ 2 files changed, 115 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 231d4f18aa0..a64b57c1388 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -3468,8 +3468,11 @@ class ProxyLogging: pipeline allows it), then runs each pipeline's steps against the assembled output through the endpoint guardrail translation, the same machinery flat post_call guardrails use at end of stream. An allow - releases the buffered chunks verbatim; a block or modify_response - terminates with the translation's block chunks or the raised error. + releases the buffered chunks as that machinery left them (the + Responses and A2A translations write guardrail output back into the + final chunk, exactly as they do for flat guardrails); a block or + modify_response terminates with the translation's block chunks or the + raised error. """ buffered: Final[list[object]] = [] # mutable-ok: accumulates the stream before the pipeline verdict async for item in response: diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py index 6c90637158e..a258442c79a 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -10,6 +10,7 @@ Covers ``_should_use_guardrail_load_balancing``, ``_execute_guardrail_hook``, from __future__ import annotations import asyncio +import json from typing import Any, Dict, List from unittest.mock import AsyncMock, MagicMock, patch @@ -1557,3 +1558,112 @@ async def test_streaming_iterator_hook_pipeline_withholds_unresolvable_response_ assert info.value.code == "500" assert "withheld" in info.value.message assert seen.get("count") is None + + +def _anthropic_sse_chunks() -> List[bytes]: + events = [ + ("message_start", {"type": "message_start", "message": {"id": "msg_1", "type": "message", "role": "assistant", "model": "m", "content": [], "stop_reason": None, "usage": {"input_tokens": 1, "output_tokens": 0}}}), + ("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}), + ("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "hello world"}}), + ("content_block_stop", {"type": "content_block_stop", "index": 0}), + ("message_delta", {"type": "message_delta", "delta": {"stop_reason": "end_turn", "stop_sequence": None}, "usage": {"output_tokens": 2}}), + ("message_stop", {"type": "message_stop"}), + ] + return [f"event: {name}\ndata: {json.dumps(payload)}\n\n".encode() for name, payload in events] + + +@pytest.mark.asyncio +async def test_streaming_iterator_hook_pipeline_modify_response_emits_translated_block( + proxy_logging, make_user_api_key_auth, monkeypatch +): + seen: Dict[str, Any] = {} + monkeypatch.setattr(litellm, "callbacks", [_unified_stream_guardrail(seen, block=True)]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + pipeline = GuardrailPipeline( + mode="post_call", + steps=[ + PipelineStep( + guardrail="gr-post", + on_pass="allow", + on_fail="modify_response", + modify_response_message="content policy block", + ) + ], + ) + data = _post_call_pipeline_data(stream=True) + data["metadata"]["_guardrail_pipelines"] = [("response-governance", pipeline)] + chunks = _anthropic_sse_chunks() + + delivered = [ + item + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/v1/messages"), + response=_async_chunk_iter(chunks), + request_data=data, + ) + ] + + raw = b"".join(delivered).decode() + assert seen["count"] == 1 + assert "content policy block" in raw + assert "hello world" not in raw + assert not any(item is chunk for item in delivered for chunk in chunks) + + +@pytest.mark.asyncio +async def test_streaming_iterator_hook_pipeline_gates_without_iterator_overrides( + proxy_logging, make_user_api_key_auth, monkeypatch +): + monkeypatch.setattr(litellm, "callbacks", []) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data(stream=True) + delivered: List[Any] = [] + + async def _drain() -> None: + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"), + response=_async_chunk_iter(_stream_chunks()), + request_data=data, + ): + delivered.append(item) + + with pytest.raises(HTTPException) as info: + await _drain() + + assert delivered == [] + assert info.value.status_code == 400 + assert info.value.detail["error"]["pipeline_context"]["step_results"] == [ + {"guardrail": "gr-post", "outcome": "error", "action": "block"} + ] + + +@pytest.mark.asyncio +async def test_per_chunk_streaming_hook_skips_pipeline_managed_guardrail( + proxy_logging, make_user_api_key_auth, monkeypatch +): + seen: Dict[str, Any] = {} + + class RecordingGuardrail(CustomGuardrail): + async def async_post_call_streaming_hook(self, user_api_key_dict, response): + seen[self.guardrail_name] = seen.get(self.guardrail_name, 0) + 1 + return None + + managed = RecordingGuardrail( + guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=True + ) + free = RecordingGuardrail( + guardrail_name="gr-free", event_hook=GuardrailEventHooks.post_call, default_on=True + ) + monkeypatch.setattr(litellm, "callbacks", [managed, free]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data(stream=True) + + result = await proxy_logging.async_post_call_streaming_hook( + data=data, + response=_stream_chunks()[0], + user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"), + ) + + assert result is not None + assert seen.get("gr-post") is None + assert seen["gr-free"] == 1 From 0c1f33dff7084559cf1612a38b5dde4ea0afe9b8 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 14:19:21 -0700 Subject: [PATCH 11/60] fix(policy_engine): fail closed on content filter MASK steps for streaming pipelines A litellm_content_filter step with a MASK action masks chat streams through its own iterator hook under guardrails.add, which pipeline-managed guardrails skip, so the pipeline path released the stream unmasked. CustomGuardrail now declares rewrites_streamed_output (mask_response_content by default, any MASK action for the content filter) and the upfront streaming check names such steps in the same 400 it gives mask_response_content and incremental_diff --- litellm/integrations/custom_guardrail.py | 3 ++ .../litellm_content_filter/content_filter.py | 7 ++++ litellm/proxy/utils.py | 14 ++++---- .../content_filter/test_content_filter.py | 36 +++++++++++++++++++ .../proxy_logging/test_guardrail_pipeline.py | 34 +++++++++++++++++- 5 files changed, 86 insertions(+), 8 deletions(-) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 8dc6881d23e..a6e3d000120 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -762,6 +762,9 @@ class CustomGuardrail(CustomLogger): def uses_apply_guardrail_interface(self) -> bool: return type(self).apply_guardrail is not CustomGuardrail.apply_guardrail + def rewrites_streamed_output(self) -> bool: + return self.mask_response_content + def _deployment_pre_call_target(self) -> "CustomLogger": if not self.uses_apply_guardrail_interface() or self.use_native_lifecycle_hooks: return self diff --git a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py index 722f96ef814..bd31882841e 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py +++ b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py @@ -1947,6 +1947,13 @@ class ContentFilterGuardrail(CustomGuardrail): exception_str=exception_str, ) + def rewrites_streamed_output(self) -> bool: + return ( + super().rewrites_streamed_output() + or any(entry["action"] == ContentFilterAction.MASK for entry in self.compiled_patterns) + or any(action == ContentFilterAction.MASK for action, _ in self.blocked_words.values()) + ) + async def async_post_call_streaming_iterator_hook( self, user_api_key_dict: UserAPIKeyAuth, diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index a64b57c1388..138f272af42 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -497,7 +497,7 @@ def _pipeline_step_rewrites_streamed_content(guardrail_name: str) -> bool: if callback is None: return False transform_mode: Final = unified_guardrail.resolve_streaming_flag(callback, "streaming_transform_mode", "block_only") - return callback.mask_response_content or transform_mode == "incremental_diff" + return callback.rewrites_streamed_output() or transform_mode == "incremental_diff" class _PipelineErrorBody(TypedDict): @@ -518,10 +518,10 @@ def _raise_for_streaming_post_call_pipelines(data: Mapping[str, object], user_ap Background responses skip the post_call hooks entirely, so a pipeline governing one would silently never execute. Streaming responses execute pipelines against the buffered stream through the endpoint guardrail - translation of the request route, releasing the buffered chunks verbatim - on allow. That needs every step's guardrail to support the unified - apply_guardrail interface and to only allow or block (a step that rewrites - streamed content, via mask_response_content or + translation of the request route, releasing the buffered chunks on allow. + That needs every step's guardrail to support the unified apply_guardrail + interface and to only allow or block (a step that rewrites streamed + content, via mask_response_content, a MASK action, or streaming_transform_mode=incremental_diff, would have its rewrite silently dropped), and needs the route to have a translation at all; anything else keeps the 400 rather than letting ungoverned output stream through. @@ -577,8 +577,8 @@ def _raise_for_streaming_post_call_pipelines(data: Mapping[str, object], user_ap "error": { "message": ( "Policies with post_call guardrail pipelines cannot govern streaming responses " - "because these pipeline guardrails rewrite streamed content (mask_response_content " - "or streaming_transform_mode=incremental_diff), which pipeline steps would release " + "because these pipeline guardrails rewrite streamed content (mask_response_content, " + "a MASK action, or streaming_transform_mode=incremental_diff), which pipeline steps would release " f"unmodified: {', '.join(rewriting_guardrails)}. Retry with stream=false, or drop " "them from the pipeline steps so guardrails.add applies them to streamed output." ), diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py index be55ac47bde..ffbedfa43ff 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py @@ -3068,3 +3068,39 @@ class TestContentFilterToolCallArguments: request_data={}, input_type="response", ) + + +class TestRewritesStreamedOutput: + def test_block_only_rules_do_not_rewrite(self): + guardrail = ContentFilterGuardrail( + guardrail_name="cf", + patterns=[ContentFilterPattern(pattern_type="prebuilt", pattern_name="us_ssn", action=ContentFilterAction.BLOCK)], + blocked_words=[BlockedWord(keyword="kumquat", action=ContentFilterAction.BLOCK)], + ) + + assert guardrail.rewrites_streamed_output() is False + + def test_mask_blocked_word_rewrites(self): + guardrail = ContentFilterGuardrail( + guardrail_name="cf", + blocked_words=[BlockedWord(keyword="persimmon", action=ContentFilterAction.MASK)], + ) + + assert guardrail.rewrites_streamed_output() is True + + def test_mask_pattern_rewrites(self): + guardrail = ContentFilterGuardrail( + guardrail_name="cf", + patterns=[ContentFilterPattern(pattern_type="prebuilt", pattern_name="us_ssn", action=ContentFilterAction.MASK)], + ) + + assert guardrail.rewrites_streamed_output() is True + + def test_mask_response_content_rewrites(self): + guardrail = ContentFilterGuardrail( + guardrail_name="cf", + blocked_words=[BlockedWord(keyword="kumquat", action=ContentFilterAction.BLOCK)], + mask_response_content=True, + ) + + assert guardrail.rewrites_streamed_output() is True diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py index a258442c79a..ee0a9a10172 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -27,7 +27,8 @@ from litellm.integrations.prometheus import PrometheusLogger from litellm.proxy._types import ProxyException from litellm.proxy.common_utils.callback_utils import add_guardrail_to_applied_guardrails_header from litellm.proxy.utils import ProxyLogging, _raise_for_streaming_post_call_pipelines -from litellm.types.guardrails import GuardrailEventHooks +from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ContentFilterGuardrail +from litellm.types.guardrails import BlockedWord, ContentFilterAction, GuardrailEventHooks from litellm.types.proxy.policy_engine.pipeline_types import ( GuardrailPipeline, PipelineStep, @@ -1461,6 +1462,37 @@ async def test_pre_call_hook_rejects_streaming_when_pipeline_guardrail_rewrites_ assert seen.get("count") is None +@pytest.mark.asyncio +@pytest.mark.parametrize("action, rejected", [(ContentFilterAction.MASK, True), (ContentFilterAction.BLOCK, False)]) +async def test_pre_call_hook_rejects_streaming_only_when_content_filter_step_masks( + proxy_logging, make_user_api_key_auth, monkeypatch, action, rejected +): + guardrail = ContentFilterGuardrail( + guardrail_name="gr-post", + event_hook=GuardrailEventHooks.post_call, + blocked_words=[BlockedWord(keyword="persimmon", action=action)], + ) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + data = _post_call_pipeline_data(stream=True) + user_api_key_dict = make_user_api_key_auth(request_route="/v1/chat/completions") + + if not rejected: + out = await proxy_logging.pre_call_hook( + user_api_key_dict=user_api_key_dict, data=data, call_type="completion", guardrails_only=True + ) + assert out is not None and out.get("stream") is True + return + + with pytest.raises(HTTPException) as info: + await proxy_logging.pre_call_hook( + user_api_key_dict=user_api_key_dict, data=data, call_type="completion", guardrails_only=True + ) + + assert info.value.status_code == 400 + assert info.value.detail["error"]["guardrails"] == ("gr-post",) + assert "a MASK action" in info.value.detail["error"]["message"] + + @pytest.mark.asyncio async def test_pre_call_hook_rejects_streaming_when_route_has_no_guardrail_translation( proxy_logging, make_user_api_key_auth, monkeypatch From 90c8031dd76c5565c25b2adfe301a4ce715a6964 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 14:32:35 -0700 Subject: [PATCH 12/60] fix(policy_engine): fail closed on content filter category MASK steps for streaming pipelines --- .../litellm_content_filter/content_filter.py | 2 ++ .../content_filter/test_content_filter.py | 20 +++++++++++++++++ .../proxy_logging/test_guardrail_pipeline.py | 22 +++++++++++++++++++ 3 files changed, 44 insertions(+) diff --git a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py index bd31882841e..85eb50c78e7 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py +++ b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py @@ -1952,6 +1952,8 @@ class ContentFilterGuardrail(CustomGuardrail): super().rewrites_streamed_output() or any(entry["action"] == ContentFilterAction.MASK for entry in self.compiled_patterns) or any(action == ContentFilterAction.MASK for action, _ in self.blocked_words.values()) + or any(action == ContentFilterAction.MASK for _, _, action in self.category_keywords.values()) + or any(action == ContentFilterAction.MASK for _, _, action in self.always_block_category_keywords.values()) ) async def async_post_call_streaming_iterator_hook( diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py index ffbedfa43ff..73020fe3e6f 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py @@ -3104,3 +3104,23 @@ class TestRewritesStreamedOutput: ) assert guardrail.rewrites_streamed_output() is True + + @pytest.mark.parametrize("action, expected", [("MASK", True), ("BLOCK", False)]) + def test_category_keywords_follow_the_category_action(self, action, expected): + guardrail = ContentFilterGuardrail( + guardrail_name="cf", + categories=[{"category": "bias_gender", "enabled": True, "action": action}], + ) + + assert guardrail.category_keywords and not guardrail.always_block_category_keywords + assert guardrail.rewrites_streamed_output() is expected + + @pytest.mark.parametrize("action, expected", [("MASK", True), ("BLOCK", False)]) + def test_always_block_category_keywords_follow_the_category_action(self, action, expected): + guardrail = ContentFilterGuardrail( + guardrail_name="cf", + categories=[{"category": "age_discrimination", "enabled": True, "action": action}], + ) + + assert guardrail.always_block_category_keywords and not guardrail.category_keywords + assert guardrail.rewrites_streamed_output() is expected diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py index ee0a9a10172..895a986f348 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -1493,6 +1493,28 @@ async def test_pre_call_hook_rejects_streaming_only_when_content_filter_step_mas assert "a MASK action" in info.value.detail["error"]["message"] +@pytest.mark.asyncio +async def test_pre_call_hook_rejects_streaming_when_content_filter_category_masks( + proxy_logging, make_user_api_key_auth, monkeypatch +): + guardrail = ContentFilterGuardrail( + guardrail_name="gr-post", + event_hook=GuardrailEventHooks.post_call, + categories=[{"category": "bias_gender", "enabled": True, "action": "MASK"}], + ) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + data = _post_call_pipeline_data(stream=True) + user_api_key_dict = make_user_api_key_auth(request_route="/v1/chat/completions") + + with pytest.raises(HTTPException) as info: + await proxy_logging.pre_call_hook( + user_api_key_dict=user_api_key_dict, data=data, call_type="completion", guardrails_only=True + ) + + assert info.value.status_code == 400 + assert info.value.detail["error"]["guardrails"] == ("gr-post",) + + @pytest.mark.asyncio async def test_pre_call_hook_rejects_streaming_when_route_has_no_guardrail_translation( proxy_logging, make_user_api_key_auth, monkeypatch From 2247fbc66df9769c3e3661b457b0b132513c1bd5 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 15:44:08 -0700 Subject: [PATCH 13/60] fix(policy_engine): withhold streams when a pipeline guardrail rewrites output at runtime --- .../proxy/policy_engine/pipeline_executor.py | 114 +++++++++++++++--- litellm/proxy/utils.py | 60 ++++++--- .../policy_engine/test_pipeline_executor.py | 63 +++++++++- .../proxy_logging/test_guardrail_pipeline.py | 103 +++++++++++++++- 4 files changed, 302 insertions(+), 38 deletions(-) diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index c422a7c0964..acd2c2c973a 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -6,8 +6,11 @@ pass/fail actions (allow, block, next, modify_response) and data forwarding. """ import time +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final, Literal +from pydantic import BaseModel + import litellm from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_guardrail import ( @@ -24,8 +27,10 @@ from litellm.types.proxy.policy_engine.pipeline_types import ( PipelineStep, PipelineStepResult, ) +from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.guardrail_translation.base_translation import ( BaseTranslation, ) @@ -36,6 +41,90 @@ except ImportError: HTTPException = None +class UndeliverableStreamRewrite(Exception): + def __init__(self, guardrail_name: str) -> None: + super().__init__( + f"Guardrail '{guardrail_name}' rewrote the streamed response, which streaming pipelines cannot deliver" + ) + self.guardrail_name: Final = guardrail_name + + +def _tool_call_shape(tool_call: object) -> tuple[object, object]: + plain: Final = tool_call.model_dump() if isinstance(tool_call, BaseModel) else tool_call + function: Final = plain.get("function") if isinstance(plain, Mapping) else None + if not isinstance(function, Mapping): + return (None, None) + return (function.get("name"), function.get("arguments")) + + +def _rewrote_texts(sent: Sequence[str] | None, returned: Sequence[str] | None) -> bool: + return sent is not None and returned is not None and list(returned) != list(sent) + + +def _rewrote_tool_calls(sent: Sequence[object] | None, returned: Sequence[object] | None) -> bool: + if sent is None or returned is None: + return False + return [_tool_call_shape(tool_call) for tool_call in returned] != [ + _tool_call_shape(tool_call) for tool_call in sent + ] + + +class _StreamRewriteObserver(CustomGuardrail): + """Stand-in handed to the endpoint translation in place of a streaming pipeline step's + guardrail. Translations cannot rewrite every buffered chunk consistently, so the gate + withholds the stream whenever the guardrail returned different output than it was given, + which for guardrails like Bedrock's ANONYMIZED action is only known at runtime.""" + + def __init__(self, inner: CustomGuardrail) -> None: + super().__init__(guardrail_name=inner.guardrail_name) + self.inner: Final = inner + self.rewrote = False + + def structured_messages_cover_full_request(self) -> bool: + return self.inner.structured_messages_cover_full_request() + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, # mutable-ok: matches CustomGuardrail.apply_guardrail + input_type: Literal["request", "response"], + logging_obj: "LiteLLMLoggingObj | None" = None, + ) -> GenericGuardrailAPIInputs: + outputs: Final = await self.inner.apply_guardrail( + inputs=inputs, request_data=request_data, input_type=input_type, logging_obj=logging_obj + ) + self.rewrote = ( + self.rewrote + or _rewrote_texts(inputs.get("texts"), outputs.get("texts")) + or _rewrote_tool_calls(inputs.get("tool_calls"), outputs.get("tool_calls")) + ) + return outputs + + +def _prepare_hook_input( + step: PipelineStep, + callback: CustomLogger, + data: dict, # mutable-ok: same request-payload shape the hooks mutate + raw_request_snapshot: dict | None, # mutable-ok: same request-payload shape as data +) -> tuple[dict, bool]: # mutable-ok: returns that same request-payload dict + """Inject the step's guardrail name into metadata so should_run_guardrail() allows it, + and pick the payload the step scans: a scan_raw_request step evaluates the pristine + pre-pipeline snapshot instead of `data` (which earlier pass_data steps in this same + pipeline may have already rewritten), same reason the normal sequential/parallel + guardrail loops do this.""" + if "metadata" not in data: + data["metadata"] = {} + data["metadata"]["guardrails"] = [step.guardrail] + + scans_raw_request: Final = getattr(callback, "scan_raw_request", False) + hook_input: Final[dict] = ( # mutable-ok: same request-payload shape as data + independent_snapshot(raw_request_snapshot) if scans_raw_request and raw_request_snapshot is not None else data + ) + if hook_input is not data: + hook_input.setdefault("metadata", {})["guardrails"] = [step.guardrail] + return hook_input, scans_raw_request + + class PipelineExecutor: """Executes guardrail pipelines with ordered, conditional step logic.""" @@ -195,23 +284,7 @@ class PipelineExecutor: return ("error", None, f"Guardrail '{step.guardrail}' not found", None) try: - # Inject guardrail name into metadata so should_run_guardrail() allows it - if "metadata" not in data: - data["metadata"] = {} - data["metadata"]["guardrails"] = [step.guardrail] - - # A scan_raw_request step evaluates the pristine pre-pipeline - # snapshot instead of `data` (which earlier pass_data steps in - # this same pipeline may have already rewritten), same reason - # the normal sequential/parallel guardrail loops do this. - scans_raw_request: Final = getattr(callback, "scan_raw_request", False) - hook_input: Final[dict] = ( # mutable-ok: same request-payload shape as data - independent_snapshot(raw_request_snapshot) - if scans_raw_request and raw_request_snapshot is not None - else data - ) - if hook_input is not data: - hook_input.setdefault("metadata", {})["guardrails"] = [step.guardrail] + hook_input, scans_raw_request = _prepare_hook_input(step, callback, data, raw_request_snapshot) # Use unified_guardrail path if callback implements apply_guardrail target: CustomLogger = callback @@ -239,13 +312,16 @@ class PipelineExecutor: f"Guardrail '{step.guardrail}' does not support streaming pipeline execution", None, ) + observer: Final = _StreamRewriteObserver(callback) await endpoint_translation.process_output_streaming_response( responses_so_far=streaming_chunks, - guardrail_to_apply=callback, + guardrail_to_apply=observer, litellm_logging_obj=data.get("litellm_logging_obj"), user_api_key_dict=user_api_key_dict, request_data=hook_input, ) + if observer.rewrote: + raise UndeliverableStreamRewrite(step.guardrail) response = None elif mode == "post_call": response = await target.async_post_call_success_hook( @@ -269,6 +345,8 @@ class PipelineExecutor: return ("pass", {"response": response}, None, None) return ("pass", response if isinstance(response, dict) else None, None, None) + except UndeliverableStreamRewrite: + raise except Exception as e: if CustomGuardrail._is_guardrail_intervention(e): error_msg: Final = _extract_error_message(e) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 138f272af42..6c990222e51 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -155,7 +155,7 @@ from litellm.proxy.hooks.sensitive_data_routing import ( ) from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup from litellm.proxy.management_helpers.key_settings_audit import with_settings_updated_at -from litellm.proxy.policy_engine.pipeline_executor import PipelineExecutor +from litellm.proxy.policy_engine.pipeline_executor import PipelineExecutor, UndeliverableStreamRewrite from litellm.repositories.budget_repository import BudgetRepository from litellm.repositories.config_repository import ConfigRepository from litellm.repositories.table_repositories import ( @@ -511,6 +511,23 @@ class _PipelineErrorDetail(TypedDict): error: ReadOnly[_PipelineErrorBody] +def _undeliverable_stream_rewrite_error(policy_name: str, guardrail_name: str) -> HTTPException: + detail: Final[_PipelineErrorDetail] = { + "error": { + "message": ( + f"Streaming response withheld by policy pipeline '{policy_name}' because guardrail " + f"'{guardrail_name}' rewrote the streamed output, and streaming pipelines cannot deliver " + "rewrites. Retry with stream=false, or drop it from the pipeline steps so guardrails.add " + "applies it to streamed output." + ), + "type": "guardrail_pipeline_error", + "policies": (policy_name,), + "guardrails": (guardrail_name,), + } + } + return HTTPException(status_code=400, detail=detail) + + def _raise_for_streaming_post_call_pipelines(data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth) -> None: """ Reject up front the requests whose post_call pipelines could never run. @@ -3468,11 +3485,12 @@ class ProxyLogging: pipeline allows it), then runs each pipeline's steps against the assembled output through the endpoint guardrail translation, the same machinery flat post_call guardrails use at end of stream. An allow - releases the buffered chunks as that machinery left them (the - Responses and A2A translations write guardrail output back into the - final chunk, exactly as they do for flat guardrails); a block or - modify_response terminates with the translation's block chunks or the - raised error. + releases the buffered chunks verbatim; a step whose guardrail rewrote + the output withholds the stream with a 400 instead, since no + translation rewrites every buffered chunk consistently and some + rewrites (Bedrock's ANONYMIZED action, for one) are only decided at + runtime; a block or modify_response terminates with the translation's + block chunks or the raised error. """ buffered: Final[list[object]] = [] # mutable-ok: accumulates the stream before the pipeline verdict async for item in response: @@ -3495,16 +3513,26 @@ class ProxyLogging: call_type, endpoint_translation = resolved for policy_name, pipeline in pipelines: - result: PipelineExecutionResult = await PipelineExecutor.execute_steps( - steps=pipeline.steps, - mode="post_call", - data=request_data, - user_api_key_dict=user_api_key_dict, - call_type=call_type, - policy_name=policy_name, - streaming_chunks=buffered, - endpoint_translation=endpoint_translation, - ) + try: + result: PipelineExecutionResult = await PipelineExecutor.execute_steps( + steps=pipeline.steps, + mode="post_call", + data=request_data, + user_api_key_dict=user_api_key_dict, + call_type=call_type, + policy_name=policy_name, + streaming_chunks=buffered, + endpoint_translation=endpoint_translation, + ) + except UndeliverableStreamRewrite as rewrite: + async for error_chunk in unified_guardrail.emit_streaming_http_error( + _undeliverable_stream_rewrite_error(policy_name, rewrite.guardrail_name), + call_type, + buffered, + request_data, + ): + yield error_chunk + return try: ProxyLogging._handle_pipeline_result( result, data=request_data, policy_name=policy_name, original_response=buffered diff --git a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py index 054a5af4148..52fd8777a19 100644 --- a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py +++ b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py @@ -13,7 +13,7 @@ from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.proxy.guardrails.guardrail_hooks.custom_code.custom_code_guardrail import ( CustomCodeGuardrail, ) -from litellm.proxy.policy_engine.pipeline_executor import PipelineExecutor +from litellm.proxy.policy_engine.pipeline_executor import PipelineExecutor, UndeliverableStreamRewrite from litellm.types.proxy.policy_engine.pipeline_types import ( GuardrailPipeline, PipelineStep, @@ -811,3 +811,64 @@ async def test_pipeline_step_keeps_native_hook_when_opted_out(monkeypatch): assert outcome == "pass" assert guardrail.native_pre_call_ran is True assert "guardrail_to_apply" not in data + + +class _TextReturningGuardrail(CustomGuardrail): + def __init__(self, returned_texts): + super().__init__(guardrail_name="masker", event_hook="post_call", default_on=True) + self.returned_texts = returned_texts + + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + return {**inputs, "texts": self.returned_texts} + + +class _TextTranslation: + def __init__(self): + self.seen_guardrail_names = [] + + async def process_output_streaming_response( + self, responses_so_far, guardrail_to_apply, litellm_logging_obj=None, user_api_key_dict=None, request_data=None + ): + self.seen_guardrail_names.append(guardrail_to_apply.guardrail_name) + await guardrail_to_apply.apply_guardrail( + inputs={"texts": ["hello world"]}, + request_data=request_data or {}, + input_type="response", + logging_obj=litellm_logging_obj, + ) + return responses_so_far + + +async def _run_streaming_step(returned_texts, translation): + return await PipelineExecutor.execute_steps( + steps=[PipelineStep(guardrail="masker", on_pass="allow", on_fail="next", on_error="next")], + mode="post_call", + data={"model": "m"}, + user_api_key_dict=MagicMock(), + call_type="completion", + policy_name="p", + streaming_chunks=[object()], + endpoint_translation=translation, + ) + + +@pytest.mark.asyncio +async def test_streaming_step_rewrite_escapes_execute_steps_regardless_of_step_actions(monkeypatch): + monkeypatch.setattr(litellm, "callbacks", [_TextReturningGuardrail(["hello [MASKED]"])]) + translation = _TextTranslation() + + with pytest.raises(UndeliverableStreamRewrite) as info: + await _run_streaming_step(["hello [MASKED]"], translation) + + assert info.value.guardrail_name == "masker" + assert translation.seen_guardrail_names == ["masker"] + + +@pytest.mark.asyncio +async def test_streaming_step_unchanged_texts_in_another_container_allow(monkeypatch): + monkeypatch.setattr(litellm, "callbacks", [_TextReturningGuardrail(("hello world",))]) + + result = await _run_streaming_step(("hello world",), _TextTranslation()) + + assert result.terminal_action == "allow" + assert [step.outcome for step in result.step_results] == ["pass"] diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py index 895a986f348..74bd2e483cc 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -11,7 +11,7 @@ from __future__ import annotations import asyncio import json -from typing import Any, Dict, List +from typing import Any, Callable, Dict, List from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -878,10 +878,12 @@ async def test_process_prompt_template_aresponses_swaps_model_and_merges_input(p # --------------------------------------------------------------------------- -def _post_call_pipeline_data(guardrail: str = "gr-post", **extra: Any) -> Dict[str, Any]: +def _post_call_pipeline_data( + guardrail: str = "gr-post", step: PipelineStep | None = None, **extra: Any +) -> Dict[str, Any]: pipeline = GuardrailPipeline( mode="post_call", - steps=[PipelineStep(guardrail=guardrail, on_pass="allow", on_fail="block")], + steps=[step or PipelineStep(guardrail=guardrail, on_pass="allow", on_fail="block")], ) return { "model": "m", @@ -1587,6 +1589,101 @@ async def test_streaming_iterator_hook_pipeline_block_withholds_all_chunks( assert "output blocked" in str(info.value.detail) +def _rewriting_stream_guardrail(transform: Callable[[Dict[str, Any]], Dict[str, Any]]) -> CustomGuardrail: + class RewritingStreamGuardrail(CustomGuardrail): + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + return {**inputs, **transform(inputs)} + + return RewritingStreamGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False) + + +def _tool_call_stream_chunks() -> List[Any]: + tool_call = { + "index": 0, + "id": "call_1", + "type": "function", + "function": {"name": "lookup", "arguments": '{"ssn": "123"}'}, + } + return [ + litellm.ModelResponseStream( + choices=[{"index": 0, "delta": {"tool_calls": [tool_call]}, "finish_reason": None}] + ), + litellm.ModelResponseStream(choices=[{"index": 0, "delta": {}, "finish_reason": "tool_calls"}]), + ] + + +def _echoed_tool_call_dicts(arguments: str) -> List[Dict[str, Any]]: + return [{"id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": arguments}}] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("on_fail, on_error", [("block", None), ("next", "next")]) +@pytest.mark.parametrize( + "make_chunks, transform", + [ + (_stream_chunks, lambda inputs: {"texts": ["hello [MASKED]"]}), + (_tool_call_stream_chunks, lambda inputs: {"tool_calls": _echoed_tool_call_dicts('{"ssn": "[MASKED]"}')}), + ], + ids=["texts", "tool_calls"], +) +async def test_streaming_iterator_hook_pipeline_withholds_runtime_rewrite( + proxy_logging, make_user_api_key_auth, monkeypatch, make_chunks, transform, on_fail, on_error +): + monkeypatch.setattr(litellm, "callbacks", [_rewriting_stream_guardrail(transform)]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + step = PipelineStep(guardrail="gr-post", on_pass="allow", on_fail=on_fail, on_error=on_error) + data = _post_call_pipeline_data(step=step, stream=True) + delivered: List[Any] = [] + + async def _drain() -> None: + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"), + response=_async_chunk_iter(make_chunks()), + request_data=data, + ): + delivered.append(item) + + with pytest.raises(HTTPException) as info: + await _drain() + + error = info.value.detail["error"] + assert delivered == [] + assert info.value.status_code == 400 + assert error["type"] == "guardrail_pipeline_error" + assert error["policies"] == ("response-governance",) + assert error["guardrails"] == ("gr-post",) + assert "stream=false" in error["message"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "make_chunks, transform", + [ + (_stream_chunks, lambda inputs: {"texts": tuple(inputs["texts"])}), + (_tool_call_stream_chunks, lambda inputs: {"tool_calls": _echoed_tool_call_dicts('{"ssn": "123"}')}), + ], + ids=["texts_as_tuple", "tool_calls_as_dicts"], +) +async def test_streaming_iterator_hook_pipeline_releases_stream_echoed_in_another_shape( + proxy_logging, make_user_api_key_auth, monkeypatch, make_chunks, transform +): + monkeypatch.setattr(litellm, "callbacks", [_rewriting_stream_guardrail(transform)]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data(stream=True) + chunks = make_chunks() + + delivered = [ + item + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"), + response=_async_chunk_iter(chunks), + request_data=data, + ) + ] + + assert [id(item) for item in delivered] == [id(chunk) for chunk in chunks] + + @pytest.mark.asyncio async def test_streaming_iterator_hook_pipeline_withholds_unresolvable_response_shape( proxy_logging, make_user_api_key_auth, monkeypatch From badefa395cbfea283e2bac3d7ba545638a0792d0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 21:30:31 -0700 Subject: [PATCH 14/60] fix(policy_engine): snapshot guardrail inputs before apply_guardrail so in-place stream rewrites are withheld --- .../proxy/policy_engine/pipeline_executor.py | 22 ++++++++++--------- .../policy_engine/test_pipeline_executor.py | 22 +++++++++++++++++++ 2 files changed, 34 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index 21f3aca3f58..4c192a50096 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -57,16 +57,16 @@ def _tool_call_shape(tool_call: object) -> tuple[object, object]: return (function.get("name"), function.get("arguments")) -def _rewrote_texts(sent: Sequence[str] | None, returned: Sequence[str] | None) -> bool: - return sent is not None and returned is not None and tuple(returned) != tuple(sent) +def _text_snapshot(texts: Sequence[str] | None) -> tuple[str, ...] | None: + return None if texts is None else tuple(texts) -def _rewrote_tool_calls(sent: Sequence[object] | None, returned: Sequence[object] | None) -> bool: - if sent is None or returned is None: - return False - return tuple(_tool_call_shape(tool_call) for tool_call in returned) != tuple( - _tool_call_shape(tool_call) for tool_call in sent - ) +def _tool_call_shapes(tool_calls: Sequence[object] | None) -> tuple[tuple[object, object], ...] | None: + return None if tool_calls is None else tuple(_tool_call_shape(tool_call) for tool_call in tool_calls) + + +def _rewrote(sent: tuple[object, ...] | None, returned: tuple[object, ...] | None) -> bool: + return sent is not None and returned is not None and returned != sent class _StreamRewriteObserver(CustomGuardrail): @@ -90,13 +90,15 @@ class _StreamRewriteObserver(CustomGuardrail): input_type: Literal["request", "response"], logging_obj: "LiteLLMLoggingObj | None" = None, ) -> GenericGuardrailAPIInputs: + sent_texts: Final = _text_snapshot(inputs.get("texts")) + sent_tool_shapes: Final = _tool_call_shapes(inputs.get("tool_calls")) outputs: Final = await self.inner.apply_guardrail( inputs=inputs, request_data=request_data, input_type=input_type, logging_obj=logging_obj ) self.rewrote = ( self.rewrote - or _rewrote_texts(inputs.get("texts"), outputs.get("texts")) - or _rewrote_tool_calls(inputs.get("tool_calls"), outputs.get("tool_calls")) + or _rewrote(sent_texts, _text_snapshot(outputs.get("texts"))) + or _rewrote(sent_tool_shapes, _tool_call_shapes(outputs.get("tool_calls"))) ) return outputs diff --git a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py index 52fd8777a19..ef5d206f1d8 100644 --- a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py +++ b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py @@ -872,3 +872,25 @@ async def test_streaming_step_unchanged_texts_in_another_container_allow(monkeyp assert result.terminal_action == "allow" assert [step.outcome for step in result.step_results] == ["pass"] + + +class _InPlaceMutatingGuardrail(CustomGuardrail): + """Rewrites like bedrock/presidio do: rebinds inputs["texts"] on the dict it was handed + and returns that same dict, so a post-call comparison against inputs sees no change.""" + + def __init__(self): + super().__init__(guardrail_name="masker", event_hook="post_call", default_on=True) + + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + inputs["texts"] = ["hello [MASKED]"] + return inputs + + +@pytest.mark.asyncio +async def test_streaming_step_in_place_rewrite_still_withholds_stream(monkeypatch): + monkeypatch.setattr(litellm, "callbacks", [_InPlaceMutatingGuardrail()]) + + with pytest.raises(UndeliverableStreamRewrite) as info: + await _run_streaming_step(["hello [MASKED]"], _TextTranslation()) + + assert info.value.guardrail_name == "masker" From bcee01a7a7a3a29c5f6e54a0045ff3688d2dbdef Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 21:52:52 -0700 Subject: [PATCH 15/60] fix(policy_engine): merge guardrail metadata writes back on block and modify_response so failure spend records keep guardrail cost and status --- .../proxy/policy_engine/pipeline_executor.py | 2 + litellm/proxy/utils.py | 7 +++- .../proxy_logging/test_guardrail_pipeline.py | 38 +++++++++++++++++++ 3 files changed, 46 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index 4c192a50096..0c3ceb53707 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -236,6 +236,7 @@ class PipelineExecutor: step_results=step_results, error_message=error_detail, original_exception=original_exception, + modified_data=working_data if working_data != data else None, ) if action == "modify_response": @@ -243,6 +244,7 @@ class PipelineExecutor: terminal_action="modify_response", step_results=step_results, modify_response_message=step.modify_response_message or error_detail, + modified_data=working_data if working_data != data else None, ) # action == "next" → continue to next step diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index f9d018bc452..dbcea834d50 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1838,7 +1838,9 @@ class ProxyLogging: payload (already sent upstream) must stay untouched; a replacement response carried in ``modified_data`` is adopted by the caller, and metadata-bucket writes (applied guardrails, guardrail logging info) - are merged back so headers and spend logs still see them. On the + are merged back so headers and spend logs still see them, on block + and modify_response too, so failure spend records keep guardrail + cost and status. On the streaming path it is the buffered chunk list, carried into ``ModifyResponseException.original_response`` for usage reporting. """ @@ -1850,6 +1852,9 @@ class ProxyLogging: _merge_pipeline_metadata_writes(data, result.modified_data) return data + if result.modified_data is not None: + _merge_pipeline_metadata_writes(data, result.modified_data) + if result.terminal_action == "block": original_exception: Final = result.original_exception if original_exception is not None and not _exception_changes_request_flow(original_exception): diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py index 74bd2e483cc..d9b3578c966 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -1197,6 +1197,44 @@ async def test_post_call_pipeline_guardrail_metadata_writes_reach_request_data( assert slg_entries[0]["guardrail_name"] == "gr-post" +@pytest.mark.asyncio +async def test_post_call_pipeline_block_keeps_guardrail_metadata_writes( + proxy_logging, make_user_api_key_auth, monkeypatch +): + class BlockingWriterGuardrail(CustomGuardrail): + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name="gr-post") + self.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response={"verdict": "fail"}, + request_data=data, + guardrail_status="guardrail_intervened", + ) + raise HTTPException(status_code=400, detail={"error": "output blocked"}) + + monkeypatch.setattr( + litellm, + "callbacks", + [ + BlockingWriterGuardrail( + guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False + ) + ], + ) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data() + + with pytest.raises(HTTPException): + await proxy_logging.post_call_success_hook( + data=data, response=litellm.ModelResponse(), user_api_key_dict=make_user_api_key_auth() + ) + + assert data["metadata"]["applied_guardrails"] == ["gr-post"] + slg_entries = data["metadata"]["standard_logging_guardrail_information"] + assert len(slg_entries) == 1 + assert slg_entries[0]["guardrail_name"] == "gr-post" + assert slg_entries[0]["guardrail_status"] == "guardrail_intervened" + + @pytest.mark.asyncio async def test_post_call_pipeline_managed_parallel_guardrail_runs_exactly_once( proxy_logging, make_user_api_key_auth, monkeypatch From 673d1743a66363022777f0b3b261142ef77964ab Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 22:12:59 -0700 Subject: [PATCH 16/60] fix(policy_engine): apply post_call pipeline text rewrites on streams Buffered streams governed by post_call policy pipelines now deliver text rewrites back into the stream per surface (chat SSE, responses SSE, anthropic messages SSE) instead of rejecting the request with a 400 upfront. Rewrites chain across pipeline steps; tool-call rewrites and translations without stream write-back still withhold the stream. --- .../chat/guardrail_translation/handler.py | 74 +++++- .../guardrail_translation/base_translation.py | 15 +- .../chat/guardrail_translation/handler.py | 91 ++++++- .../guardrail_translation/handler.py | 83 ++++++- .../proxy/policy_engine/pipeline_executor.py | 90 +++++-- litellm/proxy/utils.py | 61 ++--- .../test_anthropic_guardrail_handler.py | 65 +++++ .../test_openai_guardrail_handler.py | 55 +++++ ...test_openai_responses_guardrail_handler.py | 83 +++++++ .../policy_engine/test_pipeline_executor.py | 2 + .../proxy_logging/test_guardrail_pipeline.py | 222 ++++++++++++++---- 11 files changed, 711 insertions(+), 130 deletions(-) diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index b9ca18c7843..89c8431dfe4 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -13,9 +13,10 @@ Pattern Overview: """ import json -from collections.abc import Mapping, Sequence +from collections.abc import Iterator, Mapping, Sequence from copy import deepcopy from dataclasses import dataclass +from itertools import chain, repeat from typing import TYPE_CHECKING, Any, Final, cast from typing_extensions import assert_never @@ -120,6 +121,8 @@ class AnthropicMessagesHandler(BaseTranslation): them through guardrail rewrites; downstream provider handling is out of scope. """ + delivers_ended_stream_text_rewrites = True + def __init__(self): super().__init__() self.adapter = LiteLLMAnthropicMessagesAdapter() @@ -931,11 +934,14 @@ class AnthropicMessagesHandler(BaseTranslation): litellm_logging_obj: "LiteLLMLoggingObj | None" = None, user_api_key_dict: "UserAPIKeyAuth | None" = None, request_data: dict | None = None, + deliver_ended_stream_rewrites: bool = False, ) -> list[Any]: """ Process output streaming response by applying guardrails to text content. Get the string so far, check the apply guardrail to the string so far, and return the list of responses so far. + With ``deliver_ended_stream_rewrites``, an ended stream whose guardrail rewrote the text gets the rewrite + written back across the buffered chunks (full rewritten text in the first ``text_delta``, the rest blanked). """ from litellm.integrations.custom_guardrail import ModifyResponseException @@ -982,6 +988,15 @@ class AnthropicMessagesHandler(BaseTranslation): responses_so_far, request_data ) raise + guardrailed_texts: Final = _guardrailed_inputs.get("texts") + if ( + deliver_ended_stream_rewrites + and isinstance(string_so_far, str) + and string_so_far + and guardrailed_texts + and guardrailed_texts[0] != string_so_far + ): + self._write_ended_stream_text_rewrite(responses_so_far, guardrailed_texts[0]) else: verbose_proxy_logger.debug("Skipping output guardrail - model response has no choices") return responses_so_far @@ -1093,6 +1108,63 @@ class AnthropicMessagesHandler(BaseTranslation): inputs["model"] = response_model return inputs + @staticmethod + def _write_ended_stream_text_rewrite( + responses_so_far: list[Any], # mutable-ok: rewrites the caller's buffered chunks in place + rewritten_text: str, + ) -> None: + """Deliver an ended-stream guardrail text rewrite by rewriting the + buffered chunks in place: the first ``text_delta`` carries the full + rewritten text and every later one is blanked, leaving the surrounding + message and content-block framing untouched. Handles both chunk formats + this stream carries (parsed event dicts and raw SSE bytes).""" + replacements: Final = chain((rewritten_text,), repeat("")) + for idx, item in enumerate(responses_so_far): + if isinstance(item, dict): + delta = item.get("delta") + if item.get("type") == "content_block_delta" and isinstance(delta, dict): + if delta.get("type") == "text_delta": + delta["text"] = next(replacements) + elif isinstance(item, (bytes, bytearray)): + responses_so_far[idx] = ( # rebind-ok: delivers the rewrite into the caller's buffer + AnthropicMessagesHandler._rewrite_sse_text_deltas(bytes(item), replacements) + ) + + @staticmethod + def _rewrite_sse_text_deltas(sse_bytes: bytes, replacements: "Iterator[str]") -> bytes: + """Rewrite every ``text_delta`` data line in one SSE chunk with the next + replacement text, leaving all other events and framing byte-identical.""" + try: + decoded: Final = sse_bytes.decode("utf-8") + except UnicodeDecodeError: + return sse_bytes + return "\n\n".join( + AnthropicMessagesHandler._rewrite_sse_block(block, replacements) for block in decoded.split("\n\n") + ).encode("utf-8") + + @staticmethod + def _rewrite_sse_block(block: str, replacements: "Iterator[str]") -> str: + return "\n".join(AnthropicMessagesHandler._rewrite_sse_line(line, replacements) for line in block.split("\n")) + + @staticmethod + def _rewrite_sse_line(line: str, replacements: "Iterator[str]") -> str: + if not line.startswith("data:"): + return line + try: + data: Final[str | int | float | bool | None | Sequence[object] | Mapping[str, object]] = json.loads( + line[len("data:") :].strip() + ) + except json.JSONDecodeError: + return line + if not isinstance(data, dict) or data.get("type") != "content_block_delta": + return line + delta: Final = data.get("delta") + if not isinstance(delta, dict) or delta.get("type") != "text_delta": + return line + return "data: " + json.dumps( + {**data, "delta": {**delta, "text": next(replacements)}} # mutable-ok: json.dumps needs plain dicts + ) + def get_streaming_string_so_far(self, responses_so_far: list[Any]) -> str: """ Parse streaming responses and extract accumulated text content. diff --git a/litellm/llms/base_llm/guardrail_translation/base_translation.py b/litellm/llms/base_llm/guardrail_translation/base_translation.py index ba96ab3dc99..4b0cc0fd97c 100644 --- a/litellm/llms/base_llm/guardrail_translation/base_translation.py +++ b/litellm/llms/base_llm/guardrail_translation/base_translation.py @@ -1,6 +1,6 @@ from abc import ABC, abstractmethod from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Any, Final, Optional +from typing import TYPE_CHECKING, Any, ClassVar, Final, Optional if TYPE_CHECKING: from litellm.integrations.custom_guardrail import ( @@ -33,6 +33,13 @@ class StreamTransformSink: class BaseTranslation(ABC): + delivers_ended_stream_text_rewrites: ClassVar[bool] = False + """Whether ``process_output_streaming_response`` accepts + ``deliver_ended_stream_rewrites=True`` and, on an ended (fully buffered) + stream, writes guardrail text rewrites back across ``responses_so_far`` so + a buffered pipeline can release rewritten chunks instead of withholding the + stream. Tool-call rewrites stay undeliverable everywhere.""" + @staticmethod def transform_user_api_key_dict_to_metadata( user_api_key_dict: Any | None, @@ -113,6 +120,7 @@ class BaseTranslation(ABC): user_api_key_dict: Optional["UserAPIKeyAuth"] = None, request_data: dict | None = None, stream_transform_sink: StreamTransformSink | None = None, + deliver_ended_stream_rewrites: bool = False, ) -> Any: """ Process output streaming response with guardrails. @@ -120,6 +128,11 @@ class BaseTranslation(ABC): Optional to override in subclasses. ``stream_transform_sink`` is the out-parameter used by handlers that support streaming text transformations (see ``StreamTransformSink``); base handlers ignore it. + ``deliver_ended_stream_rewrites`` is passed True only when the caller + holds the whole buffered stream and the subclass declares + ``delivers_ended_stream_text_rewrites``: the handler then writes + guardrail text rewrites back across ``responses_so_far`` instead of + discarding them. """ return responses_so_far diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 54673c77f80..1358cf7c37a 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -61,6 +61,8 @@ class OpenAIChatCompletionsHandler(BaseTranslation): Methods can be overridden to customize behavior for different message formats. """ + delivers_ended_stream_text_rewrites = True + def get_structured_messages(self, data: dict) -> list[AllMessageValues] | None: """ Convert chat completions request data to OpenAI-spec structured messages. @@ -440,6 +442,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): user_api_key_dict: Any | None = None, request_data: dict | None = None, stream_transform_sink: StreamTransformSink | None = None, + deliver_ended_stream_rewrites: bool = False, ) -> list["ModelResponseStream"]: """ Process output streaming responses by applying guardrails to text content. @@ -454,6 +457,10 @@ class OpenAIChatCompletionsHandler(BaseTranslation): accumulated text (``responses_so_far`` is left untouched so it stays a correct raw accumulator across rounds) and the guardrailed text plus requested holdback are reported per choice on the sink. + deliver_ended_stream_rewrites: When True and the buffered stream has + ended, guardrail text rewrites are written back across + ``responses_so_far`` (full rewritten text in each choice's first + content-carrying chunk, the rest blanked) instead of discarded. Returns: The (unmodified) list of responses. @@ -479,6 +486,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): litellm_logging_obj=litellm_logging_obj, user_api_key_dict=user_api_key_dict, request_data=request_data, + deliver_ended_stream_rewrites=deliver_ended_stream_rewrites, ) async def _process_streaming_block_only( @@ -489,10 +497,12 @@ class OpenAIChatCompletionsHandler(BaseTranslation): litellm_logging_obj: "LiteLLMLoggingObj | None", user_api_key_dict: Any | None, request_data: dict | None, + deliver_ended_stream_rewrites: bool = False, ) -> list["ModelResponseStream"]: """Block-only streaming path: run the guardrail so an in-flight BLOCK can terminate the stream. Text rewrites are not propagated to the client here - (see ``_process_streaming_transform`` for the incremental_diff path).""" + (see ``_process_streaming_transform`` for the incremental_diff path) unless + ``deliver_ended_stream_rewrites`` opts the ended-stream branch in.""" # check if the stream has ended has_stream_ended = False for chunk in responses_so_far: @@ -501,20 +511,14 @@ class OpenAIChatCompletionsHandler(BaseTranslation): break if has_stream_ended: - # convert to model response - model_response: Final = cast( - ModelResponse, - stream_chunk_builder(chunks=responses_so_far, logging_obj=litellm_logging_obj), - ) - # run process_output_response - await self.process_output_response( - response=model_response, + await self._process_ended_stream( + responses_so_far=responses_so_far, guardrail_to_apply=guardrail_to_apply, litellm_logging_obj=litellm_logging_obj, user_api_key_dict=user_api_key_dict, request_data=request_data, + deliver_ended_stream_rewrites=deliver_ended_stream_rewrites, ) - return responses_so_far # Step 0: Check if any response has text content to process @@ -591,6 +595,38 @@ class OpenAIChatCompletionsHandler(BaseTranslation): return responses_so_far + async def _process_ended_stream( + self, + *, + responses_so_far: list["ModelResponseStream"], # mutable-ok: rewrites the caller's buffered chunks in place + guardrail_to_apply: "CustomGuardrail", + litellm_logging_obj: "LiteLLMLoggingObj | None", + user_api_key_dict: object, + request_data: dict[str, object] | None, # mutable-ok: same request-payload shape the hooks take + deliver_ended_stream_rewrites: bool, + ) -> None: + """Ended-stream path: rebuild the full response, run the non-streaming + output guardrail against it, and (when opted in) write any text rewrite + back across the buffered chunks.""" + model_response: Final = cast( + ModelResponse, + stream_chunk_builder(chunks=responses_so_far, logging_obj=litellm_logging_obj), + ) + pre_guardrail_texts: Final = self._string_choice_contents(model_response) + await self.process_output_response( + response=model_response, + guardrail_to_apply=guardrail_to_apply, + litellm_logging_obj=litellm_logging_obj, + user_api_key_dict=user_api_key_dict, + request_data=request_data, + ) + if deliver_ended_stream_rewrites: + await self._write_ended_stream_text_rewrites( + responses_so_far=responses_so_far, + guardrailed_response=model_response, + pre_guardrail_texts=pre_guardrail_texts, + ) + @staticmethod def _accumulate_string_content_by_choice_index( responses_so_far: list["ModelResponseStream"], @@ -922,6 +958,41 @@ class OpenAIChatCompletionsHandler(BaseTranslation): if "name" in func_dict: existing_tool_call.function.name = func_dict["name"] + @staticmethod + def _string_choice_contents(response: "ModelResponse") -> tuple[str | None, ...]: + return tuple( + choice.message.content if isinstance(choice.message.content, str) else None for choice in response.choices + ) + + async def _write_ended_stream_text_rewrites( + self, + responses_so_far: list["ModelResponseStream"], # mutable-ok: rewrites the caller's buffered chunks in place + guardrailed_response: "ModelResponse", + pre_guardrail_texts: tuple[str | None, ...], + ) -> None: + """Write ended-stream guardrail text rewrites back across the buffered + chunks: each rewritten choice's full text lands in its first + content-carrying chunk and the rest are blanked, the same shape the + in-flight write-back uses. Chunks carrying only finish_reason or usage + stay untouched.""" + post_guardrail_texts: Final = self._string_choice_contents(guardrailed_response) + changed: Final = tuple( + (choice_idx, after) + for choice_idx, (before, after) in enumerate(zip(pre_guardrail_texts, post_guardrail_texts)) + if before is not None and after is not None and after != before + ) + if not changed: + return + await self._apply_guardrail_responses_to_output_streaming( + responses=responses_so_far, + guardrailed_texts=[ + after for _choice_idx, after in changed + ], # mutable-ok: the callee's signature predates this change and takes lists + task_mappings=[ + (choice_idx, None) for choice_idx, _after in changed + ], # mutable-ok: the callee's signature predates this change and takes lists + ) + async def _apply_guardrail_responses_to_output_streaming( self, responses: list["ModelResponseStream"], diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 7c5d8ac99ad..0f475aa04c8 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -28,7 +28,9 @@ Output: response.output is List[GenericResponseOutputItem] where each has: - text: str """ -from collections.abc import Sequence +from collections.abc import Mapping, Sequence +from itertools import chain, repeat +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Union, cast from openai.types.responses.response_function_tool_call import ResponseFunctionToolCall @@ -91,6 +93,8 @@ class OpenAIResponsesHandler(BaseTranslation): Methods can be overridden to customize behavior for different message formats. """ + delivers_ended_stream_text_rewrites = True + def get_structured_messages(self, data: dict) -> list[AllMessageValues] | None: """ Convert Responses API request data to OpenAI-spec structured messages. @@ -482,6 +486,7 @@ class OpenAIResponsesHandler(BaseTranslation): litellm_logging_obj: "LiteLLMLoggingObj | None" = None, user_api_key_dict: "UserAPIKeyAuth | None" = None, request_data: dict | None = None, + deliver_ended_stream_rewrites: bool = False, ) -> list[Any]: """ Process output streaming response by applying guardrails to text content. @@ -493,7 +498,11 @@ class OpenAIResponsesHandler(BaseTranslation): For ``response.completed`` events (the normal end-of-stream signal) we use the same per-item extraction + task-mapping approach as ``process_output_response`` so that unmasking / blocking works correctly - for every output item. + for every output item. With ``deliver_ended_stream_rewrites`` the earlier + text-carrying events (``response.output_text.delta`` / ``.done``, + ``response.content_part.done``, ``response.output_item.done``) are synced + to the rewritten completed response too, so a client reading deltas sees + the rewrite instead of the raw model output. """ if not responses_so_far: return responses_so_far @@ -562,6 +571,19 @@ class OpenAIResponsesHandler(BaseTranslation): responses=guardrailed_texts, task_mappings=task_mappings, ) + if deliver_ended_stream_rewrites: + rewrites_by_position: Final = MappingProxyType( + { + task_mappings[task_idx]: rewritten + for task_idx, rewritten in enumerate(guardrailed_texts) + if task_idx < len(texts_to_check) and rewritten != texts_to_check[task_idx] + } + ) + if rewrites_by_position: + self._sync_stream_events_with_rewrites( + stream_events=responses_so_far[:-1], + rewrites_by_position=rewrites_by_position, + ) return responses_so_far @@ -607,6 +629,63 @@ class OpenAIResponsesHandler(BaseTranslation): ) return responses_so_far + @staticmethod + def _write_event_field(event: object, field: str, value: str) -> None: + if isinstance(event, dict): + event[field] = value # rebind-ok: delivering the rewrite means editing the buffered event in place + else: + setattr(event, field, value) + + def _sync_stream_events_with_rewrites( + self, + stream_events: Sequence[Any], + rewrites_by_position: Mapping[tuple[int, int], str], + ) -> None: + """Sync pre-completion stream events with the rewritten completed + response, keyed by ``(output_index, content_index)``: the first + ``output_text.delta`` for a rewritten item carries the full rewritten + text and the rest are blanked, while ``output_text.done``, + ``content_part.done``, and ``output_item.done`` events carry the full + rewritten text, so every event a client may read agrees with the + rewritten ``response.completed`` payload.""" + delta_replacements: Final = MappingProxyType( + {position: chain((rewritten,), repeat("")) for position, rewritten in rewrites_by_position.items()} + ) + for event in stream_events: + if not (isinstance(event, dict) or hasattr(event, "get")): + continue + event_type = event.get("type") + output_index = event.get("output_index") + content_index = event.get("content_index") + if event_type == "response.output_item.done" and isinstance(output_index, int): + self._sync_output_item_done_event(event.get("item"), output_index, rewrites_by_position) + continue + if not isinstance(output_index, int) or not isinstance(content_index, int): + continue + position = (output_index, content_index) + if event_type == "response.output_text.delta" and position in delta_replacements: + self._write_event_field(event, "delta", next(delta_replacements[position])) + elif event_type == "response.output_text.done" and position in rewrites_by_position: + self._write_event_field(event, "text", rewrites_by_position[position]) + elif event_type == "response.content_part.done" and position in rewrites_by_position: + part = event.get("part") + if isinstance(part, dict) or hasattr(part, "text"): + self._write_event_field(part, "text", rewrites_by_position[position]) + + @staticmethod + def _sync_output_item_done_event( + item: object, + output_index: int, + rewrites_by_position: Mapping[tuple[int, int], str], + ) -> None: + content: Final = item.get("content") if isinstance(item, dict) else getattr(item, "content", None) + if not isinstance(content, list): + return + for (item_idx, content_idx), rewritten in rewrites_by_position.items(): + if item_idx != output_index or content_idx >= len(content): + continue + OpenAIResponsesHandler._write_event_field(content[content_idx], "text", rewritten) + def _check_streaming_has_ended(self, responses_so_far: Sequence[ResponsesStreamChunk]) -> bool: """ Check if the streaming has ended. diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index acd2c2c973a..1264cdfc14a 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -34,6 +34,7 @@ if TYPE_CHECKING: from litellm.llms.base_llm.guardrail_translation.base_translation import ( BaseTranslation, ) + from litellm.proxy._types import UserAPIKeyAuth try: from fastapi.exceptions import HTTPException @@ -44,7 +45,8 @@ except ImportError: class UndeliverableStreamRewrite(Exception): def __init__(self, guardrail_name: str) -> None: super().__init__( - f"Guardrail '{guardrail_name}' rewrote the streamed response, which streaming pipelines cannot deliver" + f"Guardrail '{guardrail_name}' rewrote the streamed response in a way this endpoint's " + "streaming pipeline cannot deliver" ) self.guardrail_name: Final = guardrail_name @@ -57,28 +59,31 @@ def _tool_call_shape(tool_call: object) -> tuple[object, object]: return (function.get("name"), function.get("arguments")) -def _rewrote_texts(sent: Sequence[str] | None, returned: Sequence[str] | None) -> bool: - return sent is not None and returned is not None and list(returned) != list(sent) +def _text_snapshot(texts: Sequence[str] | None) -> tuple[str, ...] | None: + return None if texts is None else tuple(texts) -def _rewrote_tool_calls(sent: Sequence[object] | None, returned: Sequence[object] | None) -> bool: - if sent is None or returned is None: - return False - return [_tool_call_shape(tool_call) for tool_call in returned] != [ - _tool_call_shape(tool_call) for tool_call in sent - ] +def _tool_call_shapes(tool_calls: Sequence[object] | None) -> tuple[tuple[object, object], ...] | None: + return None if tool_calls is None else tuple(_tool_call_shape(tool_call) for tool_call in tool_calls) + + +def _rewrote(sent: tuple[object, ...] | None, returned: tuple[object, ...] | None) -> bool: + return sent is not None and returned is not None and returned != sent class _StreamRewriteObserver(CustomGuardrail): """Stand-in handed to the endpoint translation in place of a streaming pipeline step's - guardrail. Translations cannot rewrite every buffered chunk consistently, so the gate - withholds the stream whenever the guardrail returned different output than it was given, - which for guardrails like Bedrock's ANONYMIZED action is only known at runtime.""" + guardrail. It records whether the guardrail returned different output than it was given, + which for guardrails like Bedrock's ANONYMIZED action is only known at runtime. Text + rewrites are deliverable on translations that write them back across the buffered chunks + (``delivers_ended_stream_text_rewrites``); tool-call rewrites and text rewrites on any + other translation make the gate withhold the stream.""" def __init__(self, inner: CustomGuardrail) -> None: super().__init__(guardrail_name=inner.guardrail_name) self.inner: Final = inner - self.rewrote = False + self.rewrote_texts = False + self.rewrote_tool_calls = False def structured_messages_cover_full_request(self) -> bool: return self.inner.structured_messages_cover_full_request() @@ -90,13 +95,14 @@ class _StreamRewriteObserver(CustomGuardrail): input_type: Literal["request", "response"], logging_obj: "LiteLLMLoggingObj | None" = None, ) -> GenericGuardrailAPIInputs: + sent_texts: Final = _text_snapshot(inputs.get("texts")) + sent_tool_shapes: Final = _tool_call_shapes(inputs.get("tool_calls")) outputs: Final = await self.inner.apply_guardrail( inputs=inputs, request_data=request_data, input_type=input_type, logging_obj=logging_obj ) - self.rewrote = ( - self.rewrote - or _rewrote_texts(inputs.get("texts"), outputs.get("texts")) - or _rewrote_tool_calls(inputs.get("tool_calls"), outputs.get("tool_calls")) + self.rewrote_texts = self.rewrote_texts or _rewrote(sent_texts, _text_snapshot(outputs.get("texts"))) + self.rewrote_tool_calls = self.rewrote_tool_calls or _rewrote( + sent_tool_shapes, _tool_call_shapes(outputs.get("tool_calls")) ) return outputs @@ -250,6 +256,41 @@ class PipelineExecutor: modified_data=working_data if working_data != data else None, ) + @staticmethod + async def _run_streaming_step( + step: PipelineStep, + callback: CustomGuardrail, + endpoint_translation: "BaseTranslation", + streaming_chunks: list[object], # mutable-ok: shared buffered-stream chunks the translation rewrites in place + hook_input: dict[str, object], # mutable-ok: same request-payload shape as data + user_api_key_dict: "UserAPIKeyAuth | None", + litellm_logging_obj: "LiteLLMLoggingObj | None", + ) -> None: + """Run one streaming post_call step through the endpoint translation, delivering + text rewrites on translations that support ended-stream write-back and raising + ``UndeliverableStreamRewrite`` for any rewrite that cannot reach the client.""" + observer: Final = _StreamRewriteObserver(callback) + deliver_rewrites: Final = type(endpoint_translation).delivers_ended_stream_text_rewrites + if deliver_rewrites: + await endpoint_translation.process_output_streaming_response( + responses_so_far=streaming_chunks, + guardrail_to_apply=observer, + litellm_logging_obj=litellm_logging_obj, + user_api_key_dict=user_api_key_dict, + request_data=hook_input, + deliver_ended_stream_rewrites=True, + ) + else: + await endpoint_translation.process_output_streaming_response( + responses_so_far=streaming_chunks, + guardrail_to_apply=observer, + litellm_logging_obj=litellm_logging_obj, + user_api_key_dict=user_api_key_dict, + request_data=hook_input, + ) + if observer.rewrote_tool_calls or (observer.rewrote_texts and not deliver_rewrites): + raise UndeliverableStreamRewrite(step.guardrail) + @staticmethod async def _run_step( step: PipelineStep, @@ -312,16 +353,15 @@ class PipelineExecutor: f"Guardrail '{step.guardrail}' does not support streaming pipeline execution", None, ) - observer: Final = _StreamRewriteObserver(callback) - await endpoint_translation.process_output_streaming_response( - responses_so_far=streaming_chunks, - guardrail_to_apply=observer, - litellm_logging_obj=data.get("litellm_logging_obj"), + await PipelineExecutor._run_streaming_step( + step=step, + callback=callback, + endpoint_translation=endpoint_translation, + streaming_chunks=streaming_chunks, + hook_input=hook_input, user_api_key_dict=user_api_key_dict, - request_data=hook_input, + litellm_logging_obj=data.get("litellm_logging_obj"), ) - if observer.rewrote: - raise UndeliverableStreamRewrite(step.guardrail) response = None elif mode == "post_call": response = await target.async_post_call_success_hook( diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 6c990222e51..7f3c3aecd87 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -492,14 +492,6 @@ def _pipeline_step_supports_unified_streaming(guardrail_name: str) -> bool: return callback is not None and PipelineExecutor.supports_unified_execution(callback) -def _pipeline_step_rewrites_streamed_content(guardrail_name: str) -> bool: - callback: Final = PipelineExecutor.find_guardrail_callback(guardrail_name) - if callback is None: - return False - transform_mode: Final = unified_guardrail.resolve_streaming_flag(callback, "streaming_transform_mode", "block_only") - return callback.rewrites_streamed_output() or transform_mode == "incremental_diff" - - class _PipelineErrorBody(TypedDict): message: ReadOnly[str] type: ReadOnly[str] @@ -516,9 +508,10 @@ def _undeliverable_stream_rewrite_error(policy_name: str, guardrail_name: str) - "error": { "message": ( f"Streaming response withheld by policy pipeline '{policy_name}' because guardrail " - f"'{guardrail_name}' rewrote the streamed output, and streaming pipelines cannot deliver " - "rewrites. Retry with stream=false, or drop it from the pipeline steps so guardrails.add " - "applies it to streamed output." + f"'{guardrail_name}' rewrote the streamed output in a way this endpoint's streaming " + "pipeline cannot deliver (a tool-call rewrite, or a text rewrite on a route without " + "stream write-back). Retry with stream=false, or drop it from the pipeline steps so " + "guardrails.add applies it to streamed output." ), "type": "guardrail_pipeline_error", "policies": (policy_name,), @@ -535,13 +528,13 @@ def _raise_for_streaming_post_call_pipelines(data: Mapping[str, object], user_ap Background responses skip the post_call hooks entirely, so a pipeline governing one would silently never execute. Streaming responses execute pipelines against the buffered stream through the endpoint guardrail - translation of the request route, releasing the buffered chunks on allow. - That needs every step's guardrail to support the unified apply_guardrail - interface and to only allow or block (a step that rewrites streamed - content, via mask_response_content, a MASK action, or - streaming_transform_mode=incremental_diff, would have its rewrite silently - dropped), and needs the route to have a translation at all; anything else - keeps the 400 rather than letting ungoverned output stream through. + translation of the request route, releasing the buffered chunks on allow + (rewritten in place when a guardrail rewrote text and the translation + delivers ended-stream rewrites; a rewrite the translation cannot deliver + fails closed at runtime instead). That needs every step's guardrail to + support the unified apply_guardrail interface, and needs the route to have + a translation at all; anything else keeps the 400 rather than letting + ungoverned output stream through. """ is_stream: Final = data.get("stream") is True is_background: Final = data.get("background") is True @@ -586,25 +579,6 @@ def _raise_for_streaming_post_call_pipelines(data: Mapping[str, object], user_ap } } raise HTTPException(status_code=400, detail=unsupported_detail) - rewriting_guardrails: Final = tuple( - guardrail for guardrail in step_guardrails if _pipeline_step_rewrites_streamed_content(guardrail) - ) - if rewriting_guardrails: - rewriting_detail: Final[_PipelineErrorDetail] = { - "error": { - "message": ( - "Policies with post_call guardrail pipelines cannot govern streaming responses " - "because these pipeline guardrails rewrite streamed content (mask_response_content, " - "a MASK action, or streaming_transform_mode=incremental_diff), which pipeline steps would release " - f"unmodified: {', '.join(rewriting_guardrails)}. Retry with stream=false, or drop " - "them from the pipeline steps so guardrails.add applies them to streamed output." - ), - "type": "guardrail_pipeline_error", - "policies": post_call_policies, - "guardrails": rewriting_guardrails, - } - } - raise HTTPException(status_code=400, detail=rewriting_detail) route: Final = user_api_key_dict.request_route if not route or resolve_endpoint_translation(user_api_key_dict, None) is not None: return @@ -3485,12 +3459,13 @@ class ProxyLogging: pipeline allows it), then runs each pipeline's steps against the assembled output through the endpoint guardrail translation, the same machinery flat post_call guardrails use at end of stream. An allow - releases the buffered chunks verbatim; a step whose guardrail rewrote - the output withholds the stream with a 400 instead, since no - translation rewrites every buffered chunk consistently and some - rewrites (Bedrock's ANONYMIZED action, for one) are only decided at - runtime; a block or modify_response terminates with the translation's - block chunks or the raised error. + releases the buffered chunks: verbatim when no guardrail rewrote the + output, rewritten in place when one rewrote text and the translation + delivers ended-stream rewrites (later steps then re-scan the rewritten + chunks, so rewrites chain). A rewrite the translation cannot deliver + (a tool-call rewrite, or a text rewrite on a route without write-back) + withholds the stream with a 400; a block or modify_response terminates + with the translation's block chunks or the raised error. """ buffered: Final[list[object]] = [] # mutable-ok: accumulates the stream before the pipeline verdict async for item in response: diff --git a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py index af3ccd65b11..ba26da50bc8 100644 --- a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py +++ b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py @@ -263,6 +263,71 @@ class TestAnthropicMessagesHandlerStreamingOutputProcessing: # Should return the responses unchanged assert result == responses_so_far + @staticmethod + def _ended_sse_chunks() -> list: + events = [ + ("message_start", {"type": "message_start", "message": {"id": "msg_1", "type": "message", "role": "assistant", "model": "claude-sonnet-4-5", "content": [], "stop_reason": None, "usage": {"input_tokens": 1, "output_tokens": 0}}}), + ("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}), + ("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "hello "}}), + ("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "world"}}), + ("content_block_stop", {"type": "content_block_stop", "index": 0}), + ("message_delta", {"type": "message_delta", "delta": {"stop_reason": "end_turn", "stop_sequence": None}, "usage": {"output_tokens": 2}}), + ("message_stop", {"type": "message_stop"}), + ] + return [f"event: {name}\ndata: {json.dumps(payload)}\n\n".encode() for name, payload in events] + + @staticmethod + def _masking_guardrail() -> CustomGuardrail: + class MaskWorld(CustomGuardrail): + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + return {**inputs, "texts": [text.replace("world", "[MASKED]") for text in inputs.get("texts", [])]} + + return MaskWorld(guardrail_name="test") + + @staticmethod + def _delta_texts(chunks: list) -> list: + texts = [] + for chunk in chunks: + for line in chunk.decode().split("\n"): + if not line.startswith("data:"): + continue + data = json.loads(line[len("data:") :].strip()) + if data.get("type") == "content_block_delta": + texts.append(data["delta"]["text"]) + return texts + + @pytest.mark.asyncio + async def test_deliver_ended_stream_rewrites_writes_text_back_into_sse_chunks(self): + handler = AnthropicMessagesHandler() + chunks = self._ended_sse_chunks() + + result = await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=self._masking_guardrail(), + litellm_logging_obj=MagicMock(), + deliver_ended_stream_rewrites=True, + ) + + assert result is chunks + assert self._delta_texts(chunks) == ["hello [MASKED]", ""] + raw = b"".join(chunks).decode() + assert "event: message_start" in raw and "event: message_stop" in raw + assert '"stop_reason": "end_turn"' in raw + + @pytest.mark.asyncio + async def test_ended_stream_rewrite_leaves_chunks_untouched_by_default(self): + handler = AnthropicMessagesHandler() + chunks = self._ended_sse_chunks() + original = [bytes(chunk) for chunk in chunks] + + await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=self._masking_guardrail(), + litellm_logging_obj=MagicMock(), + ) + + assert chunks == original + class TestAnthropicMessagesHandlerInputProcessing: """Test input processing preserves litellm_metadata for dynamic guardrails.""" diff --git a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index a29e0be4655..26442a4a6ed 100644 --- a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -1073,6 +1073,61 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput: # Should return the responses assert result == responses_so_far + @staticmethod + def _ended_stream_chunks() -> list: + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + return [ + ModelResponseStream( + id="chatcmpl-123", + created=1234567890, + model="gpt-4", + object="chat.completion.chunk", + choices=[StreamingChoices(index=0, delta=Delta(content="Hello"), finish_reason=None)], + ), + ModelResponseStream( + id="chatcmpl-123", + created=1234567890, + model="gpt-4", + object="chat.completion.chunk", + choices=[StreamingChoices(index=0, delta=Delta(content=" world"), finish_reason="stop")], + ), + ] + + @pytest.mark.asyncio + async def test_deliver_ended_stream_rewrites_writes_text_back_into_chunks(self): + handler = OpenAIChatCompletionsHandler() + guardrail = MockGuardrail(guardrail_name="test") + chunks = self._ended_stream_chunks() + + result = await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=guardrail, + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + assert result is chunks + assert chunks[0].choices[0].delta.content == "HELLO WORLD" + assert chunks[1].choices[0].delta.content in (None, "") + assert chunks[1].choices[0].finish_reason == "stop" + + @pytest.mark.asyncio + async def test_ended_stream_rewrite_leaves_chunks_untouched_by_default(self): + handler = OpenAIChatCompletionsHandler() + guardrail = MockGuardrail(guardrail_name="test") + chunks = self._ended_stream_chunks() + + await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=guardrail, + litellm_logging_obj=None, + ) + + assert chunks[0].choices[0].delta.content == "Hello" + assert chunks[1].choices[0].delta.content == " world" + assert chunks[1].choices[0].finish_reason == "stop" + class TestGetStructuredMessages: """Test the get_structured_messages method.""" diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py index 447175b09a6..4f95e08cb71 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -1104,6 +1104,89 @@ class TestOpenAIResponsesHandlerStreamingOutputProcessing: output_text = result[-1]["response"]["output"][0]["content"][0]["text"] assert output_text == original_text + @staticmethod + def _ended_stream_events() -> List[dict]: + content = [{"type": "output_text", "text": "hello world"}] + item = { + "type": "message", + "id": "msg_123", + "status": "completed", + "role": "assistant", + "content": content, + } + return [ + {"type": "response.output_text.delta", "output_index": 0, "content_index": 0, "delta": "hello "}, + {"type": "response.output_text.delta", "output_index": 0, "content_index": 0, "delta": "world"}, + {"type": "response.output_text.done", "output_index": 0, "content_index": 0, "text": "hello world"}, + { + "type": "response.content_part.done", + "output_index": 0, + "content_index": 0, + "part": {"type": "output_text", "text": "hello world"}, + }, + {"type": "response.output_item.done", "output_index": 0, "item": {**item, "content": [dict(c) for c in content]}}, + { + "type": "response.completed", + "response": { + "id": "resp_123", + "model": "gpt-4o", + "output": [{**item, "content": [dict(c) for c in content]}], + "status": "completed", + }, + }, + ] + + @staticmethod + def _masking_guardrail() -> CustomGuardrail: + class MaskWorld(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + texts = inputs.get("texts", []) + return {**inputs, "texts": [t.replace("world", "[MASKED]") for t in texts]} + + return MaskWorld(guardrail_name="test-mask") + + @pytest.mark.asyncio + async def test_deliver_ended_stream_rewrites_syncs_all_stream_events(self): + handler = OpenAIResponsesHandler() + events = self._ended_stream_events() + + result = await handler.process_output_streaming_response( + responses_so_far=events, + guardrail_to_apply=self._masking_guardrail(), + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + assert result is events + assert events[0]["delta"] == "hello [MASKED]" + assert events[1]["delta"] == "" + assert events[2]["text"] == "hello [MASKED]" + assert events[3]["part"]["text"] == "hello [MASKED]" + assert events[4]["item"]["content"][0]["text"] == "hello [MASKED]" + assert events[5]["response"]["output"][0]["content"][0]["text"] == "hello [MASKED]" + + @pytest.mark.asyncio + async def test_ended_stream_rewrite_leaves_delta_events_untouched_by_default(self): + handler = OpenAIResponsesHandler() + events = self._ended_stream_events() + + await handler.process_output_streaming_response( + responses_so_far=events, + guardrail_to_apply=self._masking_guardrail(), + litellm_logging_obj=None, + ) + + assert events[0]["delta"] == "hello " + assert events[1]["delta"] == "world" + assert events[2]["text"] == "hello world" + assert events[5]["response"]["output"][0]["content"][0]["text"] == "hello [MASKED]" + class TestGetStructuredMessages: """Test the get_structured_messages method for Responses API handler.""" diff --git a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py index 52fd8777a19..908c9f12c9e 100644 --- a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py +++ b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py @@ -823,6 +823,8 @@ class _TextReturningGuardrail(CustomGuardrail): class _TextTranslation: + delivers_ended_stream_text_rewrites = False + def __init__(self): self.seen_guardrail_names = [] diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py index 74bd2e483cc..73270ec5671 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -24,7 +24,7 @@ from litellm.integrations.custom_guardrail import ( ModifyResponseException, ) from litellm.integrations.prometheus import PrometheusLogger -from litellm.proxy._types import ProxyException +from litellm.proxy._types import ProxyException, UserAPIKeyAuth from litellm.proxy.common_utils.callback_utils import add_guardrail_to_applied_guardrails_header from litellm.proxy.utils import ProxyLogging, _raise_for_streaming_post_call_pipelines from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ContentFilterGuardrail @@ -1441,7 +1441,7 @@ async def test_pre_call_hook_rejects_streaming_when_pipeline_guardrail_lacks_uni ("guardrail_config", {"streaming_transform_mode": "incremental_diff"}), ], ) -async def test_pre_call_hook_rejects_streaming_when_pipeline_guardrail_rewrites_streamed_content( +async def test_pre_call_hook_allows_streaming_when_pipeline_guardrail_rewrites_streamed_content( proxy_logging, make_user_api_key_auth, monkeypatch, rewrite_attribute, value ): seen: Dict[str, Any] = {} @@ -1450,24 +1450,21 @@ async def test_pre_call_hook_rejects_streaming_when_pipeline_guardrail_rewrites_ monkeypatch.setattr(litellm, "callbacks", [guardrail]) data = _post_call_pipeline_data(stream=True) - with pytest.raises(HTTPException) as info: - await proxy_logging.pre_call_hook( - user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"), - data=data, - call_type="completion", - guardrails_only=True, - ) + out = await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"), + data=data, + call_type="completion", + guardrails_only=True, + ) - assert info.value.status_code == 400 - assert info.value.detail["error"]["guardrails"] == ("gr-post",) - assert "rewrite streamed content" in info.value.detail["error"]["message"] - assert seen.get("count") is None + assert out is not None + assert out.get("stream") is True @pytest.mark.asyncio -@pytest.mark.parametrize("action, rejected", [(ContentFilterAction.MASK, True), (ContentFilterAction.BLOCK, False)]) -async def test_pre_call_hook_rejects_streaming_only_when_content_filter_step_masks( - proxy_logging, make_user_api_key_auth, monkeypatch, action, rejected +@pytest.mark.parametrize("action", [ContentFilterAction.MASK, ContentFilterAction.BLOCK]) +async def test_pre_call_hook_allows_streaming_when_content_filter_step_masks_or_blocks( + proxy_logging, make_user_api_key_auth, monkeypatch, action ): guardrail = ContentFilterGuardrail( guardrail_name="gr-post", @@ -1478,25 +1475,15 @@ async def test_pre_call_hook_rejects_streaming_only_when_content_filter_step_mas data = _post_call_pipeline_data(stream=True) user_api_key_dict = make_user_api_key_auth(request_route="/v1/chat/completions") - if not rejected: - out = await proxy_logging.pre_call_hook( - user_api_key_dict=user_api_key_dict, data=data, call_type="completion", guardrails_only=True - ) - assert out is not None and out.get("stream") is True - return + out = await proxy_logging.pre_call_hook( + user_api_key_dict=user_api_key_dict, data=data, call_type="completion", guardrails_only=True + ) - with pytest.raises(HTTPException) as info: - await proxy_logging.pre_call_hook( - user_api_key_dict=user_api_key_dict, data=data, call_type="completion", guardrails_only=True - ) - - assert info.value.status_code == 400 - assert info.value.detail["error"]["guardrails"] == ("gr-post",) - assert "a MASK action" in info.value.detail["error"]["message"] + assert out is not None and out.get("stream") is True @pytest.mark.asyncio -async def test_pre_call_hook_rejects_streaming_when_content_filter_category_masks( +async def test_pre_call_hook_allows_streaming_when_content_filter_category_masks( proxy_logging, make_user_api_key_auth, monkeypatch ): guardrail = ContentFilterGuardrail( @@ -1508,13 +1495,11 @@ async def test_pre_call_hook_rejects_streaming_when_content_filter_category_mask data = _post_call_pipeline_data(stream=True) user_api_key_dict = make_user_api_key_auth(request_route="/v1/chat/completions") - with pytest.raises(HTTPException) as info: - await proxy_logging.pre_call_hook( - user_api_key_dict=user_api_key_dict, data=data, call_type="completion", guardrails_only=True - ) + out = await proxy_logging.pre_call_hook( + user_api_key_dict=user_api_key_dict, data=data, call_type="completion", guardrails_only=True + ) - assert info.value.status_code == 400 - assert info.value.detail["error"]["guardrails"] == ("gr-post",) + assert out is not None and out.get("stream") is True @pytest.mark.asyncio @@ -1618,17 +1603,10 @@ def _echoed_tool_call_dicts(arguments: str) -> List[Dict[str, Any]]: @pytest.mark.asyncio @pytest.mark.parametrize("on_fail, on_error", [("block", None), ("next", "next")]) -@pytest.mark.parametrize( - "make_chunks, transform", - [ - (_stream_chunks, lambda inputs: {"texts": ["hello [MASKED]"]}), - (_tool_call_stream_chunks, lambda inputs: {"tool_calls": _echoed_tool_call_dicts('{"ssn": "[MASKED]"}')}), - ], - ids=["texts", "tool_calls"], -) -async def test_streaming_iterator_hook_pipeline_withholds_runtime_rewrite( - proxy_logging, make_user_api_key_auth, monkeypatch, make_chunks, transform, on_fail, on_error +async def test_streaming_iterator_hook_pipeline_withholds_runtime_tool_call_rewrite( + proxy_logging, make_user_api_key_auth, monkeypatch, on_fail, on_error ): + transform = lambda inputs: {"tool_calls": _echoed_tool_call_dicts('{"ssn": "[MASKED]"}')} # noqa: E731 monkeypatch.setattr(litellm, "callbacks", [_rewriting_stream_guardrail(transform)]) monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) step = PipelineStep(guardrail="gr-post", on_pass="allow", on_fail=on_fail, on_error=on_error) @@ -1638,7 +1616,7 @@ async def test_streaming_iterator_hook_pipeline_withholds_runtime_rewrite( async def _drain() -> None: async for item in proxy_logging.async_post_call_streaming_iterator_hook( user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"), - response=_async_chunk_iter(make_chunks()), + response=_async_chunk_iter(_tool_call_stream_chunks()), request_data=data, ): delivered.append(item) @@ -1655,6 +1633,81 @@ async def test_streaming_iterator_hook_pipeline_withholds_runtime_rewrite( assert "stream=false" in error["message"] +@pytest.mark.asyncio +async def test_streaming_iterator_hook_pipeline_delivers_runtime_text_rewrite( + proxy_logging, make_user_api_key_auth, monkeypatch +): + transform = lambda inputs: {"texts": ["hello [MASKED]"]} # noqa: E731 + monkeypatch.setattr(litellm, "callbacks", [_rewriting_stream_guardrail(transform)]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data(stream=True) + chunks = _stream_chunks() + + delivered = [ + item + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"), + response=_async_chunk_iter(chunks), + request_data=data, + ) + ] + + assert [id(item) for item in delivered] == [id(chunk) for chunk in chunks] + assert delivered[0].choices[0].delta.content == "hello [MASKED]" + assert delivered[1].choices[0].delta.content in (None, "") + assert delivered[1].choices[0].finish_reason == "stop" + + +@pytest.mark.asyncio +async def test_streaming_iterator_hook_pipeline_chains_text_rewrites_across_steps( + proxy_logging, make_user_api_key_auth, monkeypatch +): + second_step_saw: Dict[str, Any] = {} + + class FirstMask(CustomGuardrail): + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + return {**inputs, "texts": [text.replace("world", "[MASKED]") for text in inputs["texts"]]} + + class SecondMask(CustomGuardrail): + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + second_step_saw["texts"] = list(inputs["texts"]) + return {**inputs, "texts": [text.replace("hello", "[GREETING]") for text in inputs["texts"]]} + + monkeypatch.setattr( + litellm, + "callbacks", + [ + FirstMask(guardrail_name="gr-first", event_hook=GuardrailEventHooks.post_call, default_on=False), + SecondMask(guardrail_name="gr-second", event_hook=GuardrailEventHooks.post_call, default_on=False), + ], + ) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + pipeline = GuardrailPipeline( + mode="post_call", + steps=[ + PipelineStep(guardrail="gr-first", on_pass="next", on_fail="block"), + PipelineStep(guardrail="gr-second", on_pass="allow", on_fail="block"), + ], + ) + data = _post_call_pipeline_data(stream=True) + data["metadata"]["_guardrail_pipelines"] = [("response-governance", pipeline)] + chunks = _stream_chunks() + + delivered = [ + item + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"), + response=_async_chunk_iter(chunks), + request_data=data, + ) + ] + + assert second_step_saw["texts"] == ["hello [MASKED]"] + assert delivered[0].choices[0].delta.content == "[GREETING] [MASKED]" + assert delivered[1].choices[0].delta.content in (None, "") + assert delivered[1].choices[0].finish_reason == "stop" + + @pytest.mark.asyncio @pytest.mark.parametrize( "make_chunks, transform", @@ -1761,6 +1814,79 @@ async def test_streaming_iterator_hook_pipeline_modify_response_emits_translated assert not any(item is chunk for item in delivered for chunk in chunks) +@pytest.mark.asyncio +async def test_streaming_iterator_hook_pipeline_delivers_text_rewrite_on_anthropic_sse( + proxy_logging, make_user_api_key_auth, monkeypatch +): + transform = lambda inputs: {"texts": ["hello [MASKED]"]} # noqa: E731 + monkeypatch.setattr(litellm, "callbacks", [_rewriting_stream_guardrail(transform)]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data(stream=True) + chunks = _anthropic_sse_chunks() + + delivered = [ + item + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/v1/messages"), + response=_async_chunk_iter(chunks), + request_data=data, + ) + ] + + raw = b"".join(delivered).decode() + assert "hello [MASKED]" in raw + assert "hello world" not in raw + assert raw.count("event: content_block_delta") == 1 + for expected_event in ("message_start", "content_block_start", "content_block_stop", "message_delta", "message_stop"): + assert f"event: {expected_event}" in raw + + +@pytest.mark.asyncio +async def test_pipeline_executor_withholds_text_rewrite_when_translation_lacks_write_back(monkeypatch): + from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation + from litellm.proxy.policy_engine.pipeline_executor import PipelineExecutor, UndeliverableStreamRewrite + + class NoWriteBackTranslation(BaseTranslation): + async def process_input_messages(self, data, guardrail_to_apply, litellm_logging_obj): + return data + + async def process_output_response(self, response, guardrail_to_apply, litellm_logging_obj, **kwargs): + return response + + async def process_output_streaming_response( + self, + responses_so_far, + guardrail_to_apply, + litellm_logging_obj, + user_api_key_dict=None, + request_data=None, + stream_transform_sink=None, + deliver_ended_stream_rewrites=False, + ): + assert deliver_ended_stream_rewrites is False + await guardrail_to_apply.apply_guardrail( + inputs={"texts": ["hello world"]}, + request_data=request_data or {}, + input_type="response", + ) + return responses_so_far + + transform = lambda inputs: {"texts": ["hello [MASKED]"]} # noqa: E731 + monkeypatch.setattr(litellm, "callbacks", [_rewriting_stream_guardrail(transform)]) + + with pytest.raises(UndeliverableStreamRewrite): + await PipelineExecutor.execute_steps( + steps=[PipelineStep(guardrail="gr-post", on_pass="allow", on_fail="block")], + mode="post_call", + data={"metadata": {}}, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + call_type="acompletion", + policy_name="response-governance", + streaming_chunks=_stream_chunks(), + endpoint_translation=NoWriteBackTranslation(), + ) + + @pytest.mark.asyncio async def test_streaming_iterator_hook_pipeline_gates_without_iterator_overrides( proxy_logging, make_user_api_key_auth, monkeypatch From 85fea1a6752ac5f0c4c0415c5d3da4f6bc542911 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 1 Sep 2026 17:15:11 -0700 Subject: [PATCH 17/60] fix(policy_engine): move mutable-ok suppressions onto the flagged lines --- litellm/llms/openai/chat/guardrail_translation/handler.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index fcb447d5277..02206a36b0e 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -989,12 +989,8 @@ class OpenAIChatCompletionsHandler(BaseTranslation): return await self._apply_guardrail_responses_to_output_streaming( responses=responses_so_far, - guardrailed_texts=[ - after for _choice_idx, after in changed - ], # mutable-ok: the callee's signature predates this change and takes lists - task_mappings=[ - (choice_idx, None) for choice_idx, _after in changed - ], # mutable-ok: the callee's signature predates this change and takes lists + guardrailed_texts=[after for _choice_idx, after in changed], # mutable-ok: callee takes lists + task_mappings=[(choice_idx, None) for choice_idx, _after in changed], # mutable-ok: callee takes lists ) async def _apply_guardrail_responses_to_output_streaming( From 4fbe4ce2e225940b009d390543bd76694eece343 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 1 Sep 2026 17:29:52 -0700 Subject: [PATCH 18/60] fix(guardrail_translation): deliver stream rewrites on incomplete and failed responses terminals --- .../guardrail_translation/handler.py | 60 ++++++++++++------- ...test_openai_responses_guardrail_handler.py | 56 +++++++++++++++++ 2 files changed, 93 insertions(+), 23 deletions(-) diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index ac894401628..3c902f1b829 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -86,6 +86,15 @@ class ResponsesStreamChunk(TypedDict, total=False): text: ReadOnly[str] +_TERMINAL_ENVELOPE_EVENT_TYPES: Final = frozenset( + { + ResponsesAPIStreamEvents.RESPONSE_COMPLETED.value, + ResponsesAPIStreamEvents.RESPONSE_FAILED.value, + ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE.value, + } +) + + def _next_stream_sequence_number(responses_so_far: Sequence[Any] | None) -> int: sequence_numbers: Final = ( item.get("sequence_number") if isinstance(item, dict) else getattr(item, "sequence_number", None) @@ -507,14 +516,17 @@ class OpenAIResponsesHandler(BaseTranslation): chunk, apply the guardrail, then write the result back in-place so the caller sees the modified content (e.g. PII tokens replaced). - For ``response.completed`` events (the normal end-of-stream signal) we - use the same per-item extraction + task-mapping approach as - ``process_output_response`` so that unmasking / blocking works correctly - for every output item. With ``deliver_ended_stream_rewrites`` the earlier - text-carrying events (``response.output_text.delta`` / ``.done``, + For terminal envelope events (``response.completed``, and equally + ``response.incomplete`` / ``response.failed``, whose envelopes carry the + partial output) we use the same per-item extraction + task-mapping + approach as ``process_output_response`` so that unmasking / blocking + works correctly for every output item. With + ``deliver_ended_stream_rewrites`` the earlier text-carrying events + (``response.output_text.delta`` / ``.done``, ``response.content_part.done``, ``response.output_item.done``) are synced - to the rewritten completed response too, so a client reading deltas sees - the rewrite instead of the raw model output. + to the rewritten envelope too, so a client reading deltas sees the + rewrite instead of the raw model output; a rewrite observed where no + write-back is possible fails closed instead of releasing raw output. """ if not responses_so_far: return responses_so_far @@ -526,14 +538,16 @@ class OpenAIResponsesHandler(BaseTranslation): return responses_so_far # ------------------------------------------------------------------ # - # Case 1: response.completed — full response is available in the # - # final chunk; iterate output items, apply guardrail, write back. # + # Case 1: terminal envelope events (completed/incomplete/failed). # + # the accumulated response is available in the final chunk; iterate # + # output items, apply guardrail, write back. Falls through to the # + # string fallback when the envelope yields nothing to check. # # ------------------------------------------------------------------ # - if final_chunk.get("type") == "response.completed": + if final_chunk.get("type") in _TERMINAL_ENVELOPE_EVENT_TYPES: response_obj: Final[ResponseOutputEnvelope] = final_chunk.get("response") or {} - if not hasattr(response_obj, "get"): - return responses_so_far - outputs: Final[Sequence[object]] = response_obj.get("output") or [] + outputs: Final[Sequence[object]] = ( + (response_obj.get("output") or []) if hasattr(response_obj, "get") else [] + ) texts_to_check: Final[list[str]] = [] tool_calls_to_check: Final[list[ChatCompletionToolCallChunk]] = [] @@ -596,8 +610,7 @@ class OpenAIResponsesHandler(BaseTranslation): stream_events=responses_so_far[:-1], rewrites_by_position=rewrites_by_position, ) - - return responses_so_far + return responses_so_far # ------------------------------------------------------------------ # # Case 2: response.output_item.done — extract tool calls only. # @@ -623,7 +636,8 @@ class OpenAIResponsesHandler(BaseTranslation): # ------------------------------------------------------------------ # # Fallback: apply guardrail to the accumulated text string. # # No structured write-back is possible here; guardrails that only # - # need to block/flag (not rewrite) still work correctly. # + # need to block/flag (not rewrite) still work correctly, and a # + # rewrite a caller expects delivered fails closed instead. # # ------------------------------------------------------------------ # string_so_far: Final = self.get_streaming_string_so_far(responses_so_far) if string_so_far: @@ -633,12 +647,17 @@ class OpenAIResponsesHandler(BaseTranslation): ) if response_model: fallback_inputs["model"] = response_model - await guardrail_to_apply.apply_guardrail( + fallback_outputs: Final = await guardrail_to_apply.apply_guardrail( inputs=fallback_inputs, request_data=request_data if request_data is not None else {}, input_type="response", logging_obj=litellm_logging_obj, ) + fallback_texts: Final = fallback_outputs.get("texts") + if deliver_ended_stream_rewrites and fallback_texts and tuple(fallback_texts) != (string_so_far,): + from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite + + raise UndeliverableStreamRewrite(guardrail_to_apply.guardrail_name or "unknown") return responses_so_far @staticmethod @@ -704,12 +723,7 @@ class OpenAIResponsesHandler(BaseTranslation): """ if not responses_so_far: return False - terminal_types: Final = { - ResponsesAPIStreamEvents.RESPONSE_COMPLETED.value, - ResponsesAPIStreamEvents.RESPONSE_FAILED.value, - ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE.value, - } - return responses_so_far[-1].get("type") in terminal_types + return responses_so_far[-1].get("type") in _TERMINAL_ENVELOPE_EVENT_TYPES def build_stream_error_items( self, diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py index 4f95e08cb71..dfd6352f9ec 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -1171,6 +1171,62 @@ class TestOpenAIResponsesHandlerStreamingOutputProcessing: assert events[4]["item"]["content"][0]["text"] == "hello [MASKED]" assert events[5]["response"]["output"][0]["content"][0]["text"] == "hello [MASKED]" + @pytest.mark.asyncio + @pytest.mark.parametrize("terminal_type", ["response.incomplete", "response.failed"]) + async def test_deliver_ended_stream_rewrites_syncs_non_completed_terminals(self, terminal_type): + handler = OpenAIResponsesHandler() + events = self._ended_stream_events() + events[-1]["type"] = terminal_type + events[-1]["response"]["status"] = terminal_type.split(".")[-1] + + result = await handler.process_output_streaming_response( + responses_so_far=events, + guardrail_to_apply=self._masking_guardrail(), + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + assert result is events + assert events[0]["delta"] == "hello [MASKED]" + assert events[1]["delta"] == "" + assert events[2]["text"] == "hello [MASKED]" + assert events[3]["part"]["text"] == "hello [MASKED]" + assert events[4]["item"]["content"][0]["text"] == "hello [MASKED]" + assert events[5]["response"]["output"][0]["content"][0]["text"] == "hello [MASKED]" + + @pytest.mark.asyncio + async def test_fallback_rewrite_with_delivery_expected_fails_closed(self): + from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite + + handler = OpenAIResponsesHandler() + events = [ + {"type": "response.output_text.delta", "output_index": 0, "content_index": 0, "delta": "hello "}, + {"type": "response.output_text.done", "output_index": 0, "content_index": 0, "text": "hello world"}, + ] + + with pytest.raises(UndeliverableStreamRewrite): + await handler.process_output_streaming_response( + responses_so_far=events, + guardrail_to_apply=self._masking_guardrail(), + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + @pytest.mark.asyncio + async def test_fallback_rewrite_without_delivery_expected_does_not_raise(self): + handler = OpenAIResponsesHandler() + events = [ + {"type": "response.output_text.done", "output_index": 0, "content_index": 0, "text": "hello world"}, + ] + + result = await handler.process_output_streaming_response( + responses_so_far=events, + guardrail_to_apply=self._masking_guardrail(), + litellm_logging_obj=None, + ) + + assert result is events + @pytest.mark.asyncio async def test_ended_stream_rewrite_leaves_delta_events_untouched_by_default(self): handler = OpenAIResponsesHandler() From c9435b5ff3f72c6d9c5ebfa30175498bc5928daf Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 1 Sep 2026 18:00:42 -0700 Subject: [PATCH 19/60] fix(guardrails): key stream rewrites by choice index and scan delta-only responses buffers Chat streaming write-backs now match chunks by the choice's index field instead of its list position, delivering rewrites to the right choice on n>1 streams; an ended-stream rewrite on a multi-choice buffer fails closed since stream_chunk_builder collapses the choices. The Responses fallback joins output_text.delta events when delivery is expected, so a delta-only buffer is guardrail-checked instead of released raw. --- .../chat/guardrail_translation/handler.py | 50 ++++++---- .../guardrail_translation/handler.py | 31 +++++-- .../test_openai_guardrail_handler.py | 93 +++++++++++++++++++ ...test_openai_responses_guardrail_handler.py | 71 ++++++++++++++ 4 files changed, 222 insertions(+), 23 deletions(-) diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 02206a36b0e..182dec81937 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -620,6 +620,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): responses_so_far=responses_so_far, guardrailed_response=model_response, pre_guardrail_texts=pre_guardrail_texts, + guardrail_name=guardrail_to_apply.guardrail_name or "unknown", ) def build_stream_error_items( @@ -747,8 +748,8 @@ class OpenAIChatCompletionsHandler(BaseTranslation): """ combined_texts: Final[dict[tuple[int, int | None], str]] = {} - for response_idx, response in enumerate(responses_so_far): - for choice_idx, choice in enumerate(response.choices): + for response in responses_so_far: + for choice in response.choices: if isinstance(choice, litellm.StreamingChoices): content = choice.delta.content elif isinstance(choice, litellm.Choices): @@ -761,7 +762,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): if isinstance(content, str): # String content - accumulate for this choice - str_key: tuple[int, int | None] = (choice_idx, None) + str_key: tuple[int, int | None] = (choice.index, None) if str_key not in combined_texts: combined_texts[str_key] = "" combined_texts[str_key] += content @@ -772,7 +773,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): text_str = content_item.get("text") if text_str: list_key: tuple[int, int | None] = ( - choice_idx, + choice.index, content_idx, ) if list_key not in combined_texts: @@ -973,24 +974,38 @@ class OpenAIChatCompletionsHandler(BaseTranslation): responses_so_far: list["ModelResponseStream"], # mutable-ok: rewrites the caller's buffered chunks in place guardrailed_response: "ModelResponse", pre_guardrail_texts: tuple[str | None, ...], + guardrail_name: str, ) -> None: """Write ended-stream guardrail text rewrites back across the buffered - chunks: each rewritten choice's full text lands in its first + chunks: the full rewritten text lands in the choice's first content-carrying chunk and the rest are blanked, the same shape the in-flight write-back uses. Chunks carrying only finish_reason or usage - stay untouched.""" + stay untouched. A rewrite on a stream carrying more than one distinct + choice index fails closed.""" post_guardrail_texts: Final = self._string_choice_contents(guardrailed_response) changed: Final = tuple( - (choice_idx, after) - for choice_idx, (before, after) in enumerate(zip(pre_guardrail_texts, post_guardrail_texts)) + after + for before, after in zip(pre_guardrail_texts, post_guardrail_texts) if before is not None and after is not None and after != before ) if not changed: return + stream_choice_indices: Final = frozenset( + choice.index for response in responses_so_far for choice in response.choices + ) + if len(stream_choice_indices) != 1: + # stream_chunk_builder collapses every choice into one index-0 + # choice, so a rewrite of the rebuilt response cannot be attributed + # back to a single choice on an n>1 stream: withhold the stream + # rather than deliver the rewrite on the wrong choice + from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite + + raise UndeliverableStreamRewrite(guardrail_name) + target_choice_index: Final = next(iter(stream_choice_indices)) await self._apply_guardrail_responses_to_output_streaming( responses=responses_so_far, - guardrailed_texts=[after for _choice_idx, after in changed], # mutable-ok: callee takes lists - task_mappings=[(choice_idx, None) for choice_idx, _after in changed], # mutable-ok: callee takes lists + guardrailed_texts=list(changed), # mutable-ok: callee takes lists + task_mappings=[(target_choice_index, None) for _ in changed], # mutable-ok: callee takes lists ) async def _apply_guardrail_responses_to_output_streaming( @@ -1008,7 +1023,8 @@ class OpenAIChatCompletionsHandler(BaseTranslation): Args: responses: List of ModelResponseStream objects to modify guardrailed_texts: List of guardrailed text responses (combined from all chunks) - task_mappings: List of tuples (choice_idx, content_idx) + task_mappings: List of tuples (choice_idx, content_idx), where choice_idx + is the choice's ``index`` field, not its position in a chunk's list Override this method to customize how responses are applied to streaming responses. """ @@ -1024,9 +1040,11 @@ class OpenAIChatCompletionsHandler(BaseTranslation): # Key: (choice_idx, content_idx), Value: boolean (True if already set) already_set: Final[dict[tuple[int, int | None], bool]] = {} - # Iterate through all responses and update content - for response_idx, response in enumerate(responses): - for choice_idx_in_response, choice in enumerate(response.choices): + # Iterate through all responses and update content, matching each chunk's + # choice by its index field: on n>1 streams a chunk usually carries one + # choice at list position 0 whose index names the logical choice. + for response in responses: + for choice in response.choices: if isinstance(choice, litellm.StreamingChoices): content = choice.delta.content elif isinstance(choice, litellm.Choices): @@ -1039,7 +1057,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): if isinstance(content, str): # String content - str_key: tuple[int, int | None] = (choice_idx_in_response, None) + str_key: tuple[int, int | None] = (choice.index, None) if str_key in guardrail_map: if str_key not in already_set: # First chunk - set the complete guardrailed text @@ -1060,7 +1078,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): for content_idx, content_item in enumerate(content): if "text" in content_item: list_key: tuple[int, int | None] = ( - choice_idx_in_response, + choice.index, content_idx, ) if list_key in guardrail_map: diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 3c902f1b829..ebd1070a8e8 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -634,14 +634,20 @@ class OpenAIResponsesHandler(BaseTranslation): return responses_so_far # ------------------------------------------------------------------ # - # Fallback: apply guardrail to the accumulated text string. # - # No structured write-back is possible here; guardrails that only # - # need to block/flag (not rewrite) still work correctly, and a # - # rewrite a caller expects delivered fails closed instead. # + # Fallback: apply guardrail to the accumulated text string. When a # + # caller expects rewrites delivered and only output_text.delta events # + # carried the text (a stream cut off before any .done or terminal # + # envelope), the delta text is scanned instead so nothing escapes # + # unchecked. No structured write-back is possible here; guardrails # + # that only need to block/flag (not rewrite) still work correctly, # + # and a rewrite a caller expects delivered fails closed instead. # # ------------------------------------------------------------------ # string_so_far: Final = self.get_streaming_string_so_far(responses_so_far) - if string_so_far: - fallback_inputs: Final = GenericGuardrailAPIInputs(texts=[string_so_far]) + text_to_check: Final = string_so_far or ( + self._delta_text_so_far(responses_so_far) if deliver_ended_stream_rewrites else "" + ) + if text_to_check: + fallback_inputs: Final = GenericGuardrailAPIInputs(texts=[text_to_check]) response_model = ( final_chunk.get("response", {}).get("model") if isinstance(final_chunk.get("response"), dict) else None ) @@ -654,7 +660,7 @@ class OpenAIResponsesHandler(BaseTranslation): logging_obj=litellm_logging_obj, ) fallback_texts: Final = fallback_outputs.get("texts") - if deliver_ended_stream_rewrites and fallback_texts and tuple(fallback_texts) != (string_so_far,): + if deliver_ended_stream_rewrites and fallback_texts and tuple(fallback_texts) != (text_to_check,): from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite raise UndeliverableStreamRewrite(guardrail_to_apply.guardrail_name or "unknown") @@ -754,6 +760,17 @@ class OpenAIResponsesHandler(BaseTranslation): """ return "".join([response.get("text", "") for response in responses_so_far]) + @staticmethod + def _delta_text_so_far(responses_so_far: Sequence[ResponsesStreamChunk]) -> str: + """Accumulate the text carried by ``response.output_text.delta`` events, + for buffers where no ``.done`` event or terminal envelope repeats it.""" + deltas: Final = ( + response.get("delta") + for response in responses_so_far + if response.get("type") == ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA.value + ) + return "".join(delta for delta in deltas if isinstance(delta, str)) + def _has_text_content(self, response: "ResponsesAPIResponse") -> bool: """ Check if response has any text content to process. diff --git a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index 26442a4a6ed..b177d4a73d4 100644 --- a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -1128,6 +1128,99 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput: assert chunks[1].choices[0].delta.content == " world" assert chunks[1].choices[0].finish_reason == "stop" + @staticmethod + def _two_choice_stream_chunks() -> list: + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + def chunk(index: int, content: str, finish_reason: Optional[str] = None) -> ModelResponseStream: + return ModelResponseStream( + id="chatcmpl-123", + created=1234567890, + model="gpt-4", + object="chat.completion.chunk", + choices=[StreamingChoices(index=index, delta=Delta(content=content), finish_reason=finish_reason)], + ) + + return [ + chunk(0, "safe "), + chunk(1, "hello "), + chunk(0, "text", "stop"), + chunk(1, "world", "stop"), + ] + + @staticmethod + def _world_masking_guardrail() -> CustomGuardrail: + class MaskWorld(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + texts = inputs.get("texts", []) + return {**inputs, "texts": [t.replace("world", "[MASKED]") for t in texts]} + + return MaskWorld(guardrail_name="test-mask") + + @pytest.mark.asyncio + async def test_deliver_ended_stream_rewrite_on_multi_choice_stream_fails_closed(self): + from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite + + handler = OpenAIChatCompletionsHandler() + chunks = self._two_choice_stream_chunks() + + with pytest.raises(UndeliverableStreamRewrite): + await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=self._world_masking_guardrail(), + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + @pytest.mark.asyncio + async def test_deliver_ended_stream_clean_multi_choice_stream_released_untouched(self): + handler = OpenAIChatCompletionsHandler() + chunks = self._two_choice_stream_chunks() + + result = await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=MockPassThroughGuardrail(guardrail_name="test"), + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + assert result is chunks + assert [c.choices[0].delta.content for c in chunks] == ["safe ", "hello ", "text", "world"] + + @pytest.mark.asyncio + async def test_deliver_ended_stream_rewrite_lands_on_nonzero_choice_index(self): + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + handler = OpenAIChatCompletionsHandler() + + def chunk(content: str, finish_reason: Optional[str]) -> ModelResponseStream: + return ModelResponseStream( + id="chatcmpl-123", + created=1234567890, + model="gpt-4", + object="chat.completion.chunk", + choices=[StreamingChoices(index=1, delta=Delta(content=content), finish_reason=finish_reason)], + ) + + chunks = [chunk("hello ", None), chunk("world", "stop")] + + result = await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=self._world_masking_guardrail(), + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + assert result is chunks + assert chunks[0].choices[0].delta.content == "hello [MASKED]" + assert chunks[1].choices[0].delta.content in (None, "") + class TestGetStructuredMessages: """Test the get_structured_messages method.""" diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py index dfd6352f9ec..ae3700b1123 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -1212,6 +1212,77 @@ class TestOpenAIResponsesHandlerStreamingOutputProcessing: deliver_ended_stream_rewrites=True, ) + @staticmethod + def _recording_guardrail() -> "tuple[CustomGuardrail, List[List[str]]]": + seen: List[List[str]] = [] + + class Recorder(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + seen.append(list(inputs.get("texts", []))) + return inputs + + return Recorder(guardrail_name="recorder"), seen + + @pytest.mark.asyncio + async def test_fallback_delta_only_rewrite_with_delivery_expected_fails_closed(self): + from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite + + handler = OpenAIResponsesHandler() + events = [ + {"type": "response.output_text.delta", "output_index": 0, "content_index": 0, "delta": "hello "}, + {"type": "response.output_text.delta", "output_index": 0, "content_index": 0, "delta": "world"}, + ] + + with pytest.raises(UndeliverableStreamRewrite): + await handler.process_output_streaming_response( + responses_so_far=events, + guardrail_to_apply=self._masking_guardrail(), + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + @pytest.mark.asyncio + async def test_fallback_scans_delta_text_when_delivery_expected(self): + handler = OpenAIResponsesHandler() + guardrail, seen = self._recording_guardrail() + events = [ + {"type": "response.output_text.delta", "output_index": 0, "content_index": 0, "delta": "hello "}, + {"type": "response.output_text.delta", "output_index": 0, "content_index": 0, "delta": "there"}, + ] + + result = await handler.process_output_streaming_response( + responses_so_far=events, + guardrail_to_apply=guardrail, + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + assert result is events + assert seen == [["hello there"]] + + @pytest.mark.asyncio + async def test_fallback_ignores_delta_text_without_delivery_expected(self): + handler = OpenAIResponsesHandler() + guardrail, seen = self._recording_guardrail() + events = [ + {"type": "response.output_text.delta", "output_index": 0, "content_index": 0, "delta": "hello world"}, + ] + + result = await handler.process_output_streaming_response( + responses_so_far=events, + guardrail_to_apply=guardrail, + litellm_logging_obj=None, + ) + + assert result is events + assert seen == [] + @pytest.mark.asyncio async def test_fallback_rewrite_without_delivery_expected_does_not_raise(self): handler = OpenAIResponsesHandler() From c09db7c7a3d77d3ccd6a1a2163778b6f22f29a8f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 1 Sep 2026 22:01:04 -0700 Subject: [PATCH 20/60] Fail closed on rewrites for buffers that never reached their terminal event An Anthropic buffer without a stop_reason only ran the flat text scan, so a rewrite there was dropped while the executor trusted the translation to have delivered it. A Responses buffer ending at response.output_item.done returned after the tool-call scan without ever checking the text. Both now reach the flat scan and raise UndeliverableStreamRewrite when a caller expects the rewrite delivered, matching the existing Responses no-envelope fallback. --- .../chat/guardrail_translation/handler.py | 8 +++- .../guardrail_translation/handler.py | 7 ++- .../test_anthropic_guardrail_handler.py | 46 +++++++++++++++++++ ...test_openai_responses_guardrail_handler.py | 46 +++++++++++++++++++ 4 files changed, 104 insertions(+), 3 deletions(-) diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index 091ad8ecba9..fc276dd7fe6 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -1020,7 +1020,8 @@ class AnthropicMessagesHandler(BaseTranslation): Get the string so far, check the apply guardrail to the string so far, and return the list of responses so far. With ``deliver_ended_stream_rewrites``, an ended stream whose guardrail rewrote the text gets the rewrite - written back across the buffered chunks (full rewritten text in the first ``text_delta``, the rest blanked). + written back across the buffered chunks (full rewritten text in the first ``text_delta``, the rest blanked); + a rewrite on a stream that never reported a ``stop_reason`` has no write-back and fails closed instead. """ from litellm.integrations.custom_guardrail import ModifyResponseException @@ -1098,6 +1099,11 @@ class AnthropicMessagesHandler(BaseTranslation): if e.original_response is None: e.original_response = self._build_streaming_usage_response(responses_so_far, request_data) raise + unended_texts: Final = _guardrailed_inputs.get("texts") + if deliver_ended_stream_rewrites and unended_texts and tuple(unended_texts) != (string_so_far,): + from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite + + raise UndeliverableStreamRewrite(guardrail_to_apply.guardrail_name or "unknown") return responses_so_far def _prepare_request_data( diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 8b808348c79..724a0a1d2f0 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -638,7 +638,9 @@ class OpenAIResponsesHandler(BaseTranslation): return responses_so_far # ------------------------------------------------------------------ # - # Case 2: response.output_item.done — extract tool calls only. # + # Case 2: response.output_item.done — extract tool calls only, then # + # fall through to the text fallback when a caller expects rewrites # + # delivered, so a buffer truncated here still fails closed on text. # # ------------------------------------------------------------------ # if final_chunk.get("type") == "response.output_item.done": model_response_stream: Final = ( @@ -656,7 +658,8 @@ class OpenAIResponsesHandler(BaseTranslation): input_type="response", logging_obj=litellm_logging_obj, ) - return responses_so_far + if not deliver_ended_stream_rewrites: + return responses_so_far # ------------------------------------------------------------------ # # Fallback: apply guardrail to the accumulated text string. # diff --git a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py index f60655cd7ec..274c351ebc7 100644 --- a/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py +++ b/tests/test_litellm/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py @@ -328,6 +328,52 @@ class TestAnthropicMessagesHandlerStreamingOutputProcessing: assert chunks == original + @pytest.mark.asyncio + async def test_unended_stream_rewrite_with_delivery_expected_fails_closed(self): + from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite + + handler = AnthropicMessagesHandler() + chunks = self._ended_sse_chunks()[:-2] + + with pytest.raises(UndeliverableStreamRewrite): + await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=self._masking_guardrail(), + litellm_logging_obj=MagicMock(), + deliver_ended_stream_rewrites=True, + ) + + @pytest.mark.asyncio + async def test_unended_stream_without_rewrite_is_released_with_delivery_expected(self): + handler = AnthropicMessagesHandler() + chunks = self._ended_sse_chunks()[:-2] + original = [bytes(chunk) for chunk in chunks] + + result = await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=MockPassThroughGuardrail(guardrail_name="test"), + litellm_logging_obj=MagicMock(), + deliver_ended_stream_rewrites=True, + ) + + assert result is chunks + assert chunks == original + + @pytest.mark.asyncio + async def test_unended_stream_rewrite_without_delivery_expected_does_not_raise(self): + handler = AnthropicMessagesHandler() + chunks = self._ended_sse_chunks()[:-2] + original = [bytes(chunk) for chunk in chunks] + + result = await handler.process_output_streaming_response( + responses_so_far=chunks, + guardrail_to_apply=self._masking_guardrail(), + litellm_logging_obj=MagicMock(), + ) + + assert result is chunks + assert chunks == original + class TestAnthropicMessagesHandlerInputProcessing: """Test input processing preserves litellm_metadata for dynamic guardrails.""" diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py index 155b9ef9810..bd07e924d12 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -1248,6 +1248,52 @@ class TestOpenAIResponsesHandlerStreamingOutputProcessing: deliver_ended_stream_rewrites=True, ) + @pytest.mark.asyncio + async def test_output_item_done_last_rewrite_with_delivery_expected_fails_closed(self): + from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite + + handler = OpenAIResponsesHandler() + events = self._ended_stream_events()[:-1] + + with pytest.raises(UndeliverableStreamRewrite): + await handler.process_output_streaming_response( + responses_so_far=events, + guardrail_to_apply=self._masking_guardrail(), + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + @pytest.mark.asyncio + async def test_output_item_done_last_scans_text_with_delivery_expected(self): + handler = OpenAIResponsesHandler() + events = self._ended_stream_events()[:-1] + guardrail = MockRecordingGuardrail(guardrail_name="test") + + result = await handler.process_output_streaming_response( + responses_so_far=events, + guardrail_to_apply=guardrail, + litellm_logging_obj=None, + deliver_ended_stream_rewrites=True, + ) + + assert result is events + assert [inputs.get("texts") for inputs in guardrail.seen_inputs] == [["hello world"]] + + @pytest.mark.asyncio + async def test_output_item_done_last_without_delivery_expected_skips_text(self): + handler = OpenAIResponsesHandler() + events = self._ended_stream_events()[:-1] + guardrail = MockRecordingGuardrail(guardrail_name="test") + + result = await handler.process_output_streaming_response( + responses_so_far=events, + guardrail_to_apply=guardrail, + litellm_logging_obj=None, + ) + + assert result is events + assert guardrail.seen_inputs == [] + @pytest.mark.asyncio async def test_fallback_rewrite_without_delivery_expected_does_not_raise(self): handler = OpenAIResponsesHandler() From 2a2c49ad4cd711aded30da0384c1a200ddfca5d7 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Thu, 3 Sep 2026 16:09:44 -0400 Subject: [PATCH 21/60] feat(ui): surface batch results on the logs page --- .../LogDetailContent.test.tsx | 78 ++++++++++++++++ .../LogDetailsDrawer/LogDetailContent.tsx | 61 +++++++++++++ .../RequestLogsTableColumns.test.tsx | 53 +++++++++++ .../view_logs/RequestLogsTableColumns.tsx | 29 +++++- .../components/view_logs/TypeBadges.test.tsx | 14 ++- .../src/components/view_logs/TypeBadges.tsx | 26 ++++++ .../view_logs/batchLogUtils.test.ts | 90 +++++++++++++++++++ .../src/components/view_logs/batchLogUtils.ts | 67 ++++++++++++++ .../src/components/view_logs/constants.ts | 3 + 9 files changed, 418 insertions(+), 3 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/view_logs/batchLogUtils.test.ts create mode 100644 ui/litellm-dashboard/src/components/view_logs/batchLogUtils.ts diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.test.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.test.tsx index a679dc49427..91778d2a98a 100644 --- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.test.tsx @@ -113,6 +113,84 @@ describe("LogDetailContent", () => { expect(screen.getAllByText("$0.00200000").length).toBeGreaterThanOrEqual(1); }); + it("shows reasoning tokens in Metrics when the usage breakout carries them", () => { + render( + , + ); + + expect(screen.getByText("Reasoning Tokens")).toBeInTheDocument(); + expect(screen.getByText("224")).toBeInTheDocument(); + }); + + it("hides the reasoning metric when the breakout is absent or zero", () => { + render( + , + ); + + expect(screen.queryByText("Reasoning Tokens")).not.toBeInTheDocument(); + }); + + describe("Batch Results section", () => { + const batchCostEntry = (metadata: Record) => + createLogEntry({ + request_id: "batch_abc123_batch_cost", + call_type: "aretrieve_batch", + metadata: { status: "success", ...metadata }, + }); + + it("renders batch id, per-request outcome counts, and batch models for a batch cost row", () => { + render( + , + ); + + const section = screen.getByText("Batch Results").closest('[data-slot="card"]') as HTMLElement; + expect(within(section).getByText("batch_abc123")).toBeInTheDocument(); + expect(within(section).getByText("2")).toBeInTheDocument(); + expect(within(section).getByText("1")).toBeInTheDocument(); + expect(within(section).getByText("gemini-2.5-flash")).toBeInTheDocument(); + }); + + it("still renders the batch id when a legacy row carries no counts", () => { + render(); + + expect(screen.getByText("Batch Results")).toBeInTheDocument(); + expect(screen.getByText("batch_abc123")).toBeInTheDocument(); + expect(screen.queryByText("Successful Requests")).not.toBeInTheDocument(); + }); + + it("never renders for a non-batch call type", () => { + render( + , + ); + + expect(screen.queryByText("Batch Results")).not.toBeInTheDocument(); + }); + }); + it("should show Input Tokens and Output Tokens for anthropic_messages when uncached text_tokens exist", () => { render( + {/* Batch Results */} + {isBatchCallType(logEntry.call_type) && } + {/* Routing */} @@ -374,6 +384,53 @@ function MetricLabel({ label, tooltip, docsUrl }: { label: string; tooltip: stri ); } +/** + * Aggregate per-request outcomes for a batch cost row: batch id, success/failure counts + * from the parsed output and error files, and the models the batch actually ran on. + */ +function BatchResultsSection({ logEntry, metadata }: { logEntry: LogEntry; metadata: Record }) { + const counts = getBatchRequestCounts(metadata); + const batchId = getBatchIdFromRequestId(logEntry.request_id); + const batchModels = getBatchModels(metadata); + if (!counts && !batchId && !batchModels) return null; + + return ( +
+ + + Batch Results + + + + {batchId && ( + + + + )} + {counts && ( + <> + + {formatNumberWithCommas(counts.successful)} + + + {counts.failed > 0 ? ( + + {formatNumberWithCommas(counts.failed)} + + ) : ( + formatNumberWithCommas(counts.failed) + )} + + + )} + {batchModels && {batchModels.join(", ")}} + + + +
+ ); +} + function MetricsSection({ logEntry, metadata }: { logEntry: LogEntry; metadata: Record }) { const completionStartTime = logEntry.completionStartTime; const ttftMs = @@ -391,6 +448,7 @@ function MetricsSection({ logEntry, metadata }: { logEntry: LogEntry; metadata: const uncachedInputTokens = getUncachedInputTextTokens(metadata); const showAnthropicMessagesInputOutput = logEntry.call_type === "anthropic_messages" && uncachedInputTokens !== undefined; + const reasoningTokens = getReasoningTokens(metadata); return (
@@ -416,6 +474,9 @@ function MetricsSection({ logEntry, metadata }: { logEntry: LogEntry; metadata: /> )} + {reasoningTokens !== undefined && reasoningTokens > 0 && ( + {formatNumberWithCommas(reasoningTokens)} + )} ${formatNumberWithCommas(logEntry.spend || 0, 8)} {logEntry.request_duration_ms != null ? (logEntry.request_duration_ms / 1000).toFixed(3) : "-"} s diff --git a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.test.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.test.tsx index 3c9e6543c1c..3e8e522afe1 100644 --- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.test.tsx @@ -134,6 +134,59 @@ describe("Type column", () => { expect(screen.getByText("MCP")).toBeInTheDocument(); }); + + it("marks a batch cost row with the Batch badge instead of LLM", () => { + renderRows([logEntry({ request_id: "batch_1_batch_cost", call_type: "aretrieve_batch" })]); + + expect(screen.getByText("Batch")).toBeInTheDocument(); + expect(screen.queryByText("LLM")).not.toBeInTheDocument(); + }); +}); + +describe("batch rows", () => { + const batchRow = (overrides: Partial): LogEntry => + logEntry({ + request_id: "batch_abc123_batch_cost", + call_type: "aretrieve_batch", + ...overrides, + }); + + it("rolls partial failures into the status badge instead of reporting blanket Success", async () => { + const user = userEvent.setup(); + renderRows([batchRow({ metadata: { batch_successful_requests: 2, batch_failed_requests: 1 } })]); + + expect(screen.queryByText("Success")).not.toBeInTheDocument(); + await user.hover(screen.getByText("2/3 succeeded")); + expect(await screen.findByText("1 of 3 batch requests failed")).toBeInTheDocument(); + }); + + it("keeps the Success badge when every batch request succeeded", () => { + renderRows([batchRow({ metadata: { batch_successful_requests: 3, batch_failed_requests: 0 } })]); + + expect(screen.getByText("Success")).toBeInTheDocument(); + }); + + it("keeps the Failure badge when the batch row itself failed, whatever the counts say", () => { + renderRows([batchRow({ metadata: { status: "failure", batch_successful_requests: 2, batch_failed_requests: 1 } })]); + + expect(screen.getByText("Failure")).toBeInTheDocument(); + expect(screen.queryByText("2/3 succeeded")).not.toBeInTheDocument(); + }); + + it("shows the provider batch id, not the synthetic _batch_cost request id", () => { + renderRows([batchRow({ metadata: { batch_successful_requests: 1, batch_failed_requests: 0 } })]); + + expect(screen.getByText("batch_abc123")).toBeInTheDocument(); + expect(screen.queryByText("batch_abc123_batch_cost")).not.toBeInTheDocument(); + expect(screen.getByText("batch cost")).toBeInTheDocument(); + }); + + it("leaves ordinary request ids untouched", () => { + renderRows([logEntry({ request_id: "chatcmpl-42" })]); + + expect(screen.getByText("chatcmpl-42")).toBeInTheDocument(); + expect(screen.queryByText("batch cost")).not.toBeInTheDocument(); + }); }); describe("Model column", () => { diff --git a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.tsx index 9d4dc4f7898..b4296abe266 100644 --- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.tsx @@ -7,9 +7,10 @@ import { CellTooltip, DateCell, IdCell, MoneyCell, StatusBadge } from "@/compone import { getSpendString } from "@/utils/dataUtils"; import { getProviderLogoAndName } from "../provider_info_helpers"; +import { getBatchIdFromRequestId, getBatchRequestCounts, isBatchCallType } from "./batchLogUtils"; import type { LogEntry } from "./columns"; import { AGENT_CALL_TYPES, MCP_CALL_TYPES } from "./constants"; -import { AgentBadge, AgentIcon, LlmBadge, McpBadge, SparkleIcon, WrenchIcon } from "./TypeBadges"; +import { AgentBadge, AgentIcon, BatchBadge, LlmBadge, McpBadge, SparkleIcon, WrenchIcon } from "./TypeBadges"; export interface RequestLogsTableColumnsDeps { onKeyHashClick: (keyHash: string) => void; @@ -66,6 +67,7 @@ export const getRequestLogsTableColumns = ({ if (sessionCount <= 1) { if (isMcp) return ; if (isAgent) return ; + if (isBatchCallType(log.call_type)) return ; return ; } @@ -106,6 +108,17 @@ export const getRequestLogsTableColumns = ({ cell: ({ row }) => { const status = readMetaString(row.original.metadata, "status") ?? "Success"; const isSuccess = status.toLowerCase() !== "failure"; + const batchCounts = isSuccess ? getBatchRequestCounts(row.original.metadata) : undefined; + if (batchCounts && batchCounts.failed > 0) { + const total = batchCounts.successful + batchCounts.failed; + return ( + + ); + } return ; }, }, @@ -122,7 +135,19 @@ export const getRequestLogsTableColumns = ({ accessorKey: "request_id", header: "Request ID", enableSorting: false, - cell: ({ row }) => , + cell: ({ row }) => { + const log = row.original; + const batchId = isBatchCallType(log.call_type) ? getBatchIdFromRequestId(log.request_id) : undefined; + if (batchId) { + return ( +
+ + batch cost +
+ ); + } + return ; + }, }, { id: "spend", diff --git a/ui/litellm-dashboard/src/components/view_logs/TypeBadges.test.tsx b/ui/litellm-dashboard/src/components/view_logs/TypeBadges.test.tsx index e3467310265..8a964ea2124 100644 --- a/ui/litellm-dashboard/src/components/view_logs/TypeBadges.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/TypeBadges.test.tsx @@ -1,6 +1,6 @@ import { render, screen } from "@testing-library/react"; import { describe, expect, it } from "vitest"; -import { LlmBadge, McpBadge, AgentBadge } from "./TypeBadges"; +import { LlmBadge, McpBadge, AgentBadge, BatchBadge } from "./TypeBadges"; describe("TypeBadges", () => { describe("LlmBadge", () => { @@ -43,4 +43,16 @@ describe("TypeBadges", () => { expect(screen.getByText("12")).toBeInTheDocument(); }); }); + + describe("BatchBadge", () => { + it("should render with default 'Batch' text when no count is provided", () => { + render(); + expect(screen.getByText("Batch")).toBeInTheDocument(); + }); + + it("should render the count when provided", () => { + render(); + expect(screen.getByText("4")).toBeInTheDocument(); + }); + }); }); diff --git a/ui/litellm-dashboard/src/components/view_logs/TypeBadges.tsx b/ui/litellm-dashboard/src/components/view_logs/TypeBadges.tsx index 1ba4365f66e..84079848487 100644 --- a/ui/litellm-dashboard/src/components/view_logs/TypeBadges.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/TypeBadges.tsx @@ -57,6 +57,25 @@ export const AgentIcon = ({ size = 12 }: { size?: number }) => ( ); +/** Stacked-layers icon for Batch API call types (Lucide Layers-style). */ +export const LayersIcon = ({ size = 12 }: { size?: number }) => ( + + + + + +); + export const LlmBadge = ({ count }: { count?: number }) => ( @@ -77,3 +96,10 @@ export const AgentBadge = ({ count }: { count?: number }) => ( {count != null ? count : "Agent"} ); + +export const BatchBadge = ({ count }: { count?: number }) => ( + + + {count != null ? count : "Batch"} + +); diff --git a/ui/litellm-dashboard/src/components/view_logs/batchLogUtils.test.ts b/ui/litellm-dashboard/src/components/view_logs/batchLogUtils.test.ts new file mode 100644 index 00000000000..4da1e389646 --- /dev/null +++ b/ui/litellm-dashboard/src/components/view_logs/batchLogUtils.test.ts @@ -0,0 +1,90 @@ +import { describe, expect, it } from "vitest"; + +import { + getBatchIdFromRequestId, + getBatchModels, + getBatchRequestCounts, + getReasoningTokens, + isBatchCallType, +} from "./batchLogUtils"; + +/** Metadata shape the batch cost poller writes on an aretrieve_batch spend row. */ +const batchCostMetadata = { + batch_models: ["gemini-2.5-flash"], + batch_successful_requests: 2, + batch_failed_requests: 1, + usage_object: { + total_tokens: 270, + prompt_tokens: 14, + completion_tokens: 256, + completion_tokens_details: { text_tokens: 32, reasoning_tokens: 224 }, + }, +}; + +describe("isBatchCallType", () => { + it("recognizes the poller's aretrieve_batch and the create call types", () => { + for (const callType of ["aretrieve_batch", "retrieve_batch", "acreate_batch", "create_batch"]) { + expect(isBatchCallType(callType)).toBe(true); + } + expect(isBatchCallType("acompletion")).toBe(false); + }); +}); + +describe("getBatchRequestCounts", () => { + it("reads both counts off a batch cost row", () => { + expect(getBatchRequestCounts(batchCostMetadata)).toEqual({ successful: 2, failed: 1 }); + }); + + it("returns undefined for a non-batch row and for null counts, so no rollup renders", () => { + expect(getBatchRequestCounts({ status: "success" })).toBeUndefined(); + expect(getBatchRequestCounts({ batch_successful_requests: null, batch_failed_requests: null })).toBeUndefined(); + expect(getBatchRequestCounts(undefined)).toBeUndefined(); + }); + + it("treats a lone present count as the other being 0, for rows logged mid-rollout", () => { + expect(getBatchRequestCounts({ batch_successful_requests: 3 })).toEqual({ successful: 3, failed: 0 }); + }); +}); + +describe("getBatchIdFromRequestId", () => { + it("strips the poller's synthetic _batch_cost suffix down to the provider batch id", () => { + expect(getBatchIdFromRequestId("batch_abc123_batch_cost")).toBe("batch_abc123"); + }); + + it("returns undefined for ordinary request ids and a bare suffix", () => { + expect(getBatchIdFromRequestId("chatcmpl-123")).toBeUndefined(); + expect(getBatchIdFromRequestId("_batch_cost")).toBeUndefined(); + }); +}); + +describe("getBatchModels", () => { + it("returns the model list from metadata.batch_models", () => { + expect(getBatchModels(batchCostMetadata)).toEqual(["gemini-2.5-flash"]); + }); + + it("returns undefined when absent, null, or empty", () => { + expect(getBatchModels({})).toBeUndefined(); + expect(getBatchModels({ batch_models: null })).toBeUndefined(); + expect(getBatchModels({ batch_models: [] })).toBeUndefined(); + }); +}); + +describe("getReasoningTokens", () => { + it("reads reasoning tokens from usage_object on a batch cost row", () => { + expect(getReasoningTokens(batchCostMetadata)).toBe(224); + }); + + it("prefers additional_usage_values, which per-request rows carry", () => { + const metadata = { + additional_usage_values: { completion_tokens_details: { reasoning_tokens: 40 } }, + usage_object: { completion_tokens_details: { reasoning_tokens: 999 } }, + }; + expect(getReasoningTokens(metadata)).toBe(40); + }); + + it("returns undefined when the breakout is null or missing", () => { + expect(getReasoningTokens({ usage_object: { completion_tokens_details: null } })).toBeUndefined(); + expect(getReasoningTokens({})).toBeUndefined(); + expect(getReasoningTokens(undefined)).toBeUndefined(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/view_logs/batchLogUtils.ts b/ui/litellm-dashboard/src/components/view_logs/batchLogUtils.ts new file mode 100644 index 00000000000..ce7793065d2 --- /dev/null +++ b/ui/litellm-dashboard/src/components/view_logs/batchLogUtils.ts @@ -0,0 +1,67 @@ +/** + * Helpers for reading batch-specific fields off a spend log row. + * + * The proxy's batch cost poller (CheckBatchCost) writes one spend log per completed batch + * with request_id "_batch_cost" and call_type "aretrieve_batch", carrying + * batch_models / batch_successful_requests / batch_failed_requests in metadata + * (see litellm/proxy/spend_tracking/spend_tracking_utils.py). + */ + +import { BATCH_CALL_TYPES } from "./constants"; + +export const BATCH_COST_REQUEST_ID_SUFFIX = "_batch_cost"; + +export interface BatchRequestCounts { + successful: number; + failed: number; +} + +export const isBatchCallType = (callType: string): boolean => BATCH_CALL_TYPES.includes(callType); + +const readMetaNumber = (metadata: Record | undefined, key: string): number | undefined => { + const value = metadata?.[key]; + return typeof value === "number" && Number.isFinite(value) ? value : undefined; +}; + +/** + * Per-request outcome counts of a batch cost row. Undefined when the row carries neither + * count (a non-batch row, or a batch logged before counts were tracked). + */ +export const getBatchRequestCounts = ( + metadata: Record | undefined, +): BatchRequestCounts | undefined => { + const successful = readMetaNumber(metadata, "batch_successful_requests"); + const failed = readMetaNumber(metadata, "batch_failed_requests"); + if (successful === undefined && failed === undefined) return undefined; + return { successful: successful ?? 0, failed: failed ?? 0 }; +}; + +/** The provider batch id behind a poller-written "_batch_cost" spend row. */ +export const getBatchIdFromRequestId = (requestId: string): string | undefined => + requestId.endsWith(BATCH_COST_REQUEST_ID_SUFFIX) && requestId.length > BATCH_COST_REQUEST_ID_SUFFIX.length + ? requestId.slice(0, -BATCH_COST_REQUEST_ID_SUFFIX.length) + : undefined; + +/** The models the batch's requests actually ran on, from metadata.batch_models. */ +export const getBatchModels = (metadata: Record | undefined): string[] | undefined => { + const models = metadata?.["batch_models"]; + if (!Array.isArray(models)) return undefined; + const names = models.filter((model): model is string => typeof model === "string" && model !== ""); + return names.length > 0 ? names : undefined; +}; + +/** + * Reasoning tokens aggregated across the row's completion usage. Read from the same two + * metadata containers the drawer already uses for prompt-token details: per-request rows + * carry additional_usage_values, batch cost rows carry usage_object. + */ +export const getReasoningTokens = (metadata: Record | undefined): number | undefined => { + const readDetails = (container: unknown): number | undefined => { + if (typeof container !== "object" || container === null) return undefined; + const details = (container as Record)["completion_tokens_details"]; + if (typeof details !== "object" || details === null) return undefined; + const reasoning = (details as Record)["reasoning_tokens"]; + return typeof reasoning === "number" && Number.isFinite(reasoning) ? reasoning : undefined; + }; + return readDetails(metadata?.["additional_usage_values"]) ?? readDetails(metadata?.["usage_object"]); +}; diff --git a/ui/litellm-dashboard/src/components/view_logs/constants.ts b/ui/litellm-dashboard/src/components/view_logs/constants.ts index 1c17f398e35..51889ae7d04 100644 --- a/ui/litellm-dashboard/src/components/view_logs/constants.ts +++ b/ui/litellm-dashboard/src/components/view_logs/constants.ts @@ -21,6 +21,9 @@ export const MCP_CALL_TYPES = ["call_mcp_tool", "list_mcp_tools"]; /** Call types that represent agent/A2A requests (e.g. asend_message). */ export const AGENT_CALL_TYPES = ["asend_message"]; +/** Call types that represent Batch API operations (creation and retrieval, sync and async). */ +export const BATCH_CALL_TYPES = ["acreate_batch", "create_batch", "aretrieve_batch", "retrieve_batch"]; + export const QUICK_SELECT_OPTIONS: { label: string; value: number; unit: string }[] = [ { label: "Last Minute", value: 1, unit: "minutes" }, { label: "Last 15 Minutes", value: 15, unit: "minutes" }, From c276813cb46f57b1934d5e18bedfc53a09e74df0 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Thu, 3 Sep 2026 17:23:53 -0400 Subject: [PATCH 22/60] feat(batches): enrich batch cost rows with breakdown, identity, session, and org spend --- .../proxy/common_utils/check_batch_cost.py | 48 ++++++++++---- litellm/batches/batch_utils.py | 52 +++++++++------- litellm/litellm_core_utils/litellm_logging.py | 15 +++++ .../spend_tracking/spend_tracking_utils.py | 27 +++++++- .../test_batches_logging_unit_tests.py | 60 ++++++++++++++++-- .../proxy_unit_tests/test_check_batch_cost.py | 45 ++++++++++++++ .../test_litellm/batches/test_batch_utils.py | 34 +++++++--- .../test_spend_tracking_utils.py | 62 +++++++++++++++++++ .../view_logs/RequestLogsTableColumns.tsx | 4 +- 9 files changed, 295 insertions(+), 52 deletions(-) diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py index 354a6ed2fd0..be6e52681fb 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -112,35 +112,39 @@ class CheckBatchCost: verbose_proxy_logger.error(f"CheckBatchCost: could not look up user {user_id} for batch {batch_id}: {e}") return {} - async def _get_key_alias(self, batch_id: str, api_key: str | None) -> str | None: - """Resolve the creating virtual key's alias from its hashed token.""" + async def _get_key_attribution(self, batch_id: str, api_key: str | None) -> tuple[str | None, str | None]: + """Resolve the creating virtual key's (alias, org_id) from its hashed token.""" if not api_key: - return None + return None, None try: key_row: prisma_models.LiteLLM_VerificationToken | None = ( await self.prisma_client.db.litellm_verificationtoken.find_unique( where={"token": api_key} ) ) - return getattr(key_row, "key_alias", None) if key_row is not None else None + if key_row is None: + return None, None + return getattr(key_row, "key_alias", None), getattr(key_row, "org_id", None) except Exception as e: verbose_proxy_logger.error(f"CheckBatchCost: could not look up key alias for batch {batch_id}: {e}") - return None + return None, None - async def _get_team_alias(self, team_id: str | None) -> str | None: - """Resolve a team's alias from its id.""" + async def _get_team_attribution(self, team_id: str | None) -> tuple[str | None, str | None]: + """Resolve a team's (alias, organization_id) from its id.""" if not team_id: - return None + return None, None try: team_row: prisma_models.LiteLLM_TeamTable | None = ( await self.prisma_client.db.litellm_teamtable.find_unique( where={"team_id": team_id} ) ) - return getattr(team_row, "team_alias", None) if team_row is not None else None + if team_row is None: + return None, None + return getattr(team_row, "team_alias", None), getattr(team_row, "organization_id", None) except Exception as e: verbose_proxy_logger.error(f"CheckBatchCost: could not look up team alias for team {team_id}: {e}") - return None + return None, None async def _build_creator_attribution_metadata( self, job: "LiteLLM_ManagedObjectTable", batch_id: str @@ -153,6 +157,10 @@ class CheckBatchCost: user_api_key_alias; when it has no alias, or the key has since been rotated or deleted, the field keeps the creating user's alias that _get_user_info filled in, because a resolvable name is more useful on the spend row than a null. + + user_api_key_org_id must be resolved here too: the spend update writer reads it + off this metadata to increment organization spend, so leaving it out silently + drops batch cost from org accounting for keys and teams that belong to one. """ api_key = getattr(job, "api_key", None) team_id = getattr(job, "team_id", None) @@ -165,12 +173,15 @@ class CheckBatchCost: **(await self._get_user_info(batch_id, job.created_by)), } - key_alias = await self._get_key_alias(batch_id, api_key) + key_alias, key_org_id = await self._get_key_attribution(batch_id, api_key) if key_alias is not None: metadata["user_api_key_alias"] = key_alias - team_alias = await self._get_team_alias(team_id) + team_alias, team_org_id = await self._get_team_attribution(team_id) if team_alias is not None: metadata["user_api_key_team_alias"] = team_alias + org_id: Final = key_org_id or team_org_id + if org_id is not None: + metadata["user_api_key_org_id"] = org_id if isinstance(request_tags, list) and request_tags: metadata["tags"] = [tag for tag in request_tags if isinstance(tag, str)] @@ -804,6 +815,7 @@ class CheckBatchCost: function_id=str(uuid.uuid4()), ) + deployment_api_base: Final = deployment_info.litellm_params.api_base logging_obj.update_environment_variables( litellm_params={ # set the user-agent header so that S3 callback consumers can easily identify CheckBatchCost callbacks @@ -812,9 +824,17 @@ class CheckBatchCost: "user-agent": CHECK_BATCH_COST_USER_AGENT, } }, - "metadata": await self._build_creator_attribution_metadata(job, batch_id), + **({"api_base": deployment_api_base} if deployment_api_base else {}), + "metadata": { + **(await self._build_creator_attribution_metadata(job, batch_id)), + # spend logs read the deployment identity off these metadata keys, so + # without them the batch cost row carries no model_id or model_group + "model_info": {"id": model_id}, + "model_group": deployment_info.model_name, + }, }, optional_params={}, + custom_llm_provider=str(llm_provider) if llm_provider else None, ) if not await self._claim_job_for_costing(job): @@ -832,6 +852,8 @@ class CheckBatchCost: batch_models=batch_result.models, batch_successful_requests=batch_result.successful_requests, batch_failed_requests=batch_result.failed_requests, + batch_prompt_cost=batch_result.prompt_cost, + batch_completion_cost=batch_result.completion_cost, ) except Exception: await self._release_job_claim(job) diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 3831f57a10d..a688e9f69fb 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -10,7 +10,7 @@ from litellm._logging import verbose_logger from litellm.litellm_core_utils.get_litellm_params import AWS_CREDENTIAL_KWARGS_KEYS from litellm.litellm_core_utils.llm_cost_calc.utils import parse_prompt_tokens_details from litellm.types.llms.openai import Batch -from litellm.types.utils import CallTypes, ModelInfo, Usage +from litellm.types.utils import ModelInfo, Usage from litellm.utils import token_counter @@ -23,6 +23,8 @@ class BatchCostUsageResult: models: list[str] successful_requests: int failed_requests: int + prompt_cost: float = 0.0 + completion_cost: float = 0.0 async def calculate_batch_cost_and_usage( @@ -130,7 +132,8 @@ class _LineOutcome(Enum): @dataclass(frozen=True, slots=True) class _BatchOutputLineStats: - cost: float + prompt_cost: float + completion_cost: float prompt_tokens: int completion_tokens: int total_tokens: int @@ -193,15 +196,16 @@ def _compute_output_line_stats( raw_model: Final = response_body.get("model") response_model: Final = raw_model if isinstance(raw_model, str) and raw_model else None completion_details: Final = usage.completion_tokens_details + line_prompt_cost, line_completion_cost = _output_line_cost( + usage=usage, + custom_llm_provider=custom_llm_provider, + model_name=model_name, + response_model=response_model, + model_info=model_info, + ) return _BatchOutputLineStats( - cost=_output_line_cost( - response_body=response_body, - usage=usage, - custom_llm_provider=custom_llm_provider, - model_name=model_name, - response_model=response_model, - model_info=model_info, - ), + prompt_cost=line_prompt_cost, + completion_cost=line_completion_cost, prompt_tokens=usage.prompt_tokens, completion_tokens=usage.completion_tokens, total_tokens=usage.total_tokens, @@ -213,31 +217,24 @@ def _compute_output_line_stats( def _output_line_cost( - response_body: Mapping[str, Any], usage: Usage, custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"], model_name: str | None, response_model: str | None, model_info: ModelInfo | None, -) -> float: +) -> tuple[float, float]: + """(prompt_cost, completion_cost) for one output line, priced at batch rates.""" from litellm.cost_calculator import batch_cost_calculator - if model_info is None and custom_llm_provider not in ("anthropic", "bedrock"): - return litellm.completion_cost( - completion_response=response_body, - custom_llm_provider=custom_llm_provider, - call_type=CallTypes.aretrieve_batch.value, - ) cost_model: Final = ( model_name if custom_llm_provider == "bedrock" and model_name else response_model or model_name or "" ) - prompt_cost, completion_cost = batch_cost_calculator( + return batch_cost_calculator( usage=usage, model=cost_model, custom_llm_provider=custom_llm_provider, model_info=model_info, ) - return prompt_cost + completion_cost def _aggregate_batch_cost_usage_models( @@ -270,7 +267,9 @@ def _aggregate_batch_cost_usage_models( **cache_token_params, ) batch_models: Final = [model_name] if model_name else [stats.model for stats in line_stats if stats.model] - total_cost: Final = sum((stats.cost for stats in line_stats), 0.0) + total_prompt_cost: Final = sum((stats.prompt_cost for stats in line_stats), 0.0) + total_completion_cost: Final = sum((stats.completion_cost for stats in line_stats), 0.0) + total_cost: Final = total_prompt_cost + total_completion_cost verbose_logger.debug( "batch output aggregate: cost=%s usage=%s models=%s successful=%d failed=%d", total_cost, @@ -285,6 +284,8 @@ def _aggregate_batch_cost_usage_models( models=batch_models, successful_requests=successful_requests, failed_requests=failed_requests, + prompt_cost=total_prompt_cost, + completion_cost=total_completion_cost, ) @@ -309,7 +310,8 @@ def calculate_vertex_ai_batch_cost_and_usage( """ from litellm.cost_calculator import batch_cost_calculator - total_cost = 0.0 + total_prompt_cost = 0.0 # rebind-ok: loop accumulator, matches total_tokens below + total_completion_cost = 0.0 # rebind-ok: loop accumulator, matches total_tokens below total_tokens = 0 prompt_tokens = 0 completion_tokens = 0 @@ -341,7 +343,8 @@ def calculate_vertex_ai_batch_cost_and_usage( model=actual_model_name, custom_llm_provider="vertex_ai", ) - total_cost += p_cost + c_cost + total_prompt_cost += p_cost + total_completion_cost += c_cost except Exception as e: verbose_logger.debug("vertex_ai batch cost calculation error for line: %s", str(e)) @@ -349,6 +352,7 @@ def calculate_vertex_ai_batch_cost_and_usage( completion_tokens += _completion total_tokens += _total + total_cost: Final = total_prompt_cost + total_completion_cost verbose_logger.info( "vertex_ai batch cost: cost=%s, prompt=%d, completion=%d, total=%d, successful=%d, failed=%d", total_cost, @@ -369,6 +373,8 @@ def calculate_vertex_ai_batch_cost_and_usage( models=[actual_model_name], successful_requests=successful_requests, failed_requests=failed_requests, + prompt_cost=total_prompt_cost, + completion_cost=total_completion_cost, ) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index f54eeca5178..3a616e7961d 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -2904,6 +2904,15 @@ class Logging(LiteLLMLoggingBaseClass): result._hidden_params["batch_successful_requests"] = batch_successful_requests # pyright: ignore[reportPrivateUsage] # rebind-ok: same result._hidden_params pattern as response_cost/batch_models above result._hidden_params["batch_failed_requests"] = batch_failed_requests # pyright: ignore[reportPrivateUsage] # rebind-ok: same pattern as above result.usage = batch_usage + batch_prompt_cost: Final = kwargs.get("batch_prompt_cost", None) + batch_completion_cost: Final = kwargs.get("batch_completion_cost", None) + if batch_prompt_cost is not None and batch_completion_cost is not None: + self.set_cost_breakdown( + input_cost=batch_prompt_cost, + output_cost=batch_completion_cost, + total_cost=batch_cost, + cost_for_built_in_tools_cost_usd_dollar=0.0, + ) elif should_compute_batch_data: batch_result: Final = await _handle_completed_batch( @@ -2919,6 +2928,12 @@ class Logging(LiteLLMLoggingBaseClass): result._hidden_params["batch_successful_requests"] = batch_result.successful_requests # pyright: ignore[reportPrivateUsage] # rebind-ok: same pattern as above result._hidden_params["batch_failed_requests"] = batch_result.failed_requests # pyright: ignore[reportPrivateUsage] # rebind-ok: same pattern as above result.usage = batch_result.usage + self.set_cost_breakdown( + input_cost=batch_result.prompt_cost, + output_cost=batch_result.completion_cost, + total_cost=batch_result.cost, + cost_for_built_in_tools_cost_usd_dollar=0.0, + ) start_time, end_time, result = self._success_handler_helper_fn( start_time=start_time, diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index a37c3ba4405..10db82d1b64 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -582,6 +582,7 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs metadata=metadata, standard_logging_payload=standard_logging_payload, omit_when_missing=_omits_session_id_when_missing(metadata), + batch_trace_session_id=_get_batch_trace_session_id(call_type=call_type, request_id=id), ), request_duration_ms=_get_request_duration_ms(start_time, end_time), status=_get_status_for_spend_log( @@ -620,20 +621,44 @@ def _omits_session_id_when_missing(metadata: Mapping[str, object] | None) -> boo return general_settings.get("missing_session_id") == "omit" +_BATCH_TRACE_CALL_TYPES: Final = frozenset( + { + CallTypes.create_batch.value, + CallTypes.acreate_batch.value, + CallTypes.retrieve_batch.value, + CallTypes.aretrieve_batch.value, + } +) + + +def _get_batch_trace_session_id(call_type: str | None, request_id: str | None) -> str | None: + """A batch's create row and its poller-written cost row both derive their request id + from the same batch id (the cost row appends BATCH_COST_REQUEST_ID_SUFFIX), so using + that id as the session groups the batch lifecycle into one trace on the logs UI. The + poller builds its own logging context, so per-request trace ids can never link them.""" + if call_type not in _BATCH_TRACE_CALL_TYPES or not request_id: + return None + return request_id.removesuffix(BATCH_COST_REQUEST_ID_SUFFIX) + + def _get_session_id_for_spend_log( kwargs: Mapping[str, object], metadata: Mapping[str, object] | None, standard_logging_payload: StandardLoggingPayload | None, omit_when_missing: bool, + batch_trace_session_id: str | None = None, ) -> str | None: """Under `omit` only `metadata.session_id`, the key Langfuse reads, counts as a session; `litellm_session_id` may - be a copied trace id.""" + be a copied trace id. Batch call types carry a deterministic session derived from the batch id, which outranks + the per-request trace ids because those differ between the create call and the cost poller's row.""" if omit_when_missing: session_id: Final = metadata.get("session_id") if metadata else None return str(session_id) if session_id else None from litellm._uuid import uuid + if batch_trace_session_id is not None: + return batch_trace_session_id if standard_logging_payload is not None and standard_logging_payload.get("trace_id") is not None: return str(standard_logging_payload.get("trace_id")) if kwargs.get("litellm_trace_id") is not None: diff --git a/tests/batches_tests/test_batches_logging_unit_tests.py b/tests/batches_tests/test_batches_logging_unit_tests.py index 5bde40d90b0..73adb391481 100644 --- a/tests/batches_tests/test_batches_logging_unit_tests.py +++ b/tests/batches_tests/test_batches_logging_unit_tests.py @@ -144,18 +144,20 @@ def test_get_batch_job_total_usage_from_file_content(sample_file_content_dict): @pytest.mark.asyncio async def test_batch_cost_calculator(sample_file_content_dict): """ - mock litellm.completion_cost to return 0.5 + mock batch_cost_calculator to return (0.3, 0.2) per line we know sample_file_content_dict has 2 successful responses - so we expect the cost to be 0.5 * 2 = 1.0 + so we expect the cost to be (0.3 + 0.2) * 2 = 1.0, split 0.6 / 0.4 """ - with patch("litellm.completion_cost", return_value=0.5): + with patch("litellm.cost_calculator.batch_cost_calculator", return_value=(0.3, 0.2)): result = _aggregate_batch_cost_usage_models( entries=sample_file_content_dict, custom_llm_provider="openai", ) - assert result.cost == 1.0 # 0.5 * 2 successful responses + assert result.cost == pytest.approx(1.0) # (0.3 + 0.2) * 2 successful responses + assert result.prompt_cost == pytest.approx(0.6) + assert result.completion_cost == pytest.approx(0.4) def test_get_response_from_batch_job_output_file(sample_file_content_dict): @@ -402,6 +404,56 @@ async def test_batch_retrieve_cost_tracking_with_explicit_cost_data(): assert mock_batch.usage == explicit_usage +@pytest.mark.asyncio +async def test_batch_retrieve_explicit_cost_split_sets_cost_breakdown(): + """The poller passes the batch's prompt/completion cost split so the spend row's + cost_breakdown carries real input/output costs; without it the UI's Cost Breakdown + card renders blank for every batch. Regression for the split being dropped.""" + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.types.utils import CallTypes, LiteLLMBatch + + mock_batch = LiteLLMBatch( + id="batch-breakdown-1", + object="batch", + endpoint="/v1/chat/completions", + errors=None, + input_file_id="file-input-1", + completion_window="24h", + status="completed", + output_file_id="file-output-1", + created_at=1234567890, + ) + mock_batch._hidden_params = {} + + logging_obj = Logging( + model="gpt-5-mini", + messages=[{"role": "user", "content": "test"}], + stream=False, + call_type=CallTypes.aretrieve_batch.value, + litellm_call_id="test-call-breakdown", + function_id="test-function", + start_time=time.time(), + dynamic_success_callbacks=[], + ) + logging_obj.custom_llm_provider = "openai" + + await logging_obj.async_success_handler( + result=mock_batch, + start_time=time.time(), + end_time=time.time() + 1, + batch_cost=0.10, + batch_usage=litellm.Usage(prompt_tokens=200, completion_tokens=100, total_tokens=300), + batch_models=["gpt-5-mini"], + batch_prompt_cost=0.06, + batch_completion_cost=0.04, + ) + + assert logging_obj.cost_breakdown is not None + assert logging_obj.cost_breakdown["input_cost"] == 0.06 + assert logging_obj.cost_breakdown["output_cost"] == 0.04 + assert logging_obj.cost_breakdown["total_cost"] == 0.10 + + @pytest.mark.asyncio async def test_batch_retrieve_cost_tracking_with_unified_file_id_incomplete_batch(): """ diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index ff5e8f89d64..9ee0a36f00b 100644 --- a/tests/proxy_unit_tests/test_check_batch_cost.py +++ b/tests/proxy_unit_tests/test_check_batch_cost.py @@ -2553,6 +2553,51 @@ class TestBatchCostAttribution: assert metadata["user_api_key_alias"] == "prod-key" + @pytest.mark.asyncio + async def test_org_id_comes_from_the_creating_key(self): + """The spend update writer increments organization spend from user_api_key_org_id, + so an org-scoped key's batch cost must carry the key's org id.""" + from types import SimpleNamespace + + instance = self._instance( + key_row=SimpleNamespace(key_alias="prod-key", org_id="org-42"), + team_row=SimpleNamespace(team_alias="Team Alpha", organization_id="org-team"), + ) + + metadata = await instance._build_creator_attribution_metadata(self._job(), "batch-1") + + assert metadata["user_api_key_org_id"] == "org-42" + + @pytest.mark.asyncio + async def test_org_id_falls_back_to_the_team_organization(self): + """A key with no org of its own still books batch spend against its team's + organization, matching how the request path resolves org attribution.""" + from types import SimpleNamespace + + instance = self._instance( + key_row=SimpleNamespace(key_alias="prod-key", org_id=None), + team_row=SimpleNamespace(team_alias="Team Alpha", organization_id="org-team"), + ) + + metadata = await instance._build_creator_attribution_metadata(self._job(), "batch-1") + + assert metadata["user_api_key_org_id"] == "org-team" + + @pytest.mark.asyncio + async def test_no_org_leaves_the_key_unset(self): + """Without any org the key is absent entirely, so the spend writer's org update + stays skipped instead of matching an empty-string organization.""" + from types import SimpleNamespace + + instance = self._instance( + key_row=SimpleNamespace(key_alias="prod-key", org_id=None), + team_row=SimpleNamespace(team_alias="Team Alpha", organization_id=None), + ) + + metadata = await instance._build_creator_attribution_metadata(self._job(), "batch-1") + + assert "user_api_key_org_id" not in metadata + class TestPollPageStarvation: """LIT-5462 regression: a row that can never be costed used to keep its slot in the diff --git a/tests/test_litellm/batches/test_batch_utils.py b/tests/test_litellm/batches/test_batch_utils.py index c86c7c4df03..1fd78870481 100644 --- a/tests/test_litellm/batches/test_batch_utils.py +++ b/tests/test_litellm/batches/test_batch_utils.py @@ -489,7 +489,9 @@ def test_aggregate_counts_successful_and_failed_requests(monkeypatch): def test_aggregate_returns_batch_cost_usage_result_dataclass(monkeypatch): - monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 1.0) + import litellm.cost_calculator as cc + + monkeypatch.setattr(cc, "batch_cost_calculator", lambda **kw: (0.4, 0.6)) result = bu._aggregate_batch_cost_usage_models( entries=[_success_row(usage=_usage(10, 5))], custom_llm_provider="openai" ) @@ -500,6 +502,7 @@ def test_aggregate_returns_batch_cost_usage_result_dataclass(monkeypatch): 1, 0, ) + assert (result.prompt_cost, result.completion_cost) == (0.4, 0.6) # =========================================================================== # @@ -507,15 +510,17 @@ def test_aggregate_returns_batch_cost_usage_result_dataclass(monkeypatch): # =========================================================================== # -def test_cost_from_content_completion_cost_path(monkeypatch): - # model_info is None -> litellm.completion_cost per successful row. +def test_cost_without_model_info_prices_each_row_by_its_response_model(monkeypatch): + # model_info is None -> batch_cost_calculator per successful row, model from the response body. + import litellm.cost_calculator as cc + calls = [] - def _completion_cost(**kw): + def _batch_cost(**kw): calls.append(kw) - return 0.5 + return (0.3, 0.2) - monkeypatch.setattr(litellm, "completion_cost", _completion_cost) + monkeypatch.setattr(cc, "batch_cost_calculator", _batch_cost) rows = [ _success_row(usage=_usage(10, 5)), _failed_row(), # excluded -> not costed @@ -524,8 +529,10 @@ def test_cost_from_content_completion_cost_path(monkeypatch): result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai") - assert result.cost == 1.0 # 2 successful * 0.5 + assert result.cost == pytest.approx(1.0) # 2 successful * (0.3 + 0.2) + assert (result.prompt_cost, result.completion_cost) == (pytest.approx(0.6), pytest.approx(0.4)) assert len(calls) == 2 # failed row not costed + assert all(call["model"] == "gpt-4o" and call["model_info"] is None for call in calls) assert result.successful_requests == 2 assert result.failed_requests == 1 @@ -578,7 +585,9 @@ def test_aggregate_consumes_entries_in_a_single_pass(monkeypatch): """A one-shot generator: any implementation that iterates the entries twice (e.g. separate cost and usage passes) sees nothing on the second pass and returns wrong totals for at least one of cost/usage/models.""" - monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 0.5) + import litellm.cost_calculator as cc + + monkeypatch.setattr(cc, "batch_cost_calculator", lambda **kw: (0.25, 0.25)) one_shot = (row for row in [_success_row(usage=_usage(10, 5)), _failed_row(), _success_row(usage=_usage(20, 10))]) result = bu._aggregate_batch_cost_usage_models(entries=one_shot, custom_llm_provider="openai") @@ -753,12 +762,15 @@ def test_vertex_cost_error_in_line_is_swallowed(monkeypatch): @pytest.mark.asyncio async def test_calculate_batch_cost_and_usage_orchestration(monkeypatch): + import litellm.cost_calculator as cc + rows = [_success_row(model="gpt-4o", usage=_usage(10, 5))] - monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 2.5) + monkeypatch.setattr(cc, "batch_cost_calculator", lambda **kw: (1.5, 1.0)) result = await bu.calculate_batch_cost_and_usage(file_content_dictionary=rows, custom_llm_provider="openai") assert result.cost == 2.5 + assert (result.prompt_cost, result.completion_cost) == (1.5, 1.0) assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (10, 5, 15) assert result.models == ["gpt-4o"] @@ -1107,8 +1119,10 @@ async def test_handle_completed_batch_orchestration(monkeypatch): async def fake_fetch(batch, custom_llm_provider, litellm_params=None): return _vertex_jsonl(rows) + import litellm.cost_calculator as cc + monkeypatch.setattr(bu, "_fetch_batch_output_file_content", fake_fetch) - monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 3.3) + monkeypatch.setattr(cc, "batch_cost_calculator", lambda **kw: (2.0, 1.3)) result = await bu._handle_completed_batch(_batch("of"), custom_llm_provider="openai") diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py index 323930eee60..be274befd25 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py @@ -128,6 +128,68 @@ def test_legacy_policy_keeps_trace_id_fallback(): assert len(str(generated)) == 36 +def test_batch_lifecycle_rows_derive_the_same_session_from_the_batch_id(): + """The create call's request id IS the batch id and the poller's cost row appends + _batch_cost to it, so deriving the session from the request id lands both rows in one + trace on the logs UI even though the poller builds a fresh logging context per cycle.""" + from litellm.proxy.spend_tracking.spend_tracking_utils import _get_batch_trace_session_id + + create_session: Final = _get_batch_trace_session_id(call_type="acreate_batch", request_id="batch-uid-1") + cost_session: Final = _get_batch_trace_session_id( + call_type="aretrieve_batch", request_id="batch-uid-1_batch_cost" + ) + assert create_session == cost_session == "batch-uid-1" + + +def test_non_batch_call_types_derive_no_batch_session(): + from litellm.proxy.spend_tracking.spend_tracking_utils import _get_batch_trace_session_id + + assert _get_batch_trace_session_id(call_type="acompletion", request_id="chatcmpl-1") is None + + +def test_batch_session_outranks_the_per_request_trace_id(): + """Each batch lifecycle call carries its own auto-generated trace id, so letting the + trace id win would scatter the rows across sessions again.""" + session_id: Final = _get_session_id_for_spend_log( + kwargs={"litellm_trace_id": "trace-abc"}, + metadata={"trace_id": "trace-abc"}, + standard_logging_payload=_TRACE_ONLY_STANDARD_LOGGING, + omit_when_missing=False, + batch_trace_session_id="batch-uid-1", + ) + assert session_id == "batch-uid-1" + + +def test_omit_policy_still_suppresses_batch_sessions(): + session_id: Final = _get_session_id_for_spend_log( + kwargs={}, + metadata=None, + standard_logging_payload=None, + omit_when_missing=True, + batch_trace_session_id="batch-uid-1", + ) + assert session_id is None + + +def test_get_logging_payload_groups_batch_create_and_cost_rows_in_one_session(): + def _payload(call_type: str) -> SpendLogsPayload: + return get_logging_payload( + kwargs={ + "call_type": call_type, + "model": "gpt-4o-mini", + "litellm_params": {"metadata": {"user_api_key": "test-key"}}, + }, + response_obj=litellm.ModelResponse(id="batch-uid-1", choices=[], usage=litellm.Usage()), + start_time=datetime.datetime.now(timezone.utc), + end_time=datetime.datetime.now(timezone.utc), + ) + + create_payload: Final = _payload("acreate_batch") + cost_payload: Final = _payload("aretrieve_batch") + assert cost_payload["request_id"] == "batch-uid-1_batch_cost" + assert create_payload["session_id"] == cost_payload["session_id"] == "batch-uid-1" + + @pytest.mark.parametrize( ("request_metadata", "expected"), [ diff --git a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.tsx index b4296abe266..a444acb9517 100644 --- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.tsx @@ -64,10 +64,12 @@ export const getRequestLogsTableColumns = ({ const sessionAgentCount = log.session_agent_count ?? (isAgent ? sessionCount : 0); const sessionMcpCount = log.mcp_tool_call_count ?? (isMcp ? sessionCount : 0); + if (isBatchCallType(log.call_type)) { + return 1 ? sessionCount : undefined} />; + } if (sessionCount <= 1) { if (isMcp) return ; if (isAgent) return ; - if (isBatchCallType(log.call_type)) return ; return ; } From 06c860b60a5c4a2d14a02dead9c809b8b7eaf22f Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Thu, 3 Sep 2026 18:09:18 -0400 Subject: [PATCH 23/60] fix(ui): type batch results metadata as unknown instead of any --- .../components/view_logs/LogDetailsDrawer/LogDetailContent.tsx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.tsx index 8e0f041b88d..710d82b2f9d 100644 --- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.tsx @@ -388,7 +388,7 @@ function MetricLabel({ label, tooltip, docsUrl }: { label: string; tooltip: stri * Aggregate per-request outcomes for a batch cost row: batch id, success/failure counts * from the parsed output and error files, and the models the batch actually ran on. */ -function BatchResultsSection({ logEntry, metadata }: { logEntry: LogEntry; metadata: Record }) { +function BatchResultsSection({ logEntry, metadata }: { logEntry: LogEntry; metadata: Record }) { const counts = getBatchRequestCounts(metadata); const batchId = getBatchIdFromRequestId(logEntry.request_id); const batchModels = getBatchModels(metadata); From 538cd2c3b00d31fcb516c8326de87883e0e27f94 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Thu, 3 Sep 2026 18:16:31 -0400 Subject: [PATCH 24/60] fix(batches): narrow batch cost kwargs before the breakdown and drop node access in test --- litellm/litellm_core_utils/litellm_logging.py | 6 +++++- .../LogDetailsDrawer/LogDetailContent.test.tsx | 12 +++++++----- 2 files changed, 12 insertions(+), 6 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 3a616e7961d..03486f4f729 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -2906,7 +2906,11 @@ class Logging(LiteLLMLoggingBaseClass): result.usage = batch_usage batch_prompt_cost: Final = kwargs.get("batch_prompt_cost", None) batch_completion_cost: Final = kwargs.get("batch_completion_cost", None) - if batch_prompt_cost is not None and batch_completion_cost is not None: + if ( + isinstance(batch_prompt_cost, float) + and isinstance(batch_completion_cost, float) + and isinstance(batch_cost, float) + ): self.set_cost_breakdown( input_cost=batch_prompt_cost, output_cost=batch_completion_cost, diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.test.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.test.tsx index 91778d2a98a..a893e9bffd0 100644 --- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.test.tsx @@ -163,11 +163,13 @@ describe("LogDetailContent", () => { />, ); - const section = screen.getByText("Batch Results").closest('[data-slot="card"]') as HTMLElement; - expect(within(section).getByText("batch_abc123")).toBeInTheDocument(); - expect(within(section).getByText("2")).toBeInTheDocument(); - expect(within(section).getByText("1")).toBeInTheDocument(); - expect(within(section).getByText("gemini-2.5-flash")).toBeInTheDocument(); + expect(screen.getByText("Batch Results")).toBeInTheDocument(); + expect(screen.getByText("batch_abc123")).toBeInTheDocument(); + expect(screen.getByText("Successful Requests")).toBeInTheDocument(); + expect(screen.getByText("2")).toBeInTheDocument(); + expect(screen.getByText("Failed Requests")).toBeInTheDocument(); + expect(screen.getByText("1")).toBeInTheDocument(); + expect(screen.getByText("gemini-2.5-flash")).toBeInTheDocument(); }); it("still renders the batch id when a legacy row carries no counts", () => { From 7fb4b427cec0db7d3bc85dd6c2055ca987f2390b Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Thu, 3 Sep 2026 18:36:05 -0400 Subject: [PATCH 25/60] feat(batches): snapshot the creating key's org on the managed object row --- .../proxy/common_utils/check_batch_cost.py | 61 +++++++++++++------ .../proxy/hooks/managed_files.py | 1 + .../migration.sql | 4 ++ .../litellm_proxy_extras/schema.prisma | 1 + litellm/models/managed_files.py | 1 + schema.prisma | 1 + .../proxy_unit_tests/test_check_batch_cost.py | 21 ++++++- ..._batch_update_db_managed_output_file_id.py | 4 +- 8 files changed, 74 insertions(+), 20 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260903230000_add_org_id_to_managed_object_table/migration.sql diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py index be6e52681fb..595a8d04bed 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -112,39 +112,66 @@ class CheckBatchCost: verbose_proxy_logger.error(f"CheckBatchCost: could not look up user {user_id} for batch {batch_id}: {e}") return {} - async def _get_key_attribution(self, batch_id: str, api_key: str | None) -> tuple[str | None, str | None]: - """Resolve the creating virtual key's (alias, org_id) from its hashed token.""" + async def _get_key_alias(self, batch_id: str, api_key: str | None) -> str | None: + """Resolve the creating virtual key's alias from its hashed token.""" if not api_key: - return None, None + return None try: key_row: prisma_models.LiteLLM_VerificationToken | None = ( await self.prisma_client.db.litellm_verificationtoken.find_unique( where={"token": api_key} ) ) - if key_row is None: - return None, None - return getattr(key_row, "key_alias", None), getattr(key_row, "org_id", None) + return getattr(key_row, "key_alias", None) if key_row is not None else None except Exception as e: verbose_proxy_logger.error(f"CheckBatchCost: could not look up key alias for batch {batch_id}: {e}") - return None, None + return None - async def _get_team_attribution(self, team_id: str | None) -> tuple[str | None, str | None]: - """Resolve a team's (alias, organization_id) from its id.""" + async def _get_team_alias(self, team_id: str | None) -> str | None: + """Resolve a team's alias from its id.""" if not team_id: - return None, None + return None try: team_row: prisma_models.LiteLLM_TeamTable | None = ( await self.prisma_client.db.litellm_teamtable.find_unique( where={"team_id": team_id} ) ) - if team_row is None: - return None, None - return getattr(team_row, "team_alias", None), getattr(team_row, "organization_id", None) + return getattr(team_row, "team_alias", None) if team_row is not None else None except Exception as e: verbose_proxy_logger.error(f"CheckBatchCost: could not look up team alias for team {team_id}: {e}") - return None, None + return None + + async def _get_org_id(self, job: "LiteLLM_ManagedObjectTable", batch_id: str) -> str | None: + """Organization to bill the batch against, snapshotted on the row at creation + like team_id. Rows created before the org_id column existed carry None, so they + fall back to the creating key's org (or its team's) as resolved today.""" + org_id = getattr(job, "org_id", None) + if org_id: + return org_id + api_key = getattr(job, "api_key", None) + team_id = getattr(job, "team_id", None) + try: + if api_key: + key_row: prisma_models.LiteLLM_VerificationToken | None = ( + await self.prisma_client.db.litellm_verificationtoken.find_unique( + where={"token": api_key} + ) + ) + key_org_id = getattr(key_row, "org_id", None) if key_row is not None else None + if key_org_id: + return key_org_id + if team_id: + team_row: prisma_models.LiteLLM_TeamTable | None = ( + await self.prisma_client.db.litellm_teamtable.find_unique( + where={"team_id": team_id} + ) + ) + return getattr(team_row, "organization_id", None) if team_row is not None else None + return None + except Exception as e: + verbose_proxy_logger.error(f"CheckBatchCost: could not resolve org for batch {batch_id}: {e}") + return None async def _build_creator_attribution_metadata( self, job: "LiteLLM_ManagedObjectTable", batch_id: str @@ -173,13 +200,13 @@ class CheckBatchCost: **(await self._get_user_info(batch_id, job.created_by)), } - key_alias, key_org_id = await self._get_key_attribution(batch_id, api_key) + key_alias = await self._get_key_alias(batch_id, api_key) if key_alias is not None: metadata["user_api_key_alias"] = key_alias - team_alias, team_org_id = await self._get_team_attribution(team_id) + team_alias = await self._get_team_alias(team_id) if team_alias is not None: metadata["user_api_key_team_alias"] = team_alias - org_id: Final = key_org_id or team_org_id + org_id: Final = await self._get_org_id(job, batch_id) if org_id is not None: metadata["user_api_key_org_id"] = org_id if isinstance(request_tags, list) and request_tags: diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 5cfcf6129f0..4f0ca787280 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -349,6 +349,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): "file_purpose": file_purpose, "created_by": resolve_resource_owner_id(user_api_key_dict), "team_id": user_api_key_dict.team_id, + "org_id": user_api_key_dict.org_id, "updated_by": user_api_key_dict.user_id, "status": file_object.status, **attribution_columns, diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260903230000_add_org_id_to_managed_object_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260903230000_add_org_id_to_managed_object_table/migration.sql new file mode 100644 index 00000000000..bbe980bb66f --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260903230000_add_org_id_to_managed_object_table/migration.sql @@ -0,0 +1,4 @@ +-- Add org_id column to LiteLLM_ManagedObjectTable +-- Snapshots the creating key's organization at submission time, like team_id, +-- so CheckBatchCost can bill organization spend hours later without re-resolving +ALTER TABLE "LiteLLM_ManagedObjectTable" ADD COLUMN IF NOT EXISTS "org_id" TEXT; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 7604ceadf7a..e386446d5bd 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -1034,6 +1034,7 @@ model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use t created_at DateTime @default(now()) created_by String? team_id String? + org_id String? // creating key's organization at submission time; CheckBatchCost bills org spend against it api_key String? request_tags Json? @default("[]") updated_at DateTime @updatedAt diff --git a/litellm/models/managed_files.py b/litellm/models/managed_files.py index 23d70ef5c48..c90f9b535ea 100644 --- a/litellm/models/managed_files.py +++ b/litellm/models/managed_files.py @@ -32,6 +32,7 @@ class LiteLLM_ManagedObjectTable(LiteLLMPydanticObjectBase): file_object: LiteLLMBatch | LiteLLMFineTuningJob | ResponsesAPIResponse created_by: str | None = None team_id: str | None = None + org_id: str | None = None class LiteLLM_ManagedVectorStoreTable(LiteLLMPydanticObjectBase): diff --git a/schema.prisma b/schema.prisma index 7604ceadf7a..e386446d5bd 100644 --- a/schema.prisma +++ b/schema.prisma @@ -1034,6 +1034,7 @@ model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use t created_at DateTime @default(now()) created_by String? team_id String? + org_id String? // creating key's organization at submission time; CheckBatchCost bills org spend against it api_key String? request_tags Json? @default("[]") updated_at DateTime @updatedAt diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index 9ee0a36f00b..92b946e19b1 100644 --- a/tests/proxy_unit_tests/test_check_batch_cost.py +++ b/tests/proxy_unit_tests/test_check_batch_cost.py @@ -2553,10 +2553,27 @@ class TestBatchCostAttribution: assert metadata["user_api_key_alias"] == "prod-key" + @pytest.mark.asyncio + async def test_org_id_snapshotted_on_the_row_wins(self): + """The org_id column captures the creating key's organization at submission time, + like team_id, so a key later moved to another org still bills the original one.""" + from types import SimpleNamespace + + instance = self._instance( + key_row=SimpleNamespace(key_alias="prod-key", org_id="org-moved-to"), + team_row=SimpleNamespace(team_alias="Team Alpha", organization_id="org-team"), + ) + + metadata = await instance._build_creator_attribution_metadata( + self._job(org_id="org-at-creation"), "batch-1" + ) + + assert metadata["user_api_key_org_id"] == "org-at-creation" + @pytest.mark.asyncio async def test_org_id_comes_from_the_creating_key(self): - """The spend update writer increments organization spend from user_api_key_org_id, - so an org-scoped key's batch cost must carry the key's org id.""" + """The spend update writer increments organization spend from user_api_key_org_id. + A legacy row without the org_id column falls back to the creating key's org.""" from types import SimpleNamespace instance = self._instance( diff --git a/tests/test_litellm/enterprise/proxy/test_batch_update_db_managed_output_file_id.py b/tests/test_litellm/enterprise/proxy/test_batch_update_db_managed_output_file_id.py index ebd33aa2e53..63884e0a779 100644 --- a/tests/test_litellm/enterprise/proxy/test_batch_update_db_managed_output_file_id.py +++ b/tests/test_litellm/enterprise/proxy/test_batch_update_db_managed_output_file_id.py @@ -390,7 +390,7 @@ async def test_store_unified_object_id_persists_key_and_tags_on_create(): """Regression (spend loss): the batch create persists the creating key hash and tags so CheckBatchCost can write an attributed spend row instead of a blank one the DB drops.""" instance, store = _in_memory_managed_files() - creator = UserAPIKeyAuth(user_id="alice", team_id="team-alpha", api_key="hash-alice") + creator = UserAPIKeyAuth(user_id="alice", team_id="team-alpha", api_key="hash-alice", org_id="org-acme") await instance.store_unified_object_id( unified_object_id="unified-b", @@ -407,6 +407,7 @@ async def test_store_unified_object_id_persists_key_and_tags_on_create(): assert row["api_key"] == "hash-alice" assert row["created_by"] == "alice" assert row["team_id"] == "team-alpha" + assert row["org_id"] == "org-acme" assert row["request_tags"].data == ["env:prod"] @@ -471,6 +472,7 @@ async def test_store_unified_object_id_attribution_columns_are_write_once(): upsert_data = instance.prisma_client.db.litellm_managedobjecttable.upsert.call_args.kwargs["data"] assert "api_key" not in upsert_data["update"] assert "request_tags" not in upsert_data["update"] + assert "org_id" not in upsert_data["update"] @pytest.mark.asyncio From 9859d1e64ca1e6930d2f3e440a51222e834c7a5f Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" <41898282+github-actions[bot]@users.noreply.github.com> Date: Thu, 3 Sep 2026 22:36:28 +0000 Subject: [PATCH 26/60] chore: sync schema.prisma copies from root --- litellm/proxy/schema.prisma | 1 + 1 file changed, 1 insertion(+) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 7604ceadf7a..e386446d5bd 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -1034,6 +1034,7 @@ model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use t created_at DateTime @default(now()) created_by String? team_id String? + org_id String? // creating key's organization at submission time; CheckBatchCost bills org spend against it api_key String? request_tags Json? @default("[]") updated_at DateTime @updatedAt From 0d7976116eff0a9090f846e0066067bd867e67fb Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Thu, 3 Sep 2026 18:51:18 -0400 Subject: [PATCH 27/60] fix(batches): resolve team-scoped keys' org at creation for the snapshot --- .../proxy/hooks/managed_files.py | 20 +++++++++++++- ..._batch_update_db_managed_output_file_id.py | 27 +++++++++++++++++++ 2 files changed, 46 insertions(+), 1 deletion(-) diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 4f0ca787280..2ed446ca44e 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -277,6 +277,24 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): ) verbose_logger.debug(f"LiteLLM Managed File object with id={file_id} stored in db: {result}") + async def _resolve_creator_org_id(self, user_api_key_dict: UserAPIKeyAuth) -> Optional[str]: + """Organization to snapshot on the managed object row, like team_id. A key that + belongs to an org only through its team carries no org_id on the auth object, so + resolve the team's organization at creation time; costing then bills the org the + batch was submitted under even if the key or team moves before it completes.""" + if user_api_key_dict.org_id: + return user_api_key_dict.org_id + if not user_api_key_dict.team_id: + return None + try: + team_row = await self.prisma_client.db.litellm_teamtable.find_unique( + where={"team_id": user_api_key_dict.team_id} + ) + return getattr(team_row, "organization_id", None) if team_row is not None else None + except Exception as e: + verbose_logger.warning(f"could not resolve org for managed object attribution: {e}") + return None + async def store_unified_object_id( self, unified_object_id: str, @@ -349,7 +367,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): "file_purpose": file_purpose, "created_by": resolve_resource_owner_id(user_api_key_dict), "team_id": user_api_key_dict.team_id, - "org_id": user_api_key_dict.org_id, + "org_id": await self._resolve_creator_org_id(user_api_key_dict), "updated_by": user_api_key_dict.user_id, "status": file_object.status, **attribution_columns, diff --git a/tests/test_litellm/enterprise/proxy/test_batch_update_db_managed_output_file_id.py b/tests/test_litellm/enterprise/proxy/test_batch_update_db_managed_output_file_id.py index 63884e0a779..b5865ab4a13 100644 --- a/tests/test_litellm/enterprise/proxy/test_batch_update_db_managed_output_file_id.py +++ b/tests/test_litellm/enterprise/proxy/test_batch_update_db_managed_output_file_id.py @@ -375,6 +375,7 @@ def _in_memory_managed_files(): table.upsert = AsyncMock(side_effect=_upsert) prisma = MagicMock() prisma.db.litellm_managedobjecttable = table + prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=None) cache = MagicMock() cache.async_set_cache = AsyncMock() @@ -411,6 +412,32 @@ async def test_store_unified_object_id_persists_key_and_tags_on_create(): assert row["request_tags"].data == ["env:prod"] +@pytest.mark.asyncio +async def test_store_unified_object_id_resolves_org_through_the_team(): + """Most keys belong to an org only through their team, so the auth object carries no + org_id. The create resolves the team's organization so org spend is snapshotted at + submission time instead of never being billed.""" + from types import SimpleNamespace + + instance, store = _in_memory_managed_files() + instance.prisma_client.db.litellm_teamtable.find_unique = AsyncMock( + return_value=SimpleNamespace(organization_id="org-via-team") + ) + creator = UserAPIKeyAuth(user_id="alice", team_id="team-alpha", api_key="hash-alice") + + await instance.store_unified_object_id( + unified_object_id="unified-b", + file_object=_build_batch_response(batch_id="b", status="validating"), + litellm_parent_otel_span=None, + model_object_id="b", + file_purpose="batch", + user_api_key_dict=creator, + persist_attribution=True, + ) + + assert store["unified-b"]["org_id"] == "org-via-team" + + @pytest.mark.asyncio async def test_store_unified_object_id_omits_key_and_tags_without_persist_attribution(): """Regression (spend redirect): a caller that is not the batch create (a poll, or the From 35371a34c1c873a0afafe9dffc1873fcee42928e Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Thu, 3 Sep 2026 19:09:45 -0400 Subject: [PATCH 28/60] fix(batches): keep team org attribution when the key lookup fails --- .../proxy/common_utils/check_batch_cost.py | 24 ++++++++++++------- .../proxy_unit_tests/test_check_batch_cost.py | 17 +++++++++++++ 2 files changed, 32 insertions(+), 9 deletions(-) diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py index 595a8d04bed..bd7f3c93ff3 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -151,8 +151,8 @@ class CheckBatchCost: return org_id api_key = getattr(job, "api_key", None) team_id = getattr(job, "team_id", None) - try: - if api_key: + if api_key: + try: key_row: prisma_models.LiteLLM_VerificationToken | None = ( await self.prisma_client.db.litellm_verificationtoken.find_unique( where={"token": api_key} @@ -161,16 +161,22 @@ class CheckBatchCost: key_org_id = getattr(key_row, "org_id", None) if key_row is not None else None if key_org_id: return key_org_id - if team_id: - team_row: prisma_models.LiteLLM_TeamTable | None = ( - await self.prisma_client.db.litellm_teamtable.find_unique( - where={"team_id": team_id} - ) + except Exception as e: + verbose_proxy_logger.error( + f"CheckBatchCost: could not resolve the key's org for batch {batch_id}, " + f"still trying the team's: {e}" ) - return getattr(team_row, "organization_id", None) if team_row is not None else None + if not team_id: return None + try: + team_row: prisma_models.LiteLLM_TeamTable | None = ( + await self.prisma_client.db.litellm_teamtable.find_unique( + where={"team_id": team_id} + ) + ) + return getattr(team_row, "organization_id", None) if team_row is not None else None except Exception as e: - verbose_proxy_logger.error(f"CheckBatchCost: could not resolve org for batch {batch_id}: {e}") + verbose_proxy_logger.error(f"CheckBatchCost: could not resolve the team's org for batch {batch_id}: {e}") return None async def _build_creator_attribution_metadata( diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index 92b946e19b1..59dfb67c486 100644 --- a/tests/proxy_unit_tests/test_check_batch_cost.py +++ b/tests/proxy_unit_tests/test_check_batch_cost.py @@ -2600,6 +2600,23 @@ class TestBatchCostAttribution: assert metadata["user_api_key_org_id"] == "org-team" + @pytest.mark.asyncio + async def test_key_lookup_failure_still_bills_the_team_org(self): + """A key-table error while resolving a legacy row's org must not drop the team's + organization: the two lookups fail independently, so org spend still lands.""" + from types import SimpleNamespace + + instance = self._instance( + team_row=SimpleNamespace(team_alias="Team Alpha", organization_id="org-team"), + ) + instance.prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( + side_effect=Exception("db down") + ) + + metadata = await instance._build_creator_attribution_metadata(self._job(), "batch-1") + + assert metadata["user_api_key_org_id"] == "org-team" + @pytest.mark.asyncio async def test_no_org_leaves_the_key_unset(self): """Without any org the key is absent entirely, so the spend writer's org update From 8debf8294cbe5e8d66d470b8dda1254668aa65df Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 12:42:39 -0700 Subject: [PATCH 29/60] fix(batches): resolve a legacy row's org from the key's organization_id and keep the Batch label on grouped rows --- .../proxy/common_utils/check_batch_cost.py | 2 +- tests/proxy_unit_tests/test_check_batch_cost.py | 8 ++++---- .../view_logs/RequestLogsTableColumns.test.tsx | 13 +++++++++++++ .../view_logs/RequestLogsTableColumns.tsx | 2 +- .../src/components/view_logs/TypeBadges.test.tsx | 7 +------ .../src/components/view_logs/TypeBadges.tsx | 4 ++-- 6 files changed, 22 insertions(+), 14 deletions(-) diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py index ce80ba6493f..6b3ccfd3694 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -158,7 +158,7 @@ class CheckBatchCost: where={"token": api_key} ) ) - key_org_id = getattr(key_row, "org_id", None) if key_row is not None else None + key_org_id = getattr(key_row, "organization_id", None) if key_row is not None else None if key_org_id: return key_org_id except Exception as e: diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index 9ce91a8aee4..f9fb782b052 100644 --- a/tests/proxy_unit_tests/test_check_batch_cost.py +++ b/tests/proxy_unit_tests/test_check_batch_cost.py @@ -2561,7 +2561,7 @@ class TestBatchCostAttribution: from types import SimpleNamespace instance = self._instance( - key_row=SimpleNamespace(key_alias="prod-key", org_id="org-moved-to"), + key_row=SimpleNamespace(key_alias="prod-key", organization_id="org-moved-to"), team_row=SimpleNamespace(team_alias="Team Alpha", organization_id="org-team"), ) @@ -2578,7 +2578,7 @@ class TestBatchCostAttribution: from types import SimpleNamespace instance = self._instance( - key_row=SimpleNamespace(key_alias="prod-key", org_id="org-42"), + key_row=SimpleNamespace(key_alias="prod-key", organization_id="org-42"), team_row=SimpleNamespace(team_alias="Team Alpha", organization_id="org-team"), ) @@ -2593,7 +2593,7 @@ class TestBatchCostAttribution: from types import SimpleNamespace instance = self._instance( - key_row=SimpleNamespace(key_alias="prod-key", org_id=None), + key_row=SimpleNamespace(key_alias="prod-key", organization_id=None), team_row=SimpleNamespace(team_alias="Team Alpha", organization_id="org-team"), ) @@ -2625,7 +2625,7 @@ class TestBatchCostAttribution: from types import SimpleNamespace instance = self._instance( - key_row=SimpleNamespace(key_alias="prod-key", org_id=None), + key_row=SimpleNamespace(key_alias="prod-key", organization_id=None), team_row=SimpleNamespace(team_alias="Team Alpha", organization_id=None), ) diff --git a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.test.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.test.tsx index 3e8e522afe1..9f0e659cb1f 100644 --- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.test.tsx @@ -141,6 +141,19 @@ describe("Type column", () => { expect(screen.getByText("Batch")).toBeInTheDocument(); expect(screen.queryByText("LLM")).not.toBeInTheDocument(); }); + + it("keeps the Batch label on the grouped create-plus-cost session instead of a row count", () => { + const groupedCostRow: Partial = { + request_id: "batch_1_batch_cost", + call_type: "aretrieve_batch", + session_id: "batch_1", + session_total_count: 2, + }; + renderRows([logEntry(groupedCostRow)]); + + expect(screen.getByText("Batch")).toBeInTheDocument(); + expect(screen.queryByText("2")).not.toBeInTheDocument(); + }); }); describe("batch rows", () => { diff --git a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.tsx index a444acb9517..1ec1087a1a4 100644 --- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.tsx @@ -65,7 +65,7 @@ export const getRequestLogsTableColumns = ({ const sessionMcpCount = log.mcp_tool_call_count ?? (isMcp ? sessionCount : 0); if (isBatchCallType(log.call_type)) { - return 1 ? sessionCount : undefined} />; + return ; } if (sessionCount <= 1) { if (isMcp) return ; diff --git a/ui/litellm-dashboard/src/components/view_logs/TypeBadges.test.tsx b/ui/litellm-dashboard/src/components/view_logs/TypeBadges.test.tsx index 8a964ea2124..9a3b53685ca 100644 --- a/ui/litellm-dashboard/src/components/view_logs/TypeBadges.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/TypeBadges.test.tsx @@ -45,14 +45,9 @@ describe("TypeBadges", () => { }); describe("BatchBadge", () => { - it("should render with default 'Batch' text when no count is provided", () => { + it("should render 'Batch'", () => { render(); expect(screen.getByText("Batch")).toBeInTheDocument(); }); - - it("should render the count when provided", () => { - render(); - expect(screen.getByText("4")).toBeInTheDocument(); - }); }); }); diff --git a/ui/litellm-dashboard/src/components/view_logs/TypeBadges.tsx b/ui/litellm-dashboard/src/components/view_logs/TypeBadges.tsx index 84079848487..db64bfbbe73 100644 --- a/ui/litellm-dashboard/src/components/view_logs/TypeBadges.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/TypeBadges.tsx @@ -97,9 +97,9 @@ export const AgentBadge = ({ count }: { count?: number }) => ( ); -export const BatchBadge = ({ count }: { count?: number }) => ( +export const BatchBadge = () => ( - {count != null ? count : "Batch"} + Batch ); From 217cb7da653e025f7c6d10eaae8e291b826fd76f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 20:30:33 -0700 Subject: [PATCH 30/60] fix(managed_files): resolve the creator org through the cached team lookup Batch creation snapshotted the team's organization with a direct litellm_teamtable query on every create. Go through get_team_object instead, which serves the team auth already cached and only falls back to the database when the team was never cached. --- .../proxy/hooks/managed_files.py | 13 +++- ..._batch_update_db_managed_output_file_id.py | 68 ++++++++++++++----- 2 files changed, 61 insertions(+), 20 deletions(-) diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 850e5aadf0e..68782c5516d 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -288,11 +288,18 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): return user_api_key_dict.org_id if not user_api_key_dict.team_id: return None + from litellm.proxy.auth.auth_checks import get_team_object + from litellm.proxy.proxy_server import proxy_logging_obj, user_api_key_cache + try: - team_row = await self.prisma_client.db.litellm_teamtable.find_unique( - where={"team_id": user_api_key_dict.team_id} + team: Final = await get_team_object( + team_id=user_api_key_dict.team_id, + prisma_client=self.prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=user_api_key_dict.parent_otel_span, + proxy_logging_obj=proxy_logging_obj, ) - return getattr(team_row, "organization_id", None) if team_row is not None else None + return team.organization_id except Exception as e: verbose_logger.warning(f"could not resolve org for managed object attribution: {e}") return None diff --git a/tests/test_litellm/enterprise/proxy/test_batch_update_db_managed_output_file_id.py b/tests/test_litellm/enterprise/proxy/test_batch_update_db_managed_output_file_id.py index b5865ab4a13..d3e668b8987 100644 --- a/tests/test_litellm/enterprise/proxy/test_batch_update_db_managed_output_file_id.py +++ b/tests/test_litellm/enterprise/proxy/test_batch_update_db_managed_output_file_id.py @@ -413,29 +413,63 @@ async def test_store_unified_object_id_persists_key_and_tags_on_create(): @pytest.mark.asyncio -async def test_store_unified_object_id_resolves_org_through_the_team(): +async def test_store_unified_object_id_resolves_org_through_the_cached_team(): """Most keys belong to an org only through their team, so the auth object carries no - org_id. The create resolves the team's organization so org spend is snapshotted at - submission time instead of never being billed.""" - from types import SimpleNamespace + org_id. The create reads the team that auth already cached, so org spend is snapshotted + at submission time without a database query in the request path.""" + from litellm.models.team import LiteLLM_TeamTableCachedObj + from litellm.proxy.proxy_server import user_api_key_cache + + instance, store = _in_memory_managed_files() + creator = UserAPIKeyAuth(user_id="alice", team_id="team-cached", api_key="hash-alice") + await user_api_key_cache.async_set_cache( + key="team_id:team-cached", + value=LiteLLM_TeamTableCachedObj(team_id="team-cached", organization_id="org-via-team"), + model_type=LiteLLM_TeamTableCachedObj, + ) + try: + await instance.store_unified_object_id( + unified_object_id="unified-b", + file_object=_build_batch_response(batch_id="b", status="validating"), + litellm_parent_otel_span=None, + model_object_id="b", + file_purpose="batch", + user_api_key_dict=creator, + persist_attribution=True, + ) + finally: + user_api_key_cache.delete_cache(key="team_id:team-cached") + + assert store["unified-b"]["org_id"] == "org-via-team" + instance.prisma_client.db.litellm_teamtable.find_unique.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_store_unified_object_id_resolves_org_from_the_db_when_the_team_is_not_cached(): + """A team no request has run under yet is absent from the auth cache; its organization + still comes back from the table so the org is billed rather than dropped.""" + from litellm.models.team import LiteLLM_TeamTable + from litellm.proxy.proxy_server import user_api_key_cache instance, store = _in_memory_managed_files() instance.prisma_client.db.litellm_teamtable.find_unique = AsyncMock( - return_value=SimpleNamespace(organization_id="org-via-team") + return_value=LiteLLM_TeamTable(team_id="team-uncached", organization_id="org-via-db") ) - creator = UserAPIKeyAuth(user_id="alice", team_id="team-alpha", api_key="hash-alice") + creator = UserAPIKeyAuth(user_id="alice", team_id="team-uncached", api_key="hash-alice") + try: + await instance.store_unified_object_id( + unified_object_id="unified-b", + file_object=_build_batch_response(batch_id="b", status="validating"), + litellm_parent_otel_span=None, + model_object_id="b", + file_purpose="batch", + user_api_key_dict=creator, + persist_attribution=True, + ) + finally: + user_api_key_cache.delete_cache(key="team_id:team-uncached") - await instance.store_unified_object_id( - unified_object_id="unified-b", - file_object=_build_batch_response(batch_id="b", status="validating"), - litellm_parent_otel_span=None, - model_object_id="b", - file_purpose="batch", - user_api_key_dict=creator, - persist_attribution=True, - ) - - assert store["unified-b"]["org_id"] == "org-via-team" + assert store["unified-b"]["org_id"] == "org-via-db" @pytest.mark.asyncio From 814c151b02b41e4fb333f4f505203d7db7c4343a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 01:35:09 -0700 Subject: [PATCH 31/60] fix(batches): mask api base credentials on batch cost rows --- .../proxy/common_utils/check_batch_cost.py | 4 +- litellm/litellm_core_utils/litellm_logging.py | 16 ++-- .../proxy_unit_tests/test_check_batch_cost.py | 84 +++++++++++++++++++ 3 files changed, 94 insertions(+), 10 deletions(-) diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py index 6b3ccfd3694..64a86387c39 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -685,7 +685,7 @@ class CheckBatchCost: from litellm.files.main import afile_content from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging - from litellm.litellm_core_utils.litellm_logging import deployment_pricing_model_info + from litellm.litellm_core_utils.litellm_logging import deployment_pricing_model_info, mask_api_base_credentials from litellm.proxy.openai_files_endpoints.common_utils import ( _is_base64_encoded_unified_file_id, ) @@ -858,7 +858,7 @@ class CheckBatchCost: "user-agent": CHECK_BATCH_COST_USER_AGENT, } }, - **({"api_base": deployment_api_base} if deployment_api_base else {}), + **({"api_base": mask_api_base_credentials(deployment_api_base)} if deployment_api_base else {}), "metadata": { **(await self._build_creator_attribution_metadata(job, batch_id)), # spend logs read the deployment identity off these metadata keys, so diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 73c06b891fb..7e800139228 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -419,6 +419,13 @@ def _provider_response_id(source: object) -> str | None: return candidate if isinstance(candidate, str) and candidate else None +def mask_api_base_credentials(api_base: str) -> str: + if "key=" not in api_base: + return api_base + key_end: Final = api_base.find("key=") + 4 + return api_base[:key_end] + "*" * 5 + api_base[-4:] + + class Logging(LiteLLMLoggingBaseClass): global \ supabaseClient, \ @@ -1160,14 +1167,7 @@ class Logging(LiteLLMLoggingBaseClass): return data def _get_masked_api_base(self, api_base: str) -> str: - if "key=" in api_base: - # Find the position of "key=" in the string - key_index: Final = api_base.find("key=") + 4 - # Mask the last 5 characters after "key=" - masked_api_base = api_base[:key_index] + "*" * 5 + api_base[-4:] - else: - masked_api_base = api_base - return str(masked_api_base) + return str(mask_api_base_credentials(api_base)) def _pre_call(self, input, api_key, model=None, additional_args={}): """ diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index f9fb782b052..36177b44930 100644 --- a/tests/proxy_unit_tests/test_check_batch_cost.py +++ b/tests/proxy_unit_tests/test_check_batch_cost.py @@ -583,6 +583,90 @@ class TestCheckBatchCost: assert passed_model_info["input_cost_per_token_batches"] == 2e-06 assert passed_model_info["output_cost_per_token_batches"] == 4e-06 + @pytest.mark.asyncio + async def test_poller_masks_api_base_credentials_before_logging( + self, check_batch_cost_instance, mock_prisma_client, mock_llm_router + ): + """Request rows mask `key=` query credentials out of api_base before it is + logged, but the poller skips that pre-call step, so an unmasked deployment + api_base would land verbatim on the batch cost row: regression test for the + poller masking the same way. + """ + import base64 + from unittest.mock import patch + + import httpx + import respx + + from litellm.litellm_core_utils.litellm_logging import Logging + + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=1) + mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() + mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + + mock_job = MagicMock() + mock_job.id = "job-masked-api-base-1" + mock_job.unified_object_id = base64.urlsafe_b64encode( + b"litellm_proxy;model_id:model-123;llm_batch_id:batch-456" + ).decode() + mock_job.created_by = "user-1" + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=[mock_job]) + + mock_response = MagicMock() + mock_response.status = "completed" + mock_response.output_file_id = "file-output-123" + mock_response.error_file_id = None + mock_response.model_dump_json.return_value = '{"id":"batch-1","status":"completed"}' + mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response) + mock_llm_router.get_deployment_credentials_with_provider = MagicMock(return_value={"api_key": "sk-test"}) + + mock_deployment = MagicMock() + mock_deployment.litellm_params.custom_llm_provider = "openai" + mock_deployment.litellm_params.model = "gpt-5.4-mini" + mock_deployment.litellm_params.api_base = "https://gateway.example.com/v1?key=AIzaSyVERYSECRET7890" + mock_deployment.model_info.model_dump.return_value = {} + mock_llm_router.get_deployment = MagicMock(return_value=mock_deployment) + + output_line = json.dumps( + { + "custom_id": "req-1", + "response": { + "status_code": 200, + "body": { + "id": "chatcmpl-1", + "object": "chat.completion", + "model": "gpt-5.4-mini", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "hi"}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + }, + }, + "error": None, + } + ) + + with ( + respx.mock(assert_all_called=True) as provider, + patch.object( # test-quality-ok: the poller builds Logging inline, the only seam to the row it logs + Logging, "async_success_handler", autospec=True + ) as success_handler, + ): + provider.get("https://api.openai.com/v1/files/file-output-123/content").mock( + return_value=httpx.Response(200, content=f"{output_line}\n".encode()) + ) + await check_batch_cost_instance.check_batch_cost() + + cost_row_calls = [call for call in success_handler.await_args_list if "batch_cost" in call.kwargs] + assert len(cost_row_calls) == 1 + logged_api_base = cost_row_calls[0].args[0].litellm_params["api_base"] + assert logged_api_base == "https://gateway.example.com/v1?key=*****7890" + assert "VERYSECRET" not in logged_api_base + @pytest.mark.asyncio async def test_primary_path_completion_update_includes_batch_processed( self, check_batch_cost_instance, mock_prisma_client, mock_llm_router From 95d721ffe60e6fa8409e2ff3e57ab9a85093fcc7 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 02:06:40 -0700 Subject: [PATCH 32/60] refactor(batches): drop docstrings restating the org fallbacks --- .../litellm_enterprise/proxy/common_utils/check_batch_cost.py | 3 --- enterprise/litellm_enterprise/proxy/hooks/managed_files.py | 4 ---- 2 files changed, 7 deletions(-) diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py index 64a86387c39..13e9e5093a8 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -143,9 +143,6 @@ class CheckBatchCost: return None async def _get_org_id(self, job: "LiteLLM_ManagedObjectTable", batch_id: str) -> str | None: - """Organization to bill the batch against, snapshotted on the row at creation - like team_id. Rows created before the org_id column existed carry None, so they - fall back to the creating key's org (or its team's) as resolved today.""" org_id = getattr(job, "org_id", None) if org_id: return org_id diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 6474475f1bf..486904d0abe 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -281,10 +281,6 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): verbose_logger.debug(f"LiteLLM Managed File object with id={file_id} stored in db: {result}") async def _resolve_creator_org_id(self, user_api_key_dict: UserAPIKeyAuth) -> Optional[str]: - """Organization to snapshot on the managed object row, like team_id. A key that - belongs to an org only through its team carries no org_id on the auth object, so - resolve the team's organization at creation time; costing then bills the org the - batch was submitted under even if the key or team moves before it completes.""" if user_api_key_dict.org_id: return user_api_key_dict.org_id if not user_api_key_dict.team_id: From 4cc0180eab7482a59722e68999def58e11077a72 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 7 Sep 2026 17:06:34 -0700 Subject: [PATCH 33/60] fix(router): keep unresolved drop_params strings so DB rows and env refs survive The drop_params validator collapsed every string it did not recognize to None. A pre-fix DB row holds the flag as ciphertext, so a partial PATCH rebuilt the deployment without it and dropped the key from the stored row, and /model/new turned an os.environ/ reference into nothing before the loader could resolve it. The validator now returns the raw value when it is not a boolean flag, the field admits strings the way timeout already does, and the flag set follows pydantic's lax bool parsing instead of a hand-rolled true/false pair --- litellm/litellm_core_utils/core_helpers.py | 15 +++++----- litellm/types/router.py | 7 +++-- .../litellm_core_utils/test_core_helpers.py | 11 +++++-- .../test_model_management_endpoints.py | 27 ++++++++++++++++- .../proxy/proxy_server/test_proxy_config.py | 30 +++++++++++++++++++ tests/test_litellm/types/test_router.py | 23 ++++++++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 4 +-- 7 files changed, 101 insertions(+), 16 deletions(-) diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index 671b2ddce99..bd7a1b8f384 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -5,6 +5,7 @@ from collections.abc import Iterable, Mapping from typing import TYPE_CHECKING, Any, Final, Literal import httpx +from pydantic import TypeAdapter, ValidationError from litellm._logging import verbose_logger from litellm.types.llms.openai import AllMessageValues, OpenAIChatCompletionFinishReason @@ -37,16 +38,16 @@ def safe_divide_seconds(seconds: float, denominator: float, default: float | Non return float(seconds / denominator) +_DROP_PARAMS_BOOL: Final = TypeAdapter(bool) + + def normalize_drop_params(value: object) -> bool | None: if isinstance(value, bool): return value - if isinstance(value, str): - lowered: Final = value.strip().lower() - if lowered == "true": - return True - if lowered == "false": - return False - return None + try: + return _DROP_PARAMS_BOOL.validate_python(value.strip() if isinstance(value, str) else value) + except ValidationError: + return None def safe_divide( diff --git a/litellm/types/router.py b/litellm/types/router.py index e2a04f9da77..47aea6430c3 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -315,7 +315,7 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams): timeout: float | str | httpx.Timeout | None = None # if str, pass in as os.environ/ stream_timeout: float | str | None = None # timeout when making stream=True calls, if str, pass in as os.environ/ max_retries: int | None = None - drop_params: bool | None = None + drop_params: bool | str | None = None organization: str | None = None # for openai orgs configurable_clientside_auth_params: CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS = None litellm_credential_name: str | None = None @@ -408,8 +408,9 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams): @field_validator("drop_params", mode="before") @classmethod - def coerce_drop_params(cls, value: object) -> bool | None: - return normalize_drop_params(value) + def coerce_drop_params(cls, value: object) -> object: + normalized: Final = normalize_drop_params(value) + return value if normalized is None else normalized def __contains__(self, key) -> bool: # Define custom behavior for the 'in' operator diff --git a/tests/test_litellm/litellm_core_utils/test_core_helpers.py b/tests/test_litellm/litellm_core_utils/test_core_helpers.py index ab43141af23..8797f7c3591 100644 --- a/tests/test_litellm/litellm_core_utils/test_core_helpers.py +++ b/tests/test_litellm/litellm_core_utils/test_core_helpers.py @@ -268,11 +268,16 @@ class TestRedactNestedMatchAndRegexKeys: (" TRUE ", True), ("false", False), ("False", False), + ("yes", True), + ("off", False), + ("1", True), + (1, True), + (0, False), (None, None), - ("yes", None), ("", None), - (1, None), - (0, None), + ("os.environ/DROP_PARAMS", None), + ("v2:gcm:not-a-flag", None), + (2, None), ], ) def test_normalize_drop_params(value, expected): diff --git a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py index d90338f8480..c02f886fc31 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py @@ -30,7 +30,7 @@ from litellm.proxy.management_endpoints.model_management_endpoints import ( ) from litellm.proxy.utils import PrismaClient from litellm.router import Router -from litellm.types.router import Deployment, LiteLLM_Params, updateDeployment +from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo, updateDeployment, updateLiteLLMParams async def _passthrough_row(update_data): @@ -3070,6 +3070,31 @@ class TestUpdateDBModelBlocked: assert "blocked" not in result +class TestUpdateDBModelKeepsLegacyDropParams: + def test_partial_patch_keeps_encrypted_string_drop_params(self, monkeypatch): + from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper + from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model + + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-1234") + legacy_row = Deployment( + model_name="gpt-5-nano", + litellm_params=LiteLLM_Params( + model="openai/gpt-5-nano", + api_key=encrypt_value_helper(value="sk-old"), + drop_params=encrypt_value_helper(value="true"), + ), + model_info=ModelInfo(id="legacy-row"), + ) + + result = update_db_model( + db_model=legacy_row, + updated_patch=updateDeployment(litellm_params=updateLiteLLMParams(api_key="sk-new")), + ) + + stored = json.loads(result["litellm_params"]) + assert decrypt_value_helper(value=stored["drop_params"], key="drop_params") == "true" + + def _build_db_model_with_pricing(): """Wildcard deployment with custom pricing in litellm_params; Deployment.__init__ mirrors SPECIAL_MODEL_INFO_PARAMS into model_info, so both blobs hold the rate.""" diff --git a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py index 2babfe432f3..a4be9574f89 100644 --- a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py @@ -19,6 +19,7 @@ import pytest import litellm from litellm.proxy._types import CommonProxyErrors +from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper from litellm.proxy.proxy_server import ( ProxyConfig, _is_remote_module_url, @@ -2401,6 +2402,35 @@ def test_ProxyConfig__add_deployment_resolves_env_refs_on_arbitrary_field(monkey assert deployment.litellm_params.some_future_field == "resolved-custom-value" +@pytest.mark.parametrize( + "stored_drop_params", + ["true", "os.environ/DROP_PARAMS_FLAG"], +) +def test_ProxyConfig__add_deployment_turns_stored_drop_params_string_into_bool(monkeypatch, stored_drop_params): + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-1234") + monkeypatch.setenv("DROP_PARAMS_FLAG", "true") + fake_router = MagicMock() + fake_router.upsert_deployment = MagicMock(return_value=True) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", fake_router) + pc = ProxyConfig() + db_model = SimpleNamespace( + model_id="model-1", + model_name="gpt-5-nano", + model_info={"id": "model-1"}, + litellm_params={ + "model": encrypt_value_helper(value="openai/gpt-5-nano"), + "drop_params": encrypt_value_helper(value=stored_drop_params), + }, + blocked=False, + ) + + added = pc._add_deployment(db_models=[db_model]) + deployment = fake_router.upsert_deployment.call_args.kwargs["deployment"] + + assert added == 1 + assert deployment.litellm_params.drop_params is True + + # --------------------------------------------------------------------------- # ProxyConfig.decrypt_model_list_from_db # --------------------------------------------------------------------------- diff --git a/tests/test_litellm/types/test_router.py b/tests/test_litellm/types/test_router.py index accd3b32a0d..2f4a29473c3 100644 --- a/tests/test_litellm/types/test_router.py +++ b/tests/test_litellm/types/test_router.py @@ -1,8 +1,10 @@ import pytest +from pydantic import ValidationError from litellm.types.router import ( SPECIAL_MODEL_INFO_PARAMS, Deployment, + GenericLiteLLMParams, LiteLLM_Params, ModelInfo, ) @@ -89,3 +91,24 @@ def test_pricing_strings_are_coerced_to_float(): def test_invalid_pricing_is_rejected(): with pytest.raises(ValueError, match='validation error for ModelInfo'): ModelInfo(id="x", input_cost_per_token="free") + + +@pytest.mark.parametrize( + "value, expected", + [ + (True, True), + ("true", True), + (" False ", False), + ("yes", True), + (None, None), + ("os.environ/DROP_PARAMS", "os.environ/DROP_PARAMS"), + ("v2:gcm:ciphertext-from-a-pre-fix-row", "v2:gcm:ciphertext-from-a-pre-fix-row"), + ], +) +def test_drop_params_coerces_flags_and_keeps_unresolved_strings(value, expected): + assert GenericLiteLLMParams(drop_params=value).drop_params == expected + + +def test_drop_params_rejects_non_flag_non_string_values(): + with pytest.raises(ValidationError): + GenericLiteLLMParams(drop_params=2) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index f7552bfb397..7b5edf6acfd 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -29397,7 +29397,7 @@ export interface components { /** Default Api Key Tpm Limit */ default_api_key_tpm_limit?: number | null; /** Drop Params */ - drop_params?: boolean | null; + drop_params?: boolean | string | null; /** Gcs Bucket Name */ gcs_bucket_name?: string | null; /** Google Maps Grounding Cost Per Query */ @@ -39570,7 +39570,7 @@ export interface components { /** Default Api Key Tpm Limit */ default_api_key_tpm_limit?: number | null; /** Drop Params */ - drop_params?: boolean | null; + drop_params?: boolean | string | null; /** Gcs Bucket Name */ gcs_bucket_name?: string | null; /** Google Maps Grounding Cost Per Query */ From 1067697c7b471200cd26ee93834f71b09d013415 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 7 Sep 2026 17:11:56 -0700 Subject: [PATCH 34/60] refactor(utils): gate the triton branch on bool(drop_params) like every other provider --- litellm/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/utils.py b/litellm/utils.py index 59db4162410..50282674868 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -4302,7 +4302,7 @@ def get_optional_params( non_default_params=non_default_params, optional_params=optional_params, model=model, - drop_params=drop_params if drop_params is not None else False, + drop_params=bool(drop_params), ) elif custom_llm_provider == "maritalk": From 192ea9ec80ed97d771f9258320ac948080e87d2e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 7 Sep 2026 17:55:37 -0700 Subject: [PATCH 35/60] fix(policy_engine): fail open on streaming shapes post_call pipelines cannot govern yet A post_call pipeline now releases the original stream instead of refusing the request on every shape it has no handler for: a background request, a pipeline guardrail without the unified apply_guardrail interface, a route with no endpoint translation, a buffered stream no translation resolves, and a rewrite the translation cannot write back (tool-call edits, text edits on translations without write-back, n>1 chat, an unended Anthropic stream, a Responses dump with no event envelope). Each case logs a warning naming the policy and guardrail. Real blocks and writable text masks are unchanged. --- .../chat/guardrail_translation/handler.py | 3 +- .../guardrail_translation/base_translation.py | 5 +- .../chat/guardrail_translation/handler.py | 5 +- .../guardrail_translation/handler.py | 7 +- .../proxy/policy_engine/pipeline_executor.py | 66 ++-- litellm/proxy/utils.py | 226 ++++++-------- .../policy_engine/test_pipeline_executor.py | 137 ++++++++- .../proxy_logging/test_guardrail_pipeline.py | 282 +++++++++++------- 8 files changed, 437 insertions(+), 294 deletions(-) diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index da30d85e26b..e486be12fe2 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -1027,7 +1027,8 @@ class AnthropicMessagesHandler(BaseTranslation): Get the string so far, check the apply guardrail to the string so far, and return the list of responses so far. With ``deliver_ended_stream_rewrites``, an ended stream whose guardrail rewrote the text gets the rewrite written back across the buffered chunks (full rewritten text in the first ``text_delta``, the rest blanked); - a rewrite on a stream that never reported a ``stop_reason`` has no write-back and fails closed instead. + a rewrite on a stream that never reported a ``stop_reason`` has no write-back and is reported as + undeliverable, so the pipeline executor discards it and releases the original chunks. """ from litellm.integrations.custom_guardrail import ModifyResponseException diff --git a/litellm/llms/base_llm/guardrail_translation/base_translation.py b/litellm/llms/base_llm/guardrail_translation/base_translation.py index 770fac6e443..afd8e0f67f7 100644 --- a/litellm/llms/base_llm/guardrail_translation/base_translation.py +++ b/litellm/llms/base_llm/guardrail_translation/base_translation.py @@ -56,8 +56,9 @@ class BaseTranslation(ABC): """Whether ``process_output_streaming_response`` accepts ``deliver_ended_stream_rewrites=True`` and, on an ended (fully buffered) stream, writes guardrail text rewrites back across ``responses_so_far`` so - a buffered pipeline can release rewritten chunks instead of withholding the - stream. Tool-call rewrites stay undeliverable everywhere.""" + a buffered pipeline can release rewritten chunks. Tool-call rewrites, and + text rewrites on every other translation, are undeliverable: the pipeline + executor discards them and releases the original chunks.""" @staticmethod def transform_user_api_key_dict_to_metadata( diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 4b7b1cc2700..80292aef2cf 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -1015,7 +1015,8 @@ class OpenAIChatCompletionsHandler(BaseTranslation): content-carrying chunk and the rest are blanked, the same shape the in-flight write-back uses. Chunks carrying only finish_reason or usage stay untouched. A rewrite on a stream carrying more than one distinct - choice index fails closed.""" + choice index is reported as undeliverable, so the pipeline executor + discards it and releases the original chunks.""" post_guardrail_texts: Final = self._string_choice_contents(guardrailed_response) changed: Final = tuple( after @@ -1030,7 +1031,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): if len(stream_choice_indices) != 1: # stream_chunk_builder collapses every choice into one index-0 # choice, so a rewrite of the rebuilt response cannot be attributed - # back to a single choice on an n>1 stream: withhold the stream + # back to a single choice on an n>1 stream: report it undeliverable # rather than deliver the rewrite on the wrong choice from litellm.proxy.policy_engine.pipeline_executor import UndeliverableStreamRewrite diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index b543c7f1e17..b0f79552bc5 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -699,7 +699,8 @@ class OpenAIResponsesHandler(BaseTranslation): ``response.content_part.done``, ``response.output_item.done``) are synced to the rewritten envelope too, so a client reading deltas sees the rewrite instead of the raw model output; a rewrite observed where no - write-back is possible fails closed instead of releasing raw output. + write-back is possible is reported as undeliverable, so the pipeline + executor discards it and releases the original events. """ if not responses_so_far: return responses_so_far @@ -788,7 +789,7 @@ class OpenAIResponsesHandler(BaseTranslation): # ------------------------------------------------------------------ # # Case 2: response.output_item.done — extract tool calls only, then # # fall through to the text fallback when a caller expects rewrites # - # delivered, so a buffer truncated here still fails closed on text. # + # delivered, so a truncated buffer still reports text undeliverable. # # ------------------------------------------------------------------ # if final_chunk.get("type") == "response.output_item.done": model_response_stream: Final = ( @@ -813,7 +814,7 @@ class OpenAIResponsesHandler(BaseTranslation): # Fallback: apply guardrail to the accumulated text string. # # No structured write-back is possible here; guardrails that only # # need to block/flag (not rewrite) still work correctly, and a # - # rewrite a caller expects delivered fails closed instead. # + # rewrite a caller expects delivered is reported undeliverable. # # ------------------------------------------------------------------ # string_so_far: Final = self.get_streaming_string_so_far(responses_so_far) if string_so_far: diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index 160938b723c..970ac20487c 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -5,6 +5,7 @@ Runs guardrails sequentially per pipeline step definitions, handling pass/fail actions (allow, block, next, modify_response) and data forwarding. """ +import copy import time from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final, Literal @@ -77,7 +78,7 @@ class _StreamRewriteObserver(CustomGuardrail): which for guardrails like Bedrock's ANONYMIZED action is only known at runtime. Text rewrites are deliverable on translations that write them back across the buffered chunks (``delivers_ended_stream_text_rewrites``); tool-call rewrites and text rewrites on any - other translation make the gate withhold the stream.""" + other translation are discarded by the executor, which releases the original chunks.""" def __init__(self, inner: CustomGuardrail) -> None: super().__init__(guardrail_name=inner.guardrail_name) @@ -133,6 +134,19 @@ def _prepare_hook_input( return hook_input, scans_raw_request +def _release_original_chunks( + guardrail_name: str, + streaming_chunks: list[object], # mutable-ok: shared buffered-stream chunks, restored in place + originals: Sequence[object], +) -> None: + streaming_chunks[:] = originals # rebind-ok: the caller's buffer is the stream the client receives + verbose_proxy_logger.warning( + "Pipeline: guardrail '%s' rewrote the streamed response in a way this endpoint's streaming " + "pipeline cannot deliver yet; the rewrite was discarded and the original stream released", + guardrail_name, + ) + + class PipelineExecutor: """Executes guardrail pipelines with ordered, conditional step logic.""" @@ -263,29 +277,37 @@ class PipelineExecutor: litellm_logging_obj: "LiteLLMLoggingObj | None", ) -> None: """Run one streaming post_call step through the endpoint translation, delivering - text rewrites on translations that support ended-stream write-back and raising - ``UndeliverableStreamRewrite`` for any rewrite that cannot reach the client.""" + text rewrites on translations that support ended-stream write-back. A rewrite that + cannot reach the client yet (a tool-call rewrite, a text rewrite on a translation + without write-back, or one the translation refused with + ``UndeliverableStreamRewrite``) is discarded: the buffered chunks go back to the + originals and the step passes, so the client gets the stream the merge base sent.""" observer: Final = _StreamRewriteObserver(callback) deliver_rewrites: Final = type(endpoint_translation).delivers_ended_stream_text_rewrites - if deliver_rewrites: - await endpoint_translation.process_output_streaming_response( - responses_so_far=streaming_chunks, - guardrail_to_apply=observer, - litellm_logging_obj=litellm_logging_obj, - user_api_key_dict=user_api_key_dict, - request_data=hook_input, - deliver_ended_stream_rewrites=True, - ) - else: - await endpoint_translation.process_output_streaming_response( - responses_so_far=streaming_chunks, - guardrail_to_apply=observer, - litellm_logging_obj=litellm_logging_obj, - user_api_key_dict=user_api_key_dict, - request_data=hook_input, - ) + originals: Final = copy.deepcopy(streaming_chunks) + try: + if deliver_rewrites: + await endpoint_translation.process_output_streaming_response( + responses_so_far=streaming_chunks, + guardrail_to_apply=observer, + litellm_logging_obj=litellm_logging_obj, + user_api_key_dict=user_api_key_dict, + request_data=hook_input, + deliver_ended_stream_rewrites=True, + ) + else: + await endpoint_translation.process_output_streaming_response( + responses_so_far=streaming_chunks, + guardrail_to_apply=observer, + litellm_logging_obj=litellm_logging_obj, + user_api_key_dict=user_api_key_dict, + request_data=hook_input, + ) + except UndeliverableStreamRewrite: + _release_original_chunks(step.guardrail, streaming_chunks, originals) + return if observer.rewrote_tool_calls or (observer.rewrote_texts and not deliver_rewrites): - raise UndeliverableStreamRewrite(step.guardrail) + _release_original_chunks(step.guardrail, streaming_chunks, originals) @staticmethod async def _run_step( @@ -386,8 +408,6 @@ class PipelineExecutor: ) # mutable-ok: modified-data contract is a plain dict return ("pass", response if isinstance(response, dict) else None, None, None) - except UndeliverableStreamRewrite: - raise except Exception as e: if CustomGuardrail._is_guardrail_intervention(e): error_msg: Final = _extract_error_message(e) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index ca8e48b5dfa..d7bf832bca4 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -19,7 +19,7 @@ from email.mime.text import MIMEText from types import MappingProxyType from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, Protocol, TypeVar, Union, cast, overload -from typing_extensions import NotRequired, ReadOnly, TypedDict +from typing_extensions import ReadOnly, TypedDict from litellm import _custom_logger_compatible_callbacks_literal from litellm.constants import ( @@ -155,7 +155,7 @@ from litellm.proxy.hooks.sensitive_data_routing import ( ) from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup, add_guardrails_from_auth_metadata from litellm.proxy.management_helpers.key_settings_audit import with_settings_updated_at -from litellm.proxy.policy_engine.pipeline_executor import PipelineExecutor, UndeliverableStreamRewrite +from litellm.proxy.policy_engine.pipeline_executor import PipelineExecutor from litellm.repositories.budget_repository import BudgetRepository from litellm.repositories.config_repository import ConfigRepository from litellm.repositories.table_repositories import ( @@ -518,108 +518,72 @@ def _pipeline_step_supports_unified_streaming(guardrail_name: str) -> bool: return callback is not None and PipelineExecutor.supports_unified_execution(callback) -class _PipelineErrorBody(TypedDict): - message: ReadOnly[str] - type: ReadOnly[str] - policies: ReadOnly[tuple[str, ...]] - guardrails: NotRequired[ReadOnly[tuple[str, ...]]] - - -class _PipelineErrorDetail(TypedDict): - error: ReadOnly[_PipelineErrorBody] - - -def _undeliverable_stream_rewrite_error(policy_name: str, guardrail_name: str) -> HTTPException: - detail: Final[_PipelineErrorDetail] = { - "error": { - "message": ( - f"Streaming response withheld by policy pipeline '{policy_name}' because guardrail " - f"'{guardrail_name}' rewrote the streamed output in a way this endpoint's streaming " - "pipeline cannot deliver (a tool-call rewrite, or a text rewrite on a route without " - "stream write-back). Retry with stream=false, or drop it from the pipeline steps so " - "guardrails.add applies it to streamed output." - ), - "type": "guardrail_pipeline_error", - "policies": (policy_name,), - "guardrails": (guardrail_name,), - } - } - return HTTPException(status_code=400, detail=detail) - - -def _raise_for_streaming_post_call_pipelines(data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth) -> None: - """ - Reject up front the requests whose post_call pipelines could never run. - - Background responses skip the post_call hooks entirely, so a pipeline - governing one would silently never execute. Streaming responses execute - pipelines against the buffered stream through the endpoint guardrail - translation of the request route, releasing the buffered chunks on allow - (rewritten in place when a guardrail rewrote text and the translation - delivers ended-stream rewrites; a rewrite the translation cannot deliver - fails closed at runtime instead). That needs every step's guardrail to - support the unified apply_guardrail interface, and needs the route to have - a translation at all; anything else keeps the 400 rather than letting - ungoverned output stream through. - """ - is_stream: Final = data.get("stream") is True - is_background: Final = data.get("background") is True - if not is_stream and not is_background: - return - post_call_pipelines: Final = tuple( +def _post_call_pipelines(data: Mapping[str, object]) -> tuple[tuple[str, "GuardrailPipeline"], ...]: + return tuple( (policy_name, pipeline) for policy_name, pipeline in _policy_pipelines(data) if pipeline.mode == "post_call" ) + + +def _warn_background_skips_post_call_pipelines(data: Mapping[str, object]) -> None: + if data.get("background") is not True: + return + policy_names: Final = tuple(policy_name for policy_name, _pipeline in _post_call_pipelines(data)) + if not policy_names: + return + verbose_proxy_logger.warning( + "Policies with post_call guardrail pipelines do not run on background responses yet; " + "the response is released ungoverned by them: %s", + ", ".join(policy_names), + ) + + +def _pipeline_is_streamable(policy_name: str, pipeline: "GuardrailPipeline") -> bool: + unsupported: Final = tuple( + dict.fromkeys( + step.guardrail for step in pipeline.steps if not _pipeline_step_supports_unified_streaming(step.guardrail) + ) + ) + if not unsupported: + return True + verbose_proxy_logger.warning( + "Policy '%s' has post_call pipeline guardrails without the unified apply_guardrail interface, " + "which streaming pipelines need; the stream is released ungoverned by it: %s", + policy_name, + ", ".join(unsupported), + ) + return False + + +def _streamable_post_call_pipelines( + request_data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth +) -> tuple[tuple[str, "GuardrailPipeline"], ...]: + """ + The post_call pipelines a streaming response can be gated through. + + Streaming pipelines scan the buffered stream through the endpoint guardrail + translation of the request route, so every step's guardrail needs the + unified apply_guardrail interface and the route needs a translation. A + pipeline that cannot be run that way yet is left out and the stream is + released the way it was before pipelines ran on streams at all, with a + warning naming what went ungoverned. + """ + post_call_pipelines: Final = _post_call_pipelines(request_data) if not post_call_pipelines: - return - post_call_policies: Final = tuple(policy_name for policy_name, _pipeline in post_call_pipelines) - if is_background: - background_detail: Final[_PipelineErrorDetail] = { - "error": { - "message": ( - "Policies with post_call guardrail pipelines cannot govern background " - f"responses: {', '.join(post_call_policies)}. Retry with background=false." - ), - "type": "guardrail_pipeline_error", - "policies": post_call_policies, - } - } - raise HTTPException(status_code=400, detail=background_detail) - step_guardrails: Final = tuple( - dict.fromkeys(step.guardrail for _policy_name, pipeline in post_call_pipelines for step in pipeline.steps) - ) - unsupported_guardrails: Final = tuple( - guardrail for guardrail in step_guardrails if not _pipeline_step_supports_unified_streaming(guardrail) - ) - if unsupported_guardrails: - unsupported_detail: Final[_PipelineErrorDetail] = { - "error": { - "message": ( - "Policies with post_call guardrail pipelines cannot govern streaming responses " - "because these pipeline guardrails do not support the unified apply_guardrail " - f"interface: {', '.join(unsupported_guardrails)}. Retry with stream=false, or drop " - "them from the pipeline steps so guardrails.add scans them on streamed output." - ), - "type": "guardrail_pipeline_error", - "policies": post_call_policies, - "guardrails": unsupported_guardrails, - } - } - raise HTTPException(status_code=400, detail=unsupported_detail) + return () route: Final = user_api_key_dict.request_route - if not route or resolve_endpoint_translation(user_api_key_dict, None) is not None: - return - route_detail: Final[_PipelineErrorDetail] = { - "error": { - "message": ( - "Policies with post_call guardrail pipelines cannot govern streaming responses on " - f"route {route} because it has no endpoint guardrail translation to scan the stream " - f"through: {', '.join(post_call_policies)}. Retry with stream=false." - ), - "type": "guardrail_pipeline_error", - "policies": post_call_policies, - } - } - raise HTTPException(status_code=400, detail=route_detail) + if route and resolve_endpoint_translation(user_api_key_dict, None) is None: + verbose_proxy_logger.warning( + "Policies with post_call guardrail pipelines cannot scan streaming responses on route %s yet " + "(no endpoint guardrail translation); the stream is released ungoverned by them: %s", + route, + ", ".join(policy_name for policy_name, _pipeline in post_call_pipelines), + ) + return () + return tuple( + (policy_name, pipeline) + for policy_name, pipeline in post_call_pipelines + if _pipeline_is_streamable(policy_name, pipeline) + ) def _prompt_block_text(block: object) -> str: @@ -1990,7 +1954,7 @@ class ProxyLogging: ) try: - _raise_for_streaming_post_call_pipelines(data, user_api_key_dict) + _warn_background_skips_post_call_pipelines(data) # Execute guardrail pipelines before the normal callback loop data, _ = await self._maybe_execute_pipelines( # rebind-ok: pipeline edits feed the callback loop below @@ -3371,11 +3335,7 @@ class ProxyLogging: 1. /chat/completions """ caps: Final = ProxyLogging._callback_capabilities() - post_call_pipelines: Final = tuple( - (policy_name, pipeline) - for policy_name, pipeline in _policy_pipelines(request_data) - if pipeline.mode == "post_call" - ) + post_call_pipelines: Final = _streamable_post_call_pipelines(request_data, user_api_key_dict) # Fast path: no real overrides. Internal proxy CustomLogger callbacks # (e.g. _PROXY_MaxBudgetLimiter, ManagedFiles) inherit the default # ``async for chunk: yield chunk`` body, so wrapping the iterator @@ -3486,9 +3446,11 @@ class ProxyLogging: output, rewritten in place when one rewrote text and the translation delivers ended-stream rewrites (later steps then re-scan the rewritten chunks, so rewrites chain). A rewrite the translation cannot deliver - (a tool-call rewrite, or a text rewrite on a route without write-back) - withholds the stream with a 400; a block or modify_response terminates - with the translation's block chunks or the raised error. + yet (a tool-call rewrite, or a text rewrite on a route without + write-back) is discarded by the executor and the original chunks are + released, as is a buffered shape no translation resolves; a block or + modify_response terminates with the translation's block chunks or the + raised error. """ buffered: Final[list[object]] = [] # mutable-ok: accumulates the stream before the pipeline verdict async for item in response: @@ -3498,39 +3460,27 @@ class ProxyLogging: resolved: Final = resolve_endpoint_translation(user_api_key_dict, buffered[0]) if resolved is None: - policy_names: Final = tuple(policy_name for policy_name, _pipeline in pipelines) - raise ProxyException( - message=( - "Policy pipelines could not govern this streaming response shape; " - f"the response was withheld: {', '.join(policy_names)}." - ), - type="guardrail_pipeline_error", - param=None, - code=500, + verbose_proxy_logger.warning( + "Policies with post_call guardrail pipelines cannot scan this streaming response shape yet; " + "the stream is released ungoverned by them: %s", + ", ".join(policy_name for policy_name, _pipeline in pipelines), ) + for buffered_item in buffered: + yield buffered_item + return call_type, endpoint_translation = resolved for policy_name, pipeline in pipelines: - try: - result: PipelineExecutionResult = await PipelineExecutor.execute_steps( - steps=pipeline.steps, - mode="post_call", - data=request_data, - user_api_key_dict=user_api_key_dict, - call_type=call_type, - policy_name=policy_name, - streaming_chunks=buffered, - endpoint_translation=endpoint_translation, - ) - except UndeliverableStreamRewrite as rewrite: - async for error_chunk in unified_guardrail.emit_streaming_http_error( - _undeliverable_stream_rewrite_error(policy_name, rewrite.guardrail_name), - call_type, - buffered, - request_data, - ): - yield error_chunk - return + result: PipelineExecutionResult = await PipelineExecutor.execute_steps( + steps=pipeline.steps, + mode="post_call", + data=request_data, + user_api_key_dict=user_api_key_dict, + call_type=call_type, + policy_name=policy_name, + streaming_chunks=buffered, + endpoint_translation=endpoint_translation, + ) try: ProxyLogging._handle_pipeline_result( result, data=request_data, policy_name=policy_name, original_response=buffered diff --git a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py index d399b4f01cb..0685bc6aa1e 100644 --- a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py +++ b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py @@ -4,6 +4,7 @@ Tests for the pipeline executor. Uses mock guardrails to validate pipeline execution without external services. """ +import logging from unittest.mock import MagicMock import pytest @@ -941,7 +942,55 @@ class _TextTranslation: return responses_so_far -async def _run_streaming_step(returned_texts, translation): +class _WritingTranslation: + """Writes the guardrail's text (and tool-call) outputs back into the buffered chunks the way the + chat/Responses/Messages handlers do on an ended stream.""" + + delivers_ended_stream_text_rewrites = True + + async def process_output_streaming_response( + self, + responses_so_far, + guardrail_to_apply, + litellm_logging_obj=None, + user_api_key_dict=None, + request_data=None, + deliver_ended_stream_rewrites=False, + ): + assert deliver_ended_stream_rewrites is True + outputs = await guardrail_to_apply.apply_guardrail( + inputs={"texts": [responses_so_far[0]["text"]], "tool_calls": [dict(responses_so_far[0]["tool_call"])]}, + request_data=request_data or {}, + input_type="response", + logging_obj=litellm_logging_obj, + ) + responses_so_far[0]["text"] = outputs["texts"][0] + responses_so_far[0]["tool_call"] = outputs["tool_calls"][0] + return responses_so_far + + +class _RefusingTranslation: + delivers_ended_stream_text_rewrites = True + + async def process_output_streaming_response( + self, + responses_so_far, + guardrail_to_apply, + litellm_logging_obj=None, + user_api_key_dict=None, + request_data=None, + deliver_ended_stream_rewrites=False, + ): + responses_so_far[0]["text"] = "half-written" + raise UndeliverableStreamRewrite(guardrail_to_apply.guardrail_name) + + +def _chunk(): + return {"text": "hello world", "tool_call": {"function": {"name": "lookup", "arguments": '{"ssn": "123"}'}}} + + +async def _run_streaming_step(translation, streaming_chunks=None): + chunks = [object()] if streaming_chunks is None else streaming_chunks return await PipelineExecutor.execute_steps( steps=[PipelineStep(guardrail="masker", on_pass="allow", on_fail="next", on_error="next")], mode="post_call", @@ -949,31 +998,41 @@ async def _run_streaming_step(returned_texts, translation): user_api_key_dict=MagicMock(), call_type="completion", policy_name="p", - streaming_chunks=[object()], + streaming_chunks=chunks, endpoint_translation=translation, ) +def _assert_passed_with_discard_warning(result, caplog): + assert result.terminal_action == "allow" + assert [step.outcome for step in result.step_results] == ["pass"] + assert any("'masker'" in record.getMessage() and "discarded" in record.getMessage() for record in caplog.records) + + @pytest.mark.asyncio -async def test_streaming_step_rewrite_escapes_execute_steps_regardless_of_step_actions(monkeypatch): +async def test_streaming_step_discards_text_rewrite_when_translation_lacks_write_back(monkeypatch, caplog): monkeypatch.setattr(litellm, "callbacks", [_TextReturningGuardrail(["hello [MASKED]"])]) translation = _TextTranslation() + chunks = [_chunk()] - with pytest.raises(UndeliverableStreamRewrite) as info: - await _run_streaming_step(["hello [MASKED]"], translation) + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + result = await _run_streaming_step(translation, chunks) - assert info.value.guardrail_name == "masker" + _assert_passed_with_discard_warning(result, caplog) + assert chunks == [_chunk()] assert translation.seen_guardrail_names == ["masker"] @pytest.mark.asyncio -async def test_streaming_step_unchanged_texts_in_another_container_allow(monkeypatch): +async def test_streaming_step_unchanged_texts_in_another_container_allow(monkeypatch, caplog): monkeypatch.setattr(litellm, "callbacks", [_TextReturningGuardrail(("hello world",))]) - result = await _run_streaming_step(("hello world",), _TextTranslation()) + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + result = await _run_streaming_step(_TextTranslation()) assert result.terminal_action == "allow" assert [step.outcome for step in result.step_results] == ["pass"] + assert not any("discarded" in record.getMessage() for record in caplog.records) class _InPlaceMutatingGuardrail(CustomGuardrail): @@ -989,10 +1048,64 @@ class _InPlaceMutatingGuardrail(CustomGuardrail): @pytest.mark.asyncio -async def test_streaming_step_in_place_rewrite_still_withholds_stream(monkeypatch): +async def test_streaming_step_in_place_rewrite_is_discarded_without_write_back(monkeypatch, caplog): monkeypatch.setattr(litellm, "callbacks", [_InPlaceMutatingGuardrail()]) + chunks = [_chunk()] - with pytest.raises(UndeliverableStreamRewrite) as info: - await _run_streaming_step(["hello [MASKED]"], _TextTranslation()) + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + result = await _run_streaming_step(_TextTranslation(), chunks) - assert info.value.guardrail_name == "masker" + _assert_passed_with_discard_warning(result, caplog) + assert chunks == [_chunk()] + + +class _TextAndToolCallRewritingGuardrail(CustomGuardrail): + def __init__(self, rewrite_tool_call): + super().__init__(guardrail_name="masker", event_hook="post_call", default_on=True) + self.rewrite_tool_call = rewrite_tool_call + + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + tool_calls = ( + [{"function": {"name": "lookup", "arguments": '{"ssn": "[MASKED]"}'}}] + if self.rewrite_tool_call + else inputs["tool_calls"] + ) + return {**inputs, "texts": ["hello [MASKED]"], "tool_calls": tool_calls} + + +@pytest.mark.asyncio +async def test_streaming_step_delivers_text_rewrite_through_writing_translation(monkeypatch, caplog): + monkeypatch.setattr(litellm, "callbacks", [_TextAndToolCallRewritingGuardrail(rewrite_tool_call=False)]) + chunks = [_chunk()] + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + result = await _run_streaming_step(_WritingTranslation(), chunks) + + assert result.terminal_action == "allow" + assert chunks[0]["text"] == "hello [MASKED]" + assert chunks[0]["tool_call"]["function"]["arguments"] == '{"ssn": "123"}' + assert not any("discarded" in record.getMessage() for record in caplog.records) + + +@pytest.mark.asyncio +async def test_streaming_step_discards_tool_call_rewrite_and_restores_written_text(monkeypatch, caplog): + monkeypatch.setattr(litellm, "callbacks", [_TextAndToolCallRewritingGuardrail(rewrite_tool_call=True)]) + chunks = [_chunk()] + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + result = await _run_streaming_step(_WritingTranslation(), chunks) + + _assert_passed_with_discard_warning(result, caplog) + assert chunks == [_chunk()] + + +@pytest.mark.asyncio +async def test_streaming_step_restores_chunks_when_translation_refuses_the_rewrite(monkeypatch, caplog): + monkeypatch.setattr(litellm, "callbacks", [_TextReturningGuardrail(["hello [MASKED]"])]) + chunks = [_chunk()] + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + result = await _run_streaming_step(_RefusingTranslation(), chunks) + + _assert_passed_with_discard_warning(result, caplog) + assert chunks == [_chunk()] diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py index 2e73bebb07c..ad2b4a0efbb 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -11,6 +11,7 @@ from __future__ import annotations import asyncio import json +import logging from typing import Any, Callable, Dict, List from unittest.mock import AsyncMock, MagicMock, patch @@ -24,9 +25,9 @@ from litellm.integrations.custom_guardrail import ( ModifyResponseException, ) from litellm.integrations.prometheus import PrometheusLogger -from litellm.proxy._types import ProxyException, UserAPIKeyAuth +from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.callback_utils import add_guardrail_to_applied_guardrails_header -from litellm.proxy.utils import ProxyLogging, _raise_for_streaming_post_call_pipelines +from litellm.proxy.utils import ProxyLogging, _streamable_post_call_pipelines from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ContentFilterGuardrail from litellm.types.guardrails import BlockedWord, ContentFilterAction, GuardrailEventHooks from litellm.types.proxy.policy_engine.pipeline_types import ( @@ -1384,76 +1385,87 @@ async def test_pre_call_pipeline_managed_parallel_guardrail_runs_exactly_once( assert seen["count"] == 1 +def _warnings(caplog: pytest.LogCaptureFixture) -> List[str]: + return [record.getMessage() for record in caplog.records if record.levelno >= logging.WARNING] + + @pytest.mark.asyncio -async def test_pre_call_hook_rejects_streaming_request_with_post_call_pipeline( - proxy_logging, make_user_api_key_auth, monkeypatch +async def test_streaming_request_whose_pipeline_guardrail_is_missing_streams_verbatim( + proxy_logging, make_user_api_key_auth, monkeypatch, caplog ): monkeypatch.setattr(litellm, "callbacks", []) data = _post_call_pipeline_data(stream=True) + chunks = _stream_chunks() + delivered: List[Any] = [] - with pytest.raises(HTTPException) as info: - await proxy_logging.pre_call_hook( + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + out = await proxy_logging.pre_call_hook( user_api_key_dict=make_user_api_key_auth(), data=data, call_type="completion", guardrails_only=True, ) + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"), + response=_async_chunk_iter(chunks), + request_data=data, + ): + delivered.append(item) - assert info.value.status_code == 400 - assert info.value.detail["error"]["policies"] == ("response-governance",) - assert info.value.detail["error"]["guardrails"] == ("gr-post",) - assert "stream=false" in info.value.detail["error"]["message"] + assert out is not None + assert out.get("stream") is True + assert [item is chunk for item, chunk in zip(delivered, chunks)] == [True, True] + assert len(delivered) == 2 + assert any("response-governance" in message and "gr-post" in message for message in _warnings(caplog)) @pytest.mark.asyncio -async def test_pre_call_hook_rejects_background_request_with_post_call_pipeline( - proxy_logging, make_user_api_key_auth, monkeypatch +async def test_pre_call_hook_accepts_background_request_with_post_call_pipeline( + proxy_logging, make_user_api_key_auth, monkeypatch, caplog ): monkeypatch.setattr(litellm, "callbacks", []) data = _post_call_pipeline_data(background=True) - with pytest.raises(HTTPException) as info: - await proxy_logging.pre_call_hook( + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + out = await proxy_logging.pre_call_hook( user_api_key_dict=make_user_api_key_auth(), data=data, call_type="aresponses", guardrails_only=True, ) - assert info.value.status_code == 400 - assert info.value.detail["error"]["policies"] == ("response-governance",) - assert "background=false" in info.value.detail["error"]["message"] + assert out is not None + assert out.get("background") is True + assert any("response-governance" in message and "background" in message for message in _warnings(caplog)) -def test_raise_for_streaming_post_call_pipelines_ignores_non_streaming_and_pre_call(make_user_api_key_auth): - post_call = GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="g", on_fail="block")]) - pre_call = GuardrailPipeline(mode="pre_call", steps=[PipelineStep(guardrail="g", on_fail="block")]) - auth = make_user_api_key_auth(request_route="/custom/stream") +@pytest.mark.asyncio +async def test_pre_call_hook_stays_quiet_on_background_request_without_post_call_pipeline( + proxy_logging, make_user_api_key_auth, monkeypatch, caplog +): + seen: Dict[str, Any] = {} + monkeypatch.setattr(litellm, "callbacks", [_unified_stream_guardrail(seen)]) + pre_call = GuardrailPipeline(mode="pre_call", steps=[PipelineStep(guardrail="gr-post", on_fail="block")]) + data = { + "model": "m", + "messages": [{"role": "user", "content": "hi"}], + "background": True, + "metadata": { + "_guardrail_pipelines": [("request-governance", pre_call)], + "_pipeline_managed_guardrails": {"gr-post"}, + }, + } - assert ( - _raise_for_streaming_post_call_pipelines( - {"stream": False, "metadata": {"_guardrail_pipelines": [("p", post_call)]}}, auth + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + out = await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data=data, + call_type="aresponses", + guardrails_only=True, ) - is None - ) - assert ( - _raise_for_streaming_post_call_pipelines( - {"background": False, "metadata": {"_guardrail_pipelines": [("p", post_call)]}}, auth - ) - is None - ) - assert ( - _raise_for_streaming_post_call_pipelines({"metadata": {"_guardrail_pipelines": [("p", post_call)]}}, auth) - is None - ) - assert ( - _raise_for_streaming_post_call_pipelines( - {"stream": True, "metadata": {"_guardrail_pipelines": [("p", pre_call)]}}, auth - ) - is None - ) - assert _raise_for_streaming_post_call_pipelines({"stream": True}, auth) is None - assert _raise_for_streaming_post_call_pipelines({"background": True}, auth) is None + + assert out is not None + assert not any("background" in message for message in _warnings(caplog)) # --------------------------------------------------------------------------- @@ -1485,6 +1497,56 @@ async def _async_chunk_iter(chunks: List[Any]): yield chunk +def test_streamable_post_call_pipelines_keeps_supported_and_drops_unsupported( + make_user_api_key_auth, monkeypatch, caplog +): + class NativeOnlyGuardrail(CustomGuardrail): + pass + + supported = _unified_stream_guardrail({}) + native_only = NativeOnlyGuardrail(guardrail_name="gr-native", event_hook=GuardrailEventHooks.post_call) + monkeypatch.setattr(litellm, "callbacks", [supported, native_only]) + governed = GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="gr-post", on_fail="block")]) + ungoverned = GuardrailPipeline( + mode="post_call", + steps=[PipelineStep(guardrail="gr-post", on_fail="next"), PipelineStep(guardrail="gr-native", on_fail="block")], + ) + pre_call = GuardrailPipeline(mode="pre_call", steps=[PipelineStep(guardrail="gr-native", on_fail="block")]) + data = {"metadata": {"_guardrail_pipelines": [("governed", governed), ("ungoverned", ungoverned), ("req", pre_call)]}} + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + streamable = _streamable_post_call_pipelines(data, make_user_api_key_auth(request_route="/v1/chat/completions")) + + assert streamable == (("governed", governed),) + assert any("'ungoverned'" in message and "gr-native" in message for message in _warnings(caplog)) + assert not any("'governed'" in message for message in _warnings(caplog)) + + +def test_streamable_post_call_pipelines_is_empty_on_route_without_translation( + make_user_api_key_auth, monkeypatch, caplog +): + monkeypatch.setattr(litellm, "callbacks", [_unified_stream_guardrail({})]) + governed = GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="gr-post", on_fail="block")]) + data = {"metadata": {"_guardrail_pipelines": [("governed", governed)]}} + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + streamable = _streamable_post_call_pipelines(data, make_user_api_key_auth(request_route="/custom/stream")) + + assert streamable == () + assert any("/custom/stream" in message and "governed" in message for message in _warnings(caplog)) + + +def test_streamable_post_call_pipelines_is_empty_without_post_call_pipelines(make_user_api_key_auth, caplog): + pre_call = GuardrailPipeline(mode="pre_call", steps=[PipelineStep(guardrail="g", on_fail="block")]) + auth = make_user_api_key_auth(request_route="/custom/stream") + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + assert _streamable_post_call_pipelines({"metadata": {"_guardrail_pipelines": [("p", pre_call)]}}, auth) == () + assert _streamable_post_call_pipelines({"stream": True}, auth) == () + + assert _warnings(caplog) == [] + + @pytest.mark.asyncio @pytest.mark.parametrize("request_route", [None, "/v1/chat/completions"]) async def test_pre_call_hook_allows_streaming_when_pipeline_guardrail_supports_unified( @@ -1507,21 +1569,25 @@ async def test_pre_call_hook_allows_streaming_when_pipeline_guardrail_supports_u @pytest.mark.asyncio @pytest.mark.parametrize("native_lifecycle", [False, True]) -async def test_pre_call_hook_rejects_streaming_when_pipeline_guardrail_lacks_unified_support( - proxy_logging, make_user_api_key_auth, monkeypatch, native_lifecycle +async def test_streaming_iterator_hook_releases_stream_when_pipeline_guardrail_lacks_unified_support( + proxy_logging, make_user_api_key_auth, monkeypatch, native_lifecycle, caplog ): + seen: Dict[str, Any] = {} if native_lifecycle: class NativeOnlyGuardrail(CustomGuardrail): use_native_lifecycle_hooks = True async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + seen["count"] = seen.get("count", 0) + 1 return inputs else: class NativeOnlyGuardrail(CustomGuardrail): - pass + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + seen["count"] = seen.get("count", 0) + 1 + return response monkeypatch.setattr( litellm, @@ -1529,18 +1595,29 @@ async def test_pre_call_hook_rejects_streaming_when_pipeline_guardrail_lacks_uni [NativeOnlyGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False)], ) data = _post_call_pipeline_data(stream=True) + chunks = _stream_chunks() + delivered: List[Any] = [] - with pytest.raises(HTTPException) as info: - await proxy_logging.pre_call_hook( + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + out = await proxy_logging.pre_call_hook( user_api_key_dict=make_user_api_key_auth(), data=data, call_type="completion", guardrails_only=True, ) + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"), + response=_async_chunk_iter(chunks), + request_data=data, + ): + delivered.append(item) - assert info.value.status_code == 400 - assert info.value.detail["error"]["guardrails"] == ("gr-post",) - assert "apply_guardrail" in info.value.detail["error"]["message"] + assert out is not None + assert out.get("stream") is True + assert [item is chunk for item, chunk in zip(delivered, chunks)] == [True, True] + assert len(delivered) == 2 + assert seen.get("count") is None + assert any("'response-governance'" in message and "gr-post" in message for message in _warnings(caplog)) @pytest.mark.asyncio @@ -1614,25 +1691,34 @@ async def test_pre_call_hook_allows_streaming_when_content_filter_category_masks @pytest.mark.asyncio -async def test_pre_call_hook_rejects_streaming_when_route_has_no_guardrail_translation( - proxy_logging, make_user_api_key_auth, monkeypatch +async def test_streaming_iterator_hook_releases_stream_when_route_has_no_guardrail_translation( + proxy_logging, make_user_api_key_auth, monkeypatch, caplog ): seen: Dict[str, Any] = {} monkeypatch.setattr(litellm, "callbacks", [_unified_stream_guardrail(seen)]) data = _post_call_pipeline_data(stream=True) + chunks = _stream_chunks() + delivered: List[Any] = [] - with pytest.raises(HTTPException) as info: - await proxy_logging.pre_call_hook( + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + out = await proxy_logging.pre_call_hook( user_api_key_dict=make_user_api_key_auth(request_route="/custom/stream"), data=data, call_type="completion", guardrails_only=True, ) + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/custom/stream"), + response=_async_chunk_iter(chunks), + request_data=data, + ): + delivered.append(item) - assert info.value.status_code == 400 - assert info.value.detail["error"]["policies"] == ("response-governance",) - assert "/custom/stream" in info.value.detail["error"]["message"] + assert out is not None + assert [item is chunk for item, chunk in zip(delivered, chunks)] == [True, True] + assert len(delivered) == 2 assert seen.get("count") is None + assert any("/custom/stream" in message and "response-governance" in message for message in _warnings(caplog)) @pytest.mark.asyncio @@ -1714,8 +1800,8 @@ def _echoed_tool_call_dicts(arguments: str) -> List[Dict[str, Any]]: @pytest.mark.asyncio @pytest.mark.parametrize("on_fail, on_error", [("block", None), ("next", "next")]) -async def test_streaming_iterator_hook_pipeline_withholds_runtime_tool_call_rewrite( - proxy_logging, make_user_api_key_auth, monkeypatch, on_fail, on_error +async def test_streaming_iterator_hook_pipeline_releases_originals_on_runtime_tool_call_rewrite( + proxy_logging, make_user_api_key_auth, monkeypatch, on_fail, on_error, caplog ): transform = lambda inputs: {"tool_calls": _echoed_tool_call_dicts('{"ssn": "[MASKED]"}')} # noqa: E731 monkeypatch.setattr(litellm, "callbacks", [_rewriting_stream_guardrail(transform)]) @@ -1724,7 +1810,7 @@ async def test_streaming_iterator_hook_pipeline_withholds_runtime_tool_call_rewr data = _post_call_pipeline_data(step=step, stream=True) delivered: List[Any] = [] - async def _drain() -> None: + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): async for item in proxy_logging.async_post_call_streaming_iterator_hook( user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"), response=_async_chunk_iter(_tool_call_stream_chunks()), @@ -1732,16 +1818,10 @@ async def test_streaming_iterator_hook_pipeline_withholds_runtime_tool_call_rewr ): delivered.append(item) - with pytest.raises(HTTPException) as info: - await _drain() - - error = info.value.detail["error"] - assert delivered == [] - assert info.value.status_code == 400 - assert error["type"] == "guardrail_pipeline_error" - assert error["policies"] == ("response-governance",) - assert error["guardrails"] == ("gr-post",) - assert "stream=false" in error["message"] + assert len(delivered) == 2 + assert delivered[0].choices[0].delta.tool_calls[0].function.arguments == '{"ssn": "123"}' + assert delivered[1].choices[0].finish_reason == "tool_calls" + assert any("'gr-post'" in message and "discarded" in message for message in _warnings(caplog)) @pytest.mark.asyncio @@ -1849,30 +1929,28 @@ async def test_streaming_iterator_hook_pipeline_releases_stream_echoed_in_anothe @pytest.mark.asyncio -async def test_streaming_iterator_hook_pipeline_withholds_unresolvable_response_shape( - proxy_logging, make_user_api_key_auth, monkeypatch +async def test_streaming_iterator_hook_pipeline_releases_originals_on_unresolvable_response_shape( + proxy_logging, make_user_api_key_auth, monkeypatch, caplog ): seen: Dict[str, Any] = {} monkeypatch.setattr(litellm, "callbacks", [_unified_stream_guardrail(seen)]) monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) data = _post_call_pipeline_data(stream=True) + chunks = [object(), object()] delivered: List[Any] = [] - async def _drain() -> None: + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): async for item in proxy_logging.async_post_call_streaming_iterator_hook( user_api_key_dict=make_user_api_key_auth(), - response=_async_chunk_iter([object(), object()]), + response=_async_chunk_iter(chunks), request_data=data, ): delivered.append(item) - with pytest.raises(ProxyException) as info: - await _drain() - - assert delivered == [] - assert info.value.code == "500" - assert "withheld" in info.value.message + assert [item is chunk for item, chunk in zip(delivered, chunks)] == [True, True] + assert len(delivered) == 2 assert seen.get("count") is None + assert any("response-governance" in message and "shape" in message for message in _warnings(caplog)) def _anthropic_sse_chunks() -> List[bytes]: @@ -1953,9 +2031,9 @@ async def test_streaming_iterator_hook_pipeline_delivers_text_rewrite_on_anthrop @pytest.mark.asyncio -async def test_pipeline_executor_withholds_text_rewrite_when_translation_lacks_write_back(monkeypatch): +async def test_pipeline_executor_discards_text_rewrite_when_translation_lacks_write_back(monkeypatch, caplog): from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation - from litellm.proxy.policy_engine.pipeline_executor import PipelineExecutor, UndeliverableStreamRewrite + from litellm.proxy.policy_engine.pipeline_executor import PipelineExecutor class NoWriteBackTranslation(BaseTranslation): async def process_input_messages(self, data, guardrail_to_apply, litellm_logging_obj): @@ -1984,45 +2062,23 @@ async def test_pipeline_executor_withholds_text_rewrite_when_translation_lacks_w transform = lambda inputs: {"texts": ["hello [MASKED]"]} # noqa: E731 monkeypatch.setattr(litellm, "callbacks", [_rewriting_stream_guardrail(transform)]) + chunks = _stream_chunks() - with pytest.raises(UndeliverableStreamRewrite): - await PipelineExecutor.execute_steps( + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + result = await PipelineExecutor.execute_steps( steps=[PipelineStep(guardrail="gr-post", on_pass="allow", on_fail="block")], mode="post_call", data={"metadata": {}}, user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), call_type="acompletion", policy_name="response-governance", - streaming_chunks=_stream_chunks(), + streaming_chunks=chunks, endpoint_translation=NoWriteBackTranslation(), ) - -@pytest.mark.asyncio -async def test_streaming_iterator_hook_pipeline_gates_without_iterator_overrides( - proxy_logging, make_user_api_key_auth, monkeypatch -): - monkeypatch.setattr(litellm, "callbacks", []) - monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) - data = _post_call_pipeline_data(stream=True) - delivered: List[Any] = [] - - async def _drain() -> None: - async for item in proxy_logging.async_post_call_streaming_iterator_hook( - user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"), - response=_async_chunk_iter(_stream_chunks()), - request_data=data, - ): - delivered.append(item) - - with pytest.raises(HTTPException) as info: - await _drain() - - assert delivered == [] - assert info.value.status_code == 400 - assert info.value.detail["error"]["pipeline_context"]["step_results"] == [ - {"guardrail": "gr-post", "outcome": "error", "action": "block"} - ] + assert result.terminal_action == "allow" + assert [chunk.choices[0].delta.content for chunk in chunks] == ["hello ", "world"] + assert any("'gr-post'" in message and "discarded" in message for message in _warnings(caplog)) @pytest.mark.asyncio From 2f397fa12812afba555dbdeae407597e95968123 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 7 Sep 2026 18:06:29 -0700 Subject: [PATCH 36/60] fix(drop_params): honor string flags in litellm_settings and responses, and fail open on non-flag values --- litellm/proxy/proxy_server.py | 3 ++ litellm/responses/main.py | 5 ++- litellm/types/router.py | 6 ++- .../proxy/proxy_server/test_proxy_config.py | 43 +++++++++++++++++++ .../test_responses_api_request_body.py | 4 +- tests/test_litellm/types/test_router.py | 7 ++- 6 files changed, 59 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 0915b8dd1b9..76acd414976 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -279,6 +279,7 @@ from litellm.litellm_core_utils.audio_utils.utils import resolve_speech_media_ty from litellm.litellm_core_utils.core_helpers import ( _get_parent_otel_span_from_kwargs, get_litellm_metadata_from_kwargs, + normalize_drop_params, ) from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -5508,6 +5509,8 @@ class ProxyConfig: parse_budget_reset_time(value) setattr(litellm, key, value) + elif key == "drop_params": + litellm.drop_params = bool(normalize_drop_params(value)) else: verbose_proxy_logger.debug( "%s setting litellm.%s=%s%s", diff --git a/litellm/responses/main.py b/litellm/responses/main.py index 5e74b7324b4..52ca6ebb07a 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -17,6 +17,7 @@ from litellm.completion_extras.litellm_responses_transformation.transformation i from litellm.constants import request_timeout from litellm.integrations.anthropic_cache_control_hook import CARRY_UNMATCHED_MESSAGE_POINTS from litellm.litellm_core_utils.asyncify import run_async_function +from litellm.litellm_core_utils.core_helpers import normalize_drop_params from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.prompt_templates.common_utils import ( update_responses_input_with_model_file_ids, @@ -1253,7 +1254,7 @@ def responses( responses_api_provider_config=responses_api_provider_config, response_api_optional_params=response_api_optional_params, allowed_openai_params=allowed_openai_params, - drop_params=request_drop_params if isinstance(request_drop_params, bool) else None, + drop_params=normalize_drop_params(request_drop_params), ) litellm_logging_obj.update_from_kwargs( @@ -2081,7 +2082,7 @@ def compact_responses( responses_api_provider_config=responses_api_provider_config, response_api_optional_params=response_api_optional_params, allowed_openai_params=None, - drop_params=request_drop_params if isinstance(request_drop_params, bool) else None, + drop_params=normalize_drop_params(request_drop_params), ) # Pre Call logging diff --git a/litellm/types/router.py b/litellm/types/router.py index 47aea6430c3..e8aec027a5e 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -408,9 +408,11 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams): @field_validator("drop_params", mode="before") @classmethod - def coerce_drop_params(cls, value: object) -> object: + def coerce_drop_params(cls, value: object) -> bool | str | None: normalized: Final = normalize_drop_params(value) - return value if normalized is None else normalized + if normalized is not None: + return normalized + return value if isinstance(value, str) else None def __contains__(self, key) -> bool: # Define custom behavior for the 'in' operator diff --git a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py index a4be9574f89..c57cd387d13 100644 --- a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py @@ -2431,6 +2431,49 @@ def test_ProxyConfig__add_deployment_turns_stored_drop_params_string_into_bool(m assert deployment.litellm_params.drop_params is True +def test_ProxyConfig__add_deployment_keeps_loading_rows_after_a_non_flag_drop_params(monkeypatch): + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-1234") + fake_router = MagicMock() + fake_router.upsert_deployment = MagicMock(return_value=True) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", fake_router) + pc = ProxyConfig() + + def db_model(model_id, drop_params): + return SimpleNamespace( + model_id=model_id, + model_name="gpt-5-nano", + model_info={"id": model_id}, + litellm_params={ + "model": encrypt_value_helper(value="openai/gpt-5-nano"), + "drop_params": encrypt_value_helper(value=drop_params), + }, + blocked=False, + ) + + added = pc._add_deployment(db_models=[db_model("bad-row", 2), db_model("good-after", "true")]) + deployments = [call.kwargs["deployment"] for call in fake_router.upsert_deployment.call_args_list] + + assert added == 2 + assert [d.litellm_params.drop_params for d in deployments] == [None, True] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("configured, expected", [("true", True), ("false", False)]) +async def test_ProxyConfig_load_config_turns_litellm_settings_drop_params_string_into_bool( + tmp_path, monkeypatch, configured, expected +): + f = tmp_path / "c.yaml" + f.write_text(f'model_list: []\nlitellm_settings:\n drop_params: "{configured}"\n') + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False) + monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False) + monkeypatch.setattr(litellm, "drop_params", not expected) + + await ProxyConfig().load_config(router=None, config_file_path=str(f)) + + assert litellm.drop_params is expected + + # --------------------------------------------------------------------------- # ProxyConfig.decrypt_model_list_from_db # --------------------------------------------------------------------------- diff --git a/tests/test_litellm/responses/test_responses_api_request_body.py b/tests/test_litellm/responses/test_responses_api_request_body.py index 3e60906ec6d..5fced458208 100644 --- a/tests/test_litellm/responses/test_responses_api_request_body.py +++ b/tests/test_litellm/responses/test_responses_api_request_body.py @@ -246,8 +246,10 @@ async def test_aresponses_keeps_include_obfuscation_in_stream_options(): @pytest.mark.asyncio +@pytest.mark.parametrize("drop_params", [True, "true"]) async def test_aresponses_request_level_drop_params_drops_bedrock_mantle_service_tier( monkeypatch, + drop_params, ): """ Request-level drop_params=True (as the proxy injects for agentic CLIs) must @@ -271,7 +273,7 @@ async def test_aresponses_request_level_drop_params_drops_bedrock_mantle_service aws_region_name="us-east-1", input="hi", service_tier="priority", - drop_params=True, + drop_params=drop_params, ) mock_post.assert_called_once() diff --git a/tests/test_litellm/types/test_router.py b/tests/test_litellm/types/test_router.py index 2f4a29473c3..47fc08167e3 100644 --- a/tests/test_litellm/types/test_router.py +++ b/tests/test_litellm/types/test_router.py @@ -1,5 +1,4 @@ import pytest -from pydantic import ValidationError from litellm.types.router import ( SPECIAL_MODEL_INFO_PARAMS, @@ -109,6 +108,6 @@ def test_drop_params_coerces_flags_and_keeps_unresolved_strings(value, expected) assert GenericLiteLLMParams(drop_params=value).drop_params == expected -def test_drop_params_rejects_non_flag_non_string_values(): - with pytest.raises(ValidationError): - GenericLiteLLMParams(drop_params=2) +@pytest.mark.parametrize("value", [2, 2.5, [], {}]) +def test_drop_params_ignores_non_flag_non_string_values(value): + assert GenericLiteLLMParams(drop_params=value).drop_params is None From b7c2decb7db01b54463d94112f7943032d9309da Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 7 Sep 2026 19:00:47 -0700 Subject: [PATCH 37/60] fix(drop_params): honor string values in litellm_params and the LITELLM_DROP_PARAMS env var get_litellm_params normalizes drop_params once, so a client-body string and router_settings.default_litellm_params reach the anthropic, bedrock, and azure_ai gates as a bool. LITELLM_DROP_PARAMS=false now means off. A value that is neither a flag nor a string logs one warning and counts as unset, both in the deployment validator and in litellm_settings. --- litellm/__init__.py | 3 +- litellm/litellm_core_utils/core_helpers.py | 2 +- .../litellm_core_utils/get_litellm_params.py | 5 +-- litellm/proxy/proxy_server.py | 9 ++++- litellm/types/router.py | 7 +++- .../test_get_litellm_params.py | 8 +++++ .../proxy/proxy_server/test_proxy_config.py | 34 +++++++++++++++++++ .../test_litellm/test_drop_params_env_var.py | 17 ++++++++++ tests/test_litellm/types/test_router.py | 15 ++++++-- 9 files changed, 92 insertions(+), 8 deletions(-) create mode 100644 tests/test_litellm/test_drop_params_env_var.py diff --git a/litellm/__init__.py b/litellm/__init__.py index 62477dd6264..36f143376ff 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -47,6 +47,7 @@ from typing import ( ) from litellm.types.integrations.datadog import DatadogInitParams from litellm.types.integrations.newrelic import NewRelicInitParams +from litellm.litellm_core_utils.core_helpers import normalize_drop_params from litellm._logging import ( set_verbose, _turn_on_debug, @@ -238,7 +239,7 @@ token: Optional[str] = ( ) telemetry = True max_tokens: int = DEFAULT_MAX_TOKENS # OpenAI Defaults -drop_params = bool(os.getenv("LITELLM_DROP_PARAMS", False)) +drop_params = bool(normalize_drop_params(os.getenv("LITELLM_DROP_PARAMS"))) modify_params = bool(os.getenv("LITELLM_MODIFY_PARAMS", False)) use_chat_completions_url_for_anthropic_messages: bool = bool( os.getenv("LITELLM_USE_CHAT_COMPLETIONS_URL_FOR_ANTHROPIC_MESSAGES", False) diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index bd7a1b8f384..c1f1076c710 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -42,7 +42,7 @@ _DROP_PARAMS_BOOL: Final = TypeAdapter(bool) def normalize_drop_params(value: object) -> bool | None: - if isinstance(value, bool): + if value is None or isinstance(value, bool): return value try: return _DROP_PARAMS_BOOL.validate_python(value.strip() if isinstance(value, str) else value) diff --git a/litellm/litellm_core_utils/get_litellm_params.py b/litellm/litellm_core_utils/get_litellm_params.py index 1fd79db15a6..92b32d32dc0 100644 --- a/litellm/litellm_core_utils/get_litellm_params.py +++ b/litellm/litellm_core_utils/get_litellm_params.py @@ -2,6 +2,7 @@ from collections.abc import Mapping, MutableMapping from types import MappingProxyType from typing import Final +from litellm.litellm_core_utils.core_helpers import normalize_drop_params from litellm.llms.openai.data_residency import infer_openai_data_residency AWS_CREDENTIAL_KWARGS_KEYS: Final = frozenset( @@ -113,7 +114,7 @@ def get_litellm_params( custom_prompt_dict: dict | None = None, litellm_metadata: dict | None = None, disable_add_transform_inline_image_block: bool | None = None, - drop_params: bool | None = None, + drop_params: bool | str | None = None, prompt_id: str | None = None, prompt_variables: dict | None = None, async_call: bool | None = None, @@ -175,7 +176,7 @@ def get_litellm_params( "custom_prompt_dict": custom_prompt_dict, "litellm_metadata": litellm_metadata, "disable_add_transform_inline_image_block": disable_add_transform_inline_image_block, - "drop_params": drop_params, + "drop_params": normalize_drop_params(drop_params), "prompt_id": prompt_id, "prompt_variables": prompt_variables, "async_call": async_call, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 76acd414976..8a30274cc9f 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -5510,7 +5510,7 @@ class ProxyConfig: parse_budget_reset_time(value) setattr(litellm, key, value) elif key == "drop_params": - litellm.drop_params = bool(normalize_drop_params(value)) + litellm.drop_params = _drop_params_from_litellm_settings(value) else: verbose_proxy_logger.debug( "%s setting litellm.%s=%s%s", @@ -16915,6 +16915,13 @@ def _redact_config_param_value_for_logging(param_name: str | None, param_value: return param_value +def _drop_params_from_litellm_settings(value: object) -> bool: + normalized: Final = normalize_drop_params(value) + if normalized is None and value is not None: + verbose_proxy_logger.warning("litellm_settings.drop_params=%r is not a flag value, treating it as off", value) + return bool(normalized) + + def _redact_general_setting_value(field_name: str, value: JsonValue, is_full_admin: bool) -> JsonValue: if is_full_admin: return value diff --git a/litellm/types/router.py b/litellm/types/router.py index e8aec027a5e..5c9eab30f3d 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -12,6 +12,7 @@ import httpx from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator from typing_extensions import Protocol, ReadOnly, Required, TypedDict, runtime_checkable +from litellm._logging import verbose_logger from litellm._uuid import uuid from litellm.litellm_core_utils.core_helpers import normalize_drop_params @@ -412,7 +413,11 @@ class GenericLiteLLMParams(CredentialLiteLLMParams, CustomPricingLiteLLMParams): normalized: Final = normalize_drop_params(value) if normalized is not None: return normalized - return value if isinstance(value, str) else None + if isinstance(value, str): + return value + if value is not None: + verbose_logger.warning("drop_params=%r is not a flag value, treating it as unset", value) + return None def __contains__(self, key) -> bool: # Define custom behavior for the 'in' operator diff --git a/tests/test_litellm/litellm_core_utils/test_get_litellm_params.py b/tests/test_litellm/litellm_core_utils/test_get_litellm_params.py index fb4cb494bee..f026ff57719 100644 --- a/tests/test_litellm/litellm_core_utils/test_get_litellm_params.py +++ b/tests/test_litellm/litellm_core_utils/test_get_litellm_params.py @@ -215,3 +215,11 @@ class TestMetadataFallsBackToLitellmMetadata: assert result["metadata"] is not litellm_metadata result["metadata"].pop("trace_id") assert litellm_metadata == {"trace_id": "trace-1"} + + +@pytest.mark.parametrize( + "value, expected", + [("true", True), ("false", False), (" TRUE ", True), (True, True), (None, None), ("os.environ/DROP_PARAMS", None)], +) +def test_drop_params_strings_reach_litellm_params_as_flags(value, expected): + assert get_litellm_params(drop_params=value)["drop_params"] is expected diff --git a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py index c57cd387d13..774c63d754a 100644 --- a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py @@ -9,6 +9,7 @@ Pins covered: from __future__ import annotations import json +import logging import os import re from types import SimpleNamespace @@ -2474,6 +2475,39 @@ async def test_ProxyConfig_load_config_turns_litellm_settings_drop_params_string assert litellm.drop_params is expected +@pytest.mark.asyncio +async def test_ProxyConfig_load_config_resolves_a_litellm_settings_drop_params_env_ref(tmp_path, monkeypatch): + f = tmp_path / "c.yaml" + f.write_text("model_list: []\nlitellm_settings:\n drop_params: os.environ/DROP_PARAMS_FROM_ENV\n") + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False) + monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False) + monkeypatch.setenv("DROP_PARAMS_FROM_ENV", "true") + monkeypatch.setattr(litellm, "drop_params", False) + + await ProxyConfig().load_config(router=None, config_file_path=str(f)) + + assert litellm.drop_params is True + + +@pytest.mark.asyncio +async def test_ProxyConfig_load_config_warns_and_turns_off_a_non_flag_litellm_settings_drop_params( + tmp_path, monkeypatch, caplog +): + f = tmp_path / "c.yaml" + f.write_text("model_list: []\nlitellm_settings:\n drop_params: ture\n") + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False) + monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False) + monkeypatch.setattr(litellm, "drop_params", True) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + await ProxyConfig().load_config(router=None, config_file_path=str(f)) + + assert litellm.drop_params is False + assert "litellm_settings.drop_params='ture' is not a flag value, treating it as off" in caplog.text + + # --------------------------------------------------------------------------- # ProxyConfig.decrypt_model_list_from_db # --------------------------------------------------------------------------- diff --git a/tests/test_litellm/test_drop_params_env_var.py b/tests/test_litellm/test_drop_params_env_var.py new file mode 100644 index 00000000000..a1ef3648f95 --- /dev/null +++ b/tests/test_litellm/test_drop_params_env_var.py @@ -0,0 +1,17 @@ +import os +import subprocess +import sys + +import pytest + + +@pytest.mark.parametrize("configured, expected", [("false", "False"), ("true", "True")]) +def test_litellm_drop_params_env_var_is_parsed_as_a_flag(configured, expected): + result = subprocess.run( + [sys.executable, "-c", "import litellm; print(litellm.drop_params)"], + env={**os.environ, "LITELLM_DROP_PARAMS": configured}, + capture_output=True, + text=True, + check=True, + ) + assert result.stdout.strip() == expected diff --git a/tests/test_litellm/types/test_router.py b/tests/test_litellm/types/test_router.py index 47fc08167e3..fd933a9d993 100644 --- a/tests/test_litellm/types/test_router.py +++ b/tests/test_litellm/types/test_router.py @@ -1,3 +1,5 @@ +import logging + import pytest from litellm.types.router import ( @@ -109,5 +111,14 @@ def test_drop_params_coerces_flags_and_keeps_unresolved_strings(value, expected) @pytest.mark.parametrize("value", [2, 2.5, [], {}]) -def test_drop_params_ignores_non_flag_non_string_values(value): - assert GenericLiteLLMParams(drop_params=value).drop_params is None +def test_drop_params_ignores_non_flag_non_string_values_with_a_warning(value, caplog): + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + assert GenericLiteLLMParams(drop_params=value).drop_params is None + assert f"drop_params={value!r} is not a flag value" in caplog.text + + +@pytest.mark.parametrize("value", [True, "true", None, "os.environ/DROP_PARAMS", "v2:gcm:ciphertext-from-a-pre-fix-row"]) +def test_drop_params_flags_and_strings_log_nothing(value, caplog): + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + GenericLiteLLMParams(drop_params=value) + assert caplog.text == "" From d594b9385eb8eae760c809e43ac9174bceadfd46 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 7 Sep 2026 21:05:01 -0700 Subject: [PATCH 38/60] fix(drop_params): warn when a deployment or env drop_params value is not a flag A deployment drop_params string that is not a flag value (a typo like ture) stayed silently off. The router now logs one warning per deployment. LITELLM_DROP_PARAMS and litellm_settings.drop_params share the same helper, so a non-flag value there warns as well instead of flipping silently from on to off --- litellm/__init__.py | 4 +-- litellm/litellm_core_utils/core_helpers.py | 8 +++++ litellm/proxy/proxy_server.py | 11 ++----- litellm/router.py | 6 ++++ .../litellm_core_utils/test_core_helpers.py | 17 ++++++++++ .../test_litellm/test_drop_params_env_var.py | 19 +++++++++-- tests/test_litellm/test_router.py | 33 +++++++++++++++++++ 7 files changed, 84 insertions(+), 14 deletions(-) diff --git a/litellm/__init__.py b/litellm/__init__.py index 36f143376ff..f7d4dce87d2 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -47,7 +47,7 @@ from typing import ( ) from litellm.types.integrations.datadog import DatadogInitParams from litellm.types.integrations.newrelic import NewRelicInitParams -from litellm.litellm_core_utils.core_helpers import normalize_drop_params +from litellm.litellm_core_utils.core_helpers import drop_params_flag from litellm._logging import ( set_verbose, _turn_on_debug, @@ -239,7 +239,7 @@ token: Optional[str] = ( ) telemetry = True max_tokens: int = DEFAULT_MAX_TOKENS # OpenAI Defaults -drop_params = bool(normalize_drop_params(os.getenv("LITELLM_DROP_PARAMS"))) +drop_params = drop_params_flag(os.getenv("LITELLM_DROP_PARAMS"), "LITELLM_DROP_PARAMS", verbose_logger) modify_params = bool(os.getenv("LITELLM_MODIFY_PARAMS", False)) use_chat_completions_url_for_anthropic_messages: bool = bool( os.getenv("LITELLM_USE_CHAT_COMPLETIONS_URL_FOR_ANTHROPIC_MESSAGES", False) diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index c1f1076c710..a3dfac81cc1 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -1,6 +1,7 @@ # What is this? ## Helper utilities import copy +import logging from collections.abc import Iterable, Mapping from typing import TYPE_CHECKING, Any, Final, Literal @@ -50,6 +51,13 @@ def normalize_drop_params(value: object) -> bool | None: return None +def drop_params_flag(value: object, source: str, logger: logging.Logger) -> bool: + normalized: Final = normalize_drop_params(value) + if normalized is None and value is not None: + logger.warning("%s=%r is not a flag value, treating it as off", source, value) + return bool(normalized) + + def safe_divide( numerator: float, denominator: float, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 8a30274cc9f..83b99822b4c 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -278,8 +278,8 @@ from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.audio_utils.utils import resolve_speech_media_type from litellm.litellm_core_utils.core_helpers import ( _get_parent_otel_span_from_kwargs, + drop_params_flag, get_litellm_metadata_from_kwargs, - normalize_drop_params, ) from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -5510,7 +5510,7 @@ class ProxyConfig: parse_budget_reset_time(value) setattr(litellm, key, value) elif key == "drop_params": - litellm.drop_params = _drop_params_from_litellm_settings(value) + litellm.drop_params = drop_params_flag(value, "litellm_settings.drop_params", verbose_proxy_logger) else: verbose_proxy_logger.debug( "%s setting litellm.%s=%s%s", @@ -16915,13 +16915,6 @@ def _redact_config_param_value_for_logging(param_name: str | None, param_value: return param_value -def _drop_params_from_litellm_settings(value: object) -> bool: - normalized: Final = normalize_drop_params(value) - if normalized is None and value is not None: - verbose_proxy_logger.warning("litellm_settings.drop_params=%r is not a flag value, treating it as off", value) - return bool(normalized) - - def _redact_general_setting_value(field_name: str, value: JsonValue, is_full_admin: bool) -> JsonValue: if is_full_admin: return value diff --git a/litellm/router.py b/litellm/router.py index 95cabfad4bd..4bbdba08e13 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -9366,6 +9366,12 @@ class Router: #### VALIDATE MODEL ######## # Check if this is a prompt management model before validating as LLM provider litellm_model: Final = deployment.litellm_params.model + if isinstance(deployment.litellm_params.drop_params, str): + verbose_router_logger.warning( + "model=%s drop_params=%r is not a flag value, treating it as unset", + deployment.model_name, + deployment.litellm_params.drop_params, + ) is_prompt_management_model = False if "/" in litellm_model: diff --git a/tests/test_litellm/litellm_core_utils/test_core_helpers.py b/tests/test_litellm/litellm_core_utils/test_core_helpers.py index 8797f7c3591..e3880175759 100644 --- a/tests/test_litellm/litellm_core_utils/test_core_helpers.py +++ b/tests/test_litellm/litellm_core_utils/test_core_helpers.py @@ -1,9 +1,12 @@ """Tests for litellm_core_utils.core_helpers module.""" +import logging + import pytest from litellm.litellm_core_utils.core_helpers import ( _FINISH_REASON_MAP, + drop_params_flag, get_or_create_metadata_bucket, map_finish_reason, normalize_drop_params, @@ -284,6 +287,20 @@ def test_normalize_drop_params(value, expected): assert normalize_drop_params(value) is expected +@pytest.mark.parametrize("value, expected", [("true", True), ("off", False), (None, False)]) +def test_drop_params_flag_returns_a_bool_without_a_warning(value, expected, caplog): + with caplog.at_level(logging.WARNING, logger="drop-params-test"): + assert drop_params_flag(value, "LITELLM_DROP_PARAMS", logging.getLogger("drop-params-test")) is expected + assert caplog.text == "" + + +@pytest.mark.parametrize("value", ["temperature", "ture", 2]) +def test_drop_params_flag_treats_non_flag_values_as_off_with_a_warning(value, caplog): + with caplog.at_level(logging.WARNING, logger="drop-params-test"): + assert drop_params_flag(value, "LITELLM_DROP_PARAMS", logging.getLogger("drop-params-test")) is False + assert f"LITELLM_DROP_PARAMS={value!r} is not a flag value, treating it as off" in caplog.text + + class TestIsExpectedClientError: def test_status_ranges(self): from litellm.litellm_core_utils.core_helpers import is_expected_client_error diff --git a/tests/test_litellm/test_drop_params_env_var.py b/tests/test_litellm/test_drop_params_env_var.py index a1ef3648f95..339298df3d1 100644 --- a/tests/test_litellm/test_drop_params_env_var.py +++ b/tests/test_litellm/test_drop_params_env_var.py @@ -5,13 +5,26 @@ import sys import pytest -@pytest.mark.parametrize("configured, expected", [("false", "False"), ("true", "True")]) -def test_litellm_drop_params_env_var_is_parsed_as_a_flag(configured, expected): - result = subprocess.run( +def _import_litellm_with(configured: str) -> subprocess.CompletedProcess[str]: + return subprocess.run( [sys.executable, "-c", "import litellm; print(litellm.drop_params)"], env={**os.environ, "LITELLM_DROP_PARAMS": configured}, capture_output=True, text=True, check=True, ) + + +@pytest.mark.parametrize("configured, expected", [("false", "False"), ("true", "True")]) +def test_litellm_drop_params_env_var_is_parsed_as_a_flag(configured, expected): + result = _import_litellm_with(configured) + assert result.stdout.strip() == expected + assert "is not a flag value" not in result.stderr + + +def test_litellm_drop_params_env_var_non_flag_value_is_off_with_a_warning(): + result = _import_litellm_with("temperature") + + assert result.stdout.strip() == "False" + assert "LITELLM_DROP_PARAMS='temperature' is not a flag value, treating it as off" in result.stderr diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 201a4d84ebe..0c0bad8080d 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -14424,3 +14424,36 @@ async def test_router_deployment_drop_params_string_true_is_honored(monkeypatch) temperature=0.1, ) assert response.choices[0].message.content == "Hello, world!" + + +@pytest.mark.parametrize("value", ["ture", "enabled"]) +def test_router_warns_when_a_deployment_drop_params_string_is_not_a_flag(value, caplog): + with caplog.at_level(logging.WARNING, logger="LiteLLM Router"): + router = Router( + model_list=[ + { + "model_name": "gpt-5-nano", + "litellm_params": {"model": "openai/gpt-5-nano", "api_key": "sk-fake", "drop_params": value}, + } + ] + ) + + deployment = router.get_deployment_by_model_group_name(model_group_name="gpt-5-nano") + assert deployment is not None + assert deployment.litellm_params.drop_params == value + assert f"model=gpt-5-nano drop_params={value!r} is not a flag value, treating it as unset" in caplog.text + + +@pytest.mark.parametrize("value", [True, "true", "off", None]) +def test_router_stays_quiet_when_a_deployment_drop_params_is_a_flag(value, caplog): + with caplog.at_level(logging.WARNING, logger="LiteLLM Router"): + Router( + model_list=[ + { + "model_name": "gpt-5-nano", + "litellm_params": {"model": "openai/gpt-5-nano", "api_key": "sk-fake", "drop_params": value}, + } + ] + ) + + assert "is not a flag value" not in caplog.text From 69d2ac1edb83336723d4c9ce5024b93612230e5d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 7 Sep 2026 21:46:00 -0700 Subject: [PATCH 39/60] fix(policy_engine): run iterator-hook guardrails whose post_call pipeline cannot stream The streaming loop skipped every guardrail stepped by a post_call pipeline, even when the pipeline was dropped from the stream for lacking the unified apply_guardrail interface, so a default_on guardrail that only implements async_post_call_streaming_iterator_hook stopped governing streams it governed on the merge base. The skip set now comes from the pipelines that will gate the stream --- litellm/proxy/utils.py | 26 +++++++------- .../proxy_logging/test_guardrail_pipeline.py | 36 +++++++++++++++++++ 2 files changed, 50 insertions(+), 12 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index d7bf832bca4..8b36061c6be 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -447,14 +447,15 @@ def _policy_pipelines(data: Mapping[str, object]) -> tuple[tuple[str, "Guardrail ) +def _pipeline_step_guardrail_names(pipelines: Sequence[tuple[str, "GuardrailPipeline"]]) -> frozenset[str]: + return frozenset(step.guardrail for _policy_name, pipeline in pipelines for step in pipeline.steps) + + def _pipeline_managed_guardrail_names( data: Mapping[str, object], mode: Literal["pre_call", "post_call"] ) -> frozenset[str]: - return frozenset( - step.guardrail - for _policy_name, pipeline in _policy_pipelines(data) - if pipeline.mode == mode - for step in pipeline.steps + return _pipeline_step_guardrail_names( + tuple((policy_name, pipeline) for policy_name, pipeline in _policy_pipelines(data) if pipeline.mode == mode) ) @@ -547,7 +548,7 @@ def _pipeline_is_streamable(policy_name: str, pipeline: "GuardrailPipeline") -> return True verbose_proxy_logger.warning( "Policy '%s' has post_call pipeline guardrails without the unified apply_guardrail interface, " - "which streaming pipelines need; the stream is released ungoverned by it: %s", + "which streaming pipelines need; the stream skips the pipeline and its guardrails run on their own: %s", policy_name, ", ".join(unsupported), ) @@ -563,9 +564,9 @@ def _streamable_post_call_pipelines( Streaming pipelines scan the buffered stream through the endpoint guardrail translation of the request route, so every step's guardrail needs the unified apply_guardrail interface and the route needs a translation. A - pipeline that cannot be run that way yet is left out and the stream is - released the way it was before pipelines ran on streams at all, with a - warning naming what went ungoverned. + pipeline that cannot be run that way yet is left out and its guardrails + run on the stream on their own, the way they did before pipelines ran on + streams at all, with a warning naming the pipeline. """ post_call_pipelines: Final = _post_call_pipelines(request_data) if not post_call_pipelines: @@ -574,7 +575,8 @@ def _streamable_post_call_pipelines( if route and resolve_endpoint_translation(user_api_key_dict, None) is None: verbose_proxy_logger.warning( "Policies with post_call guardrail pipelines cannot scan streaming responses on route %s yet " - "(no endpoint guardrail translation); the stream is released ungoverned by them: %s", + "(no endpoint guardrail translation); the stream skips the pipelines and their guardrails run " + "on their own: %s", route, ", ".join(policy_name for policy_name, _pipeline in post_call_pipelines), ) @@ -3361,10 +3363,10 @@ class ProxyLogging: current_response = response stream_needs_translation: Final = ProxyLogging._stream_requires_guardrail_translation(user_api_key_dict) - pipeline_managed_names: Final = _pipeline_managed_guardrail_names(request_data, "post_call") + pipeline_gated_names: Final = _pipeline_step_guardrail_names(post_call_pipelines) for resolved_callback, kind in caps.iterator_overrides: if isinstance(resolved_callback, CustomGuardrail): - if resolved_callback.guardrail_name in pipeline_managed_names: + if resolved_callback.guardrail_name in pipeline_gated_names: continue if ( resolved_callback.should_run_guardrail(data=request_data, event_type=GuardrailEventHooks.post_call) diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py index ad2b4a0efbb..73dd6746b29 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -1620,6 +1620,42 @@ async def test_streaming_iterator_hook_releases_stream_when_pipeline_guardrail_l assert any("'response-governance'" in message and "gr-post" in message for message in _warnings(caplog)) +@pytest.mark.asyncio +async def test_streaming_iterator_hook_runs_iterator_hook_guardrail_whose_pipeline_cannot_stream( + proxy_logging, make_user_api_key_auth, monkeypatch, caplog +): + seen: Dict[str, Any] = {} + + class IteratorHookGuardrail(CustomGuardrail): + async def async_post_call_streaming_iterator_hook(self, user_api_key_dict, response, request_data): + seen["count"] = seen.get("count", 0) + 1 + async for item in response: + item.choices[0].delta.content = f"[governed] {item.choices[0].delta.content}" + yield item + + monkeypatch.setattr( + litellm, + "callbacks", + [IteratorHookGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=True)], + ) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + data = _post_call_pipeline_data(stream=True) + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + delivered = [ + item + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/v1/chat/completions"), + response=_async_chunk_iter(_stream_chunks()), + request_data=data, + ) + ] + + assert seen["count"] == 1 + assert [item.choices[0].delta.content for item in delivered] == ["[governed] hello ", "[governed] world"] + assert any("'response-governance'" in message and "gr-post" in message for message in _warnings(caplog)) + + @pytest.mark.asyncio @pytest.mark.parametrize( "rewrite_attribute, value", From 91fc1b201027b3d21a7ced292ac19a51661cba0e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 7 Sep 2026 21:48:19 -0700 Subject: [PATCH 40/60] fix(policy_engine): record a streaming pipeline step once and in the applied guardrails header CustomGuardrail.__init_subclass__ wrapped _StreamRewriteObserver.apply_guardrail in log_guardrail_information, so every streaming step recorded a second standard_logging_guardrail_information entry and span next to the inner guardrail's own. The observer's method now carries the marker that skips the wrapper. The step also adds the guardrail to the applied guardrails header the way the non-streaming unified path does, so streamed spend rows name the guardrail that scanned them --- .../proxy/policy_engine/pipeline_executor.py | 27 +++++++++--- .../policy_engine/test_pipeline_executor.py | 44 +++++++++++++++++++ 2 files changed, 65 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index 970ac20487c..9bc10949e9f 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -7,19 +7,21 @@ pass/fail actions (allow, block, next, modify_response) and data forwarding. import copy import time -from collections.abc import Mapping, Sequence -from typing import TYPE_CHECKING, Any, Final, Literal +from collections.abc import Callable, Mapping, Sequence +from typing import TYPE_CHECKING, Any, Final, Literal, TypeVar from pydantic import BaseModel import litellm from litellm._logging import verbose_proxy_logger +from litellm.constants import LOGS_GUARDRAIL_INFORMATION_MARKER from litellm.integrations.custom_guardrail import ( CustomGuardrail, ModifyResponseException, ) from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.core_helpers import independent_snapshot +from litellm.proxy.common_utils.callback_utils import add_guardrail_to_applied_guardrails_header from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( UnifiedLLMGuardrails, ) @@ -72,13 +74,23 @@ def _rewrote(sent: tuple[object, ...] | None, returned: tuple[object, ...] | Non return sent is not None and returned is not None and returned != sent +_GuardrailMethodT = TypeVar("_GuardrailMethodT", bound=Callable[..., object]) + + +def _logged_by_inner_guardrail(method: _GuardrailMethodT) -> _GuardrailMethodT: + vars(method)[LOGS_GUARDRAIL_INFORMATION_MARKER] = True # rebind-ok: stamps the method the class body just defined + return method + + class _StreamRewriteObserver(CustomGuardrail): """Stand-in handed to the endpoint translation in place of a streaming pipeline step's guardrail. It records whether the guardrail returned different output than it was given, which for guardrails like Bedrock's ANONYMIZED action is only known at runtime. Text rewrites are deliverable on translations that write them back across the buffered chunks (``delivers_ended_stream_text_rewrites``); tool-call rewrites and text rewrites on any - other translation are discarded by the executor, which releases the original chunks.""" + other translation are discarded by the executor, which releases the original chunks. + The inner guardrail's ``apply_guardrail`` already records the guardrail information + and span, so the observer's stays out of ``log_guardrail_information``.""" def __init__(self, inner: CustomGuardrail) -> None: super().__init__(guardrail_name=inner.guardrail_name) @@ -89,6 +101,7 @@ class _StreamRewriteObserver(CustomGuardrail): def structured_messages_cover_full_request(self) -> bool: return self.inner.structured_messages_cover_full_request() + @_logged_by_inner_guardrail async def apply_guardrail( self, inputs: GenericGuardrailAPIInputs, @@ -305,9 +318,11 @@ class PipelineExecutor: ) except UndeliverableStreamRewrite: _release_original_chunks(step.guardrail, streaming_chunks, originals) - return - if observer.rewrote_tool_calls or (observer.rewrote_texts and not deliver_rewrites): - _release_original_chunks(step.guardrail, streaming_chunks, originals) + else: + if observer.rewrote_tool_calls or (observer.rewrote_texts and not deliver_rewrites): + _release_original_chunks(step.guardrail, streaming_chunks, originals) + if not callback.records_own_guardrail_information: + add_guardrail_to_applied_guardrails_header(request_data=hook_input, guardrail_name=step.guardrail) @staticmethod async def _run_step( diff --git a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py index 0685bc6aa1e..54ef1f79f4d 100644 --- a/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py +++ b/tests/test_litellm/proxy/policy_engine/test_pipeline_executor.py @@ -1099,6 +1099,50 @@ async def test_streaming_step_discards_tool_call_rewrite_and_restores_written_te assert chunks == [_chunk()] +class _BlockingStreamGuardrail(CustomGuardrail): + def __init__(self): + super().__init__(guardrail_name="masker", event_hook="post_call", default_on=True) + + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + raise HTTPException(status_code=400, detail={"error": "output blocked"}) + + +def _recorded_guardrail_statuses(result): + return [ + entry["guardrail_status"] + for entry in result.modified_data["metadata"]["standard_logging_guardrail_information"] + ] + + +@pytest.mark.asyncio +async def test_streaming_step_records_guardrail_information_once_on_mask(monkeypatch): + monkeypatch.setattr(litellm, "callbacks", [_TextReturningGuardrail(["hello [MASKED]"])]) + + result = await _run_streaming_step(_WritingTranslation(), [_chunk()]) + + assert result.terminal_action == "allow" + assert _recorded_guardrail_statuses(result) == ["success"] + + +@pytest.mark.asyncio +async def test_streaming_step_records_the_guardrail_in_the_applied_guardrails_header(monkeypatch): + monkeypatch.setattr(litellm, "callbacks", [_TextReturningGuardrail(["hello [MASKED]"])]) + + result = await _run_streaming_step(_WritingTranslation(), [_chunk()]) + + assert result.modified_data["metadata"]["applied_guardrails"] == ["masker"] + + +@pytest.mark.asyncio +async def test_streaming_step_records_guardrail_information_once_on_block(monkeypatch): + monkeypatch.setattr(litellm, "callbacks", [_BlockingStreamGuardrail()]) + + result = await _run_streaming_step(_WritingTranslation(), [_chunk()]) + + assert [step.outcome for step in result.step_results] == ["fail"] + assert _recorded_guardrail_statuses(result) == ["guardrail_intervened"] + + @pytest.mark.asyncio async def test_streaming_step_restores_chunks_when_translation_refuses_the_rewrite(monkeypatch, caplog): monkeypatch.setattr(litellm, "callbacks", [_TextReturningGuardrail(["hello [MASKED]"])]) From d08a177bc75b45769e5e22efa2bbb4b0ea27ee86 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 7 Sep 2026 21:51:34 -0700 Subject: [PATCH 41/60] fix(policy_engine): keep a policy-added guardrail's other stages when a pipeline steps it A policy that both adds a guardrail and steps it in a post_call pipeline used to drop the guardrail from the request's guardrail list outright, so its pre_call stage never ran. The per-hook loops already skip guardrails by pipeline mode, so the mode-agnostic subtraction only lost coverage --- litellm/proxy/litellm_pre_call_utils.py | 7 +--- .../proxy/test_litellm_pre_call_utils.py | 42 +++++++++++++++++++ 2 files changed, 44 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 56512570448..e4bce4378f0 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -3079,10 +3079,9 @@ def _apply_resolved_guardrails_to_metadata( if metadata_variable_name not in data: data[metadata_variable_name] = {} - # Track pipeline-managed guardrails to exclude from independent execution - pipeline_managed_guardrails: set = set() + # Record the pipelines and the guardrails they step; the hook loops skip those per pipeline mode if pipelines: - pipeline_managed_guardrails = PolicyResolver.get_pipeline_managed_guardrails(pipelines) + pipeline_managed_guardrails: Final = PolicyResolver.get_pipeline_managed_guardrails(pipelines) data[metadata_variable_name]["_guardrail_pipelines"] = pipelines data[metadata_variable_name]["_pipeline_managed_guardrails"] = pipeline_managed_guardrails verbose_proxy_logger.debug( @@ -3099,10 +3098,8 @@ def _apply_resolved_guardrails_to_metadata( existing_guardrails = [] # Combine existing guardrails with policy-resolved guardrails (no duplicates) - # Exclude pipeline-managed guardrails from the flat list combined = set(existing_guardrails) combined.update(resolved_guardrails) - combined -= pipeline_managed_guardrails data[metadata_variable_name]["guardrails"] = list(combined) verbose_proxy_logger.debug("Policy engine: added guardrails to request metadata: %s", list(combined)) diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index 7070617ce3e..78fe2e6df88 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -4148,6 +4148,48 @@ async def test_add_guardrails_from_policy_engine(): attachment_registry._initialized = False +@pytest.mark.asyncio +async def test_add_guardrails_from_policy_engine_keeps_a_policy_added_guardrail_its_pipeline_also_steps(): + from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry + from litellm.proxy.policy_engine.policy_registry import get_policy_registry + from litellm.types.proxy.policy_engine import ( + GuardrailPipeline, + PipelineStep, + Policy, + PolicyAttachment, + PolicyGuardrails, + ) + + data = {"model": "gpt-4", "messages": [{"role": "user", "content": "Hello"}], "metadata": {}} + policy_registry = get_policy_registry() + policy_registry._policies = { + "response-governance": Policy( + guardrails=PolicyGuardrails(add=["pii_blocker"]), + pipeline=GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="pii_blocker")]), + ), + } + policy_registry._initialized = True + attachment_registry = get_attachment_registry() + attachment_registry._attachments = [PolicyAttachment(policy="response-governance", scope="*")] + attachment_registry._initialized = True + + try: + await add_guardrails_from_policy_engine( + data=data, + metadata_variable_name="metadata", + user_api_key_dict=UserAPIKeyAuth(api_key="test-key"), + ) + finally: + policy_registry._policies = {} + policy_registry._initialized = False + attachment_registry._attachments = [] + attachment_registry._initialized = False + + assert data["metadata"]["guardrails"] == ["pii_blocker"] + assert data["metadata"]["_pipeline_managed_guardrails"] == {"pii_blocker"} + assert [pipeline.mode for _policy_name, pipeline in data["metadata"]["_guardrail_pipelines"]] == ["post_call"] + + @pytest.mark.asyncio async def test_add_guardrails_from_policy_engine_accepts_dynamic_policies_and_pops_from_data(): """ From 08b60c409a0b556efb6bbe6470402a4530666870 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 7 Sep 2026 21:53:26 -0700 Subject: [PATCH 42/60] refactor(guardrails): drop the unused rewrites_streamed_output hook Nothing calls it since the streaming pipeline detects rewrites at run time through the stream observer, so the base method and the content filter's override were dead code with dead tests --- litellm/integrations/custom_guardrail.py | 3 - .../litellm_content_filter/content_filter.py | 9 --- .../content_filter/test_content_filter.py | 56 ------------------- 3 files changed, 68 deletions(-) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 883329c9fa8..2d66a280663 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -773,9 +773,6 @@ class CustomGuardrail(CustomLogger): def uses_apply_guardrail_interface(self) -> bool: return type(self).apply_guardrail is not CustomGuardrail.apply_guardrail - def rewrites_streamed_output(self) -> bool: - return self.mask_response_content - def _deployment_pre_call_target(self) -> "CustomLogger": if not self.uses_apply_guardrail_interface() or self.use_native_lifecycle_hooks: return self diff --git a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py index 85eb50c78e7..722f96ef814 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py +++ b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py @@ -1947,15 +1947,6 @@ class ContentFilterGuardrail(CustomGuardrail): exception_str=exception_str, ) - def rewrites_streamed_output(self) -> bool: - return ( - super().rewrites_streamed_output() - or any(entry["action"] == ContentFilterAction.MASK for entry in self.compiled_patterns) - or any(action == ContentFilterAction.MASK for action, _ in self.blocked_words.values()) - or any(action == ContentFilterAction.MASK for _, _, action in self.category_keywords.values()) - or any(action == ContentFilterAction.MASK for _, _, action in self.always_block_category_keywords.values()) - ) - async def async_post_call_streaming_iterator_hook( self, user_api_key_dict: UserAPIKeyAuth, diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py index 73020fe3e6f..be55ac47bde 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py @@ -3068,59 +3068,3 @@ class TestContentFilterToolCallArguments: request_data={}, input_type="response", ) - - -class TestRewritesStreamedOutput: - def test_block_only_rules_do_not_rewrite(self): - guardrail = ContentFilterGuardrail( - guardrail_name="cf", - patterns=[ContentFilterPattern(pattern_type="prebuilt", pattern_name="us_ssn", action=ContentFilterAction.BLOCK)], - blocked_words=[BlockedWord(keyword="kumquat", action=ContentFilterAction.BLOCK)], - ) - - assert guardrail.rewrites_streamed_output() is False - - def test_mask_blocked_word_rewrites(self): - guardrail = ContentFilterGuardrail( - guardrail_name="cf", - blocked_words=[BlockedWord(keyword="persimmon", action=ContentFilterAction.MASK)], - ) - - assert guardrail.rewrites_streamed_output() is True - - def test_mask_pattern_rewrites(self): - guardrail = ContentFilterGuardrail( - guardrail_name="cf", - patterns=[ContentFilterPattern(pattern_type="prebuilt", pattern_name="us_ssn", action=ContentFilterAction.MASK)], - ) - - assert guardrail.rewrites_streamed_output() is True - - def test_mask_response_content_rewrites(self): - guardrail = ContentFilterGuardrail( - guardrail_name="cf", - blocked_words=[BlockedWord(keyword="kumquat", action=ContentFilterAction.BLOCK)], - mask_response_content=True, - ) - - assert guardrail.rewrites_streamed_output() is True - - @pytest.mark.parametrize("action, expected", [("MASK", True), ("BLOCK", False)]) - def test_category_keywords_follow_the_category_action(self, action, expected): - guardrail = ContentFilterGuardrail( - guardrail_name="cf", - categories=[{"category": "bias_gender", "enabled": True, "action": action}], - ) - - assert guardrail.category_keywords and not guardrail.always_block_category_keywords - assert guardrail.rewrites_streamed_output() is expected - - @pytest.mark.parametrize("action, expected", [("MASK", True), ("BLOCK", False)]) - def test_always_block_category_keywords_follow_the_category_action(self, action, expected): - guardrail = ContentFilterGuardrail( - guardrail_name="cf", - categories=[{"category": "age_discrimination", "enabled": True, "action": action}], - ) - - assert guardrail.always_block_category_keywords and not guardrail.category_keywords - assert guardrail.rewrites_streamed_output() is expected From 704013dbb6f83962e4cac97c79ac0dd2109cbabf Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 7 Sep 2026 22:13:23 -0700 Subject: [PATCH 43/60] fix(init): keep non-flag LITELLM_DROP_PARAMS values on with a warning The merge base read the variable by truthiness, so any non-empty value turned the global flag on. Parsing it as a flag made a value such as temperature or enabled silently turn it off, and the only docs for the variable describe it as a list of parameter names, so keep those values on and log a warning that asks for true or false. A blank value stays off without a warning --- litellm/__init__.py | 4 +-- litellm/litellm_core_utils/core_helpers.py | 16 +++++++++++ .../litellm_core_utils/test_core_helpers.py | 28 +++++++++++++++++++ .../test_litellm/test_drop_params_env_var.py | 11 +++++--- 4 files changed, 53 insertions(+), 6 deletions(-) diff --git a/litellm/__init__.py b/litellm/__init__.py index f7d4dce87d2..fc6dc35fe55 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -47,7 +47,7 @@ from typing import ( ) from litellm.types.integrations.datadog import DatadogInitParams from litellm.types.integrations.newrelic import NewRelicInitParams -from litellm.litellm_core_utils.core_helpers import drop_params_flag +from litellm.litellm_core_utils.core_helpers import drop_params_env_flag from litellm._logging import ( set_verbose, _turn_on_debug, @@ -239,7 +239,7 @@ token: Optional[str] = ( ) telemetry = True max_tokens: int = DEFAULT_MAX_TOKENS # OpenAI Defaults -drop_params = drop_params_flag(os.getenv("LITELLM_DROP_PARAMS"), "LITELLM_DROP_PARAMS", verbose_logger) +drop_params = drop_params_env_flag(os.environ, verbose_logger) modify_params = bool(os.getenv("LITELLM_MODIFY_PARAMS", False)) use_chat_completions_url_for_anthropic_messages: bool = bool( os.getenv("LITELLM_USE_CHAT_COMPLETIONS_URL_FOR_ANTHROPIC_MESSAGES", False) diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index a3dfac81cc1..eacc3e4860a 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -58,6 +58,22 @@ def drop_params_flag(value: object, source: str, logger: logging.Logger) -> bool return bool(normalized) +DROP_PARAMS_ENV_VAR: Final = "LITELLM_DROP_PARAMS" + + +def drop_params_env_flag(environ: Mapping[str, str], logger: logging.Logger) -> bool: + configured: Final = environ.get(DROP_PARAMS_ENV_VAR, "").strip() + if configured == "": + return False + normalized: Final = normalize_drop_params(configured) + if normalized is None: + logger.warning( + "%s=%r is not a flag value, treating it as on. Set it to true or false", DROP_PARAMS_ENV_VAR, configured + ) + return True + return normalized + + def safe_divide( numerator: float, denominator: float, diff --git a/tests/test_litellm/litellm_core_utils/test_core_helpers.py b/tests/test_litellm/litellm_core_utils/test_core_helpers.py index e3880175759..e937be47441 100644 --- a/tests/test_litellm/litellm_core_utils/test_core_helpers.py +++ b/tests/test_litellm/litellm_core_utils/test_core_helpers.py @@ -6,6 +6,7 @@ import pytest from litellm.litellm_core_utils.core_helpers import ( _FINISH_REASON_MAP, + drop_params_env_flag, drop_params_flag, get_or_create_metadata_bucket, map_finish_reason, @@ -301,6 +302,33 @@ def test_drop_params_flag_treats_non_flag_values_as_off_with_a_warning(value, ca assert f"LITELLM_DROP_PARAMS={value!r} is not a flag value, treating it as off" in caplog.text +@pytest.mark.parametrize( + "environ, expected", + [ + ({}, False), + ({"LITELLM_DROP_PARAMS": ""}, False), + ({"LITELLM_DROP_PARAMS": " "}, False), + ({"LITELLM_DROP_PARAMS": "true"}, True), + ({"LITELLM_DROP_PARAMS": " False "}, False), + ({"LITELLM_DROP_PARAMS": "0"}, False), + ], +) +def test_drop_params_env_flag_reads_a_flag_without_a_warning(environ, expected, caplog): + with caplog.at_level(logging.WARNING, logger="drop-params-test"): + assert drop_params_env_flag(environ, logging.getLogger("drop-params-test")) is expected + assert caplog.text == "" + + +@pytest.mark.parametrize("configured", ["temperature", "temperature,top_p", "enabled"]) +def test_drop_params_env_flag_keeps_a_non_flag_value_on_with_a_warning(configured, caplog): + with caplog.at_level(logging.WARNING, logger="drop-params-test"): + assert drop_params_env_flag({"LITELLM_DROP_PARAMS": configured}, logging.getLogger("drop-params-test")) is True + assert ( + f"LITELLM_DROP_PARAMS={configured!r} is not a flag value, treating it as on. Set it to true or false" + in caplog.text + ) + + class TestIsExpectedClientError: def test_status_ranges(self): from litellm.litellm_core_utils.core_helpers import is_expected_client_error diff --git a/tests/test_litellm/test_drop_params_env_var.py b/tests/test_litellm/test_drop_params_env_var.py index 339298df3d1..1e0b7801ef1 100644 --- a/tests/test_litellm/test_drop_params_env_var.py +++ b/tests/test_litellm/test_drop_params_env_var.py @@ -15,7 +15,7 @@ def _import_litellm_with(configured: str) -> subprocess.CompletedProcess[str]: ) -@pytest.mark.parametrize("configured, expected", [("false", "False"), ("true", "True")]) +@pytest.mark.parametrize("configured, expected", [("false", "False"), ("true", "True"), ("", "False")]) def test_litellm_drop_params_env_var_is_parsed_as_a_flag(configured, expected): result = _import_litellm_with(configured) @@ -23,8 +23,11 @@ def test_litellm_drop_params_env_var_is_parsed_as_a_flag(configured, expected): assert "is not a flag value" not in result.stderr -def test_litellm_drop_params_env_var_non_flag_value_is_off_with_a_warning(): +def test_litellm_drop_params_env_var_non_flag_value_stays_on_with_a_warning(): result = _import_litellm_with("temperature") - assert result.stdout.strip() == "False" - assert "LITELLM_DROP_PARAMS='temperature' is not a flag value, treating it as off" in result.stderr + assert result.stdout.strip() == "True" + assert ( + "LITELLM_DROP_PARAMS='temperature' is not a flag value, treating it as on. Set it to true or false" + in result.stderr + ) From 3448175184d8b68dbcf9ec007d62f85651bb0f41 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 11:44:49 -0700 Subject: [PATCH 44/60] fix(responses bridge): keep mid-conversation system messages in input Only the leading run of system messages (before the first non-system message) is joined into the Responses `instructions` field. A system message that arrives after a user, assistant, or tool turn now becomes a system input item at its position, whether its content is a string or a list, so a client that re-sends the same reminder as a string on the next request produces byte-identical input and `instructions` stays stable. Claude Code >= 2.1.237 appends such reminders after every user turn, and folding them into `instructions` made Azure treat every request as a cold prompt (cached_tokens 0 on every request of a session). Fixes #40198 --- .../transformation.py | 9 +- ...responses_transformation_transformation.py | 140 ++++++++++++++++++ 2 files changed, 147 insertions(+), 2 deletions(-) diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index 87350b5479c..4fe069b0b7d 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -370,7 +370,12 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): and isinstance(tool_call.get("custom"), dict) ) - for msg in messages: + leading_system_count: Final = next( + (index for index, msg in enumerate(messages) if msg.get("role") != "system"), + len(messages), + ) + + for index, msg in enumerate(messages): role = msg.get("role") content = msg.get("content", "") tool_calls = msg.get("tool_calls") @@ -378,7 +383,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): if role == "system": # Extract system message as instructions - if isinstance(content, str): + if isinstance(content, str) and index < leading_system_count: if instructions: # Concatenate multiple system prompts with a space instructions = f"{instructions} {content}" diff --git a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py index d4d47b145d1..e4ada0a9b31 100644 --- a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py +++ b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py @@ -4147,3 +4147,143 @@ def test_streaming_final_chunk_carries_provider_metadata(): assert chunks[-1]["content_filters"] == content_filters assert "background" not in chunks[-1] assert all("service_tier" not in chunk for chunk in chunks[:-1]) + + +def _system_input_item(text: str) -> dict[str, object]: + return {"type": "message", "role": "system", "content": [{"type": "input_text", "text": text}]} + + +def test_mid_conversation_system_string_stays_in_input_after_a_user_turn(): + handler: Final = LiteLLMResponsesTransformationHandler() + + input_items, instructions = handler.convert_chat_completion_messages_to_responses_api( + [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "Read the file."}, + {"role": "system", "content": "14982391 tokens left"}, + {"role": "user", "content": "Now summarize it."}, + ] + ) + + assert instructions == "You are a helpful assistant." + assert input_items == [ + {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "Read the file."}]}, + _system_input_item("14982391 tokens left"), + {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "Now summarize it."}]}, + ] + + +def test_leading_system_strings_still_join_instructions_without_a_following_turn(): + handler: Final = LiteLLMResponsesTransformationHandler() + + input_items, instructions = handler.convert_chat_completion_messages_to_responses_api( + [ + {"role": "system", "content": "Be brief."}, + {"role": "system", "content": "Answer in French."}, + ] + ) + + assert instructions == "Be brief. Answer in French." + assert input_items == [] + + +def test_mid_conversation_system_reminder_as_string_and_as_text_block_produce_identical_input_items(): + handler: Final = LiteLLMResponsesTransformationHandler() + reminder: Final = "14982391 tokens left" + + as_string, string_instructions = handler.convert_chat_completion_messages_to_responses_api( + [{"role": "user", "content": "Read the file."}, {"role": "system", "content": reminder}] + ) + as_block, block_instructions = handler.convert_chat_completion_messages_to_responses_api( + [ + {"role": "user", "content": "Read the file."}, + { + "role": "system", + "content": [{"type": "text", "text": reminder, "cache_control": {"type": "ephemeral"}}], + }, + ] + ) + + assert string_instructions is None + assert block_instructions is None + assert json.dumps(as_string) == json.dumps(as_block) + assert as_string[1] == _system_input_item(reminder) + + +def test_claude_code_shaped_history_keeps_a_byte_stable_input_prefix_across_requests(): + handler: Final = LiteLLMResponsesTransformationHandler() + top_level_system: Final = [{"type": "text", "text": "You are Claude Code.", "cache_control": {"type": "ephemeral"}}] + first_reminder: Final = "27k chars of deferred tools" + second_reminder: Final = "14982391 tokens left" + first_request_messages: Final = [ + {"role": "system", "content": top_level_system}, + {"role": "user", "content": "Read inventory.py."}, + { + "role": "system", + "content": [{"type": "text", "text": first_reminder, "cache_control": {"type": "ephemeral"}}], + }, + ] + second_request_messages: Final = [ + {"role": "system", "content": top_level_system}, + {"role": "user", "content": "Read inventory.py."}, + {"role": "system", "content": first_reminder}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "Read", "arguments": '{"file_path": "inventory.py"}'}, + } + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "ITEMS = []"}, + { + "role": "system", + "content": [{"type": "text", "text": second_reminder, "cache_control": {"type": "ephemeral"}}], + }, + ] + + first_request: Final = handler.transform_request( + model="gpt-5.6-luna", + messages=first_request_messages, + optional_params={}, + litellm_params={}, + headers={}, + litellm_logging_obj=Mock(), + ) + second_request: Final = handler.transform_request( + model="gpt-5.6-luna", + messages=second_request_messages, + optional_params={}, + litellm_params={}, + headers={}, + litellm_logging_obj=Mock(), + ) + + assert "instructions" not in first_request + assert "instructions" not in second_request + assert first_request["input"][0] == _system_input_item("You are Claude Code.") + assert json.dumps(second_request["input"][: len(first_request["input"])]) == json.dumps(first_request["input"]) + assert second_request["input"][len(first_request["input"]) :] == [ + {"type": "function_call", "call_id": "call_1", "name": "Read", "arguments": '{"file_path": "inventory.py"}'}, + {"type": "function_call_output", "call_id": "call_1", "output": [{"type": "input_text", "text": "ITEMS = []"}]}, + _system_input_item(second_reminder), + ] + + +def test_system_string_after_a_developer_message_stays_in_input_in_client_order(): + handler: Final = LiteLLMResponsesTransformationHandler() + + input_items, instructions = handler.convert_chat_completion_messages_to_responses_api( + [ + {"role": "developer", "content": "Always answer in French."}, + {"role": "system", "content": "Be brief."}, + {"role": "user", "content": "Bonjour"}, + ] + ) + + assert instructions is None + assert [item["role"] for item in input_items] == ["developer", "system", "user"] + assert input_items[1] == _system_input_item("Be brief.") From 35d1d40a6784c222714e880c7d07729bc90018b2 Mon Sep 17 00:00:00 2001 From: yujonglee Date: Tue, 8 Sep 2026 12:07:25 -0700 Subject: [PATCH 45/60] fix(ocr): run post-call logging hooks (#40154) --- litellm/integrations/custom_guardrail.py | 10 +-- litellm/llms/custom_httpx/llm_http_handler.py | 12 ++++ .../test_secret_detection.py | 2 +- .../integrations/test_custom_guardrail.py | 33 ++++++++++ .../custom_httpx/test_llm_http_handler.py | 64 +++++++++++++++++++ .../test_proxy_logging_hook_detection.py | 10 +-- 6 files changed, 121 insertions(+), 10 deletions(-) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 2d66a280663..e76d4b05bb2 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -773,7 +773,7 @@ class CustomGuardrail(CustomLogger): def uses_apply_guardrail_interface(self) -> bool: return type(self).apply_guardrail is not CustomGuardrail.apply_guardrail - def _deployment_pre_call_target(self) -> "CustomLogger": + def _deployment_hook_target(self) -> "CustomLogger": if not self.uses_apply_guardrail_interface() or self.use_native_lifecycle_hooks: return self try: @@ -802,7 +802,7 @@ class CustomGuardrail(CustomLogger): # CHECK IF GUARDRAIL REJECTS THE REQUEST if call_type == CallTypes.completion or call_type == CallTypes.acompletion: - target: Final = self._deployment_pre_call_target() + target: Final = self._deployment_hook_target() if target is not self: kwargs["guardrail_to_apply"] = self result: Final = await target.async_pre_call_hook( @@ -845,7 +845,9 @@ class CustomGuardrail(CustomLogger): return None # CHECK IF GUARDRAIL REJECTS THE REQUEST - result: Final = await self.async_post_call_success_hook( + target: Final = self._deployment_hook_target() + hook_request_data: Final = {**request_data, "guardrail_to_apply": self} if target is not self else request_data + result: Final = await target.async_post_call_success_hook( user_api_key_dict=UserAPIKeyAuth( user_id=request_data.get("user_api_key_user_id"), team_id=request_data.get("user_api_key_team_id"), @@ -853,7 +855,7 @@ class CustomGuardrail(CustomLogger): api_key=request_data.get("user_api_key_hash"), request_route=request_data.get("user_api_key_request_route"), ), - data=request_data, + data=hook_request_data, response=response, ) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 2f561809940..8b6d3a0ceb4 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -1741,6 +1741,12 @@ class BaseLLMHTTPHandler: except Exception as e: raise self._handle_error(e=e, provider_config=provider_config) + logging_obj.post_call( + api_key=api_key, + original_response=response.text, + additional_args={"complete_input_dict": data}, + ) + return self._transform_ocr_response( provider_config=provider_config, model=model, @@ -1804,6 +1810,12 @@ class BaseLLMHTTPHandler: except Exception as e: raise self._handle_error(e=e, provider_config=provider_config) + logging_obj.post_call( + api_key=api_key, + original_response=response.text, + additional_args={"complete_input_dict": data}, + ) + # Use async response transform for async operations return await provider_config.async_transform_ocr_response( model=model, diff --git a/tests/test_litellm/enterprise/enterprise_callbacks/test_secret_detection.py b/tests/test_litellm/enterprise/enterprise_callbacks/test_secret_detection.py index f46df5baadf..c1c3569c0e2 100644 --- a/tests/test_litellm/enterprise/enterprise_callbacks/test_secret_detection.py +++ b/tests/test_litellm/enterprise/enterprise_callbacks/test_secret_detection.py @@ -262,7 +262,7 @@ def test_proxied_traffic_stays_on_native_hooks(): never sees ``data["prompt"]``.""" guardrail = _guardrail() assert guardrail.uses_apply_guardrail_interface() is True - assert guardrail._deployment_pre_call_target() is guardrail + assert guardrail._deployment_hook_target() is guardrail @pytest.mark.asyncio diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/test_litellm/integrations/test_custom_guardrail.py index 885bd1d4d72..359635335b3 100644 --- a/tests/test_litellm/integrations/test_custom_guardrail.py +++ b/tests/test_litellm/integrations/test_custom_guardrail.py @@ -2610,3 +2610,36 @@ class TestCustomGuardrailPostCallSuccessDeploymentHook: ) assert result is replacement + + @pytest.mark.asyncio + async def test_apply_guardrail_interface_modifies_deployment_response(self): + from litellm.types.guardrails import GuardrailEventHooks + from litellm.types.utils import ModelResponse + + class ReplacingGuardrail(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict[str, object], + input_type: Literal["request", "response"], + logging_obj: Optional["LiteLLMLoggingObj"] = None, + ) -> GenericGuardrailAPIInputs: + assert input_type == "response" + return {**inputs, "texts": ["filtered response"]} + + guardrail = ReplacingGuardrail( + guardrail_name="test-guardrail", + event_hook=GuardrailEventHooks.post_call, + ) + response = ModelResponse(choices=[{"message": {"role": "assistant", "content": "original response"}}]) + request_data = {"guardrails": ["test-guardrail"]} + + result = await guardrail.async_post_call_success_deployment_hook( + request_data=request_data, + response=response, + call_type=CallTypes.acompletion, + ) + + assert result is response + assert response.choices[0].message.content == "filtered response" + assert request_data == {"guardrails": ["test-guardrail"]} diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index e16855da8cb..98a5f4b2db5 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py @@ -29,6 +29,7 @@ from litellm.llms.custom_httpx.llm_http_handler import ( _rust_responses_websocket_enabled, ) from litellm.llms.azure.videos.transformation import AzureVideoConfig +from litellm.llms.mistral.ocr.transformation import MistralOCRConfig from litellm.llms.openai.videos.transformation import OpenAIVideoConfig from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.router import GenericLiteLLMParams @@ -37,6 +38,69 @@ from litellm.types.utils import ImageObject, ImageResponse, ModelResponse, Trans _ACTIVE_KEY = "_code_interpreter_interception_active" _SANDBOX_KEY = "_code_interpreter_interception_sandbox_key" +OCR_RESPONSE = { + "pages": [{"index": 0, "markdown": "OCR output", "images": []}], + "model": "mistral-ocr-latest", + "usage_info": {"pages_processed": 1}, +} + + +def _ocr_sync_client() -> HTTPHandler: + client = HTTPHandler() + client.client = httpx.Client(transport=httpx.MockTransport(lambda _request: httpx.Response(200, json=OCR_RESPONSE))) + return client + + +def _ocr_async_client() -> AsyncHTTPHandler: + client = AsyncHTTPHandler() + client.client = httpx.AsyncClient( + transport=httpx.MockTransport(lambda _request: httpx.Response(200, json=OCR_RESPONSE)) + ) + return client + + +def test_ocr_calls_post_call_with_raw_provider_response(): + logging_obj = Mock() + + response = BaseLLMHTTPHandler().ocr( + model="mistral-ocr-latest", + document={"type": "document_url", "document_url": "https://example.com/document.pdf"}, + optional_params={}, + timeout=5, + logging_obj=logging_obj, + api_key="test-key", + api_base="https://api.mistral.ai/v1/ocr", + custom_llm_provider="mistral", + client=_ocr_sync_client(), + provider_config=MistralOCRConfig(), + ) + + assert response.pages[0].markdown == "OCR output" + logging_obj.post_call.assert_called_once() + assert json.loads(logging_obj.post_call.call_args.kwargs["original_response"]) == OCR_RESPONSE + + +@pytest.mark.asyncio +async def test_async_ocr_calls_post_call_with_raw_provider_response(): + logging_obj = Mock() + + response = await BaseLLMHTTPHandler().async_ocr( + model="mistral-ocr-latest", + document={"type": "document_url", "document_url": "https://example.com/document.pdf"}, + optional_params={}, + timeout=5, + logging_obj=logging_obj, + api_key="test-key", + api_base="https://api.mistral.ai/v1/ocr", + custom_llm_provider="mistral", + client=_ocr_async_client(), + provider_config=MistralOCRConfig(), + ) + + assert response.pages[0].markdown == "OCR output" + logging_obj.post_call.assert_called_once() + assert json.loads(logging_obj.post_call.call_args.kwargs["original_response"]) == OCR_RESPONSE + def test_prepare_fake_stream_request(): # Initialize the BaseLLMHTTPHandler diff --git a/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py b/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py index 9f1321aec2c..22930a26974 100644 --- a/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py +++ b/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py @@ -637,14 +637,14 @@ def test_callback_capabilities_excludes_opted_out_guardrail_from_iterator_overri assert [cb for cb, _ in caps.iterator_overrides if cb is opted_out] == [] -def test_deployment_pre_call_target_stays_native_when_opted_out(): +def test_deployment_hook_target_stays_native_when_opted_out(): """Model-level guardrails resolve their target here rather than through ProxyLogging.""" - assert _KeepsNativeHooks()._deployment_pre_call_target() is not None + assert _KeepsNativeHooks()._deployment_hook_target() is not None opted_out = _KeepsNativeHooks() - assert opted_out._deployment_pre_call_target() is opted_out - assert _AppliesGuardrail()._deployment_pre_call_target() is not None + assert opted_out._deployment_hook_target() is opted_out + assert _AppliesGuardrail()._deployment_hook_target() is not None routed = _AppliesGuardrail() - assert routed._deployment_pre_call_target() is not routed + assert routed._deployment_hook_target() is not routed @pytest.mark.asyncio From 01c68c199bc0afddf5a4d806cb03eca5e7745295 Mon Sep 17 00:00:00 2001 From: yujonglee Date: Tue, 8 Sep 2026 12:18:09 -0700 Subject: [PATCH 46/60] fix(guardrails): allow framework-supported logging-only mode (#40267) --- litellm/integrations/custom_guardrail.py | 10 +- .../integrations/test_custom_guardrail.py | 100 +++++++++++++++++- 2 files changed, 107 insertions(+), 3 deletions(-) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index e76d4b05bb2..37d6a7e793d 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -601,6 +601,12 @@ class CustomGuardrail(CustomLogger): event_hook: GuardrailEventHooks | list[GuardrailEventHooks] | Mode | None, supported_event_hooks: list[GuardrailEventHooks], ) -> None: + allowed_hooks: Final = frozenset(supported_event_hooks) | ( + frozenset((GuardrailEventHooks.logging_only,)) + if self.uses_apply_guardrail_interface() and not self.use_native_lifecycle_hooks + else frozenset() + ) + def _validate_event_hook_list_is_in_supported_event_hooks( event_hook: list[GuardrailEventHooks] | list[str], supported_event_hooks: list[GuardrailEventHooks], @@ -608,7 +614,7 @@ class CustomGuardrail(CustomLogger): for hook in event_hook: if isinstance(hook, str): hook = GuardrailEventHooks(hook) - if hook not in supported_event_hooks: + if hook not in allowed_hooks: raise ValueError(f"Event hook {hook} is not in the supported event hooks {supported_event_hooks}") if event_hook is None: @@ -629,7 +635,7 @@ class CustomGuardrail(CustomLogger): default_list = event_hook.default if isinstance(event_hook.default, list) else [event_hook.default] _validate_event_hook_list_is_in_supported_event_hooks(default_list, supported_event_hooks) elif isinstance(event_hook, GuardrailEventHooks): - if event_hook not in supported_event_hooks: + if event_hook not in allowed_hooks: raise ValueError(f"Event hook {event_hook} is not in the supported event hooks {supported_event_hooks}") @staticmethod diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/test_litellm/integrations/test_custom_guardrail.py index 359635335b3..cd8d609cf71 100644 --- a/tests/test_litellm/integrations/test_custom_guardrail.py +++ b/tests/test_litellm/integrations/test_custom_guardrail.py @@ -1,5 +1,5 @@ import asyncio -from typing import TYPE_CHECKING, Literal, Optional +from typing import TYPE_CHECKING, ClassVar, Final, Literal, Optional from unittest.mock import AsyncMock import pytest @@ -10,6 +10,7 @@ from litellm.integrations.custom_guardrail import ( log_guardrail_information, ) from litellm.proxy._types import CallTypes, UserAPIKeyAuth +from litellm.types.guardrails import GuardrailEventHooks, Mode from litellm.types.utils import GenericGuardrailAPIInputs, GuardrailTracingDetail if TYPE_CHECKING: @@ -2378,11 +2379,108 @@ def _logged_call(messages: list | str) -> tuple[dict, object]: return kwargs, response +class _NativeApplyGuardrail(_InheritedApplyGuardrail): + use_native_lifecycle_hooks: ClassVar[bool] = True + + +@pytest.mark.parametrize("guardrail_type", (CustomGuardrail, _NativeApplyGuardrail, _InheritedApplyGuardrail)) +@pytest.mark.parametrize( + "event_hook", + ( + GuardrailEventHooks.logging_only, + "logging_only", + [GuardrailEventHooks.pre_call, GuardrailEventHooks.logging_only], + ["pre_call", "logging_only"], + Mode(tags={"audit": "logging_only"}, default="pre_call"), + Mode(tags={"audit": ["pre_call", "logging_only"]}), + Mode(tags={"enforce": "pre_call"}, default="logging_only"), + Mode(tags={}, default=["pre_call", "logging_only"]), + ), +) +def test_logging_only_requires_framework_support_or_explicit_declaration( + guardrail_type: type[CustomGuardrail], + event_hook: GuardrailEventHooks | str | list[GuardrailEventHooks] | list[str] | Mode, +) -> None: + supported: Final = [GuardrailEventHooks.pre_call] + if guardrail_type is _InheritedApplyGuardrail: + guardrail: Final = guardrail_type(event_hook=event_hook, supported_event_hooks=supported) + assert guardrail.event_hook == event_hook + assert supported == [GuardrailEventHooks.pre_call] + else: + with pytest.raises(ValueError, match=r"logging_only.*not in the supported event hooks"): + guardrail_type(event_hook=event_hook, supported_event_hooks=supported) + + explicitly_supported: Final = guardrail_type( + event_hook=event_hook, + supported_event_hooks=[GuardrailEventHooks.pre_call, GuardrailEventHooks.logging_only], + ) + assert explicitly_supported.event_hook == event_hook + + +@pytest.mark.parametrize( + "event_hook", + ( + GuardrailEventHooks.post_call, + "post_call", + [GuardrailEventHooks.logging_only, GuardrailEventHooks.post_call], + ["logging_only", "post_call"], + Mode(tags={"enforce": "post_call"}, default="logging_only"), + Mode(tags={"enforce": ["logging_only", "post_call"]}), + Mode(tags={"audit": "logging_only"}, default="post_call"), + Mode(tags={}, default=["logging_only", "post_call"]), + ), +) +def test_framework_logging_only_does_not_allow_other_unsupported_modes( + event_hook: GuardrailEventHooks | str | list[GuardrailEventHooks] | list[str] | Mode, +) -> None: + with pytest.raises(ValueError, match=r"post_call.*not in the supported event hooks"): + _InheritedApplyGuardrail(event_hook=event_hook, supported_event_hooks=[GuardrailEventHooks.pre_call]) + + class TestLoggingOnlyApplyGuardrail: """LIT-4876 regression: a guardrail in mode logging_only that implements only apply_guardrail must still run against the logged request and response and record guardrail_information, instead of inheriting the CustomLogger no-op.""" + @pytest.mark.parametrize( + "event_hook", + ( + GuardrailEventHooks.logging_only, + "logging_only", + [GuardrailEventHooks.pre_call, GuardrailEventHooks.logging_only], + ["pre_call", "logging_only"], + Mode(tags={"audit": "logging_only"}, default="pre_call"), + Mode(tags={"audit": ["pre_call", "logging_only"]}), + Mode(tags={"enforce": "pre_call"}, default="logging_only"), + Mode(tags={}, default=["pre_call", "logging_only"]), + ), + ) + @pytest.mark.asyncio + async def test_content_filter_accepts_logging_only_and_records_detection( + self, event_hook: GuardrailEventHooks | str | list[GuardrailEventHooks] | list[str] | Mode + ) -> None: + from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ( + ContentFilterGuardrail, + ) + from litellm.types.guardrails import BlockedWord, ContentFilterAction, GuardrailEventHooks + + guardrail: Final = ContentFilterGuardrail( + guardrail_name="content-review", + event_hook=event_hook, + default_on=True, + blocked_words=[BlockedWord(keyword="hello", action=ContentFilterAction.BLOCK)], + ) + kwargs, response = _logged_call([{"role": "user", "content": "hello there"}]) + + out_kwargs, out_response = await guardrail.async_logging_hook(kwargs, response, CallTypes.acompletion.value) + + assert out_response is response + assert out_kwargs["messages"] == kwargs["messages"] + assert ( + out_kwargs["standard_logging_object"]["guardrail_information"][0]["guardrail_status"] + == "guardrail_intervened" + ) + @pytest.mark.asyncio async def test_runs_apply_guardrail_observe_only_and_records_verdict(self): guardrail = _ApplyOnlyObserver() From d36e032241490a5a75063425a99da4ccdc85548a Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 8 Sep 2026 13:00:14 -0700 Subject: [PATCH 47/60] fix(proxy): initialize string success/failure callbacks at startup after config load (#38226) * fix(proxy): eagerly initialize string callbacks Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): materialize string callbacks after load_config so later litellm_settings keys are applied Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): drop casts when snapshotting string callbacks so LIT006 stays at base Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: yucheng --- litellm/proxy/utils.py | 11 +++++ tests/test_litellm/proxy/test_proxy_server.py | 41 +++++++++++++++++++ 2 files changed, 52 insertions(+) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index ddf31cb1d8a..75d7972c621 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -177,6 +177,9 @@ from litellm.types.mcp import ( ) from litellm.types.proxy.policy_engine.pipeline_types import PipelineExecutionResult from litellm.types.utils import LLMResponseTypes, LoggedLiteLLMParams +from litellm.utils import ( + _add_custom_logger_callback_to_specific_event, # pyright: ignore[reportPrivateUsage] # only string-to-logger helper +) if TYPE_CHECKING: from mcp.types import CallToolResult @@ -857,6 +860,14 @@ class ProxyLogging: litellm.logging_callback_manager.add_litellm_async_success_callback(callback) litellm.logging_callback_manager.add_litellm_async_failure_callback(callback) + # Runs after load_config applied every litellm_settings key: logger __init__s read e.g. s3_callback_params + success_callbacks: Final = tuple(cb for cb in litellm.success_callback if isinstance(cb, str)) + failure_callbacks: Final = tuple(cb for cb in litellm.failure_callback if isinstance(cb, str)) + for callback in success_callbacks: + _add_custom_logger_callback_to_specific_event(callback, "success") + for callback in failure_callbacks: + _add_custom_logger_callback_to_specific_event(callback, "failure") + async def update_request_status(self, litellm_call_id: str, status: Literal["success", "fail"]): # only use this if slack alerting is being used if self.alerting is None: diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 937e4f15741..6a901c5e7de 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -3166,6 +3166,47 @@ async def test_custom_ui_sso_sign_in_handler_config_loading(): os.unlink(config_file_path) +@pytest.mark.asyncio +async def test_startup_initializes_string_callbacks_after_all_litellm_settings_load(tmp_path, monkeypatch): + from litellm.integrations.s3_v2 import S3Logger + from litellm.litellm_core_utils import litellm_logging + from litellm.proxy.proxy_server import ProxyConfig + from litellm.proxy.utils import ProxyLogging + + config_file = tmp_path / "config.yaml" + config_file.write_text( + "model_list: []\n" + "litellm_settings:\n" + " success_callback:\n" + " - s3_v2\n" + " failure_callback:\n" + " - s3_v2\n" + " s3_callback_params:\n" + " s3_bucket_name: ordering-regression-bucket\n" + " s3_region_name: us-west-2\n" + ) + + monkeypatch.setattr(litellm, "success_callback", []) + monkeypatch.setattr(litellm, "_async_success_callback", []) + monkeypatch.setattr(litellm, "failure_callback", []) + monkeypatch.setattr(litellm, "_async_failure_callback", []) + monkeypatch.setattr(litellm, "callbacks", []) + monkeypatch.setattr(litellm, "s3_callback_params", None) + monkeypatch.setattr(litellm_logging, "_in_memory_loggers", []) + + await ProxyConfig().load_config(router=MagicMock(), config_file_path=str(config_file)) + ProxyLogging(user_api_key_cache=MagicMock())._init_litellm_callbacks(llm_router=None) + + success_loggers = [cb for cb in litellm._async_success_callback if isinstance(cb, S3Logger)] + failure_loggers = [cb for cb in litellm._async_failure_callback if isinstance(cb, S3Logger)] + assert len(success_loggers) == 1 + assert len(failure_loggers) == 1 + assert success_loggers[0].s3_bucket_name == "ordering-regression-bucket" + assert success_loggers[0].s3_region_name == "us-west-2" + assert "s3_v2" not in litellm.success_callback + assert "s3_v2" not in litellm.failure_callback + + @pytest.mark.asyncio async def test_load_config_max_budget_env_var_coerced_to_float(tmp_path, monkeypatch): """ From 6a425a5cc59ce67f3860ac04f752cf61340e1028 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 8 Sep 2026 13:08:26 -0700 Subject: [PATCH 48/60] fix(responses): record spend for native Responses API WebSocket sessions (#38856) * fix(responses): record spend for native Responses API WebSocket sessions Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(responses): bill usage from response.incomplete WebSocket turns A turn cut short by max_output_tokens ends in response.incomplete, which OpenAI bills but the processor only read response.completed, so those sessions still logged zero spend Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(responses): hoist websocket usage test imports to module scope Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(responses): price websocket sessions through the standard cost path The realtime completion_cost branch skips cost_discount_config and cost_margin_config, so a native Responses WebSocket session was priced differently from the same usage over HTTP /v1/responses. Drop the explicit widening so the LiteLLMRealtimeStreamLoggingObject built by normalize_logging_result flows through the generic usage path, and pin WS == HTTP cost under a 50% provider discount in the regression test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * ci: rerun proxy-infra after flaky test_check_migration process tree test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: yucheng --- litellm/cost_calculator.py | 40 ++++++++ litellm/litellm_core_utils/litellm_logging.py | 12 +++ litellm/types/utils.py | 1 + .../test_litellm_logging.py | 92 ++++++++++++++++++- 4 files changed, 144 insertions(+), 1 deletion(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 9a9d2ceda03..c9c7df4d75e 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -2414,6 +2414,46 @@ class RealtimeAPITokenUsageProcessor(BaseTokenUsageProcessor): ) +_RESPONSES_WS_BILLABLE_EVENT_TYPES: Final = frozenset({"response.completed", "response.incomplete"}) + + +class _ResponsesWsEventResponse(BaseModel): + usage: Mapping[str, object] | None = None + + +class _ResponsesWsEvent(BaseModel): + type: str = "" + response: _ResponsesWsEventResponse | None = None + + +class ResponsesWebSocketTokenUsageProcessor(BaseTokenUsageProcessor): + @staticmethod + def collect_usage_from_responses_ws_results( + results: Sequence[Mapping[str, object]], + ) -> tuple[Usage, ...]: + events: Final = tuple(_ResponsesWsEvent.model_validate(result) for result in results) + return tuple( + ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( # pyright: ignore[reportPrivateUsage] # same shared transform the realtime processor uses + event.response.usage + ) + for event in events + if event.type in _RESPONSES_WS_BILLABLE_EVENT_TYPES + and event.response is not None + and event.response.usage is not None + ) + + @staticmethod + def collect_and_combine_usage_from_responses_ws_results( + results: Sequence[Mapping[str, object]], + ) -> Usage: + collected_usage_objects: Final = ResponsesWebSocketTokenUsageProcessor.collect_usage_from_responses_ws_results( + results + ) + return ResponsesWebSocketTokenUsageProcessor.combine_usage_objects( + list(collected_usage_objects) # mutable-ok: combine_usage_objects requires a list parameter + ) + + _TRANSCRIPTION_COMPLETED_EVENT_TYPE: Final = "conversation.item.input_audio_transcription.completed" diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index ca2cca5360f..e3c94866985 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -48,6 +48,7 @@ from litellm.constants import ( ) from litellm.cost_calculator import ( RealtimeAPITokenUsageProcessor, + ResponsesWebSocketTokenUsageProcessor, _select_model_name_for_cost_calc, ) from litellm.exceptions import ( @@ -2028,6 +2029,17 @@ class Logging(LiteLLMLoggingBaseClass): results=result, ) + elif self.call_type == CallTypes.aresponses_websocket.value and isinstance(result, list): # pyright: ignore[reportUnknownMemberType] # Logging.call_type is untyped + combined_ws_usage: Final = ( + ResponsesWebSocketTokenUsageProcessor.collect_and_combine_usage_from_responses_ws_results( + results=result # pyright: ignore[reportUnknownArgumentType] # raw event dicts from the WS stream + ) + ) + logging_result = LiteLLMRealtimeStreamLoggingObject( + usage=combined_ws_usage, + results=result, # pyright: ignore[reportUnknownArgumentType] # raw event dicts from the WS stream + ) + elif ( self.call_type == CallTypes.llm_passthrough_route.value or self.call_type == CallTypes.allm_passthrough_route.value diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 9b5fb08a45f..40a6d77482f 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -580,6 +580,7 @@ CallTypesLiteral = Literal[ "search", "asearch", "_arealtime", + "_aresponses_websocket", "create_batch", "acreate_batch", "create_file", diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 0fdca755685..b6366bc803e 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -22,7 +22,13 @@ from litellm.litellm_core_utils.litellm_logging import ( _get_status_fields, set_callbacks, ) -from litellm.types.utils import ModelResponse, TextCompletionResponse +from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse +from litellm.types.utils import ( + CallTypes, + LiteLLMRealtimeStreamLoggingObject, + ModelResponse, + TextCompletionResponse, +) @pytest.fixture @@ -6393,6 +6399,90 @@ async def test_prompt_hook_injection_marker_recorded_for_every_surface(logging_o assert pre_choice["metadata"]["litellm_gateway_injected_cache"] == "" +def _responses_ws_logging_obj() -> LitellmLogging: + return LitellmLogging( + model="gpt-4o", + messages=[], + stream=False, + call_type=CallTypes.aresponses_websocket.value, + start_time=time.time(), + litellm_call_id="responses-ws-usage-test", + function_id="responses-ws-usage-test", + ) + + +def test_normalize_logging_result_extracts_usage_for_responses_websocket(monkeypatch): + """LIT-6512: native /v1/responses WebSocket sessions logged $0 spend because the usage + carried by stored response.completed events was never extracted. The session must cost + exactly what the same usage costs over HTTP /v1/responses, discounts included.""" + monkeypatch.setattr(litellm, "cost_discount_config", {"openai": 0.5}) + logging_obj = _responses_ws_logging_obj() + events = [ + {"type": "response.created", "response": {}}, + { + "type": "response.completed", + "response": {"usage": {"input_tokens": 100, "output_tokens": 40, "total_tokens": 140}}, + }, + { + "type": "response.completed", + "response": {"usage": {"input_tokens": 60, "output_tokens": 10, "total_tokens": 70}}, + }, + ] + + normalized = logging_obj.normalize_logging_result(result=events) + + assert isinstance(normalized, LiteLLMRealtimeStreamLoggingObject) + assert normalized.usage.prompt_tokens == 160 + assert normalized.usage.completion_tokens == 50 + + ws_cost = litellm.completion_cost( + completion_response=normalized, + model="gpt-4o", + call_type=CallTypes.aresponses_websocket.value, + custom_llm_provider="openai", + ) + http_cost = litellm.completion_cost( + completion_response=ResponsesAPIResponse( + id="resp-6512", + created_at=1700000000, + output=[], + usage=ResponseAPIUsage(input_tokens=160, output_tokens=50, total_tokens=210), + ), + model="gpt-4o", + call_type=CallTypes.aresponses.value, + custom_llm_provider="openai", + ) + assert ws_cost > 0 + assert ws_cost == http_cost + + +def test_normalize_logging_result_bills_incomplete_responses_websocket_turns(): + """LIT-6512: a turn cut short by max_output_tokens ends in response.incomplete, which + OpenAI bills, so its usage counts toward the session like a completed turn.""" + events = [ + { + "type": "response.created", + "response": {"usage": {"input_tokens": 999, "output_tokens": 999, "total_tokens": 1998}}, + }, + { + "type": "response.incomplete", + "response": {"usage": {"input_tokens": 15, "output_tokens": 16, "total_tokens": 31}}, + }, + { + "type": "response.completed", + "response": {"usage": {"input_tokens": 40, "output_tokens": 4, "total_tokens": 44}}, + }, + {"type": "response.failed", "response": {"usage": None}}, + ] + + normalized = _responses_ws_logging_obj().normalize_logging_result(result=events) + + assert isinstance(normalized, LiteLLMRealtimeStreamLoggingObject) + assert normalized.usage.prompt_tokens == 55 + assert normalized.usage.completion_tokens == 20 + assert normalized.usage.total_tokens == 75 + + def test_get_standard_logging_object_payload_reads_overhead_from_logging_obj_for_dict_results(logging_obj): """LIT-5466: /v1/messages returns a plain dict with no _hidden_params, so the overhead recorded on the logging object must reach hidden_params.litellm_overhead_time_ms (SpendLogs).""" From 9d164fa34150ac7dc0551d9efec16cdb27848c5c Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Tue, 8 Sep 2026 13:24:15 -0700 Subject: [PATCH 49/60] bump: litellm-proxy-extras 0.4.94 -> 0.4.95, litellm 1.101.0 -> 1.102.0 --- litellm-proxy-extras/pyproject.toml | 4 ++-- pyproject.toml | 6 +++--- uv.lock | 6 +++--- 3 files changed, 8 insertions(+), 8 deletions(-) diff --git a/litellm-proxy-extras/pyproject.toml b/litellm-proxy-extras/pyproject.toml index 82d31fec373..91b4e4a7ba1 100644 --- a/litellm-proxy-extras/pyproject.toml +++ b/litellm-proxy-extras/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm-proxy-extras" -version = "0.4.94" +version = "0.4.95" description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package." readme = "README.md" requires-python = ">=3.9" @@ -26,7 +26,7 @@ required-version = ">=0.10.9" module-root = "" [tool.commitizen] -version = "0.4.94" +version = "0.4.95" version_files = [ "pyproject.toml:^version", "../pyproject.toml:litellm-proxy-extras==", diff --git a/pyproject.toml b/pyproject.toml index af35c77d259..706b14b8a04 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm" -version = "1.101.0" +version = "1.102.0" description = "Library to easily interface with LLM API providers" readme = "README.md" requires-python = ">=3.10, <3.15" @@ -67,7 +67,7 @@ proxy = [ "azure-identity>=1.25.2,<2.0", "azure-storage-blob>=12.28.0,<13.0", "mcp>=1.28.1,<2.0", - "litellm-proxy-extras==0.4.94", + "litellm-proxy-extras==0.4.95", "litellm-enterprise==0.1.65", "RestrictedPython>=8.5,<9.0", "rich>=13.9.4,<14.0", @@ -328,7 +328,7 @@ members = ["enterprise", "litellm-proxy-extras"] profile = "black" [tool.commitizen] -version = "1.101.0" +version = "1.102.0" version_files = [ "pyproject.toml:^version", ] diff --git a/uv.lock b/uv.lock index 89205cd9527..12dc2c69d23 100644 --- a/uv.lock +++ b/uv.lock @@ -10,7 +10,7 @@ resolution-markers = [ ] [options] -exclude-newer = "2026-09-02T16:58:34.594994Z" +exclude-newer = "2026-09-05T20:24:07.535116Z" exclude-newer-span = "P3D" [manifest] @@ -4358,7 +4358,7 @@ wheels = [ [[package]] name = "litellm" -version = "1.101.0" +version = "1.102.0" source = { editable = "." } dependencies = [ { name = "aiohttp" }, @@ -4776,7 +4776,7 @@ source = { editable = "enterprise" } [[package]] name = "litellm-proxy-extras" -version = "0.4.94" +version = "0.4.95" source = { editable = "litellm-proxy-extras" } [[package]] From 4a3a78c256bab366597a76d44930b323bce6335b Mon Sep 17 00:00:00 2001 From: tin-berri Date: Tue, 8 Sep 2026 13:28:57 -0700 Subject: [PATCH 50/60] feat(complexity_router): rebalance heuristic weights in the dashboard and grade custom dimensions by match count (#40205) The Advanced scoring editor now lists built-in and custom dimensions together. Editing any weight holds it and rescales the others proportionally so the vector totals 1.00, and Save stores those explicit values. The backend scores exactly what is stored, with no runtime normalization, so routers nobody edits keep their weights. CustomDimension gains an opt-in scoring_mode. match_count scores 0, 0.5 or 1 by distinct matcher hits; the default stays binary. The tuning fingerprint omits a binary scoring_mode, so routers written before this change keep their recorded baseline and the upgrade does not consume the free heuristic-v1 tuning slot. --- .../complexity_router/README.md | 12 +- .../complexity_router/complexity_router.py | 38 ++++- .../complexity_router/config.py | 13 +- .../auto_router_tuning_baseline.py | 12 +- .../router_strategy/test_complexity_router.py | 148 +++++++++++++----- .../test_auto_router_tuning_baseline.py | 73 ++++++++- .../add_model/ClassificationMethodConfig.tsx | 5 +- .../add_model/ComplexityRouterConfig.test.tsx | 6 +- .../add_model/ComplexityRouterConfig.tsx | 6 + .../add_model/CustomDimensionRows.tsx | 136 ++++++++++++++++ ...euristicScoringConfig.integration.test.tsx | 90 +++++++++++ .../add_model/HeuristicScoringConfig.test.tsx | 14 +- .../add_model/HeuristicScoringConfig.tsx | 85 ++++++++-- .../add_model/add_auto_router_tab.tsx | 4 + .../build_complexity_router_config.test.ts | 47 +++++- .../build_complexity_router_config.ts | 17 +- .../add_model/custom_dimensions.test.ts | 41 +++++ .../components/add_model/custom_dimensions.ts | 54 +++++++ .../add_model/heuristic_scoring_knobs.test.ts | 119 ++++++++++++++ .../add_model/heuristic_scoring_knobs.ts | 97 +++++++++++- .../src/components/add_model/tier_rows.ts | 1 + ...d_updated_complexity_router_config.test.ts | 49 ++++++ .../edit_auto_router_modal.tsx | 17 +- .../src/lib/autorouter_presets.test.ts | 28 ++++ .../src/lib/autorouter_presets.ts | 12 ++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 9 +- 26 files changed, 1047 insertions(+), 86 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/add_model/CustomDimensionRows.tsx create mode 100644 ui/litellm-dashboard/src/components/add_model/HeuristicScoringConfig.integration.test.tsx create mode 100644 ui/litellm-dashboard/src/components/add_model/custom_dimensions.test.ts create mode 100644 ui/litellm-dashboard/src/components/add_model/custom_dimensions.ts diff --git a/litellm/router_strategy/complexity_router/README.md b/litellm/router_strategy/complexity_router/README.md index 88ed374dd3f..d605b43e42a 100644 --- a/litellm/router_strategy/complexity_router/README.md +++ b/litellm/router_strategy/complexity_router/README.md @@ -207,17 +207,27 @@ custom_dimensions: - name: sqlMigration weight: 0.7 patterns: ['\b(create|alter|drop)\s{1,4}table\b'] + - name: dataPipeline + weight: 0.4 + scoring_mode: match_count + keywords: [airflow, dbt, snowflake] ``` Each dimension contributes its weight once when any matcher hits the current ask. Repeated matches do not increase it. The built-in score and tier boundaries are unchanged, and the total score is not renormalized. Keywords use the existing case-insensitive word-boundary and CJK rules. Regexes search the first 2048 characters case-insensitively and compile during configuration validation and router initialization, never per request +`scoring_mode` is optional and defaults to `binary`, the behavior above. `match_count` grades the dimension by how many distinct matchers hit: none scores 0 and emits no signal, one scores half the weight, two or more score the full weight. Repeated occurrences of one matcher never raise the count, keywords are distinct case-insensitively, patterns are distinct by source, and a keyword and a pattern are always distinct from each other. Matching stops as soon as the selected mode's maximum is reached, so a binary dimension still stops at its first hit. Existing configurations without the field keep binary scoring and the same tuning fingerprint, so the field only counts as a tuning change when set to `match_count` + +### Weights through the API versus the dashboard + +The API and YAML store exactly the weights written. A `dimension_weights` map and inline custom weights are read literally, missing recognized built-in names score zero, and nothing renormalizes the vector, so a total other than 1 is legal and scores accordingly. The dashboard's heuristic scoring editor is the one place that rebalances: editing one weight there holds it and redistributes the remainder across the other active dimensions in the draft, then Save sends the resulting explicit values, which the backend stores and scores as written. Opening a router, applying a preset, editing matchers, changing `scoring_mode`, or saving unrelated fields never normalizes existing weights + Only `heuristic`, `heuristic_first` and `hybrid` accept custom dimensions. Each name must be a unique ASCII identifier starting with a letter, at most 64 characters, and cannot reuse a built-in dimension name or a key in `dimension_weights`. Set its weight inline, greater than zero and at most one Patterns are checked at configuration time against a grammar whose worst case stays a few milliseconds on 2048 characters. Every quantifier needs an explicit upper bound of at most 64 and must repeat a single character or character class, so `\s{1,4}` is accepted while `\s+`, `(a|aa){0,12}` and `(?:ab){0,64}` are refused. Backreferences, lookarounds, atomic groups and possessive quantifiers are refused as well. Each pattern is then costed: alternation branches and repeat lengths multiply the ways the engine can retry, and every later piece of the pattern is charged once per path that can reach it, so `a?a?a?a?a?a?a?a?` followed by a long fixed tail is refused even though each quantifier is small. The budget is 2048 work units per pattern and 8192 across the router. An invalid or over-budget pattern fails the write with a message naming the pattern and the rule it broke Limits are 16 dimensions, 32 combined keywords/patterns per dimension, 256 characters per matcher and 4096 matcher characters per dimension. Matching runs inline on the request path with no timeout and no worker thread, because the grammar is what bounds the cost. These are routing hints, not security enforcement rules -The existing heuristic-v1 tuning quota covers custom dimensions and their weights: one changed router without an auto-router license, unlimited with the entitlement. Omitting `custom_dimensions` preserves existing scoring. Routing decisions and spend logs include signals such as `custom (sqlMigration)` without recording the configured pattern or matched text. The field is configured through YAML or the model API; this change adds no dashboard editor +The existing heuristic-v1 tuning quota covers custom dimensions, their weights and their scoring mode: one changed router without an auto-router license, unlimited with the entitlement. Omitting `custom_dimensions` preserves existing scoring. Routing decisions and spend logs include signals such as `custom (sqlMigration)` without recording the configured pattern or matched text. The field is configured through YAML, the model API, or the dashboard's heuristic scoring editor ## Usage diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index c8644f52c57..ae1f82fc40b 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -20,7 +20,7 @@ import random import re import time from collections.abc import Callable, Iterator, Mapping, Sequence -from itertools import accumulate, islice, takewhile +from itertools import accumulate, chain, islice, takewhile from threading import Lock from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, cast @@ -82,6 +82,7 @@ from .config import ( ClassificationRubric, ComplexityRouterConfig, ComplexityTier, + CustomDimension, TierDefinition, ) from .stall_detector import detect_stalled_task @@ -879,6 +880,15 @@ class DimensionScore: self.signal = signal +class _CustomDimensionMatchers(NamedTuple): + """One custom dimension's distinct matchers and the number of hits that saturates its score.""" + + dimension: CustomDimension + keywords: tuple[str, ...] + patterns: tuple[re.Pattern[str], ...] + saturation: int + + class KeywordOverride(NamedTuple): """A keyword_tier_rules match: the winning tier and, on the lexical path, the keyword that fired.""" @@ -1121,7 +1131,12 @@ class ComplexityRouter(CustomLogger): ) self.simple_keywords = self.config.simple_keywords or DEFAULT_SIMPLE_KEYWORDS self._custom_dimensions = tuple( - (dimension, tuple(re.compile(pattern, re.IGNORECASE) for pattern in dimension.patterns)) + _CustomDimensionMatchers( + dimension, + tuple(dict.fromkeys(keyword.lower() for keyword in dimension.keywords)), + tuple(re.compile(pattern, re.IGNORECASE) for pattern in dict.fromkeys(dimension.patterns)), + 2 if dimension.scoring_mode == "match_count" else 1, + ) for dimension in self.config.custom_dimensions ) if self.config.has_custom_tiers: @@ -1325,15 +1340,26 @@ class ComplexityRouter(CustomLogger): score: Final = score_high if match_count >= high_threshold else score_low return DimensionScore(name, score, f"{signal_label} ({detail})"), match_count + def _count_custom_hits(self, matchers: _CustomDimensionMatchers, user_text: str, scanned: str) -> int: + hits: Final = chain( + (self._keyword_matches(user_text, keyword) for keyword in matchers.keywords), + (pattern.search(scanned) is not None for pattern in matchers.patterns), + ) + return sum(islice((1 for hit in hits if hit), matchers.saturation)) + def _score_custom_dimensions(self, prompt: str, user_text: str) -> tuple[tuple[DimensionScore, float], ...]: if not self._custom_dimensions: return () scanned: Final = prompt[:CUSTOM_PATTERN_SCAN_CHARS] return tuple( - (DimensionScore(dimension.name, 1.0, f"custom ({dimension.name})"), dimension.weight) - for dimension, patterns in self._custom_dimensions - if any(self._keyword_matches(user_text, keyword) for keyword in dimension.keywords) - or any(pattern.search(scanned) is not None for pattern in patterns) + ( + DimensionScore( + matchers.dimension.name, hits / matchers.saturation, f"custom ({matchers.dimension.name})" + ), + matchers.dimension.weight, + ) + for matchers in self._custom_dimensions + if (hits := self._count_custom_hits(matchers, user_text, scanned)) ) def _score_multi_step(self, text: str) -> DimensionScore: diff --git a/litellm/router_strategy/complexity_router/config.py b/litellm/router_strategy/complexity_router/config.py index 3ec9f9b5394..37375d9727d 100644 --- a/litellm/router_strategy/complexity_router/config.py +++ b/litellm/router_strategy/complexity_router/config.py @@ -667,6 +667,14 @@ class CustomDimension(BaseModel): weight: float = Field(gt=0, le=1, allow_inf_nan=False) keywords: tuple[Annotated[str, Field(min_length=1, max_length=256)], ...] = Field(default=(), max_length=32) patterns: tuple[Annotated[str, Field(min_length=1, max_length=256)], ...] = Field(default=(), max_length=32) + scoring_mode: Literal["binary", "match_count"] = Field( + default="binary", + description=( + "'binary' scores 1 when any matcher hits. 'match_count' scores 0.5 when one distinct matcher hits and 1 " + "when two or more do; repeated occurrences of one matcher never raise it. Keywords are distinct " + "case-insensitively, patterns by source, and a keyword and a pattern are always distinct from each other." + ), + ) @model_validator(mode="after") def _validate_matchers(self) -> "CustomDimension": @@ -794,8 +802,9 @@ class ComplexityRouterConfig(BaseModel): default=(), max_length=16, description=( - "Named binary dimensions added to the heuristic-v1 score. Each contributes its inline weight once " - "when any keyword matches the current ask or a case-insensitive regex matches its first 2048 characters. " + "Named dimensions added to the heuristic-v1 score. Each contributes its inline weight once " + "when any keyword matches the current ask or a case-insensitive regex matches its first 2048 characters; " + "scoring_mode 'match_count' instead grades half weight for one distinct matcher and full for two or more. " "Regex quantifiers repeat one character or class at most 64 times. Unbounded quantifiers, repeated groups, " "backreferences and lookarounds are rejected. Conservative work limits include alternation paths, " "repeat lengths and subsequent matching: 2048 units per pattern, 8192 across the router. " diff --git a/litellm/router_utils/auto_router_tuning_baseline.py b/litellm/router_utils/auto_router_tuning_baseline.py index 74f7b82389a..9699ab886b9 100644 --- a/litellm/router_utils/auto_router_tuning_baseline.py +++ b/litellm/router_utils/auto_router_tuning_baseline.py @@ -52,7 +52,17 @@ def tuning_fingerprint(complexity_router_config: object) -> str | None: supplied: Final = ((_TUNING_FIELD_SET - frozenset(("tier_model_configs",))) & frozenset(raw)) | ( frozenset(("tier_model_configs",)) if validated.tier_model_configs else frozenset() ) - payload: Final = validated.model_dump(mode="json", include=supplied) + payload: Final = validated.model_dump( + mode="json", + include=supplied, + exclude={ + "custom_dimensions": { + index: {"scoring_mode"} + for index, dimension in enumerate(validated.custom_dimensions) + if dimension.scoring_mode == "binary" + } + }, + ) return hashlib.sha256(json.dumps(payload, sort_keys=True, separators=(",", ":")).encode()).hexdigest() diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index 5b1d8562abd..b0f2a65bb9b 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -826,6 +826,8 @@ class TestCustomDimensions: pytest.param({"keywords": ["x"] * 32, "patterns": ["y"]}, {}, id="combined-matcher-count"), pytest.param({"keywords": ["x" * 256] * 17}, {}, id="matcher-character-budget"), pytest.param({"unknown": True}, {}, id="extra-field"), + pytest.param({"scoring_mode": "graded"}, {}, id="unknown-scoring-mode"), + pytest.param({"scoring_mode": None}, {}, id="null-scoring-mode"), ], ) def test_custom_dimension_invalid_configuration_rejected( @@ -878,43 +880,125 @@ class TestCustomDimensions: ) @pytest.mark.asyncio - @pytest.mark.parametrize("current_ask", ("Hello!", "orbitmesh")) + @pytest.mark.parametrize("scoring_mode", ("binary", "match_count")) + @pytest.mark.parametrize("current_ask", ("Hello!", "orbitmesh", "orbitmesh fluxgate")) async def test_custom_dimensions_public_hook_scores_only_current_ask( - self, mock_router_instance: MagicMock, current_ask: str + self, mock_router_instance: MagicMock, current_ask: str, scoring_mode: str ) -> None: router: Final = ComplexityRouter( "test-router", mock_router_instance, { "tiers": {"SIMPLE": "cheap", "MEDIUM": "mid", "COMPLEX": "strong", "REASONING": "top"}, - "custom_dimensions": [{"name": "internalFrameworks", "weight": 0.7, "keywords": ["orbitmesh"]}], + "dimension_weights": {}, + "custom_dimensions": [ + { + "name": "internalFrameworks", + "weight": 0.8, + "keywords": ["orbitmesh", "fluxgate"], + "scoring_mode": scoring_mode, + } + ], }, ) result: Final = await router.async_pre_routing_hook( model="test-router", request_kwargs={}, messages=[ - {"role": "system", "content": "orbitmesh"}, - {"role": "user", "content": "orbitmesh"}, - {"role": "assistant", "content": "orbitmesh is ready"}, + {"role": "system", "content": "orbitmesh fluxgate"}, + {"role": "user", "content": "orbitmesh fluxgate"}, + {"role": "assistant", "content": "orbitmesh fluxgate is ready"}, {"role": "user", "content": current_ask}, - {"role": "tool", "tool_call_id": "previous", "content": "orbitmesh"}, + {"role": "tool", "tool_call_id": "previous", "content": "orbitmesh fluxgate"}, ], ) assert result is not None assert result.routing_decision is not None - assert ("custom (internalFrameworks)" in result.routing_decision["signals"]) is (current_ask == "orbitmesh") - assert result.model == ("top" if current_ask == "orbitmesh" else "cheap") + expected_score: Final = ( + 0.0 + if current_ask == "Hello!" + else 0.4 + if scoring_mode == "match_count" and current_ask == "orbitmesh" + else 0.8 + ) + assert result.routing_decision["score"] == expected_score + assert ("custom (internalFrameworks)" in result.routing_decision["signals"]) is (expected_score > 0) + assert result.model == ("cheap" if expected_score == 0 else "strong" if expected_score == 0.4 else "top") assert "orbitmesh" not in " ".join(result.routing_decision["signals"]) - def test_custom_patterns_scan_only_the_first_2048_characters(self, mock_router_instance: MagicMock) -> None: + @pytest.mark.parametrize("scoring_mode", ("binary", "match_count")) + def test_custom_patterns_scan_only_the_first_2048_characters( + self, mock_router_instance: MagicMock, scoring_mode: str + ) -> None: router: Final = ComplexityRouter( "test-router", mock_router_instance, - {"custom_dimensions": [{"name": "late", "weight": 0.7, "patterns": [r"zzz{1,3}"]}]}, + { + "custom_dimensions": [ + { + "name": "late", + "weight": 0.7, + "patterns": [r"zzz{1,3}", r"yyy{1,3}"], + "scoring_mode": scoring_mode, + } + ] + }, ) + baseline: Final = ComplexityRouter("test-router", mock_router_instance) assert "custom (late)" in router.classify("a" * 2040 + " zzz")[2] assert "custom (late)" not in router.classify("a" * 2048 + " zzz")[2] + second_hit_past_the_bound: Final = "yyy " + "a" * 2044 + " zzz" + contribution: Final = ( + router.classify(second_hit_past_the_bound)[1] - baseline.classify(second_hit_past_the_bound)[1] + ) + assert contribution == pytest.approx(0.7 if scoring_mode == "binary" else 0.35) + + @pytest.mark.parametrize( + "prompt,expected_score", + [ + pytest.param("Hello!", 0.0, id="no-hit"), + pytest.param("orbitmesh orbitmesh ORBITMESH again", 0.5, id="one-keyword-repeated"), + pytest.param("create table a; CREATE TABLE b; create table c", 0.5, id="one-pattern-repeated"), + pytest.param("orbitmesh and fluxgate", 1.0, id="two-keywords"), + pytest.param("orbitmesh then create table t", 1.0, id="keyword-plus-pattern"), + pytest.param("create table a; alter table b", 1.0, id="two-patterns"), + pytest.param("orbitmesh fluxgate create table a alter table b", 1.0, id="all-matchers"), + ], + ) + def test_match_count_grades_distinct_matchers( + self, mock_router_instance: MagicMock, prompt: str, expected_score: float + ) -> None: + dimension: Final = { + "name": "graded", + "weight": 0.6, + "keywords": ["orbitmesh", "ORBITMESH", "fluxgate"], + "patterns": [r"\bcreate\s{1,4}table\b", r"\bcreate\s{1,4}table\b", r"\balter\s{1,4}table\b"], + } + baseline: Final = ComplexityRouter("test-router", mock_router_instance) + binary: Final = ComplexityRouter("test-router", mock_router_instance, {"custom_dimensions": [dimension]}) + graded: Final = ComplexityRouter( + "test-router", + mock_router_instance, + {"custom_dimensions": [{**dimension, "scoring_mode": "match_count"}]}, + ) + _, baseline_score, baseline_signals = baseline.classify(prompt) + _, binary_score, binary_signals = binary.classify(prompt) + _, graded_score, graded_signals = graded.classify(prompt) + assert graded_score == pytest.approx(baseline_score + 0.6 * expected_score) + assert binary_score == pytest.approx(baseline_score + (0.6 if expected_score else 0.0)) + expected_signals: Final = [*baseline_signals, *(["custom (graded)"] if expected_score else [])] + assert graded_signals == expected_signals + assert binary_signals == expected_signals + + def test_scoring_mode_round_trips_and_defaults_to_binary(self) -> None: + dimension: Final = {"name": "graded", "weight": 0.6, "keywords": ["orbitmesh"]} + legacy: Final = ComplexityRouterConfig.model_validate({"custom_dimensions": [dimension]}) + graded: Final = ComplexityRouterConfig.model_validate( + {"custom_dimensions": [{**dimension, "scoring_mode": "match_count"}]} + ) + assert legacy.custom_dimensions[0].scoring_mode == "binary" + assert graded.model_dump(mode="json")["custom_dimensions"][0]["scoring_mode"] == "match_count" + assert ComplexityRouterConfig.model_validate(graded.model_dump(mode="json")) == graded def test_custom_dimensions_router_wide_regex_work_is_capped(self) -> None: heavy: Final = {"weight": 0.5, "patterns": ["a?" * 8 + "z"]} @@ -1511,6 +1595,7 @@ class TestRouterComplexityDeploymentMethods: def test_the_shipped_rubric_and_default_prompt_stay_free(self) -> None: """Only an operator-written prompt is gated: picking a shipped rubric preset, or writing no prompt at all, leaves a router unmetered, so several of them register under a ceiling of one.""" + def rubric(model_name: str, model_id: str, preset: str | None) -> dict[str, object]: llm_config: dict[str, object] = {"model": "gpt-4o-mini"} if preset is not None: @@ -1647,6 +1732,7 @@ class TestRouterComplexityDeploymentMethods: def test_renaming_built_in_tiers_is_not_a_custom_tier_set(self) -> None: """tier_labels renames the built-in ladder without defining one, so it stays ungated: two such routers register under a ceiling of one.""" + def labeled(model_name: str, model_id: str) -> dict[str, object]: row = self._router_row(model_name, model_id, "heuristic") row["litellm_params"]["complexity_router_config"]["tier_labels"] = {"SIMPLE": "Cheap", "MEDIUM": "Standard"} @@ -2520,9 +2606,7 @@ class TestLLMClassifier: assert outcome.classifier_cost == pytest.approx(1.35e-05) @pytest.mark.asyncio - async def test_aclassify_timeout_does_not_inherit_router_retries_or_fallbacks( - self, llm_classifier_config - ): + async def test_aclassify_timeout_does_not_inherit_router_retries_or_fallbacks(self, llm_classifier_config): real_router = Router( model_list=[ { @@ -2565,9 +2649,7 @@ class TestLLMClassifier: assert real_router.total_calls["openai/mock-backup-classifier"] == 0 @pytest.mark.asyncio - async def test_aclassify_enforces_total_classifier_deadline( - self, mock_router_instance, llm_classifier_config - ): + async def test_aclassify_enforces_total_classifier_deadline(self, mock_router_instance, llm_classifier_config): cancelled = asyncio.Event() async def slow_classifier(**_kwargs: object) -> None: @@ -12414,9 +12496,7 @@ class TestTierHealthFailover: llm_provider="", ) filtered = (*cooling, *blocked, *excluded) - healthy = [ - {"model_name": model, "model_info": {"id": i}} for i in ids_by_model[model] if i not in filtered - ] + healthy = [{"model_name": model, "model_info": {"id": i}} for i in ids_by_model[model] if i not in filtered] if not healthy: raise RouterRateLimitError( model=model, cooldown_time=60.0, enable_pre_call_checks=False, cooldown_list=[] @@ -12845,9 +12925,7 @@ class TestTierHealthFailover: assert all(probed is not request_kwargs for probed in router.litellm_router_instance.probed_kwargs) @pytest.mark.asyncio - async def test_a_peer_whose_every_deployment_is_over_its_rpm_is_not_a_failover_target( - self, mock_router_instance - ): + async def test_a_peer_whose_every_deployment_is_over_its_rpm_is_not_a_failover_target(self, mock_router_instance): """RPM exhaustion is its own verdict from the owner (RouterRateLimitErrorBasic). A peer in that state would be rejected downstream, so it cannot be the substitute.""" from litellm.types.router import RouterRateLimitErrorBasic @@ -12880,9 +12958,7 @@ class TestTierHealthFailover: assert {r.model for r in results} == {"live-c"} @pytest.mark.asyncio - async def test_the_probe_forwards_input_so_window_checks_run_on_input_only_surfaces( - self, mock_router_instance - ): + async def test_the_probe_forwards_input_so_window_checks_run_on_input_only_surfaces(self, mock_router_instance): """The Responses API carries its prompt as `input`, never as messages. The owner only runs its context-window pre-call check when one of them is present, so dropping `input` would silently skip window filtering on that whole surface.""" @@ -12908,9 +12984,7 @@ class TestTierHealthFailover: ), "the eligibility probe must forward `input` to the owner" @pytest.mark.asyncio - async def test_a_group_the_router_has_no_deployment_for_is_not_a_failover_target( - self, mock_router_instance - ): + async def test_a_group_the_router_has_no_deployment_for_is_not_a_failover_target(self, mock_router_instance): """The owner answers an unconfigured group with BadRequestError. Reading that as live would both skip failover off it and let it be chosen as a substitute.""" router = self._router( @@ -13099,9 +13173,7 @@ class TestClassifierVision: routed as default_fallback on text the request never contained. """ router = self._router(mock_router_instance, vision={"enabled": True}) - response = await router.async_pre_routing_hook( - model="m", request_kwargs={}, messages=self._turn(IMG_PART) - ) + response = await router.async_pre_routing_hook(model="m", request_kwargs={}, messages=self._turn(IMG_PART)) assert response.routing_decision["cause"] == "llm_classifier" assert response.model == "t-complex" assert [block["type"] for block in self._classifier_user_content(mock_router_instance)] == [ @@ -13112,9 +13184,7 @@ class TestClassifierVision: @pytest.mark.asyncio async def test_image_only_turn_still_falls_back_when_vision_is_off(self, mock_router_instance): router = self._router(mock_router_instance, vision={"enabled": False}) - response = await router.async_pre_routing_hook( - model="m", request_kwargs={}, messages=self._turn(IMG_PART) - ) + response = await router.async_pre_routing_hook(model="m", request_kwargs={}, messages=self._turn(IMG_PART)) assert response.routing_decision["cause"] == "default_fallback" mock_router_instance.acompletion.assert_not_awaited() @@ -13184,9 +13254,7 @@ class TestClassifierVision: makes the image the only variable; a margin loose enough to leave the score undecided would pass whether or not the guard exists. """ - router = self._router( - mock_router_instance, vision={"enabled": True}, classifier_type=classifier_type, **extra - ) + router = self._router(mock_router_instance, vision={"enabled": True}, classifier_type=classifier_type, **extra) response = await router.async_pre_routing_hook( model="m", request_kwargs={}, messages=self._turn({"type": "text", "text": "what is this"}, IMG_PART) ) @@ -13200,9 +13268,7 @@ class TestClassifierVision: self, mock_router_instance, classifier_type, extra, short_circuit_cause ): """The negative class: same router, same text, no image, and the scorer still decides.""" - router = self._router( - mock_router_instance, vision={"enabled": True}, classifier_type=classifier_type, **extra - ) + router = self._router(mock_router_instance, vision={"enabled": True}, classifier_type=classifier_type, **extra) response = await router.async_pre_routing_hook( model="m", request_kwargs={}, messages=[{"role": "user", "content": "what is this"}] ) diff --git a/tests/test_litellm/router_utils/test_auto_router_tuning_baseline.py b/tests/test_litellm/router_utils/test_auto_router_tuning_baseline.py index f686a62db76..75c115cd3ab 100644 --- a/tests/test_litellm/router_utils/test_auto_router_tuning_baseline.py +++ b/tests/test_litellm/router_utils/test_auto_router_tuning_baseline.py @@ -7,6 +7,7 @@ from typing import Final import pytest +from litellm.router_strategy.complexity_router.config import ComplexityRouterConfig from litellm.router_utils.auto_router_tuning_baseline import ( DEFAULT_TUNING_FINGERPRINT, HEURISTIC_V1_TUNING_FIELDS, @@ -21,6 +22,33 @@ from litellm.router_utils.auto_router_tuning_baseline import ( _TIERS = {"SIMPLE": "cheap", "MEDIUM": "mid", "COMPLEX": "strong", "REASONING": "top"} _ALT_TIERS = {**_TIERS, "COMPLEX": "other-strong"} +_KEYWORD_DIMENSION: Final = {"name": "internalFrameworks", "weight": 0.7, "keywords": ["orbitmesh"]} +_HISTORICAL_FINGERPRINTS: Final = ( + ({}, "44136fa355b3678a1146ad16f7e8649e94fb4fc21fe77e8310c060f61caaff8a"), + ( + {"custom_dimensions": [_KEYWORD_DIMENSION]}, + "b5c3c3f3be6341a8a16148d68d9067e03f94ed01a0bbfcde7955763042744372", + ), + ( + {"custom_dimensions": [{"name": "sqlDdl", "weight": 0.4, "patterns": [r"\bCREATE\s{1,4}TABLE\b"]}]}, + "814ce0017fc7f60a160b262f658d910e9bdf784e6139a4ba4f1e2657aa203950", + ), + ( + { + "tiers": _TIERS, + "dimension_weights": {"codePresence": 0.3}, + "custom_dimensions": [ + { + "name": "internalFrameworks", + "weight": 0.2, + "keywords": ["orbitmesh", "fluxgate"], + "patterns": [r"\bALTER\s{1,4}TABLE\b"], + } + ], + }, + "38970dc9224e265ab38c89674563d8d0537822591f9239b45251db6f5ca6cc39", + ), +) def _router( @@ -77,6 +105,27 @@ class TestTuningFingerprint: def test_explicit_empty_tier_model_configs_follow_omission(self) -> None: assert tuning_fingerprint({"tier_model_configs": {}}) == DEFAULT_TUNING_FINGERPRINT + @pytest.mark.parametrize(("config", "fingerprint"), _HISTORICAL_FINGERPRINTS) + def test_fingerprints_recorded_before_scoring_mode_existed_are_preserved( + self, config: Mapping[str, object], fingerprint: str + ) -> None: + """Literal hashes captured from the merged implementation at 9bc9104102, before CustomDimension.scoring_mode.""" + assert tuning_fingerprint(config) == fingerprint + + def test_binary_scoring_mode_hashes_like_its_absence(self) -> None: + historical: Final = tuning_fingerprint({"custom_dimensions": [_KEYWORD_DIMENSION]}) + explicit: Final = tuning_fingerprint({"custom_dimensions": [{**_KEYWORD_DIMENSION, "scoring_mode": "binary"}]}) + reserialized: Final = ComplexityRouterConfig.model_validate( + {"custom_dimensions": [_KEYWORD_DIMENSION]} + ).model_dump(mode="json", include={"custom_dimensions"}) + assert reserialized["custom_dimensions"][0]["scoring_mode"] == "binary" + assert reserialized["custom_dimensions"][0]["patterns"] == [] + assert historical == explicit == tuning_fingerprint(reserialized) + assert ( + tuning_fingerprint({"custom_dimensions": [{**_KEYWORD_DIMENSION, "scoring_mode": "match_count"}]}) + != historical + ) + def test_tier_model_overrides_change_the_fingerprint(self) -> None: plain = tuning_fingerprint({"tiers": {"SIMPLE": "x"}}) with_override = tuning_fingerprint( @@ -218,15 +267,18 @@ class TestQuota: is None ) - def test_custom_dimension_add_edit_and_revert_share_one_quota_slot(self) -> None: + @pytest.mark.parametrize( + "edit", + [ + pytest.param({"weight": 0.9}, id="weight"), + pytest.param({"scoring_mode": "match_count"}, id="scoring-mode"), + ], + ) + def test_custom_dimension_add_edit_and_revert_share_one_quota_slot(self, edit: Mapping[str, object]) -> None: baselines: Final = snapshot_tuning_baselines(()) original: Final = _router("a", {}) - config: Final = { - "custom_dimensions": [{"name": "internalFrameworks", "weight": 0.7, "keywords": ["orbitmesh"]}] - } - edited_config: Final = { - "custom_dimensions": [{"name": "internalFrameworks", "weight": 0.9, "keywords": ["orbitmesh"]}] - } + config: Final = {"custom_dimensions": [_KEYWORD_DIMENSION]} + edited_config: Final = {"custom_dimensions": [{**_KEYWORD_DIMENSION, **edit}]} added: Final = _router("a", config) edited: Final = _router("a", edited_config) second: Final = _router("b", config) @@ -240,6 +292,13 @@ class TestQuota: assert mutable_tuned_identities((original,), baselines) == frozenset() assert tuning_quota_violation(candidate=second, others=(original,), baselines=baselines, limit=1) is None + def test_graded_dimension_recorded_at_snapshot_is_its_own_baseline(self) -> None: + graded: Final = _router("a", {"custom_dimensions": [{**_KEYWORD_DIMENSION, "scoring_mode": "match_count"}]}) + baselines: Final = snapshot_tuning_baselines((graded,)) + assert mutable_tuned_identities((graded,), baselines) == frozenset() + reverted_to_binary: Final = _router("a", {"custom_dimensions": [_KEYWORD_DIMENSION]}) + assert mutable_tuned_identities((reverted_to_binary,), baselines) == {router_identity(graded)} + def test_violation_message_names_the_limit_and_remedy(self) -> None: message = tuning_limit_violation(held=2, limit=1) assert message is not None diff --git a/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx index 15cfff01766..111eeebe5a3 100644 --- a/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ClassificationMethodConfig.tsx @@ -44,8 +44,9 @@ import { } from "./ComplexityRouterConfig"; const DEFAULT_SCORING_EXPLANATION = - "The router scores each request across 7 dimensions: token count, code presence, reasoning markers, technical " + - "terms, simple indicators, multi-step patterns, and question complexity. The weighted score determines the tier:"; + "The router scores each request across 7 built-in dimensions: token count, code presence, reasoning markers, technical " + + "terms, simple indicators, multi-step patterns, and question complexity, plus any custom dimensions you add. " + + "The weighted score determines the tier:"; const HEURISTIC_V2_EXPLANATION = "The router estimates success probability for all four tiers with the bundled calibrated model, then selects " + diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx index 2970e14b335..0c12cc0ba1a 100644 --- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx @@ -1521,14 +1521,16 @@ describe("ComplexityRouterConfig tier editing", () => { renderWithProviders(); fireEvent.click(screen.getByText("Advanced: Classification Method")); expect(screen.queryByText("How Classification Works")).not.toBeInTheDocument(); - expect(screen.queryByText("scores each request across 7 dimensions", { exact: false })).not.toBeInTheDocument(); + expect( + screen.queryByText("scores each request across 7 built-in dimensions", { exact: false }), + ).not.toBeInTheDocument(); }); it("keeps the scorer card on a built-in router, whose tiers the score still decides", () => { renderWithProviders(); fireEvent.click(screen.getByText("Advanced: Classification Method")); expect(screen.getByText("How Classification Works")).toBeInTheDocument(); - expect(screen.getByText("scores each request across 7 dimensions", { exact: false })).toBeInTheDocument(); + expect(screen.getByText("scores each request across 7 built-in dimensions", { exact: false })).toBeInTheDocument(); }); it("says why a custom row is blocked instead of only reddening its border", () => { diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx index d7df6ce33bb..de80e714e66 100644 --- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx @@ -50,6 +50,7 @@ import EscalationKeywords from "./EscalationKeywords"; import KeywordTierRules, { KeywordTierRule } from "./KeywordTierRules"; import SemanticKeywordMatching from "./SemanticKeywordMatching"; import { type DimensionWeights, type TierBoundaries, type TokenThresholds } from "./heuristic_scoring_knobs"; +import { type CustomDimensionRow } from "./custom_dimensions"; import CompressionControls from "./CompressionControls"; import { type AutoRouterCompressionState, DEFAULT_AUTO_ROUTER_COMPRESSION } from "./buildAutoRouterCompression"; @@ -432,6 +433,11 @@ export interface ComplexityRouterConfigValue { tier_boundaries?: TierBoundaries; token_thresholds?: TokenThresholds; dimension_weights?: DimensionWeights; + /** + * Operator-added scoring dimensions, each carrying its own inline weight. Undefined means the router has + * none and keeps the key out of the payload; an empty array is a real "the last row was removed" state. + */ + custom_dimensions?: CustomDimensionRow[]; /** * Score floor the reasoning-marker override must clear. Undefined keeps the key out of the payload, so the * floor tracks tier_boundaries.simple_medium; an explicit 0 is a real floor that promotes on the markers alone. diff --git a/ui/litellm-dashboard/src/components/add_model/CustomDimensionRows.tsx b/ui/litellm-dashboard/src/components/add_model/CustomDimensionRows.tsx new file mode 100644 index 00000000000..34d8370aafb --- /dev/null +++ b/ui/litellm-dashboard/src/components/add_model/CustomDimensionRows.tsx @@ -0,0 +1,136 @@ +import { Trash2 } from "lucide-react"; +import { useState } from "react"; +import { Button } from "@/components/ui/button"; +import { Input } from "@/components/ui/input"; +import { Label } from "@/components/ui/label"; +import { Textarea } from "@/components/ui/textarea"; +import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; +import { Slider } from "@/components/ui/slider"; +import type { CustomDimensionRow } from "./custom_dimensions"; + +const SCORING_MODES = [ + { value: "binary", label: "Binary" }, + { value: "match_count", label: "Match count" }, +] as const; + +interface Props { + rows: CustomDimensionRow[]; + disabled: boolean; + onChange: (rows: CustomDimensionRow[]) => void; + onWeight: (id: string, weight: number) => void; + onAdd: () => void; + onRemove: (id: string) => void; +} + +export default function CustomDimensionRows({ rows, disabled, onChange, onWeight, onAdd, onRemove }: Props) { + const [draft, setDraft] = useState<{ id: string; raw: string } | null>(null); + const update = (id: string, patch: Partial) => + onChange(rows.map((row) => (row.id === id ? { ...row, ...patch } : row))); + const editWeight = (id: string, raw: string) => { + setDraft({ id, raw }); + if (raw.trim() && Number.isFinite(Number(raw))) onWeight(id, Number(raw)); + }; + return ( +
+ {rows.map((row, index) => ( +
+ Custom dimension {index + 1} +
+ +
+
+ + update(row.id, { name: event.target.value })} + /> +
+
+ + onWeight(row.id, Array.isArray(value) ? value[0] : value)} + /> + setDraft(null)} + onChange={(event) => editWeight(row.id, event.target.value)} + /> +
+
+ {(["keywords", "patterns"] as const).map((field) => ( +
+ +