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 001/109] 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 002/109] 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 003/109] 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 004/109] 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 005/109] 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 006/109] 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 007/109] 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 008/109] 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 009/109] 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 010/109] 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 011/109] 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 012/109] 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 013/109] 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 014/109] 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 015/109] 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 016/109] 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 017/109] 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 018/109] 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 019/109] 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 020/109] 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 55c7872496b5568bc259fc88c0645f702af1576e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 12:57:06 -0700 Subject: [PATCH 021/109] fix(proxy-extras): rebuild indexes left INVALID by a migration deadlock Two replicas racing prisma migrate deploy can deadlock, and the loser dies mid CREATE INDEX CONCURRENTLY, leaving the index INVALID. The retried migration's IF NOT EXISTS then skips it, so the planner never uses it. After migrations succeed, look for INVALID indexes on LiteLLM tables and have one replica (advisory try-lock) REINDEX INDEX CONCURRENTLY each of them, dropping _ccnew/_ccold leftovers of an interrupted rebuild instead. The repair never blocks startup: a failed rebuild is logged and retried on the next boot. Also encode DATABASE_URL query values with quote instead of quote_plus so options=-c%20... reaches psycopg intact. --- .../litellm_proxy_extras/utils.py | 128 +++++++++- .../test_invalid_index_repair.py | 238 ++++++++++++++++++ 2 files changed, 359 insertions(+), 7 deletions(-) create mode 100644 tests/proxy_migration_tests/test_invalid_index_repair.py diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index b8032dd0d28..df019d78078 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -6,18 +6,23 @@ import shutil import subprocess import tempfile import time +from dataclasses import dataclass from pathlib import Path -from typing import Optional +from typing import TYPE_CHECKING, Final, Optional from litellm_proxy_extras._logging import logger -from litellm_proxy_extras.replica_identity import ( - REPLICA_IDENTITY_FULL_ENV_VAR, - apply_replica_identity_full, -) from litellm_proxy_extras.prisma_toolchain import ( ensure_prisma_toolchain, prisma_command_timeout, ) +from litellm_proxy_extras.replica_identity import ( + REPLICA_IDENTITY_FULL_ENV_VAR, + apply_replica_identity_full, +) + +if TYPE_CHECKING: + import psycopg + import psycopg.sql def str_to_bool(value: Optional[str]) -> bool: @@ -40,6 +45,29 @@ def _get_prisma_env() -> dict: _MIGRATION_TS_RE = re.compile(r"^(\d{14})_") +INDEX_REPAIR_ADVISORY_LOCK_KEY: Final = int.from_bytes(b"litellm", "big") +_TRANSIENT_INDEX_SUFFIX_RE: Final = re.compile(r"_cc(?:new|old)\d*$") +_INVALID_LITELLM_INDEXES_SQL: Final = ( + "SELECT n.nspname, c.relname, pg_size_pretty(pg_table_size(t.oid)) " + "FROM pg_index i " + "JOIN pg_class c ON c.oid = i.indexrelid " + "JOIN pg_class t ON t.oid = i.indrelid " + "JOIN pg_namespace n ON n.oid = t.relnamespace " + "WHERE NOT i.indisvalid " + " AND c.relkind = 'i' " + " AND n.nspname = %s " + " AND t.relname LIKE %s " + " AND NOT EXISTS (SELECT 1 FROM pg_constraint k WHERE k.conindid = i.indexrelid) " + "ORDER BY c.relname" +) + + +@dataclass(frozen=True, slots=True) +class _InvalidIndex: + schema: str + name: str + table_size: str + _SPEND_LOGS_ALTER_RE = re.compile(r'^ALTER\s+TABLE\s+"LiteLLM_SpendLogs"\s', re.IGNORECASE) _SPEND_LOGS_ARTIFACT_DROP_RE = re.compile( r'^DROP\s+TABLE\s+"LiteLLM_SpendLogs_[^"]*"', re.IGNORECASE @@ -557,7 +585,7 @@ class ProxyExtrasDBManager: def _strip_prisma_query_params(url: str) -> str: """Remove Prisma-specific query params (connection_limit, pool_timeout, schema, etc.) from DATABASE_URL so psycopg can parse it.""" - from urllib.parse import urlparse, urlunparse, parse_qsl, urlencode + from urllib.parse import parse_qsl, quote, urlencode, urlparse, urlunparse parsed = urlparse(url) if not parsed.query: @@ -578,7 +606,7 @@ class ProxyExtrasDBManager: "target_session_attrs", } kept = [(k, v) for k, v in parse_qsl(parsed.query) if k in libpq_params] - return urlunparse(parsed._replace(query=urlencode(kept))) + return urlunparse(parsed._replace(query=urlencode(kept, quote_via=quote))) @staticmethod def _warn_if_db_ahead_of_head(migrations_dir: str) -> None: @@ -652,6 +680,91 @@ class ProxyExtrasDBManager: ", ".join(sorted_hostile[:5]) + (" ..." if len(sorted_hostile) > 5 else ""), ) + @staticmethod + def _invalid_litellm_indexes( + conn: "psycopg.Connection[tuple[str, str, str]]", schema: str + ) -> tuple[_InvalidIndex, ...]: + rows: Final = conn.execute(_INVALID_LITELLM_INDEXES_SQL, (schema, "LiteLLM\\_%")).fetchall() + return tuple(_InvalidIndex(*row) for row in rows) + + @staticmethod + def _index_repair(index: _InvalidIndex) -> tuple["psycopg.sql.Composed", str]: + from psycopg import sql + + target: Final = sql.Identifier(index.schema, index.name) + if _TRANSIENT_INDEX_SUFFIX_RE.search(index.name): + return sql.SQL("DROP INDEX CONCURRENTLY IF EXISTS {}").format(target), "Dropped leftover" + return sql.SQL("REINDEX INDEX CONCURRENTLY {}").format(target), "Rebuilt" + + @staticmethod + def _repair_index(conn: "psycopg.Connection[tuple[str, str, str]]", index: _InvalidIndex) -> None: + import psycopg + + statement, action = ProxyExtrasDBManager._index_repair(index) + try: + conn.execute(statement) + except psycopg.Error as e: + logger.warning( + "Could not repair invalid index %s.%s, will retry on the next startup. " + "If this keeps happening, run `%s` by hand as the index owner. Error: %s", + index.schema, + index.name, + statement.as_string(conn), + e, + ) + return + logger.info("%s invalid index %s.%s", action, index.schema, index.name) + + @staticmethod + def repair_invalid_indexes(lock_timeout: str = "30s") -> bool: + """Rebuild LiteLLM indexes an interrupted CREATE INDEX CONCURRENTLY left + INVALID (a migration deadlock between replicas is the usual cause; the + retried migration skips them because of IF NOT EXISTS). Never raises: + returns True when no invalid index remains, False when the repair was + skipped or failed and will be retried on the next startup.""" + database_url: Final = os.getenv("DATABASE_URL") + if not database_url: + return False + + try: + import psycopg + from psycopg import sql + except ImportError: + logger.warning( + "psycopg is not installed; skipping the invalid index check. " + "Install the litellm[extra_proxy] extra, which includes psycopg." + ) + return False + + schema: Final = ProxyExtrasDBManager._prisma_schema_param(database_url) or "public" + cleaned_url: Final = ProxyExtrasDBManager._strip_prisma_query_params(database_url) + try: + with psycopg.connect(cleaned_url, connect_timeout=10, autocommit=True) as conn: + conn.execute("SET statement_timeout = 0") + conn.execute(sql.SQL("SET lock_timeout = {}").format(sql.Literal(lock_timeout))) + found: Final = ProxyExtrasDBManager._invalid_litellm_indexes(conn, schema) + if not found: + return True + logger.warning( + "Found %d invalid index(es) left by an interrupted CREATE INDEX " + "CONCURRENTLY, rebuilding: %s", + len(found), + ", ".join(f"{index.name} (table size {index.table_size})" for index in found), + ) + lock_row: Final = conn.execute( + "SELECT pg_try_advisory_lock(%s)", (INDEX_REPAIR_ADVISORY_LOCK_KEY,) + ).fetchone() + if lock_row is None or not lock_row[0]: + logger.info("Another replica is already rebuilding the invalid indexes, skipping") + return False + for index in ProxyExtrasDBManager._invalid_litellm_indexes(conn, schema): + ProxyExtrasDBManager._repair_index(conn, index) + remaining: Final = ProxyExtrasDBManager._invalid_litellm_indexes(conn, schema) + except psycopg.Error as e: + logger.warning("Could not check for invalid indexes, will retry on the next startup. Error: %s", e) + return False + return not remaining + @staticmethod def _setup_database_v2(use_migrate: bool) -> bool: """ @@ -886,6 +999,7 @@ class ProxyExtrasDBManager: use_migrate=use_migrate, use_v2_resolver=use_v2_resolver ) if migrated: + ProxyExtrasDBManager.repair_invalid_indexes() ProxyExtrasDBManager.apply_replica_identity_full_if_requested() return migrated diff --git a/tests/proxy_migration_tests/test_invalid_index_repair.py b/tests/proxy_migration_tests/test_invalid_index_repair.py new file mode 100644 index 00000000000..787a62b361b --- /dev/null +++ b/tests/proxy_migration_tests/test_invalid_index_repair.py @@ -0,0 +1,238 @@ +import os +import subprocess +import threading +import uuid +from collections.abc import Iterator, Mapping +from pathlib import Path +from types import MappingProxyType +from typing import Final + +import pytest +from litellm_proxy_extras.utils import INDEX_REPAIR_ADVISORY_LOCK_KEY, ProxyExtrasDBManager + +psycopg = pytest.importorskip("psycopg") + +pytestmark = pytest.mark.timeout(120) + +requires_db: Final = pytest.mark.skipif( + "DATABASE_URL" not in os.environ, + reason="requires a postgres database (DATABASE_URL)", +) + +HEALTH_TABLE: Final = "LiteLLM_HealthCheckTable" +HEALTH_INDEX: Final = "LiteLLM_HealthCheckTable_model_id_model_name_checked_at_idx" +HEALTH_INDEX_COLUMNS: Final = '"model_id", "model_name", "checked_at" DESC' +LOOKALIKE_TABLE: Final = "LiteLLMLookalikeTable" +LOOKALIKE_INDEX: Final = "LiteLLMLookalikeTable_id_idx" +PARTITIONED_TABLE: Final = "LiteLLM_PartitionedTable" +PARTITIONED_INDEX: Final = "LiteLLM_PartitionedTable_id_idx" + + +def _base_url() -> str: + return os.environ["DATABASE_URL"].split("?")[0] + + +def _index_validity(schema: str) -> Mapping[str, bool]: + with psycopg.connect(_base_url(), autocommit=True) as conn: + rows = conn.execute( + "SELECT c.relname, i.indisvalid FROM pg_index i " + "JOIN pg_class c ON c.oid = i.indexrelid " + "JOIN pg_namespace n ON n.oid = c.relnamespace " + "WHERE n.nspname = %s", + (schema,), + ).fetchall() + return MappingProxyType(dict(rows)) + + +def _interrupt_concurrent_build(schema: str, table: str, statement: str) -> None: + """Abort a CONCURRENTLY build while it waits on an older snapshot, the same + spot the deadlock loser dies at, so it leaves its index INVALID.""" + with psycopg.connect(_base_url()) as pin: + pin.isolation_level = psycopg.IsolationLevel.REPEATABLE_READ + pin.execute(f'SELECT count(*) FROM "{schema}"."{table}"') + with psycopg.connect(_base_url(), autocommit=True) as builder: + builder.execute("SET statement_timeout = '1s'") + with pytest.raises(psycopg.errors.QueryCanceled): + builder.execute(statement) + + +def _leave_invalid_index(schema: str, table: str, index: str, columns: str) -> None: + _interrupt_concurrent_build( + schema, table, f'CREATE INDEX CONCURRENTLY "{index}" ON "{schema}"."{table}" ({columns})' + ) + + +def _leave_invalid_reindex_leftover(schema: str, table: str, index: str) -> None: + _interrupt_concurrent_build(schema, table, f'REINDEX INDEX CONCURRENTLY "{schema}"."{index}"') + + +@pytest.fixture +def scratch_schema(monkeypatch: pytest.MonkeyPatch) -> Iterator[str]: + schema: Final = f"invalid_index_{uuid.uuid4().hex[:8]}" + with psycopg.connect(_base_url(), autocommit=True) as conn: + conn.execute(f'CREATE SCHEMA "{schema}"') + conn.execute( + f'CREATE TABLE "{schema}"."{HEALTH_TABLE}" (model_id TEXT, model_name TEXT, checked_at TIMESTAMPTZ)' + ) + conn.execute(f'CREATE TABLE "{schema}"."{LOOKALIKE_TABLE}" (id TEXT)') + + monkeypatch.setenv("DATABASE_URL", f"{_base_url()}?schema={schema}") + yield schema + + with psycopg.connect(_base_url(), autocommit=True) as conn: + conn.execute(f'DROP SCHEMA "{schema}" CASCADE') + + +@requires_db +def test_repair_rebuilds_invalid_litellm_indexes_and_leaves_lookalike_tables_alone(scratch_schema: str) -> None: + _leave_invalid_index(scratch_schema, HEALTH_TABLE, HEALTH_INDEX, HEALTH_INDEX_COLUMNS) + _leave_invalid_index(scratch_schema, LOOKALIKE_TABLE, LOOKALIKE_INDEX, "id") + assert _index_validity(scratch_schema) == {HEALTH_INDEX: False, LOOKALIKE_INDEX: False} + + assert ProxyExtrasDBManager.repair_invalid_indexes() is True + + assert _index_validity(scratch_schema) == {HEALTH_INDEX: True, LOOKALIKE_INDEX: False} + + +@requires_db +def test_repair_drops_leftovers_of_interrupted_rebuilds(scratch_schema: str) -> None: + _leave_invalid_index(scratch_schema, HEALTH_TABLE, HEALTH_INDEX, HEALTH_INDEX_COLUMNS) + _leave_invalid_reindex_leftover(scratch_schema, HEALTH_TABLE, HEALTH_INDEX) + _leave_invalid_index(scratch_schema, HEALTH_TABLE, f"{HEALTH_TABLE}_model_id_idx_ccold", '"model_id"') + _leave_invalid_index(scratch_schema, HEALTH_TABLE, f"{HEALTH_TABLE}_model_id_idx_ccnew1", '"model_id"') + before: Final = _index_validity(scratch_schema) + assert len(before) == 4 + assert set(before.values()) == {False} + + assert ProxyExtrasDBManager.repair_invalid_indexes() is True + + assert _index_validity(scratch_schema) == {HEALTH_INDEX: True} + + +@requires_db +def test_repair_is_a_no_op_when_every_index_is_valid(scratch_schema: str) -> None: + with psycopg.connect(_base_url(), autocommit=True) as conn: + conn.execute(f'CREATE INDEX "{HEALTH_INDEX}" ON "{scratch_schema}"."{HEALTH_TABLE}" ({HEALTH_INDEX_COLUMNS})') + + assert ProxyExtrasDBManager.repair_invalid_indexes() is True + + assert _index_validity(scratch_schema) == {HEALTH_INDEX: True} + + +@requires_db +def test_repair_leaves_partitioned_parent_indexes_alone(scratch_schema: str) -> None: + with psycopg.connect(_base_url(), autocommit=True) as conn: + conn.execute(f'CREATE TABLE "{scratch_schema}"."{PARTITIONED_TABLE}" (id INT) PARTITION BY RANGE (id)') + conn.execute( + f'CREATE TABLE "{scratch_schema}"."{PARTITIONED_TABLE}_p0" ' + f'PARTITION OF "{scratch_schema}"."{PARTITIONED_TABLE}" FOR VALUES FROM (0) TO (10)' + ) + conn.execute(f'CREATE INDEX "{PARTITIONED_INDEX}" ON ONLY "{scratch_schema}"."{PARTITIONED_TABLE}" (id)') + assert _index_validity(scratch_schema) == {PARTITIONED_INDEX: False} + + assert ProxyExtrasDBManager.repair_invalid_indexes() is True + + assert _index_validity(scratch_schema) == {PARTITIONED_INDEX: False} + + +@requires_db +def test_repair_yields_to_the_replica_holding_the_repair_lock(scratch_schema: str) -> None: + _leave_invalid_index(scratch_schema, HEALTH_TABLE, HEALTH_INDEX, HEALTH_INDEX_COLUMNS) + + with psycopg.connect(_base_url(), autocommit=True) as other_replica: + other_replica.execute("SELECT pg_advisory_lock(%s)", (INDEX_REPAIR_ADVISORY_LOCK_KEY,)) + assert ProxyExtrasDBManager.repair_invalid_indexes() is False + assert _index_validity(scratch_schema) == {HEALTH_INDEX: False} + + assert ProxyExtrasDBManager.repair_invalid_indexes() is True + assert _index_validity(scratch_schema) == {HEALTH_INDEX: True} + + +@requires_db +def test_repair_gives_up_on_a_blocked_rebuild_and_finishes_it_on_the_next_startup(scratch_schema: str) -> None: + _leave_invalid_index(scratch_schema, HEALTH_TABLE, HEALTH_INDEX, HEALTH_INDEX_COLUMNS) + + with psycopg.connect(_base_url()) as pin: + pin.isolation_level = psycopg.IsolationLevel.REPEATABLE_READ + pin.execute(f'SELECT count(*) FROM "{scratch_schema}"."{HEALTH_TABLE}"') + assert ProxyExtrasDBManager.repair_invalid_indexes(lock_timeout="1s") is False + blocked: Final = _index_validity(scratch_schema) + assert blocked[HEALTH_INDEX] is False + assert [name for name in blocked if name.endswith("_ccnew")] + + assert ProxyExtrasDBManager.repair_invalid_indexes() is True + assert _index_validity(scratch_schema) == {HEALTH_INDEX: True} + + +def _hold_snapshot(schema: str, table: str, pinned: threading.Event, seconds: float) -> None: + with psycopg.connect(_base_url()) as pin: + pin.isolation_level = psycopg.IsolationLevel.REPEATABLE_READ + pin.execute(f'SELECT count(*) FROM "{schema}"."{table}"') + pinned.set() + pin.execute("SELECT pg_sleep(%s)", (seconds,)) + + +@requires_db +def test_repair_outlives_a_statement_timeout_passed_through_database_url_options( + scratch_schema: str, monkeypatch: pytest.MonkeyPatch +) -> None: + _leave_invalid_index(scratch_schema, HEALTH_TABLE, HEALTH_INDEX, HEALTH_INDEX_COLUMNS) + monkeypatch.setenv("DATABASE_URL", f"{_base_url()}?schema={scratch_schema}&options=-c%20statement_timeout%3D2000") + pinned: Final = threading.Event() + holder: Final = threading.Thread(target=_hold_snapshot, args=(scratch_schema, HEALTH_TABLE, pinned, 5.0)) + holder.start() + pinned.wait() + try: + assert ProxyExtrasDBManager.repair_invalid_indexes() is True + finally: + holder.join() + + assert _index_validity(scratch_schema) == {HEALTH_INDEX: True} + + +@requires_db +def test_repair_defaults_to_the_public_schema(monkeypatch: pytest.MonkeyPatch) -> None: + table: Final = f"LiteLLM_ScratchTable_{uuid.uuid4().hex[:8]}" + index: Final = f"{table}_id_idx" + monkeypatch.setenv("DATABASE_URL", _base_url()) + with psycopg.connect(_base_url(), autocommit=True) as conn: + conn.execute(f'CREATE TABLE public."{table}" (id TEXT)') + try: + _leave_invalid_index("public", table, index, "id") + assert _index_validity("public")[index] is False + + assert ProxyExtrasDBManager.repair_invalid_indexes() is True + + assert _index_validity("public")[index] is True + finally: + with psycopg.connect(_base_url(), autocommit=True) as conn: + conn.execute(f'DROP TABLE public."{table}"') + + +def test_repair_survives_an_unreachable_database(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@127.0.0.1:9/x?schema=whatever") + + assert ProxyExtrasDBManager.repair_invalid_indexes() is False + + +class _MigrateDeployApplied: + stdout = "Applied migration.\n" + stderr = "" + returncode = 0 + + +@requires_db +@pytest.mark.parametrize("use_v2_resolver", [True, False]) +def test_setup_database_repairs_the_index_after_a_recovered_deploy( + scratch_schema: str, monkeypatch: pytest.MonkeyPatch, tmp_path: Path, use_v2_resolver: bool +) -> None: + _leave_invalid_index(scratch_schema, HEALTH_TABLE, HEALTH_INDEX, HEALTH_INDEX_COLUMNS) + (tmp_path / "schema.prisma").write_text("// stub") + monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path)) + monkeypatch.setattr(ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None) + monkeypatch.setattr(ProxyExtrasDBManager, "_resolve_all_migrations", lambda *_, **__: True) + monkeypatch.setattr(subprocess, "run", lambda *_, **__: _MigrateDeployApplied()) + + assert ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=use_v2_resolver) is True + + assert _index_validity(scratch_schema) == {HEALTH_INDEX: True} From 2d64020d432f1c6112b836a752720a25a5ab93e9 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 15:24:42 -0700 Subject: [PATCH 022/109] fix(proxy-extras): run the index repair over DIRECT_URL and exercise it through a real migrate deploy --- .../litellm_proxy_extras/utils.py | 7 ++- .../test_invalid_index_repair.py | 63 ++++++++++++------- 2 files changed, 46 insertions(+), 24 deletions(-) diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index df019d78078..75fda59f28e 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -721,8 +721,11 @@ class ProxyExtrasDBManager: INVALID (a migration deadlock between replicas is the usual cause; the retried migration skips them because of IF NOT EXISTS). Never raises: returns True when no invalid index remains, False when the repair was - skipped or failed and will be retried on the next startup.""" - database_url: Final = os.getenv("DATABASE_URL") + skipped or failed and will be retried on the next startup. Runs over + DIRECT_URL when set: the session settings, the advisory lock and REINDEX + CONCURRENTLY all need one server session, which a transaction pooler + does not give.""" + database_url: Final = os.getenv("DIRECT_URL") or os.getenv("DATABASE_URL") if not database_url: return False diff --git a/tests/proxy_migration_tests/test_invalid_index_repair.py b/tests/proxy_migration_tests/test_invalid_index_repair.py index 787a62b361b..2e46775cd2b 100644 --- a/tests/proxy_migration_tests/test_invalid_index_repair.py +++ b/tests/proxy_migration_tests/test_invalid_index_repair.py @@ -1,9 +1,7 @@ import os -import subprocess import threading import uuid from collections.abc import Iterator, Mapping -from pathlib import Path from types import MappingProxyType from typing import Final @@ -66,16 +64,14 @@ def _leave_invalid_reindex_leftover(schema: str, table: str, index: str) -> None _interrupt_concurrent_build(schema, table, f'REINDEX INDEX CONCURRENTLY "{schema}"."{index}"') -@pytest.fixture -def scratch_schema(monkeypatch: pytest.MonkeyPatch) -> Iterator[str]: +def _scratch_schema(monkeypatch: pytest.MonkeyPatch, *table_definitions: str) -> Iterator[str]: schema: Final = f"invalid_index_{uuid.uuid4().hex[:8]}" with psycopg.connect(_base_url(), autocommit=True) as conn: conn.execute(f'CREATE SCHEMA "{schema}"') - conn.execute( - f'CREATE TABLE "{schema}"."{HEALTH_TABLE}" (model_id TEXT, model_name TEXT, checked_at TIMESTAMPTZ)' - ) - conn.execute(f'CREATE TABLE "{schema}"."{LOOKALIKE_TABLE}" (id TEXT)') + for definition in table_definitions: + conn.execute(f'CREATE TABLE "{schema}".{definition}') + monkeypatch.delenv("DIRECT_URL", raising=False) monkeypatch.setenv("DATABASE_URL", f"{_base_url()}?schema={schema}") yield schema @@ -83,6 +79,20 @@ def scratch_schema(monkeypatch: pytest.MonkeyPatch) -> Iterator[str]: conn.execute(f'DROP SCHEMA "{schema}" CASCADE') +@pytest.fixture +def scratch_schema(monkeypatch: pytest.MonkeyPatch) -> Iterator[str]: + yield from _scratch_schema( + monkeypatch, + f'"{HEALTH_TABLE}" (model_id TEXT, model_name TEXT, checked_at TIMESTAMPTZ)', + f'"{LOOKALIKE_TABLE}" (id TEXT)', + ) + + +@pytest.fixture +def empty_schema(monkeypatch: pytest.MonkeyPatch) -> Iterator[str]: + yield from _scratch_schema(monkeypatch) + + @requires_db def test_repair_rebuilds_invalid_litellm_indexes_and_leaves_lookalike_tables_alone(scratch_schema: str) -> None: _leave_invalid_index(scratch_schema, HEALTH_TABLE, HEALTH_INDEX, HEALTH_INDEX_COLUMNS) @@ -210,29 +220,38 @@ def test_repair_defaults_to_the_public_schema(monkeypatch: pytest.MonkeyPatch) - def test_repair_survives_an_unreachable_database(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("DIRECT_URL", raising=False) monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@127.0.0.1:9/x?schema=whatever") assert ProxyExtrasDBManager.repair_invalid_indexes() is False -class _MigrateDeployApplied: - stdout = "Applied migration.\n" - stderr = "" - returncode = 0 +@requires_db +def test_repair_runs_over_direct_url_when_set(scratch_schema: str) -> None: + _leave_invalid_index(scratch_schema, HEALTH_TABLE, HEALTH_INDEX, HEALTH_INDEX_COLUMNS) + direct_url: Final = os.environ["DATABASE_URL"] + with pytest.MonkeyPatch.context() as env: + env.setenv("DIRECT_URL", direct_url) + env.setenv("DATABASE_URL", "postgresql://u:p@127.0.0.1:9/x?schema=whatever") + assert ProxyExtrasDBManager.repair_invalid_indexes() is True + + assert _index_validity(scratch_schema) == {HEALTH_INDEX: True} + + +def _invalidate_deployed_index(schema: str) -> None: + with psycopg.connect(_base_url(), autocommit=True) as conn: + conn.execute(f'DROP INDEX "{schema}"."{HEALTH_INDEX}"') + _leave_invalid_index(schema, HEALTH_TABLE, HEALTH_INDEX, HEALTH_INDEX_COLUMNS) @requires_db +@pytest.mark.timeout(300) @pytest.mark.parametrize("use_v2_resolver", [True, False]) -def test_setup_database_repairs_the_index_after_a_recovered_deploy( - scratch_schema: str, monkeypatch: pytest.MonkeyPatch, tmp_path: Path, use_v2_resolver: bool -) -> None: - _leave_invalid_index(scratch_schema, HEALTH_TABLE, HEALTH_INDEX, HEALTH_INDEX_COLUMNS) - (tmp_path / "schema.prisma").write_text("// stub") - monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path)) - monkeypatch.setattr(ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None) - monkeypatch.setattr(ProxyExtrasDBManager, "_resolve_all_migrations", lambda *_, **__: True) - monkeypatch.setattr(subprocess, "run", lambda *_, **__: _MigrateDeployApplied()) +def test_setup_database_repairs_the_index_after_a_recovered_deploy(empty_schema: str, use_v2_resolver: bool) -> None: + assert ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=use_v2_resolver) is True + _invalidate_deployed_index(empty_schema) + assert _index_validity(empty_schema)[HEALTH_INDEX] is False assert ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=use_v2_resolver) is True - assert _index_validity(scratch_schema) == {HEALTH_INDEX: True} + assert _index_validity(empty_schema)[HEALTH_INDEX] is True From bfdaedf51b13f4bc8872d51a99d7f288fc0cc56f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 2 Sep 2026 16:23:50 -0700 Subject: [PATCH 023/109] fix(proxy-extras): take the repair schema from DATABASE_URL and deploy the real-migration test on a fresh database --- .../litellm_proxy_extras/utils.py | 17 +++---- .../test_invalid_index_repair.py | 47 +++++++++++-------- 2 files changed, 36 insertions(+), 28 deletions(-) diff --git a/litellm-proxy-extras/litellm_proxy_extras/utils.py b/litellm-proxy-extras/litellm_proxy_extras/utils.py index 75fda59f28e..40bc4cd1dfd 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/utils.py +++ b/litellm-proxy-extras/litellm_proxy_extras/utils.py @@ -721,12 +721,13 @@ class ProxyExtrasDBManager: INVALID (a migration deadlock between replicas is the usual cause; the retried migration skips them because of IF NOT EXISTS). Never raises: returns True when no invalid index remains, False when the repair was - skipped or failed and will be retried on the next startup. Runs over - DIRECT_URL when set: the session settings, the advisory lock and REINDEX - CONCURRENTLY all need one server session, which a transaction pooler - does not give.""" - database_url: Final = os.getenv("DIRECT_URL") or os.getenv("DATABASE_URL") - if not database_url: + skipped or failed and will be retried on the next startup. Looks in the + schema DATABASE_URL names, the only URL Prisma migrates through, but + connects over DIRECT_URL when set: the session settings, the advisory + lock and REINDEX CONCURRENTLY all need one server session, which a + transaction pooler does not give.""" + prisma_url: Final = os.getenv("DATABASE_URL") + if not prisma_url: return False try: @@ -739,8 +740,8 @@ class ProxyExtrasDBManager: ) return False - schema: Final = ProxyExtrasDBManager._prisma_schema_param(database_url) or "public" - cleaned_url: Final = ProxyExtrasDBManager._strip_prisma_query_params(database_url) + schema: Final = ProxyExtrasDBManager._prisma_schema_param(prisma_url) or "public" + cleaned_url: Final = ProxyExtrasDBManager._strip_prisma_query_params(os.getenv("DIRECT_URL") or prisma_url) try: with psycopg.connect(cleaned_url, connect_timeout=10, autocommit=True) as conn: conn.execute("SET statement_timeout = 0") diff --git a/tests/proxy_migration_tests/test_invalid_index_repair.py b/tests/proxy_migration_tests/test_invalid_index_repair.py index 2e46775cd2b..741fa7386df 100644 --- a/tests/proxy_migration_tests/test_invalid_index_repair.py +++ b/tests/proxy_migration_tests/test_invalid_index_repair.py @@ -64,12 +64,15 @@ def _leave_invalid_reindex_leftover(schema: str, table: str, index: str) -> None _interrupt_concurrent_build(schema, table, f'REINDEX INDEX CONCURRENTLY "{schema}"."{index}"') -def _scratch_schema(monkeypatch: pytest.MonkeyPatch, *table_definitions: str) -> Iterator[str]: +@pytest.fixture +def scratch_schema(monkeypatch: pytest.MonkeyPatch) -> Iterator[str]: schema: Final = f"invalid_index_{uuid.uuid4().hex[:8]}" with psycopg.connect(_base_url(), autocommit=True) as conn: conn.execute(f'CREATE SCHEMA "{schema}"') - for definition in table_definitions: - conn.execute(f'CREATE TABLE "{schema}".{definition}') + conn.execute( + f'CREATE TABLE "{schema}"."{HEALTH_TABLE}" (model_id TEXT, model_name TEXT, checked_at TIMESTAMPTZ)' + ) + conn.execute(f'CREATE TABLE "{schema}"."{LOOKALIKE_TABLE}" (id TEXT)') monkeypatch.delenv("DIRECT_URL", raising=False) monkeypatch.setenv("DATABASE_URL", f"{_base_url()}?schema={schema}") @@ -80,17 +83,22 @@ def _scratch_schema(monkeypatch: pytest.MonkeyPatch, *table_definitions: str) -> @pytest.fixture -def scratch_schema(monkeypatch: pytest.MonkeyPatch) -> Iterator[str]: - yield from _scratch_schema( - monkeypatch, - f'"{HEALTH_TABLE}" (model_id TEXT, model_name TEXT, checked_at TIMESTAMPTZ)', - f'"{LOOKALIKE_TABLE}" (id TEXT)', - ) +def fresh_database(monkeypatch: pytest.MonkeyPatch) -> Iterator[str]: + """A brand-new database, what a first deploy sees. A scratch schema would + not do: the migrations guard on pg_constraint by name across every schema, + so a LiteLLM schema already pushed into public makes them skip and then + fail, which is exactly what CI's database looks like.""" + admin_url: Final = _base_url() + name: Final = f"invalid_index_{uuid.uuid4().hex[:8]}" + with psycopg.connect(admin_url, autocommit=True) as conn: + conn.execute(f'CREATE DATABASE "{name}"') + monkeypatch.delenv("DIRECT_URL", raising=False) + monkeypatch.setenv("DATABASE_URL", f"{admin_url.rsplit('/', 1)[0]}/{name}") + yield "public" -@pytest.fixture -def empty_schema(monkeypatch: pytest.MonkeyPatch) -> Iterator[str]: - yield from _scratch_schema(monkeypatch) + with psycopg.connect(admin_url, autocommit=True) as conn: + conn.execute(f'DROP DATABASE "{name}" WITH (FORCE)') @requires_db @@ -227,12 +235,11 @@ def test_repair_survives_an_unreachable_database(monkeypatch: pytest.MonkeyPatch @requires_db -def test_repair_runs_over_direct_url_when_set(scratch_schema: str) -> None: +def test_repair_connects_over_direct_url_but_looks_in_the_schema_database_url_names(scratch_schema: str) -> None: _leave_invalid_index(scratch_schema, HEALTH_TABLE, HEALTH_INDEX, HEALTH_INDEX_COLUMNS) - direct_url: Final = os.environ["DATABASE_URL"] with pytest.MonkeyPatch.context() as env: - env.setenv("DIRECT_URL", direct_url) - env.setenv("DATABASE_URL", "postgresql://u:p@127.0.0.1:9/x?schema=whatever") + env.setenv("DIRECT_URL", f"{_base_url()}?schema=public") + env.setenv("DATABASE_URL", f"postgresql://u:p@127.0.0.1:9/x?schema={scratch_schema}") assert ProxyExtrasDBManager.repair_invalid_indexes() is True assert _index_validity(scratch_schema) == {HEALTH_INDEX: True} @@ -247,11 +254,11 @@ def _invalidate_deployed_index(schema: str) -> None: @requires_db @pytest.mark.timeout(300) @pytest.mark.parametrize("use_v2_resolver", [True, False]) -def test_setup_database_repairs_the_index_after_a_recovered_deploy(empty_schema: str, use_v2_resolver: bool) -> None: +def test_setup_database_repairs_the_index_after_a_recovered_deploy(fresh_database: str, use_v2_resolver: bool) -> None: assert ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=use_v2_resolver) is True - _invalidate_deployed_index(empty_schema) - assert _index_validity(empty_schema)[HEALTH_INDEX] is False + _invalidate_deployed_index(fresh_database) + assert _index_validity(fresh_database)[HEALTH_INDEX] is False assert ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=use_v2_resolver) is True - assert _index_validity(empty_schema)[HEALTH_INDEX] is True + assert _index_validity(fresh_database)[HEALTH_INDEX] is True From bede8b5ea46c24b20d138c79025dd67beee97763 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 3 Sep 2026 02:49:49 -0700 Subject: [PATCH 024/109] fix(proxy): stop shipping the literal string "None" as error type and param The proxy's exception tails defaulted `type` and `param` to the four-character string "None", which is neither a known OpenAI error type nor the JSON null the nullable `param` field is typed as, so a client's error handler matched nothing and fell into its generic branch. Lifts the helpers PR #39521 added for the unified LLM endpoints into litellm/proxy/common_utils/openai_error_payload.py and calls them from the file, rerank, image, realtime, anthropic, and pass-through route families, plus the shared handle_exception_on_proxy handler that the management, batches, fine-tuning, credential, SCIM, guardrail, and customer routes funnel through. The remaining families (proxy_server, auth, health, spend tracking, and management endpoints) follow in separate PRs so each slice stays QA'able on a live proxy. --- .../proxy/anthropic_endpoints/endpoints.py | 11 +- litellm/proxy/common_request_processing.py | 78 ++++-------- .../common_utils/openai_error_payload.py | 48 +++++++ litellm/proxy/image_endpoints/endpoints.py | 17 ++- .../openai_files_endpoints/files_endpoints.py | 109 ++++++++-------- .../pass_through_endpoints.py | 23 ++-- litellm/proxy/realtime_endpoints/endpoints.py | 41 +++--- litellm/proxy/rerank_endpoints/endpoints.py | 17 ++- litellm/proxy/utils.py | 5 +- .../common_utils/test_openai_error_payload.py | 117 ++++++++++++++++++ .../test_files_endpoint.py | 10 +- .../proxy/utils/helpers/test_error_helpers.py | 2 +- 12 files changed, 320 insertions(+), 158 deletions(-) create mode 100644 litellm/proxy/common_utils/openai_error_payload.py create mode 100644 tests/test_litellm/proxy/common_utils/test_openai_error_payload.py diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py index 7f0045c1d93..f26cb4d41f7 100644 --- a/litellm/proxy/anthropic_endpoints/endpoints.py +++ b/litellm/proxy/anthropic_endpoints/endpoints.py @@ -25,6 +25,11 @@ from litellm.proxy.common_request_processing import ( proxy_exception_from_http_exception, ) from litellm.proxy.common_utils.http_parsing_utils import _read_request_body +from litellm.proxy.common_utils.openai_error_payload import ( + error_status_code, + openai_error_param, + openai_error_type, +) from litellm.types.utils import TokenCountResponse router: Final = APIRouter() @@ -221,9 +226,9 @@ async def anthropic_response( error_msg: Final = f"{e}" raise ProxyException( message=getattr(e, "message", error_msg), - type=getattr(e, "type", "None"), - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", 500), + type=openai_error_type(e, error_status_code(e, 500)), + param=openai_error_param(e), + code=error_status_code(e, 500), headers=headers, ) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 6542842f5e4..99028c3645a 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -54,6 +54,12 @@ from litellm.proxy.common_utils.callback_utils import ( get_logging_caching_headers, get_remaining_tokens_and_requests_from_request_data, ) +from litellm.proxy.common_utils.openai_error_payload import ( + attribute_of, + error_status_code, + openai_error_param, + openai_error_type, +) from litellm.proxy.common_utils.sse_keepalive import ( SSE_COMMENT_PING_BYTES, coerce_keepalive_interval, @@ -463,46 +469,6 @@ def _stream_usage_tracking_updates( } -def _getattr_object(value: object, name: str, default: object = None) -> object: - return getattr(value, name, default) - - -_OPENAI_ERROR_TYPE_BY_STATUS: Final[Mapping[int, str]] = MappingProxyType( - { - status.HTTP_401_UNAUTHORIZED: "authentication_error", - status.HTTP_403_FORBIDDEN: "permission_error", - status.HTTP_429_TOO_MANY_REQUESTS: "rate_limit_error", - } -) - - -def _error_status_code(exc: object, default: int) -> int: - """The HTTP status an exception carries, or ``default`` when it carries none.""" - carried: Final = _getattr_object(exc, "status_code") - return carried if isinstance(carried, int) and not isinstance(carried, bool) else default - - -def _openai_error_type(exc: object, status_code: int) -> str: - """OpenAI types ``error.type`` as a required string, so an exception carrying none - falls back to the type its status code stands for.""" - carried: Final = _getattr_object(exc, "type") - if isinstance(carried, str): - return carried - mapped: Final = _OPENAI_ERROR_TYPE_BY_STATUS.get(status_code) - if mapped is not None: - return mapped - if status_code < status.HTTP_500_INTERNAL_SERVER_ERROR: - return "invalid_request_error" - return "internal_server_error" - - -def _openai_error_param(exc: object) -> str | None: - """OpenAI types ``error.param`` as nullable, so an exception carrying none - serializes as JSON ``null``.""" - carried: Final = _getattr_object(exc, "param") - return carried if isinstance(carried, str) else None - - class _UpstreamHttpResponse(Protocol): @property def status_code(self) -> int: ... @@ -572,15 +538,15 @@ def serialize_http_exception_detail( def proxy_exception_from_http_exception(exc: HTTPException, headers: dict[str, str]) -> ProxyException: - raw_detail: Final = _getattr_object(exc, "detail", str(exc)) + raw_detail: Final = attribute_of(exc, "detail", str(exc)) message, structured_fields = serialize_http_exception_detail(raw_detail) existing_fields: Final = getattr(exc, "provider_specific_fields", None) or {} merged_fields: Final = {**existing_fields, **structured_fields} if structured_fields else (existing_fields or None) - error_status: Final = _error_status_code(exc, status.HTTP_400_BAD_REQUEST) + error_status: Final = error_status_code(exc, status.HTTP_400_BAD_REQUEST) return ProxyException( message=message, - type=_openai_error_type(exc, error_status), - param=_openai_error_param(exc), + type=openai_error_type(exc, error_status), + param=openai_error_param(exc), code=error_status, provider_specific_fields=merged_fields, headers=headers, @@ -864,8 +830,8 @@ def sse_error_payload(exc: BaseException) -> tuple[int, Mapping[str, object]]: are byte-identical. """ # Preserve status code from HTTPException (e.g. guardrail blocks) - error_status: Final = _error_status_code(exc, status.HTTP_500_INTERNAL_SERVER_ERROR) - raw_detail: Final = _getattr_object(exc, "detail", "Error processing stream start") + error_status: Final = error_status_code(exc, status.HTTP_500_INTERNAL_SERVER_ERROR) + raw_detail: Final = attribute_of(exc, "detail", "Error processing stream start") message, structured_fields = serialize_http_exception_detail(raw_detail) existing_fields: Final = getattr(exc, "provider_specific_fields", None) or {} @@ -873,8 +839,8 @@ def sse_error_payload(exc: BaseException) -> tuple[int, Mapping[str, object]]: error_obj: Final = { "message": message, - "type": _openai_error_type(exc, error_status), - "param": _openai_error_param(exc), + "type": openai_error_type(exc, error_status), + "param": openai_error_param(exc), "code": str(error_status), } if not merged_fields: @@ -2755,10 +2721,10 @@ class ProxyBaseLLMRequestProcessing: ``ResponsesAPIResponse`` directly. Handle both shapes so the container-ownership recording path can walk ``.output`` either way. """ - completed: Final = _getattr_object(stream_response, "completed_response") + completed: Final = attribute_of(stream_response, "completed_response") if completed is None: return None - response_obj: Final = _getattr_object(completed, "response") + response_obj: Final = attribute_of(completed, "response") if response_obj is not None: return response_obj return completed @@ -3380,7 +3346,7 @@ class ProxyBaseLLMRequestProcessing: headers = getattr(e, "headers", None) or {} if not headers: # Try to get headers from e.response.headers (httpx.Response) - _response: Final = _getattr_object(e, "response") + _response: Final = attribute_of(e, "response") if _response is not None: _response_headers: Final = getattr(_response, "headers", None) if _response_headers: @@ -3451,8 +3417,8 @@ class ProxyBaseLLMRequestProcessing: _code = status.HTTP_500_INTERNAL_SERVER_ERROR raise ProxyException( message=redact_internal_details_from_client_message(getattr(e, "message", error_msg)), - type=_openai_error_type(e, _code), - param=_openai_error_param(e), + type=openai_error_type(e, _code), + param=openai_error_param(e), openai_code=getattr(e, "code", None), code=_code, provider_specific_fields=getattr(e, "provider_specific_fields", None), @@ -3662,11 +3628,11 @@ class ProxyBaseLLMRequestProcessing: if isinstance(e, HTTPException): raise e - stream_error_status: Final = _error_status_code(e, status.HTTP_500_INTERNAL_SERVER_ERROR) + stream_error_status: Final = error_status_code(e, status.HTTP_500_INTERNAL_SERVER_ERROR) proxy_exception: Final = ProxyException( message=redact_internal_details_from_client_message(getattr(e, "message", str(e))), - type=_openai_error_type(e, stream_error_status), - param=_openai_error_param(e), + type=openai_error_type(e, stream_error_status), + param=openai_error_param(e), code=stream_error_status, ) stream_completed = True diff --git a/litellm/proxy/common_utils/openai_error_payload.py b/litellm/proxy/common_utils/openai_error_payload.py new file mode 100644 index 00000000000..180ec152094 --- /dev/null +++ b/litellm/proxy/common_utils/openai_error_payload.py @@ -0,0 +1,48 @@ +"""Shapes the ``error`` object the proxy answers with so it matches OpenAI's contract: +``type`` is a required string and ``param`` is nullable, neither of which the literal +string ``"None"`` satisfies.""" + +from collections.abc import Mapping +from types import MappingProxyType +from typing import Final + +from fastapi import status + +_OPENAI_ERROR_TYPE_BY_STATUS: Final[Mapping[int, str]] = MappingProxyType( + { + status.HTTP_401_UNAUTHORIZED: "authentication_error", + status.HTTP_403_FORBIDDEN: "permission_error", + status.HTTP_429_TOO_MANY_REQUESTS: "rate_limit_error", + } +) + + +def attribute_of(value: object, name: str, default: object = None) -> object: + return getattr(value, name, default) + + +def error_status_code(exc: object, default: int) -> int: + """The HTTP status an exception carries, or ``default`` when it carries none.""" + carried: Final = attribute_of(exc, "status_code") + return carried if isinstance(carried, int) and not isinstance(carried, bool) else default + + +def openai_error_type(exc: object, status_code: int) -> str: + """OpenAI types ``error.type`` as a required string, so an exception carrying none + falls back to the type its status code stands for.""" + carried: Final = attribute_of(exc, "type") + if isinstance(carried, str): + return carried + mapped: Final = _OPENAI_ERROR_TYPE_BY_STATUS.get(status_code) + if mapped is not None: + return mapped + if status_code < status.HTTP_500_INTERNAL_SERVER_ERROR: + return "invalid_request_error" + return "internal_server_error" + + +def openai_error_param(exc: object) -> str | None: + """OpenAI types ``error.param`` as nullable, so an exception carrying none + serializes as JSON ``null``.""" + carried: Final = attribute_of(exc, "param") + return carried if isinstance(carried, str) else None diff --git a/litellm/proxy/image_endpoints/endpoints.py b/litellm/proxy/image_endpoints/endpoints.py index 83caa92ede5..b83f7e5cd60 100644 --- a/litellm/proxy/image_endpoints/endpoints.py +++ b/litellm/proxy/image_endpoints/endpoints.py @@ -16,6 +16,11 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing +from litellm.proxy.common_utils.openai_error_payload import ( + error_status_code, + openai_error_param, + openai_error_type, +) from litellm.proxy.route_llm_request import route_request from litellm.types.llms.openai import ChatCompletionUserMessage @@ -193,18 +198,18 @@ async def image_generation( if isinstance(e, HTTPException): raise ProxyException( message=getattr(e, "message", str(e)), - type=getattr(e, "type", "None"), - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), + type=openai_error_type(e, error_status_code(e, status.HTTP_400_BAD_REQUEST)), + param=openai_error_param(e), + code=error_status_code(e, status.HTTP_400_BAD_REQUEST), ) else: error_msg: Final = f"{e}" raise ProxyException( message=getattr(e, "message", error_msg), - type=getattr(e, "type", "None"), - param=getattr(e, "param", "None"), + type=openai_error_type(e, error_status_code(e, 500)), + param=openai_error_param(e), openai_code=getattr(e, "code", None), - code=getattr(e, "status_code", 500), + code=error_status_code(e, 500), ) diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index bf07f4748ef..c315d30b8f3 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -45,6 +45,11 @@ from litellm.proxy.common_utils.openai_endpoint_utils import ( get_custom_llm_provider_from_request_headers, get_custom_llm_provider_from_request_query, ) +from litellm.proxy.common_utils.openai_error_payload import ( + error_status_code, + openai_error_param, + openai_error_type, +) from litellm.proxy.openai_files_endpoints.batch_file_validation import ( check_batch_file_upload, raise_batch_file_validation_failure, @@ -296,22 +301,22 @@ async def route_create_file( if managed_files_obj is None: raise ProxyException( message="Managed files hook not found", - type="None", - param="None", + type=ProxyErrorTypes.internal_server_error.value, + param=None, code=500, ) if llm_router is None: raise ProxyException( message="LLM Router not found", - type="None", - param="None", + type=ProxyErrorTypes.internal_server_error.value, + param=None, code=500, ) if not isinstance(managed_files_obj, BaseFileEndpoints): raise ProxyException( message="Managed files hook is not a BaseFileEndpoints", - type="None", - param="None", + type=ProxyErrorTypes.internal_server_error.value, + param=None, code=500, ) # Managed files internally calls llm_router.acreate_file() which includes loadbalancing @@ -713,17 +718,17 @@ async def create_file( if isinstance(e, HTTPException): raise ProxyException( message=getattr(e, "message", str(e.detail)), - type=getattr(e, "type", "None"), - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), + type=openai_error_type(e, error_status_code(e, status.HTTP_400_BAD_REQUEST)), + param=openai_error_param(e), + code=error_status_code(e, status.HTTP_400_BAD_REQUEST), ) else: error_msg: Final = f"{e}" raise ProxyException( message=getattr(e, "message", error_msg), - type=getattr(e, "type", "None"), - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", 500), + type=openai_error_type(e, error_status_code(e, 500)), + param=openai_error_param(e), + code=error_status_code(e, 500), ) finally: for spool in spools: @@ -812,22 +817,22 @@ async def get_file_content( if managed_files_obj is None: raise ProxyException( message="Managed files hook not found", - type="None", - param="None", + type=ProxyErrorTypes.internal_server_error.value, + param=None, code=500, ) if llm_router is None: raise ProxyException( message="LLM Router not found", - type="None", - param="None", + type=ProxyErrorTypes.internal_server_error.value, + param=None, code=500, ) if not isinstance(managed_files_obj, BaseFileEndpoints): raise ProxyException( message="Managed files hook is not a BaseFileEndpoints", - type="None", - param="None", + type=ProxyErrorTypes.internal_server_error.value, + param=None, code=500, ) @@ -1021,17 +1026,17 @@ async def get_file_content( if isinstance(e, HTTPException): raise ProxyException( message=getattr(e, "message", str(e.detail)), - type=getattr(e, "type", "None"), - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), + type=openai_error_type(e, error_status_code(e, status.HTTP_400_BAD_REQUEST)), + param=openai_error_param(e), + code=error_status_code(e, status.HTTP_400_BAD_REQUEST), ) else: error_msg: Final = f"{e}" raise ProxyException( message=getattr(e, "message", error_msg), - type=getattr(e, "type", "None"), - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", 500), + type=openai_error_type(e, error_status_code(e, 500)), + param=openai_error_param(e), + code=error_status_code(e, 500), ) @@ -1151,15 +1156,15 @@ async def get_file( if managed_files_obj is None: raise ProxyException( message="Managed files hook not found", - type="None", - param="None", + type=ProxyErrorTypes.internal_server_error.value, + param=None, code=500, ) if not isinstance(managed_files_obj, BaseFileEndpoints): raise ProxyException( message="Managed files hook is not a BaseFileEndpoints", - type="None", - param="None", + type=ProxyErrorTypes.internal_server_error.value, + param=None, code=500, ) response = await managed_files_obj.afile_retrieve( @@ -1215,17 +1220,17 @@ async def get_file( if isinstance(e, HTTPException): raise ProxyException( message=getattr(e, "message", str(e.detail)), - type=getattr(e, "type", "None"), - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), + type=openai_error_type(e, error_status_code(e, status.HTTP_400_BAD_REQUEST)), + param=openai_error_param(e), + code=error_status_code(e, status.HTTP_400_BAD_REQUEST), ) else: error_msg: Final = f"{e}" raise ProxyException( message=getattr(e, "message", error_msg), - type=getattr(e, "type", "None"), - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", 500), + type=openai_error_type(e, error_status_code(e, 500)), + param=openai_error_param(e), + code=error_status_code(e, 500), ) @@ -1355,22 +1360,22 @@ async def delete_file( if managed_files_obj is None: raise ProxyException( message="Managed files hook not found", - type="None", - param="None", + type=ProxyErrorTypes.internal_server_error.value, + param=None, code=500, ) if llm_router is None: raise ProxyException( message="LLM Router not found", - type="None", - param="None", + type=ProxyErrorTypes.internal_server_error.value, + param=None, code=500, ) if not isinstance(managed_files_obj, BaseFileEndpoints): raise ProxyException( message="Managed files hook is not a BaseFileEndpoints", - type="None", - param="None", + type=ProxyErrorTypes.internal_server_error.value, + param=None, code=500, ) @@ -1427,17 +1432,17 @@ async def delete_file( if isinstance(e, HTTPException): raise ProxyException( message=getattr(e, "message", str(e.detail)), - type=getattr(e, "type", "None"), - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), + type=openai_error_type(e, error_status_code(e, status.HTTP_400_BAD_REQUEST)), + param=openai_error_param(e), + code=error_status_code(e, status.HTTP_400_BAD_REQUEST), ) else: error_msg: Final = f"{e}" raise ProxyException( message=getattr(e, "message", error_msg), - type=getattr(e, "type", "None"), - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", 500), + type=openai_error_type(e, error_status_code(e, 500)), + param=openai_error_param(e), + code=error_status_code(e, 500), ) @@ -1629,15 +1634,15 @@ async def list_files( if isinstance(e, HTTPException): raise ProxyException( message=getattr(e, "message", str(e.detail)), - type=getattr(e, "type", "None"), - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), + type=openai_error_type(e, error_status_code(e, status.HTTP_400_BAD_REQUEST)), + param=openai_error_param(e), + code=error_status_code(e, status.HTTP_400_BAD_REQUEST), ) else: error_msg: Final = f"{e}" raise ProxyException( message=getattr(e, "message", error_msg), - type=getattr(e, "type", "None"), - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", 500), + type=openai_error_type(e, error_status_code(e, 500)), + param=openai_error_param(e), + code=error_status_code(e, 500), ) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 79d5d0a016f..64ce461990a 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -77,6 +77,11 @@ from litellm.proxy.common_utils.http_parsing_utils import ( _read_request_body, _safe_get_request_headers, ) +from litellm.proxy.common_utils.openai_error_payload import ( + error_status_code, + openai_error_param, + openai_error_type, +) from litellm.proxy.common_utils.sse_keepalive import ( wrap_passthrough_sse_bytes_with_keepalive_pings, ) @@ -310,9 +315,9 @@ async def chat_completion_pass_through_endpoint( error_msg: Final = f"{e}" raise ProxyException( message=getattr(e, "message", error_msg), - type=getattr(e, "type", "None"), - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", 500), + type=openai_error_type(e, error_status_code(e, 500)), + param=openai_error_param(e), + code=error_status_code(e, 500), ) @@ -1677,18 +1682,18 @@ async def pass_through_request( if isinstance(e, HTTPException): raise ProxyException( message=getattr(e, "message", str(getattr(e, "detail", str(e)))), - type=getattr(e, "type", "None"), - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), + type=openai_error_type(e, error_status_code(e, status.HTTP_400_BAD_REQUEST)), + param=openai_error_param(e), + code=error_status_code(e, status.HTTP_400_BAD_REQUEST), headers=custom_headers, ) else: error_msg: Final = f"{e}" raise ProxyException( message=getattr(e, "message", error_msg), - type=getattr(e, "type", "None"), - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", 500), + type=openai_error_type(e, error_status_code(e, 500)), + param=openai_error_param(e), + code=error_status_code(e, 500), headers=custom_headers, ) diff --git a/litellm/proxy/realtime_endpoints/endpoints.py b/litellm/proxy/realtime_endpoints/endpoints.py index 7f9cd251a8a..9996f60098d 100644 --- a/litellm/proxy/realtime_endpoints/endpoints.py +++ b/litellm/proxy/realtime_endpoints/endpoints.py @@ -17,6 +17,11 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import ( encrypt_value_helper, ) from litellm.proxy.common_utils.http_parsing_utils import _read_request_body +from litellm.proxy.common_utils.openai_error_payload import ( + error_status_code, + openai_error_param, + openai_error_type, +) from litellm.types.realtime import ( RealtimeClientSecretRequest, RealtimeClientSecretResponse, @@ -301,15 +306,15 @@ async def create_realtime_client_secret( if isinstance(e, HTTPException): raise ProxyException( message=getattr(e, "message", str(e)), - type=getattr(e, "type", "None"), - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", http_status.HTTP_400_BAD_REQUEST), + type=openai_error_type(e, error_status_code(e, http_status.HTTP_400_BAD_REQUEST)), + param=openai_error_param(e), + code=error_status_code(e, http_status.HTTP_400_BAD_REQUEST), ) raise ProxyException( message=getattr(e, "message", str(e)), - type=getattr(e, "type", "None"), - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", 500), + type=openai_error_type(e, error_status_code(e, 500)), + param=openai_error_param(e), + code=error_status_code(e, 500), ) if upstream_resp.status_code != 200: @@ -492,15 +497,15 @@ async def proxy_realtime_calls( if isinstance(e, HTTPException): raise ProxyException( message=getattr(e, "message", str(e)), - type=getattr(e, "type", "None"), - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", http_status.HTTP_400_BAD_REQUEST), + type=openai_error_type(e, error_status_code(e, http_status.HTTP_400_BAD_REQUEST)), + param=openai_error_param(e), + code=error_status_code(e, http_status.HTTP_400_BAD_REQUEST), ) raise ProxyException( message=getattr(e, "message", str(e)), - type=getattr(e, "type", "None"), - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", 500), + type=openai_error_type(e, error_status_code(e, 500)), + param=openai_error_param(e), + code=error_status_code(e, 500), ) return Response( @@ -605,15 +610,15 @@ async def create_realtime_transcription_session( if isinstance(e, HTTPException): raise ProxyException( message=getattr(e, "detail", getattr(e, "message", str(e))), - type=getattr(e, "type", "None"), - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", http_status.HTTP_400_BAD_REQUEST), + type=openai_error_type(e, error_status_code(e, http_status.HTTP_400_BAD_REQUEST)), + param=openai_error_param(e), + code=error_status_code(e, http_status.HTTP_400_BAD_REQUEST), ) raise ProxyException( message=getattr(e, "message", str(e)), - type=getattr(e, "type", "None"), - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", 500), + type=openai_error_type(e, error_status_code(e, 500)), + param=openai_error_param(e), + code=error_status_code(e, 500), ) if upstream_resp.status_code != 200: diff --git a/litellm/proxy/rerank_endpoints/endpoints.py b/litellm/proxy/rerank_endpoints/endpoints.py index dd5803796b7..16cd7368e4a 100644 --- a/litellm/proxy/rerank_endpoints/endpoints.py +++ b/litellm/proxy/rerank_endpoints/endpoints.py @@ -11,6 +11,11 @@ from litellm._logging import verbose_proxy_logger from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing +from litellm.proxy.common_utils.openai_error_payload import ( + error_status_code, + openai_error_param, + openai_error_type, +) router: Final = APIRouter() @@ -112,15 +117,15 @@ async def rerank( if isinstance(e, HTTPException): raise ProxyException( message=getattr(e, "message", str(e)), - type=getattr(e, "type", "None"), - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), + type=openai_error_type(e, error_status_code(e, status.HTTP_400_BAD_REQUEST)), + param=openai_error_param(e), + code=error_status_code(e, status.HTTP_400_BAD_REQUEST), ) else: error_msg: Final = f"{e}" raise ProxyException( message=getattr(e, "message", error_msg), - type=getattr(e, "type", "None"), - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", 500), + type=openai_error_type(e, error_status_code(e, 500)), + param=openai_error_param(e), + code=error_status_code(e, 500), ) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index cab2bd6d9db..5c0fbd64744 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -37,6 +37,7 @@ from litellm.proxy._types import ( SpendLogsMetadata, SpendLogsPayload, ) +from litellm.proxy.common_utils.openai_error_payload import openai_error_param from litellm.proxy.spend_tracking.spend_log_error_logger import spend_log_error from litellm.types.guardrails import GuardrailEventHooks from litellm.types.proxy.model_listing import ModelInfoResponse @@ -7098,7 +7099,7 @@ def handle_exception_on_proxy(e: Exception) -> ProxyException: return ProxyException( message=getattr(e, "detail", f"error({e})"), type=ProxyErrorTypes.internal_server_error, - param=getattr(e, "param", "None"), + param=openai_error_param(e), code=getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR), ) elif isinstance(e, ProxyException): @@ -7107,7 +7108,7 @@ def handle_exception_on_proxy(e: Exception) -> ProxyException: return ProxyException( message=str(e), type=ProxyErrorTypes.internal_server_error, - param=getattr(e, "param", "None"), + param=openai_error_param(e), code=_status_code, ) diff --git a/tests/test_litellm/proxy/common_utils/test_openai_error_payload.py b/tests/test_litellm/proxy/common_utils/test_openai_error_payload.py new file mode 100644 index 00000000000..c165d7ffdb1 --- /dev/null +++ b/tests/test_litellm/proxy/common_utils/test_openai_error_payload.py @@ -0,0 +1,117 @@ +import json + +import pytest +from fastapi import HTTPException + +from litellm.proxy._types import ProxyErrorTypes, ProxyException +from litellm.proxy.common_utils.openai_error_payload import ( + error_status_code, + openai_error_param, + openai_error_type, +) + + +@pytest.mark.parametrize( + "status_code, expected_type", + [ + (400, "invalid_request_error"), + (401, "authentication_error"), + (403, "permission_error"), + (404, "invalid_request_error"), + (422, "invalid_request_error"), + (429, "rate_limit_error"), + (499, "invalid_request_error"), + (500, "internal_server_error"), + (502, "internal_server_error"), + (503, "internal_server_error"), + ], +) +def test_status_code_decides_the_type_when_the_exception_carries_none(status_code, expected_type): + """A route that raises a bare HTTPException carries no error type, so the status it + answered with is the only thing left to name the OpenAI type from.""" + assert openai_error_type(HTTPException(status_code=status_code, detail="boom"), status_code) == expected_type + + +def test_a_carried_type_wins_over_the_one_the_status_would_imply(): + """A ProxyException raised mid-request already names its own type, and relabelling a + 402 budget_exceeded as the status map's guess would lose what the client branches on.""" + carried = ProxyException( + message="Budget has been exceeded", + type=ProxyErrorTypes.budget_exceeded.value, + param=None, + code=400, + ) + + assert openai_error_type(carried, 400) == ProxyErrorTypes.budget_exceeded.value + + +@pytest.mark.parametrize("carried_type", [None, 400, {"type": "invalid_request_error"}, ["invalid_request_error"]]) +def test_a_non_string_carried_type_falls_back_to_the_status(carried_type): + """OpenAI types error.type as a string, so anything else on the exception is not one and + must not reach the wire the way the literal "None" used to.""" + + class _Carrier(Exception): + type = carried_type + + assert openai_error_type(_Carrier("boom"), 401) == "authentication_error" + + +def test_the_type_is_never_the_string_none_after_a_json_round_trip(): + """The bug this module exists for: json.dumps of a "None" default is indistinguishable + from a real type to a client's error handler.""" + payload = json.loads( + json.dumps( + { + "type": openai_error_type(HTTPException(status_code=400, detail="boom"), 400), + "param": openai_error_param(HTTPException(status_code=400, detail="boom")), + } + ) + ) + + assert payload == {"type": "invalid_request_error", "param": None} + + +def test_a_carried_param_names_the_offending_field(): + carried = ProxyException(message="Invalid purpose", type="invalid_request_error", param="purpose", code=400) + + assert openai_error_param(carried) == "purpose" + + +@pytest.mark.parametrize("exc", [HTTPException(status_code=400, detail="boom"), ValueError("boom"), None]) +def test_param_is_json_null_when_the_exception_names_no_field(exc): + assert openai_error_param(exc) is None + + +def test_a_non_string_carried_param_is_json_null(): + class _Carrier(Exception): + param = 42 + + assert openai_error_param(_Carrier("boom")) is None + + +def test_a_carried_status_code_wins_over_the_default(): + assert error_status_code(HTTPException(status_code=429, detail="slow down"), 400) == 429 + + +@pytest.mark.parametrize("default", [400, 500]) +def test_the_default_status_stands_when_the_exception_carries_none(default): + assert error_status_code(ValueError("boom"), default) == default + + +@pytest.mark.parametrize("carried_status", [True, False, "429", None, 429.0]) +def test_a_non_int_carried_status_falls_back_to_the_default(carried_status): + """True is an int in Python but not an HTTP status, and a stringified one would break + every caller that compares the code numerically.""" + + class _Carrier(Exception): + status_code = carried_status + + assert error_status_code(_Carrier("boom"), 500) == 500 + + +def test_a_status_carried_by_an_exception_drives_the_type_it_reports(): + """The two helpers compose at every call site: the status the exception carries is what + names its type, not the default the route would have used.""" + exc = HTTPException(status_code=403, detail="blocked by policy") + + assert openai_error_type(exc, error_status_code(exc, 400)) == "permission_error" diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py index bca97915347..123d54789bf 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py @@ -2920,9 +2920,9 @@ def test_unscoped_list_files_accepts_every_documented_purpose( def test_list_files_reports_a_bad_target_model_names_as_a_400( mocker: MockerFixture, monkeypatch, llm_router: Router ): - """The exception tail reports an HTTPException with its own status and error - type rather than relabelling it, so a client that branches on either keeps - reading the same thing off a bad request.""" + """The exception tail answers with the OpenAI error object a client can branch on: + the type its 400 status stands for, and a JSON null param rather than the literal + string "None" no OpenAI SDK has a case for.""" _setup_unscoped_list_files_route(mocker, monkeypatch, llm_router, _permissive_afile_list) response = _get_list_files("/v1/files?target_model_names=gpt-3.5-turbo,gpt-4o") @@ -2931,8 +2931,8 @@ def test_list_files_reports_a_bad_target_model_names_as_a_400( assert response.json() == { "error": { "message": "target_model_names on list files must be a list of one model name. Example: ['gpt-4o']", - "type": "None", - "param": "None", + "type": "invalid_request_error", + "param": None, "code": "400", } } diff --git a/tests/test_litellm/proxy/utils/helpers/test_error_helpers.py b/tests/test_litellm/proxy/utils/helpers/test_error_helpers.py index e73e3c151e0..dc30798df55 100644 --- a/tests/test_litellm/proxy/utils/helpers/test_error_helpers.py +++ b/tests/test_litellm/proxy/utils/helpers/test_error_helpers.py @@ -135,7 +135,7 @@ def test_handle_exception_on_proxy_happy_path_generic_exception_defaults_to_500( "message": "kaboom", "type": ProxyErrorTypes.internal_server_error.value, "code": "500", - "param": "None", + "param": None, } From 7a5b8bce7e42157cbfcbed306e2ea33672c7e5ea Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 3 Sep 2026 03:40:58 -0700 Subject: [PATCH 025/109] test(proxy): type the parametrized inputs of the error payload helper tests --- .../proxy/common_utils/test_openai_error_payload.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/tests/test_litellm/proxy/common_utils/test_openai_error_payload.py b/tests/test_litellm/proxy/common_utils/test_openai_error_payload.py index c165d7ffdb1..3b39ea706fe 100644 --- a/tests/test_litellm/proxy/common_utils/test_openai_error_payload.py +++ b/tests/test_litellm/proxy/common_utils/test_openai_error_payload.py @@ -26,7 +26,7 @@ from litellm.proxy.common_utils.openai_error_payload import ( (503, "internal_server_error"), ], ) -def test_status_code_decides_the_type_when_the_exception_carries_none(status_code, expected_type): +def test_status_code_decides_the_type_when_the_exception_carries_none(status_code: int, expected_type: str): """A route that raises a bare HTTPException carries no error type, so the status it answered with is the only thing left to name the OpenAI type from.""" assert openai_error_type(HTTPException(status_code=status_code, detail="boom"), status_code) == expected_type @@ -46,7 +46,7 @@ def test_a_carried_type_wins_over_the_one_the_status_would_imply(): @pytest.mark.parametrize("carried_type", [None, 400, {"type": "invalid_request_error"}, ["invalid_request_error"]]) -def test_a_non_string_carried_type_falls_back_to_the_status(carried_type): +def test_a_non_string_carried_type_falls_back_to_the_status(carried_type: object): """OpenAI types error.type as a string, so anything else on the exception is not one and must not reach the wire the way the literal "None" used to.""" @@ -78,7 +78,7 @@ def test_a_carried_param_names_the_offending_field(): @pytest.mark.parametrize("exc", [HTTPException(status_code=400, detail="boom"), ValueError("boom"), None]) -def test_param_is_json_null_when_the_exception_names_no_field(exc): +def test_param_is_json_null_when_the_exception_names_no_field(exc: Exception | None): assert openai_error_param(exc) is None @@ -94,12 +94,12 @@ def test_a_carried_status_code_wins_over_the_default(): @pytest.mark.parametrize("default", [400, 500]) -def test_the_default_status_stands_when_the_exception_carries_none(default): +def test_the_default_status_stands_when_the_exception_carries_none(default: int): assert error_status_code(ValueError("boom"), default) == default @pytest.mark.parametrize("carried_status", [True, False, "429", None, 429.0]) -def test_a_non_int_carried_status_falls_back_to_the_default(carried_status): +def test_a_non_int_carried_status_falls_back_to_the_default(carried_status: object): """True is an int in Python but not an HTTP status, and a stringified one would break every caller that compares the code numerically.""" From cedf35992bdec398e7623d28e45b5218e4eff9bf Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 3 Sep 2026 04:02:55 -0700 Subject: [PATCH 026/109] fix(proxy): give each spend-log queue monitor its own flush event `PrismaClient.spend_log_flush_requested` was an `asyncio.Event` built at import time, so it bound to whichever event loop first awaited it and every later loop got `RuntimeError: ... is bound to a different event loop` out of `_wait_for_spend_log_flush_request`. The queue monitor's blanket `except Exception` swallowed that into its error logger, so the flush silently never happened and the row sat in the worker's queue until the next poll. The monitor now creates its own Event inside the loop that awaits it and hands it to the client, and `request_spend_log_flush` signals through the client instead of the class. A request that arrives before the monitor is running is dropped and loses nothing, because the monitor reads the queue on its first pass before it ever waits. In CI this showed up as the proxy-endpoints shard flaking on test_monitor_spend_logs_queue_flushes_as_soon_as_one_is_requested whenever --dist=loadscope put the health-endpoint tests, which boot a proxy TestClient and start a monitor, on the same worker ahead of the spend-log tests. --- litellm/proxy/db/db_spend_update_writer.py | 2 +- litellm/proxy/utils.py | 22 ++++--- .../proxy/db/test_db_spend_update_writer.py | 7 +- .../prisma_and_spend/test_spend_functions.py | 66 +++++++++++++++++-- 4 files changed, 79 insertions(+), 18 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index e6880d521f1..18058f31385 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -940,7 +940,7 @@ class DBSpendUpdateWriter: await enqueue_spend_logs(prisma_client, (payload,)) if payload.get("call_type") in RESPONSES_SESSION_CALL_TYPES: - request_spend_log_flush() + request_spend_log_flush(prisma_client) else: verbose_proxy_logger.debug("prisma_client is None. Skipping writing spend logs to db.") diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index cab2bd6d9db..aa3076dbd19 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -3519,7 +3519,7 @@ class _StaleReadEngine: class PrismaClient: spend_log_transactions: list = [] _spend_log_transactions_lock = asyncio.Lock() - spend_log_flush_requested: ClassVar[asyncio.Event] = asyncio.Event() + spend_log_flush_requested: "asyncio.Event | None" = None spend_log_queue_bytes: ClassVar[int] = 0 spend_logs_queue_monitor_task: "asyncio.Task[None] | None" = None tool_usage_transactions: list["ToolUsageTransaction"] = [] @@ -6245,23 +6245,27 @@ async def enqueue_spend_logs( ) -def request_spend_log_flush() -> None: - """Wake the queue monitor now rather than leaving the rows for its next poll. +def request_spend_log_flush(prisma_client: PrismaClient) -> None: + """Wake this client's queue monitor now rather than leaving the rows for its next poll. The Responses API hands the client an id it can chain from straight away, and that lookup reads the DB, so the row cannot sit in this worker's queue for a poll interval. Repeated requests coalesce into the monitor's next pass, so the batching holds. + A request made before the monitor is running is dropped, and loses nothing: the + monitor reads the queue on its first pass, before it ever waits on a request. """ - PrismaClient.spend_log_flush_requested.set() + flush_requested: Final = prisma_client.spend_log_flush_requested + if flush_requested is not None: + flush_requested.set() -async def _wait_for_spend_log_flush_request(interval: float) -> bool: +async def _wait_for_spend_log_flush_request(flush_requested: asyncio.Event, interval: float) -> bool: """Wait out ``interval``, returning early and True when a flush was requested.""" try: - await asyncio.wait_for(PrismaClient.spend_log_flush_requested.wait(), timeout=interval) + await asyncio.wait_for(flush_requested.wait(), timeout=interval) except asyncio.TimeoutError: return False - PrismaClient.spend_log_flush_requested.clear() + flush_requested.clear() return True @@ -6681,6 +6685,8 @@ async def _monitor_spend_logs_queue( max_backoff: Final = 30.0 # Maximum backoff interval in seconds backoff_multiplier: Final = 1.5 # Exponential backoff multiplier current_interval = base_interval + flush_requested: Final = asyncio.Event() + prisma_client.spend_log_flush_requested = flush_requested # rebind-ok: the client owns its monitor's flush signal verbose_proxy_logger.info( "Starting spend logs queue monitor (threshold: %s, poll_interval: %ss)", threshold, base_interval @@ -6719,7 +6725,7 @@ async def _monitor_spend_logs_queue( # Exponential backoff when no logs to process current_interval = min(current_interval * backoff_multiplier, max_backoff) - if await _wait_for_spend_log_flush_request(current_interval): + if await _wait_for_spend_log_flush_request(flush_requested, current_interval): current_interval = base_interval except Exception as e: spend_log_error("Error in spend logs queue monitor: %s", str(e), exc=e) diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index 11ef911de3e..104dcf55e66 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -2941,11 +2941,9 @@ async def test_insert_spend_log_asks_for_an_immediate_flush_on_responses_calls(c A `previous_response_id` chained straight off the previous turn reads the DB, so a Responses row cannot sit in this worker's queue until the monitor's next poll. """ - from litellm.proxy.utils import PrismaClient - db_writer = DBSpendUpdateWriter() prisma = _tool_usage_prisma() - PrismaClient.spend_log_flush_requested.clear() + prisma.spend_log_flush_requested = asyncio.Event() await db_writer._insert_spend_log_to_db( payload={"request_id": "req-1", "call_type": call_type}, @@ -2953,8 +2951,7 @@ async def test_insert_spend_log_asks_for_an_immediate_flush_on_responses_calls(c ) assert prisma.spend_log_transactions == [{"request_id": "req-1", "call_type": call_type}] - assert PrismaClient.spend_log_flush_requested.is_set() is expects_flush - PrismaClient.spend_log_flush_requested.clear() + assert prisma.spend_log_flush_requested.is_set() is expects_flush @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py index a1eb88a7834..ed854d2c95c 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py @@ -538,10 +538,9 @@ async def test_monitor_spend_logs_queue_flushes_as_soon_as_one_is_requested( """ import litellm.constants as constants_mod import litellm.proxy.utils as utils_mod - from litellm.proxy.utils import PrismaClient, request_spend_log_flush + from litellm.proxy.utils import request_spend_log_flush monkeypatch.setattr(constants_mod, "SPEND_LOG_QUEUE_POLL_INTERVAL", 30.0, raising=False) - PrismaClient.spend_log_flush_requested.clear() mock_prisma_client.spend_log_transactions = [] mock_prisma_client.tool_usage_transactions = [] @@ -562,16 +561,75 @@ async def test_monitor_spend_logs_queue_flushes_as_soon_as_one_is_requested( try: await asyncio.sleep(0.05) assert not flushed.is_set() + assert isinstance(mock_prisma_client.spend_log_flush_requested, asyncio.Event) mock_prisma_client.spend_log_transactions.append(make_spend_log_row(request_id="r1")) - request_spend_log_flush() + request_spend_log_flush(mock_prisma_client) await asyncio.wait_for(flushed.wait(), timeout=5.0) finally: monitor.cancel() with suppress(asyncio.CancelledError): await monitor - PrismaClient.spend_log_flush_requested.clear() + + +def test_monitor_spend_logs_queue_flush_survives_an_earlier_event_loop( + mock_prisma_client: Any, + make_spend_log_row: Any, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A second monitor, started in a fresh event loop, is still woken by a flush request, + so a worker whose first loop is gone keeps flushing Responses rows instead of stalling. + """ + import litellm.constants as constants_mod + import litellm.proxy.utils as utils_mod + from litellm.proxy.utils import request_spend_log_flush + + monkeypatch.setattr(constants_mod, "SPEND_LOG_QUEUE_POLL_INTERVAL", 30.0, raising=False) + mock_prisma_client.tool_usage_transactions = [] + + async def _flush_once_under_a_monitor() -> None: + flushed: Final = asyncio.Event() + + async def _fake_job(*args: Any, **kwargs: Any) -> None: + flushed.set() + + monkeypatch.setattr(utils_mod, "update_spend_logs_job", _fake_job) + mock_prisma_client.spend_log_transactions = [] + + monitor: Final = asyncio.create_task( + _monitor_spend_logs_queue( + prisma_client=mock_prisma_client, + db_writer_client=None, + proxy_logging_obj=MagicMock(), + ) + ) + try: + await asyncio.sleep(0.05) + assert not flushed.is_set() + + mock_prisma_client.spend_log_transactions.append(make_spend_log_row(request_id="r1")) + request_spend_log_flush(mock_prisma_client) + + await asyncio.wait_for(flushed.wait(), timeout=5.0) + finally: + monitor.cancel() + with suppress(asyncio.CancelledError): + await monitor + + asyncio.run(_flush_once_under_a_monitor()) + asyncio.run(_flush_once_under_a_monitor()) + + +def test_request_spend_log_flush_is_a_no_op_before_the_monitor_starts(mock_prisma_client: Any) -> None: + """A Responses row enqueued before the monitor's first pass must not fail the request.""" + from litellm.proxy.utils import request_spend_log_flush + + mock_prisma_client.spend_log_flush_requested = None + + request_spend_log_flush(mock_prisma_client) + + assert mock_prisma_client.spend_log_flush_requested is None def test_raise_failed_update_spend_exception_emits_failure_handler() -> None: From 29bcb0eeb9b085492764cdd970bc4b28340f554b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 3 Sep 2026 04:55:57 -0700 Subject: [PATCH 027/109] test(proxy): pin that a flush requested before the monitor starts costs the row nothing The monitor reads the queue on its first pass, before it ever waits on a request, so dropping a request made while spend_log_flush_requested is still None delays nothing. Reordering the loop to wait first would turn that drop into a real delay for the Responses chaining flow, and now fails this test. --- .../prisma_and_spend/test_spend_functions.py | 38 +++++++++++++++++-- 1 file changed, 35 insertions(+), 3 deletions(-) diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py index ed854d2c95c..c8b87bd671e 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/test_spend_functions.py @@ -621,16 +621,48 @@ def test_monitor_spend_logs_queue_flush_survives_an_earlier_event_loop( asyncio.run(_flush_once_under_a_monitor()) -def test_request_spend_log_flush_is_a_no_op_before_the_monitor_starts(mock_prisma_client: Any) -> None: - """A Responses row enqueued before the monitor's first pass must not fail the request.""" +@pytest.mark.asyncio +async def test_flush_requested_before_the_monitor_starts_costs_the_row_nothing( + mock_prisma_client: Any, + make_spend_log_row: Any, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A Responses row enqueued before the monitor exists still reaches the DB on its first + pass, so dropping that early request delays nothing. + """ + import litellm.constants as constants_mod + import litellm.proxy.utils as utils_mod from litellm.proxy.utils import request_spend_log_flush + monkeypatch.setattr(constants_mod, "SPEND_LOG_QUEUE_POLL_INTERVAL", 30.0, raising=False) mock_prisma_client.spend_log_flush_requested = None + mock_prisma_client.tool_usage_transactions = [] + mock_prisma_client.spend_log_transactions = [make_spend_log_row(request_id="r1")] + + flushed: Final = asyncio.Event() + + async def _fake_job(*args: Any, **kwargs: Any) -> None: + flushed.set() + + monkeypatch.setattr(utils_mod, "update_spend_logs_job", _fake_job) request_spend_log_flush(mock_prisma_client) - assert mock_prisma_client.spend_log_flush_requested is None + monitor: Final = asyncio.create_task( + _monitor_spend_logs_queue( + prisma_client=mock_prisma_client, + db_writer_client=None, + proxy_logging_obj=MagicMock(), + ) + ) + try: + await asyncio.wait_for(flushed.wait(), timeout=5.0) + finally: + monitor.cancel() + with suppress(asyncio.CancelledError): + await monitor + def test_raise_failed_update_spend_exception_emits_failure_handler() -> None: proxy_logging = MagicMock() From 2a2c49ad4cd711aded30da0384c1a200ddfca5d7 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Thu, 3 Sep 2026 16:09:44 -0400 Subject: [PATCH 028/109] 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 029/109] 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 030/109] 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 031/109] 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 032/109] 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 033/109] 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 034/109] 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 035/109] 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 036/109] 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 037/109] 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 038/109] 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 039/109] 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 16fd14f53705dab68ee8c4b0fd2c9093c8e3cdac Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sat, 5 Sep 2026 18:03:24 -0700 Subject: [PATCH 040/109] fix(docker): ship pymongo in the proxy images for the MongoDB vector store The MongoDB Atlas vector store provider imports pymongo lazily from the opt-in `mongodb` extra, but none of the shipped images installed that extra. Any image-based deployment that configured a MongoDB vector store failed at search time with "requires the 'pymongo' package", which the user cannot fix without extending the image Adds `--extra mongodb` to every uv sync in the root Dockerfile, Dockerfile.database, Dockerfile.non_root, and the gateway component image. The backend component does not serve /vector_stores so it is left as is. The extra resolves from the existing uv.lock to pymongo 4.17.0 plus dnspython 2.8.0, no lock change needed --- Dockerfile | 2 ++ docker/Dockerfile.database | 2 ++ docker/Dockerfile.non_root | 3 +++ gateway/Dockerfile | 2 ++ 4 files changed, 9 insertions(+) diff --git a/Dockerfile b/Dockerfile index 0a92aa9a68c..1648ec69d13 100644 --- a/Dockerfile +++ b/Dockerfile @@ -67,6 +67,7 @@ RUN uv sync --frozen --no-install-project --no-install-workspace --no-default-gr --extra semantic-router \ --extra saml \ --extra bedrock-realtime \ + --extra mongodb \ --python python3.13 # Copy full source tree @@ -89,6 +90,7 @@ RUN uv sync --frozen --no-default-groups --no-editable \ --extra semantic-router \ --extra saml \ --extra bedrock-realtime \ + --extra mongodb \ --python python3.13 RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \ diff --git a/docker/Dockerfile.database b/docker/Dockerfile.database index e9ad2849bb2..cc81ad6b3d3 100644 --- a/docker/Dockerfile.database +++ b/docker/Dockerfile.database @@ -65,6 +65,7 @@ RUN uv sync --frozen --no-install-project --no-install-workspace --no-default-gr --extra semantic-router \ --extra saml \ --extra bedrock-realtime \ + --extra mongodb \ --python python3.13 # Copy full source tree @@ -87,6 +88,7 @@ RUN uv sync --frozen --no-default-groups --no-editable \ --extra semantic-router \ --extra saml \ --extra bedrock-realtime \ + --extra mongodb \ --python python3.13 RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \ diff --git a/docker/Dockerfile.non_root b/docker/Dockerfile.non_root index edf20e8bbff..358425af901 100644 --- a/docker/Dockerfile.non_root +++ b/docker/Dockerfile.non_root @@ -71,6 +71,7 @@ RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \ --extra semantic-router \ --extra saml \ --extra bedrock-realtime \ + --extra mongodb \ --python python3.13 # Copy full source tree @@ -99,6 +100,7 @@ RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \ --extra semantic-router \ --extra saml \ --extra bedrock-realtime \ + --extra mongodb \ --python python3.13 \ --no-sources-package litellm-proxy-extras; \ else \ @@ -109,6 +111,7 @@ RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \ --extra semantic-router \ --extra saml \ --extra bedrock-realtime \ + --extra mongodb \ --python python3.13; \ fi diff --git a/gateway/Dockerfile b/gateway/Dockerfile index 308d70a6b26..e42e488d57f 100644 --- a/gateway/Dockerfile +++ b/gateway/Dockerfile @@ -47,6 +47,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \ --extra extra_proxy \ --extra semantic-router \ --extra bedrock-realtime \ + --extra mongodb \ --python python3.13 # Stage 2 — copy source and install the project + workspace members. @@ -59,6 +60,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \ --extra extra_proxy \ --extra semantic-router \ --extra bedrock-realtime \ + --extra mongodb \ --python python3.13 RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \ From 48cc4efca38b46e1f561da4728954d06bbe223c7 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 20:37:17 -0700 Subject: [PATCH 041/109] fix(least-busy): share in-flight request counts across proxy workers Least-busy kept one dict of in-flight counts per model group in the router cache, which reads in-memory first, so every worker and replica routed on its own stale copy and each write overwrote the shared value. Counts now live in one key per deployment, incremented and read through Redis when the router has a Redis cache, and in the process-local cache otherwise. --- basedpyright-code-budget.json | 12 +- litellm/caching/redis_cache.py | 4 +- litellm/router_strategy/least_busy.py | 372 +++++++++--------- ruff-strict-budget.json | 10 +- .../local_testing/test_least_busy_routing.py | 46 ++- .../router_strategy/test_least_busy.py | 151 +++++++ type-discipline-budget.json | 8 +- 7 files changed, 376 insertions(+), 227 deletions(-) create mode 100644 tests/test_litellm/router_strategy/test_least_busy.py diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 0b0a61192e6..c21dc76743a 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -54,10 +54,10 @@ "limit": 0 }, "reportMissingParameterType": { - "limit": 5570 + "limit": 5551 }, "reportMissingTypeArgument": { - "limit": 15281 + "limit": 15277 }, "reportMissingTypeStubs": { "limit": 40 @@ -99,19 +99,19 @@ "limit": 0 }, "reportUnknownArgumentType": { - "limit": 44358 + "limit": 44357 }, "reportUnknownLambdaType": { "limit": 109 }, "reportUnknownMemberType": { - "limit": 38271 + "limit": 38240 }, "reportUnknownParameterType": { - "limit": 19584 + "limit": 19558 }, "reportUnknownVariableType": { - "limit": 29814 + "limit": 29781 }, "reportUnnecessaryCast": { "limit": 110 diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index aaee7188d86..d4c47914394 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -680,7 +680,7 @@ class RedisCache(BaseCache): # NON blocking - notify users Redis is throwing an exception print_verbose(f"litellm.caching.caching: set() - Got exception from REDIS : {e}") - def increment_cache(self, key, value: int, ttl: float | None = None, **kwargs) -> int: + def increment_cache(self, key, value: int, ttl: float | None = None, refresh_ttl: bool = False, **kwargs) -> int: _redis_client: Final = self.redis_client start_time = time.time() set_ttl: Final = self.get_ttl(ttl=ttl) @@ -701,7 +701,7 @@ class RedisCache(BaseCache): if set_ttl is not None: # check if key already has ttl, if not -> set ttl start_time = time.time() - current_ttl: Final = _redis_client.ttl(key) + current_ttl: Final = -1 if refresh_ttl else _redis_client.ttl(key) end_time = time.time() _duration = end_time - start_time self.service_logger_obj.service_success_hook( diff --git a/litellm/router_strategy/least_busy.py b/litellm/router_strategy/least_busy.py index 1433e8ba4d4..de6a4c4f59a 100644 --- a/litellm/router_strategy/least_busy.py +++ b/litellm/router_strategy/least_busy.py @@ -1,17 +1,96 @@ -#### What this does #### -# identifies least busy deployment -# How is this achieved? -# - Before each call, have the router print the state of requests {"deployment": "requests_in_flight"} -# - use litellm.input_callbacks to log when a request is just about to be made to a model - {"deployment-id": traffic} -# - use litellm.success + failure callbacks to log when a request completed -# - in get_available_deployment, for a given model group name -> pick based on traffic - -import random +from collections.abc import Mapping, Sequence from typing import Final +from pydantic import TypeAdapter, ValidationError +from typing_extensions import ReadOnly, TypedDict + +from litellm._logging import verbose_router_logger from litellm.caching.caching import DualCache from litellm.integrations.custom_logger import CustomLogger +IN_FLIGHT_COUNT_TTL_SECONDS: Final = 60 * 60 + + +class _ModelInfo(TypedDict, total=False): + id: ReadOnly[str | int | None] + + +class _Metadata(TypedDict, total=False): + model_group: ReadOnly[str | None] + + +class _LitellmParams(TypedDict, total=False): + metadata: ReadOnly[_Metadata | None] + model_info: ReadOnly[_ModelInfo | None] + + +class _CallKwargs(TypedDict, total=False): + litellm_params: ReadOnly[_LitellmParams | None] + + +class _DeploymentModelInfo(TypedDict): + id: ReadOnly[str | int] + + +class _Deployment(TypedDict): + model_info: ReadOnly[_DeploymentModelInfo] + + +_CALL_KWARGS: Final = TypeAdapter(_CallKwargs) +_DEPLOYMENTS: Final = TypeAdapter(list[_Deployment]) +_REDIS_COUNTS: Final = TypeAdapter(dict[str, float | None]) +_MEMORY_COUNTS: Final = TypeAdapter(tuple[float | None, ...]) + + +def _request_count_key(model_group: str, deployment_id: str) -> str: + return f"{model_group}_request_count:{deployment_id}" + + +def _deployment_ref(kwargs: Mapping[str, object]) -> tuple[str, str] | None: + try: + call: Final = _CALL_KWARGS.validate_python(kwargs) + except ValidationError: + return None + litellm_params: Final = call.get("litellm_params") + metadata: Final = litellm_params.get("metadata") if litellm_params else None + model_info: Final = litellm_params.get("model_info") if litellm_params else None + model_group: Final = metadata.get("model_group") if metadata else None + deployment_id: Final = model_info.get("id") if model_info else None + if model_group is None or deployment_id is None: + return None + return model_group, str(deployment_id) + + +def _request_count_keys(model_group: str, healthy_deployments: Sequence[Mapping[str, object]]) -> tuple[str, ...]: + return tuple( + _request_count_key(model_group, str(deployment["model_info"]["id"])) + for deployment in _DEPLOYMENTS.validate_python(healthy_deployments) + ) + + +def _as_count(value: float | None) -> int: + return 0 if value is None else int(value) + + +def _least_busy( + healthy_deployments: Sequence[Mapping[str, object]], counts: tuple[int, ...] +) -> Mapping[str, object] | None: + if not healthy_deployments: + return None + return healthy_deployments[min(range(len(healthy_deployments)), key=lambda index: counts[index])] + + +def _warn_unreadable(model_group: str, error: Exception) -> None: + verbose_router_logger.warning( + "least-busy routing could not read the in-flight counts for %s, treating every deployment as idle: %s", + model_group, + error, + ) + + +def _warn_unwritable(key: str, error: Exception) -> None: + verbose_router_logger.warning("least-busy routing could not update the in-flight count under %s: %s", key, error) + class LeastBusyLoggingHandler(CustomLogger): test_flag: bool = False @@ -21,194 +100,101 @@ class LeastBusyLoggingHandler(CustomLogger): def __init__(self, router_cache: DualCache): self.router_cache = router_cache - def log_pre_api_call(self, model, messages, kwargs): - """ - Log when a model is being used. + def log_pre_api_call(self, model: str, messages: object, kwargs: Mapping[str, object]) -> None: + self._increment(kwargs, 1) - Caching based on model group. - """ - try: - if kwargs["litellm_params"].get("metadata") is None: - pass - else: - model_group: Final = kwargs["litellm_params"]["metadata"].get("model_group", None) - id = kwargs["litellm_params"].get("model_info", {}).get("id", None) - if model_group is None or id is None: - return - elif isinstance(id, int): - id = str(id) + def log_success_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + self._increment(kwargs, -1) + if self.test_flag: + self.logged_success += 1 - request_count_api_key: Final = f"{model_group}_request_count" - # update cache - request_count_dict: Final = self.router_cache.get_cache(key=request_count_api_key) or {} - request_count_dict[id] = request_count_dict.get(id, 0) + 1 + def log_failure_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + self._increment(kwargs, -1) + if self.test_flag: + self.logged_failure += 1 - self.router_cache.set_cache(key=request_count_api_key, value=request_count_dict) - except Exception: - pass + async def async_log_success_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + await self._async_increment(kwargs, -1) + if self.test_flag: + self.logged_success += 1 - def log_success_event(self, kwargs, response_obj, start_time, end_time): - try: - if kwargs["litellm_params"].get("metadata") is None: - pass - else: - model_group: Final = kwargs["litellm_params"]["metadata"].get("model_group", None) - - id = kwargs["litellm_params"].get("model_info", {}).get("id", None) - if model_group is None or id is None: - return - elif isinstance(id, int): - id = str(id) - - request_count_api_key: Final = f"{model_group}_request_count" - # decrement count in cache - request_count_dict: Final = self.router_cache.get_cache(key=request_count_api_key) or {} - request_count_value: Final[int | None] = request_count_dict.get(id, 0) - if request_count_value is None: - return - request_count_dict[id] = request_count_value - 1 - self.router_cache.set_cache(key=request_count_api_key, value=request_count_dict) - - ### TESTING ### - if self.test_flag: - self.logged_success += 1 - except Exception: - pass - - def log_failure_event(self, kwargs, response_obj, start_time, end_time): - try: - if kwargs["litellm_params"].get("metadata") is None: - pass - else: - model_group: Final = kwargs["litellm_params"]["metadata"].get("model_group", None) - id = kwargs["litellm_params"].get("model_info", {}).get("id", None) - if model_group is None or id is None: - return - elif isinstance(id, int): - id = str(id) - - request_count_api_key: Final = f"{model_group}_request_count" - # decrement count in cache - request_count_dict: Final = self.router_cache.get_cache(key=request_count_api_key) or {} - request_count_value: Final[int | None] = request_count_dict.get(id, 0) - if request_count_value is None: - return - request_count_dict[id] = request_count_value - 1 - self.router_cache.set_cache(key=request_count_api_key, value=request_count_dict) - - ### TESTING ### - if self.test_flag: - self.logged_failure += 1 - except Exception: - pass - - async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): - try: - if kwargs["litellm_params"].get("metadata") is None: - pass - else: - model_group: Final = kwargs["litellm_params"]["metadata"].get("model_group", None) - - id = kwargs["litellm_params"].get("model_info", {}).get("id", None) - if model_group is None or id is None: - return - elif isinstance(id, int): - id = str(id) - - request_count_api_key: Final = f"{model_group}_request_count" - # decrement count in cache - request_count_dict: Final = await self.router_cache.async_get_cache(key=request_count_api_key) or {} - request_count_value: Final[int | None] = request_count_dict.get(id, 0) - if request_count_value is None: - return - request_count_dict[id] = request_count_value - 1 - await self.router_cache.async_set_cache(key=request_count_api_key, value=request_count_dict) - - ### TESTING ### - if self.test_flag: - self.logged_success += 1 - except Exception: - pass - - async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): - try: - if kwargs["litellm_params"].get("metadata") is None: - pass - else: - model_group: Final = kwargs["litellm_params"]["metadata"].get("model_group", None) - id = kwargs["litellm_params"].get("model_info", {}).get("id", None) - if model_group is None or id is None: - return - elif isinstance(id, int): - id = str(id) - - request_count_api_key: Final = f"{model_group}_request_count" - # decrement count in cache - request_count_dict: Final = await self.router_cache.async_get_cache(key=request_count_api_key) or {} - request_count_value: Final[int | None] = request_count_dict.get(id, 0) - if request_count_value is None: - return - request_count_dict[id] = request_count_value - 1 - await self.router_cache.async_set_cache(key=request_count_api_key, value=request_count_dict) - - ### TESTING ### - if self.test_flag: - self.logged_failure += 1 - except Exception: - pass - - def _get_available_deployments( - self, - healthy_deployments: list, - all_deployments: dict, - ): - """ - Helper to get deployments using least busy strategy - """ - for d in healthy_deployments: - ## if healthy deployment not yet used - if d["model_info"]["id"] not in all_deployments: - all_deployments[d["model_info"]["id"]] = 0 - # map deployment to id - # pick least busy deployment - min_traffic = float("inf") - min_deployment = None - for k, v in all_deployments.items(): - if v < min_traffic: - min_traffic = v - min_deployment = k - if min_deployment is not None: - ## check if min deployment is a string, if so, cast it to int - for m in healthy_deployments: - if m["model_info"]["id"] == min_deployment: - return m - min_deployment = random.choice(healthy_deployments) - else: - min_deployment = random.choice(healthy_deployments) - return min_deployment + async def async_log_failure_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + await self._async_increment(kwargs, -1) + if self.test_flag: + self.logged_failure += 1 def get_available_deployments( - self, - model_group: str, - healthy_deployments: list, - ): - """ - Sync helper to get deployments using least busy strategy - """ - request_count_api_key: Final = f"{model_group}_request_count" - all_deployments: Final = self.router_cache.get_cache(key=request_count_api_key) or {} - return self._get_available_deployments( - healthy_deployments=healthy_deployments, - all_deployments=all_deployments, - ) + self, model_group: str, healthy_deployments: Sequence[Mapping[str, object]] + ) -> Mapping[str, object] | None: + keys: Final = _request_count_keys(model_group, healthy_deployments) + try: + counts: Final = tuple(_as_count(value) for value in self._read_counts(keys)) + except Exception as e: + _warn_unreadable(model_group, e) + return _least_busy(healthy_deployments, (0,) * len(keys)) + return _least_busy(healthy_deployments, counts) - async def async_get_available_deployments(self, model_group: str, healthy_deployments: list): - """ - Async helper to get deployments using least busy strategy - """ - request_count_api_key: Final = f"{model_group}_request_count" - all_deployments: Final = await self.router_cache.async_get_cache(key=request_count_api_key) or {} - return self._get_available_deployments( - healthy_deployments=healthy_deployments, - all_deployments=all_deployments, - ) + async def async_get_available_deployments( + self, model_group: str, healthy_deployments: Sequence[Mapping[str, object]] + ) -> Mapping[str, object] | None: + keys: Final = _request_count_keys(model_group, healthy_deployments) + try: + counts: Final = tuple(_as_count(value) for value in await self._async_read_counts(keys)) + except Exception as e: + _warn_unreadable(model_group, e) + return _least_busy(healthy_deployments, (0,) * len(keys)) + return _least_busy(healthy_deployments, counts) + + def _increment(self, kwargs: Mapping[str, object], delta: int) -> None: + ref: Final = _deployment_ref(kwargs) + if ref is None: + return + key: Final = _request_count_key(*ref) + redis_cache: Final = self.router_cache.redis_cache + try: + if redis_cache is None: + self.router_cache.increment_cache(key, delta, local_only=True, ttl=IN_FLIGHT_COUNT_TTL_SECONDS) + else: + redis_cache.increment_cache(key, delta, ttl=IN_FLIGHT_COUNT_TTL_SECONDS, refresh_ttl=True) + except Exception as e: + _warn_unwritable(key, e) + + async def _async_increment(self, kwargs: Mapping[str, object], delta: int) -> None: + ref: Final = _deployment_ref(kwargs) + if ref is None: + return + key: Final = _request_count_key(*ref) + redis_cache: Final = self.router_cache.redis_cache + try: + if redis_cache is None: + await self.router_cache.async_increment_cache( + key, delta, local_only=True, ttl=IN_FLIGHT_COUNT_TTL_SECONDS + ) + else: + await redis_cache.async_increment(key, delta, ttl=IN_FLIGHT_COUNT_TTL_SECONDS, refresh_ttl=True) + except Exception as e: + _warn_unwritable(key, e) + + def _read_counts(self, keys: tuple[str, ...]) -> tuple[float | None, ...]: + redis_cache: Final = self.router_cache.redis_cache + if redis_cache is None: + return _MEMORY_COUNTS.validate_python(self.router_cache.batch_get_cache(list(keys), local_only=True)) + by_key: Final = _REDIS_COUNTS.validate_python(redis_cache.batch_get_cache(key_list=list(keys))) + return tuple(by_key.get(key) for key in keys) + + async def _async_read_counts(self, keys: tuple[str, ...]) -> tuple[float | None, ...]: + redis_cache: Final = self.router_cache.redis_cache + if redis_cache is None: + return _MEMORY_COUNTS.validate_python( + await self.router_cache.async_batch_get_cache(list(keys), local_only=True) + ) + by_key: Final = _REDIS_COUNTS.validate_python(await redis_cache.async_batch_get_cache(key_list=list(keys))) + return tuple(by_key.get(key) for key in keys) diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index fd7b30bc314..ceff8c6e15c 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -1,6 +1,6 @@ { "ANN001": { - "limit": 2956 + "limit": 2937 }, "ANN002": { "limit": 71 @@ -9,10 +9,10 @@ "limit": 806 }, "ANN201": { - "limit": 1979 + "limit": 1972 }, "ANN202": { - "limit": 831 + "limit": 830 }, "ANN204": { "limit": 683 @@ -57,7 +57,7 @@ "limit": 3 }, "BLE001": { - "limit": 2916 + "limit": 2915 }, "C401": { "limit": 8 @@ -189,7 +189,7 @@ "limit": 0 }, "S110": { - "limit": 217 + "limit": 212 }, "S112": { "limit": 22 diff --git a/tests/local_testing/test_least_busy_routing.py b/tests/local_testing/test_least_busy_routing.py index 18ab8bf779d..fb83d4e601f 100644 --- a/tests/local_testing/test_least_busy_routing.py +++ b/tests/local_testing/test_least_busy_routing.py @@ -33,8 +33,8 @@ def test_model_added(): } } least_busy_logger.log_pre_api_call(model="test", messages=[], kwargs=kwargs) - request_count_api_key = f"gpt-3.5-turbo_request_count" - assert test_cache.get_cache(key=request_count_api_key) is not None + request_count_api_key = "gpt-3.5-turbo_request_count:1234" + assert test_cache.get_cache(key=request_count_api_key) == 1 def test_get_available_deployments(): @@ -52,8 +52,8 @@ def test_get_available_deployments(): } } least_busy_logger.log_pre_api_call(model="test", messages=[], kwargs=kwargs) - request_count_api_key = f"{model_group}_request_count" - assert test_cache.get_cache(key=request_count_api_key) is not None + request_count_api_key = f"{model_group}_request_count:1234" + assert test_cache.get_cache(key=request_count_api_key) == 1 # test_get_available_deployments() @@ -104,15 +104,20 @@ async def test_router_get_available_deployments(async_test): router.leastbusy_logger.test_flag = True model_group = "azure-model" - request_count_dict = {1: 10, 2: 54, 3: 100} - cache_key = f"{model_group}_request_count" + request_count_dict = {"1": 10, "2": 54, "3": 100} + cache_keys = { + deployment_id: f"{model_group}_request_count:{deployment_id}" + for deployment_id in request_count_dict + } if async_test is True: - await router.cache.async_set_cache(key=cache_key, value=request_count_dict) + for deployment_id, count in request_count_dict.items(): + await router.cache.async_set_cache(key=cache_keys[deployment_id], value=count) deployment = await router.async_get_available_deployment( model=model_group, messages=None, request_kwargs={} ) else: - router.cache.set_cache(key=cache_key, value=request_count_dict) + for deployment_id, count in request_count_dict.items(): + router.cache.set_cache(key=cache_keys[deployment_id], value=count) deployment = router.get_available_deployment(model=model_group, messages=None) print(f"deployment: {deployment}") assert deployment["model_info"]["id"] == "1" @@ -124,15 +129,18 @@ async def test_router_get_available_deployments(async_test): messages=[{"role": "user", "content": "Hey, how's it going?"}], ) - return_dict = router.cache.get_cache(key=cache_key) - # wait 2 seconds time.sleep(2) + return_dict = { + deployment_id: router.cache.get_cache(key=cache_key) + for deployment_id, cache_key in cache_keys.items() + } + assert router.leastbusy_logger.logged_success == 1 - assert return_dict[1] == 10 - assert return_dict[2] == 54 - assert return_dict[3] == 100 + assert return_dict["1"] == 10 + assert return_dict["2"] == 54 + assert return_dict["3"] == 100 ## Test with Real calls ## @@ -192,9 +200,11 @@ async def test_router_atext_completion_streaming(): await asyncio.sleep(random.uniform(0, 2)) await router.atext_completion(model=model, prompt=prompt, stream=True) - cache_key = f"{model}_request_count" ## check if calls equally distributed - cache_dict = router.cache.get_cache(key=cache_key) + cache_dict = { + deployment_id: router.cache.get_cache(key=f"{model}_request_count:{deployment_id}") + for deployment_id in ("1", "2", "3") + } for k, v in cache_dict.items(): assert v == 1, f"Failed. K={k} called v={v} times, cache_dict={cache_dict}" @@ -259,8 +269,10 @@ async def test_router_completion_streaming(): await asyncio.sleep(random.uniform(0, 2)) await router.acompletion(model=model, messages=messages, stream=True) - cache_key = f"{model}_request_count" ## check if calls equally distributed - cache_dict = router.cache.get_cache(key=cache_key) + cache_dict = { + deployment_id: router.cache.get_cache(key=f"{model}_request_count:{deployment_id}") + for deployment_id in ("1", "2", "3") + } for k, v in cache_dict.items(): assert v == 1, f"Failed. K={k} called v={v} times, cache_dict={cache_dict}" diff --git a/tests/test_litellm/router_strategy/test_least_busy.py b/tests/test_litellm/router_strategy/test_least_busy.py new file mode 100644 index 00000000000..32a226a5d5b --- /dev/null +++ b/tests/test_litellm/router_strategy/test_least_busy.py @@ -0,0 +1,151 @@ +import json +from typing import Final + +import pytest + +from litellm.caching.caching import DualCache +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.router_strategy.least_busy import IN_FLIGHT_COUNT_TTL_SECONDS, LeastBusyLoggingHandler + +GROUP: Final = "least-busy-group" +DEPLOYMENT_A: Final[dict[str, object]] = {"model_info": {"id": "dep-a"}} +DEPLOYMENT_B: Final[dict[str, object]] = {"model_info": {"id": "dep-b"}} +HEALTHY: Final = [DEPLOYMENT_A, DEPLOYMENT_B] + + +def _call_kwargs(deployment_id: str) -> dict[str, object]: + return {"litellm_params": {"metadata": {"model_group": GROUP}, "model_info": {"id": deployment_id}}} + + +class SharedRedisCounters: + """Stores JSON strings and hands back a fresh object per read, the way a real Redis client does.""" + + def __init__(self) -> None: + self.encoded: dict[str, str] = {} + self.ttls: dict[str, float] = {} + + def count(self, key: str) -> object: + raw: Final = self.encoded.get(key) + return None if raw is None else json.loads(raw) + + def get_cache(self, key: str, **kwargs: object) -> object: + return self.count(key) + + def set_cache(self, key: str, value: object, **kwargs: object) -> None: + self.encoded[key] = json.dumps(value) + + async def async_get_cache(self, key: str, **kwargs: object) -> object: + return self.count(key) + + async def async_set_cache(self, key: str, value: object, **kwargs: object) -> None: + self.set_cache(key, value) + + def increment_cache(self, key: str, value: int, ttl: float | None = None, refresh_ttl: bool = False) -> int: + current: Final = self.count(key) or 0 + assert isinstance(current, int) + incremented: Final = current + value + self.encoded[key] = json.dumps(incremented) + if ttl is not None and (refresh_ttl or key not in self.ttls): + self.ttls[key] = ttl + return incremented + + async def async_increment(self, key: str, value: float, ttl: int | None = None, refresh_ttl: bool = False) -> float: + return self.increment_cache(key, int(value), ttl, refresh_ttl) + + def batch_get_cache(self, key_list: list[str], **kwargs: object) -> dict[str, object]: + return {key: self.count(key) for key in key_list} + + async def async_batch_get_cache(self, key_list: list[str], **kwargs: object) -> dict[str, object]: + return self.batch_get_cache(key_list) + + +def _worker(shared: SharedRedisCounters | None) -> LeastBusyLoggingHandler: + cache: Final = DualCache(in_memory_cache=InMemoryCache(), redis_cache=shared) # pyright: ignore[reportArgumentType] # duck-typed Redis double + return LeastBusyLoggingHandler(router_cache=cache) + + +@pytest.mark.asyncio +async def test_worker_routes_around_a_request_another_worker_started() -> None: + shared: Final = SharedRedisCounters() + streaming_worker: Final = _worker(shared) + picking_worker: Final = _worker(shared) + + picking_worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + await picking_worker.async_log_success_event(_call_kwargs("dep-a"), None, None, None) + + streaming_worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + + assert await picking_worker.async_get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B + + await streaming_worker.async_log_success_event(_call_kwargs("dep-a"), None, None, None) + + assert await picking_worker.async_get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A + + +def test_sync_pick_reads_the_shared_counts() -> None: + shared: Final = SharedRedisCounters() + streaming_worker: Final = _worker(shared) + picking_worker: Final = _worker(shared) + + picking_worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + picking_worker.log_success_event(_call_kwargs("dep-a"), None, None, None) + + streaming_worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + + assert picking_worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B + + streaming_worker.log_failure_event(_call_kwargs("dep-a"), None, None, None) + + assert picking_worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A + + +def test_redis_counts_keep_a_refreshed_ttl() -> None: + shared: Final = SharedRedisCounters() + worker: Final = _worker(shared) + + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + worker.log_success_event(_call_kwargs("dep-a"), None, None, None) + + assert shared.count(f"{GROUP}_request_count:dep-a") == 0 + assert shared.ttls == {f"{GROUP}_request_count:dep-a": IN_FLIGHT_COUNT_TTL_SECONDS} + + +@pytest.mark.asyncio +async def test_counts_stay_in_memory_without_redis() -> None: + worker: Final = _worker(None) + + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + + assert await worker.async_get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B + assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B + + await worker.async_log_success_event(_call_kwargs("dep-a"), None, None, None) + + assert await worker.async_get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A + assert worker.router_cache.get_cache(f"{GROUP}_request_count:dep-a") == 0 + + +class UnavailableRedis(SharedRedisCounters): + def batch_get_cache(self, key_list: list[str], **kwargs: object) -> dict[str, object]: + raise ConnectionError("redis is down") + + def increment_cache(self, key: str, value: int, ttl: float | None = None, refresh_ttl: bool = False) -> int: + raise ConnectionError("redis is down") + + +def test_redis_outage_never_fails_the_request() -> None: + worker: Final = _worker(UnavailableRedis()) + + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + + assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A + + +def test_calls_without_a_deployment_are_ignored() -> None: + shared: Final = SharedRedisCounters() + worker: Final = _worker(shared) + + worker.log_pre_api_call(model="m", messages=[], kwargs={"litellm_params": {"metadata": None}}) + worker.log_pre_api_call(model="m", messages=[], kwargs={}) + + assert shared.encoded == {} diff --git a/type-discipline-budget.json b/type-discipline-budget.json index e7186dfe186..e413e6db2c6 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 22180 + "limit": 22176 }, "LIT002": { - "limit": 26729 + "limit": 26721 }, "LIT003": { "limit": 261 @@ -27,10 +27,10 @@ "limit": 0 }, "LIT010": { - "limit": 16426 + "limit": 16412 }, "LIT011": { - "limit": 5506 + "limit": 5505 }, "LIT012": { "limit": 4486 From c5aa4f07185c91f7c055e614734d5e216c6d043d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 21:14:58 -0700 Subject: [PATCH 042/109] fix(least-busy): fall back to per-worker counts when Redis is unreadable Keep each worker's own in-flight counter up to date alongside the shared one, so a Redis outage routes on that worker's counts the way it did before this branch instead of treating every deployment as idle. Floor a counter at zero when a decrement finds the key gone, which happens when a request outlives the 1 hour TTL, so an expired counter cannot settle at -1 and win every pick. --- litellm/router_strategy/least_busy.py | 97 +++++++++++-------- .../router_strategy/test_least_busy.py | 39 +++++++- 2 files changed, 95 insertions(+), 41 deletions(-) diff --git a/litellm/router_strategy/least_busy.py b/litellm/router_strategy/least_busy.py index de6a4c4f59a..00d27b5f8d9 100644 --- a/litellm/router_strategy/least_busy.py +++ b/litellm/router_strategy/least_busy.py @@ -39,7 +39,7 @@ class _Deployment(TypedDict): _CALL_KWARGS: Final = TypeAdapter(_CallKwargs) _DEPLOYMENTS: Final = TypeAdapter(list[_Deployment]) _REDIS_COUNTS: Final = TypeAdapter(dict[str, float | None]) -_MEMORY_COUNTS: Final = TypeAdapter(tuple[float | None, ...]) +_MEMORY_COUNTS: Final = TypeAdapter(tuple[float | None, ...] | None) def _request_count_key(model_group: str, deployment_id: str) -> str: @@ -68,8 +68,20 @@ def _request_count_keys(model_group: str, healthy_deployments: Sequence[Mapping[ ) -def _as_count(value: float | None) -> int: - return 0 if value is None else int(value) +def _as_counts(values: Sequence[float | None]) -> tuple[int, ...]: + return tuple(0 if value is None else int(value) for value in values) + + +def _shared_counts(raw: object, keys: tuple[str, ...]) -> tuple[int, ...]: + by_key: Final = _REDIS_COUNTS.validate_python(raw) + return _as_counts([by_key.get(key) for key in keys]) + + +def _local_counts(raw: object, keys: tuple[str, ...]) -> tuple[int, ...]: + values: Final = _MEMORY_COUNTS.validate_python(raw) + if values is None or len(values) != len(keys): + return (0,) * len(keys) + return _as_counts(values) def _least_busy( @@ -82,7 +94,8 @@ def _least_busy( def _warn_unreadable(model_group: str, error: Exception) -> None: verbose_router_logger.warning( - "least-busy routing could not read the in-flight counts for %s, treating every deployment as idle: %s", + "least-busy routing could not read the shared in-flight counts for %s, " + "falling back to this worker's own counts: %s", model_group, error, ) @@ -135,23 +148,31 @@ class LeastBusyLoggingHandler(CustomLogger): self, model_group: str, healthy_deployments: Sequence[Mapping[str, object]] ) -> Mapping[str, object] | None: keys: Final = _request_count_keys(model_group, healthy_deployments) - try: - counts: Final = tuple(_as_count(value) for value in self._read_counts(keys)) - except Exception as e: - _warn_unreadable(model_group, e) - return _least_busy(healthy_deployments, (0,) * len(keys)) - return _least_busy(healthy_deployments, counts) + redis_cache: Final = self.router_cache.redis_cache + if redis_cache is not None: + try: + shared: Final = _shared_counts(redis_cache.batch_get_cache(key_list=list(keys)), keys) + except Exception as e: + _warn_unreadable(model_group, e) + else: + return _least_busy(healthy_deployments, shared) + local: Final = _local_counts(self.router_cache.batch_get_cache(list(keys), local_only=True), keys) + return _least_busy(healthy_deployments, local) async def async_get_available_deployments( self, model_group: str, healthy_deployments: Sequence[Mapping[str, object]] ) -> Mapping[str, object] | None: keys: Final = _request_count_keys(model_group, healthy_deployments) - try: - counts: Final = tuple(_as_count(value) for value in await self._async_read_counts(keys)) - except Exception as e: - _warn_unreadable(model_group, e) - return _least_busy(healthy_deployments, (0,) * len(keys)) - return _least_busy(healthy_deployments, counts) + redis_cache: Final = self.router_cache.redis_cache + if redis_cache is not None: + try: + shared: Final = _shared_counts(await redis_cache.async_batch_get_cache(key_list=list(keys)), keys) + except Exception as e: + _warn_unreadable(model_group, e) + else: + return _least_busy(healthy_deployments, shared) + local: Final = _local_counts(await self.router_cache.async_batch_get_cache(list(keys), local_only=True), keys) + return _least_busy(healthy_deployments, local) def _increment(self, kwargs: Mapping[str, object], delta: int) -> None: ref: Final = _deployment_ref(kwargs) @@ -160,10 +181,16 @@ class LeastBusyLoggingHandler(CustomLogger): key: Final = _request_count_key(*ref) redis_cache: Final = self.router_cache.redis_cache try: + local: Final = self.router_cache.increment_cache( + key, delta, local_only=True, ttl=IN_FLIGHT_COUNT_TTL_SECONDS + ) + if local < 0: + self.router_cache.set_cache(key, 0, local_only=True, ttl=IN_FLIGHT_COUNT_TTL_SECONDS) if redis_cache is None: - self.router_cache.increment_cache(key, delta, local_only=True, ttl=IN_FLIGHT_COUNT_TTL_SECONDS) - else: - redis_cache.increment_cache(key, delta, ttl=IN_FLIGHT_COUNT_TTL_SECONDS, refresh_ttl=True) + return + shared: Final = redis_cache.increment_cache(key, delta, ttl=IN_FLIGHT_COUNT_TTL_SECONDS, refresh_ttl=True) + if shared < 0: + redis_cache.set_cache(key, 0, ttl=IN_FLIGHT_COUNT_TTL_SECONDS) except Exception as e: _warn_unwritable(key, e) @@ -174,27 +201,17 @@ class LeastBusyLoggingHandler(CustomLogger): key: Final = _request_count_key(*ref) redis_cache: Final = self.router_cache.redis_cache try: + local: Final = await self.router_cache.async_increment_cache( + key, delta, local_only=True, ttl=IN_FLIGHT_COUNT_TTL_SECONDS + ) + if local is not None and local < 0: + await self.router_cache.async_set_cache(key, 0, local_only=True, ttl=IN_FLIGHT_COUNT_TTL_SECONDS) if redis_cache is None: - await self.router_cache.async_increment_cache( - key, delta, local_only=True, ttl=IN_FLIGHT_COUNT_TTL_SECONDS - ) - else: - await redis_cache.async_increment(key, delta, ttl=IN_FLIGHT_COUNT_TTL_SECONDS, refresh_ttl=True) + return + shared: Final = await redis_cache.async_increment( + key, delta, ttl=IN_FLIGHT_COUNT_TTL_SECONDS, refresh_ttl=True + ) + if shared < 0: + await redis_cache.async_set_cache(key, 0, ttl=IN_FLIGHT_COUNT_TTL_SECONDS) except Exception as e: _warn_unwritable(key, e) - - def _read_counts(self, keys: tuple[str, ...]) -> tuple[float | None, ...]: - redis_cache: Final = self.router_cache.redis_cache - if redis_cache is None: - return _MEMORY_COUNTS.validate_python(self.router_cache.batch_get_cache(list(keys), local_only=True)) - by_key: Final = _REDIS_COUNTS.validate_python(redis_cache.batch_get_cache(key_list=list(keys))) - return tuple(by_key.get(key) for key in keys) - - async def _async_read_counts(self, keys: tuple[str, ...]) -> tuple[float | None, ...]: - redis_cache: Final = self.router_cache.redis_cache - if redis_cache is None: - return _MEMORY_COUNTS.validate_python( - await self.router_cache.async_batch_get_cache(list(keys), local_only=True) - ) - by_key: Final = _REDIS_COUNTS.validate_python(await redis_cache.async_batch_get_cache(key_list=list(keys))) - return tuple(by_key.get(key) for key in keys) diff --git a/tests/test_litellm/router_strategy/test_least_busy.py b/tests/test_litellm/router_strategy/test_least_busy.py index 32a226a5d5b..55702c6fe73 100644 --- a/tests/test_litellm/router_strategy/test_least_busy.py +++ b/tests/test_litellm/router_strategy/test_least_busy.py @@ -133,14 +133,51 @@ class UnavailableRedis(SharedRedisCounters): raise ConnectionError("redis is down") -def test_redis_outage_never_fails_the_request() -> None: +@pytest.mark.asyncio +async def test_a_redis_outage_falls_back_to_this_workers_own_counts() -> None: worker: Final = _worker(UnavailableRedis()) worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B + assert await worker.async_get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B + + await worker.async_log_success_event(_call_kwargs("dep-a"), None, None, None) + assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A +def test_a_shared_counter_that_expired_mid_request_cannot_go_negative() -> None: + shared: Final = SharedRedisCounters() + worker: Final = _worker(shared) + + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + shared.encoded.clear() + worker.log_success_event(_call_kwargs("dep-a"), None, None, None) + + assert shared.count(f"{GROUP}_request_count:dep-a") == 0 + + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + + assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B + + +@pytest.mark.asyncio +async def test_a_local_counter_that_expired_mid_request_cannot_go_negative() -> None: + worker: Final = _worker(None) + in_memory: Final = worker.router_cache.in_memory_cache + + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + in_memory.delete_cache(f"{GROUP}_request_count:dep-a") + await worker.async_log_success_event(_call_kwargs("dep-a"), None, None, None) + + assert worker.router_cache.get_cache(f"{GROUP}_request_count:dep-a") == 0 + + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + + assert await worker.async_get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B + + def test_calls_without_a_deployment_are_ignored() -> None: shared: Final = SharedRedisCounters() worker: Final = _worker(shared) From a9bc2cb50b4ba443e72b7d03a33447cc8724b037 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 22:38:26 -0700 Subject: [PATCH 043/109] fix(least-busy): clamp the shared in-flight count to zero in one call A decrement whose matching increment is gone, because the counter key expired while the request was still in flight, used to recreate the key at -1, and a deployment with a negative count looks permanently idle, so it collects every pick from then on. The repair write that followed the decrement could also land after another pod's increment and erase it. The increment, the clamp at zero and the TTL refresh now run as a single Lua call, so nothing can interleave between them. --- litellm/caching/redis_cache.py | 39 ++++++++++++++++++- litellm/router_strategy/least_busy.py | 10 +---- .../router_strategy/test_least_busy.py | 20 +++++----- 3 files changed, 50 insertions(+), 19 deletions(-) diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index d4c47914394..4eccde5742b 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -20,6 +20,8 @@ from contextvars import ContextVar from datetime import timedelta from typing import TYPE_CHECKING, Any, Final, Protocol, TypeVar, cast +from pydantic import TypeAdapter + import litellm from litellm._logging import print_verbose, verbose_logger from litellm.constants import ( @@ -80,11 +82,22 @@ class _AsyncRedisCommands(Protocol): def pipeline(self, transaction: bool = True) -> "Pipeline[bytes]": ... + def eval(self, script: str, numkeys: int, *keys_and_args: str | bytes | float) -> Awaitable[object]: ... + _BREAKER_GUARD_FRAME_NAMES: Final = frozenset( {"", "wrapper", "_run_under_circuit_breaker", "_run_under_circuit_breaker_sync"} ) +_INCREMENT_WITH_FLOOR_LUA: Final = ( + "local count = redis.call('INCRBY', KEYS[1], ARGV[1]) " + "if count < 0 then redis.call('SET', KEYS[1], 0) count = 0 end " + "redis.call('EXPIRE', KEYS[1], ARGV[2]) " + "return count" +) + +_LUA_COUNT: Final = TypeAdapter(int) + def _get_call_stack_info(num_frames: int = 2) -> str: """ @@ -680,7 +693,7 @@ class RedisCache(BaseCache): # NON blocking - notify users Redis is throwing an exception print_verbose(f"litellm.caching.caching: set() - Got exception from REDIS : {e}") - def increment_cache(self, key, value: int, ttl: float | None = None, refresh_ttl: bool = False, **kwargs) -> int: + def increment_cache(self, key, value: int, ttl: float | None = None, **kwargs) -> int: _redis_client: Final = self.redis_client start_time = time.time() set_ttl: Final = self.get_ttl(ttl=ttl) @@ -701,7 +714,7 @@ class RedisCache(BaseCache): if set_ttl is not None: # check if key already has ttl, if not -> set ttl start_time = time.time() - current_ttl: Final = -1 if refresh_ttl else _redis_client.ttl(key) + current_ttl: Final = _redis_client.ttl(key) end_time = time.time() _duration = end_time - start_time self.service_logger_obj.service_success_hook( @@ -736,6 +749,20 @@ class RedisCache(BaseCache): ) raise e + def increment_with_floor(self, key: str, value: int, ttl: int) -> int: + """Add ``value`` to ``key``, clamp the result at zero, and refresh the TTL, in one Lua call. + + A counter whose key expired while a request was still in flight would otherwise be + recreated negative by that request's decrement. Clamping inside the same call is what + keeps it safe: a separate corrective write could land after another pod's increment and + erase it. Returns the resulting count. + """ + namespaced_key: Final = self.check_and_fix_namespace(key=key) + count: Final[object] = self.redis_client.eval( # pyright: ignore[reportAttributeAccessIssue] # stubs omit eval + _INCREMENT_WITH_FLOOR_LUA, 1, namespaced_key, value, ttl + ) + return _LUA_COUNT.validate_python(count) + @_redis_circuit_breaker_guard async def async_scan_iter(self, pattern: str, count: int = 100) -> list: start_time: Final = time.time() @@ -1241,6 +1268,14 @@ class RedisCache(BaseCache): result = result.decode() return float(result) + @_redis_circuit_breaker_guard + async def async_increment_with_floor(self, key: str, value: int, ttl: int) -> int: + """Async twin of ``increment_with_floor``, sharing its Lua script and its guarantees.""" + _redis_client: Final = self._async_commands() + namespaced_key: Final = self.check_and_fix_namespace(key=key) + count: Final = await _redis_client.eval(_INCREMENT_WITH_FLOOR_LUA, 1, namespaced_key, value, ttl) + return _LUA_COUNT.validate_python(count) + async def flush_cache_buffer(self): print_verbose(f"flushing to redis....reached size of buffer {len(self.redis_batch_writing_buffer)}") await self.async_set_cache_pipeline(self.redis_batch_writing_buffer) diff --git a/litellm/router_strategy/least_busy.py b/litellm/router_strategy/least_busy.py index 00d27b5f8d9..771d2bb4328 100644 --- a/litellm/router_strategy/least_busy.py +++ b/litellm/router_strategy/least_busy.py @@ -188,9 +188,7 @@ class LeastBusyLoggingHandler(CustomLogger): self.router_cache.set_cache(key, 0, local_only=True, ttl=IN_FLIGHT_COUNT_TTL_SECONDS) if redis_cache is None: return - shared: Final = redis_cache.increment_cache(key, delta, ttl=IN_FLIGHT_COUNT_TTL_SECONDS, refresh_ttl=True) - if shared < 0: - redis_cache.set_cache(key, 0, ttl=IN_FLIGHT_COUNT_TTL_SECONDS) + redis_cache.increment_with_floor(key, delta, IN_FLIGHT_COUNT_TTL_SECONDS) except Exception as e: _warn_unwritable(key, e) @@ -208,10 +206,6 @@ class LeastBusyLoggingHandler(CustomLogger): await self.router_cache.async_set_cache(key, 0, local_only=True, ttl=IN_FLIGHT_COUNT_TTL_SECONDS) if redis_cache is None: return - shared: Final = await redis_cache.async_increment( - key, delta, ttl=IN_FLIGHT_COUNT_TTL_SECONDS, refresh_ttl=True - ) - if shared < 0: - await redis_cache.async_set_cache(key, 0, ttl=IN_FLIGHT_COUNT_TTL_SECONDS) + await redis_cache.async_increment_with_floor(key, delta, IN_FLIGHT_COUNT_TTL_SECONDS) except Exception as e: _warn_unwritable(key, e) diff --git a/tests/test_litellm/router_strategy/test_least_busy.py b/tests/test_litellm/router_strategy/test_least_busy.py index 55702c6fe73..eb9591811ac 100644 --- a/tests/test_litellm/router_strategy/test_least_busy.py +++ b/tests/test_litellm/router_strategy/test_least_busy.py @@ -40,17 +40,16 @@ class SharedRedisCounters: async def async_set_cache(self, key: str, value: object, **kwargs: object) -> None: self.set_cache(key, value) - def increment_cache(self, key: str, value: int, ttl: float | None = None, refresh_ttl: bool = False) -> int: + def increment_with_floor(self, key: str, value: int, ttl: int) -> int: current: Final = self.count(key) or 0 assert isinstance(current, int) - incremented: Final = current + value + incremented: Final = max(0, current + value) self.encoded[key] = json.dumps(incremented) - if ttl is not None and (refresh_ttl or key not in self.ttls): - self.ttls[key] = ttl + self.ttls[key] = ttl return incremented - async def async_increment(self, key: str, value: float, ttl: int | None = None, refresh_ttl: bool = False) -> float: - return self.increment_cache(key, int(value), ttl, refresh_ttl) + async def async_increment_with_floor(self, key: str, value: int, ttl: int) -> int: + return self.increment_with_floor(key, value, ttl) def batch_get_cache(self, key_list: list[str], **kwargs: object) -> dict[str, object]: return {key: self.count(key) for key in key_list} @@ -102,12 +101,14 @@ def test_sync_pick_reads_the_shared_counts() -> None: def test_redis_counts_keep_a_refreshed_ttl() -> None: shared: Final = SharedRedisCounters() worker: Final = _worker(shared) + key: Final = f"{GROUP}_request_count:dep-a" + shared.ttls[key] = 5 worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) worker.log_success_event(_call_kwargs("dep-a"), None, None, None) - assert shared.count(f"{GROUP}_request_count:dep-a") == 0 - assert shared.ttls == {f"{GROUP}_request_count:dep-a": IN_FLIGHT_COUNT_TTL_SECONDS} + assert shared.count(key) == 0 + assert shared.ttls == {key: IN_FLIGHT_COUNT_TTL_SECONDS} @pytest.mark.asyncio @@ -129,7 +130,7 @@ class UnavailableRedis(SharedRedisCounters): def batch_get_cache(self, key_list: list[str], **kwargs: object) -> dict[str, object]: raise ConnectionError("redis is down") - def increment_cache(self, key: str, value: int, ttl: float | None = None, refresh_ttl: bool = False) -> int: + def increment_with_floor(self, key: str, value: int, ttl: int) -> int: raise ConnectionError("redis is down") @@ -159,6 +160,7 @@ def test_a_shared_counter_that_expired_mid_request_cannot_go_negative() -> None: worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + assert shared.count(f"{GROUP}_request_count:dep-a") == 1 assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B From caf9bbbd5a535e61da844e1a0225a4f8a6312b8d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 6 Sep 2026 00:20:05 -0700 Subject: [PATCH 044/109] fix(least-busy): keep the shared count readable, counted once, and off the loop A Redis outage read as "every deployment is idle", because batch_get_cache swallows the failure and answers with an empty dict. batch_get_counts and its async twin raise instead, so a worker that cannot reach Redis falls back to its own numbers rather than routing on zeros. The counter's TTL is now set only on a key that has none, so a +1 left behind by a worker that died mid-request ages out an hour after the key was created. It used to be refreshed on every touch, which kept that stuck count alive for as long as the group took traffic. Two least-busy groups counted the same request twice, since the pre-call list kept a selector per group while the success list deduped by class. The selector now goes on through add_litellm_input_callback, which dedupes the same way. A prompt-management model picked its deployment on the synchronous path, so the new Redis read landed on the event loop and configured routing plugins never ran. It awaits the async selector now. --- basedpyright-code-budget.json | 12 ++-- litellm/caching/redis_cache.py | 36 +++++++++- litellm/router.py | 8 ++- litellm/router_strategy/least_busy.py | 10 +-- ruff-strict-budget.json | 10 +-- .../test_litellm/caching/test_redis_cache.py | 44 ++++++++++++ .../router_strategy/test_least_busy.py | 67 +++++++++---------- .../test_router_routing_groups.py | 38 +++++++++++ .../test_router_routing_plugins.py | 30 +++++++++ type-discipline-budget.json | 8 +-- 10 files changed, 199 insertions(+), 64 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index c21dc76743a..32d3730aa32 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -54,10 +54,10 @@ "limit": 0 }, "reportMissingParameterType": { - "limit": 5551 + "limit": 5532 }, "reportMissingTypeArgument": { - "limit": 15277 + "limit": 15273 }, "reportMissingTypeStubs": { "limit": 40 @@ -105,13 +105,13 @@ "limit": 109 }, "reportUnknownMemberType": { - "limit": 38240 + "limit": 38208 }, "reportUnknownParameterType": { - "limit": 19558 + "limit": 19532 }, "reportUnknownVariableType": { - "limit": 29781 + "limit": 29751 }, "reportUnnecessaryCast": { "limit": 110 @@ -135,7 +135,7 @@ "limit": 21 }, "reportUnusedFunction": { - "limit": 138 + "limit": 137 }, "reportUnusedImport": { "limit": 542 diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index 4eccde5742b..3abc6e6f3c9 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -92,11 +92,18 @@ _BREAKER_GUARD_FRAME_NAMES: Final = frozenset( _INCREMENT_WITH_FLOOR_LUA: Final = ( "local count = redis.call('INCRBY', KEYS[1], ARGV[1]) " "if count < 0 then redis.call('SET', KEYS[1], 0) count = 0 end " - "redis.call('EXPIRE', KEYS[1], ARGV[2]) " + "if redis.call('TTL', KEYS[1]) < 0 then redis.call('EXPIRE', KEYS[1], ARGV[2]) end " "return count" ) _LUA_COUNT: Final = TypeAdapter(int) +_OPTIONAL_COUNTS: Final = TypeAdapter(tuple[int | None, ...]) + + +def _decoded_counts(values: Sequence[bytes | str | None]) -> tuple[int | None, ...]: + return _OPTIONAL_COUNTS.validate_python( + tuple(value.decode("utf-8") if isinstance(value, bytes) else value for value in values) + ) def _get_call_stack_info(num_frames: int = 2) -> str: @@ -749,13 +756,19 @@ class RedisCache(BaseCache): ) raise e + @_redis_circuit_breaker_guard_sync def increment_with_floor(self, key: str, value: int, ttl: int) -> int: - """Add ``value`` to ``key``, clamp the result at zero, and refresh the TTL, in one Lua call. + """Add ``value`` to ``key``, clamp the result at zero, and give a new key ``ttl``, in one Lua call. A counter whose key expired while a request was still in flight would otherwise be recreated negative by that request's decrement. Clamping inside the same call is what keeps it safe: a separate corrective write could land after another pod's increment and - erase it. Returns the resulting count. + erase it. + + The TTL is set only on a key that has none, so a counter expires ``ttl`` after it was + created rather than ``ttl`` after it was last touched. Refreshing it on every touch + would keep a count a dead worker never decremented alive for as long as the group + takes traffic. Returns the resulting count. """ namespaced_key: Final = self.check_and_fix_namespace(key=key) count: Final[object] = self.redis_client.eval( # pyright: ignore[reportAttributeAccessIssue] # stubs omit eval @@ -763,6 +776,23 @@ class RedisCache(BaseCache): ) return _LUA_COUNT.validate_python(count) + @_redis_circuit_breaker_guard_sync + def batch_get_counts(self, key_list: list[str]) -> tuple[int | None, ...]: + """Read integer counters for ``key_list``, in order, raising when Redis cannot answer. + + ``batch_get_cache`` swallows every failure and returns an empty dict, which the caller + cannot tell apart from "every counter is unset". A caller that has to fall back to its + own numbers when Redis is unreachable needs the failure, not a dict of zeros. + """ + namespaced_keys: Final = [self.check_and_fix_namespace(key=key) for key in key_list] + return _decoded_counts(self._run_redis_mget_operation(keys=namespaced_keys)) + + @_redis_circuit_breaker_guard + async def async_batch_get_counts(self, key_list: list[str]) -> tuple[int | None, ...]: + """Async twin of ``batch_get_counts``, raising on failure the same way.""" + namespaced_keys: Final = [self.check_and_fix_namespace(key=key) for key in key_list] + return _decoded_counts(await self._async_run_redis_mget_operation(keys=namespaced_keys)) + @_redis_circuit_breaker_guard async def async_scan_iter(self, pattern: str, count: int = 100) -> list: start_time: Final = time.time() diff --git a/litellm/router.py b/litellm/router.py index 3f450661946..f5b7924fb58 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -1212,7 +1212,7 @@ class Router: selector = LeastBusyLoggingHandler(router_cache=self.cache) if register_callbacks: if isinstance(litellm.input_callback, list): - litellm.input_callback.append(selector) + litellm.logging_callback_manager.add_litellm_input_callback(selector) else: litellm.input_callback = [selector] case RoutingStrategy.USAGE_BASED_ROUTING.value: @@ -4139,10 +4139,12 @@ class Router: } ) litellm_logging_object = cast(LiteLLMLogging, litellm_logging_object) - prompt_management_deployment: Final = self.get_available_deployment( + specific_deployment: Final = kwargs.pop("specific_deployment", None) + prompt_management_deployment: Final = await self.async_get_available_deployment( model=model, messages=[{"role": "user", "content": "prompt"}], - specific_deployment=kwargs.pop("specific_deployment", None), + specific_deployment=specific_deployment, + request_kwargs=kwargs, ) self._update_kwargs_with_deployment(deployment=prompt_management_deployment, kwargs=kwargs) diff --git a/litellm/router_strategy/least_busy.py b/litellm/router_strategy/least_busy.py index 771d2bb4328..ab6d702bbc2 100644 --- a/litellm/router_strategy/least_busy.py +++ b/litellm/router_strategy/least_busy.py @@ -38,7 +38,6 @@ class _Deployment(TypedDict): _CALL_KWARGS: Final = TypeAdapter(_CallKwargs) _DEPLOYMENTS: Final = TypeAdapter(list[_Deployment]) -_REDIS_COUNTS: Final = TypeAdapter(dict[str, float | None]) _MEMORY_COUNTS: Final = TypeAdapter(tuple[float | None, ...] | None) @@ -72,11 +71,6 @@ def _as_counts(values: Sequence[float | None]) -> tuple[int, ...]: return tuple(0 if value is None else int(value) for value in values) -def _shared_counts(raw: object, keys: tuple[str, ...]) -> tuple[int, ...]: - by_key: Final = _REDIS_COUNTS.validate_python(raw) - return _as_counts([by_key.get(key) for key in keys]) - - def _local_counts(raw: object, keys: tuple[str, ...]) -> tuple[int, ...]: values: Final = _MEMORY_COUNTS.validate_python(raw) if values is None or len(values) != len(keys): @@ -151,7 +145,7 @@ class LeastBusyLoggingHandler(CustomLogger): redis_cache: Final = self.router_cache.redis_cache if redis_cache is not None: try: - shared: Final = _shared_counts(redis_cache.batch_get_cache(key_list=list(keys)), keys) + shared: Final = _as_counts(redis_cache.batch_get_counts(list(keys))) except Exception as e: _warn_unreadable(model_group, e) else: @@ -166,7 +160,7 @@ class LeastBusyLoggingHandler(CustomLogger): redis_cache: Final = self.router_cache.redis_cache if redis_cache is not None: try: - shared: Final = _shared_counts(await redis_cache.async_batch_get_cache(key_list=list(keys)), keys) + shared: Final = _as_counts(await redis_cache.async_batch_get_counts(list(keys))) except Exception as e: _warn_unreadable(model_group, e) else: diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index ceff8c6e15c..d63a69de76f 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -1,6 +1,6 @@ { "ANN001": { - "limit": 2937 + "limit": 2918 }, "ANN002": { "limit": 71 @@ -9,10 +9,10 @@ "limit": 806 }, "ANN201": { - "limit": 1972 + "limit": 1965 }, "ANN202": { - "limit": 830 + "limit": 829 }, "ANN204": { "limit": 683 @@ -57,7 +57,7 @@ "limit": 3 }, "BLE001": { - "limit": 2915 + "limit": 2914 }, "C401": { "limit": 8 @@ -189,7 +189,7 @@ "limit": 0 }, "S110": { - "limit": 212 + "limit": 207 }, "S112": { "limit": 22 diff --git a/tests/test_litellm/caching/test_redis_cache.py b/tests/test_litellm/caching/test_redis_cache.py index 2f412e7382b..6b2df118611 100644 --- a/tests/test_litellm/caching/test_redis_cache.py +++ b/tests/test_litellm/caching/test_redis_cache.py @@ -525,6 +525,50 @@ def test_circuit_breaker_open_keeps_sync_batch_get_cache_as_a_miss(sync_batch_re assert sync_batch_redis_cache.batch_get_cache(key_list=["lit6729"]) == {} +def test_batch_get_counts_raises_where_batch_get_cache_reports_a_miss(sync_batch_redis_cache): + """A caller that must fall back when Redis is unreachable needs the failure, not zeros. + + The batch read answers a dead Redis with an empty dict, which a counting caller cannot tell + apart from "every counter is unset". Least-busy routing read that as an idle deployment and + kept sending traffic to it instead of falling back to this worker's own in-flight counts. + """ + assert sync_batch_redis_cache.batch_get_cache(key_list=["lit7039"]) == {} + + with pytest.raises(OSError, match="redis unavailable"): + sync_batch_redis_cache.batch_get_counts(["lit7039"]) + + +@pytest.mark.asyncio +async def test_async_batch_get_counts_raises_where_async_batch_get_cache_reports_a_miss(redis_no_ping: None): + """Async twin: the async batch read hides the same failure behind an empty dict.""" + failing_client = AsyncMock() + failing_client.mget.side_effect = OSError("redis unavailable") + with patch( # test-quality-ok: RedisCache.__init__ builds its client eagerly, with no injection point + "litellm._redis.get_redis_client", return_value=MagicMock() + ): + cache = RedisCache(host="127.0.0.1", port=6379) + + with patch.object(cache, "init_async_client", return_value=failing_client): + assert await cache.async_batch_get_cache(key_list=["lit7039"]) == {} + + with pytest.raises(OSError, match="redis unavailable"): + await cache.async_batch_get_counts(["lit7039"]) + + +@pytest.mark.parametrize("stored", [b"3", "3"]) +def test_batch_get_counts_reads_counters_in_order_and_keeps_unset_keys_apart(stored, redis_no_ping: None): + """Counters come back positionally, so an unset key has to stay a hole rather than shift the + rest of the row onto the wrong deployments, and a count has to survive whether the client + hands it back as bytes or as text.""" + with patch( # test-quality-ok: RedisCache.__init__ builds its client eagerly, with no injection point + "litellm._redis.get_redis_client", return_value=MagicMock() + ): + cache = RedisCache(host="127.0.0.1", port=6379) + cache.redis_client.mget.return_value = [stored, None, b"0"] + + assert cache.batch_get_counts(["dep-a", "dep-b", "dep-c"]) == (3, None, 0) + + @pytest.fixture def sync_batch_cache_with_service_logger(redis_no_ping: None) -> Iterator[tuple[RedisCache, ServiceLogging]]: service_logger = ServiceLogging(mock_testing=True) diff --git a/tests/test_litellm/router_strategy/test_least_busy.py b/tests/test_litellm/router_strategy/test_least_busy.py index eb9591811ac..9efa526fc02 100644 --- a/tests/test_litellm/router_strategy/test_least_busy.py +++ b/tests/test_litellm/router_strategy/test_least_busy.py @@ -1,4 +1,3 @@ -import json from typing import Final import pytest @@ -18,44 +17,34 @@ def _call_kwargs(deployment_id: str) -> dict[str, object]: class SharedRedisCounters: - """Stores JSON strings and hands back a fresh object per read, the way a real Redis client does.""" + """Mirrors what Redis gives the handler: increments clamped at zero, a TTL set once when + the key is created, and ordered reads that raise rather than invent a value.""" def __init__(self) -> None: - self.encoded: dict[str, str] = {} - self.ttls: dict[str, float] = {} + self.counts: dict[str, int] = {} + self.ttls: dict[str, int] = {} - def count(self, key: str) -> object: - raw: Final = self.encoded.get(key) - return None if raw is None else json.loads(raw) + def count(self, key: str) -> int | None: + return self.counts.get(key) - def get_cache(self, key: str, **kwargs: object) -> object: - return self.count(key) - - def set_cache(self, key: str, value: object, **kwargs: object) -> None: - self.encoded[key] = json.dumps(value) - - async def async_get_cache(self, key: str, **kwargs: object) -> object: - return self.count(key) - - async def async_set_cache(self, key: str, value: object, **kwargs: object) -> None: - self.set_cache(key, value) + def expire(self, key: str) -> None: + self.counts.pop(key, None) + self.ttls.pop(key, None) def increment_with_floor(self, key: str, value: int, ttl: int) -> int: - current: Final = self.count(key) or 0 - assert isinstance(current, int) - incremented: Final = max(0, current + value) - self.encoded[key] = json.dumps(incremented) - self.ttls[key] = ttl + incremented: Final = max(0, self.counts.get(key, 0) + value) + self.counts[key] = incremented + self.ttls.setdefault(key, ttl) return incremented async def async_increment_with_floor(self, key: str, value: int, ttl: int) -> int: return self.increment_with_floor(key, value, ttl) - def batch_get_cache(self, key_list: list[str], **kwargs: object) -> dict[str, object]: - return {key: self.count(key) for key in key_list} + def batch_get_counts(self, key_list: list[str]) -> tuple[int | None, ...]: + return tuple(self.counts.get(key) for key in key_list) - async def async_batch_get_cache(self, key_list: list[str], **kwargs: object) -> dict[str, object]: - return self.batch_get_cache(key_list) + async def async_batch_get_counts(self, key_list: list[str]) -> tuple[int | None, ...]: + return self.batch_get_counts(key_list) def _worker(shared: SharedRedisCounters | None) -> LeastBusyLoggingHandler: @@ -98,17 +87,25 @@ def test_sync_pick_reads_the_shared_counts() -> None: assert picking_worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A -def test_redis_counts_keep_a_refreshed_ttl() -> None: +def test_the_handler_never_pushes_a_counters_ttl_forward() -> None: + """A worker that dies mid-request leaves a +1 nobody will ever decrement. Redis expires that + stuck count an hour after the key was created, which only works while nothing writes the TTL + again: a handler that refreshed it on every touch would keep the count alive for as long as + the group takes traffic, and the deployment would read busier than it is forever.""" shared: Final = SharedRedisCounters() worker: Final = _worker(shared) key: Final = f"{GROUP}_request_count:dep-a" - shared.ttls[key] = 5 + worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) + + assert shared.ttls == {key: IN_FLIGHT_COUNT_TTL_SECONDS} + + shared.ttls[key] = 5 worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) worker.log_success_event(_call_kwargs("dep-a"), None, None, None) - assert shared.count(key) == 0 - assert shared.ttls == {key: IN_FLIGHT_COUNT_TTL_SECONDS} + assert shared.count(key) == 1 + assert shared.ttls == {key: 5} @pytest.mark.asyncio @@ -127,10 +124,10 @@ async def test_counts_stay_in_memory_without_redis() -> None: class UnavailableRedis(SharedRedisCounters): - def batch_get_cache(self, key_list: list[str], **kwargs: object) -> dict[str, object]: + def increment_with_floor(self, key: str, value: int, ttl: int) -> int: raise ConnectionError("redis is down") - def increment_with_floor(self, key: str, value: int, ttl: int) -> int: + def batch_get_counts(self, key_list: list[str]) -> tuple[int | None, ...]: raise ConnectionError("redis is down") @@ -153,7 +150,7 @@ def test_a_shared_counter_that_expired_mid_request_cannot_go_negative() -> None: worker: Final = _worker(shared) worker.log_pre_api_call(model="m", messages=[], kwargs=_call_kwargs("dep-a")) - shared.encoded.clear() + shared.expire(f"{GROUP}_request_count:dep-a") worker.log_success_event(_call_kwargs("dep-a"), None, None, None) assert shared.count(f"{GROUP}_request_count:dep-a") == 0 @@ -187,4 +184,4 @@ def test_calls_without_a_deployment_are_ignored() -> None: worker.log_pre_api_call(model="m", messages=[], kwargs={"litellm_params": {"metadata": None}}) worker.log_pre_api_call(model="m", messages=[], kwargs={}) - assert shared.encoded == {} + assert shared.counts == {} diff --git a/tests/test_litellm/router_strategy/test_router_routing_groups.py b/tests/test_litellm/router_strategy/test_router_routing_groups.py index 5599c5aad63..0ce0ed5b37e 100644 --- a/tests/test_litellm/router_strategy/test_router_routing_groups.py +++ b/tests/test_litellm/router_strategy/test_router_routing_groups.py @@ -12,6 +12,7 @@ import pytest import litellm from litellm import Router +from litellm.integrations.custom_logger import CustomLogger from litellm.types.router import RoutingGroup, RoutingStrategy @@ -435,6 +436,43 @@ def test_update_settings_unregisters_group_selectors_when_groups_removed(monkeyp assert router._group_selectors == {} +def test_two_least_busy_groups_count_a_request_once(monkeypatch): + """ + Least-busy counts a request up from the pre-call hooks on `litellm.input_callback` and + back down from the success hooks on `litellm.callbacks`. The success list drops a second + selector of the same class, so a pre-call list that kept both counted every request twice + and released it once, and the deployment's in-flight count climbed until it looked pinned. + """ + monkeypatch.setattr(litellm, "callbacks", []) + monkeypatch.setattr(litellm, "input_callback", []) + + router = _build_router( + routing_strategy="least-busy", + routing_groups=[ + { + "group_name": "fast", + "models": ["filtered-model"], + "routing_strategy": "least-busy", + } + ], + ) + kwargs = { + "litellm_params": { + "metadata": {"model_group": "filtered-model"}, + "model_info": {"id": "deploy-1"}, + } + } + + for callback in litellm.input_callback: + if isinstance(callback, CustomLogger): + callback.log_pre_api_call(model="filtered-model", messages=[], kwargs=kwargs) + for callback in litellm.callbacks: + if isinstance(callback, CustomLogger): + callback.log_success_event(kwargs, None, None, None) + + assert router.cache.get_cache("filtered-model_request_count:deploy-1") == 0 + + # --------------------------------------------------------------------------- # Direct helper coverage # --------------------------------------------------------------------------- diff --git a/tests/test_litellm/router_strategy/test_router_routing_plugins.py b/tests/test_litellm/router_strategy/test_router_routing_plugins.py index 293af36080a..c49b22ea367 100644 --- a/tests/test_litellm/router_strategy/test_router_routing_plugins.py +++ b/tests/test_litellm/router_strategy/test_router_routing_plugins.py @@ -164,6 +164,36 @@ async def test_async_completion_with_unsupported_strategy_rejects_configured_plu await router.acompletion(model="smart-router", messages=[{"role": "user", "content": "hi"}]) +@pytest.mark.asyncio +async def test_prompt_management_model_still_runs_the_plugin_pipeline(): + """ + A prompt-management model routes through its own factory, which picked the deployment + on the synchronous path. Plugins never run there, so the guard turned every such request + into an error message about the caller's own API choice, on an async call the caller made + correctly. It also read the in-flight counts with a blocking call inside the event loop. + """ + router = Router( + model_list=[ + { + "model_name": "cached-claude", + "litellm_params": { + "model": "anthropic_cache_control_hook/claude-sonnet-5", + "prompt_id": "cache-points", + }, + } + ], + routing_strategy="least-busy", + plugins=[BlockEverything()], + ) + + with pytest.raises(ValueError, match="No deployments left after routing-plugin filtering"): + await router.acompletion( + model="cached-claude", + messages=[{"role": "user", "content": "hi"}], + litellm_call_id="lit-7039", + ) + + @pytest.mark.asyncio async def test_router_without_plugins_is_unaffected(): """Regression guard: a Router with no `plugins` configured behaves exactly as before.""" diff --git a/type-discipline-budget.json b/type-discipline-budget.json index e413e6db2c6..0c0952289e2 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 22176 + "limit": 22174 }, "LIT002": { - "limit": 26721 + "limit": 26715 }, "LIT003": { "limit": 261 @@ -27,10 +27,10 @@ "limit": 0 }, "LIT010": { - "limit": 16412 + "limit": 16398 }, "LIT011": { - "limit": 5505 + "limit": 5504 }, "LIT012": { "limit": 4486 From ef5f51abca77d3081d04796ce9b3191dc2409f9b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 6 Sep 2026 01:59:15 -0700 Subject: [PATCH 045/109] fix(least-busy): count for every router, and keep the expiry through a clamp Two routers in one process shared a single handler, because the callback manager dedupes on the class name plus the handler's public attributes and the handler had none. The second router's requests were never counted. The handler now carries the id of the cache it was built on, so routers with different caches both register while the two selectors one router builds for its routing groups still collapse into one. Clamping a negative count back to zero used SET, which drops the key's TTL, so the next write started the hour over. It uses INCRBY by the negative amount now, which leaves the expiry alone. The Lua script had no test that ran it, so tests/local_testing covers both the sync and async paths against a real Redis, and the file is wired into the CircleCI job that provides one. --- .circleci/config.yml | 1 + litellm/caching/redis_cache.py | 2 +- litellm/router_strategy/least_busy.py | 1 + .../test_redis_increment_with_floor.py | 80 +++++++++++++++++++ .../test_router_routing_groups.py | 38 +++++++++ 5 files changed, 121 insertions(+), 1 deletion(-) create mode 100644 tests/local_testing/test_redis_increment_with_floor.py diff --git a/.circleci/config.yml b/.circleci/config.yml index 6e368a3debe..d10962864bf 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -1440,6 +1440,7 @@ jobs: TEST_FILES=$(printf "%s\n" \ tests/local_testing/test_dual_cache.py \ tests/local_testing/test_redis_batch_optimizations.py \ + tests/local_testing/test_redis_increment_with_floor.py \ tests/local_testing/test_router_utils.py) echo "$TEST_FILES" | circleci tests run \ --verbose \ diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index 3abc6e6f3c9..106c1580110 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -91,7 +91,7 @@ _BREAKER_GUARD_FRAME_NAMES: Final = frozenset( _INCREMENT_WITH_FLOOR_LUA: Final = ( "local count = redis.call('INCRBY', KEYS[1], ARGV[1]) " - "if count < 0 then redis.call('SET', KEYS[1], 0) count = 0 end " + "if count < 0 then count = redis.call('INCRBY', KEYS[1], -count) end " "if redis.call('TTL', KEYS[1]) < 0 then redis.call('EXPIRE', KEYS[1], ARGV[2]) end " "return count" ) diff --git a/litellm/router_strategy/least_busy.py b/litellm/router_strategy/least_busy.py index ab6d702bbc2..14e6592e1fd 100644 --- a/litellm/router_strategy/least_busy.py +++ b/litellm/router_strategy/least_busy.py @@ -106,6 +106,7 @@ class LeastBusyLoggingHandler(CustomLogger): def __init__(self, router_cache: DualCache): self.router_cache = router_cache + self.router_cache_id = str(id(router_cache)) def log_pre_api_call(self, model: str, messages: object, kwargs: Mapping[str, object]) -> None: self._increment(kwargs, 1) diff --git a/tests/local_testing/test_redis_increment_with_floor.py b/tests/local_testing/test_redis_increment_with_floor.py new file mode 100644 index 00000000000..e358d5f31e0 --- /dev/null +++ b/tests/local_testing/test_redis_increment_with_floor.py @@ -0,0 +1,80 @@ +"""Least-busy routing keeps its in-flight counters in Redis, and the clamp at zero plus the +create-once TTL both live inside a Lua script. Nothing but a real Redis runs that script, so +these are the only tests that fail when the script itself is wrong.""" + +import os +import uuid +from typing import Final + +import pytest +from dotenv import load_dotenv + +load_dotenv() + +from litellm.caching.redis_cache import RedisCache + +TTL: Final = 600 + + +@pytest.fixture +def counter(): + cache: Final = RedisCache(host=os.getenv("REDIS_HOST"), port=os.getenv("REDIS_PORT")) + key: Final = f"lit7039-{uuid.uuid4()}" + yield cache, key, cache.check_and_fix_namespace(key=key) + cache.delete_cache(key) + + +def test_a_counter_adds_every_increment_and_reads_back_what_it_holds(counter): + cache, key, _ = counter + + assert cache.increment_with_floor(key, 3, TTL) == 3 + assert cache.increment_with_floor(key, 2, TTL) == 5 + assert cache.batch_get_counts([key]) == (5,) + + +def test_a_decrement_past_zero_leaves_the_counter_at_zero(counter): + """A worker whose counter expired mid-request decrements a key that is no longer there. + Without the clamp that deployment reads negative, and least-busy pins every later request + on it until the count climbs back to zero.""" + cache, key, _ = counter + + assert cache.increment_with_floor(key, 1, TTL) == 1 + assert cache.increment_with_floor(key, -5, TTL) == 0 + assert cache.batch_get_counts([key]) == (0,) + + +def test_traffic_never_pushes_a_counters_expiry_back_out(counter): + """The TTL is what releases a count whose worker died mid-request. Rewriting it on every + touch would keep that stuck count alive for as long as the group takes traffic.""" + cache, key, namespaced_key = counter + + cache.increment_with_floor(key, 1, TTL) + assert cache.redis_client.ttl(namespaced_key) > TTL - 60 + + cache.redis_client.expire(namespaced_key, 30) + cache.increment_with_floor(key, 1, TTL) + + assert cache.redis_client.ttl(namespaced_key) <= 30 + + +def test_clamping_to_zero_keeps_the_expiry_it_already_had(counter): + cache, key, namespaced_key = counter + + cache.increment_with_floor(key, 1, TTL) + cache.redis_client.expire(namespaced_key, 30) + + assert cache.increment_with_floor(key, -5, TTL) == 0 + assert cache.redis_client.ttl(namespaced_key) <= 30 + + +@pytest.mark.asyncio +async def test_the_async_counter_behaves_the_same_way(counter): + cache, key, namespaced_key = counter + + assert await cache.async_increment_with_floor(key, 2, TTL) == 2 + assert await cache.async_batch_get_counts([key]) == (2,) + + cache.redis_client.expire(namespaced_key, 30) + + assert await cache.async_increment_with_floor(key, -9, TTL) == 0 + assert cache.redis_client.ttl(namespaced_key) <= 30 diff --git a/tests/test_litellm/router_strategy/test_router_routing_groups.py b/tests/test_litellm/router_strategy/test_router_routing_groups.py index 0ce0ed5b37e..af390c3292b 100644 --- a/tests/test_litellm/router_strategy/test_router_routing_groups.py +++ b/tests/test_litellm/router_strategy/test_router_routing_groups.py @@ -473,6 +473,44 @@ def test_two_least_busy_groups_count_a_request_once(monkeypatch): assert router.cache.get_cache("filtered-model_request_count:deploy-1") == 0 +def test_two_routers_in_one_process_each_count_their_own_requests(monkeypatch): + """ + Least-busy hangs its counting off litellm's global callback lists, and those lists keep one + logger per class unless the instances differ in a plain attribute. Two routers in one process + (a second Router, or a per-request `user_config` one) therefore have to register separately: + a second router whose selector is dropped counts nothing, reads zero for every deployment, + and sends every request to whichever one is listed first. + """ + monkeypatch.setattr(litellm, "callbacks", []) + monkeypatch.setattr(litellm, "input_callback", []) + + first = _build_router(routing_strategy="least-busy") + second = _build_router(routing_strategy="least-busy") + kwargs = { + "litellm_params": { + "metadata": {"model_group": "filtered-model"}, + "model_info": {"id": "deploy-1"}, + } + } + + for callback in litellm.input_callback: + if isinstance(callback, CustomLogger): + callback.log_pre_api_call(model="filtered-model", messages=[], kwargs=kwargs) + + assert second.cache.get_cache("filtered-model_request_count:deploy-1") == 1 + assert ( + second.get_available_deployment(model="filtered-model", messages=[])["model_info"]["id"] + == "deploy-2" + ) + + for callback in litellm.callbacks: + if isinstance(callback, CustomLogger): + callback.log_success_event(kwargs, None, None, None) + + assert first.cache.get_cache("filtered-model_request_count:deploy-1") == 0 + assert second.cache.get_cache("filtered-model_request_count:deploy-1") == 0 + + # --------------------------------------------------------------------------- # Direct helper coverage # --------------------------------------------------------------------------- From 314e8905c583792d08207cdcd1bfa8e22d985f51 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 6 Sep 2026 02:44:45 -0700 Subject: [PATCH 046/109] fix(router): let prompt-management plugins see the caller's own messages The prompt-management factory picks its deployment with a placeholder message. That was inert while the pick ran on the synchronous path, which never runs the routing plugin pipeline. Now that the pick runs the pipeline, a plugin classifying request content would score the placeholder instead of the conversation, and the narrowing it writes decides which deployments the real call may use. --- litellm/router.py | 2 +- .../test_router_routing_plugins.py | 47 +++++++++++++++++++ 2 files changed, 48 insertions(+), 1 deletion(-) diff --git a/litellm/router.py b/litellm/router.py index f5b7924fb58..398d87ce843 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -4142,7 +4142,7 @@ class Router: specific_deployment: Final = kwargs.pop("specific_deployment", None) prompt_management_deployment: Final = await self.async_get_available_deployment( model=model, - messages=[{"role": "user", "content": "prompt"}], + messages=cast(list[dict[str, str]], messages), # cast-ok: selection reads messages structurally specific_deployment=specific_deployment, request_kwargs=kwargs, ) diff --git a/tests/test_litellm/router_strategy/test_router_routing_plugins.py b/tests/test_litellm/router_strategy/test_router_routing_plugins.py index c49b22ea367..0a54addce5e 100644 --- a/tests/test_litellm/router_strategy/test_router_routing_plugins.py +++ b/tests/test_litellm/router_strategy/test_router_routing_plugins.py @@ -56,6 +56,18 @@ class BlockEverything: return context +class MessageRecorder: + """Records what each plugin pass was handed, then blocks so the request stops there.""" + + def __init__(self): + self.seen = [] + + async def run(self, context: RoutingContext) -> RoutingContext: + self.seen.append(list(context.raw_messages)) + context.candidate_models = [] + return context + + def _smart_router_model_list(): return [ { @@ -194,6 +206,41 @@ async def test_prompt_management_model_still_runs_the_plugin_pipeline(): ) +@pytest.mark.asyncio +async def test_prompt_management_plugins_see_the_callers_own_messages(): + """ + The prompt-management factory picks its deployment with a placeholder message, which was + harmless while that pick ran on the synchronous path (plugins never ran there at all). Now + that the pick runs the plugin pipeline, a plugin that classifies request content would score + the placeholder instead of the conversation, and the narrowing it produces decides which + deployments the real call is allowed to use. + """ + recorder = MessageRecorder() + router = Router( + model_list=[ + { + "model_name": "cached-claude", + "litellm_params": { + "model": "anthropic_cache_control_hook/claude-sonnet-5", + "prompt_id": "cache-points", + }, + } + ], + routing_strategy="least-busy", + plugins=[recorder], + ) + messages = [{"role": "user", "content": "wire me $40,000 to account 12345"}] + + with pytest.raises(ValueError, match="No deployments left after routing-plugin filtering"): + await router.acompletion( + model="cached-claude", + messages=messages, + litellm_call_id="lit-7039", + ) + + assert recorder.seen == [messages] + + @pytest.mark.asyncio async def test_router_without_plugins_is_unaffected(): """Regression guard: a Router with no `plugins` configured behaves exactly as before.""" 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 047/109] 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 048/109] 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 dc09d9e7cfbd62f590f73fdb1844bc2aac5578bf Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 7 Sep 2026 17:24:17 -0700 Subject: [PATCH 049/109] feat(bedrock): add TwelveLabs Marengo Embed 3.0 embeddings --- litellm/constants.py | 1 + litellm/llms/bedrock/embed/embedding.py | 2 +- .../twelvelabs_marengo_3_transformation.py | 204 +++++++++++++ .../twelvelabs_marengo_transformation.py | 51 +++- ...odel_prices_and_context_window_backup.json | 39 +++ litellm/types/llms/bedrock.py | 113 +++++++- litellm/utils.py | 2 +- model_prices_and_context_window.json | 39 +++ .../test_bedrock_async_invoke_embedding.py | 38 +++ .../bedrock/embed/test_bedrock_embedding.py | 129 +++++++++ ...est_twelvelabs_marengo_3_transformation.py | 268 ++++++++++++++++++ ..._bedrock_marengo_embed_3_model_metadata.py | 88 ++++++ 12 files changed, 961 insertions(+), 13 deletions(-) create mode 100644 litellm/llms/bedrock/embed/twelvelabs_marengo_3_transformation.py create mode 100644 tests/test_litellm/llms/bedrock/embed/test_twelvelabs_marengo_3_transformation.py create mode 100644 tests/test_litellm/test_bedrock_marengo_embed_3_model_metadata.py diff --git a/litellm/constants.py b/litellm/constants.py index d53686e5e5b..78cca3c6212 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1370,6 +1370,7 @@ bedrock_embedding_models: Final[set] = set( "cohere.embed-multilingual-v3", "cohere.embed-v4:0", "twelvelabs.marengo-embed-2-7-v1:0", + "twelvelabs.marengo-embed-3-0-v1:0", ] ) diff --git a/litellm/llms/bedrock/embed/embedding.py b/litellm/llms/bedrock/embed/embedding.py index 5fb86d476f4..ab27afcf817 100644 --- a/litellm/llms/bedrock/embed/embedding.py +++ b/litellm/llms/bedrock/embed/embedding.py @@ -474,7 +474,7 @@ class BedrockEmbedding(BaseAWSLLM): elif provider == "twelvelabs": batch_data = [] for i in input: - twelvelabs_request = TwelveLabsMarengoEmbeddingConfig()._transform_request( + twelvelabs_request = TwelveLabsMarengoEmbeddingConfig(model=model)._transform_request( input=i, inference_params=inference_params, async_invoke_route=has_async_invoke, diff --git a/litellm/llms/bedrock/embed/twelvelabs_marengo_3_transformation.py b/litellm/llms/bedrock/embed/twelvelabs_marengo_3_transformation.py new file mode 100644 index 00000000000..2ea99db47f0 --- /dev/null +++ b/litellm/llms/bedrock/embed/twelvelabs_marengo_3_transformation.py @@ -0,0 +1,204 @@ +""" +Request builder for Bedrock TwelveLabs Marengo Embed 3.0, whose payload nests the input under a key named after +``inputType`` instead of the flat 2.7 layout. + +Docs - https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-marengo-3.html +""" + +from collections.abc import Mapping +from types import MappingProxyType +from typing import Final, assert_never + +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError + +from litellm.llms.bedrock.common_utils import BedrockError +from litellm.types.llms.bedrock import ( + TWELVELABS_MARENGO_3_EMBEDDING_OPTIONS, + TWELVELABS_MARENGO_3_EMBEDDING_SCOPES, + TWELVELABS_MARENGO_3_EMBEDDING_TYPES, + TWELVELABS_MARENGO_3_INPUT_TYPES, + TwelveLabsMarengo3AudioRequest, + TwelveLabsMarengo3EmbeddingRequest, + TwelveLabsMarengo3ImageRequest, + TwelveLabsMarengo3MultiInputRequest, + TwelveLabsMarengo3NamedMediaSource, + TwelveLabsMarengo3RequestBase, + TwelveLabsMarengo3Segmentation, + TwelveLabsMarengo3TextImageRequest, + TwelveLabsMarengo3TextRequest, + TwelveLabsMarengo3TimedMediaInput, + TwelveLabsMarengo3TimedMediaOptions, + TwelveLabsMarengo3VideoRequest, + TwelveLabsMediaSource, + TwelveLabsS3Location, +) +from litellm.utils import get_base64_str + +MARENGO_3_MODEL_MARKER: Final = "marengo-embed-3" +S3_URI_PREFIX: Final = "s3://" +TIMED_MEDIA_OPTION_FIELDS: Final = MappingProxyType( + { + "startSec": True, + "endSec": True, + "segmentation": True, + "embeddingOption": True, + "embeddingType": True, + "embeddingScope": True, + } +) +TIMED_MEDIA_OPTIONS: Final = TypeAdapter(TwelveLabsMarengo3TimedMediaOptions) + + +def is_marengo_3_model(model: str | None) -> bool: + return MARENGO_3_MODEL_MARKER in (model or "") + + +class Marengo3Params(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + inputType: TWELVELABS_MARENGO_3_INPUT_TYPES | None = None + input_type: TWELVELABS_MARENGO_3_INPUT_TYPES | None = None + media_source: str | None = None + media_sources: Mapping[str, str] | None = None + bucketOwner: str | None = None + startSec: float | None = None + endSec: float | None = None + segmentation: TwelveLabsMarengo3Segmentation | None = None + embeddingOption: tuple[TWELVELABS_MARENGO_3_EMBEDDING_OPTIONS, ...] | None = None + embeddingType: tuple[TWELVELABS_MARENGO_3_EMBEDDING_TYPES, ...] | None = None + embeddingScope: tuple[TWELVELABS_MARENGO_3_EMBEDDING_SCOPES, ...] | None = None + inferenceId: str | None = None + + @property + def resolved_input_type(self) -> TWELVELABS_MARENGO_3_INPUT_TYPES: + return self.inputType or self.input_type or "text" + + def timed_media_options(self) -> TwelveLabsMarengo3TimedMediaOptions: + return TIMED_MEDIA_OPTIONS.validate_python( + self.model_dump(include=TIMED_MEDIA_OPTION_FIELDS, exclude_none=True) + ) + + +def _s3_location(uri: str, bucket_owner: str | None) -> TwelveLabsS3Location: + if bucket_owner is None: + unowned: Final[TwelveLabsS3Location] = {"uri": uri} + return unowned + owned: Final[TwelveLabsS3Location] = {"uri": uri, "bucketOwner": bucket_owner} + return owned + + +def _media_source(media: str, bucket_owner: str | None) -> TwelveLabsMediaSource: + if not media.startswith(S3_URI_PREFIX): + inline: Final[TwelveLabsMediaSource] = {"base64String": get_base64_str(media)} + return inline + remote: Final[TwelveLabsMediaSource] = {"s3Location": _s3_location(media, bucket_owner)} + return remote + + +def _named_media_source(name: str, media: str, bucket_owner: str | None) -> TwelveLabsMarengo3NamedMediaSource: + named: Final[TwelveLabsMarengo3NamedMediaSource] = { + "name": name, + "mediaType": "image", + **_media_source(media, bucket_owner), + } + return named + + +def _timed_media_input(media: str, params: Marengo3Params) -> TwelveLabsMarengo3TimedMediaInput: + timed: Final[TwelveLabsMarengo3TimedMediaInput] = { + "mediaSource": _media_source(media, params.bucketOwner), + **params.timed_media_options(), + } + return timed + + +def _validated_params(inference_params: Mapping[str, object]) -> Marengo3Params: + try: + return Marengo3Params.model_validate(inference_params) + except ValidationError as error: + raise BedrockError(status_code=400, message=f"Invalid Marengo 3.0 parameters: {error}") from error + + +def _require(value: str | None, input_type: str, param_name: str) -> str: + if value is None: + raise BedrockError(status_code=400, message=f"Input type '{input_type}' requires the '{param_name}' parameter") + return value + + +def _require_media_sources(value: Mapping[str, str] | None) -> Mapping[str, str]: + if not value: + raise BedrockError( + status_code=400, + message="Input type 'multi_input' requires a non-empty 'media_sources' mapping of name to media", + ) + return value + + +def _request_base(inference_id: str | None) -> TwelveLabsMarengo3RequestBase: + if inference_id is None: + anonymous: Final[TwelveLabsMarengo3RequestBase] = {} + return anonymous + identified: Final[TwelveLabsMarengo3RequestBase] = {"inferenceId": inference_id} + return identified + + +def build_marengo_3_request(input: str, inference_params: Mapping[str, object]) -> TwelveLabsMarengo3EmbeddingRequest: + params: Final = _validated_params(inference_params) + base: Final = _request_base(params.inferenceId) + input_type: Final = params.resolved_input_type + match input_type: + case "text": + text_request: Final[TwelveLabsMarengo3TextRequest] = { + **base, + "inputType": "text", + "text": {"inputText": input}, + } + return text_request + case "image": + image_request: Final[TwelveLabsMarengo3ImageRequest] = { + **base, + "inputType": "image", + "image": {"mediaSource": _media_source(input, params.bucketOwner)}, + } + return image_request + case "video": + video_request: Final[TwelveLabsMarengo3VideoRequest] = { + **base, + "inputType": "video", + "video": _timed_media_input(input, params), + } + return video_request + case "audio": + audio_request: Final[TwelveLabsMarengo3AudioRequest] = { + **base, + "inputType": "audio", + "audio": _timed_media_input(input, params), + } + return audio_request + case "text_image": + text_image_request: Final[TwelveLabsMarengo3TextImageRequest] = { + **base, + "inputType": "text_image", + "text_image": { + "inputText": input, + "mediaSource": _media_source( + _require(params.media_source, input_type, "media_source"), params.bucketOwner + ), + }, + } + return text_image_request + case "multi_input": + media_sources: Final = tuple( + _named_media_source(name, media, params.bucketOwner) + for name, media in _require_media_sources(params.media_sources).items() + ) + multi_input_request: Final[TwelveLabsMarengo3MultiInputRequest] = { + **base, + "inputType": "multi_input", + "multi_input": {"inputText": input, "mediaSources": media_sources} + if input + else {"mediaSources": media_sources}, + } + return multi_input_request + case _: + assert_never(input_type) diff --git a/litellm/llms/bedrock/embed/twelvelabs_marengo_transformation.py b/litellm/llms/bedrock/embed/twelvelabs_marengo_transformation.py index a39c59b0efd..79b5825d2eb 100644 --- a/litellm/llms/bedrock/embed/twelvelabs_marengo_transformation.py +++ b/litellm/llms/bedrock/embed/twelvelabs_marengo_transformation.py @@ -4,13 +4,19 @@ Transformation logic from OpenAI /v1/embeddings format to Bedrock TwelveLabs Mar Why separate file? Make it easy to see how transformation works Docs - https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-marengo.html +Marengo 3.0 docs - https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-marengo-3.html """ from typing import Final, cast +from litellm.llms.bedrock.embed.twelvelabs_marengo_3_transformation import ( + build_marengo_3_request, + is_marengo_3_model, +) from litellm.types.llms.bedrock import ( TWELVELABS_EMBEDDING_INPUT_TYPES, TwelveLabsAsyncInvokeRequest, + TwelveLabsMarengo3EmbeddingRequest, TwelveLabsMarengoEmbeddingRequest, TwelveLabsOutputDataConfig, TwelveLabsS3Location, @@ -26,10 +32,13 @@ class TwelveLabsMarengoEmbeddingConfig: Supports text, image, video, and audio inputs. - InvokeModel: text and image inputs - StartAsyncInvoke: video, audio, image, and text inputs + + Marengo 3.0 (model ids containing "marengo-embed-3") nests the input under a key named after inputType and + adds the text_image and multi_input input types; that payload is built by build_marengo_3_request. """ - def __init__(self) -> None: - pass + def __init__(self, model: str | None = None) -> None: + self.is_marengo_3: Final = is_marengo_3_model(model) def get_supported_openai_params(self) -> list[str]: return [ @@ -41,13 +50,20 @@ class TwelveLabsMarengoEmbeddingConfig: "useFixedLengthSec", "minClipSec", "input_type", + "endSec", + "segmentation", + "embeddingType", + "embeddingScope", + "inferenceId", + "media_source", + "media_sources", ] def map_openai_params(self, non_default_params: dict, optional_params: dict) -> dict: for k, v in non_default_params.items(): if k == "encoding_format": # TwelveLabs doesn't have encoding_format, but we can map it to embeddingOption - if v == "float": + if v == "float" and not self.is_marengo_3: optional_params["embeddingOption"] = ["visual-text", "visual-image"] elif k == "textTruncate": optional_params["textTruncate"] = v @@ -56,7 +72,19 @@ class TwelveLabsMarengoEmbeddingConfig: elif k == "input_type": # Map input_type to inputType for Bedrock optional_params["inputType"] = v - elif k in ["startSec", "lengthSec", "useFixedLengthSec", "minClipSec"]: + elif k in ( + "startSec", + "lengthSec", + "useFixedLengthSec", + "minClipSec", + "endSec", + "segmentation", + "embeddingType", + "embeddingScope", + "inferenceId", + "media_source", + "media_sources", + ): optional_params[k] = v return optional_params @@ -77,7 +105,7 @@ class TwelveLabsMarengoEmbeddingConfig: async_invoke_route: bool = False, model_id: str | None = None, output_s3_uri: str | None = None, - ) -> TwelveLabsMarengoEmbeddingRequest | TwelveLabsAsyncInvokeRequest: + ) -> TwelveLabsMarengoEmbeddingRequest | TwelveLabsMarengo3EmbeddingRequest | TwelveLabsAsyncInvokeRequest: """ Transform OpenAI-style input to TwelveLabs Marengo format/async-invoke format. @@ -87,20 +115,27 @@ class TwelveLabsMarengoEmbeddingConfig: - Video inputs (async-invoke only) - Audio inputs (async-invoke only) - S3 URLs for all media types (async-invoke only) + - Marengo 3.0 only: text_image and multi_input inputs (nested payload) """ - # Get input_type or default to "text" input_type: Final = cast( TWELVELABS_EMBEDDING_INPUT_TYPES, inference_params.get("inputType") or inference_params.get("input_type") or "text", ) - # Validate that async-invoke is used for video/audio if input_type in ["video", "audio"] and not async_invoke_route: raise ValueError( f"Input type '{input_type}' requires async_invoke route. " f"Use model format: 'bedrock/async_invoke/model_id'" ) + if self.is_marengo_3: + marengo_3_request: Final = build_marengo_3_request(input=input, inference_params=inference_params) + if async_invoke_route and model_id: + return self._wrap_async_invoke_request( + model_input=marengo_3_request, model_id=model_id, output_s3_uri=output_s3_uri + ) + return marengo_3_request + transformed_request: Final[TwelveLabsMarengoEmbeddingRequest] = {"inputType": input_type} if input_type == "text": @@ -154,7 +189,7 @@ class TwelveLabsMarengoEmbeddingConfig: def _wrap_async_invoke_request( self, - model_input: TwelveLabsMarengoEmbeddingRequest, + model_input: TwelveLabsMarengoEmbeddingRequest | TwelveLabsMarengo3EmbeddingRequest, model_id: str, output_s3_uri: str | None = None, ) -> TwelveLabsAsyncInvokeRequest: diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index b1ffc1583e4..cc75354a495 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -690,6 +690,45 @@ "supports_embedding_image_input": true, "supports_image_input": true }, + "twelvelabs.marengo-embed-3-0-v1:0": { + "input_cost_per_token": 7e-05, + "litellm_provider": "bedrock", + "max_input_tokens": 500, + "max_tokens": 500, + "mode": "embedding", + "output_cost_per_token": 0.0, + "output_vector_size": 512, + "supports_embedding_image_input": true, + "supports_image_input": true + }, + "us.twelvelabs.marengo-embed-3-0-v1:0": { + "input_cost_per_token": 7e-05, + "input_cost_per_video_per_second": 0.0007, + "input_cost_per_audio_per_second": 0.00014, + "input_cost_per_image": 0.0001, + "litellm_provider": "bedrock", + "max_input_tokens": 500, + "max_tokens": 500, + "mode": "embedding", + "output_cost_per_token": 0.0, + "output_vector_size": 512, + "supports_embedding_image_input": true, + "supports_image_input": true + }, + "eu.twelvelabs.marengo-embed-3-0-v1:0": { + "input_cost_per_token": 7e-05, + "input_cost_per_video_per_second": 0.0007, + "input_cost_per_audio_per_second": 0.00014, + "input_cost_per_image": 0.0001, + "litellm_provider": "bedrock", + "max_input_tokens": 500, + "max_tokens": 500, + "mode": "embedding", + "output_cost_per_token": 0.0, + "output_vector_size": 512, + "supports_embedding_image_input": true, + "supports_image_input": true + }, "twelvelabs.pegasus-1-2-v1:0": { "input_cost_per_video_per_second": 0.00049, "output_cost_per_token": 7.5e-06, diff --git a/litellm/types/llms/bedrock.py b/litellm/types/llms/bedrock.py index bed0ba3dc08..9f93886a9c6 100644 --- a/litellm/types/llms/bedrock.py +++ b/litellm/types/llms/bedrock.py @@ -1,7 +1,7 @@ import json from collections.abc import Sequence from enum import Enum -from typing import TYPE_CHECKING, Any, Final, Literal +from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias from typing_extensions import ReadOnly, Required, TypedDict, override @@ -557,7 +557,7 @@ class AmazonTitanMultimodalEmbeddingResponse(TypedDict): message: str # Specifies any errors that occur during generation. -# TwelveLabs Marengo Embed 2.7 types +# TwelveLabs Marengo Embed types TWELVELABS_EMBEDDING_INPUT_TYPES = Literal["text", "image", "video", "audio"] TWELVELABS_EMBEDDING_OPTIONS = Literal["visual-text", "visual-image", "audio"] @@ -591,6 +591,113 @@ class TwelveLabsMarengoEmbeddingResponse(TypedDict): endSec: float +TWELVELABS_MARENGO_3_INPUT_TYPES: TypeAlias = Literal["text", "image", "video", "audio", "text_image", "multi_input"] +TWELVELABS_MARENGO_3_EMBEDDING_OPTIONS: TypeAlias = Literal["visual", "audio", "transcription"] +TWELVELABS_MARENGO_3_EMBEDDING_TYPES: TypeAlias = Literal["separate_embedding", "fused_embedding"] +TWELVELABS_MARENGO_3_EMBEDDING_SCOPES: TypeAlias = Literal["clip", "asset"] + + +class TwelveLabsMarengo3FixedSegmentationConfig(TypedDict): + durationSec: ReadOnly[int] + + +class TwelveLabsMarengo3FixedSegmentation(TypedDict): + method: ReadOnly[Literal["fixed"]] + fixed: ReadOnly[TwelveLabsMarengo3FixedSegmentationConfig] + + +class TwelveLabsMarengo3DynamicSegmentationConfig(TypedDict): + minDurationSec: ReadOnly[int] + + +class TwelveLabsMarengo3DynamicSegmentation(TypedDict): + method: ReadOnly[Literal["dynamic"]] + dynamic: ReadOnly[TwelveLabsMarengo3DynamicSegmentationConfig] + + +TwelveLabsMarengo3Segmentation: TypeAlias = TwelveLabsMarengo3FixedSegmentation | TwelveLabsMarengo3DynamicSegmentation + + +class TwelveLabsMarengo3TextInput(TypedDict): + inputText: ReadOnly[str] + + +class TwelveLabsMarengo3ImageInput(TypedDict): + mediaSource: ReadOnly[TwelveLabsMediaSource] + + +class TwelveLabsMarengo3TimedMediaOptions(TypedDict, total=False): + startSec: ReadOnly[float] + endSec: ReadOnly[float] + segmentation: ReadOnly[TwelveLabsMarengo3Segmentation] + embeddingOption: ReadOnly[Sequence[TWELVELABS_MARENGO_3_EMBEDDING_OPTIONS]] + embeddingType: ReadOnly[Sequence[TWELVELABS_MARENGO_3_EMBEDDING_TYPES]] + embeddingScope: ReadOnly[Sequence[TWELVELABS_MARENGO_3_EMBEDDING_SCOPES]] + + +class TwelveLabsMarengo3TimedMediaInput(TwelveLabsMarengo3TimedMediaOptions): + mediaSource: Required[ReadOnly[TwelveLabsMediaSource]] + + +class TwelveLabsMarengo3TextImageInput(TypedDict): + inputText: ReadOnly[str] + mediaSource: ReadOnly[TwelveLabsMediaSource] + + +class TwelveLabsMarengo3NamedMediaSource(TwelveLabsMediaSource): + name: Required[ReadOnly[str]] + mediaType: Required[ReadOnly[Literal["image"]]] + + +class TwelveLabsMarengo3MultiInput(TypedDict, total=False): + inputText: ReadOnly[str] + mediaSources: Required[ReadOnly[Sequence[TwelveLabsMarengo3NamedMediaSource]]] + + +class TwelveLabsMarengo3RequestBase(TypedDict, total=False): + inferenceId: ReadOnly[str] + + +class TwelveLabsMarengo3TextRequest(TwelveLabsMarengo3RequestBase): + inputType: ReadOnly[Literal["text"]] + text: ReadOnly[TwelveLabsMarengo3TextInput] + + +class TwelveLabsMarengo3ImageRequest(TwelveLabsMarengo3RequestBase): + inputType: ReadOnly[Literal["image"]] + image: ReadOnly[TwelveLabsMarengo3ImageInput] + + +class TwelveLabsMarengo3VideoRequest(TwelveLabsMarengo3RequestBase): + inputType: ReadOnly[Literal["video"]] + video: ReadOnly[TwelveLabsMarengo3TimedMediaInput] + + +class TwelveLabsMarengo3AudioRequest(TwelveLabsMarengo3RequestBase): + inputType: ReadOnly[Literal["audio"]] + audio: ReadOnly[TwelveLabsMarengo3TimedMediaInput] + + +class TwelveLabsMarengo3TextImageRequest(TwelveLabsMarengo3RequestBase): + inputType: ReadOnly[Literal["text_image"]] + text_image: ReadOnly[TwelveLabsMarengo3TextImageInput] + + +class TwelveLabsMarengo3MultiInputRequest(TwelveLabsMarengo3RequestBase): + inputType: ReadOnly[Literal["multi_input"]] + multi_input: ReadOnly[TwelveLabsMarengo3MultiInput] + + +TwelveLabsMarengo3EmbeddingRequest: TypeAlias = ( + TwelveLabsMarengo3TextRequest + | TwelveLabsMarengo3ImageRequest + | TwelveLabsMarengo3VideoRequest + | TwelveLabsMarengo3AudioRequest + | TwelveLabsMarengo3TextImageRequest + | TwelveLabsMarengo3MultiInputRequest +) + + class TwelveLabsS3OutputDataConfig(TypedDict): s3Uri: str @@ -601,7 +708,7 @@ class TwelveLabsOutputDataConfig(TypedDict): class TwelveLabsAsyncInvokeRequest(TypedDict): modelId: str - modelInput: TwelveLabsMarengoEmbeddingRequest + modelInput: ReadOnly[TwelveLabsMarengoEmbeddingRequest | TwelveLabsMarengo3EmbeddingRequest] outputDataConfig: TwelveLabsOutputDataConfig diff --git a/litellm/utils.py b/litellm/utils.py index d0e11bc9551..b98aa821ff3 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -3623,7 +3623,7 @@ def get_optional_params_embeddings( elif "cohere.embed" in model: object = litellm.BedrockCohereEmbeddingConfig() elif "twelvelabs" in model or "marengo" in model: - object = litellm.TwelveLabsMarengoEmbeddingConfig() + object = litellm.TwelveLabsMarengoEmbeddingConfig(model=model) elif "nova" in model.lower(): object = litellm.AmazonNovaEmbeddingConfig() else: # unmapped model diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index b1ffc1583e4..cc75354a495 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -690,6 +690,45 @@ "supports_embedding_image_input": true, "supports_image_input": true }, + "twelvelabs.marengo-embed-3-0-v1:0": { + "input_cost_per_token": 7e-05, + "litellm_provider": "bedrock", + "max_input_tokens": 500, + "max_tokens": 500, + "mode": "embedding", + "output_cost_per_token": 0.0, + "output_vector_size": 512, + "supports_embedding_image_input": true, + "supports_image_input": true + }, + "us.twelvelabs.marengo-embed-3-0-v1:0": { + "input_cost_per_token": 7e-05, + "input_cost_per_video_per_second": 0.0007, + "input_cost_per_audio_per_second": 0.00014, + "input_cost_per_image": 0.0001, + "litellm_provider": "bedrock", + "max_input_tokens": 500, + "max_tokens": 500, + "mode": "embedding", + "output_cost_per_token": 0.0, + "output_vector_size": 512, + "supports_embedding_image_input": true, + "supports_image_input": true + }, + "eu.twelvelabs.marengo-embed-3-0-v1:0": { + "input_cost_per_token": 7e-05, + "input_cost_per_video_per_second": 0.0007, + "input_cost_per_audio_per_second": 0.00014, + "input_cost_per_image": 0.0001, + "litellm_provider": "bedrock", + "max_input_tokens": 500, + "max_tokens": 500, + "mode": "embedding", + "output_cost_per_token": 0.0, + "output_vector_size": 512, + "supports_embedding_image_input": true, + "supports_image_input": true + }, "twelvelabs.pegasus-1-2-v1:0": { "input_cost_per_video_per_second": 0.00049, "output_cost_per_token": 7.5e-06, diff --git a/tests/test_litellm/llms/bedrock/embed/test_bedrock_async_invoke_embedding.py b/tests/test_litellm/llms/bedrock/embed/test_bedrock_async_invoke_embedding.py index 74a55cc1ef2..00f5145269a 100644 --- a/tests/test_litellm/llms/bedrock/embed/test_bedrock_async_invoke_embedding.py +++ b/tests/test_litellm/llms/bedrock/embed/test_bedrock_async_invoke_embedding.py @@ -184,6 +184,44 @@ class TestBedrockAsyncInvokeEmbedding: request_url = mock_post.call_args.kwargs.get("url", "") assert "/async-invoke" in request_url + def test_async_invoke_marengo_3_wraps_the_nested_payload_with_the_base_model_id(self): + client = HTTPHandler() + + with patch.object(client, "post") as mock_post: + mock_response = Mock() + mock_response.status_code = 200 + mock_response.text = json.dumps(async_invoke_response) + mock_response.json = lambda: json.loads(mock_response.text) + mock_post.return_value = mock_response + + response = litellm.embedding( + model="bedrock/async_invoke/twelvelabs.marengo-embed-3-0-v1:0", + input="s3://test-bucket/clip.mp4", + client=client, + aws_region_name="us-east-1", + aws_bedrock_runtime_endpoint="https://bedrock-runtime.us-east-1.amazonaws.com", + api_key="test-bearer-token-12345", + input_type="video", + embeddingOption=["visual", "audio"], + segmentation={"method": "fixed", "fixed": {"durationSec": 6}}, + output_s3_uri="s3://test-bucket/async-invoke-output/", + ) + + assert response._hidden_params._invocation_arn == async_invoke_response["invocationArn"] + assert mock_post.call_args.kwargs["url"].endswith("/async-invoke") + assert json.loads(mock_post.call_args.kwargs["data"]) == { + "modelId": "twelvelabs.marengo-embed-3-0-v1:0", + "modelInput": { + "inputType": "video", + "video": { + "mediaSource": {"s3Location": {"uri": "s3://test-bucket/clip.mp4"}}, + "segmentation": {"method": "fixed", "fixed": {"durationSec": 6}}, + "embeddingOption": ["visual", "audio"], + }, + }, + "outputDataConfig": {"s3OutputDataConfig": {"s3Uri": "s3://test-bucket/async-invoke-output/"}}, + } + @pytest.mark.asyncio async def test_async_invoke_twelvelabs_embedding_async_with_mock(self): """Test async invoke embedding with async calls.""" diff --git a/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py b/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py index 50f8bbcf584..b37e991b0b2 100644 --- a/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py +++ b/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py @@ -1059,3 +1059,132 @@ def test_bedrock_embedding_bearer_token_never_runs_the_sigv4_credential_chain(mo assert response.data[0]["embedding"] == titan_embedding_response["embedding"] assert mock_post.call_args.kwargs["headers"]["Authorization"] == "Bearer env-bearer-token-12345" + + +marengo_3_embedding_response = {"data": [{"embedding": [0.01 * i for i in range(512)]}]} +MARENGO_3_DUCK = "data:image/png;base64,ZHVjaw==" + + +@pytest.mark.parametrize( + "model,kwargs,expected_body", + [ + ( + "bedrock/us.twelvelabs.marengo-embed-3-0-v1:0", + {"input_type": "text"}, + {"inputType": "text", "text": {"inputText": "a duck on water"}}, + ), + ( + "bedrock/twelvelabs.marengo-embed-3-0-v1:0", + {"input_type": "text"}, + {"inputType": "text", "text": {"inputText": "a duck on water"}}, + ), + ( + "bedrock/us.twelvelabs.marengo-embed-3-0-v1:0", + {"input_type": "text_image", "media_source": MARENGO_3_DUCK}, + { + "inputType": "text_image", + "text_image": {"inputText": "a duck on water", "mediaSource": {"base64String": "ZHVjaw=="}}, + }, + ), + ( + "bedrock/us.twelvelabs.marengo-embed-3-0-v1:0", + {"input_type": "multi_input", "media_sources": {"bird": MARENGO_3_DUCK}}, + { + "inputType": "multi_input", + "multi_input": { + "inputText": "a duck on water", + "mediaSources": [{"name": "bird", "mediaType": "image", "base64String": "ZHVjaw=="}], + }, + }, + ), + ], +) +def test_marengo_3_embedding_sends_the_nested_payload_and_parses_512_dims(model, kwargs, expected_body): + client = HTTPHandler() + + with patch.object(client, "post") as mock_post: + mock_response = Mock() + mock_response.status_code = 200 + mock_response.text = json.dumps(marengo_3_embedding_response) + mock_response.json = lambda: json.loads(mock_response.text) + mock_post.return_value = mock_response + + response = litellm.embedding( + model=model, + input="a duck on water", + client=client, + aws_region_name="us-east-1", + api_key="test-bearer-token-12345", + **kwargs, + ) + + assert json.loads(mock_post.call_args.kwargs["data"]) == expected_body + assert mock_post.call_args.kwargs["url"].endswith(f"/model/{model.removeprefix('bedrock/').replace(':', '%3A')}/invoke") + assert len(response.data[0]["embedding"]) == 512 + assert response.data[0]["embedding"][:2] == [0.0, 0.01] + assert response.usage.prompt_tokens == 128 + + +def test_marengo_3_image_embedding_sends_the_media_under_the_image_key(): + client = HTTPHandler() + + with patch.object(client, "post") as mock_post: + mock_response = Mock() + mock_response.status_code = 200 + mock_response.text = json.dumps(marengo_3_embedding_response) + mock_response.json = lambda: json.loads(mock_response.text) + mock_post.return_value = mock_response + + response = litellm.embedding( + model="bedrock/us.twelvelabs.marengo-embed-3-0-v1:0", + input=MARENGO_3_DUCK, + client=client, + aws_region_name="us-east-1", + api_key="test-bearer-token-12345", + input_type="image", + ) + + assert json.loads(mock_post.call_args.kwargs["data"]) == { + "inputType": "image", + "image": {"mediaSource": {"base64String": "ZHVjaw=="}}, + } + assert len(response.data[0]["embedding"]) == 512 + assert response.data[0]["embedding"][:2] == [0.0, 0.01] + + +def test_marengo_2_7_embedding_keeps_the_flat_payload(): + client = HTTPHandler() + + with patch.object(client, "post") as mock_post: + mock_response = Mock() + mock_response.status_code = 200 + mock_response.text = json.dumps(twelvelabs_embedding_response) + mock_response.json = lambda: json.loads(mock_response.text) + mock_post.return_value = mock_response + + response = litellm.embedding( + model="bedrock/us.twelvelabs.marengo-embed-2-7-v1:0", + input="a duck on water", + client=client, + aws_region_name="us-east-1", + api_key="test-bearer-token-12345", + input_type="text", + ) + + assert json.loads(mock_post.call_args.kwargs["data"]) == { + "inputType": "text", + "inputText": "a duck on water", + "textTruncate": "end", + } + assert response.data[0]["embedding"] == [0.1, 0.2, 0.3] + + +def test_marengo_3_text_image_without_media_source_is_a_bad_request(): + with pytest.raises(litellm.BadRequestError, match=r"text_image.*media_source"): + litellm.embedding( + model="bedrock/us.twelvelabs.marengo-embed-3-0-v1:0", + input="a duck on water", + aws_region_name="us-east-1", + api_key="test-bearer-token-12345", + input_type="text_image", + ) diff --git a/tests/test_litellm/llms/bedrock/embed/test_twelvelabs_marengo_3_transformation.py b/tests/test_litellm/llms/bedrock/embed/test_twelvelabs_marengo_3_transformation.py new file mode 100644 index 00000000000..0bf86352a5e --- /dev/null +++ b/tests/test_litellm/llms/bedrock/embed/test_twelvelabs_marengo_3_transformation.py @@ -0,0 +1,268 @@ +import json + +import pytest + +from litellm.llms.bedrock.common_utils import BedrockError +from litellm.llms.bedrock.embed.twelvelabs_marengo_3_transformation import ( + build_marengo_3_request, + is_marengo_3_model, +) +from litellm.llms.bedrock.embed.twelvelabs_marengo_transformation import ( + TwelveLabsMarengoEmbeddingConfig, +) + +MARENGO_3_BASE = "twelvelabs.marengo-embed-3-0-v1:0" +MARENGO_3_US = "us.twelvelabs.marengo-embed-3-0-v1:0" +MARENGO_27_US = "us.twelvelabs.marengo-embed-2-7-v1:0" +DUCK_DATA_URL = "data:image/png;base64,ZHVjaw==" +OUTPUT_S3_URI = "s3://out-bucket/marengo/" + + +@pytest.mark.parametrize( + "model,expected", + [ + (MARENGO_3_BASE, True), + (MARENGO_3_US, True), + ("eu.twelvelabs.marengo-embed-3-0-v1:0", True), + ("async_invoke/twelvelabs.marengo-embed-3-0-v1:0", True), + (MARENGO_27_US, False), + ("twelvelabs.marengo-embed-2-7-v1:0", False), + (None, False), + ], +) +def test_is_marengo_3_model(model, expected): + assert is_marengo_3_model(model) is expected + + +def wire(request: object) -> object: + return json.loads(json.dumps(request)) + + +def test_text_request_nests_input_text_under_text(): + assert build_marengo_3_request("a dog on the beach", {"input_type": "text"}) == { + "inputType": "text", + "text": {"inputText": "a dog on the beach"}, + } + + +def test_missing_input_type_defaults_to_text(): + assert build_marengo_3_request("hello", {})["inputType"] == "text" + + +def test_camel_case_input_type_wins_over_snake_case(): + request = build_marengo_3_request(DUCK_DATA_URL, {"inputType": "image", "input_type": "text"}) + assert request["inputType"] == "image" + + +def test_image_request_strips_data_url_prefix(): + assert build_marengo_3_request(DUCK_DATA_URL, {"input_type": "image"}) == { + "inputType": "image", + "image": {"mediaSource": {"base64String": "ZHVjaw=="}}, + } + + +def test_image_request_from_s3_carries_bucket_owner(): + request = build_marengo_3_request("s3://media/duck.png", {"input_type": "image", "bucketOwner": "123456789012"}) + assert request == { + "inputType": "image", + "image": {"mediaSource": {"s3Location": {"uri": "s3://media/duck.png", "bucketOwner": "123456789012"}}}, + } + + +def test_s3_media_without_bucket_owner_omits_the_key(): + request = build_marengo_3_request("s3://media/duck.png", {"input_type": "image"}) + assert request["image"]["mediaSource"] == {"s3Location": {"uri": "s3://media/duck.png"}} + + +def test_text_image_request_pairs_text_with_media_source(): + request = build_marengo_3_request( + "a duck", {"input_type": "text_image", "media_source": DUCK_DATA_URL, "output_s3_uri": OUTPUT_S3_URI} + ) + assert request == { + "inputType": "text_image", + "text_image": {"inputText": "a duck", "mediaSource": {"base64String": "ZHVjaw=="}}, + } + + +def test_text_image_request_requires_media_source(): + with pytest.raises(BedrockError, match=r"text_image.*media_source") as excinfo: + build_marengo_3_request("a duck", {"input_type": "text_image"}) + assert excinfo.value.status_code == 400 + + +def test_multi_input_request_names_each_media_source(): + request = build_marengo_3_request( + "a photo of <@bird> next to <@dog>", + { + "input_type": "multi_input", + "media_sources": {"bird": DUCK_DATA_URL, "dog": "s3://media/dog.png"}, + "bucketOwner": "123456789012", + }, + ) + assert wire(request) == { + "inputType": "multi_input", + "multi_input": { + "inputText": "a photo of <@bird> next to <@dog>", + "mediaSources": [ + {"name": "bird", "mediaType": "image", "base64String": "ZHVjaw=="}, + { + "name": "dog", + "mediaType": "image", + "s3Location": {"uri": "s3://media/dog.png", "bucketOwner": "123456789012"}, + }, + ], + }, + } + + +def test_multi_input_without_text_omits_input_text(): + request = build_marengo_3_request("", {"input_type": "multi_input", "media_sources": {"bird": DUCK_DATA_URL}}) + assert "inputText" not in request["multi_input"] + assert request["multi_input"]["mediaSources"][0]["name"] == "bird" + + +@pytest.mark.parametrize("params", [{"input_type": "multi_input"}, {"input_type": "multi_input", "media_sources": {}}]) +def test_multi_input_request_requires_media_sources(params): + with pytest.raises(BedrockError, match=r"multi_input.*media_sources") as excinfo: + build_marengo_3_request("<@bird>", params) + assert excinfo.value.status_code == 400 + + +@pytest.mark.parametrize("input_type", ["video", "audio"]) +def test_timed_media_request_nests_every_option_under_the_media_key(input_type): + request = build_marengo_3_request( + "s3://media/clip.mp4", + { + "input_type": input_type, + "startSec": 2, + "endSec": 12.5, + "segmentation": {"method": "dynamic", "dynamic": {"minDurationSec": 4}}, + "embeddingOption": ["visual", "audio"], + "embeddingType": ["fused_embedding"], + "embeddingScope": ["clip", "asset"], + "inferenceId": "req-42", + }, + ) + assert wire(request) == { + "inputType": input_type, + input_type: { + "mediaSource": {"s3Location": {"uri": "s3://media/clip.mp4"}}, + "startSec": 2.0, + "endSec": 12.5, + "segmentation": {"method": "dynamic", "dynamic": {"minDurationSec": 4}}, + "embeddingOption": ["visual", "audio"], + "embeddingType": ["fused_embedding"], + "embeddingScope": ["clip", "asset"], + }, + "inferenceId": "req-42", + } + + +def test_timed_media_request_without_options_carries_only_the_media_source(): + request = build_marengo_3_request("s3://media/clip.mp4", {"input_type": "video"}) + assert request["video"] == {"mediaSource": {"s3Location": {"uri": "s3://media/clip.mp4"}}} + + +@pytest.mark.parametrize( + "params", + [ + {"input_type": "clip"}, + {"input_type": "video", "embeddingOption": ["visual-text"]}, + {"input_type": "video", "segmentation": {"method": "fixed", "dynamic": {"minDurationSec": 4}}}, + {"input_type": "multi_input", "media_sources": ["not", "a", "mapping"]}, + ], +) +def test_invalid_marengo_3_params_are_rejected_before_the_request_is_sent(params): + with pytest.raises(BedrockError, match=r"Invalid Marengo 3\.0 parameters") as excinfo: + build_marengo_3_request("s3://media/clip.mp4", params) + assert excinfo.value.status_code == 400 + + +def test_config_sends_the_nested_payload_for_marengo_3_and_the_flat_one_for_2_7(): + nested = TwelveLabsMarengoEmbeddingConfig(model=MARENGO_3_US)._transform_request( + input="hello", inference_params={"input_type": "text"} + ) + flat = TwelveLabsMarengoEmbeddingConfig(model=MARENGO_27_US)._transform_request( + input="hello", inference_params={"input_type": "text"} + ) + assert nested == {"inputType": "text", "text": {"inputText": "hello"}} + assert flat == {"inputType": "text", "inputText": "hello", "textTruncate": "end"} + + +def test_config_without_a_model_keeps_the_2_7_payload(): + request = TwelveLabsMarengoEmbeddingConfig()._transform_request(input="hello", inference_params={}) + assert request == {"inputType": "text", "inputText": "hello", "textTruncate": "end"} + + +@pytest.mark.parametrize("input_type", ["video", "audio"]) +def test_marengo_3_video_and_audio_still_require_the_async_route(input_type): + with pytest.raises(ValueError, match=f"Input type '{input_type}' requires async_invoke route"): + TwelveLabsMarengoEmbeddingConfig(model=MARENGO_3_BASE)._transform_request( + input="s3://media/clip.mp4", inference_params={"input_type": input_type} + ) + + +def test_marengo_3_async_invoke_wraps_the_nested_payload_with_the_base_model_id(): + request = TwelveLabsMarengoEmbeddingConfig(model=MARENGO_3_BASE)._transform_request( + input="s3://media/clip.mp4", + inference_params={"input_type": "video", "embeddingOption": ["visual"], "output_s3_uri": OUTPUT_S3_URI}, + async_invoke_route=True, + model_id="async_invoke%2Ftwelvelabs.marengo-embed-3-0-v1%3A0", + output_s3_uri=OUTPUT_S3_URI, + ) + assert wire(request) == { + "modelId": MARENGO_3_BASE, + "modelInput": { + "inputType": "video", + "video": {"mediaSource": {"s3Location": {"uri": "s3://media/clip.mp4"}}, "embeddingOption": ["visual"]}, + }, + "outputDataConfig": {"s3OutputDataConfig": {"s3Uri": OUTPUT_S3_URI}}, + } + + +def test_marengo_3_async_invoke_requires_an_output_s3_uri(): + with pytest.raises(ValueError, match="output_s3_uri cannot be empty"): + TwelveLabsMarengoEmbeddingConfig(model=MARENGO_3_BASE)._transform_request( + input="hello", + inference_params={"input_type": "text"}, + async_invoke_route=True, + model_id=MARENGO_3_BASE, + output_s3_uri="", + ) + + +def test_encoding_format_float_no_longer_injects_2_7_embedding_options_for_marengo_3(): + marengo_3 = TwelveLabsMarengoEmbeddingConfig(model=MARENGO_3_US).map_openai_params( + non_default_params={"encoding_format": "float"}, optional_params={} + ) + marengo_27 = TwelveLabsMarengoEmbeddingConfig(model=MARENGO_27_US).map_openai_params( + non_default_params={"encoding_format": "float"}, optional_params={} + ) + assert marengo_3 == {} + assert marengo_27 == {"embeddingOption": ["visual-text", "visual-image"]} + + +def test_marengo_3_only_params_are_forwarded_by_map_openai_params(): + mapped = TwelveLabsMarengoEmbeddingConfig(model=MARENGO_3_US).map_openai_params( + non_default_params={ + "input_type": "text_image", + "media_source": DUCK_DATA_URL, + "media_sources": {"bird": DUCK_DATA_URL}, + "endSec": 5, + "segmentation": {"method": "fixed", "fixed": {"durationSec": 6}}, + "embeddingType": ["separate_embedding"], + "embeddingScope": ["clip"], + "inferenceId": "req-1", + }, + optional_params={}, + ) + assert mapped == { + "inputType": "text_image", + "media_source": DUCK_DATA_URL, + "media_sources": {"bird": DUCK_DATA_URL}, + "endSec": 5, + "segmentation": {"method": "fixed", "fixed": {"durationSec": 6}}, + "embeddingType": ["separate_embedding"], + "embeddingScope": ["clip"], + "inferenceId": "req-1", + } diff --git a/tests/test_litellm/test_bedrock_marengo_embed_3_model_metadata.py b/tests/test_litellm/test_bedrock_marengo_embed_3_model_metadata.py new file mode 100644 index 00000000000..300dfeb5238 --- /dev/null +++ b/tests/test_litellm/test_bedrock_marengo_embed_3_model_metadata.py @@ -0,0 +1,88 @@ +import json +from pathlib import Path + +import pytest + +import litellm +from litellm.constants import bedrock_embedding_models +from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider +from litellm.types.utils import Usage + +REPO_ROOT = Path(__file__).parents[2] +MAIN_PATH = REPO_ROOT / "model_prices_and_context_window.json" +BACKUP_PATH = REPO_ROOT / "litellm" / "model_prices_and_context_window_backup.json" + +BASE_MODEL = "twelvelabs.marengo-embed-3-0-v1:0" +PROFILE_MODELS = ("us.twelvelabs.marengo-embed-3-0-v1:0", "eu.twelvelabs.marengo-embed-3-0-v1:0") +ALL_MODELS = (BASE_MODEL, *PROFILE_MODELS) + +TEXT_REQUEST_COST = 7e-05 +IMAGE_REQUEST_COST = 0.0001 +VIDEO_COST_PER_SECOND = 0.0007 +AUDIO_COST_PER_SECOND = 0.00014 + + +def _load(path): + with open(path) as f: + return json.load(f) + + +@pytest.mark.parametrize("model", ALL_MODELS) +def test_marengo_embed_3_specs(model): + info = _load(MAIN_PATH).get(model) + assert info is not None, f"{model} missing from model_prices_and_context_window.json" + + assert info["litellm_provider"] == "bedrock" + assert info["mode"] == "embedding" + assert info["input_cost_per_token"] == TEXT_REQUEST_COST + assert info["output_cost_per_token"] == 0.0 + assert info["max_input_tokens"] == 500 + assert info["max_tokens"] == 500 + assert info["output_vector_size"] == 512 + assert info["supports_embedding_image_input"] is True + assert info["supports_image_input"] is True + assert "deprecation_date" not in info + + routed_model, provider, _, _ = get_llm_provider(model=f"bedrock/{model}") + assert routed_model == model + assert provider == "bedrock" + + +@pytest.mark.parametrize("model", PROFILE_MODELS) +def test_marengo_embed_3_inference_profiles_price_image_video_and_audio(model): + info = _load(MAIN_PATH)[model] + assert info["input_cost_per_image"] == IMAGE_REQUEST_COST + assert info["input_cost_per_video_per_second"] == VIDEO_COST_PER_SECOND + assert info["input_cost_per_audio_per_second"] == AUDIO_COST_PER_SECOND + + +@pytest.mark.parametrize("model", ALL_MODELS) +def test_marengo_embed_3_is_visible_to_callers(model, local_model_cost_map): + info = litellm.get_model_info(model=model, custom_llm_provider="bedrock") + assert info["mode"] == "embedding" + assert info["output_vector_size"] == 512 + assert info["max_input_tokens"] == 500 + + +@pytest.mark.parametrize("model", ALL_MODELS) +def test_marengo_embed_3_text_request_is_billed(model, local_model_cost_map): + usage = Usage(prompt_tokens=128, completion_tokens=0, total_tokens=128) + prompt_cost, completion_cost = litellm.cost_per_token( + model=model, usage_object=usage, custom_llm_provider="bedrock" + ) + assert prompt_cost == pytest.approx(128 * TEXT_REQUEST_COST) + assert completion_cost == 0.0 + + +def test_marengo_embed_3_is_a_known_bedrock_embedding_model(): + assert BASE_MODEL in bedrock_embedding_models + + +@pytest.mark.parametrize("model", ALL_MODELS) +def test_backup_matches_main(model): + main_cost = _load(MAIN_PATH) + backup_cost = _load(BACKUP_PATH) + + assert model in main_cost, f"{model} missing from model_prices_and_context_window.json" + assert model in backup_cost, f"{model} missing from model_prices_and_context_window_backup.json" + assert backup_cost[model] == main_cost[model], f"{model} differs between main and backup model cost maps" From e2560390770bd4387c82889eff2d606ba88c42c5 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 7 Sep 2026 17:49:03 -0700 Subject: [PATCH 050/109] fix(bedrock): import assert_never from typing_extensions for Python 3.10 --- .../llms/bedrock/embed/twelvelabs_marengo_3_transformation.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/litellm/llms/bedrock/embed/twelvelabs_marengo_3_transformation.py b/litellm/llms/bedrock/embed/twelvelabs_marengo_3_transformation.py index 2ea99db47f0..0f61d37258f 100644 --- a/litellm/llms/bedrock/embed/twelvelabs_marengo_3_transformation.py +++ b/litellm/llms/bedrock/embed/twelvelabs_marengo_3_transformation.py @@ -7,9 +7,10 @@ Docs - https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-mar from collections.abc import Mapping from types import MappingProxyType -from typing import Final, assert_never +from typing import Final from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError +from typing_extensions import assert_never from litellm.llms.bedrock.common_utils import BedrockError from litellm.types.llms.bedrock import ( 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 051/109] 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 052/109] 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 86790a7723892c338ea3fcc680296a721c7fc47f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 7 Sep 2026 18:20:28 -0700 Subject: [PATCH 053/109] fix(bedrock): bill Marengo embeddings per request instead of per estimated token AWS prices Marengo 2.7 and 3.0 text and image embeddings per request, never per token, and their responses carry no token count. The old transform estimated prompt tokens from the vector length, which billed a text request at 128 tokens times the per-token rate (0.00896 instead of 0.00007). Marengo responses now report zero tokens with query_count and image_count derived from the request batch, and all six Marengo cost-map entries price per request (with the video and audio per-second and per-image rates on the base entries). query_count is a new prompt_tokens_details field wired to input_cost_per_query in the cost calculator. --- .../litellm_core_utils/llm_cost_calc/utils.py | 10 ++ litellm/llms/bedrock/embed/embedding.py | 2 +- .../twelvelabs_marengo_transformation.py | 142 +++++++++++------- ...odel_prices_and_context_window_backup.json | 18 ++- litellm/types/utils.py | 7 +- litellm/utils.py | 2 +- model_prices_and_context_window.json | 18 ++- .../llm_cost_calc/test_llm_cost_calc_utils.py | 32 ++++ .../bedrock/embed/test_bedrock_embedding.py | 47 +++++- ..._bedrock_marengo_embed_3_model_metadata.py | 45 +++++- 10 files changed, 240 insertions(+), 83 deletions(-) diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 68dc27ec25e..46574ebae3f 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -780,6 +780,7 @@ class PromptTokensDetailsResult(TypedDict): image_count: int video_length_seconds: float audio_length_seconds: float + query_count: int def parse_prompt_tokens_details(usage: Usage) -> PromptTokensDetailsResult: @@ -828,6 +829,7 @@ def parse_prompt_tokens_details(usage: Usage) -> PromptTokensDetailsResult: ) or 0.0 ) + query_count: Final = _coerce_token_count(getattr(usage.prompt_tokens_details, "query_count", 0)) return PromptTokensDetailsResult( cache_hit_tokens=cache_hit_tokens, @@ -841,6 +843,7 @@ def parse_prompt_tokens_details(usage: Usage) -> PromptTokensDetailsResult: image_count=image_count, video_length_seconds=float(video_length_seconds), audio_length_seconds=float(audio_length_seconds), + query_count=query_count, ) @@ -978,6 +981,12 @@ def _calculate_input_cost( prompt_tokens_details["audio_length_seconds"], ) + ### QUERY COUNT COST + if prompt_tokens_details["query_count"]: + prompt_cost += calculate_cost_component( + model_info, "input_cost_per_query", prompt_tokens_details["query_count"] + ) + return prompt_cost @@ -1149,6 +1158,7 @@ def generic_cost_per_token( image_count=0, video_length_seconds=0.0, audio_length_seconds=0.0, + query_count=0, ) if usage.prompt_tokens_details: prompt_tokens_details = parse_prompt_tokens_details(usage) diff --git a/litellm/llms/bedrock/embed/embedding.py b/litellm/llms/bedrock/embed/embedding.py index ab27afcf817..69eb9b693b1 100644 --- a/litellm/llms/bedrock/embed/embedding.py +++ b/litellm/llms/bedrock/embed/embedding.py @@ -229,7 +229,7 @@ class BedrockEmbedding(BaseAWSLLM): returned_response = AmazonTitanG1Config()._transform_response(response_list=response_list, model=model) elif provider == "twelvelabs": returned_response = TwelveLabsMarengoEmbeddingConfig()._transform_response( - response_list=response_list, model=model + response_list=response_list, model=model, batch_data=batch_data ) elif provider == "nova": returned_response = AmazonNovaEmbeddingConfig()._transform_response( diff --git a/litellm/llms/bedrock/embed/twelvelabs_marengo_transformation.py b/litellm/llms/bedrock/embed/twelvelabs_marengo_transformation.py index 79b5825d2eb..ddf6dfbcc4d 100644 --- a/litellm/llms/bedrock/embed/twelvelabs_marengo_transformation.py +++ b/litellm/llms/bedrock/embed/twelvelabs_marengo_transformation.py @@ -7,14 +7,19 @@ Docs - https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-mar Marengo 3.0 docs - https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-marengo-3.html """ +from collections.abc import Mapping from typing import Final, cast +from pydantic import BaseModel, ConfigDict, TypeAdapter +from typing_extensions import assert_never + from litellm.llms.bedrock.embed.twelvelabs_marengo_3_transformation import ( build_marengo_3_request, is_marengo_3_model, ) from litellm.types.llms.bedrock import ( TWELVELABS_EMBEDDING_INPUT_TYPES, + TWELVELABS_MARENGO_3_INPUT_TYPES, TwelveLabsAsyncInvokeRequest, TwelveLabsMarengo3EmbeddingRequest, TwelveLabsMarengoEmbeddingRequest, @@ -22,7 +27,76 @@ from litellm.types.llms.bedrock import ( TwelveLabsS3Location, TwelveLabsS3OutputDataConfig, ) -from litellm.types.utils import Embedding, EmbeddingResponse, Usage +from litellm.types.utils import Embedding, EmbeddingResponse, PromptTokensDetailsWrapper, Usage + + +class MarengoEmbeddingItem(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + embedding: tuple[float, ...] + + +class MarengoInvokeResponse(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + data: tuple[MarengoEmbeddingItem, ...] = () + embedding: tuple[float, ...] | None = None + embeddings: tuple[MarengoEmbeddingItem, ...] = () + + def vectors(self) -> tuple[tuple[float, ...], ...]: + if self.data: + return tuple(item.embedding for item in self.data) + if self.embedding is not None: + return (self.embedding,) + return tuple(item.embedding for item in self.embeddings) + + +class MarengoBilledMultiInput(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + inputText: str | None = None + mediaSources: tuple[Mapping[str, object], ...] = () + + +class MarengoBilledRequest(BaseModel): + model_config = ConfigDict(extra="ignore", frozen=True) + + inputType: TWELVELABS_MARENGO_3_INPUT_TYPES | None = None + multi_input: MarengoBilledMultiInput | None = None + + +INVOKE_RESPONSES: Final = TypeAdapter(tuple[MarengoInvokeResponse, ...]) +BILLED_REQUESTS: Final = TypeAdapter(tuple[MarengoBilledRequest, ...]) + + +def _billed_units(request: MarengoBilledRequest) -> tuple[int, int]: + input_type: Final = request.inputType + match input_type: + case "text": + return (1, 0) + case "image": + return (0, 1) + case "text_image": + return (1, 1) + case "multi_input": + multi_input: Final = request.multi_input or MarengoBilledMultiInput() + return (1 if multi_input.inputText else 0, len(multi_input.mediaSources)) + case "video" | "audio" | None: + return (0, 0) + case _: + assert_never(input_type) + + +def _billed_usage(batch_data: list[dict] | None) -> Usage: + units: Final = tuple(_billed_units(request) for request in BILLED_REQUESTS.validate_python(batch_data or ())) + query_count: Final = sum(text_requests for text_requests, _ in units) + image_count: Final = sum(images for _, images in units) + details: Final = ( + PromptTokensDetailsWrapper(query_count=query_count or None, image_count=image_count or None) + if query_count or image_count + else None + ) + return Usage(prompt_tokens=0, completion_tokens=0, total_tokens=0, prompt_tokens_details=details) class TwelveLabsMarengoEmbeddingConfig: @@ -223,62 +297,16 @@ class TwelveLabsMarengoEmbeddingConfig: ), ) - def _transform_response(self, response_list: list[dict], model: str) -> EmbeddingResponse: - """ - Transform TwelveLabs response to OpenAI format. - Handles the actual TwelveLabs response format: {"data": [{"embedding": [...]}]} - """ - embeddings: Final[list[Embedding]] = [] - total_tokens = 0 - - for response in response_list: - # TwelveLabs response format has a "data" field containing the embeddings - if "data" in response and isinstance(response["data"], list): - for item in response["data"]: - if "embedding" in item: - # Single embedding response - embedding = Embedding( - embedding=item["embedding"], - index=len(embeddings), - object="embedding", - ) - embeddings.append(embedding) - - # Estimate token count (rough approximation) - if "inputTextTokenCount" in item: - total_tokens += item["inputTextTokenCount"] - else: - # Rough estimate: 1 token per 4 characters for text, or use embedding size - total_tokens += len(item["embedding"]) // 4 - elif "embedding" in response: - # Direct embedding response (fallback for other formats) - embedding = Embedding( - embedding=response["embedding"], - index=len(embeddings), - object="embedding", - ) - embeddings.append(embedding) - - # Estimate token count (rough approximation) - if "inputTextTokenCount" in response: - total_tokens += response["inputTextTokenCount"] - else: - # Rough estimate: 1 token per 4 characters for text - total_tokens += len(response.get("inputText", "")) // 4 - elif "embeddings" in response: - # Multiple embeddings response (from video/audio) - for i, emb in enumerate(response["embeddings"]): - embedding = Embedding( - embedding=emb["embedding"], - index=len(embeddings), - object="embedding", - ) - embeddings.append(embedding) - total_tokens += len(emb["embedding"]) // 4 # Rough estimate - - usage: Final = Usage(prompt_tokens=total_tokens, total_tokens=total_tokens) - - return EmbeddingResponse(data=embeddings, model=model, usage=usage) + def _transform_response( + self, response_list: list[dict], model: str, batch_data: list[dict] | None = None + ) -> EmbeddingResponse: + vectors: Final = tuple( + vector for response in INVOKE_RESPONSES.validate_python(response_list) for vector in response.vectors() + ) + embeddings: Final = [ + Embedding(embedding=list(vector), index=index, object="embedding") for index, vector in enumerate(vectors) + ] + return EmbeddingResponse(data=embeddings, model=model, usage=_billed_usage(batch_data)) def _transform_async_invoke_response(self, response: dict, model: str) -> EmbeddingResponse: """ diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index cc75354a495..7784ed2a6ac 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -650,7 +650,10 @@ }, "twelvelabs.marengo-embed-2-7-v1:0": { "deprecation_date": "2026-11-30", - "input_cost_per_token": 7e-05, + "input_cost_per_query": 7e-05, + "input_cost_per_video_per_second": 0.0007, + "input_cost_per_audio_per_second": 0.00014, + "input_cost_per_image": 0.0001, "litellm_provider": "bedrock", "max_input_tokens": 77, "max_tokens": 77, @@ -662,7 +665,7 @@ }, "us.twelvelabs.marengo-embed-2-7-v1:0": { "deprecation_date": "2026-11-30", - "input_cost_per_token": 7e-05, + "input_cost_per_query": 7e-05, "input_cost_per_video_per_second": 0.0007, "input_cost_per_audio_per_second": 0.00014, "input_cost_per_image": 0.0001, @@ -677,7 +680,7 @@ }, "eu.twelvelabs.marengo-embed-2-7-v1:0": { "deprecation_date": "2026-11-30", - "input_cost_per_token": 7e-05, + "input_cost_per_query": 7e-05, "input_cost_per_video_per_second": 0.0007, "input_cost_per_audio_per_second": 0.00014, "input_cost_per_image": 0.0001, @@ -691,7 +694,10 @@ "supports_image_input": true }, "twelvelabs.marengo-embed-3-0-v1:0": { - "input_cost_per_token": 7e-05, + "input_cost_per_query": 7e-05, + "input_cost_per_video_per_second": 0.0007, + "input_cost_per_audio_per_second": 0.00014, + "input_cost_per_image": 0.0001, "litellm_provider": "bedrock", "max_input_tokens": 500, "max_tokens": 500, @@ -702,7 +708,7 @@ "supports_image_input": true }, "us.twelvelabs.marengo-embed-3-0-v1:0": { - "input_cost_per_token": 7e-05, + "input_cost_per_query": 7e-05, "input_cost_per_video_per_second": 0.0007, "input_cost_per_audio_per_second": 0.00014, "input_cost_per_image": 0.0001, @@ -716,7 +722,7 @@ "supports_image_input": true }, "eu.twelvelabs.marengo-embed-3-0-v1:0": { - "input_cost_per_token": 7e-05, + "input_cost_per_query": 7e-05, "input_cost_per_video_per_second": 0.0007, "input_cost_per_audio_per_second": 0.00014, "input_cost_per_image": 0.0001, diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 9b5fb08a45f..c55eb6831c7 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -272,7 +272,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): input_cost_per_token_above_272k_tokens_flex: float | None input_cost_per_token_above_512k_tokens: float | None # MiniMax-M3: prompts >512K priced at 2x input input_cost_per_character_above_128k_tokens: float | None # only for vertex ai models - input_cost_per_query: float | None # only for rerank models + input_cost_per_query: float | None # per-request pricing: rerank, search, and Bedrock Marengo embeddings input_cost_per_image: float | None # only for vertex ai models input_cost_per_image_token: float | None # for gpt-image-1 and similar models input_cost_per_video_token: float | None # for gemini omni models with video input @@ -1693,6 +1693,9 @@ class PromptTokensDetailsWrapper( audio_length_seconds: float | None = None """Length of audio sent to the model. Used for multimodal embeddings priced per audio-second.""" + query_count: int | None = None + """Number of billable requests sent to the model. Used for embeddings priced per request, such as Bedrock Marengo.""" + cache_write_tokens: int | None = None """Number of cache write (creation) tokens sent to the model. OpenAI naming (prompt_tokens_details.cache_write_tokens); this is the canonical field.""" @@ -1734,6 +1737,8 @@ class PromptTokensDetailsWrapper( del self.video_length_seconds if self.audio_length_seconds is None: del self.audio_length_seconds + if self.query_count is None: + del self.query_count if self.web_search_requests is None: del self.web_search_requests if self.google_maps_grounding_requests is None: diff --git a/litellm/utils.py b/litellm/utils.py index b98aa821ff3..2ed1ad84e9c 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -6033,7 +6033,7 @@ def get_model_info( input_cost_per_character_above_128k_tokens: Optional[ float ] # only for vertex ai models - input_cost_per_query: Optional[float] # only for rerank models + input_cost_per_query: Optional[float] # per-request pricing: rerank, search, and Bedrock Marengo embeddings input_cost_per_image: Optional[float] # only for vertex ai models input_cost_per_audio_token: Optional[float] input_cost_per_audio_per_second: Optional[float] # only for vertex ai models diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index cc75354a495..7784ed2a6ac 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -650,7 +650,10 @@ }, "twelvelabs.marengo-embed-2-7-v1:0": { "deprecation_date": "2026-11-30", - "input_cost_per_token": 7e-05, + "input_cost_per_query": 7e-05, + "input_cost_per_video_per_second": 0.0007, + "input_cost_per_audio_per_second": 0.00014, + "input_cost_per_image": 0.0001, "litellm_provider": "bedrock", "max_input_tokens": 77, "max_tokens": 77, @@ -662,7 +665,7 @@ }, "us.twelvelabs.marengo-embed-2-7-v1:0": { "deprecation_date": "2026-11-30", - "input_cost_per_token": 7e-05, + "input_cost_per_query": 7e-05, "input_cost_per_video_per_second": 0.0007, "input_cost_per_audio_per_second": 0.00014, "input_cost_per_image": 0.0001, @@ -677,7 +680,7 @@ }, "eu.twelvelabs.marengo-embed-2-7-v1:0": { "deprecation_date": "2026-11-30", - "input_cost_per_token": 7e-05, + "input_cost_per_query": 7e-05, "input_cost_per_video_per_second": 0.0007, "input_cost_per_audio_per_second": 0.00014, "input_cost_per_image": 0.0001, @@ -691,7 +694,10 @@ "supports_image_input": true }, "twelvelabs.marengo-embed-3-0-v1:0": { - "input_cost_per_token": 7e-05, + "input_cost_per_query": 7e-05, + "input_cost_per_video_per_second": 0.0007, + "input_cost_per_audio_per_second": 0.00014, + "input_cost_per_image": 0.0001, "litellm_provider": "bedrock", "max_input_tokens": 500, "max_tokens": 500, @@ -702,7 +708,7 @@ "supports_image_input": true }, "us.twelvelabs.marengo-embed-3-0-v1:0": { - "input_cost_per_token": 7e-05, + "input_cost_per_query": 7e-05, "input_cost_per_video_per_second": 0.0007, "input_cost_per_audio_per_second": 0.00014, "input_cost_per_image": 0.0001, @@ -716,7 +722,7 @@ "supports_image_input": true }, "eu.twelvelabs.marengo-embed-3-0-v1:0": { - "input_cost_per_token": 7e-05, + "input_cost_per_query": 7e-05, "input_cost_per_video_per_second": 0.0007, "input_cost_per_audio_per_second": 0.00014, "input_cost_per_image": 0.0001, diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 59f0938e338..65a6dd2a4ca 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -2658,6 +2658,7 @@ def test_cache_writing_cost_with_zero_creation_tokens_and_ephemeral_details(): "image_count": 0, "video_length_seconds": 0.0, "audio_length_seconds": 0.0, + "query_count": 0, } model_info: ModelInfo = {} @@ -3239,6 +3240,37 @@ def test_image_count_prevents_text_tokens_fallback(_local_model_cost_map): assert completion_cost == 0.0 +def test_query_count_bills_input_cost_per_query(_local_model_cost_map): + usage = Usage( + prompt_tokens=0, + completion_tokens=0, + total_tokens=0, + prompt_tokens_details=PromptTokensDetailsWrapper(query_count=3, image_count=1), + ) + + prompt_cost, completion_cost = generic_cost_per_token( + model="us.twelvelabs.marengo-embed-3-0-v1:0", + usage=usage, + custom_llm_provider="bedrock", + ) + + assert prompt_cost == pytest.approx(3 * 7e-05 + 1e-04) + assert completion_cost == 0.0 + + +def test_query_count_is_free_without_a_per_query_price(_local_model_cost_map): + usage = Usage( + prompt_tokens=0, + completion_tokens=0, + total_tokens=0, + prompt_tokens_details=PromptTokensDetailsWrapper(query_count=1), + ) + + prompt_cost, _ = generic_cost_per_token(model="text-embedding-3-small", usage=usage, custom_llm_provider="openai") + + assert prompt_cost == 0.0 + + # --------------------------------------------------------------------------- # Data-residency (OpenAI regional processing) tests # --------------------------------------------------------------------------- diff --git a/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py b/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py index b37e991b0b2..c29a87cd0cf 100644 --- a/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py +++ b/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py @@ -5,6 +5,7 @@ from unittest.mock import Mock, patch import pytest import litellm +from litellm.llms.bedrock.embed.twelvelabs_marengo_transformation import TwelveLabsMarengoEmbeddingConfig from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler # Mock responses for different embedding models @@ -1066,17 +1067,19 @@ MARENGO_3_DUCK = "data:image/png;base64,ZHVjaw==" @pytest.mark.parametrize( - "model,kwargs,expected_body", + "model,kwargs,expected_body,expected_usage_details", [ ( "bedrock/us.twelvelabs.marengo-embed-3-0-v1:0", {"input_type": "text"}, {"inputType": "text", "text": {"inputText": "a duck on water"}}, + {"query_count": 1}, ), ( "bedrock/twelvelabs.marengo-embed-3-0-v1:0", {"input_type": "text"}, {"inputType": "text", "text": {"inputText": "a duck on water"}}, + {"query_count": 1}, ), ( "bedrock/us.twelvelabs.marengo-embed-3-0-v1:0", @@ -1085,6 +1088,7 @@ MARENGO_3_DUCK = "data:image/png;base64,ZHVjaw==" "inputType": "text_image", "text_image": {"inputText": "a duck on water", "mediaSource": {"base64String": "ZHVjaw=="}}, }, + {"query_count": 1, "image_count": 1}, ), ( "bedrock/us.twelvelabs.marengo-embed-3-0-v1:0", @@ -1096,10 +1100,13 @@ MARENGO_3_DUCK = "data:image/png;base64,ZHVjaw==" "mediaSources": [{"name": "bird", "mediaType": "image", "base64String": "ZHVjaw=="}], }, }, + {"query_count": 1, "image_count": 1}, ), ], ) -def test_marengo_3_embedding_sends_the_nested_payload_and_parses_512_dims(model, kwargs, expected_body): +def test_marengo_3_embedding_sends_the_nested_payload_and_parses_512_dims( + model, kwargs, expected_body, expected_usage_details +): client = HTTPHandler() with patch.object(client, "post") as mock_post: @@ -1122,7 +1129,9 @@ def test_marengo_3_embedding_sends_the_nested_payload_and_parses_512_dims(model, assert mock_post.call_args.kwargs["url"].endswith(f"/model/{model.removeprefix('bedrock/').replace(':', '%3A')}/invoke") assert len(response.data[0]["embedding"]) == 512 assert response.data[0]["embedding"][:2] == [0.0, 0.01] - assert response.usage.prompt_tokens == 128 + assert response.usage.prompt_tokens == 0 + assert response.usage.total_tokens == 0 + assert response.usage.prompt_tokens_details.model_dump(exclude_none=True) == expected_usage_details def test_marengo_3_image_embedding_sends_the_media_under_the_image_key(): @@ -1150,6 +1159,8 @@ def test_marengo_3_image_embedding_sends_the_media_under_the_image_key(): } assert len(response.data[0]["embedding"]) == 512 assert response.data[0]["embedding"][:2] == [0.0, 0.01] + assert response.usage.prompt_tokens == 0 + assert response.usage.prompt_tokens_details.model_dump(exclude_none=True) == {"image_count": 1} def test_marengo_2_7_embedding_keeps_the_flat_payload(): @@ -1177,6 +1188,36 @@ def test_marengo_2_7_embedding_keeps_the_flat_payload(): "textTruncate": "end", } assert response.data[0]["embedding"] == [0.1, 0.2, 0.3] + assert response.usage.prompt_tokens == 0 + assert response.usage.prompt_tokens_details.model_dump(exclude_none=True) == {"query_count": 1} + + +def test_marengo_usage_counts_text_requests_and_images_across_a_batch(): + duck = {"mediaType": "image", "base64String": "ZHVjaw=="} + response = TwelveLabsMarengoEmbeddingConfig()._transform_response( + response_list=[marengo_3_embedding_response, marengo_3_embedding_response, marengo_3_embedding_response], + model="us.twelvelabs.marengo-embed-3-0-v1:0", + batch_data=[ + {"inputType": "text", "text": {"inputText": "a duck"}}, + {"inputType": "image", "image": {"mediaSource": {"base64String": "ZHVjaw=="}}}, + {"inputType": "multi_input", "multi_input": {"mediaSources": [{"name": "a", **duck}, {"name": "b", **duck}]}}, + ], + ) + + assert [item["index"] for item in response.data] == [0, 1, 2] + assert response.usage.prompt_tokens == 0 + assert response.usage.total_tokens == 0 + assert response.usage.prompt_tokens_details.model_dump(exclude_none=True) == {"query_count": 1, "image_count": 3} + + +def test_marengo_usage_without_request_data_bills_nothing(): + response = TwelveLabsMarengoEmbeddingConfig()._transform_response( + response_list=[marengo_3_embedding_response], model="us.twelvelabs.marengo-embed-3-0-v1:0" + ) + + assert len(response.data[0]["embedding"]) == 512 + assert response.usage.prompt_tokens == 0 + assert response.usage.prompt_tokens_details is None def test_marengo_3_text_image_without_media_source_is_a_bad_request(): diff --git a/tests/test_litellm/test_bedrock_marengo_embed_3_model_metadata.py b/tests/test_litellm/test_bedrock_marengo_embed_3_model_metadata.py index 300dfeb5238..0bb99339435 100644 --- a/tests/test_litellm/test_bedrock_marengo_embed_3_model_metadata.py +++ b/tests/test_litellm/test_bedrock_marengo_embed_3_model_metadata.py @@ -6,7 +6,7 @@ import pytest import litellm from litellm.constants import bedrock_embedding_models from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider -from litellm.types.utils import Usage +from litellm.types.utils import PromptTokensDetailsWrapper, Usage REPO_ROOT = Path(__file__).parents[2] MAIN_PATH = REPO_ROOT / "model_prices_and_context_window.json" @@ -15,6 +15,12 @@ BACKUP_PATH = REPO_ROOT / "litellm" / "model_prices_and_context_window_backup.js BASE_MODEL = "twelvelabs.marengo-embed-3-0-v1:0" PROFILE_MODELS = ("us.twelvelabs.marengo-embed-3-0-v1:0", "eu.twelvelabs.marengo-embed-3-0-v1:0") ALL_MODELS = (BASE_MODEL, *PROFILE_MODELS) +MARENGO_2_7_MODELS = ( + "twelvelabs.marengo-embed-2-7-v1:0", + "us.twelvelabs.marengo-embed-2-7-v1:0", + "eu.twelvelabs.marengo-embed-2-7-v1:0", +) +PER_REQUEST_MODELS = (*ALL_MODELS, *MARENGO_2_7_MODELS) TEXT_REQUEST_COST = 7e-05 IMAGE_REQUEST_COST = 0.0001 @@ -34,7 +40,7 @@ def test_marengo_embed_3_specs(model): assert info["litellm_provider"] == "bedrock" assert info["mode"] == "embedding" - assert info["input_cost_per_token"] == TEXT_REQUEST_COST + assert info["input_cost_per_query"] == TEXT_REQUEST_COST assert info["output_cost_per_token"] == 0.0 assert info["max_input_tokens"] == 500 assert info["max_tokens"] == 500 @@ -48,9 +54,11 @@ def test_marengo_embed_3_specs(model): assert provider == "bedrock" -@pytest.mark.parametrize("model", PROFILE_MODELS) -def test_marengo_embed_3_inference_profiles_price_image_video_and_audio(model): +@pytest.mark.parametrize("model", PER_REQUEST_MODELS) +def test_marengo_prices_are_per_request_not_per_token(model): info = _load(MAIN_PATH)[model] + assert "input_cost_per_token" not in info + assert info["input_cost_per_query"] == TEXT_REQUEST_COST assert info["input_cost_per_image"] == IMAGE_REQUEST_COST assert info["input_cost_per_video_per_second"] == VIDEO_COST_PER_SECOND assert info["input_cost_per_audio_per_second"] == AUDIO_COST_PER_SECOND @@ -64,13 +72,34 @@ def test_marengo_embed_3_is_visible_to_callers(model, local_model_cost_map): assert info["max_input_tokens"] == 500 -@pytest.mark.parametrize("model", ALL_MODELS) -def test_marengo_embed_3_text_request_is_billed(model, local_model_cost_map): +@pytest.mark.parametrize("model", PER_REQUEST_MODELS) +@pytest.mark.parametrize( + "details,expected_cost", + [ + (PromptTokensDetailsWrapper(query_count=1), TEXT_REQUEST_COST), + (PromptTokensDetailsWrapper(image_count=1), IMAGE_REQUEST_COST), + (PromptTokensDetailsWrapper(query_count=1, image_count=1), TEXT_REQUEST_COST + IMAGE_REQUEST_COST), + (PromptTokensDetailsWrapper(query_count=1, image_count=2), TEXT_REQUEST_COST + 2 * IMAGE_REQUEST_COST), + (PromptTokensDetailsWrapper(video_length_seconds=10), 10 * VIDEO_COST_PER_SECOND), + (PromptTokensDetailsWrapper(audio_length_seconds=10), 10 * AUDIO_COST_PER_SECOND), + ], +) +def test_marengo_requests_are_billed_per_request(model, details, expected_cost, local_model_cost_map): + usage = Usage(prompt_tokens=0, completion_tokens=0, total_tokens=0, prompt_tokens_details=details) + prompt_cost, completion_cost = litellm.cost_per_token( + model=model, usage_object=usage, custom_llm_provider="bedrock" + ) + assert prompt_cost == pytest.approx(expected_cost) + assert completion_cost == 0.0 + + +@pytest.mark.parametrize("model", PER_REQUEST_MODELS) +def test_marengo_token_counts_bill_nothing(model, local_model_cost_map): usage = Usage(prompt_tokens=128, completion_tokens=0, total_tokens=128) prompt_cost, completion_cost = litellm.cost_per_token( model=model, usage_object=usage, custom_llm_provider="bedrock" ) - assert prompt_cost == pytest.approx(128 * TEXT_REQUEST_COST) + assert prompt_cost == 0.0 assert completion_cost == 0.0 @@ -78,7 +107,7 @@ def test_marengo_embed_3_is_a_known_bedrock_embedding_model(): assert BASE_MODEL in bedrock_embedding_models -@pytest.mark.parametrize("model", ALL_MODELS) +@pytest.mark.parametrize("model", PER_REQUEST_MODELS) def test_backup_matches_main(model): main_cost = _load(MAIN_PATH) backup_cost = _load(BACKUP_PATH) From 80fea089b63e6b89e989f3a109b96a26bac90224 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 7 Sep 2026 18:40:31 -0700 Subject: [PATCH 054/109] fix(bedrock): reject Marengo 2.7-only and misplaced media params on 3.0 unless drop_params Marengo 3.0 requests now get a 400 naming any textTruncate, lengthSec, useFixedLengthSec, or minClipSec parameter, and any video or audio option sent with a text, image, text_image, or multi_input request, instead of silently dropping them. drop_params (global, per deployment, or per request) drops them instead. Pydantic validation errors name the field and the reason, and the 3.0 marker is the exact "marengo-embed-3-" model id segment. --- litellm/llms/bedrock/embed/embedding.py | 3 +- .../twelvelabs_marengo_3_transformation.py | 44 +++++++-- .../twelvelabs_marengo_transformation.py | 46 ++++++---- ...est_twelvelabs_marengo_3_transformation.py | 90 +++++++++++++++++++ 4 files changed, 158 insertions(+), 25 deletions(-) diff --git a/litellm/llms/bedrock/embed/embedding.py b/litellm/llms/bedrock/embed/embedding.py index 69eb9b693b1..987c7cbf981 100644 --- a/litellm/llms/bedrock/embed/embedding.py +++ b/litellm/llms/bedrock/embed/embedding.py @@ -35,7 +35,7 @@ from .amazon_titan_multimodal_transformation import ( ) from .amazon_titan_v2_transformation import AmazonTitanV2Config from .cohere_transformation import BedrockCohereEmbeddingConfig -from .twelvelabs_marengo_transformation import TwelveLabsMarengoEmbeddingConfig +from .twelvelabs_marengo_transformation import TwelveLabsMarengoEmbeddingConfig, drop_params_enabled if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -480,6 +480,7 @@ class BedrockEmbedding(BaseAWSLLM): async_invoke_route=has_async_invoke, model_id=modelId, output_s3_uri=inference_params.get("output_s3_uri"), + drop_params=drop_params_enabled(litellm_params), ) batch_data.append(twelvelabs_request) elif provider == "nova": diff --git a/litellm/llms/bedrock/embed/twelvelabs_marengo_3_transformation.py b/litellm/llms/bedrock/embed/twelvelabs_marengo_3_transformation.py index 0f61d37258f..f4f9cb03dab 100644 --- a/litellm/llms/bedrock/embed/twelvelabs_marengo_3_transformation.py +++ b/litellm/llms/bedrock/embed/twelvelabs_marengo_3_transformation.py @@ -35,7 +35,7 @@ from litellm.types.llms.bedrock import ( ) from litellm.utils import get_base64_str -MARENGO_3_MODEL_MARKER: Final = "marengo-embed-3" +MARENGO_3_MODEL_MARKER: Final = "marengo-embed-3-" S3_URI_PREFIX: Final = "s3://" TIMED_MEDIA_OPTION_FIELDS: Final = MappingProxyType( { @@ -48,6 +48,9 @@ TIMED_MEDIA_OPTION_FIELDS: Final = MappingProxyType( } ) TIMED_MEDIA_OPTIONS: Final = TypeAdapter(TwelveLabsMarengo3TimedMediaOptions) +TIMED_INPUT_TYPES: Final = frozenset({"video", "audio"}) +MARENGO_2_7_ONLY_PARAMS: Final = ("textTruncate", "lengthSec", "useFixedLengthSec", "minClipSec") +MARENGO_2_7_ONLY_FIELDS: Final = MappingProxyType({name: True for name in MARENGO_2_7_ONLY_PARAMS}) def is_marengo_3_model(model: str | None) -> bool: @@ -69,15 +72,23 @@ class Marengo3Params(BaseModel): embeddingType: tuple[TWELVELABS_MARENGO_3_EMBEDDING_TYPES, ...] | None = None embeddingScope: tuple[TWELVELABS_MARENGO_3_EMBEDDING_SCOPES, ...] | None = None inferenceId: str | None = None + textTruncate: object = None + lengthSec: object = None + useFixedLengthSec: object = None + minClipSec: object = None @property def resolved_input_type(self) -> TWELVELABS_MARENGO_3_INPUT_TYPES: return self.inputType or self.input_type or "text" def timed_media_options(self) -> TwelveLabsMarengo3TimedMediaOptions: - return TIMED_MEDIA_OPTIONS.validate_python( - self.model_dump(include=TIMED_MEDIA_OPTION_FIELDS, exclude_none=True) - ) + return TIMED_MEDIA_OPTIONS.validate_python(self.given_timed_media_options()) + + def given_timed_media_options(self) -> dict[str, object]: + return self.model_dump(include=TIMED_MEDIA_OPTION_FIELDS, exclude_none=True) + + def given_2_7_only_params(self) -> dict[str, object]: + return self.model_dump(include=MARENGO_2_7_ONLY_FIELDS, exclude_none=True) def _s3_location(uri: str, bucket_owner: str | None) -> TwelveLabsS3Location: @@ -113,11 +124,23 @@ def _timed_media_input(media: str, params: Marengo3Params) -> TwelveLabsMarengo3 return timed +def _describe(error: ValidationError) -> str: + return "; ".join( + f"{'.'.join(str(part) for part in problem['loc'])}: {problem['msg']}" for problem in error.errors() + ) + + def _validated_params(inference_params: Mapping[str, object]) -> Marengo3Params: try: return Marengo3Params.model_validate(inference_params) except ValidationError as error: - raise BedrockError(status_code=400, message=f"Invalid Marengo 3.0 parameters: {error}") from error + raise BedrockError(status_code=400, message=f"Invalid Marengo 3.0 parameters: {_describe(error)}") from error + + +def _reject_unless_dropped(given: Mapping[str, object], drop_params: bool, reason: str) -> None: + if not given or drop_params: + return + raise BedrockError(status_code=400, message=f"{reason} {', '.join(given)}; set drop_params to drop them") def _require(value: str | None, input_type: str, param_name: str) -> str: @@ -143,10 +166,19 @@ def _request_base(inference_id: str | None) -> TwelveLabsMarengo3RequestBase: return identified -def build_marengo_3_request(input: str, inference_params: Mapping[str, object]) -> TwelveLabsMarengo3EmbeddingRequest: +def build_marengo_3_request( + input: str, inference_params: Mapping[str, object], drop_params: bool = False +) -> TwelveLabsMarengo3EmbeddingRequest: params: Final = _validated_params(inference_params) base: Final = _request_base(params.inferenceId) input_type: Final = params.resolved_input_type + _reject_unless_dropped( + params.given_2_7_only_params(), drop_params, "Marengo 3.0 does not accept the Marengo 2.7 parameters" + ) + if input_type not in TIMED_INPUT_TYPES: + _reject_unless_dropped( + params.given_timed_media_options(), drop_params, f"Input type '{input_type}' does not accept" + ) match input_type: case "text": text_request: Final[TwelveLabsMarengo3TextRequest] = { diff --git a/litellm/llms/bedrock/embed/twelvelabs_marengo_transformation.py b/litellm/llms/bedrock/embed/twelvelabs_marengo_transformation.py index ddf6dfbcc4d..35163ecf848 100644 --- a/litellm/llms/bedrock/embed/twelvelabs_marengo_transformation.py +++ b/litellm/llms/bedrock/embed/twelvelabs_marengo_transformation.py @@ -13,7 +13,9 @@ from typing import Final, cast from pydantic import BaseModel, ConfigDict, TypeAdapter from typing_extensions import assert_never +import litellm from litellm.llms.bedrock.embed.twelvelabs_marengo_3_transformation import ( + MARENGO_2_7_ONLY_PARAMS, build_marengo_3_request, is_marengo_3_model, ) @@ -99,6 +101,25 @@ def _billed_usage(batch_data: list[dict] | None) -> Usage: return Usage(prompt_tokens=0, completion_tokens=0, total_tokens=0, prompt_tokens_details=details) +MARENGO_SHARED_PARAMS: Final = ( + "encoding_format", + "embeddingOption", + "startSec", + "input_type", + "endSec", + "segmentation", + "embeddingType", + "embeddingScope", + "inferenceId", + "media_source", + "media_sources", +) + + +def drop_params_enabled(litellm_params: Mapping[str, object]) -> bool: + return litellm.drop_params is True or litellm_params.get("drop_params") is True + + class TwelveLabsMarengoEmbeddingConfig: """ Reference - https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-marengo.html @@ -115,23 +136,9 @@ class TwelveLabsMarengoEmbeddingConfig: self.is_marengo_3: Final = is_marengo_3_model(model) def get_supported_openai_params(self) -> list[str]: - return [ - "encoding_format", - "textTruncate", - "embeddingOption", - "startSec", - "lengthSec", - "useFixedLengthSec", - "minClipSec", - "input_type", - "endSec", - "segmentation", - "embeddingType", - "embeddingScope", - "inferenceId", - "media_source", - "media_sources", - ] + if self.is_marengo_3: + return list(MARENGO_SHARED_PARAMS) + return [*MARENGO_SHARED_PARAMS, *MARENGO_2_7_ONLY_PARAMS] def map_openai_params(self, non_default_params: dict, optional_params: dict) -> dict: for k, v in non_default_params.items(): @@ -179,6 +186,7 @@ class TwelveLabsMarengoEmbeddingConfig: async_invoke_route: bool = False, model_id: str | None = None, output_s3_uri: str | None = None, + drop_params: bool = False, ) -> TwelveLabsMarengoEmbeddingRequest | TwelveLabsMarengo3EmbeddingRequest | TwelveLabsAsyncInvokeRequest: """ Transform OpenAI-style input to TwelveLabs Marengo format/async-invoke format. @@ -203,7 +211,9 @@ class TwelveLabsMarengoEmbeddingConfig: ) if self.is_marengo_3: - marengo_3_request: Final = build_marengo_3_request(input=input, inference_params=inference_params) + marengo_3_request: Final = build_marengo_3_request( + input=input, inference_params=inference_params, drop_params=drop_params + ) if async_invoke_route and model_id: return self._wrap_async_invoke_request( model_input=marengo_3_request, model_id=model_id, output_s3_uri=output_s3_uri diff --git a/tests/test_litellm/llms/bedrock/embed/test_twelvelabs_marengo_3_transformation.py b/tests/test_litellm/llms/bedrock/embed/test_twelvelabs_marengo_3_transformation.py index 0bf86352a5e..5033256d089 100644 --- a/tests/test_litellm/llms/bedrock/embed/test_twelvelabs_marengo_3_transformation.py +++ b/tests/test_litellm/llms/bedrock/embed/test_twelvelabs_marengo_3_transformation.py @@ -2,13 +2,16 @@ import json import pytest +import litellm from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.bedrock.embed.twelvelabs_marengo_3_transformation import ( + MARENGO_2_7_ONLY_PARAMS, build_marengo_3_request, is_marengo_3_model, ) from litellm.llms.bedrock.embed.twelvelabs_marengo_transformation import ( TwelveLabsMarengoEmbeddingConfig, + drop_params_enabled, ) MARENGO_3_BASE = "twelvelabs.marengo-embed-3-0-v1:0" @@ -27,6 +30,7 @@ OUTPUT_S3_URI = "s3://out-bucket/marengo/" ("async_invoke/twelvelabs.marengo-embed-3-0-v1:0", True), (MARENGO_27_US, False), ("twelvelabs.marengo-embed-2-7-v1:0", False), + ("twelvelabs.marengo-embed-30-v1:0", False), (None, False), ], ) @@ -266,3 +270,89 @@ def test_marengo_3_only_params_are_forwarded_by_map_openai_params(): "embeddingScope": ["clip"], "inferenceId": "req-1", } + + +@pytest.mark.parametrize( + "params,problem", + [ + ( + {"input_type": "clip"}, + "input_type: Input should be 'text', 'image', 'video', 'audio', 'text_image' or 'multi_input'", + ), + ({"input_type": "video", "embeddingOption": "visual"}, "embeddingOption: Input should be a valid tuple"), + ( + {"input_type": "multi_input", "media_sources": ["not", "a", "mapping"]}, + "media_sources: Input should be a valid dictionary", + ), + ], +) +def test_invalid_marengo_3_params_name_the_field_and_the_reason(params, problem): + with pytest.raises(BedrockError) as excinfo: + build_marengo_3_request("s3://media/clip.mp4", params) + assert excinfo.value.message == f"Invalid Marengo 3.0 parameters: {problem}" + + +MARENGO_2_7_ONLY_VALUES = {"textTruncate": "end", "lengthSec": 5, "useFixedLengthSec": True, "minClipSec": 2} + + +@pytest.mark.parametrize("name", MARENGO_2_7_ONLY_PARAMS) +def test_marengo_2_7_only_params_are_rejected_on_3_0_unless_dropped(name): + params = {"input_type": "text", name: MARENGO_2_7_ONLY_VALUES[name]} + with pytest.raises(BedrockError) as excinfo: + build_marengo_3_request("hello", params) + assert excinfo.value.status_code == 400 + assert excinfo.value.message == ( + f"Marengo 3.0 does not accept the Marengo 2.7 parameters {name}; set drop_params to drop them" + ) + assert build_marengo_3_request("hello", params, drop_params=True) == { + "inputType": "text", + "text": {"inputText": "hello"}, + } + + +def test_marengo_2_7_only_params_are_advertised_only_for_2_7(): + marengo_3 = TwelveLabsMarengoEmbeddingConfig(model=MARENGO_3_US).get_supported_openai_params() + marengo_27 = TwelveLabsMarengoEmbeddingConfig(model=MARENGO_27_US).get_supported_openai_params() + assert set(MARENGO_2_7_ONLY_PARAMS).isdisjoint(marengo_3) + assert set(MARENGO_2_7_ONLY_PARAMS) <= set(marengo_27) + assert set(marengo_3) <= set(marengo_27) + + +def test_drop_params_comes_from_the_call_or_the_global(monkeypatch): + monkeypatch.setattr(litellm, "drop_params", False) + assert drop_params_enabled({}) is False + assert drop_params_enabled({"drop_params": True}) is True + monkeypatch.setattr(litellm, "drop_params", True) + assert drop_params_enabled({}) is True + + +def test_config_drops_marengo_2_7_only_params_only_when_asked(): + config = TwelveLabsMarengoEmbeddingConfig(model=MARENGO_3_US) + with pytest.raises(BedrockError, match=r"Marengo 2\.7 parameters textTruncate"): + config._transform_request("hello", {"textTruncate": "end"}) + assert config._transform_request("hello", {"textTruncate": "end"}, drop_params=True) == { + "inputType": "text", + "text": {"inputText": "hello"}, + } + + +@pytest.mark.parametrize( + "params", + [ + {"input_type": "text"}, + {"input_type": "image"}, + {"input_type": "text_image", "media_source": DUCK_DATA_URL}, + {"input_type": "multi_input", "media_sources": {"bird": DUCK_DATA_URL}}, + ], +) +def test_timed_media_options_are_rejected_on_untimed_input_types_unless_dropped(params): + timed = {**params, "startSec": 0, "embeddingOption": ["visual"]} + with pytest.raises(BedrockError) as excinfo: + build_marengo_3_request(DUCK_DATA_URL, timed) + assert excinfo.value.status_code == 400 + assert excinfo.value.message == ( + f"Input type '{params['input_type']}' does not accept startSec, embeddingOption; set drop_params to drop them" + ) + assert build_marengo_3_request(DUCK_DATA_URL, timed, drop_params=True) == build_marengo_3_request( + DUCK_DATA_URL, params + ) From a601c00afdc5ddda50cb7552aef0dc5f3d9dcf0c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 7 Sep 2026 18:47:25 -0700 Subject: [PATCH 055/109] fix(bedrock): pass litellm_params into the Bedrock embedding call so drop_params reaches Marengo 3.0 --- litellm/main.py | 2 +- ...est_twelvelabs_marengo_3_transformation.py | 34 +++++++++++++++++++ 2 files changed, 35 insertions(+), 1 deletion(-) diff --git a/litellm/main.py b/litellm/main.py index 75b7f7f10a5..56f9cb2c0d0 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -6545,7 +6545,7 @@ def embedding( client=client, timeout=timeout, aembedding=aembedding, - litellm_params={}, + litellm_params=litellm_params_dict, api_base=api_base, print_verbose=print_verbose, extra_headers=headers, diff --git a/tests/test_litellm/llms/bedrock/embed/test_twelvelabs_marengo_3_transformation.py b/tests/test_litellm/llms/bedrock/embed/test_twelvelabs_marengo_3_transformation.py index 5033256d089..d8d29cac35e 100644 --- a/tests/test_litellm/llms/bedrock/embed/test_twelvelabs_marengo_3_transformation.py +++ b/tests/test_litellm/llms/bedrock/embed/test_twelvelabs_marengo_3_transformation.py @@ -1,9 +1,11 @@ import json +from unittest.mock import Mock, patch import pytest import litellm from litellm.llms.bedrock.common_utils import BedrockError +from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.llms.bedrock.embed.twelvelabs_marengo_3_transformation import ( MARENGO_2_7_ONLY_PARAMS, build_marengo_3_request, @@ -356,3 +358,35 @@ def test_timed_media_options_are_rejected_on_untimed_input_types_unless_dropped( assert build_marengo_3_request(DUCK_DATA_URL, timed, drop_params=True) == build_marengo_3_request( DUCK_DATA_URL, params ) + + +def _embed_marengo_3_us(client: HTTPHandler, **params: object): + return litellm.embedding( + model=f"bedrock/{MARENGO_3_US}", + input="hello", + client=client, + aws_region_name="us-east-1", + aws_bedrock_runtime_endpoint="https://bedrock-runtime.us-east-1.amazonaws.com", + api_key="test-bearer-token", + **params, + ) + + +def test_per_request_drop_params_reaches_the_marengo_3_builder(monkeypatch): + monkeypatch.setattr(litellm, "drop_params", False) + client = HTTPHandler() + with patch.object(client, "post") as mock_post: + mock_response = Mock() + mock_response.status_code = 200 + mock_response.text = json.dumps({"data": [{"embedding": [0.1, 0.2]}]}) + mock_response.json = lambda: json.loads(mock_response.text) + mock_post.return_value = mock_response + + with pytest.raises(litellm.BadRequestError, match=r"Marengo 2\.7 parameters textTruncate"): + _embed_marengo_3_us(client, textTruncate="end") + assert mock_post.call_count == 0 + + response = _embed_marengo_3_us(client, textTruncate="end", drop_params=True) + + assert response.data[0]["embedding"] == [0.1, 0.2] + assert json.loads(mock_post.call_args.kwargs["data"]) == {"inputType": "text", "text": {"inputText": "hello"}} From 6076e9f61103f9f8054b8dfaa00539e47fc9163f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 7 Sep 2026 18:54:31 -0700 Subject: [PATCH 056/109] chore(cost_calc): drop the query count section label comment --- litellm/litellm_core_utils/llm_cost_calc/utils.py | 1 - 1 file changed, 1 deletion(-) diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 46574ebae3f..c05d4c29a5e 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -981,7 +981,6 @@ def _calculate_input_cost( prompt_tokens_details["audio_length_seconds"], ) - ### QUERY COUNT COST if prompt_tokens_details["query_count"]: prompt_cost += calculate_cost_component( model_info, "input_cost_per_query", prompt_tokens_details["query_count"] 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 057/109] 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 058/109] 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 059/109] 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 060/109] 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 061/109] 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 062/109] 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 063/109] 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 5c037299f413c38609cbb7f5a582662e44b197d4 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 7 Sep 2026 23:02:27 -0700 Subject: [PATCH 064/109] feat: move MongoDB vector search to an optional sidecar --- .github/workflows/_test-unit-base.yml | 2 +- Dockerfile | 2 - docker/Dockerfile.database | 2 - docker/Dockerfile.non_root | 3 - gateway/Dockerfile | 2 - .../base_llm/vector_store/transformation.py | 3 + litellm/llms/custom_httpx/llm_http_handler.py | 17 +- litellm/llms/mongodb/common_utils.py | 303 --- .../mongodb/vector_stores/transformation.py | 449 ++--- pyproject.toml | 1 - .../test_mongodb_transformation.py | 1691 ++--------------- .../_components/VectorStoreForm.test.tsx | 12 +- .../_components/VectorStoreForm.tsx | 4 +- .../vector_store_providers.test.tsx | 11 +- .../src/components/vector_store_providers.tsx | 27 +- uv.lock | 79 +- 16 files changed, 412 insertions(+), 2196 deletions(-) delete mode 100644 litellm/llms/mongodb/common_utils.py diff --git a/.github/workflows/_test-unit-base.yml b/.github/workflows/_test-unit-base.yml index 75b0f93fd77..62790e23143 100644 --- a/.github/workflows/_test-unit-base.yml +++ b/.github/workflows/_test-unit-base.yml @@ -113,7 +113,7 @@ jobs: if: steps.changes.outputs.decision != 'skip' timeout-minutes: 8 run: | - .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml --extra mongodb + .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml uv run --no-sync python -c 'import os, sys; print(sys.version); assert f"{sys.version_info.major}.{sys.version_info.minor}" == os.environ["UV_PYTHON"]' - name: Cache Prisma binaries diff --git a/Dockerfile b/Dockerfile index 1648ec69d13..0a92aa9a68c 100644 --- a/Dockerfile +++ b/Dockerfile @@ -67,7 +67,6 @@ RUN uv sync --frozen --no-install-project --no-install-workspace --no-default-gr --extra semantic-router \ --extra saml \ --extra bedrock-realtime \ - --extra mongodb \ --python python3.13 # Copy full source tree @@ -90,7 +89,6 @@ RUN uv sync --frozen --no-default-groups --no-editable \ --extra semantic-router \ --extra saml \ --extra bedrock-realtime \ - --extra mongodb \ --python python3.13 RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \ diff --git a/docker/Dockerfile.database b/docker/Dockerfile.database index cc81ad6b3d3..e9ad2849bb2 100644 --- a/docker/Dockerfile.database +++ b/docker/Dockerfile.database @@ -65,7 +65,6 @@ RUN uv sync --frozen --no-install-project --no-install-workspace --no-default-gr --extra semantic-router \ --extra saml \ --extra bedrock-realtime \ - --extra mongodb \ --python python3.13 # Copy full source tree @@ -88,7 +87,6 @@ RUN uv sync --frozen --no-default-groups --no-editable \ --extra semantic-router \ --extra saml \ --extra bedrock-realtime \ - --extra mongodb \ --python python3.13 RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \ diff --git a/docker/Dockerfile.non_root b/docker/Dockerfile.non_root index 358425af901..edf20e8bbff 100644 --- a/docker/Dockerfile.non_root +++ b/docker/Dockerfile.non_root @@ -71,7 +71,6 @@ RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \ --extra semantic-router \ --extra saml \ --extra bedrock-realtime \ - --extra mongodb \ --python python3.13 # Copy full source tree @@ -100,7 +99,6 @@ RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \ --extra semantic-router \ --extra saml \ --extra bedrock-realtime \ - --extra mongodb \ --python python3.13 \ --no-sources-package litellm-proxy-extras; \ else \ @@ -111,7 +109,6 @@ RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \ --extra semantic-router \ --extra saml \ --extra bedrock-realtime \ - --extra mongodb \ --python python3.13; \ fi diff --git a/gateway/Dockerfile b/gateway/Dockerfile index e42e488d57f..308d70a6b26 100644 --- a/gateway/Dockerfile +++ b/gateway/Dockerfile @@ -47,7 +47,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \ --extra extra_proxy \ --extra semantic-router \ --extra bedrock-realtime \ - --extra mongodb \ --python python3.13 # Stage 2 — copy source and install the project + workspace members. @@ -60,7 +59,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \ --extra extra_proxy \ --extra semantic-router \ --extra bedrock-realtime \ - --extra mongodb \ --python python3.13 RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \ diff --git a/litellm/llms/base_llm/vector_store/transformation.py b/litellm/llms/base_llm/vector_store/transformation.py index c8d2b7fe522..07b60cb4b72 100644 --- a/litellm/llms/base_llm/vector_store/transformation.py +++ b/litellm/llms/base_llm/vector_store/transformation.py @@ -121,6 +121,9 @@ class RouterVectorStoreEmbeddingExecutor: class BaseVectorStoreConfig: + def validate_create_vector_store(self) -> None: + return None + def get_supported_openai_params(self, model: str) -> list[VECTOR_STORE_OPENAI_PARAMS]: return [] diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 2f561809940..1dc4b198890 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -9814,7 +9814,7 @@ class BaseLLMHTTPHandler: vector_store_search_optional_params=vector_store_search_optional_params, api_base=api_base, litellm_logging_obj=logging_obj, - litellm_params=dict(litellm_params), + litellm_params={**dict(litellm_params), "timeout": timeout}, extra_body=extra_body, embedding_executor=embedding_executor, ) @@ -9859,6 +9859,10 @@ class BaseLLMHTTPHandler: data=request_data, timeout=timeout, ) + except httpx.TimeoutException: + raise vector_store_provider_config.get_error_class( + error_message="Vector store search exceeded the caller timeout.", status_code=408, headers={} + ) from None except Exception as e: raise self._handle_error(e=e, provider_config=vector_store_provider_config) @@ -9943,7 +9947,7 @@ class BaseLLMHTTPHandler: vector_store_search_optional_params=vector_store_search_optional_params, api_base=api_base, litellm_logging_obj=logging_obj, - litellm_params=dict(litellm_params), + litellm_params={**dict(litellm_params), "timeout": timeout}, extra_body=extra_body, embedding_executor=embedding_executor, ) @@ -9988,7 +9992,12 @@ class BaseLLMHTTPHandler: url=url, headers=headers, data=request_data, + timeout=timeout, ) + except httpx.TimeoutException: + raise vector_store_provider_config.get_error_class( + error_message="Vector store search exceeded the caller timeout.", status_code=408, headers={} + ) from None except Exception as e: raise self._handle_error(e=e, provider_config=vector_store_provider_config) @@ -10018,6 +10027,8 @@ class BaseLLMHTTPHandler: else: async_httpx_client = client + vector_store_provider_config.validate_create_vector_store() + headers: Final = vector_store_provider_config.validate_environment( headers=extra_headers or {}, litellm_params=litellm_params ) @@ -10088,6 +10099,8 @@ class BaseLLMHTTPHandler: else: sync_httpx_client = client + vector_store_provider_config.validate_create_vector_store() + headers: Final = vector_store_provider_config.validate_environment( headers=extra_headers or {}, litellm_params=litellm_params ) diff --git a/litellm/llms/mongodb/common_utils.py b/litellm/llms/mongodb/common_utils.py deleted file mode 100644 index 02c0b359407..00000000000 --- a/litellm/llms/mongodb/common_utils.py +++ /dev/null @@ -1,303 +0,0 @@ -"""Shared helpers for the MongoDB integrations. pymongo lives in the optional ``mongodb`` extra, -so every import of it is deferred to call time.""" - -import asyncio -import threading -import weakref -from asyncio import AbstractEventLoop -from collections import OrderedDict -from collections.abc import Callable, Mapping -from dataclasses import dataclass -from types import MappingProxyType -from typing import TYPE_CHECKING, Final, TypeAlias, TypeVar - -from litellm.exceptions import BadRequestError, ServiceUnavailableError, Timeout - -if TYPE_CHECKING: - from pymongo import AsyncMongoClient, MongoClient - -PYMONGO_INSTALL_HINT: Final = ( - "The MongoDB vector store requires the 'pymongo' package. " - "Run 'pip install litellm[mongodb]' (or 'pip install pymongo') to install it." -) - -MONGODB_PROVIDER: Final = "mongodb" - - -def config_error(message: str) -> BadRequestError: - """400 rather than the 500 a bare ValueError becomes once litellm.exception_type wraps it.""" - return BadRequestError(message=message, model=None, llm_provider=MONGODB_PROVIDER) - - -def timeout_error(message: str) -> Timeout: - return Timeout(message=message, model=None, llm_provider=MONGODB_PROVIDER) - - -def unavailable_error(message: str) -> ServiceUnavailableError: - """litellm only retries 408, 409, 429 and 5xx, so a 400 here would make a failover permanent.""" - return ServiceUnavailableError(message=message, model=None, llm_provider=MONGODB_PROVIDER) - - -DEFAULT_CONNECT_TIMEOUT_MS: Final = 10_000 -DEFAULT_SOCKET_TIMEOUT_MS: Final = 30_000 -DEFAULT_SERVER_SELECTION_TIMEOUT_MS: Final = 10_000 - -_MAX_CACHED_CLIENTS: Final = 32 - -_APP_NAME: Final = "litellm" - - -@dataclass(frozen=True, slots=True) -class MongoClientKey: - connection_string: str - connect_timeout_ms: int - socket_timeout_ms: int - server_selection_timeout_ms: int - - -SyncClientFactory: TypeAlias = Callable[..., "MongoClient"] -AsyncClientFactory: TypeAlias = Callable[..., "AsyncMongoClient"] - -_K = TypeVar("_K") -_V = TypeVar("_V") - -_AsyncClientCacheKey: TypeAlias = tuple[MongoClientKey, int] -# CPython recycles id() aggressively, so the id alone would hand a new loop a closed loop's client -_AsyncClientEntry: TypeAlias = tuple["weakref.ref[AbstractEventLoop]", "AsyncMongoClient"] - -_SyncClientCache: TypeAlias = "OrderedDict[MongoClientKey, MongoClient]" -_AsyncClientCache: TypeAlias = "OrderedDict[_AsyncClientCacheKey, _AsyncClientEntry]" - -_sync_clients: Final[_SyncClientCache] = OrderedDict() # mutable-ok: process-level client cache -_async_clients: Final[_AsyncClientCache] = OrderedDict() # mutable-ok: same cache, per loop -# async searches reach the sync client through executor threads, so both caches are shared state -_cache_lock: Final = threading.Lock() - - -def _store_bounded(cache: "OrderedDict[_K, _V]", cache_key: "_K", value: "_V") -> None: - """Eviction only drops this cache's reference; an in-flight search keeps its client alive.""" - with _cache_lock: - cache[cache_key] = value # mutable-ok: an LRU cache is mutable state by definition - cache.move_to_end(cache_key) - while len(cache) > _MAX_CACHED_CLIENTS: - cache.popitem(last=False) - - -def _mark_used(cache: "OrderedDict[_K, _V]", cache_key: "_K") -> None: - with _cache_lock: - if cache_key in cache: - cache.move_to_end(cache_key) - - -def import_sync_mongo_client() -> "type[MongoClient]": - try: - from pymongo import MongoClient as SyncMongoClient - except ImportError as e: - raise config_error(PYMONGO_INSTALL_HINT) from e - return SyncMongoClient - - -def import_async_mongo_client() -> "type[AsyncMongoClient]": - try: - from pymongo import AsyncMongoClient as AsyncMongoClientClass - except ImportError as e: - raise config_error(PYMONGO_INSTALL_HINT) from e - return AsyncMongoClientClass - - -def _client_kwargs(key: MongoClientKey) -> Mapping[str, object]: - return MappingProxyType( - { - "connectTimeoutMS": key.connect_timeout_ms, - "socketTimeoutMS": key.socket_timeout_ms, - "serverSelectionTimeoutMS": key.server_selection_timeout_ms, - "appname": _APP_NAME, - } - ) - - -def get_sync_client(key: MongoClientKey, client_class: SyncClientFactory | None = None) -> "MongoClient": - cached: Final = _sync_clients.get(key) - if cached is not None: - _mark_used(_sync_clients, key) - return cached - build: Final = client_class if client_class is not None else import_sync_mongo_client() - client: Final = build(key.connection_string, **_client_kwargs(key)) - _store_bounded(_sync_clients, key, client) - return client - - -def _purge_dead_loops() -> None: - """A cached client holds its loop alive, so a closed loop's entry would pin that client and its - sockets for the life of the process.""" - with _cache_lock: - for stale in tuple( - cache_key - for cache_key, (loop_ref, _) in _async_clients.items() - if (cached_loop := loop_ref()) is None or cached_loop.is_closed() - ): - del _async_clients[stale] - - -def get_async_client(key: MongoClientKey, client_class: AsyncClientFactory | None = None) -> "AsyncMongoClient": - """Async clients bind to the loop that created them, so the cache is keyed per loop.""" - loop: Final = asyncio.get_running_loop() - loop_key: Final = (key, id(loop)) - cached: Final = _async_clients.get(loop_key) - if cached is not None and cached[0]() is loop: - _mark_used(_async_clients, loop_key) - return cached[1] - _purge_dead_loops() - build: Final = client_class if client_class is not None else import_async_mongo_client() - client: Final = build(key.connection_string, **_client_kwargs(key)) - _store_bounded(_async_clients, loop_key, (weakref.ref(loop), client)) - return client - - -def reset_client_cache() -> None: - with _cache_lock: - _sync_clients.clear() - _async_clients.clear() - - -_AUTHENTICATION_FAILED_CODE: Final = 18 -_UNAUTHORIZED_CODE: Final = 13 -# Atlas reports a rejected user as code 8000 "AtlasError" where a self-managed mongod reports 18 -_AUTHENTICATION_MESSAGE_MARKERS: Final = ("bad auth", "authentication failed", "not authorized") -_RESOLUTION_TIMEOUT_MARKERS: Final = ("resolution lifetime expired", "dns operation timed out") -_UNKNOWN_HOSTNAME_MARKERS: Final = ("dns query name does not exist", "name or service not known") -_CREDENTIAL_ESCAPING_MARKERS: Final = ("must be escaped according to rfc 3986", "bad database name") - - -def _index_hint(index_name: str, database: str, collection: str) -> str: - return ( - f"No queryable MongoDB Vector Search index named '{index_name}' was found on " - f"'{database}.{collection}'. Confirm the index exists on that exact collection, that its " - "status is READY rather than still building, and that the vector store id matches the index name." - ) - - -def missing_index_error(index_name: str, database: str, collection: str) -> BadRequestError: - """$vectorSearch against a missing index, database or collection returns zero documents rather - than failing, so an empty result set is checked against the catalogue and reported as this.""" - return config_error( - f"{_index_hint(index_name, database, collection)} A vector search against a database, " - "collection or index that does not exist returns no results rather than an error, so this " - "was reported as an empty result set by MongoDB." - ) - - -def index_not_ready_error(index_name: str, database: str, collection: str, status: str) -> BadRequestError: - return config_error( - f"The MongoDB Vector Search index '{index_name}' on '{database}.{collection}' is not queryable " - f"yet; its status is {status}. Searches against it return no results until the build finishes." - ) - - -def translate_mongo_error(error: Exception, index_name: str, database: str, collection: str) -> Exception: - """Returns the exception to raise, so callers keep the driver error as ``__cause__``.""" - try: - from pymongo.errors import ( - ConfigurationError, - ConnectionFailure, - ExecutionTimeout, - InvalidOperation, - NetworkTimeout, - OperationFailure, - ServerSelectionTimeoutError, - ) - except ImportError: - return error - - if isinstance(error, ServerSelectionTimeoutError): - return timeout_error( - "Could not reach the MongoDB deployment before the timeout. On Atlas this is usually the " - "project's IP access list not containing this host, or a paused cluster. On a self-managed " - "deployment it is usually the host or port in the URI, or a firewall between this process " - f"and mongod. Either way it can also be an unresolvable hostname. Driver detail: {error}" - ) - # ExecutionTimeout subclasses OperationFailure, so it has to be matched before it - if isinstance(error, (NetworkTimeout, ExecutionTimeout)): - return timeout_error( - f"The MongoDB vector search against '{database}.{collection}' timed out before returning. " - f"Driver detail: {error}" - ) - # ServerSelectionTimeoutError and NetworkTimeout also subclass ConnectionFailure, so this only - # sees what those branches left - if isinstance(error, ConnectionFailure): - return unavailable_error( - f"The connection to '{database}.{collection}' was dropped or refused. That is usually a " - "replica set failover or a restarted node, so the search is worth retrying. If it keeps " - "happening: on Atlas the usual cause is a connection string with no username and password, " - "or a TLS failure, so confirm the URI is the one Atlas shows under Connect, Drivers; on a " - "self-managed deployment, check that mongod is listening on the host and port in the URI. " - f"Driver detail: {error}" - ) - if isinstance(error, OperationFailure): - code: Final = error.code - detail: Final = str(error).lower() - if code in (_AUTHENTICATION_FAILED_CODE, _UNAUTHORIZED_CODE) or any( - marker in detail for marker in _AUTHENTICATION_MESSAGE_MARKERS - ): - return config_error( - "MongoDB rejected the credentials in mongodb_connection_string, or the database user " - f"lacks read access to '{database}.{collection}'. Driver detail: {error.details}" - ) - if "dimension" in detail: - return config_error( - "The query embedding does not match the vector dimensions the index was built for. " - "litellm_embedding_model must be the same model that produced the stored vectors. " - f"Driver detail: {error}" - ) - if "is not indexed as vector" in detail: - return config_error( - "mongodb_embedding_field names a field the MongoDB Vector Search index does not cover. " - f"It must match the 'path' the index '{index_name}' was created on. Driver detail: {error}" - ) - if "index" in detail and ("not found" in detail or "does not exist" in detail or "unknown" in detail): - return config_error(f"{_index_hint(index_name, database, collection)} Driver detail: {error}") - return config_error( - f"MongoDB rejected the vector search against '{database}.{collection}' using index " - f"'{index_name}'. Driver detail: {error}" - ) - if isinstance(error, ConfigurationError): - configuration_detail: Final = str(error).lower() - if any(marker in configuration_detail for marker in _RESOLUTION_TIMEOUT_MARKERS): - return timeout_error( - "The DNS lookup for the cluster in mongodb_connection_string did not finish in time. " - "A mongodb+srv:// URI needs an SRV lookup before any connection is attempted, so this " - f"is DNS or the configured timeout, not MongoDB. Driver detail: {error}" - ) - if any(marker in configuration_detail for marker in _UNKNOWN_HOSTNAME_MARKERS): - return config_error( - "The hostname in mongodb_connection_string does not exist in DNS. On Atlas, check the " - "cluster name against the URI shown under Connect, Drivers. On a self-managed deployment, " - f"check that the hostname resolves from this process. Driver detail: {error}" - ) - if any(marker in configuration_detail for marker in _CREDENTIAL_ESCAPING_MARKERS): - return config_error( - "mongodb_connection_string could not be parsed. A username or password containing " - "'@', '/', ':' or '%' has to be percent-encoded per RFC 3986, so 'p@ss/word' becomes " - "'p%40ss%2Fword'. If the credentials are already encoded, check the database name in " - f"the URI path instead. Driver detail: {error}" - ) - return config_error( - f"mongodb_connection_string is not a usable MongoDB connection string. Driver detail: {error}" - ) - if isinstance(error, InvalidOperation): - return config_error(f"The MongoDB client was already closed or is unusable. Driver detail: {error}") - # An unreadable tlsCAFile or tlsCertificateKeyFile raises OSError, not a PyMongoError - if isinstance(error, OSError) and error.filename: - return config_error( - f"'{error.filename}', named by a TLS option in mongodb_connection_string, could not be read. " - "Check that tlsCAFile and tlsCertificateKeyFile point at files this process can open; inside " - f"a container that is the path in the container, not on the host. Driver detail: {error}" - ) - # pymongo raises a plain ValueError, not a PyMongoError, for an unusable port - if isinstance(error, ValueError): - return config_error( - "The host and port in mongodb_connection_string could not be parsed. If the port is a " - "number between 0 and 65535, the cause is usually an unescaped ':' in the password, which " - f"has to be percent-encoded per RFC 3986 as '%3A'. Driver detail: {error}" - ) - return error diff --git a/litellm/llms/mongodb/vector_stores/transformation.py b/litellm/llms/mongodb/vector_stores/transformation.py index 3382c931c96..92965ab745b 100644 --- a/litellm/llms/mongodb/vector_stores/transformation.py +++ b/litellm/llms/mongodb/vector_stores/transformation.py @@ -1,37 +1,28 @@ -"""MongoDB Vector Search has no HTTP query API, so this is a direct provider that runs the -``$vectorSearch`` aggregation through pymongo. ``vector_store_id`` is the search index name.""" - -from collections.abc import Callable, Mapping, Sequence +from collections.abc import Mapping, Sequence +from math import isfinite from types import MappingProxyType -from typing import TYPE_CHECKING, Final, NoReturn +from typing import TYPE_CHECKING, Final, Literal, NoReturn +from urllib.parse import quote, urlsplit import httpx -from pydantic import BaseModel, ConfigDict +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError +from litellm.exceptions import AuthenticationError, BadRequestError, ServiceUnavailableError, Timeout +from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.vector_store.transformation import ( - BaseDirectVectorStoreConfig, + BaseQueryEmbeddingVectorStoreConfig, LiteLLMVectorStoreEmbeddingExecutor, VectorStoreEmbeddingExecutor, ) -from litellm.llms.mongodb.common_utils import ( - DEFAULT_CONNECT_TIMEOUT_MS, - DEFAULT_SERVER_SELECTION_TIMEOUT_MS, - DEFAULT_SOCKET_TIMEOUT_MS, - MongoClientKey, - config_error, - get_async_client, - get_sync_client, - index_not_ready_error, - missing_index_error, - translate_mongo_error, -) +from litellm.secret_managers.main import get_secret_str +from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import EmbeddingResponse from litellm.types.vector_stores import ( + BaseVectorStoreAuthCredentials, VectorStoreCreateOptionalRequestParams, - VectorStoreResultContent, + VectorStoreIndexEndpoints, VectorStoreSearchOptionalRequestParams, VectorStoreSearchResponse, - VectorStoreSearchResult, ) if TYPE_CHECKING: @@ -39,26 +30,45 @@ if TYPE_CHECKING: DEFAULT_EMBEDDING_FIELD_NAME: Final = "embedding" DEFAULT_TEXT_FIELD_NAME: Final = "text" -SCORE_FIELD_NAME: Final = "score" - DEFAULT_MAX_NUM_RESULTS: Final = 10 MIN_MAX_NUM_RESULTS: Final = 1 MAX_MAX_NUM_RESULTS: Final = 50 - NUM_CANDIDATES_MULTIPLIER: Final = 10 MIN_NUM_CANDIDATES: Final = 100 MAX_NUM_CANDIDATES: Final = 10_000 - MAX_QUERY_CHARACTERS: Final = 32_000 - _EMPTY_EMBEDDING_CONFIG: Final = MappingProxyType({}) - _SEARCH_ONLY_MESSAGE: Final = ( "MongoDB vector store is search-only. Create the collection and its MongoDB Vector Search " "index in MongoDB directly, then register it here by index name." ) +def config_error(message: str) -> BadRequestError: + return BadRequestError(message=message, model=None, llm_provider="mongodb") + + +class _Content(BaseModel): + model_config = ConfigDict(frozen=True, strict=True) + type: Literal["text"] + text: str + + +class _Result(BaseModel): + model_config = ConfigDict(frozen=True, strict=True, allow_inf_nan=False) + score: float | None + content: list[_Content] + file_id: str | None + filename: str | None + + +class _SearchResponse(BaseModel): + model_config = ConfigDict(frozen=True, strict=True) + object: Literal["vector_store.search_results.page"] + search_query: str + data: list[_Result] + + class _MongoDBSearchParams(BaseModel): """Typed view over the vector store's litellm_params; unrelated keys are ignored.""" @@ -66,7 +76,6 @@ class _MongoDBSearchParams(BaseModel): litellm_embedding_model: str | None = None litellm_embedding_config: Mapping[str, object] | None = None - mongodb_connection_string: str | None = None mongodb_database: str | None = None mongodb_collection: str | None = None mongodb_text_field: str | None = None @@ -91,21 +100,6 @@ class _MongoDBSearchParams(BaseModel): ) return self.litellm_embedding_model - def require_connection_string(self) -> str: - if not self.mongodb_connection_string: - raise config_error( - "mongodb_connection_string is required in litellm_params for the MongoDB vector store. " - "Example: mongodb+srv://:@.mongodb.net for Atlas, or " - "mongodb://:@:27017 for a self-managed deployment" - ) - scheme: Final = self.mongodb_connection_string.split("://", 1)[0].lower() - if scheme not in ("mongodb", "mongodb+srv"): - raise config_error( - "mongodb_connection_string must start with 'mongodb://' or 'mongodb+srv://', " - f"got '{self.mongodb_connection_string.split('://', 1)[0]}://'" - ) - return self.mongodb_connection_string - def require_database(self) -> str: if not self.mongodb_database: raise config_error( @@ -127,30 +121,28 @@ _MONGODB_PARAM_PREFIX: Final = "mongodb_" _KNOWN_MONGODB_PARAMS: Final = frozenset( name for name in _MongoDBSearchParams.model_fields if name.startswith(_MONGODB_PARAM_PREFIX) ) +_RESPONSE_ADAPTER: Final = TypeAdapter(VectorStoreSearchResponse) -class MongoDBVectorStoreConfig(BaseDirectVectorStoreConfig): - def __init__( - self, - embedding_executor: VectorStoreEmbeddingExecutor | None = None, - sync_client_factory: Callable[[MongoClientKey], object] | None = None, - async_client_factory: Callable[[MongoClientKey], object] | None = None, - ) -> None: - super().__init__() - self.embedding_executor: Final[VectorStoreEmbeddingExecutor] = ( - embedding_executor if embedding_executor is not None else LiteLLMVectorStoreEmbeddingExecutor() - ) - self.sync_client_factory: Final[Callable[[MongoClientKey], object]] = ( - sync_client_factory if sync_client_factory is not None else get_sync_client - ) - self.async_client_factory: Final[Callable[[MongoClientKey], object]] = ( - async_client_factory if async_client_factory is not None else get_async_client - ) +class MongoDBVectorStoreConfig(BaseQueryEmbeddingVectorStoreConfig): + def __init__(self, embedding_executor: VectorStoreEmbeddingExecutor | None = None) -> None: + self.embedding_executor: Final = embedding_executor or LiteLLMVectorStoreEmbeddingExecutor() + + def get_auth_credentials(self, litellm_params: Mapping[str, object]) -> BaseVectorStoreAuthCredentials: + return BaseVectorStoreAuthCredentials() + + def get_vector_store_endpoints_by_type(self) -> VectorStoreIndexEndpoints: + return VectorStoreIndexEndpoints(read=[], write=[]) # mutable-ok: the TypedDict declares list fields @staticmethod def _reject_unknown_params(litellm_params: Mapping[str, object]) -> None: """Without this a mistyped mongodb_collection reads as 'mongodb_collection is required', naming a key the reader can see they have set.""" + if litellm_params.get("mongodb_connection_string") is not None: + raise config_error( + "MongoDB vector stores now use the BETA sidecar. Move mongodb_connection_string to " + "MONGODB_CONNECTION_STRING in the sidecar, remove it from LiteLLM, and configure api_base and api_key." + ) unknown: Final = sorted( key for key in litellm_params if key.startswith(_MONGODB_PARAM_PREFIX) and key not in _KNOWN_MONGODB_PARAMS ) @@ -191,239 +183,182 @@ class MongoDBVectorStoreConfig(BaseDirectVectorStoreConfig): return configured return min(max(limit * NUM_CANDIDATES_MULTIPLIER, MIN_NUM_CANDIDATES), MAX_NUM_CANDIDATES) - @staticmethod - def _timeout_ms(timeout: float | httpx.Timeout | None) -> tuple[int, int]: - """The connect and socket budgets pymongo is built with, in that order.""" - if isinstance(timeout, httpx.Timeout): - return ( - int((timeout.connect or DEFAULT_CONNECT_TIMEOUT_MS / 1000) * 1000), - int((timeout.read or DEFAULT_SOCKET_TIMEOUT_MS / 1000) * 1000), + def validate_environment( + self, headers: dict[str, object], litellm_params: GenericLiteLLMParams | None + ) -> dict[str, object]: + if litellm_params is None: + raise config_error("Configure api_base and api_key for the MongoDB BETA sidecar.") + self._reject_unknown_params(dict(litellm_params)) + api_key: Final = litellm_params.api_key or get_secret_str("MONGODB_SIDECAR_API_KEY") + if not api_key: + raise config_error("MongoDB sidecar api_key is required. Set api_key or MONGODB_SIDECAR_API_KEY.") + return {**headers, "Authorization": f"Bearer {api_key}", "Content-Type": "application/json"} + + def get_complete_url(self, api_base: str | None, litellm_params: dict[str, object]) -> str: + if not api_base: + raise config_error("MongoDB sidecar api_base is required, for example http://mongodb-sidecar:8080.") + try: + parsed: Final = urlsplit(api_base) + valid: Final = parsed.scheme in ("http", "https") and bool(parsed.hostname) and parsed.port != 0 + except ValueError: + raise config_error("MongoDB sidecar api_base must be a valid HTTP or HTTPS URL.") from None + if not valid or parsed.username or parsed.password or parsed.query or parsed.fragment: + raise config_error( + "MongoDB sidecar api_base must be an HTTP or HTTPS URL without credentials, query, or fragment." ) - if timeout is None: - return DEFAULT_CONNECT_TIMEOUT_MS, DEFAULT_SOCKET_TIMEOUT_MS - return min(int(float(timeout) * 1000), DEFAULT_CONNECT_TIMEOUT_MS), int(float(timeout) * 1000) + return api_base.rstrip("/") + + @staticmethod + def _timeout_ms(value: object) -> int: + seconds: Final = value.read if isinstance(value, httpx.Timeout) else value + if seconds is None: + return 30_000 + if not isinstance(seconds, (int, float)) or not isfinite(seconds) or seconds <= 0: + raise config_error("MongoDB search timeout must be a positive finite number.") + return max(1, min(int(seconds * 1000), 30_000)) @classmethod - def _client_key(cls, params: _MongoDBSearchParams, timeout: float | httpx.Timeout | None) -> MongoClientKey: - connect_ms, socket_ms = cls._timeout_ms(timeout) - return MongoClientKey( - connection_string=params.require_connection_string(), - connect_timeout_ms=connect_ms, - socket_timeout_ms=socket_ms, - server_selection_timeout_ms=min(connect_ms, DEFAULT_SERVER_SELECTION_TIMEOUT_MS), - ) + def _params( + cls, + litellm_params: Mapping[str, object], + optional_params: VectorStoreSearchOptionalRequestParams, + extra_body: Mapping[str, object] | None, + ) -> _MongoDBSearchParams: + cls._reject_unknown_params(litellm_params) + if extra_body: + raise config_error("MongoDB vector store does not support extra_body overrides.") + for unsupported in ("filters", "ranking_options", "rewrite_query"): + if optional_params.get(unsupported) is not None: + raise config_error(f"MongoDB vector store does not support the {unsupported} parameter.") + try: + params: Final = _MongoDBSearchParams.model_validate(litellm_params) + except ValidationError: + raise config_error( + "Invalid MongoDB vector-store configuration. Check the database, collection, fields, and candidate count." + ) from None + params.require_database() + params.require_collection() + params.require_embedding_model() + cls._num_candidates(cls._limit(optional_params), params.mongodb_num_candidates) + cls._timeout_ms(litellm_params.get("timeout")) + return params @classmethod - def _pipeline( + def _request( cls, vector_store_id: str, - query_vector: Sequence[float], + query_text: str, params: _MongoDBSearchParams, - vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, - ) -> Sequence[Mapping[str, object]]: - if vector_store_search_optional_params.get("filters") is not None: + optional_params: VectorStoreSearchOptionalRequestParams, + api_base: str, + embedding_response: EmbeddingResponse, + timeout: object, + ) -> tuple[str, dict[str, object]]: + if not embedding_response.data: raise config_error( - "MongoDB vector store does not support the filters parameter yet. " - "Restrict the collection or the MongoDB Vector Search index definition instead." + "The embedding model returned no embedding for the search query. Check litellm_embedding_model." ) - if vector_store_search_optional_params.get("ranking_options") is not None: - raise config_error( - "MongoDB vector store does not support the ranking_options parameter yet. " - "Every result already carries the vectorSearchScore, so filter or re-rank " - "on that rather than having the threshold silently ignored." - ) - if vector_store_search_optional_params.get("rewrite_query") is not None: - raise config_error( - "MongoDB vector store does not support the rewrite_query parameter. The query is " - "embedded exactly as sent; rewrite it before calling if you need that." - ) - limit: Final = cls._limit(vector_store_search_optional_params) - search: Final = MappingProxyType( - { - "index": vector_store_id, - "path": params.embedding_field, - "queryVector": tuple(query_vector), - "numCandidates": cls._num_candidates(limit, params.mongodb_num_candidates), - "limit": limit, - } - ) - projection: Final = MappingProxyType( - {params.text_field: 1, SCORE_FIELD_NAME: MappingProxyType({"$meta": "vectorSearchScore"})} - ) - return [ # mutable-ok: pymongo rejects any non-list pipeline in common.validate_list - MappingProxyType({"$vectorSearch": search}), - MappingProxyType({"$project": projection}), - ] + vector: Final = embedding_response.data[0]["embedding"] + if not vector or any(not isinstance(value, (float, int)) or not isfinite(value) for value in vector): + raise config_error("The embedding model must return a non-empty, finite query vector.") + limit: Final = cls._limit(optional_params) + return f"{api_base}/v1/vector_stores/{quote(vector_store_id, safe='')}/search", { + "query": query_text, + "query_vector": tuple(vector), + "mongodb_database": params.require_database(), + "mongodb_collection": params.require_collection(), + "mongodb_embedding_field": params.embedding_field, + "mongodb_text_field": params.text_field, + "mongodb_num_candidates": cls._num_candidates(limit, params.mongodb_num_candidates), + "max_num_results": limit, + "timeout_ms": cls._timeout_ms(timeout), + } - @classmethod - def _field_value(cls, document: Mapping[str, object], dotted_path: str) -> str | None: - """None means absent, which is what separates a mistyped field from genuinely empty text.""" - head, _, rest = dotted_path.partition(".") - if head not in document: - return None - value: Final = document[head] - if not rest: - return None if value is None else str(value) - return cls._field_value(value, rest) if isinstance(value, Mapping) else None - - @classmethod - def _to_result(cls, document: Mapping[str, object], text_field: str) -> VectorStoreSearchResult: - document_id: Final = document.get("_id") - identifier: Final = None if document_id is None else str(document_id) - content: Final = [ # mutable-ok: VectorStoreSearchResult declares a list of content parts - VectorStoreResultContent(text=cls._field_value(document, text_field) or "", type="text") - ] - raw_score: Final = document.get(SCORE_FIELD_NAME) - return VectorStoreSearchResult( - score=float(raw_score) if isinstance(raw_score, (int, float)) else None, - content=content, - file_id=identifier, - filename=identifier, - ) - - @classmethod - def _raise_for_missing_text_field( - cls, documents: Sequence[Mapping[str, object]], text_field: str, database: str, collection: str - ) -> None: - """$vectorSearch matches documents carrying no text, so a mistyped mongodb_text_field - returns well-scored results with empty content instead of failing.""" - if documents and all(cls._field_value(document, text_field) is None for document in documents): - raise config_error( - f"None of the {len(documents)} matched documents in '{database}.{collection}' has a " - f"'{text_field}' field, so every result would carry empty text. Set mongodb_text_field " - "to the field holding the readable text; it accepts a dotted path such as metadata.body." - ) - - @classmethod - def _to_response( - cls, documents: Sequence[Mapping[str, object]], query_text: str, text_field: str - ) -> VectorStoreSearchResponse: - return VectorStoreSearchResponse( - object="vector_store.search_results.page", - search_query=query_text, - data=[ # mutable-ok: VectorStoreSearchResponse declares data as a list - cls._to_result(document, text_field) for document in documents - ], - ) - - @staticmethod - def _raise_for_unusable_index( - catalogue: Sequence[Mapping[str, object]], index_name: str, database: str, collection: str - ) -> None: - """mongod returns zero documents both for a query that matched nothing and for a missing - database, collection or index, so the catalogue decides which one happened.""" - if not catalogue: - raise missing_index_error(index_name, database, collection) - entry: Final = catalogue[0] - if not entry.get("queryable"): - raise index_not_ready_error(index_name, database, collection, str(entry.get("status") or "unknown")) - - @staticmethod - def _embedding_vector(embedding_response: EmbeddingResponse) -> Sequence[float]: - data: Final = embedding_response.data - if not data: - raise config_error( - "The embedding model returned no embedding for the search query, so there is nothing " - "to search MongoDB with. Check the embedding deployment named by litellm_embedding_model." - ) - return data[0]["embedding"] - - def execute_search_vector_store_request( + def transform_search_vector_store_request( self, vector_store_id: str, query: str | Sequence[str], vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, + api_base: str, litellm_logging_obj: "LiteLLMLoggingObj", litellm_params: Mapping[str, object], + extra_body: Mapping[str, object] | None = None, embedding_executor: VectorStoreEmbeddingExecutor | None = None, - timeout: float | httpx.Timeout | None = None, - ) -> VectorStoreSearchResponse: - self._reject_unknown_params(litellm_params) - params: Final = _MongoDBSearchParams.model_validate(litellm_params) + ) -> tuple[str, dict[str, object]]: + params: Final = self._params(litellm_params, vector_store_search_optional_params, extra_body) query_text: Final = self._query_text(query) - key: Final = self._client_key(params, timeout) - database: Final = params.require_database() - collection: Final = params.require_collection() - - embedding_response: Final = (embedding_executor or self.embedding_executor).embed( - params.require_embedding_model(), + response: Final = (embedding_executor or self.embedding_executor).embed( + params.require_embedding_model(), query_text, params.litellm_embedding_config or _EMPTY_EMBEDDING_CONFIG + ) + return self._request( + vector_store_id, query_text, - params.litellm_embedding_config or _EMPTY_EMBEDDING_CONFIG, - ) - pipeline: Final = self._pipeline( - vector_store_id, self._embedding_vector(embedding_response), params, vector_store_search_optional_params + params, + vector_store_search_optional_params, + api_base, + response, + litellm_params.get("timeout"), ) - try: - client: Final = self.sync_client_factory(key) - target: Final = client[database][collection] # pyright: ignore[reportIndexIssue] # factory is typed as returning object so injected doubles are accepted - documents: Final = tuple(target.aggregate(pipeline)) - except Exception as e: - raise translate_mongo_error(e, index_name=vector_store_id, database=database, collection=collection) from e - if not documents: - try: - catalogue: Final = tuple(target.list_search_indexes(vector_store_id)) - except Exception as e: - raise translate_mongo_error( - e, index_name=vector_store_id, database=database, collection=collection - ) from e - self._raise_for_unusable_index(catalogue, vector_store_id, database, collection) - self._raise_for_missing_text_field(documents, params.text_field, database, collection) - return self._to_response(documents, query_text, params.text_field) - - async def aexecute_search_vector_store_request( + async def atransform_search_vector_store_request( self, vector_store_id: str, query: str | Sequence[str], vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, + api_base: str, litellm_logging_obj: "LiteLLMLoggingObj", litellm_params: Mapping[str, object], + extra_body: Mapping[str, object] | None = None, embedding_executor: VectorStoreEmbeddingExecutor | None = None, - timeout: float | httpx.Timeout | None = None, - ) -> VectorStoreSearchResponse: - self._reject_unknown_params(litellm_params) - params: Final = _MongoDBSearchParams.model_validate(litellm_params) + ) -> tuple[str, dict[str, object]]: + params: Final = self._params(litellm_params, vector_store_search_optional_params, extra_body) query_text: Final = self._query_text(query) - key: Final = self._client_key(params, timeout) - database: Final = params.require_database() - collection: Final = params.require_collection() - - embedding_response: Final = await (embedding_executor or self.embedding_executor).aembed( - params.require_embedding_model(), + response: Final = await (embedding_executor or self.embedding_executor).aembed( + params.require_embedding_model(), query_text, params.litellm_embedding_config or _EMPTY_EMBEDDING_CONFIG + ) + return self._request( + vector_store_id, query_text, - params.litellm_embedding_config or _EMPTY_EMBEDDING_CONFIG, - ) - pipeline: Final = self._pipeline( - vector_store_id, self._embedding_vector(embedding_response), params, vector_store_search_optional_params + params, + vector_store_search_optional_params, + api_base, + response, + litellm_params.get("timeout"), ) + def transform_search_vector_store_response( + self, response: httpx.Response, litellm_logging_obj: "LiteLLMLoggingObj" + ) -> VectorStoreSearchResponse: try: - client: Final = self.async_client_factory(key) - target: Final = client[database][collection] # pyright: ignore[reportIndexIssue] # factory is typed as returning object so injected doubles are accepted - cursor: Final = await target.aggregate(pipeline) - documents: Final = [ # mutable-ok: an async comprehension cannot build a tuple directly - document async for document in cursor - ] - except Exception as e: - raise translate_mongo_error(e, index_name=vector_store_id, database=database, collection=collection) from e - if not documents: - try: - index_cursor: Final = await target.list_search_indexes(vector_store_id) - catalogue: Final = [ # mutable-ok: an async comprehension cannot build a tuple directly - entry async for entry in index_cursor - ] - except Exception as e: - raise translate_mongo_error( - e, index_name=vector_store_id, database=database, collection=collection - ) from e - self._raise_for_unusable_index(catalogue, vector_store_id, database, collection) - self._raise_for_missing_text_field(documents, params.text_field, database, collection) - return self._to_response(documents, query_text, params.text_field) + validated: Final = _SearchResponse.model_validate_json(response.content) + return _RESPONSE_ADAPTER.validate_python(validated.model_dump()) + except ValidationError: + raise ServiceUnavailableError( + message="MongoDB sidecar returned an invalid search response. Check the sidecar version and deployment.", + model=None, + llm_provider="mongodb", + ) from None + + def get_error_class( + self, error_message: str, status_code: int, headers: dict[str, object] | httpx.Headers + ) -> BaseLLMException: + if status_code == 400: + raise config_error(error_message) + if status_code == 401: + raise AuthenticationError(message="MongoDB sidecar rejected api_key.", model=None, llm_provider="mongodb") + if status_code == 408: + raise Timeout(message=error_message, model=None, llm_provider="mongodb") + raise ServiceUnavailableError( + message="MongoDB sidecar is unavailable. Check its address, health, and logs.", + model=None, + llm_provider="mongodb", + ) + + def validate_create_vector_store(self) -> NoReturn: + raise config_error(_SEARCH_ONLY_MESSAGE) def transform_create_vector_store_request( - self, - vector_store_create_optional_params: VectorStoreCreateOptionalRequestParams, - api_base: str, + self, vector_store_create_optional_params: VectorStoreCreateOptionalRequestParams, api_base: str ) -> NoReturn: raise config_error(_SEARCH_ONLY_MESSAGE) diff --git a/pyproject.toml b/pyproject.toml index af35c77d259..b7eecfb2109 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -114,7 +114,6 @@ caching = ["diskcache>=5.6.3,<6.0"] mcp = ["mcp>=1.28.1,<2.0"] # Driver for the MongoDB Atlas vector store; Atlas Vector Search has no HTTP query API. # The floor is 4.9 because that is the release AsyncMongoClient landed in. -mongodb = ["pymongo>=4.9,<5.0"] # SAML SSO for the admin UI. python3-saml pulls in xmlsec/lxml, whose wheels # bundle the native libxmlsec1/libxml2 libraries, so no system packages are # required. Kept out of the base `proxy` extra so it stays optional. diff --git a/tests/test_litellm/llms/mongodb/vector_stores/test_mongodb_transformation.py b/tests/test_litellm/llms/mongodb/vector_stores/test_mongodb_transformation.py index f5d31c0da54..bca2b544673 100644 --- a/tests/test_litellm/llms/mongodb/vector_stores/test_mongodb_transformation.py +++ b/tests/test_litellm/llms/mongodb/vector_stores/test_mongodb_transformation.py @@ -1,1537 +1,182 @@ -import asyncio -import gc -import sys -import threading -import weakref -from types import SimpleNamespace -from unittest.mock import MagicMock, patch +import json +from collections.abc import Mapping +from typing import Final +from unittest.mock import MagicMock import httpx import pytest import litellm -from litellm.exceptions import BadRequestError, ServiceUnavailableError, Timeout -from litellm.llms.mongodb.common_utils import ( - _MAX_CACHED_CLIENTS, - _async_clients, - _sync_clients, - MongoClientKey, - index_not_ready_error, - missing_index_error, - get_async_client, - get_sync_client, - reset_client_cache, - translate_mongo_error, -) -from litellm.llms.mongodb.vector_stores.transformation import ( - MongoDBVectorStoreConfig, - _MongoDBSearchParams, -) -from litellm.types.utils import LlmProviders -from litellm.utils import ProviderConfigManager +from litellm.llms.custom_httpx.http_handler import HTTPHandler +from litellm.llms.mongodb.vector_stores.transformation import MongoDBVectorStoreConfig +from litellm.types.utils import EmbeddingResponse +from litellm.types.vector_stores import VectorStoreSearchOptionalRequestParams -CONNECTION_STRING = "mongodb+srv://user:pw@cluster.example.mongodb.net" -INDEX = "movies_vector_index" - -BASE_PARAMS = { - "litellm_embedding_model": "openai/text-embedding-ada-002", - "mongodb_connection_string": CONNECTION_STRING, - "mongodb_database": "sample_mflix", - "mongodb_collection": "embedded_movies", +BASE_PARAMS: Final = { + "api_base": "https://sidecar.example/prefix", + "api_key": "test-sidecar-key", + "litellm_embedding_model": "embedding-alias", + "mongodb_database": "policies", + "mongodb_collection": "documents", +} +RESULT: Final = { + "object": "vector_store.search_results.page", + "search_query": "travel policy", + "data": [ + {"score": 0.9, "file_id": "123", "filename": "123", "content": [{"type": "text", "text": "Use code BLUE-42"}]} + ], } -READY_INDEX = [{"name": INDEX, "status": "READY", "queryable": True}] +class RecordingEmbeddingExecutor: + def __init__(self) -> None: + self.call: Final = MagicMock(return_value=EmbeddingResponse(data=[{"embedding": [0.1, 0.2, 0.3]}])) + + def embed(self, model: str, query: str, configuration: Mapping[str, object]) -> EmbeddingResponse: + return self.call(model, query, configuration) + + async def aembed(self, model: str, query: str, configuration: Mapping[str, object]) -> EmbeddingResponse: + return self.call(model, query, configuration) -class RecordingClient: - """Stands in for pymongo's client class so the cache tests inject a fake rather than - patching the importer, and so they can assert what the client was actually built with.""" - - def __init__(self, connection_string, **kwargs): - self.connection_string = connection_string - self.kwargs = kwargs - - -class FakeCollection: - def __init__(self, documents, error=None, search_indexes=None): - self.documents = documents - self.error = error - self.search_indexes = READY_INDEX if search_indexes is None else search_indexes - self.pipeline = None - self.listed_indexes = [] - - def aggregate(self, pipeline): - self.pipeline = pipeline - if self.error is not None: - raise self.error - return iter(self.documents) - - def list_search_indexes(self, name): - self.listed_indexes.append(name) - return iter(self.search_indexes) - - -class FakeAsyncCollection(FakeCollection): - async def aggregate(self, pipeline): - self.pipeline = pipeline - if self.error is not None: - raise self.error - - async def cursor(): - for document in self.documents: - yield document - - return cursor() - - async def list_search_indexes(self, name): - self.listed_indexes.append(name) - - async def cursor(): - for entry in self.search_indexes: - yield entry - - return cursor() - - -class FakeDatabase: - def __init__(self, collection): - self.collection = collection - self.requested_collection = None - - def __getitem__(self, name): - self.requested_collection = name - return self.collection - - -class FakeClient: - def __init__(self, collection): - self.database = FakeDatabase(collection) - self.requested_database = None - - def __getitem__(self, name): - self.requested_database = name - return self.database - - -class FakeEmbeddingExecutor: - def __init__(self, embedding): - self.embedding = embedding - self.captured = None - - def _respond(self, model, query, configuration): - self.captured = SimpleNamespace(model=model, query=query, configuration=configuration) - return SimpleNamespace(data=[{"embedding": self.embedding}] if self.embedding is not None else []) - - def embed(self, model, query, configuration): - return self._respond(model, query, configuration) - - async def aembed(self, model, query, configuration): - return self._respond(model, query, configuration) - - -def _config(documents=(), embedding=(0.1, 0.2, 0.3), error=None, search_indexes=None): - collection = FakeCollection(list(documents), error, search_indexes) - client = FakeClient(collection) - config = MongoDBVectorStoreConfig( - embedding_executor=FakeEmbeddingExecutor(list(embedding) if embedding is not None else None), - sync_client_factory=lambda key: client, - ) - return config, client, collection - - -def _async_config(documents=(), embedding=(0.1, 0.2, 0.3), error=None, search_indexes=None): - collection = FakeAsyncCollection(list(documents), error, search_indexes) - client = FakeClient(collection) - config = MongoDBVectorStoreConfig( - embedding_executor=FakeEmbeddingExecutor(list(embedding) if embedding is not None else None), - async_client_factory=lambda key: client, - ) - return config, client, collection - - -def _search(config, query="a lone astronaut", optional_params=None, litellm_params=None, timeout=None): - return config.execute_search_vector_store_request( - vector_store_id=INDEX, - query=query, - vector_store_search_optional_params=optional_params or {}, - litellm_logging_obj=MagicMock(), - litellm_params={**BASE_PARAMS, **(litellm_params or {})}, - timeout=timeout, - ) - - -async def _asearch(config, query="a lone astronaut", optional_params=None, litellm_params=None): - return await config.aexecute_search_vector_store_request( - vector_store_id=INDEX, - query=query, - vector_store_search_optional_params=optional_params or {}, - litellm_logging_obj=MagicMock(), - litellm_params={**BASE_PARAMS, **(litellm_params or {})}, - ) - - -def _stage(collection, name): - return next(stage[name] for stage in collection.pipeline if name in stage) - - -def test_search_builds_vector_search_stage_against_the_named_index(): - config, client, collection = _config() - - _search(config, optional_params={"max_num_results": 5}) - - assert client.requested_database == "sample_mflix" - assert client.database.requested_collection == "embedded_movies" - assert _stage(collection, "$vectorSearch") == { - "index": INDEX, - "path": "embedding", - "queryVector": (0.1, 0.2, 0.3), - "numCandidates": 100, - "limit": 5, +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("limit,candidates", [(None, 100), (1, 100), (50, 500)]) +@pytest.mark.asyncio +async def test_search_preserves_embedding_and_http_contract( + asynchronous: bool, limit: int | None, candidates: int +) -> None: + executor: Final = RecordingEmbeddingExecutor() + config: Final = MongoDBVectorStoreConfig(executor) + params: Final = { + **BASE_PARAMS, + "mongodb_text_field": "metadata.body", + "mongodb_embedding_field": "stored_vector", + "litellm_embedding_config": {"dimensions": 3}, + "timeout": 0.75, } - - -def test_the_pipeline_reaches_pymongo_as_a_list(): - """pymongo's common.validate_list rejects any other sequence with - 'pipeline must be a list, not ', so the outer container is part of the contract.""" - config, _, collection = _config() - - _search(config) - - assert isinstance(collection.pipeline, list) - - -def test_search_projects_the_text_field_and_the_similarity_score(): - config, _, collection = _config() - - _search(config) - - assert _stage(collection, "$project") == {"text": 1, "score": {"$meta": "vectorSearchScore"}} - - -def test_search_defaults_to_ten_results(): - config, _, collection = _config() - - _search(config) - - assert _stage(collection, "$vectorSearch")["limit"] == 10 - - -def test_search_honors_custom_field_names(): - config, _, collection = _config() - - _search( - config, - litellm_params={"mongodb_embedding_field": "plot_embedding", "mongodb_text_field": "plot"}, - ) - - assert _stage(collection, "$vectorSearch")["path"] == "plot_embedding" - assert _stage(collection, "$project") == {"plot": 1, "score": {"$meta": "vectorSearchScore"}} - - -def test_num_candidates_scales_with_the_requested_limit(): - config, _, collection = _config() - - _search(config, optional_params={"max_num_results": 40}) - - assert _stage(collection, "$vectorSearch")["numCandidates"] == 400 - - -def test_num_candidates_can_be_overridden(): - config, _, collection = _config() - - _search(config, optional_params={"max_num_results": 5}, litellm_params={"mongodb_num_candidates": 250}) - - assert _stage(collection, "$vectorSearch")["numCandidates"] == 250 - - -@pytest.mark.parametrize("configured", [4, 10_001]) -def test_num_candidates_below_the_limit_or_above_the_ceiling_is_rejected(configured): - config, _, _ = _config() - - with pytest.raises(BadRequestError, match="mongodb_num_candidates"): - _search(config, optional_params={"max_num_results": 5}, litellm_params={"mongodb_num_candidates": configured}) - - -def test_list_query_is_joined_into_one_embedding_input(): - config, _, _ = _config() - - _search(config, query=["deep", "space", "rescue"]) - - assert config.embedding_executor.captured.query == "deep space rescue" - - -def test_embedding_config_is_expanded_into_the_embedding_call(): - config, _, _ = _config() - - _search(config, litellm_params={"litellm_embedding_config": {"api_base": "https://example.test", "timeout": 7}}) - - captured = config.embedding_executor.captured - assert captured.configuration == {"api_base": "https://example.test", "timeout": 7} - assert captured.model == "openai/text-embedding-ada-002" - - -def test_response_maps_documents_to_openai_shaped_results(): - documents = [ - {"_id": "abc123", "text": "an astronaut adrift", "score": 0.94}, - {"_id": "def456", "text": "a robot dog", "score": 0.81}, - ] - config, _, _ = _config(documents=documents) - - response = _search(config) - - assert response["object"] == "vector_store.search_results.page" - assert response["search_query"] == "a lone astronaut" - assert [result["score"] for result in response["data"]] == [0.94, 0.81] - assert [result["content"][0]["text"] for result in response["data"]] == ["an astronaut adrift", "a robot dog"] - assert [result["file_id"] for result in response["data"]] == ["abc123", "def456"] - assert [result["filename"] for result in response["data"]] == ["abc123", "def456"] - assert response["data"][0]["content"][0]["type"] == "text" - - -def test_response_reads_a_dotted_text_field_path(): - config, _, _ = _config(documents=[{"_id": 1, "metadata": {"body": "nested text"}, "score": 0.5}]) - - response = _search(config, litellm_params={"mongodb_text_field": "metadata.body"}) - - assert response["data"][0]["content"][0]["text"] == "nested text" - - -def test_a_dotted_path_resolves_three_levels_deep(): - config, _, _ = _config(documents=[{"_id": 1, "a": {"b": {"c": "deep text"}}, "score": 0.5}]) - - response = _search(config, litellm_params={"mongodb_text_field": "a.b.c"}) - - assert response["data"][0]["content"][0]["text"] == "deep text" - - -def test_a_dotted_path_that_runs_through_a_scalar_counts_as_absent(): - """Walking 'plot.nope' when plot is a string must report the misconfiguration, not - stringify the scalar and hand the model text from the wrong field.""" - config, _, _ = _config(documents=[{"_id": 1, "plot": "a plain string", "score": 0.5}]) - - with pytest.raises(BadRequestError, match=r"has a 'plot\.nope' field"): - _search(config, litellm_params={"mongodb_text_field": "plot.nope"}) - - -def test_a_non_string_text_field_is_stringified(): - config, _, _ = _config(documents=[{"_id": 1, "year": 1979, "score": 0.5}]) - - response = _search(config, litellm_params={"mongodb_text_field": "year"}) - - assert response["data"][0]["content"][0]["text"] == "1979" - - -def test_a_null_text_field_counts_as_absent(): - config, _, _ = _config(documents=[{"_id": 1, "text": None, "score": 0.5}]) - - with pytest.raises(BadRequestError, match="has a 'text' field"): - _search(config) - - -def test_response_tolerates_a_sparse_document_missing_the_text_field(): - config, _, _ = _config(documents=[{"_id": 1, "score": 0.5}, {"_id": 2, "text": "has text", "score": 0.4}]) - - response = _search(config) - - assert response["data"][0]["content"][0]["text"] == "" - assert response["data"][1]["content"][0]["text"] == "has text" - - -def test_a_present_but_empty_text_field_is_not_treated_as_a_misconfiguration(): - config, _, _ = _config(documents=[{"_id": 1, "text": "", "score": 0.5}]) - - response = _search(config) - - assert response["data"][0]["content"][0]["text"] == "" - - -def test_matches_that_all_lack_the_text_field_name_the_setting_to_fix(): - """Atlas matches on the vector, so a mistyped mongodb_text_field returns confidently - scored results whose content is empty and hands the model an empty context.""" - config, _, _ = _config(documents=[{"_id": 1, "score": 0.9}, {"_id": 2, "score": 0.8}]) - - with pytest.raises(BadRequestError, match="mongodb_text_field"): - _search(config) - - -def test_response_tolerates_a_document_missing_a_score(): - config, _, _ = _config(documents=[{"_id": 1, "text": "no score"}]) - - response = _search(config) - - assert response["data"][0]["score"] is None - - -def test_response_stringifies_a_non_string_document_id(): - config, _, _ = _config(documents=[{"_id": 12345, "text": "numeric id", "score": 0.5}]) - - response = _search(config) - - assert response["data"][0]["file_id"] == "12345" - - -def test_search_requires_an_embedding_model(): - config, _, _ = _config() - - with pytest.raises(BadRequestError, match="litellm_embedding_model is required"): - config.execute_search_vector_store_request( - vector_store_id=INDEX, - query="q", - vector_store_search_optional_params={}, + kwargs: Final = { + "vector_store_id": "exact index", + "query": ["travel", "policy"], + "vector_store_search_optional_params": {"max_num_results": limit}, + "api_base": BASE_PARAMS["api_base"], + "litellm_logging_obj": MagicMock(), + "litellm_params": params, + } + if asynchronous: + url, body = await config.atransform_search_vector_store_request(**kwargs) + else: + url, body = config.transform_search_vector_store_request(**kwargs) + assert url == "https://sidecar.example/prefix/v1/vector_stores/exact%20index/search" + assert body == { + "query": "travel policy", + "query_vector": (0.1, 0.2, 0.3), + "mongodb_database": "policies", + "mongodb_collection": "documents", + "mongodb_text_field": "metadata.body", + "mongodb_embedding_field": "stored_vector", + "mongodb_num_candidates": candidates, + "max_num_results": limit or 10, + "timeout_ms": 750, + } + executor.call.assert_called_once_with("embedding-alias", "travel policy", {"dimensions": 3}) + assert config.transform_search_vector_store_response(httpx.Response(200, json=RESULT), MagicMock()) == RESULT + + +@pytest.mark.parametrize( + "query,overrides,options", + [ + ("", {}, {}), + (" ", {}, {}), + ("x" * 32_001, {}, {}), + ("travel", {"litellm_embedding_model": None}, {}), + ("travel", {"mongodb_database": None}, {}), + ("travel", {"mongodb_collection": None}, {}), + ("travel", {"mongodb_connection_string": "mongodb://obsolete-secret"}, {}), + ("travel", {"mongodb_filter": {"private": True}}, {}), + ("travel", {"mongodb_num_candidates": 9}, {}), + ("travel", {"mongodb_num_candidates": 10_001}, {}), + ("travel", {}, {"max_num_results": 0}), + ("travel", {}, {"max_num_results": 51}), + ("travel", {}, {"filters": {}}), + ("travel", {}, {"ranking_options": {}}), + ("travel", {}, {"rewrite_query": False}), + ], +) +def test_invalid_search_is_rejected_before_embedding( + query: str, overrides: Mapping[str, object], options: VectorStoreSearchOptionalRequestParams +) -> None: + executor: Final = RecordingEmbeddingExecutor() + config: Final = MongoDBVectorStoreConfig(executor) + with pytest.raises(litellm.BadRequestError) as error: + config.transform_search_vector_store_request( + vector_store_id="policy_index", + query=query, + vector_store_search_optional_params=options, + api_base=BASE_PARAMS["api_base"], litellm_logging_obj=MagicMock(), - litellm_params={k: v for k, v in BASE_PARAMS.items() if k != "litellm_embedding_model"}, + litellm_params={**BASE_PARAMS, **overrides}, ) - - -def test_missing_embedding_model_message_names_the_field_being_searched(): - config, _, _ = _config() - - with pytest.raises(BadRequestError, match=r"embedded_movies\.embedding"): - config.execute_search_vector_store_request( - vector_store_id=INDEX, - query="q", - vector_store_search_optional_params={}, - litellm_logging_obj=MagicMock(), - litellm_params={k: v for k, v in BASE_PARAMS.items() if k != "litellm_embedding_model"}, - ) - - -def test_search_requires_a_connection_string(): - config, _, _ = _config() - - with pytest.raises(BadRequestError, match="mongodb_connection_string is required"): - _search(config, litellm_params={"mongodb_connection_string": None}) - - -@pytest.mark.parametrize("connection_string", ["postgres://host/db", "https://cluster.mongodb.net", "redis://host"]) -def test_search_rejects_a_non_mongodb_connection_scheme(connection_string): - config, _, _ = _config() - - with pytest.raises(BadRequestError, match="must start with 'mongodb://' or 'mongodb\\+srv://'"): - _search(config, litellm_params={"mongodb_connection_string": connection_string}) - - -def test_search_accepts_the_plain_mongodb_scheme(): - config, _, collection = _config() - - _search(config, litellm_params={"mongodb_connection_string": "mongodb://localhost:27017"}) - - assert collection.pipeline is not None - - -def test_search_requires_a_database(): - config, _, _ = _config() - - with pytest.raises(BadRequestError, match="mongodb_database is required"): - _search(config, litellm_params={"mongodb_database": None}) - - -def test_search_requires_a_collection(): - config, _, _ = _config() - - with pytest.raises(BadRequestError, match="mongodb_collection is required"): - _search(config, litellm_params={"mongodb_collection": None}) - - -def test_search_rejects_filters_rather_than_silently_ignoring_them(): - config, _, _ = _config() - - with pytest.raises(BadRequestError, match="does not support the filters parameter"): - _search(config, optional_params={"filters": {"genre": "sci-fi"}}) - - -@pytest.mark.asyncio -async def test_async_search_rejects_filters_rather_than_silently_ignoring_them(): - config, _, _ = _async_config() - - with pytest.raises(BadRequestError, match="does not support the filters parameter"): - await _asearch(config, optional_params={"filters": {"genre": "sci-fi"}}) - - -def test_search_rejects_ranking_options_rather_than_silently_ignoring_them(): - """A score_threshold that is quietly dropped is worse than an error: the caller asked for - results above 0.9, gets results scoring 0.5, and nothing says the threshold never ran.""" - config, _, _ = _config() - - with pytest.raises(BadRequestError, match="does not support the ranking_options parameter"): - _search(config, optional_params={"ranking_options": {"score_threshold": 0.9}}) - - -def test_search_rejects_rewrite_query_rather_than_silently_ignoring_it(): - config, _, _ = _config() - - with pytest.raises(BadRequestError, match="does not support the rewrite_query parameter"): - _search(config, optional_params={"rewrite_query": True}) - - -@pytest.mark.asyncio -async def test_async_search_rejects_ranking_options_rather_than_silently_ignoring_them(): - config, _, _ = _async_config() - - with pytest.raises(BadRequestError, match="does not support the ranking_options parameter"): - await _asearch(config, optional_params={"ranking_options": {"score_threshold": 0.9}}) - - -@pytest.mark.parametrize("query", ["", " ", "\n\t", []]) -def test_search_rejects_an_empty_query(query): - config, _, _ = _config() - - with pytest.raises(BadRequestError, match="query must not be empty"): - _search(config, query=query) - - -def test_search_rejects_an_oversized_query(): - config, _, _ = _config() - - with pytest.raises(BadRequestError, match="at most 32000 characters"): - _search(config, query="x" * 32_001) - - -def test_search_accepts_a_query_at_the_size_ceiling(): - config, _, collection = _config() - - _search(config, query="x" * 32_000) - - assert collection.pipeline is not None - - -@pytest.mark.parametrize("max_num_results", [0, -1, 51, 1000]) -def test_search_rejects_out_of_range_max_num_results(max_num_results): - config, _, _ = _config() - - with pytest.raises(BadRequestError, match="max_num_results must be between 1 and 50"): - _search(config, optional_params={"max_num_results": max_num_results}) - - -@pytest.mark.parametrize("max_num_results", [1, 50]) -def test_search_allows_max_num_results_at_the_bounds(max_num_results): - config, _, collection = _config() - - _search(config, optional_params={"max_num_results": max_num_results}) - - assert _stage(collection, "$vectorSearch")["limit"] == max_num_results - - -def test_search_treats_an_explicit_null_max_num_results_as_the_default(): - config, _, collection = _config() - - _search(config, optional_params={"max_num_results": None}) - - assert _stage(collection, "$vectorSearch")["limit"] == 10 - - -def test_search_fails_when_the_embedding_model_returns_nothing(): - config, _, _ = _config(embedding=None) - - with pytest.raises(BadRequestError, match="returned no embedding"): - _search(config) - - -def test_validation_runs_before_any_connection_is_opened(): - opened = [] - config = MongoDBVectorStoreConfig( - embedding_executor=FakeEmbeddingExecutor([0.1]), - sync_client_factory=lambda key: opened.append(key) or FakeClient(FakeCollection([])), - ) - - with pytest.raises(BadRequestError, match="query must not be empty"): - _search(config, query="") - - assert opened == [] - - -def test_create_vector_store_is_not_supported_and_says_why(): - """litellm.exception_type only passes its own exception types through untouched, so a - NotImplementedError here reaches the caller as APIConnectionError, which the proxy serves - as a 500 with a traceback. Refusing an unsupported operation is a client error.""" - config = MongoDBVectorStoreConfig() - - with pytest.raises(BadRequestError, match="search-only"): - config.transform_create_vector_store_request({}, "https://example.test") - - with pytest.raises(BadRequestError, match="search-only"): - config.transform_create_vector_store_response(httpx.Response(200)) - - -def test_the_create_refusal_survives_the_public_sdk_error_wrapper(): - import litellm - - with pytest.raises(BadRequestError) as raised: - litellm.vector_stores.create(custom_llm_provider="mongodb", name="anything") - - assert "search-only" in str(raised.value) - - -def test_provider_config_manager_returns_the_mongodb_config(): - config = ProviderConfigManager.get_provider_vector_stores_config(LlmProviders.MONGODB) - - assert isinstance(config, MongoDBVectorStoreConfig) - - -@pytest.mark.asyncio -async def test_async_search_builds_the_same_pipeline_and_maps_the_response(): - documents = [{"_id": "abc123", "text": "an astronaut adrift", "score": 0.94}] - config, client, collection = _async_config(documents=documents) - - response = await _asearch(config, optional_params={"max_num_results": 3}) - - assert client.requested_database == "sample_mflix" - assert client.database.requested_collection == "embedded_movies" - assert _stage(collection, "$vectorSearch")["limit"] == 3 - assert _stage(collection, "$vectorSearch")["queryVector"] == (0.1, 0.2, 0.3) - assert response["data"][0]["content"][0]["text"] == "an astronaut adrift" - assert response["data"][0]["score"] == 0.94 - - -@pytest.mark.asyncio -async def test_async_search_requires_an_embedding_model(): - config, _, _ = _async_config() - - with pytest.raises(BadRequestError, match="litellm_embedding_model is required"): - await config.aexecute_search_vector_store_request( - vector_store_id=INDEX, - query="q", - vector_store_search_optional_params={}, - litellm_logging_obj=MagicMock(), - litellm_params={k: v for k, v in BASE_PARAMS.items() if k != "litellm_embedding_model"}, - ) - - -class TestClientCache: - def setup_method(self): - reset_client_cache() - - def teardown_method(self): - reset_client_cache() - - def _key(self, connection_string=CONNECTION_STRING, socket_timeout_ms=30_000): - return MongoClientKey( - connection_string=connection_string, - connect_timeout_ms=10_000, - socket_timeout_ms=socket_timeout_ms, - server_selection_timeout_ms=10_000, - ) - - def test_the_same_connection_reuses_one_client(self): - first = get_sync_client(self._key(), RecordingClient) - second = get_sync_client(self._key(), RecordingClient) - - assert first is second - assert first.connection_string == CONNECTION_STRING - assert first.kwargs["socketTimeoutMS"] == 30_000 - assert first.kwargs["connectTimeoutMS"] == 10_000 - assert first.kwargs["appname"] == "litellm" - - def test_a_different_connection_gets_its_own_client(self): - first = get_sync_client(self._key(), RecordingClient) - second = get_sync_client(self._key(connection_string="mongodb://other.example.test"), RecordingClient) - - assert first is not second - assert second.connection_string == "mongodb://other.example.test" - - def test_a_different_timeout_gets_its_own_client(self): - first = get_sync_client(self._key(), RecordingClient) - second = get_sync_client(self._key(socket_timeout_ms=5_000), RecordingClient) - - assert first is not second - assert second.kwargs["socketTimeoutMS"] == 5_000 - - @pytest.mark.asyncio - async def test_async_clients_are_cached_per_event_loop(self): - first = get_async_client(self._key(), RecordingClient) - second = get_async_client(self._key(), RecordingClient) - - assert first is second - assert first.connection_string == CONNECTION_STRING - - - def _fill_cache(self): - for slot in range(_MAX_CACHED_CLIENTS): - get_sync_client(self._key(f"mongodb://cold-{slot}:27017"), RecordingClient) - - def test_a_store_added_after_the_cache_filled_is_still_cached(self): - """Rebuilding a client costs an SRV lookup, a TLS handshake and topology discovery, so a - store that misses the cache on every single search pays that on every search.""" - self._fill_cache() - latecomer = self._key("mongodb://latecomer:27017") - - first = get_sync_client(latecomer, RecordingClient) - - assert get_sync_client(latecomer, RecordingClient) is first - - def test_the_cache_evicts_the_least_recently_used_client(self): - self._fill_cache() - oldest = self._key("mongodb://cold-0:27017") - newest = self._key(f"mongodb://cold-{_MAX_CACHED_CLIENTS - 1}:27017") - kept = get_sync_client(newest, RecordingClient) - - get_sync_client(self._key("mongodb://latecomer:27017"), RecordingClient) - - assert get_sync_client(newest, RecordingClient) is kept - assert oldest not in _sync_clients - - def test_concurrent_searches_never_trip_over_an_eviction(self): - """Async searches run the sync client through executor threads, so a key can be evicted - between the lookup and the reordering that follows it.""" - errors = [] - churn = _MAX_CACHED_CLIENTS + 2 - - def hammer(offset): - try: - for step in range(3_000): - get_sync_client(self._key(f"mongodb://h-{(step + offset) % churn}:27017"), RecordingClient) - except Exception as e: - errors.append(repr(e)) - - previous = sys.getswitchinterval() - sys.setswitchinterval(1e-9) - try: - threads = [threading.Thread(target=hammer, args=(offset,)) for offset in range(16)] - for thread in threads: - thread.start() - for thread in threads: - thread.join() - finally: - sys.setswitchinterval(previous) - - assert errors == [] - - def test_the_cache_never_grows_past_its_cap(self): - for slot in range(_MAX_CACHED_CLIENTS * 3): - get_sync_client(self._key(f"mongodb://host-{slot}:27017"), RecordingClient) - - assert len(_sync_clients) == _MAX_CACHED_CLIENTS - - def test_a_new_loop_never_inherits_a_closed_loop_client(self): - """CPython recycles id() so aggressively that a fresh event loop almost always lands on - the id of one already collected: measured at 37 of 40 rounds. Keying the cache on the id - alone therefore hands the new loop an AsyncMongoClient bound to a closed loop, and every - operation on it raises "Event loop is closed".""" - - class LoopAgnosticClient: - """Holds no reference to the loop, unlike pymongo's, whose own reference happens to - keep ids from being recycled and hides the bug until the cache fills.""" - - def __init__(self, *args, **kwargs): - self.built_on = None - - key = self._key() - clients_handed_out = [] - - async def fetch(): - return get_async_client(key, LoopAgnosticClient) - - for _ in range(20): - loop = asyncio.new_event_loop() - client = loop.run_until_complete(fetch()) - clients_handed_out.append((client, client.built_on, loop.is_closed())) - client.built_on = weakref.ref(loop) - loop.close() - del loop - gc.collect() - - stale = [ - handed_out - for client, built_on, _ in clients_handed_out - if built_on is not None and (built_on() is None or built_on().is_closed()) - for handed_out in (client,) - ] - assert stale == [], f"{len(stale)} of 20 loops were handed a client built on a closed loop" - - def test_the_cache_releases_clients_built_on_closed_loops(self): - """pymongo's AsyncMongoClient keeps a reference to the loop it was built on, so an entry - for a closed loop holds that client, and its sockets, for the life of the process. A - script calling asyncio.run per search fills the cache to its cap that way: measured live - against Atlas at 32 pinned clients and 212 open descriptors after 40 loops.""" - - class LoopHoldingClient: - def __init__(self, *args, **kwargs): - self.loop = asyncio.get_running_loop() - - key = self._key() - - async def fetch(): - return get_async_client(key, LoopHoldingClient) - - for _ in range(_MAX_CACHED_CLIENTS + 8): - loop = asyncio.new_event_loop() - loop.run_until_complete(fetch()) - loop.close() - - assert len(_async_clients) == 1, f"{len(_async_clients)} closed-loop clients are still cached" - - -class TestClientKeyDerivation: - def test_no_timeout_uses_the_bounded_defaults(self): - key = MongoDBVectorStoreConfig._client_key(_MongoDBSearchParams.model_validate(BASE_PARAMS), None) - - assert key.connect_timeout_ms == 10_000 - assert key.socket_timeout_ms == 30_000 - assert key.server_selection_timeout_ms == 10_000 - - def test_a_numeric_timeout_bounds_the_connect_phase(self): - key = MongoDBVectorStoreConfig._client_key(_MongoDBSearchParams.model_validate(BASE_PARAMS), 3.0) - - assert key.socket_timeout_ms == 3_000 - assert key.connect_timeout_ms == 3_000 - - def test_a_short_timeout_also_shortens_server_selection(self): - """Server selection runs before the connect attempt, so leaving it at the 10s default - would let a caller asking for a 3s budget block for 10s before anything is tried.""" - key = MongoDBVectorStoreConfig._client_key(_MongoDBSearchParams.model_validate(BASE_PARAMS), 3.0) - - assert key.server_selection_timeout_ms == 3_000 - - def test_a_generous_timeout_does_not_raise_server_selection_above_the_default(self): - key = MongoDBVectorStoreConfig._client_key(_MongoDBSearchParams.model_validate(BASE_PARAMS), 120.0) - - assert key.socket_timeout_ms == 120_000 - assert key.server_selection_timeout_ms == 10_000 - - def test_an_httpx_timeout_maps_connect_and_read_separately(self): - key = MongoDBVectorStoreConfig._client_key( - _MongoDBSearchParams.model_validate(BASE_PARAMS), httpx.Timeout(connect=2.0, read=45.0, write=5.0, pool=5.0) - ) - - assert key.connect_timeout_ms == 2_000 - assert key.socket_timeout_ms == 45_000 - - -class TestErrorTranslation: - def _translate(self, error): - return translate_mongo_error(error, index_name=INDEX, database="sample_mflix", collection="embedded_movies") - - def test_server_selection_timeout_points_at_the_atlas_access_list(self): - from pymongo.errors import ServerSelectionTimeoutError - - translated = self._translate(ServerSelectionTimeoutError("no servers")) - - assert "IP access list" in str(translated) - assert "paused cluster" in str(translated) - - def test_authentication_failure_points_at_the_connection_string_credentials(self): - from pymongo.errors import OperationFailure - - translated = self._translate(OperationFailure("auth failed", code=18)) - - assert "rejected the credentials" in str(translated) - - def test_a_dropped_connection_stays_retryable(self): - """A replica set failover reaches the driver as AutoReconnect. litellm only retries 408, - 409, 429 and 5xx, so classifying it as a client error would turn one failover into a - permanently failed search.""" - from pymongo.errors import AutoReconnect - - translated = self._translate(AutoReconnect("connection closed")) - - assert litellm._should_retry(translated.status_code) - assert "dropped or refused" in str(translated) - - def test_a_dropped_connection_still_names_the_misconfigurations_behind_it(self): - """Atlas answers a URI with no credentials by closing the connection rather than failing - auth, so the retryable message still has to name that.""" - from pymongo.errors import AutoReconnect - - translated = self._translate(AutoReconnect("connection closed")) - - assert "no username and password" in str(translated) - assert "mongod is listening" in str(translated) - - def test_the_retryable_classification_survives_the_public_sdk_error_wrapper(self): - """litellm.exception_type only passes its own exception types through; anything else becomes - an APIConnectionError and a 500, which would drop the retryable classification.""" - from pymongo.errors import AutoReconnect - - translated = self._translate(AutoReconnect("connection closed")) - - wrapped = litellm.exception_type( - model=None, - original_exception=translated, - custom_llm_provider="mongodb", - completion_kwargs={}, - extra_kwargs={}, - ) - - assert isinstance(wrapped, ServiceUnavailableError) - assert litellm._should_retry(wrapped.status_code) - - def test_a_pool_wait_queue_timeout_stays_retryable(self): - from pymongo.errors import WaitQueueTimeoutError - - translated = self._translate(WaitQueueTimeoutError("timed out waiting for a connection")) - - assert litellm._should_retry(translated.status_code) - - def test_server_selection_timeout_still_wins_over_the_connection_branch(self): - from pymongo.errors import ServerSelectionTimeoutError - - translated = self._translate(ServerSelectionTimeoutError("no servers")) - - assert isinstance(translated, Timeout) - assert "dropped or refused" not in str(translated) - - def test_network_timeout_still_wins_over_the_connection_branch(self): - from pymongo.errors import NetworkTimeout - - translated = self._translate(NetworkTimeout("socket timed out")) - - assert isinstance(translated, Timeout) - assert "dropped or refused" not in str(translated) - - def test_an_unescaped_password_character_is_a_400_not_a_500(self): - """pymongo's URI parser raises a plain ValueError, not a PyMongoError, for an unusable port, - which is also what an unescaped ':' in a password produces. It must not be a 500.""" - translated = self._translate(ValueError("Port contains non-digit characters")) - - assert isinstance(translated, BadRequestError) - assert "percent-encoded" in str(translated) - - def test_unauthorized_points_at_the_database_user_permissions(self): - from pymongo.errors import OperationFailure - - translated = self._translate(OperationFailure("not authorized", code=13)) - - assert "sample_mflix.embedded_movies" in str(translated) - - def test_code_13_alone_is_enough_without_a_recognisable_message(self): - """The other unauthorized case carries "not authorized", which the message markers also - match, so it cannot tell whether the code is still being checked at all.""" - from pymongo.errors import OperationFailure - - translated = self._translate(OperationFailure("user lacks privileges on this namespace", code=13)) - - assert "rejected the credentials" in str(translated) - assert "sample_mflix.embedded_movies" in str(translated) - - def test_a_missing_index_names_the_index_and_the_collection(self): - from pymongo.errors import OperationFailure - - translated = self._translate(OperationFailure("Index not found for name movies_vector_index", code=27)) - - assert INDEX in str(translated) - assert "READY" in str(translated) - - def test_a_dimension_mismatch_points_at_the_embedding_model(self): - from pymongo.errors import OperationFailure - - translated = self._translate(OperationFailure("queryVector has 1536 dimensions, index expects 2048")) - - assert "litellm_embedding_model must be the same model" in str(translated) - - def test_an_unrecognised_operation_failure_still_names_the_target(self): - from pymongo.errors import OperationFailure - - translated = self._translate(OperationFailure("something else entirely")) - - assert "sample_mflix.embedded_movies" in str(translated) - assert INDEX in str(translated) - - def test_a_configuration_error_points_at_the_connection_string(self): - from pymongo.errors import ConfigurationError - - translated = self._translate(ConfigurationError("bad uri")) - - assert "not a usable MongoDB connection string" in str(translated) - - def test_a_non_driver_error_is_returned_unchanged(self): - original = RuntimeError("unrelated") - - assert self._translate(original) is original - - def test_search_surfaces_a_translated_driver_error(self): - from pymongo.errors import ServerSelectionTimeoutError - - config, _, _ = _config(error=ServerSelectionTimeoutError("no servers")) - - with pytest.raises(Timeout, match="IP access list"): - _search(config) - - @pytest.mark.asyncio - async def test_async_search_surfaces_a_translated_driver_error(self): - from pymongo.errors import OperationFailure - - config, _, _ = _async_config(error=OperationFailure("auth failed", code=18)) - - with pytest.raises(BadRequestError, match="rejected the credentials"): - await _asearch(config) - - -class TestMissingDriver: - def test_the_sync_import_names_the_extra_to_install(self): - from litellm.llms.mongodb.common_utils import import_sync_mongo_client - - with patch.dict(sys.modules, {"pymongo": None}): - with pytest.raises(BadRequestError, match=r"pip install litellm\[mongodb\]"): - import_sync_mongo_client() - - def test_the_async_import_names_the_extra_to_install(self): - from litellm.llms.mongodb.common_utils import import_async_mongo_client - - with patch.dict(sys.modules, {"pymongo": None}): - with pytest.raises(BadRequestError, match=r"pip install litellm\[mongodb\]"): - import_async_mongo_client() - - def test_error_translation_degrades_gracefully_without_the_driver(self): - original = RuntimeError("boom") - - with patch.dict(sys.modules, {"pymongo.errors": None}): - assert translate_mongo_error(original, INDEX, "db", "col") is original - - -class TestEmptyResultsAreDisambiguated: - """$vectorSearch returns zero documents for a missing database, collection or index just as it - does for a query that matched nothing, so an empty result set is checked against the index - catalogue before it is reported as 'no matches'.""" - - def test_a_missing_index_becomes_an_error_rather_than_an_empty_page(self): - config, _, collection = _config(documents=[], search_indexes=[]) - - with pytest.raises(BadRequestError, match="No queryable MongoDB Vector Search index"): - _search(config) - - assert collection.listed_indexes == [INDEX] - - def test_the_missing_index_error_explains_why_mongodb_reported_no_results(self): - config, _, _ = _config(documents=[], search_indexes=[]) - - with pytest.raises(BadRequestError, match="returns no results rather than an error"): - _search(config) - - def test_an_index_still_building_becomes_an_error_naming_its_status(self): - config, _, _ = _config( - documents=[], search_indexes=[{"name": INDEX, "status": "PENDING", "queryable": False}] - ) - - with pytest.raises(BadRequestError, match="not queryable yet; its status is PENDING"): - _search(config) - - def test_a_genuine_no_match_against_a_ready_index_returns_an_empty_page(self): - config, _, collection = _config(documents=[]) - - response = _search(config) - - assert response["data"] == [] - assert response["object"] == "vector_store.search_results.page" - assert collection.listed_indexes == [INDEX] - - def test_the_catalogue_is_not_consulted_when_the_search_returned_hits(self): - config, _, collection = _config(documents=[{"_id": 1, "text": "hit", "score": 0.9}]) - - _search(config) - - assert collection.listed_indexes == [] - - @pytest.mark.asyncio - async def test_async_missing_index_becomes_an_error_rather_than_an_empty_page(self): - config, _, collection = _async_config(documents=[], search_indexes=[]) - - with pytest.raises(BadRequestError, match="No queryable MongoDB Vector Search index"): - await _asearch(config) - - assert collection.listed_indexes == [INDEX] - - @pytest.mark.asyncio - async def test_async_index_still_building_becomes_an_error_naming_its_status(self): - config, _, _ = _async_config( - documents=[], search_indexes=[{"name": INDEX, "status": "PENDING", "queryable": False}] - ) - - with pytest.raises(BadRequestError, match="not queryable yet; its status is PENDING"): - await _asearch(config) - - @pytest.mark.asyncio - async def test_async_genuine_no_match_returns_an_empty_page(self): - config, _, _ = _async_config(documents=[]) - - response = await _asearch(config) - - assert response["data"] == [] - - @pytest.mark.asyncio - async def test_async_catalogue_is_not_consulted_when_the_search_returned_hits(self): - config, _, collection = _async_config(documents=[{"_id": 1, "text": "hit", "score": 0.9}]) - - await _asearch(config) - - assert collection.listed_indexes == [] - - def test_a_failure_while_checking_the_catalogue_is_translated_too(self): - from pymongo.errors import OperationFailure - - class ExplodingCollection(FakeCollection): - def list_search_indexes(self, name): - raise OperationFailure("not authorized", code=13) - - collection = ExplodingCollection([], None, []) - config = MongoDBVectorStoreConfig( - embedding_executor=FakeEmbeddingExecutor([0.1]), - sync_client_factory=lambda key: FakeClient(collection), - ) - - with pytest.raises(BadRequestError, match="lacks read access"): - _search(config) - - -class TestAtlasPlanExecutorErrors: - """Atlas reports a wrong vector path and a dimension mismatch through the same error code, so - each one has to be told apart by its message or both come back as a generic index failure.""" - - def _translate(self, message): - from pymongo.errors import OperationFailure - - return translate_mongo_error( - OperationFailure(message, code=8), - index_name=INDEX, - database="sample_mflix", - collection="embedded_movies", - ) - - def test_a_wrong_vector_path_points_at_the_embedding_field_setting(self): - translated = self._translate( - "PlanExecutor error during aggregation :: caused by :: nope is not indexed as vector" - ) - - assert "mongodb_embedding_field names a field" in str(translated) - - def test_a_dimension_mismatch_is_not_reported_as_a_wrong_path(self): - translated = self._translate( - "PlanExecutor error during aggregation :: caused by :: vector field is indexed with " - "1536 dimensions but queried with 3072" - ) - - assert "does not match the vector dimensions" in str(translated) - assert "mongodb_embedding_field" not in str(translated) - - -class TestErrorsCarryTheRightHttpStatus: - """litellm.exception_type passes a litellm exception through untouched but wraps anything - else into APIConnectionError, which the proxy serves as a 500 with a Python traceback in the - body. A misconfigured connection string is the caller's to fix, so it has to arrive as a 400. - """ - - @pytest.mark.parametrize( - "invoke", - [ - pytest.param(lambda: _search(_config()[0], query=" "), id="empty-query"), - pytest.param( - lambda: _search(_config()[0], optional_params={"max_num_results": 999}), - id="max-num-results-out-of-range", - ), - pytest.param( - lambda: _search(_config()[0], optional_params={"filters": {"genre": "Action"}}), - id="unsupported-filters", - ), - pytest.param( - lambda: _search(_config()[0], litellm_params={"mongodb_connection_string": "postgres://host/db"}), - id="wrong-uri-scheme", - ), - pytest.param( - lambda: _search(_config()[0], litellm_params={"mongodb_database": None}), id="missing-database" - ), - pytest.param( - lambda: _search(_config()[0], litellm_params={"litellm_embedding_model": None}), - id="missing-embedding-model", - ), - ], - ) - def test_configuration_failures_are_400(self, invoke): - with pytest.raises(BadRequestError) as excinfo: - invoke() - assert excinfo.value.status_code == 400 - assert excinfo.value.llm_provider == "mongodb" - - def test_missing_index_is_400(self): - error = missing_index_error("idx", "db", "coll") - assert error.status_code == 400 - assert error.llm_provider == "mongodb" - - def test_index_still_building_is_400(self): - error = index_not_ready_error("idx", "db", "coll", "PENDING") - assert error.status_code == 400 - - def test_unreachable_deployment_is_a_timeout_not_a_bad_request(self): - from pymongo.errors import ServerSelectionTimeoutError - - translated = translate_mongo_error( - ServerSelectionTimeoutError("no servers"), index_name="idx", database="db", collection="coll" - ) - assert isinstance(translated, Timeout) - assert translated.status_code == 408 - - def test_query_execution_timeout_is_a_timeout(self): - from pymongo.errors import ExecutionTimeout - - translated = translate_mongo_error( - ExecutionTimeout("too slow"), index_name="idx", database="db", collection="coll" - ) - assert isinstance(translated, Timeout) - assert translated.status_code == 408 - - def test_unrecognised_errors_are_not_relabelled_as_bad_requests(self): - original = RuntimeError("something else entirely") - assert ( - translate_mongo_error(original, index_name="idx", database="db", collection="coll") - is original - ) - - -def test_atlas_rejected_credentials_are_named_even_though_the_code_is_8000(): - """Atlas answers a wrong password with code 8000 "AtlasError", not the 18 that a - self-hosted deployment returns, so a code-only check reports it as a generic - rejected search and never tells the caller to look at their connection string.""" - from pymongo.errors import OperationFailure - - error = OperationFailure( - "bad auth : authentication failed", - code=8000, - details={"ok": 0, "errmsg": "bad auth : authentication failed", "code": 8000, "codeName": "AtlasError"}, - ) - translated = translate_mongo_error(error, index_name="idx", database="sample_mflix", collection="embedded_movies") - - assert isinstance(translated, BadRequestError) - assert "mongodb_connection_string" in str(translated) - assert "sample_mflix.embedded_movies" in str(translated) - - -def test_a_rejected_search_that_is_not_an_auth_failure_keeps_the_generic_message(): - from pymongo.errors import OperationFailure - - error = OperationFailure("PlanExecutor error", code=8, details={"errmsg": "PlanExecutor error"}) - translated = translate_mongo_error(error, index_name="idx", database="db", collection="coll") - - assert "mongodb_connection_string" not in str(translated) - - -class TestUnrecognisedParameters: - """litellm_params carries plenty of keys this provider does not own, so the params model has - to ignore extras. That turns a mistyped mongodb_collection into 'mongodb_collection is - required', pointing the reader at a key they can see they have set.""" - - def test_a_mistyped_parameter_is_named(self): - config, _, _ = _config() - - with pytest.raises(BadRequestError, match="mongodb_collectoin"): - _search(config, litellm_params={"mongodb_collectoin": "embedded_movies"}) - - def test_the_supported_names_are_listed(self): - config, _, _ = _config() - - with pytest.raises(BadRequestError, match="mongodb_connection_string"): - _search(config, litellm_params={"mongodb_databse": "sample_mflix"}) - - def test_unrelated_litellm_params_are_still_ignored(self): - config, _, _ = _config(documents=[{"_id": 1, "text": "hit", "score": 0.9}]) - - response = _search( - config, - litellm_params={"use_litellm_proxy": False, "use_in_pass_through": False, "vector_store_id": "x"}, - ) - - assert len(response["data"]) == 1 - - @pytest.mark.asyncio - async def test_the_async_path_rejects_them_too(self): - config, _, _ = _async_config() - - with pytest.raises(BadRequestError, match="mongodb_collectoin"): - await _asearch(config, litellm_params={"mongodb_collectoin": "embedded_movies"}) - - -class TestClientConstructionFailures: - """Building the client parses the URI and, for mongodb+srv://, performs a DNS SRV lookup, so it - fails on exactly the inputs a user is most likely to get wrong. Constructing it outside the - translation boundary let those escape as raw pymongo errors, which litellm.exception_type then - wrapped into a 500 with a traceback in the body.""" - - def _config_that_fails_to_connect(self, error): - def factory(_key): - raise error - - return MongoDBVectorStoreConfig( - embedding_executor=FakeEmbeddingExecutor([0.1, 0.2, 0.3]), sync_client_factory=factory - ) - - def _async_config_that_fails_to_connect(self, error): - def factory(_key): - raise error - - return MongoDBVectorStoreConfig( - embedding_executor=FakeEmbeddingExecutor([0.1, 0.2, 0.3]), async_client_factory=factory - ) - - def test_a_malformed_uri_is_a_bad_request_not_a_500(self): - from pymongo.errors import InvalidURI - - config = self._config_that_fails_to_connect(InvalidURI("Invalid URI scheme")) - - with pytest.raises(BadRequestError, match="not a usable MongoDB connection string"): - _search(config) - - def test_an_unresolvable_cluster_name_says_so(self): - from pymongo.errors import ConfigurationError - - config = self._config_that_fails_to_connect(ConfigurationError("The DNS query name does not exist")) - - with pytest.raises(BadRequestError, match="does not exist in DNS"): - _search(config) - - def test_a_dns_lookup_that_ran_out_of_time_is_a_timeout(self): - from pymongo.errors import ConfigurationError - - config = self._config_that_fails_to_connect( - ConfigurationError("The resolution lifetime expired after 0.291 seconds") - ) - - with pytest.raises(Timeout, match="did not finish in time"): - _search(config) - - @pytest.mark.asyncio - async def test_the_async_path_translates_them_too(self): - from pymongo.errors import InvalidURI - - config = self._async_config_that_fails_to_connect(InvalidURI("Invalid URI scheme")) - - with pytest.raises(BadRequestError, match="not a usable MongoDB connection string"): - await _asearch(config) - - -class TestSelfManagedDeploymentsAreFirstClass: - """mongod serves $vectorSearch identically whether mongot runs under Atlas or beside a - self-managed deployment, so an operator without an Atlas account has to be able to act on - every message. Guidance that only names Atlas remedies sends them looking for an IP access - list and a paused cluster that do not exist in their deployment.""" - - def _config_that_fails_to_connect(self, error): - def factory(_key): - raise error - - return MongoDBVectorStoreConfig( - embedding_executor=FakeEmbeddingExecutor([0.1, 0.2, 0.3]), sync_client_factory=factory - ) - - def test_a_plain_mongodb_uri_without_srv_or_credentials_is_accepted(self): - params = _MongoDBSearchParams.model_validate( - {**BASE_PARAMS, "mongodb_connection_string": "mongodb://mongod.internal:27017"} - ) - - assert params.require_connection_string() == "mongodb://mongod.internal:27017" - - def test_an_unreachable_deployment_names_a_self_managed_remedy(self): - from pymongo.errors import ServerSelectionTimeoutError - - config = self._config_that_fails_to_connect(ServerSelectionTimeoutError("connection refused")) - - with pytest.raises(Timeout) as excinfo: - _search(config) - - assert "self-managed" in str(excinfo.value) - assert "host or port" in str(excinfo.value) - - def test_a_refused_connection_names_a_self_managed_remedy(self): - from pymongo.errors import ConnectionFailure - - config = self._config_that_fails_to_connect(ConnectionFailure("connection closed")) - - with pytest.raises(ServiceUnavailableError) as excinfo: - _search(config) - - assert "self-managed" in str(excinfo.value) - assert "mongod is listening" in str(excinfo.value) - - def test_an_unresolvable_hostname_names_a_self_managed_remedy(self): - from pymongo.errors import ConfigurationError - - config = self._config_that_fails_to_connect(ConfigurationError("The DNS query name does not exist")) - - with pytest.raises(BadRequestError) as excinfo: - _search(config) - - assert "self-managed" in str(excinfo.value) - - def test_the_missing_index_message_does_not_claim_atlas(self): - message = str(missing_index_error(INDEX, "sample_mflix", "embedded_movies")) - - assert "MongoDB Vector Search index" in message - assert "Atlas" not in message - - def test_the_not_ready_message_does_not_claim_atlas(self): - message = str(index_not_ready_error(INDEX, "sample_mflix", "embedded_movies", "PENDING")) - - assert "MongoDB Vector Search index" in message - assert "Atlas" not in message - - def test_the_search_only_refusal_does_not_claim_atlas(self): - config = MongoDBVectorStoreConfig() - - with pytest.raises(BadRequestError) as excinfo: - config.transform_create_vector_store_request({}, api_base="") - - assert "Atlas" not in str(excinfo.value) - - def test_a_dimension_mismatch_does_not_claim_atlas(self): - from pymongo.errors import OperationFailure - - error = OperationFailure("vector field is indexed with 128 dimensions but queried with 256") - translated = translate_mongo_error(error, index_name=INDEX, database="db", collection="c") - - assert "Atlas" not in str(translated) - assert "dimensions the index was built for" in str(translated) - - def test_an_uncovered_embedding_field_does_not_claim_atlas(self): - from pymongo.errors import OperationFailure - - error = OperationFailure("embedding is not indexed as vector") - translated = translate_mongo_error(error, index_name=INDEX, database="db", collection="c") - - assert "MongoDB Vector Search index does not cover" in str(translated) - assert "Atlas" not in str(translated) - - def test_a_self_managed_auth_failure_is_still_recognised_by_code_18(self): - from pymongo.errors import OperationFailure - - error = OperationFailure("Authentication failed.", code=18, details={"code": 18}) - translated = translate_mongo_error(error, index_name=INDEX, database="db", collection="c") - - assert isinstance(translated, BadRequestError) - assert "rejected the credentials" in str(translated) - - -class TestUnescapedCredentialsAreDiagnosed: - """Self-managed deployments usually carry a generated password, so '@', '/', ':' and '%' in one - are routine. pymongo reports those as a port, a database name or an RFC 3986 complaint, none of - which points the operator at their password, so each has to be named for what it is. The errors - here come from pymongo's real parser rather than a synthetic stand-in.""" - - @staticmethod - def _real_parse_error(uri): - from pymongo import MongoClient - - try: - MongoClient(uri, serverSelectionTimeoutMS=1) - except Exception as e: - return e - raise AssertionError(f"expected {uri!r} to fail parsing") - - def _translated(self, uri): - return translate_mongo_error( - self._real_parse_error(uri), index_name=INDEX, database="db", collection="c" - ) - - @pytest.mark.parametrize( - "uri", - [ - "mongodb://user:pa@ss@host:27017/", - "mongodb://user:pa:ss@host:27017/", - "mongodb://user:pa%ss@host:27017/", - "mongodb://user@x:pw@host:27017/", - ], - ) - def test_rfc_3986_complaints_tell_the_operator_to_encode_the_password(self, uri): - translated = self._translated(uri) - - assert isinstance(translated, BadRequestError) - assert "percent-encoded per RFC 3986" in str(translated) - - @pytest.mark.parametrize( - "uri", - ["mongodb://user:pa/ss@host:27017/", "mongodb://user/x:pw@host:27017/"], - ) - def test_a_slash_in_the_credentials_is_not_reported_as_a_database_name(self, uri): - translated = self._translated(uri) - - assert isinstance(translated, BadRequestError) - assert "percent-encoded per RFC 3986" in str(translated) - - def test_an_unusable_port_names_the_host_and_port_not_the_database(self): - translated = self._translated("mongodb://host:99999/") - - assert isinstance(translated, BadRequestError) - assert "host and port" in str(translated) - - def test_a_genuinely_bad_database_name_still_mentions_the_uri_path(self): - translated = self._translated("mongodb://host:27017/has space") - - assert isinstance(translated, BadRequestError) - assert "database name in the URI path" in str(translated) - - -class TestUnreadableTlsFilesAreDiagnosed: - """A private CA is how self-managed deployments present TLS, so tlsCAFile and - tlsCertificateKeyFile are on-prem options in practice. pymongo opens those files itself and - lets OSError out, which is not a PyMongoError, so before this they reached the caller as a 500 - with a traceback. The errors here come from pymongo's real TLS setup.""" - - @staticmethod - def _real_tls_error(uri): - from pymongo import MongoClient - - try: - MongoClient(uri, serverSelectionTimeoutMS=1500).admin.command("ping") - except Exception as e: - return e - raise AssertionError(f"expected {uri!r} to fail") - - def _translated(self, uri): - return translate_mongo_error(self._real_tls_error(uri), index_name=INDEX, database="db", collection="c") - - @pytest.mark.parametrize( - "path", - ["/nonexistent-directory-for-tests/ca.pem", "/tmp"], - ) - def test_an_unreadable_ca_file_is_a_400_naming_the_path(self, path): - translated = self._translated(f"mongodb://localhost:27717/?tls=true&tlsCAFile={path}") - - assert isinstance(translated, BadRequestError) - assert path in str(translated) - assert "tlsCAFile" in str(translated) - - def test_an_unreadable_client_certificate_is_a_400_naming_the_path(self): - path = "/nonexistent-directory-for-tests/client.pem" - translated = self._translated(f"mongodb://localhost:27717/?tls=true&tlsCertificateKeyFile={path}") - - assert isinstance(translated, BadRequestError) - assert path in str(translated) - - def test_an_oserror_carrying_no_filename_is_left_for_the_other_branches(self): - translated = translate_mongo_error(OSError("socket hung up"), index_name=INDEX, database="db", collection="c") - - assert not isinstance(translated, BadRequestError) - - -class TestTheCallerSuppliedEmbeddingExecutorIsUsed: - """litellm.vector_stores.search always hands a direct provider an embedding_executor, so the - provider has to accept it and route the query through it rather than its own default.""" - - def test_the_supplied_executor_produces_the_query_vector(self): - config, _, collection = _config(embedding=(0.9, 0.9, 0.9), search_indexes=READY_INDEX) - caller = FakeEmbeddingExecutor([0.4, 0.5, 0.6]) - - config.execute_search_vector_store_request( - vector_store_id=INDEX, - query="a lone astronaut", - vector_store_search_optional_params={}, - litellm_logging_obj=MagicMock(), - litellm_params=BASE_PARAMS, - embedding_executor=caller, - ) - - assert caller.captured.query == "a lone astronaut" - assert _stage(collection, "$vectorSearch")["queryVector"] == (0.4, 0.5, 0.6) - - @pytest.mark.asyncio - async def test_the_supplied_executor_produces_the_query_vector_on_the_async_path(self): - config, _, collection = _async_config(embedding=(0.9, 0.9, 0.9), search_indexes=READY_INDEX) - caller = FakeEmbeddingExecutor([0.4, 0.5, 0.6]) - - await config.aexecute_search_vector_store_request( - vector_store_id=INDEX, - query="a lone astronaut", - vector_store_search_optional_params={}, - litellm_logging_obj=MagicMock(), - litellm_params=BASE_PARAMS, - embedding_executor=caller, - ) - - assert caller.captured.query == "a lone astronaut" - assert _stage(collection, "$vectorSearch")["queryVector"] == (0.4, 0.5, 0.6) + assert "obsolete-secret" not in str(error.value) + executor.call.assert_not_called() + + +@pytest.mark.parametrize( + "status,body,error_type", + [ + (400, {"error": {"message": "Index is not queryable"}}, litellm.BadRequestError), + (401, {}, litellm.AuthenticationError), + (408, {}, litellm.Timeout), + (503, {}, litellm.ServiceUnavailableError), + (200, {}, litellm.ServiceUnavailableError), + (200, {**RESULT, "data": [{"score": "wrong"}]}, litellm.ServiceUnavailableError), + (0, {}, litellm.Timeout), + (-1, {}, litellm.BadRequestError), + (200, RESULT, None), + ], +) +def test_public_sdk_preserves_http_errors_response_and_timeout( + status: int, body: Mapping[str, object], error_type: type[Exception] | None +) -> None: + executor: Final = RecordingEmbeddingExecutor() + if status == -1: + with pytest.raises(litellm.BadRequestError, match="search-only"): + litellm.vector_stores.create(custom_llm_provider="mongodb") + executor.call.assert_not_called() + return + + def respond(request: httpx.Request) -> httpx.Response: + assert request.url == "https://sidecar.example/prefix/v1/vector_stores/policy_index/search" + assert request.headers["authorization"] == "Bearer test-sidecar-key" + assert request.extensions["timeout"]["read"] == 0.75 + payload: Final = json.loads(request.content) + assert payload["timeout_ms"] == 750 + assert payload["query_vector"] == [0.1, 0.2, 0.3] + if status == 0: + raise httpx.ReadTimeout("timed out", request=request) + return httpx.Response(status, json=body) + + with httpx.Client(transport=httpx.MockTransport(respond)) as transport: + client: Final = HTTPHandler(client=transport) + if error_type is not None: + with pytest.raises(error_type): + litellm.vector_stores.search( + vector_store_id="policy_index", + query="travel policy", + custom_llm_provider="mongodb", + _direct_vector_store_embedding_executor=executor, + client=client, + timeout=0.75, + **BASE_PARAMS, + ) + else: + result: Final = litellm.vector_stores.search( + vector_store_id="policy_index", + query="travel policy", + custom_llm_provider="mongodb", + _direct_vector_store_embedding_executor=executor, + client=client, + timeout=0.75, + **BASE_PARAMS, + ) + assert result == RESULT + executor.call.assert_called_once_with("embedding-alias", "travel policy", {}) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreForm.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreForm.test.tsx index 84a9314ecce..8da7098b695 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreForm.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreForm.test.tsx @@ -69,10 +69,11 @@ describe("VectorStoreForm", () => { }); }); -const MONGODB_URI = "mongodb+srv://user:pass@cluster0.mongodb.net"; +const MONGODB_SIDECAR_URL = "http://mongodb-sidecar:8080"; const MONGODB_REQUIRED_FORM_VALUES = { - mongodb_connection_string: MONGODB_URI, + api_base: MONGODB_SIDECAR_URL, + api_key: "sidecar-test-key", mongodb_database: "sample_mflix", mongodb_collection: "embedded_movies", embedding_model: "text-embedding-ada-002", @@ -127,7 +128,8 @@ describe("buildVectorStoreLitellmParams", () => { mongodb_num_candidates: "200", }; const expected = { - mongodb_connection_string: MONGODB_URI, + api_base: MONGODB_SIDECAR_URL, + api_key: "sidecar-test-key", mongodb_database: "sample_mflix", mongodb_collection: "embedded_movies", mongodb_embedding_field: "plot_embedding", @@ -142,6 +144,7 @@ describe("buildVectorStoreLitellmParams", () => { it("sends only mongodb fields when an earlier provider left values in the form", () => { const formValues = { ...MONGODB_REQUIRED_FORM_VALUES, + mongodb_connection_string: "mongodb://obsolete-credentials", valkey_host: "left-over-from-valkey.example.com", valkey_port: "6379", aws_region_name: "us-west-2", @@ -152,7 +155,8 @@ describe("buildVectorStoreLitellmParams", () => { expect(params).not.toHaveProperty("valkey_host"); expect(params).not.toHaveProperty("valkey_port"); expect(params).not.toHaveProperty("aws_region_name"); - expect(params.mongodb_connection_string).toBe(MONGODB_URI); + expect(params.api_base).toBe(MONGODB_SIDECAR_URL); + expect(params).not.toHaveProperty("mongodb_connection_string"); }); it("omits a blank mongodb_num_candidates so litellm picks its own candidate count", () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreForm.tsx b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreForm.tsx index 61da25874a5..67ef1b795ba 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreForm.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/VectorStoreForm.tsx @@ -70,7 +70,6 @@ const PROVIDER_FIELD_NAMES = [ "vector_bucket_name", "index_name", "aws_region_name", - "mongodb_connection_string", "mongodb_database", "mongodb_collection", "mongodb_embedding_field", @@ -107,7 +106,6 @@ const vectorStoreShape = { vector_bucket_name: optionalText, index_name: optionalText, aws_region_name: optionalText, - mongodb_connection_string: optionalText, mongodb_database: optionalText, mongodb_collection: optionalText, mongodb_embedding_field: optionalText, @@ -142,7 +140,7 @@ const VECTOR_STORE_ID_PLACEHOLDERS: Record = { vertex_rag_engine: '6917529027641081856 (corpus ID from Vertex AI / "RAG Engine" console)', "vertex_ai/search_api": 'my-datastore_1234567890 (data store ID from Vertex AI / "Agent Search" console)', valkey: "my-search-index (FT index name in Valkey)", - mongodb: "my-vector-index (Atlas Vector Search index name)", + mongodb: "my-vector-index (MongoDB Vector Search index name)", }; const VERTEX_SEARCH_API_WITH_ENGINE_PLACEHOLDER = "Any identifier you'll use to reference this in LiteLLM"; diff --git a/ui/litellm-dashboard/src/components/vector_store_providers.test.tsx b/ui/litellm-dashboard/src/components/vector_store_providers.test.tsx index 8e3a3aa3402..32d0f5dccc0 100644 --- a/ui/litellm-dashboard/src/components/vector_store_providers.test.tsx +++ b/ui/litellm-dashboard/src/components/vector_store_providers.test.tsx @@ -35,7 +35,8 @@ describe("getVectorStoreProviderLogoAndName", () => { }); expect(vectorStoreProviderMap.MongoDB).toBe("mongodb"); expect(getProviderSpecificFields("mongodb").map((field) => field.name)).toEqual([ - "mongodb_connection_string", + "api_base", + "api_key", "mongodb_database", "mongodb_collection", "embedding_model", @@ -45,12 +46,10 @@ describe("getVectorStoreProviderLogoAndName", () => { ]); }); - it("hides the mongodb connection string, which carries the database password", () => { - const connectionString = getProviderSpecificFields("mongodb").find( - (field) => field.name === "mongodb_connection_string", - ); + it("hides the mongodb sidecar API key", () => { + const apiKey = getProviderSpecificFields("mongodb").find((field) => field.name === "api_key"); - expect(connectionString).toMatchObject({ type: "password", required: true }); + expect(apiKey).toMatchObject({ type: "password", required: true }); }); it("picks the mongodb embedding model from the proxy's models rather than a fixed list", () => { diff --git a/ui/litellm-dashboard/src/components/vector_store_providers.tsx b/ui/litellm-dashboard/src/components/vector_store_providers.tsx index a75f10771a8..6a8b2f405d2 100644 --- a/ui/litellm-dashboard/src/components/vector_store_providers.tsx +++ b/ui/litellm-dashboard/src/components/vector_store_providers.tsx @@ -14,7 +14,7 @@ export enum VectorStoreProviders { OpenAI = "OpenAI", Azure = "Azure OpenAI", Milvus = "Milvus", - MongoDB = "MongoDB Atlas", + MongoDB = "MongoDB (BETA)", Valkey = "Valkey", } @@ -175,18 +175,25 @@ export const vectorStoreProviderFields: Record ], mongodb: [ { - name: "mongodb_connection_string", - label: "Connection String", - tooltip: - "The full MongoDB connection string for your Atlas cluster, including the database user and password. Copy it from Atlas under Connect, Drivers (e.g. mongodb+srv://user:password@cluster.mongodb.net)", - placeholder: "mongodb+srv://user:password@cluster.mongodb.net", + name: "api_base", + label: "Sidecar URL", + tooltip: "The URL of your separately deployed MongoDB sidecar. Configure MongoDB credentials in the sidecar", + placeholder: "http://mongodb-sidecar:8080", + required: true, + type: "text", + }, + { + name: "api_key", + label: "Sidecar API Key", + tooltip: "The MONGODB_SIDECAR_API_KEY configured in your MongoDB sidecar", + placeholder: "Enter sidecar API key", required: true, type: "password", }, { name: "mongodb_database", label: "Database", - tooltip: "The Atlas database holding the collection you want to search", + tooltip: "The MongoDB database holding the collection you want to search", placeholder: "sample_mflix", required: true, type: "text", @@ -194,7 +201,7 @@ export const vectorStoreProviderFields: Record { name: "mongodb_collection", label: "Collection", - tooltip: "The collection your Atlas Vector Search index was built on", + tooltip: "The collection your MongoDB Vector Search index was built on", placeholder: "embedded_movies", required: true, type: "text", @@ -212,7 +219,7 @@ export const vectorStoreProviderFields: Record name: "mongodb_embedding_field", label: "Vector Field Name", tooltip: - "The field in each document that holds its embedding. It must match the path your Atlas Vector Search index was created on (default: embedding)", + "The field in each document that holds its embedding. It must match the path your MongoDB Vector Search index was created on (default: embedding)", placeholder: "embedding", required: false, type: "text", @@ -232,7 +239,7 @@ export const vectorStoreProviderFields: Record name: "mongodb_num_candidates", label: "Candidates Considered", tooltip: - "How many nearest neighbours Atlas examines before returning the top results. Higher is more accurate and slower. Leave blank to let LiteLLM scale it with the requested result count", + "How many nearest neighbours MongoDB examines before returning the top results. Higher is more accurate and slower. Leave blank to let LiteLLM scale it with the requested result count", placeholder: "100", required: false, type: "text", diff --git a/uv.lock b/uv.lock index 89205cd9527..ce5d96f4d9c 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-05T05:15:35.833796Z" exclude-newer-span = "P3D" [manifest] @@ -4415,9 +4415,6 @@ mcp = [ mlflow = [ { name = "mlflow" }, ] -mongodb = [ - { name = "pymongo" }, -] proxy = [ { name = "apscheduler" }, { name = "azure-identity" }, @@ -4649,7 +4646,6 @@ requires-dist = [ { name = "pydantic", specifier = ">=2.10.0,<3.0.0" }, { name = "pydantic-settings", specifier = ">=2.14.1,<3.0" }, { name = "pyjwt", marker = "extra == 'proxy'", specifier = ">=2.13.0,<3.0" }, - { name = "pymongo", marker = "extra == 'mongodb'", specifier = ">=4.9,<5.0" }, { name = "pynacl", marker = "extra == 'proxy'", specifier = ">=1.6.2,<2.0" }, { name = "pypdf", marker = "extra == 'proxy-runtime'", specifier = ">=6.16.1,<7.0" }, { name = "pyroscope-io", marker = "sys_platform != 'win32' and extra == 'proxy'", specifier = ">=0.8.16,<1.0" }, @@ -4676,7 +4672,7 @@ requires-dist = [ { name = "uvloop", marker = "sys_platform != 'win32' and extra == 'proxy'", specifier = ">=0.22.1,<1.0" }, { name = "websockets", marker = "extra == 'proxy'", specifier = ">=15.0.1,<16.0" }, ] -provides-extras = ["proxy", "cli", "extra-proxy", "utils", "caching", "mcp", "mongodb", "saml", "semantic-router", "mlflow", "grpc", "stt-nvidia-riva", "google", "bedrock-realtime", "proxy-runtime"] +provides-extras = ["proxy", "cli", "extra-proxy", "utils", "caching", "mcp", "saml", "semantic-router", "mlflow", "grpc", "stt-nvidia-riva", "google", "bedrock-realtime", "proxy-runtime"] [package.metadata.requires-dev] ci = [ @@ -7620,77 +7616,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/d5/6f/9ac2548e290764781f9e7e2aaf0685b086379dabfb29ca38536985471eaf/pylint-4.0.5-py3-none-any.whl", hash = "sha256:00f51c9b14a3b3ae08cff6b2cdd43f28165c78b165b628692e428fb1f8dc2cf2", size = 536694, upload-time = "2026-02-20T09:07:31.028Z" }, ] -[[package]] -name = "pymongo" -version = "4.17.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "dnspython" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/ca/64/50be6fbac9c79fe2e4c17401a467da2d8764d82833d83cec325afe5cab32/pymongo-4.17.0.tar.gz", hash = "sha256:70ffa08ba641468cc068cf46c06b34f01a8ce3489f6411309fcb5ceabe6b2fc0", size = 2523370, upload-time = "2026-04-20T16:39:53.524Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/c9/77/28ebbf69772a4341d530831c7a006cdb06877ac23075cb53b0a227df4fe1/pymongo-4.17.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:47b021363cd923ace5edc7a1d63c0ff8a6d9d43859b8a1ba23645f5afae63221", size = 819234, upload-time = "2026-04-20T16:37:20.888Z" }, - { url = "https://files.pythonhosted.org/packages/88/cf/5a70cee503ff9a2fea20607607f14d189f4d975960ac0945ec306ee7b695/pymongo-4.17.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:422fa50d7d7f5c22ea0953554396c9ef95684a2d775f860bd75a7b510538dfca", size = 819969, upload-time = "2026-04-20T16:37:24.187Z" }, - { url = "https://files.pythonhosted.org/packages/23/d5/07b7e27e662c58d872efd104a0e8055eb6569aa1b6d4da436f3fdee7f897/pymongo-4.17.0-cp310-cp310-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:addd0498ebbdc6354227f6ed457ed9fce442d48a3bb30d5b5bad33e104996561", size = 1244510, upload-time = "2026-04-20T16:37:26.069Z" }, - { url = "https://files.pythonhosted.org/packages/fb/be/7cac5b1e89bd5a8e395067648241390321593a7c29243e36f91343c02a90/pymongo-4.17.0-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c5c8e180cb2cabe37300e1e36c60aa4f2ff956cc579f0142135a5d2cba252243", size = 1263245, upload-time = "2026-04-20T16:37:28.003Z" }, - { url = "https://files.pythonhosted.org/packages/2e/20/40e8e99824c1fda18261411e65ce3b0cd3d9a6ed3c056cdd0a569adc870b/pymongo-4.17.0-cp310-cp310-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:bd835cdb37a1adec359dd072c24f8bb14809e2644fde86fab4ee2fc9719b9483", size = 1304113, upload-time = "2026-04-20T16:37:30.048Z" }, - { url = "https://files.pythonhosted.org/packages/3a/94/fb7e25441dd66f2069a9b172380849b0eaa5881c18b3db217bf64a6d393c/pymongo-4.17.0-cp310-cp310-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:c4979e7e8887862bbb44d203f00cc8263a3f27237876fa691b6beba23e40e6d8", size = 1297046, upload-time = "2026-04-20T16:37:32.054Z" }, - { url = "https://files.pythonhosted.org/packages/4f/c9/7352e0c20fe772541556e4d283c05e07ec48f8b0d2737ad930ac4a1b6655/pymongo-4.17.0-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:77aa4bc164b4de60d5db193b322f0f5b6ead716e831031bfdef8e8bd92205556", size = 1265708, upload-time = "2026-04-20T16:37:33.934Z" }, - { url = "https://files.pythonhosted.org/packages/8d/e4/3df15494c2015ed297958517f0e4f6493e21b00990748068a973e66d45e0/pymongo-4.17.0-cp310-cp310-win32.whl", hash = "sha256:48bbc576677b50af043df870d84ded67cc3a9b4aa7553201beef4da5dc050a0a", size = 805533, upload-time = "2026-04-20T16:37:35.744Z" }, - { url = "https://files.pythonhosted.org/packages/22/fa/b4e71bb8cb82ad7d21bb4e8c476f2d573ba68b20368aac36ef06e4a196b4/pymongo-4.17.0-cp310-cp310-win_amd64.whl", hash = "sha256:e46767f28dea610e02edf6c5d956ce615c3c7790ea396660b9b1efd5c5ead2e0", size = 815677, upload-time = "2026-04-20T16:37:37.808Z" }, - { url = "https://files.pythonhosted.org/packages/22/e2/0a4bba644f1cda3970ea1012149eeae3594ebfeed3f81fdaf32b61d90c95/pymongo-4.17.0-cp310-cp310-win_arm64.whl", hash = "sha256:757f2a4c0c2c46cab87df0333681ce69e86c9d5b45bc5203ceba5410b3489e59", size = 807293, upload-time = "2026-04-20T16:37:39.707Z" }, - { url = "https://files.pythonhosted.org/packages/c4/e2/336d86f221cf1b56b2ed9330d4a3b98f9f38f0b37829ae9a9184617d5419/pymongo-4.17.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:4141e6c6a339789b2974efa00ecd9409101672d77a0e3ee2cc3839eedf8ec4df", size = 874668, upload-time = "2026-04-20T16:37:41.39Z" }, - { url = "https://files.pythonhosted.org/packages/34/8e/75d3c6c935d187ab59c61e9c15d9aab3f274b563eaf1706e8cae5f508dec/pymongo-4.17.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:e68c76b84e0c132d9dbf9307f12ff8185702328187a87b9aca8c941303873433", size = 875294, upload-time = "2026-04-20T16:37:43.432Z" }, - { url = "https://files.pythonhosted.org/packages/5f/ec/62e855744489dbcd54fd778aae4d80fa4c4819e8fb228ca0cf6f21a03997/pymongo-4.17.0-cp311-cp311-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:ba2195d4f386f839a52a23ea1cfd60ffaaba78a3d7841db51b7e433001139918", size = 1496233, upload-time = "2026-04-20T16:37:45.518Z" }, - { url = "https://files.pythonhosted.org/packages/82/e8/93e4e5e5ce8fdf8929dabeefe24aafa5ce046028eed0dfa8eeb936e72c49/pymongo-4.17.0-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:8446ff4bfcb6ec2a2e50998c860986a1e992136f998b7f53e7a717fb8aa5a0b9", size = 1522927, upload-time = "2026-04-20T16:37:47.492Z" }, - { url = "https://files.pythonhosted.org/packages/f7/ca/425dc1d21e0f17bdea0072fc463f662f7fa06d2852af52975c9eced3c07c/pymongo-4.17.0-cp311-cp311-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:2a0d5ac205728c86e0a02192f1aa5f865b0d7d51f8df6101c01a69a7fc620d72", size = 1583468, upload-time = "2026-04-20T16:37:49.221Z" }, - { url = "https://files.pythonhosted.org/packages/b3/9d/f08b07eeffda1a43c1759f0fa625e88ae12360996eb56d42aad832fa7dff/pymongo-4.17.0-cp311-cp311-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:485c8a8eaa4c739f00a331fc73757898ee7c092c214a79e63866ff76aaf282ff", size = 1572787, upload-time = "2026-04-20T16:37:51.061Z" }, - { url = "https://files.pythonhosted.org/packages/e9/c2/6855a07aafa7b894929af23675b6fb9634800ce43122b76a62f6eeb8da2a/pymongo-4.17.0-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b2dfcc795f5b9fedbe179a11fdf6051581479d196582a3fe819a92a00e9b9969", size = 1526184, upload-time = "2026-04-20T16:37:53.358Z" }, - { url = "https://files.pythonhosted.org/packages/4e/05/c952bac7db71c1942ea3559fcd308b49754cc5004b455935fb4000d1f37b/pymongo-4.17.0-cp311-cp311-win32.whl", hash = "sha256:c2292144505fb12156b981bd440f3dc994a883da06ac726c0c8692ccdbc1c510", size = 852621, upload-time = "2026-04-20T16:37:55.28Z" }, - { url = "https://files.pythonhosted.org/packages/11/c0/c04da9f4c0c6252404598f4e394b862a58a9e866822a70ae261c8a018fdf/pymongo-4.17.0-cp311-cp311-win_amd64.whl", hash = "sha256:2e190827834fce70ecdf9d46796c6dbc0ce08ea87dc2ff5bc6f3f5579b605cb9", size = 867852, upload-time = "2026-04-20T16:37:57.233Z" }, - { url = "https://files.pythonhosted.org/packages/1d/b2/c7b4870fbeef471e947d3e014676f5910d02e0197074d692ebcf24ec049a/pymongo-4.17.0-cp311-cp311-win_arm64.whl", hash = "sha256:a8f9c40a09bb7d4b9fc8b1da65ecf6efa79bda5cb2756f39d9b6940fac1d19ae", size = 855019, upload-time = "2026-04-20T16:37:58.983Z" }, - { url = "https://files.pythonhosted.org/packages/98/90/60bcb508840135d5ee46b51b1a950f548338aa8145a8366dbe6639ae51ac/pymongo-4.17.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:d53ffa94b2340dbf6b055e09a0090618c60482c158ecfc9565642fc996bf0944", size = 930529, upload-time = "2026-04-20T16:38:00.936Z" }, - { url = "https://files.pythonhosted.org/packages/a6/e9/313840f1e52c6dfac47f704428cbfbce59956ebe7633bffc92b03f74f0ad/pymongo-4.17.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:6fe0de9d0f6791abce3471230b32b4817bf89d27b1182b6a550e1ec0fa72aa9a", size = 930665, upload-time = "2026-04-20T16:38:02.915Z" }, - { url = "https://files.pythonhosted.org/packages/78/35/9d3565ea45b1606f635c1e2cd2563c28d66caafdc50f7ad7d979fcd1b363/pymongo-4.17.0-cp312-cp312-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:e537e95514dae1aaa718f481ec03151a0f0394bcd05f1322896d8fc1330cb729", size = 1762369, upload-time = "2026-04-20T16:38:05.375Z" }, - { url = "https://files.pythonhosted.org/packages/95/ee/149b0d4b1a11c38bff6f14c23d5814c9b0843fd6dc38ad40596bdb1a62d2/pymongo-4.17.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:37a8385c29881b43eab31f584100fa0eaddedd5607adf010147ba1810118be90", size = 1798044, upload-time = "2026-04-20T16:38:07.195Z" }, - { url = "https://files.pythonhosted.org/packages/7b/d4/4cee4a7b8d8f6f0550ef6cd2fea42455c5ed619a220cb6ba4fb40d6a5bc8/pymongo-4.17.0-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:f3ee3d241ed77a4fc99ce3cff3b289c3ebce37f61fdd7349d3592c23b82c8784", size = 1878567, upload-time = "2026-04-20T16:38:09.121Z" }, - { url = "https://files.pythonhosted.org/packages/45/ef/7fe366c84952619ee2f69973566c214775e083dd4df465751912153e4b72/pymongo-4.17.0-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:9eb5d63a3c518cb0804ed678f5e2b875af032d89a7cf57a57360322cf6a4d222", size = 1864881, upload-time = "2026-04-20T16:38:10.896Z" }, - { url = "https://files.pythonhosted.org/packages/2f/35/b577d82c6d1be7aee7ac7e249bc86f7847998345042e5f8360de238e177b/pymongo-4.17.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:8e97e03fa13327c87e3fdc5656acd01e71817f0c1dc3221cd8f30de136bf4ec3", size = 1800349, upload-time = "2026-04-20T16:38:13.589Z" }, - { url = "https://files.pythonhosted.org/packages/b8/69/dafcf04f66e130ddd91aeb92e7a692480eda46dcd04ec1dbe82c06619e10/pymongo-4.17.0-cp312-cp312-win32.whl", hash = "sha256:6877214bff5f06f6884a9fc8d9016a4a7a5f51f537f5c51ac3a576f93e7dfb32", size = 900518, upload-time = "2026-04-20T16:38:15.541Z" }, - { url = "https://files.pythonhosted.org/packages/11/35/5c9262a459f988b4eb2605f70815240b77a0d4131136c4326d18f1822b89/pymongo-4.17.0-cp312-cp312-win_amd64.whl", hash = "sha256:9828485f72f63c7d802e0ec41f71906f633c2692621ab3af55ca990186b091b1", size = 920335, upload-time = "2026-04-20T16:38:17.665Z" }, - { url = "https://files.pythonhosted.org/packages/8d/da/e9c7265ee176faccf4e52c4797837e794d93569a1046f6b19a4acc36e5ad/pymongo-4.17.0-cp312-cp312-win_arm64.whl", hash = "sha256:1195370a77baf003b59b10e91ecc4706297197f0dd9d29c840cc556dc08f7cee", size = 903289, upload-time = "2026-04-20T16:38:19.33Z" }, - { url = "https://files.pythonhosted.org/packages/2a/6b/c1206879708b94e82fcd8b9653440ec271f79a3674d122192df383047f5a/pymongo-4.17.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:809ec74de3b9148ae43fa8df9faf53470f511c8d384f13b99d6f671f2a379f15", size = 985829, upload-time = "2026-04-20T16:38:21.031Z" }, - { url = "https://files.pythonhosted.org/packages/cb/cf/bb044ed85160e5c40f568c7c4f4e8ea16f40764ff5d302e5befbe8f6f814/pymongo-4.17.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:a431b737816bf4cddd4fa0fcef04e424ad36b7692734a64150f872fb8f3208be", size = 985899, upload-time = "2026-04-20T16:38:23.409Z" }, - { url = "https://files.pythonhosted.org/packages/74/0a/f6dfd5ea3901e5d6888da8de8ba728971a1d447debab681cfc56f90d1208/pymongo-4.17.0-cp313-cp313-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:e4fab10f8403169ce92f3cea921609d9ee81107306caae06c08f592d4b8ad2b5", size = 2028569, upload-time = "2026-04-20T16:38:25.343Z" }, - { url = "https://files.pythonhosted.org/packages/4a/c5/081f59a1c02ae8c0dc73ae58e563838c44eec81aeafa7d0b93a637841c9b/pymongo-4.17.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:20323b0b1c1d33770ad1fc68d429c757734ce9ad3594421c3d6618f10572b1b9", size = 2072916, upload-time = "2026-04-20T16:38:27.291Z" }, - { url = "https://files.pythonhosted.org/packages/31/42/6e41d434297ffe8b30d9c3717916591a4a7be9075a0dcc2fafdfaaaa62ed/pymongo-4.17.0-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:5a5de048e6da5c18e27cc2437e8c15b3b0cdc8385c15b41178b0caa3322a09c2", size = 2173234, upload-time = "2026-04-20T16:38:29.474Z" }, - { url = "https://files.pythonhosted.org/packages/3d/cf/1e4a7db352ef9485831c7268dfe8402f0117b32a9ad54b16e810699e3617/pymongo-4.17.0-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:dff3de1294fbbc1db0ba6b511f77b8e540601d092538a31312e99c8a91a78b1e", size = 2156784, upload-time = "2026-04-20T16:38:32.134Z" }, - { url = "https://files.pythonhosted.org/packages/12/10/6195be29962a61ebb5f4bd9e4c7519890b172f7968a0a0d880398c6ddb02/pymongo-4.17.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:faf03e4c2aafd6de626dbd30ba246d369ae33f47f10629d1bbe40f72115027a6", size = 2074446, upload-time = "2026-04-20T16:38:34.004Z" }, - { url = "https://files.pythonhosted.org/packages/37/48/33410b8819837ed370c738587306bdf060b59cef11823be212f4a07703c5/pymongo-4.17.0-cp313-cp313-win32.whl", hash = "sha256:c9786665926a09630c5d420c79762cfadbff35a9438bcbc4c81a9fb5ab9228b7", size = 948435, upload-time = "2026-04-20T16:38:35.922Z" }, - { url = "https://files.pythonhosted.org/packages/6f/77/c0ed522f798a286b99acaa7914ed8d9c80ab091f97f57c59ffed72906e5e/pymongo-4.17.0-cp313-cp313-win_amd64.whl", hash = "sha256:5960519b4d7168f1ecdd3ea10c81b2aedeb9423651aca953cfbc8e76705d3b38", size = 972847, upload-time = "2026-04-20T16:38:37.888Z" }, - { url = "https://files.pythonhosted.org/packages/97/f0/c39480a2db385fde23861d0c8acda41cdaf1d43e46579db72c5c013a2e81/pymongo-4.17.0-cp313-cp313-win_arm64.whl", hash = "sha256:0ff6bd2f735ab5356541e3e57d5b7dbfbc3f2ee1ccb10b6b0f82d58af69d1d8e", size = 951575, upload-time = "2026-04-20T16:38:40.544Z" }, - { url = "https://files.pythonhosted.org/packages/da/49/2b0250762a89737ed6f9cea238331baca061b89a8ddd10dd17fee52c3970/pymongo-4.17.0-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:ff5aa3f1c7e3f08eb0e7a016c91ba468b1850ccfd63d9b1f12f56350f4974cef", size = 1040945, upload-time = "2026-04-20T16:38:42.783Z" }, - { url = "https://files.pythonhosted.org/packages/89/1c/7a9b5447a08be20e84b6e5b17330917e8d6d9507daa3cd099a9309f11ad7/pymongo-4.17.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:e816db649ba5d7de0568cf3a9f287a9dc9aad21cf0ca667ab156a7ef47fca0b0", size = 1041187, upload-time = "2026-04-20T16:38:45.358Z" }, - { url = "https://files.pythonhosted.org/packages/78/a1/71704f61632dfc90407a5834fe5f6132854937c4a3648f6c05c351d85a45/pymongo-4.17.0-cp314-cp314-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:12c4fded3a9f1d6a687e36ebd384ac6d00b9b00de1969aa74048e7051ec2a713", size = 2294806, upload-time = "2026-04-20T16:38:47.734Z" }, - { url = "https://files.pythonhosted.org/packages/ad/b9/aff42be75108b96c2469b1d9329b912c15108f3e7ef32fdc86da8423c330/pymongo-4.17.0-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:2db66aa8dd253a0fc1fad3b0d23d5b3993f7ebde02fbbd7727128debf2853675", size = 2348231, upload-time = "2026-04-20T16:38:50.371Z" }, - { url = "https://files.pythonhosted.org/packages/f2/30/44c115b8ba1479942c15fd9480eb29a7da0ba68acd56983423ba0deb4a94/pymongo-4.17.0-cp314-cp314-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:3987e96e7c7be4083d42e8ac2cc6c0d5b78db9973c90fce42ae800b616ca6b20", size = 2467614, upload-time = "2026-04-20T16:38:52.665Z" }, - { url = "https://files.pythonhosted.org/packages/d2/84/21ee95c8bf0ca7acae7ec7eb365d740bf8fc0156c194baf2c3bdfcb85ec0/pymongo-4.17.0-cp314-cp314-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:cee36b3c0d0354f880fa7a7fdcdaf2bb5e542c2281e25c1bfadf8cfe21eba7d2", size = 2445970, upload-time = "2026-04-20T16:38:55.175Z" }, - { url = "https://files.pythonhosted.org/packages/06/89/081d7f1809d5ca09d1e47e49f2111b245f5694de3a7af32cd3a353a6f43f/pymongo-4.17.0-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:320b34457b20bbcc79997801f95d25ce00472915ca5241167242b42c4359e027", size = 2348605, upload-time = "2026-04-20T16:38:57.557Z" }, - { url = "https://files.pythonhosted.org/packages/ea/c3/0d949f9d3f2a341c1f635c398c16615e96f89f51ff424ed81e914cf1a4de/pymongo-4.17.0-cp314-cp314-win32.whl", hash = "sha256:df4a644af9ae132d4bfdb2e9516ea51a615fd881caddfbfbd071cf1354844479", size = 1004119, upload-time = "2026-04-20T16:39:00.309Z" }, - { url = "https://files.pythonhosted.org/packages/f7/55/5c3a3db1048054c695c75c5964cc8bedc2247fdb5a75ef6fab4ec8bb013e/pymongo-4.17.0-cp314-cp314-win_amd64.whl", hash = "sha256:c797f8a80957134f6dd9690367a0f8f5906d672119af2c6aa55f0c527b656bed", size = 1032314, upload-time = "2026-04-20T16:39:02.665Z" }, - { url = "https://files.pythonhosted.org/packages/e0/19/e235f39906134cb0ffd5574c5a59c355ef5380f0499644ab94994afbb109/pymongo-4.17.0-cp314-cp314-win_arm64.whl", hash = "sha256:68fca71e05ee5da23a8d73cee8379dfb3d26e609a377cae731d742771ed96946", size = 1007627, upload-time = "2026-04-20T16:39:04.678Z" }, - { url = "https://files.pythonhosted.org/packages/1e/e0/c4c1a86791415b14c684fa0908f9da96de91594a3fd1fa1b8dc689fbb800/pymongo-4.17.0-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:b4384700cffc3f1dd98e088bc0072dedf6d7d68a230bb4b972665cf69c071c1e", size = 1099151, upload-time = "2026-04-20T16:39:06.969Z" }, - { url = "https://files.pythonhosted.org/packages/81/4b/69c67f3e23fd9b23b9bedc7ebd23754881cc9d5c5d5b2a9811e96b07f475/pymongo-4.17.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:93641192644fa1ee0f34030e774fd31022a27ad11ba22cb1716142231524f8bd", size = 1099346, upload-time = "2026-04-20T16:39:08.996Z" }, - { url = "https://files.pythonhosted.org/packages/a2/19/a5208f62f9508a26d73acc69bd3821b8c8adae253679a3c26d2f9652f0d5/pymongo-4.17.0-cp314-cp314t-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:75bc3aa5b94fdb7138d357ec6ca61cd97e0c79f4f7f0bd3efe9639b15cc50942", size = 2619034, upload-time = "2026-04-20T16:39:11.049Z" }, - { url = "https://files.pythonhosted.org/packages/77/27/426cba1ec5973082a56d4150798529bfdf4151c31391ed1fbbecb23ef2ac/pymongo-4.17.0-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:50e8f8e23c6df7c6d6929f5e734980b227706e73ee847517c9ba5af90f7fc466", size = 2689939, upload-time = "2026-04-20T16:39:13.617Z" }, - { url = "https://files.pythonhosted.org/packages/ef/2e/f70993d1255e33f6ee59a4ec4371cc65bff7a7e3fda7d55c3386f25287e8/pymongo-4.17.0-cp314-cp314t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:15d3f3d732aecac1f8d481bde4029755615639bd3076f258a2147210aec8515a", size = 2824994, upload-time = "2026-04-20T16:39:16.057Z" }, - { url = "https://files.pythonhosted.org/packages/b3/eb/87b0e988ba889e1fcc3430c2cfc166b251872c813e92b43174298bee17ff/pymongo-4.17.0-cp314-cp314t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:6c5f62862d0f87be481fa1fe8cb811994486773c94a2b61e509285e3f2890763", size = 2801745, upload-time = "2026-04-20T16:39:18.476Z" }, - { url = "https://files.pythonhosted.org/packages/67/4c/3f83412d086f682d4d468761d66ddc49cf161e786ea74073045eb4491c60/pymongo-4.17.0-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:64837adbbd72073301af51bb0fc80e3d7707fe5527cea1033ba0320f0b2f881b", size = 2684636, upload-time = "2026-04-20T16:39:20.878Z" }, - { url = "https://files.pythonhosted.org/packages/9e/d8/b75f6f4ab6c8beb50b0270a4f1e2530b5774f5e116563440e1677ca1820f/pymongo-4.17.0-cp314-cp314t-win32.whl", hash = "sha256:b93b22eedc62598cf5ee9d8c8007a8e9121c50fd88137012d8985500e9dc3151", size = 1056356, upload-time = "2026-04-20T16:39:22.996Z" }, - { url = "https://files.pythonhosted.org/packages/e4/5e/648c8a238eef18a25ed8a169ea6542d4a860bbec3e95b3d9badac2935c71/pymongo-4.17.0-cp314-cp314t-win_amd64.whl", hash = "sha256:3689ea34f6b647c7d1e7bdc60fcfb214b2789ed1359a7fb96569c69f50e5f18f", size = 1090964, upload-time = "2026-04-20T16:39:24.989Z" }, - { url = "https://files.pythonhosted.org/packages/dc/cb/d9780b66939c4fc1f024bcc7be23a2abcfe06a9745ca8fa76dc73395482e/pymongo-4.17.0-cp314-cp314t-win_arm64.whl", hash = "sha256:9543d8f84c2e5608565c08ac679774811e6730770d8a645439b073422a4276fb", size = 1058526, upload-time = "2026-04-20T16:39:27.924Z" }, -] - [[package]] name = "pynacl" version = "1.6.2" From 1a6aa98230571db22ccd37e5db2c6011a4f0c4c4 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Mon, 7 Sep 2026 23:29:42 -0700 Subject: [PATCH 065/109] fix(spend): compare auto-router targets by deployment identity (#40206) Preserve deployment identity through savings calculation, with canonical model fallback only when either ID is absent. Cover negotiated rates, unchanged deployments, alias/base-model cache accounting and missing IDs. Fixes #38811. Based on the deployment-identity approach proposed by @QuantumBreakz in #38834. Co-authored-by: Claude Code --- litellm/proxy/spend_tracking/savings.py | 15 ++- .../proxy/spend_tracking/test_savings.py | 98 +++++++++++++++++++ 2 files changed, 108 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/spend_tracking/savings.py b/litellm/proxy/spend_tracking/savings.py index c541b9b40e5..950fcca2039 100644 --- a/litellm/proxy/spend_tracking/savings.py +++ b/litellm/proxy/spend_tracking/savings.py @@ -299,6 +299,8 @@ def compute_autorouter_savings( selected_info: ModelInfo | None = None, baseline_info: ModelInfo | None = None, cost_breakdown: Mapping[str, object] | None = None, + baseline_deployment_id: str | None = None, + selected_deployment_id: str | None = None, ) -> float: """Net dollars the router saved, or cost, by serving this request on ``selected_model``. @@ -334,11 +336,12 @@ def compute_autorouter_savings( selected: Final = _resolve_model(selected_model, selected_provider) if baseline is None or selected is None: return 0.0 - # Same model is only the same cost when it is also the same deployment. Two - # deployments of one model can carry different negotiated rates, and routing from - # the dear one to the cheap one is a real saving that short-circuiting on the model - # name alone reports as zero. - if baseline == selected: + same_target: Final = ( + baseline_deployment_id == selected_deployment_id + if baseline_deployment_id and selected_deployment_id + else baseline == selected + ) + if same_target: return 0.0 basis: Final = _pricing_basis(cost_breakdown) effective_baseline_info: Final = baseline_info if baseline_info is not None else _model_info(baseline) @@ -517,6 +520,8 @@ def autorouter_savings_for_request( selected_info=_effective_model_info(router_instance, model_id, model or ""), baseline_info=_effective_model_info(router_instance, baseline_id, baseline_model or ""), cost_breakdown=cost_breakdown, + baseline_deployment_id=baseline_id, + selected_deployment_id=model_id, ) classifier_cost: Final = classifier_cost_from_decision(decision) return gross if classifier_cost is None else gross - classifier_cost diff --git a/tests/test_litellm/proxy/spend_tracking/test_savings.py b/tests/test_litellm/proxy/spend_tracking/test_savings.py index cc8fdeb0160..e466edab131 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_savings.py +++ b/tests/test_litellm/proxy/spend_tracking/test_savings.py @@ -1162,6 +1162,104 @@ def test_prompt_caching_prices_at_the_deployment_rate_not_the_public_one(): assert result.prompt_caching > at_public_rates.prompt_caching +@pytest.mark.parametrize( + "baseline_id, selected_id, selected_multiplier, billed_input, classifier_cost, expected", + [ + ("baseline", "selected", 0.1, None, 0.0, 0.0135), + ("baseline", "selected", 2.0, None, 0.0, -0.015), + ("baseline", "selected", 1.0, None, 0.0, 0.0), + ("baseline", "selected", 0.1, 0.004, 0.001, 0.01), + ("baseline", "baseline", 0.1, 0.004, 0.001, -0.001), + (None, "selected", 0.1, None, 0.0, 0.0), + ("baseline", None, 0.1, None, 0.0, 0.0), + (None, None, 0.1, None, 0.0, 0.0), + ("", "selected", 0.1, None, 0.0, 0.0), + ("baseline", "", 0.1, None, 0.0, 0.0), + ], +) +def test_autorouter_savings_distinguishes_priced_deployments( + baseline_id: str | None, + selected_id: str | None, + selected_multiplier: float, + billed_input: float | None, + classifier_cost: float, + expected: float, +) -> None: + router: Final = Router( + model_list=[ + { + "model_name": name, + "litellm_params": { + "model": "anthropic/claude-opus-5", + "api_key": "test-key", + "input_cost_per_token": 1e-5 * multiplier, + "output_cost_per_token": 5e-5 * multiplier, + }, + "model_info": {"id": name}, + } + for name, multiplier in (("baseline", 1.0), ("selected", selected_multiplier)) + ] + ) + result: Final = compute_savings_spend( + model="claude-opus-5", + custom_llm_provider="anthropic", + compression_saved_tokens=0, + gateway_injected_cache=False, + model_id=selected_id, + llm_router=lambda: router, + routing_decision={ + "savings_baseline_model": "anthropic/claude-opus-5", + "savings_baseline_deployment_id": baseline_id, + "conversation_continuing": False, + "classifier_cost": classifier_cost, + }, + usage_object={"prompt_tokens": 1000, "completion_tokens": 100, "total_tokens": 1100}, + cost_breakdown=None if billed_input is None else {"input_cost": billed_input, "output_cost": 0.0}, + ) + assert result.autorouter == pytest.approx(expected) + + +@pytest.mark.parametrize("selected_model", ["azure/contract-deployment", "contract-deployment"]) +def test_autorouter_savings_recognizes_one_deployment_under_its_base_model(selected_model: str) -> None: + router: Final = Router( + model_list=[ + { + "model_name": "contract", + "litellm_params": { + "model": "azure/contract-deployment", + "api_key": "test-key", + "api_base": "https://example.openai.azure.com", + "input_cost_per_token": 0.0001, + "output_cost_per_token": 0.0002, + "cache_read_input_token_cost": 0.00001, + }, + "model_info": {"id": "contract", "base_model": "azure/gpt-5.5"}, + } + ] + ) + result: Final = compute_savings_spend( + model=selected_model, + custom_llm_provider="azure", + compression_saved_tokens=0, + gateway_injected_cache=False, + model_id="contract", + llm_router=lambda: router, + routing_decision={ + "savings_baseline_model": "azure/gpt-5.5", + "savings_baseline_deployment_id": "contract", + "conversation_continuing": True, + }, + usage_object={ + "prompt_tokens": 21000, + "completion_tokens": 100, + "total_tokens": 21100, + "prompt_tokens_details": {"text_tokens": 1000, "cached_tokens": 0, "cache_creation_tokens": 20000}, + }, + cost_breakdown={"input_cost": 2.1, "output_cost": 0.02}, + ) + assert result.autorouter == 0.0 + + def test_a_recorded_baseline_deployment_prices_at_its_configured_rate(): """A hardest-tier deployment with a negotiated rate is what the traffic would really have cost; pricing its model publicly misstates the saving.""" From 9a9b4c4c2538bf0df6d68eaabe48cefbb4824e7d Mon Sep 17 00:00:00 2001 From: tin-berri Date: Mon, 7 Sep 2026 23:37:42 -0700 Subject: [PATCH 066/109] feat(ui): show auto-router classification rate (#40192) --- .../AutoRouterBenchmarksTab.test.tsx | 12 ++++----- .../_components/AutoRouterBenchmarksTab.tsx | 26 ++++++++++++++----- .../_components/costOptimizationUtils.test.ts | 18 +++++++++++++ .../_components/costOptimizationUtils.ts | 7 +++++ ...KeyAutoRouterUsageTab.integration.test.tsx | 2 +- 5 files changed, 52 insertions(+), 13 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.test.tsx index 006da4f2725..2820a9dce83 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.test.tsx @@ -187,10 +187,10 @@ describe("AutoRouterBenchmarksTab", () => { }); it.each([ - { spend: 20665.28, classifier_cost: 342.18, turns: 140815, llm: "$20,323.10", cost: "$342.18" }, - { spend: 0, classifier_cost: 0, turns: 0, llm: "$0.00", cost: "$0.00" }, - { spend: 0.002, classifier_cost: 0.0004, turns: 100, llm: "$0.0016", cost: "$0.0004" }, - ])("shows total classification cost across $turns turns without a per-turn rate", ({ llm, cost, ...values }) => { + { spend: 20665.28, classifier_cost: 342.18, turns: 140815, llm: "$20,323.10", cost: "$342.18", rate: "$2.43" }, + { spend: 0, classifier_cost: 0, turns: 0, llm: "$0.00", cost: "$0.00", rate: "$0.00" }, + { spend: 0.002, classifier_cost: 0.0004, turns: 100, llm: "$0.0016", cost: "$0.0004", rate: "$0.0040" }, + ])("shows total classification cost and its rate across $turns turns", ({ llm, cost, rate, ...values }) => { const stats = totals({ ...values, saved_spend: 10126.28, baseline_spend: values.spend + 10126.28 }); mockHook({ data: response([group(stats)], stats) }); renderTab(); @@ -201,7 +201,7 @@ describe("AutoRouterBenchmarksTab", () => { .map((node) => node.textContent) .slice(1, 3), ).toEqual([llm, cost]); - expect(screen.queryByText(/1K turns/)).not.toBeInTheDocument(); + expect(screen.getByText(`(${rate} / 1K turns)`)).toBeInTheDocument(); expect(screen.getAllByText("$10,126.28").length).toBeGreaterThan(0); }); @@ -237,7 +237,7 @@ describe("AutoRouterBenchmarksTab", () => { expect(terms).toEqual([ "Actual auto-router spend", "LLM spend", - "Classification cost", + "Classification cost($2.00 / 1K turns)", "Estimated spend at highest-tier model", ]); expect(values).toEqual(["$359.86", "$353.71", "$6.15", "$2,534.45"]); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.tsx index 33ad1bfe555..ce5ab1c6776 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.tsx @@ -30,7 +30,7 @@ import { type BenchmarkView, type BucketRow, } from "./autoRouterBenchmarks"; -import { formatRangeLabel, usd } from "./costOptimizationUtils"; +import { classificationRatePer1kTurns, formatRangeLabel, usd } from "./costOptimizationUtils"; import ShadowEvalSection from "./ShadowEvalSection"; import TierTurnsChart from "./TierTurnsChart"; import { useAutoRouterBenchmarks } from "./useAutoRouterBenchmarks"; @@ -52,9 +52,17 @@ const Metric: React.FC<{ label: string; value: string; hint?: string }> = ({ lab ); -const SpendRow: React.FC<{ label: string; value: string; subdued?: boolean }> = ({ label, value, subdued }) => ( +const SpendRow: React.FC<{ label: string; value: string; hint?: string; subdued?: boolean }> = ({ + label, + value, + hint, + subdued, +}) => (
-
{label}
+
+ {label} + {hint && {hint}} +
@@ -99,6 +107,11 @@ const HeroCard: React.FC<{ view: BenchmarkView }> = ({ view }) => { subdued label="Classification cost" value={stats.classifier_cost == null ? "Unavailable" : usd(stats.classifier_cost)} + hint={ + stats.classifier_cost == null + ? undefined + : classificationRatePer1kTurns(stats.classifier_cost, stats.turns) + } />
{stats.classifier_cost == null && ( @@ -277,9 +290,10 @@ const BenchmarksBody: React.FC = ({ isPending, error, data,

Compares your actual routed spend with the estimated cost of using only the most expensive model configured in the auto-router. It accounts for both the cache savings from staying on one model and the added cache costs from - switching models. Savings are net of recorded LLM classification cost, which is included in actual spend. The - range counts whole sessions that overlap it, so totals can differ slightly from the Overall tab, which buckets - savings by UTC day. + switching models. Savings are net of recorded LLM classification cost, which is included in actual spend. + Classification cost per 1K turns is averaged over all auto-router turns, including those that skip + classification. The range counts whole sessions that overlap it, so totals can differ slightly from the Overall + tab, which buckets savings by UTC day.

diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/costOptimizationUtils.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/costOptimizationUtils.test.ts index 9a2d7a0b0ec..5d2c48e6440 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/costOptimizationUtils.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/costOptimizationUtils.test.ts @@ -7,6 +7,7 @@ import { SAVINGS_DRIVERS, SAVINGS_SERIES, buildDailyToolSeries, + classificationRatePer1kTurns, computeCacheLeakage, formatRangeLabel, isAnthropicModel, @@ -401,6 +402,23 @@ describe("usd", () => { }); }); +describe("classificationRatePer1kTurns", () => { + it("normalizes total classification cost to one thousand turns", () => { + expect(classificationRatePer1kTurns(342.18, 140815)).toBe("($2.43 / 1K turns)"); + expect(classificationRatePer1kTurns(0.0004, 100)).toBe("($0.0040 / 1K turns)"); + }); + + it("shows a floor instead of rounding a real cost down to zero", () => { + expect(classificationRatePer1kTurns(0.00001, 1000)).toBe("(<$0.0001 / 1K turns)"); + expect(classificationRatePer1kTurns(0.0001, 1000)).toBe("($0.0001 / 1K turns)"); + }); + + it("reports zero when there are no turns or no classification cost", () => { + expect(classificationRatePer1kTurns(0, 0)).toBe("($0.00 / 1K turns)"); + expect(classificationRatePer1kTurns(0, 100)).toBe("($0.00 / 1K turns)"); + }); +}); + describe("savings driver colours", () => { it("keeps a driver's colour when a driver above it is filtered out", () => { // Charts colour by position in the data they are given, and the donut is given diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/costOptimizationUtils.ts b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/costOptimizationUtils.ts index 7019b0d3301..464c779aa2b 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/costOptimizationUtils.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/costOptimizationUtils.ts @@ -10,6 +10,13 @@ export const usd = (value: number): string => { return `${value < 0 ? "-" : ""}$${formatNumberWithCommas(magnitude, decimals)}`; }; +export const classificationRatePer1kTurns = (classifierCost: number, turns: number): string => { + if (turns <= 0) return `(${usd(0)} / 1K turns)`; + const rate = (classifierCost * 1000) / turns; + if (rate > 0 && rate < 0.0001) return "(<$0.0001 / 1K turns)"; + return `(${usd(rate)} / 1K turns)`; +}; + export const pct = (ratio: number): string => `${formatNumberWithCommas(ratio * 100, 1)}%`; export const shortDate = (iso: string): string => diff --git a/ui/litellm-dashboard/src/components/templates/KeyAutoRouterUsageTab.integration.test.tsx b/ui/litellm-dashboard/src/components/templates/KeyAutoRouterUsageTab.integration.test.tsx index be95c0e600d..1f6a67b27ac 100644 --- a/ui/litellm-dashboard/src/components/templates/KeyAutoRouterUsageTab.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/KeyAutoRouterUsageTab.integration.test.tsx @@ -89,7 +89,7 @@ describe("KeyAutoRouterUsageTab", () => { expect(screen.getByText("$1.00")).toBeInTheDocument(); expect(screen.getByText("Classification cost")).toBeInTheDocument(); expect(screen.getByText("$0.2500")).toBeInTheDocument(); - expect(screen.queryByText(/1K turns/)).not.toBeInTheDocument(); + expect(screen.getByText("($62.50 / 1K turns)")).toBeInTheDocument(); expect(screen.getByText("Estimated spend at highest-tier model")).toBeInTheDocument(); expect(screen.getByText("$10.00")).toBeInTheDocument(); expect(screen.getByText("Auto-router prompt caching")).toBeInTheDocument(); From 1af7a403c66e037bec2e0ae6ea455a1c10b17b1b Mon Sep 17 00:00:00 2001 From: tin-berri Date: Mon, 7 Sep 2026 23:40:02 -0700 Subject: [PATCH 067/109] feat(mcp): start the named server's OAuth directly for a resource-scoped gateway flow (#39933) An aggregate gateway DCR authorize whose RFC 8707 resource resolves to exactly one gateway-managed oauth2 server sealed that server into the flow and then sent the browser to the generic connect grid anyway, so the user had to find the server the client had already named and click Connect. The connect URL now carries only the flow handle. GET /authorize/flow classifies the sealed flow as unscoped, interactive, M2M, or stale, and returns the matching state to the page. Interactive flows require a live per-user vendor credential before minting and do not burn the flow on an early submit. M2M flows use the gateway's configured service credential and finish without an interactive OAuth trip. Stale flows fail closed instead of becoming unscoped. The existing explicit Finish action and a new Cancel path preserve deliberate user intent. --- litellm/proxy/_experimental/mcp_server/db.py | 21 +- .../mcp_server/discoverable_endpoints.py | 69 +++-- .../mcp_server/gateway_dcr_flow.py | 187 +++++++++---- litellm/proxy/_lazy_openapi_snapshot.json | 40 +++ .../mcp_server/test_discoverable_endpoints.py | 52 ++++ .../mcp_server/test_gateway_dcr_flow.py | 254 ++++++++++++++++-- .../src/app/chat/integrations/page.tsx | 30 +-- .../src/app/connect/page.test.tsx | 90 +------ ui/litellm-dashboard/src/app/connect/page.tsx | 23 +- .../chat/ConnectFlowBanner.test.tsx | 69 +++-- .../src/components/chat/ConnectFlowBanner.tsx | 123 ++++++--- .../chat/ConnectFlowSurface.test.tsx | 118 ++++++++ .../components/chat/ConnectFlowSurface.tsx | 59 ++++ .../src/components/chat/MCPAppsPanel.test.tsx | 9 + .../src/components/chat/MCPAppsPanel.tsx | 33 ++- .../src/components/networking.tsx | 18 ++ .../src/lib/http/client.test.ts | 9 + ui/litellm-dashboard/src/lib/http/client.ts | 6 +- ui/litellm-dashboard/src/lib/http/schema.d.ts | 48 ++++ 19 files changed, 950 insertions(+), 308 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/chat/ConnectFlowSurface.test.tsx create mode 100644 ui/litellm-dashboard/src/components/chat/ConnectFlowSurface.tsx diff --git a/litellm/proxy/_experimental/mcp_server/db.py b/litellm/proxy/_experimental/mcp_server/db.py index 082a90fdcfb..7379126983a 100644 --- a/litellm/proxy/_experimental/mcp_server/db.py +++ b/litellm/proxy/_experimental/mcp_server/db.py @@ -4,7 +4,7 @@ import hashlib import json from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence from datetime import datetime, timedelta, timezone -from typing import TYPE_CHECKING, Any, Final, Protocol, TypedDict, cast +from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypedDict, cast from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid @@ -125,6 +125,9 @@ class OAuthCredentialPayload(_OAuthCredentialAccessToken, total=False): server_id: str +OAuthGrantState = Literal["valid", "refreshable", "absent"] + + class _OAuthTokenRefreshResponse(TypedDict, total=False): access_token: str refresh_token: str @@ -1465,6 +1468,15 @@ def is_oauth_credential_expired(cred: OAuthCredentialPayload, buffer_seconds: in return False +def oauth_grant_state(cred: OAuthCredentialPayload | None) -> OAuthGrantState: + """Classify local grant readiness without attempting a refresh or checking upstream revocation.""" + if not cred or not cred.get("access_token"): + return "absent" + if not is_oauth_credential_expired(cred, buffer_seconds=MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS): + return "valid" + return "refreshable" if cred.get("refresh_token") else "absent" + + async def get_user_oauth_credential( prisma_client: PrismaClient, user_id: str, @@ -1727,12 +1739,11 @@ async def resolve_valid_user_oauth_token( dict it already holds. ``prisma_client`` is fetched lazily and only when a refresh actually happens, so the valid-token path never requires a DB handle. """ - if not cred or not cred.get("access_token"): + grant: Final = oauth_grant_state(cred) + if cred is None or grant == "absent": return None - if not is_oauth_credential_expired(cred, buffer_seconds=MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS): + if grant == "valid": return cred - if not cred.get("refresh_token"): - return None if prisma_client is None: from litellm.proxy.utils import get_prisma_client_or_throw diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index e10bfd41ed6..cab4b6c161a 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -43,9 +43,11 @@ from litellm.proxy._experimental.mcp_server.faults import ( render_token_fault, ) from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import ( + VendorCredentialState, aggregate_authorize, aggregate_token, complete_connect_flow, + describe_connect_flow, introspect_gateway_token, is_gateway_dcr_client_id, is_proxy_api_resource, @@ -798,21 +800,7 @@ def _bridge_access_denied_redirect(redirect_uri: str, state: str, mcp_server: MC return RedirectResponse(_append_query_params(redirect_uri, params), status_code=302) -async def _bridge_authorize_access_denial( - litellm_user_id: str, - mcp_server: MCPServer, - redirect_uri: str, - state: str, -) -> RedirectResponse | None: - """The denial redirect for a signed-in user who cannot reach the target server, or None to proceed. - - Admits the user exactly as MCP egress will (the same ``reload_admitted_user`` constructor and the - same ``get_allowed_mcp_servers`` resolver), so an envelope is minted only when the resulting - session can actually list and call the server's tools. Without this gate the flow completes, the - client shows connected, and every tool request fail-closes to an empty list with nothing telling - the operator why. An availability fault (5xx, e.g. a DB outage's 503) propagates; an unknown or - deactivated user denies like a missing grant, fail closed. - """ +async def _user_can_reach_mcp_server(user_id: str, server_id: str) -> bool: from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( MCPRequestHandler, ) @@ -821,13 +809,22 @@ async def _bridge_authorize_access_denial( ) try: - admitted: Final = await MCPRequestHandler.reload_admitted_user(litellm_user_id) + admitted: Final = await MCPRequestHandler.reload_admitted_user(user_id) except HTTPException as exc: if exc.status_code >= 500: raise - return _bridge_access_denied_redirect(redirect_uri, state, mcp_server) - allowed_server_ids: Final = await global_mcp_server_manager.get_allowed_mcp_servers(admitted) - if mcp_server.server_id in allowed_server_ids: + return False + return server_id in await global_mcp_server_manager.get_allowed_mcp_servers(admitted) + + +async def _bridge_authorize_access_denial( + litellm_user_id: str, + mcp_server: MCPServer, + redirect_uri: str, + state: str, +) -> RedirectResponse | None: + """The denial redirect for a signed-in user who cannot reach the target server, or None to proceed.""" + if await _user_can_reach_mcp_server(litellm_user_id, mcp_server.server_id): return None return _bridge_access_denied_redirect(redirect_uri, state, mcp_server) @@ -1910,6 +1907,38 @@ async def token_endpoint( ) +async def _vendor_credential_state(user_id: str, server_id: str) -> VendorCredentialState: + """Whether the gateway itself can see a live vendor credential for this user and server. + + The one reading of "authorized" the connect page displays and the finish step enforces, so + the button a user sees and the grant they get cannot disagree. A read fault is neither, and + fails the scoped grant closed.""" + from litellm.proxy._experimental.mcp_server.db import ( # noqa: PLC0415 # circular import at module load + get_user_oauth_credential, + oauth_grant_state, + ) + from litellm.proxy.proxy_server import prisma_client # noqa: PLC0415 # circular import at module load + + if prisma_client is None: + return "unavailable" + try: + credential: Final = await get_user_oauth_credential(prisma_client, user_id, server_id) + except Exception: # noqa: BLE001 # a credential-read fault must fail the scoped grant closed + return "unavailable" + return "absent" if oauth_grant_state(credential) == "absent" else "present" + + +@router.get("/authorize/flow") +async def authorize_flow(request: Request, flow: str) -> Response: + return await describe_connect_flow( + request=request, + flow_handle=flow, + session_user_id=_session_cookie_user_id(request), + lookup_vendor_credential=_vendor_credential_state, + lookup_server_reachability=_user_can_reach_mcp_server, + ) + + @router.post("/authorize/complete") async def authorize_complete( request: Request, @@ -1934,6 +1963,8 @@ async def authorize_complete( delivery=delivery, team_id=team_id, decision=decision, + lookup_vendor_credential=_vendor_credential_state, + lookup_server_reachability=_user_can_reach_mcp_server, ) diff --git a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py index c7b0045dde5..3d94fa345d0 100644 --- a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py +++ b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py @@ -152,10 +152,8 @@ _AUTH_CODE_DEBUG_KEY: Final = "gateway_authorization_code" ReloadUserFailure = Literal["unresolvable", "unavailable", "faulted", "no_active_key"] ReloadUser = Callable[[str], Awaitable[ReloadUserFailure | None]] -"""Injected live-user revalidation (the token endpoint's mirror of admission): -``None`` means the user is active; ``unavailable`` is a retryable DB outage; ``faulted`` is -a DB fault retrying will not clear (still 503, worded so nobody just waits); anything else -fails the grant closed.""" +VendorCredentialState = Literal["present", "absent", "unavailable"] +"""The per-user vendor credential read has three outcomes: present, absent, or unavailable.""" _DB_UNAVAILABLE_DESCRIPTION: Final = "the gateway database is unavailable; retry" _DB_FAULTED_DESCRIPTION: Final = ( @@ -195,6 +193,16 @@ class ConsentTeam(BaseModel): team_alias: str | None = None +class LookupVendorCredential(Protocol): + """Injected read of a user's vendor credential for one server.""" + + def __call__(self, user_id: str, server_id: str, /) -> Awaitable[VendorCredentialState]: ... + + +class LookupServerReachability(Protocol): + def __call__(self, user_id: str, server_id: str, /) -> Awaitable[bool]: ... + + class LookupConsentTeams(Protocol): """Injected lookup of the teams a signed-in user may bind a proxy-API credential to.""" @@ -205,6 +213,14 @@ async def _refuse_proxy_credential(user_id: str, team_id: str | None) -> ProxyCr return "unresolvable" +async def _unavailable_vendor_credential(user_id: str, server_id: str) -> VendorCredentialState: + return "unavailable" + + +async def _unreachable_server(user_id: str, server_id: str) -> bool: + return False + + class GatewayDcrClient(BaseModel): """The registration record sealed into a gateway DCR ``client_id``. @@ -449,7 +465,10 @@ def aggregate_authorize( A per-server RFC 8707 ``resource`` naming a gateway-managed oauth2 server scopes the flow to that one server: the scope is sealed into the flow, carried into the code, and - bound into the session token, while the connect page interlude runs exactly as before. + bound into the session token. The connect URL carries only the flow handle; the page + learns the client origin, the scoped server, and whether its vendor OAuth is done from + :func:`describe_connect_flow`, which reads the sealed flow, so nothing a link can carry + steers which server the page authorizes or names on the confirmation. Validation failures respond directly with 400 and never redirect: per RFC 6749 section 4.1.2.1 an unvalidated redirect URI must not receive an error redirect, and @@ -474,10 +493,7 @@ def aggregate_authorize( resource_server_id=scoped_server.server_id if scoped_server is not None else None, audience=None, ) - connect_url: Final = _append_query_params( - f"{base_url}/ui/connect", - (("connect_flow", handle), ("connect_client", _origin_only(redirect_uri))), - ) + connect_url: Final = _append_query_params(f"{base_url}/ui/connect", (("connect_flow", handle),)) response: Final = RedirectResponse(connect_url, status_code=303) _set_flow_cookie(response, request, handle, flow) return response @@ -684,6 +700,99 @@ def _origin_only(url: str) -> str: return f"{parsed.scheme}://{parsed.netloc}" if parsed.netloc else "" +def _open_flow_for( + request: Request, flow_handle: str, session_user_id: str | None, now: datetime +) -> _ConnectFlow | Response: + sealed_flow: Final = request.cookies.get(_flow_cookie_name(flow_handle)) + if sealed_flow is None: + return _oauth_error(400, "invalid_request", "unknown or expired connect flow") + flow: Final = _open_sealed(sealed_flow, _UNPREFIXED, _ConnectFlow, _CONNECT_FLOW_DEBUG_KEY) + if flow is None or now.timestamp() >= flow.exp: + return _oauth_error(400, "invalid_request", "unknown or expired connect flow") + if session_user_id is None: + return _oauth_error(401, "login_required", "sign in to LiteLLM to finish connecting") + if session_user_id != flow.user_id: + return _oauth_error(403, "access_denied", "the signed-in user does not match this connect flow") + return flow + + +async def _flow_target( + flow: _ConnectFlow, lookup_server_reachability: LookupServerReachability +) -> tuple[Literal["unscoped", "interactive", "m2m", "stale"], MCPServer | None]: + if flow.resource_server_id is None: + return "unscoped", None + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( # noqa: PLC0415 # import cycle + MCPServerManager, + global_mcp_server_manager, + ) + + server: Final = global_mcp_server_manager.get_mcp_server_by_id(flow.resource_server_id) + if ( + server is None + or not server.is_gateway_managed_oauth2 + or not await lookup_server_reachability(flow.user_id, server.server_id) + ): + return "stale", None + state: Final = "m2m" if MCPServerManager.effective_oauth2_flow(server) == "client_credentials" else "interactive" + return state, server + + +class ConnectFlowDescription(TypedDict): + """What the connect page is allowed to know about one in-flight flow.""" + + state: ReadOnly[Literal["unscoped", "interactive", "m2m", "stale"]] + client_origin: ReadOnly[str] + server_id: ReadOnly[str | None] + server_name: ReadOnly[str | None] + connected: ReadOnly[bool | None] + + +async def _describe_opened_flow( + flow: _ConnectFlow, + lookup_vendor_credential: LookupVendorCredential, + lookup_server_reachability: LookupServerReachability, +) -> ConnectFlowDescription | Response: + state, server = await _flow_target(flow, lookup_server_reachability) + if state == "interactive" and server is not None: + credential: Final = await lookup_vendor_credential(flow.user_id, server.server_id) + if credential == "unavailable": + return _oauth_error(503, "temporarily_unavailable", _DB_UNAVAILABLE_DESCRIPTION) + interactive_description: Final[ConnectFlowDescription] = { + "state": state, + "client_origin": _origin_only(flow.redirect_uri), + "server_id": server.server_id, + "server_name": server.server_name or server.alias or server.name, + "connected": credential == "present", + } + return interactive_description + described: Final[ConnectFlowDescription] = { + "state": state, + "client_origin": _origin_only(flow.redirect_uri), + "server_id": None if server is None else server.server_id, + "server_name": None if server is None else (server.server_name or server.alias or server.name), + "connected": state == "m2m" or None, + } + return described + + +async def describe_connect_flow( + request: Request, + flow_handle: str, + session_user_id: str | None, + lookup_vendor_credential: LookupVendorCredential, + lookup_server_reachability: LookupServerReachability, +) -> Response: + opened: Final = _open_flow_for(request, flow_handle, session_user_id, datetime.now(timezone.utc)) + if isinstance(opened, Response): + return opened + described: Final = await _describe_opened_flow(opened, lookup_vendor_credential, lookup_server_reachability) + return ( + described + if isinstance(described, Response) + else JSONResponse(content=described, headers=TOKEN_NO_CACHE_HEADERS) + ) + + async def complete_connect_flow( request: Request, flow_handle: str, @@ -692,56 +801,34 @@ async def complete_connect_flow( delivery: str | None = None, team_id: str | None = None, decision: str | None = None, + lookup_vendor_credential: LookupVendorCredential = _unavailable_vendor_credential, + lookup_server_reachability: LookupServerReachability = _unreachable_server, ) -> Response: - """The deliberate finish step of the connect flow: mint the gateway authorization - code and send the browser back to the client. + """Mint the code only after a deliberate POST by the sealed user. - Reached by POST so a cross-site GET cannot trigger it, and bound to the HttpOnly - per-flow cookie plus an exact match between the signed-in user and the user sealed - into the flow: a link crafted by another party dies here with ``access_denied`` - instead of minting a code for the victim's identity. The flow is single-use (an atomic - claim on its ``jti``), so a double-submit cannot mint two codes from one sign-in. - - ``delivery`` chooses how the code reaches the client. Default (absent or - ``"redirect"``) is the 303 to the client's registered redirect URI. ``"manual"`` - renders the callback URL on a page instead, for a client whose redirect URI is a - loopback host but which runs on a DIFFERENT machine than the browser (EC2/SSH box, - container): the 303 would dereference the browser machine's loopback and the code - would never arrive, so the user carries it over by pasting the URL into the client or - fetching it from the client machine's terminal. Manual delivery is honored only for - loopback redirect URIs; a routable redirect URI works from any browser by - construction, so those flows always redirect. The user who sees the page is exactly - the user the 303 would have carried the code to, and the same user already sees the - code today in the dead redirect's address bar, so the page exposes the code to no new - party. Unknown ``delivery`` values are rejected rather than defaulted: a client that - asked for manual delivery and got a dead redirect instead would silently lose its - code. - - ``decision`` and ``team_id`` come from the native-client consent page. ``"deny"`` - burns the flow and sends the client ``error=access_denied`` so it stops waiting; - ``team_id`` is sealed into the code only for proxy-API flows, where it picks which of - the user's teams the minted credential is attributed to. + A scoped flow additionally requires its sealed server to have a live vendor credential + before a code can be minted. The check happens before the single-use claim, so a + premature submit can be retried after authorization; denial deliberately bypasses it. """ if delivery not in (None, "redirect", "manual"): return _oauth_error(400, "invalid_request", "delivery must be 'redirect' or 'manual'") if decision not in (None, "approve", "deny"): return _oauth_error(400, "invalid_request", "decision must be 'approve' or 'deny'") - sealed_flow: Final = request.cookies.get(_flow_cookie_name(flow_handle)) - if sealed_flow is None: - return _oauth_error(400, "invalid_request", "unknown or expired connect flow") - flow: Final = _open_sealed(sealed_flow, _UNPREFIXED, _ConnectFlow, _CONNECT_FLOW_DEBUG_KEY) - if flow is None: - return _oauth_error(400, "invalid_request", "unknown or expired connect flow") now: Final = datetime.now(timezone.utc) - if now.timestamp() >= flow.exp: - return _oauth_error(400, "invalid_request", "the connect flow has expired; restart the connection") - if session_user_id is None: - return _oauth_error(401, "login_required", "sign in to LiteLLM to finish connecting") - if session_user_id != flow.user_id: - return _oauth_error(403, "access_denied", "the signed-in user does not match this connect flow") + opened: Final = _open_flow_for(request, flow_handle, session_user_id, now) + if isinstance(opened, Response): + return opened + if decision != "deny": + described: Final = await _describe_opened_flow(opened, lookup_vendor_credential, lookup_server_reachability) + if isinstance(described, Response): + return described + if described["state"] == "stale": + return _oauth_error(400, "invalid_request", "the requested MCP server is no longer available") + if described["connected"] is False: + return _oauth_error(400, "invalid_request", "authorize the requested MCP server before finishing") flow_refusal: Final = _claim_refusal( await _SingleUseGuard(cache).claim( - f"{_USED_FLOW_CACHE_PREFIX}{flow.jti}", CONNECT_FLOW_TTL_SECONDS + _CLAIM_TTL_BUFFER_SECONDS + f"{_USED_FLOW_CACHE_PREFIX}{opened.jti}", CONNECT_FLOW_TTL_SECONDS + _CLAIM_TTL_BUFFER_SECONDS ), replayed=_oauth_error( 400, "invalid_request", "this connect flow was already completed; restart the connection" @@ -750,7 +837,7 @@ async def complete_connect_flow( if flow_refusal is not None: return flow_refusal response: Final = ( - _denied_flow_response(flow) if decision == "deny" else _approved_flow_response(flow, delivery, team_id, now) + _denied_flow_response(opened) if decision == "deny" else _approved_flow_response(opened, delivery, team_id, now) ) path, secure = _cookie_path_and_secure(request) response.delete_cookie(key=_flow_cookie_name(flow_handle), path=path, secure=secure, httponly=True, samesite="lax") diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index dba2b2428aa..71475320c2c 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -19886,6 +19886,46 @@ ] } }, + "/authorize/flow": { + "get": { + "operationId": "authorize_flow_authorize_flow_get", + "parameters": [ + { + "in": "query", + "name": "flow", + "required": true, + "schema": { + "title": "Flow", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "summary": "Authorize Flow", + "tags": [ + "mcp_discoverable" + ] + } + }, "/callback": { "get": { "description": "OAuth 2.0 authorization response handler for MCP loopback clients.\n\nAccepts either:\n\n- A successful authorization response (``code`` + ``state``), which is\n forwarded back to the validated client ``redirect_uri`` with the\n original (un-wrapped) ``state``.\n- An error response (``error``[+``error_description``/``error_uri``]), per\n RFC 6749 \u00a74.1.2.1. When ``state`` is present and decodes to a trusted\n ``redirect_uri``, the error params are propagated back to the client so\n its OAuth library can surface them. Otherwise we render an HTML error\n page so the user is not left on an opaque 422 / blank screen.", diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index be4206a1faf..763200c3709 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -4,6 +4,7 @@ import hashlib import json import time from base64 import urlsafe_b64encode +from datetime import datetime, timedelta, timezone from typing import TYPE_CHECKING from unittest.mock import AsyncMock, MagicMock, patch @@ -18,6 +19,57 @@ if TYPE_CHECKING: from litellm.types.mcp_server.mcp_server_manager import MCPServer +def _stored_grant(access_token="access-token", refresh_token=None, expires_in_seconds=None, expires_at=None): + credential = {"type": "oauth2", "access_token": access_token} + if refresh_token is not None: + credential["refresh_token"] = refresh_token + if expires_in_seconds is not None: + credential["expires_at"] = (datetime.now(timezone.utc) + timedelta(seconds=expires_in_seconds)).isoformat() + if expires_at is not None: + credential["expires_at"] = expires_at + return credential + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("fields", "egress_has_token"), + [ + (None, False), + ({"access_token": "", "refresh_token": "refresh-token"}, False), + ({}, True), + ({"expires_at": "never"}, True), + ({"expires_in_seconds": 600}, True), + ({"expires_in_seconds": 30}, False), + ({"expires_in_seconds": -300}, False), + ({"expires_in_seconds": -300, "refresh_token": ""}, False), + ({"expires_in_seconds": 30, "refresh_token": "refresh-token"}, True), + ({"expires_in_seconds": -300, "refresh_token": "refresh-token"}, True), + ], +) +async def test_vendor_credential_state_agrees_with_egress_token_resolution(monkeypatch, fields, egress_has_token): + from litellm.proxy._experimental.mcp_server import db as mcp_db + from litellm.proxy._experimental.mcp_server import discoverable_endpoints + + monkeypatch.setattr(mcp_db, "MCP_PER_USER_TOKEN_EXPIRY_BUFFER_SECONDS", 60) + credential = _stored_grant(**fields) if fields is not None else None + read = AsyncMock(return_value=credential) + refresh = AsyncMock(return_value=_stored_grant(access_token="fresh-token", expires_in_seconds=3600)) + prisma = MagicMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma) + monkeypatch.setattr(mcp_db, "get_user_oauth_credential", read) + monkeypatch.setattr(mcp_db, "refresh_user_oauth_token", refresh) + + connect = await discoverable_endpoints._vendor_credential_state("user-1", "server-1") + read.assert_awaited_once_with(prisma, "user-1", "server-1") + refresh.assert_not_awaited() + egress = await mcp_db.resolve_valid_user_oauth_token( + user_id="user-1", server=MagicMock(), cred=credential, prisma_client=prisma + ) + + assert (egress is not None) is egress_has_token + assert connect == ("present" if egress_has_token else "absent") + + # Fixture to mock IP address check for all MCP tests # This prevents tests from failing due to IP-based access control @pytest.fixture(autouse=True) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py index 1670370f082..73a52a8d2e8 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_gateway_dcr_flow.py @@ -26,6 +26,7 @@ from litellm.proxy._experimental.mcp_server.gateway_dcr_flow import ( aggregate_authorize, aggregate_token, complete_connect_flow, + describe_connect_flow, introspect_gateway_token, is_gateway_dcr_client_id, is_proxy_api_resource, @@ -230,7 +231,7 @@ async def test_authorize_with_session_hands_browser_to_connect_page_with_flow_co assert location.path == "/ui/connect" params = parse_qs(location.query) handle = params["connect_flow"][0] - assert params["connect_client"] == ["https://claude.ai"] + assert set(params) == {"connect_flow"} set_cookie = response.headers["set-cookie"] assert f"{CONNECT_FLOW_COOKIE_PREFIX}{handle}" in set_cookie assert "HttpOnly" in set_cookie @@ -889,14 +890,60 @@ def _opened_principal(payload): return admitted.principal -async def _finish_connect_page(response): +class _VendorCredential: + def __init__(self, state="present"): + self.calls = [] + self.state = state + + async def __call__(self, user_id, server_id): + self.calls.append((user_id, server_id)) + return self.state + + +class _ServerReachability: + def __init__(self, reachable=True): + self.calls = [] + self.reachable = reachable + + async def __call__(self, user_id, server_id): + self.calls.append((user_id, server_id)) + return self.reachable + + +async def _complete_page(response, scoped_server=None, vendor=None, reachable=None, cache=None, **overrides): + from unittest.mock import patch + handle, cookies = _flow_cookie_from(response) - completed = await complete_connect_flow( - request=_request("/authorize/complete", cookies=cookies, method="POST"), - flow_handle=handle, - session_user_id="u1", - cache=DualCache(), - ) + with patch(_MANAGER_PATCH) as manager: + manager.get_mcp_server_by_id.return_value = scoped_server + return await complete_connect_flow( + request=_request("/authorize/complete", cookies=cookies, method="POST"), + flow_handle=handle, + session_user_id="u1", + cache=cache or DualCache(), + lookup_vendor_credential=vendor or _VendorCredential(), + lookup_server_reachability=reachable or _ServerReachability(), + **overrides, + ) + + +async def _describe_page(response, scoped_server=None, vendor=None, reachable=None, session_user_id="u1", cookies=None): + from unittest.mock import patch + + handle, flow_cookies = _flow_cookie_from(response) + with patch(_MANAGER_PATCH) as manager: + manager.get_mcp_server_by_id.return_value = scoped_server + return await describe_connect_flow( + request=_request("/authorize/flow", cookies=flow_cookies if cookies is None else cookies), + flow_handle=handle, + session_user_id=session_user_id, + lookup_vendor_credential=vendor or _VendorCredential(), + lookup_server_reachability=reachable or _ServerReachability(), + ) + + +async def _finish_connect_page(response, scoped_server=None): + completed = await _complete_page(response, scoped_server=scoped_server) return parse_qs(urlparse(completed.headers["location"]).query)["code"][0] @@ -910,23 +957,43 @@ def _sealed_wire_json(sealed, prefix, debug_key): @pytest.mark.asyncio async def test_scoped_authorize_runs_connect_page_with_sealed_scope(): - """LIT-4917: a per-server RFC 8707 resource naming a gateway-managed oauth2 server - seals that server into the flow. The connect page interlude runs exactly as before - (the scope restricts, it never skips consent), and the code minted at the finish step - and the session pair it redeems for are both scoped.""" + """LIT-4917 plus LIT-7075: a per-server RFC 8707 resource naming a gateway-managed oauth2 + server seals that server into the flow. The connect URL carries only the handle; the page + learns the scoped server and its vendor state from describe_connect_flow, and the finish + step refuses to mint a scoped code until that vendor credential exists, without burning + the flow. The code minted afterwards and the session pair it redeems for are both scoped.""" from unittest.mock import patch client_id = (await _register([REDIRECT_URI]))["client_id"] + github = _scoped_mcp_server() with patch(_MANAGER_PATCH) as manager: - manager.get_mcp_server_by_name.return_value = _scoped_mcp_server() + manager.get_mcp_server_by_name.return_value = github response = _scoped_authorize(client_id, SCOPED_RESOURCE) assert response.status_code == 303 - assert "/ui/connect" in response.headers["location"] + location = urlparse(response.headers["location"]) + assert location.path == "/ui/connect" + assert set(parse_qs(location.query)) == {"connect_flow"} _, cookies = _flow_cookie_from(response) assert ( _sealed_wire_json(next(iter(cookies.values())), "", "gateway_connect_flow")["resource_server_id"] == "github-id" ) - code = await _finish_connect_page(response) + described = await _describe_page(response, scoped_server=github, vendor=_VendorCredential("absent")) + assert json.loads(described.body) == { + "state": "interactive", + "client_origin": "https://claude.ai", + "server_id": "github-id", + "server_name": "github", + "connected": False, + } + cache = DualCache() + premature = await _complete_page(response, scoped_server=github, vendor=_VendorCredential("absent"), cache=cache) + assert premature.status_code == 400 + assert "authorize the requested MCP server" in json.loads(premature.body)["error_description"] + present = _VendorCredential("present") + completed = await _complete_page(response, scoped_server=github, vendor=present, cache=cache) + assert completed.status_code == 303 + assert present.calls == [("u1", "github-id")] + code = parse_qs(urlparse(completed.headers["location"]).query)["code"][0] assert ( _sealed_wire_json(code, GATEWAY_AUTH_CODE_PREFIX, "gateway_authorization_code")["resource_server_id"] == "github-id" @@ -952,9 +1019,11 @@ async def test_scoped_authorize_runs_connect_page_with_sealed_scope(): ) async def test_unscoped_resources_leave_flow_and_token_byte_identical(resource, resolves): """Every resource shape outside 'exactly one gateway-managed server' keeps today's flow: - connect page interlude, and NONE of the minted artifacts carry the scope key on the - wire, not the flow cookie, not the code, not the session JWT, so an unscoped flow - started on a new pod completes on a pod whose strict models predate the claim.""" + the generic connect grid (describe names no server, the finish step never consults the + vendor credential), and NONE of the minted + artifacts carry the scope key on the wire, not the flow cookie, not the code, not the + session JWT, so an unscoped flow started on a new pod completes on a pod whose strict + models predate the claim.""" import base64 from unittest.mock import patch @@ -963,10 +1032,18 @@ async def test_unscoped_resources_leave_flow_and_token_byte_identical(resource, manager.get_mcp_server_by_name.return_value = None if resolves is None else _scoped_mcp_server() response = _scoped_authorize(client_id, resource) assert response.status_code == 303 - assert "/ui/connect" in response.headers["location"] + location = urlparse(response.headers["location"]) + assert location.path == "/ui/connect" + assert set(parse_qs(location.query)) == {"connect_flow"} _, cookies = _flow_cookie_from(response) assert "resource_server_id" not in _sealed_wire_json(next(iter(cookies.values())), "", "gateway_connect_flow") - code = await _finish_connect_page(response) + vendor = _VendorCredential("absent") + described = await _describe_page(response, vendor=vendor) + assert json.loads(described.body)["state"] == "unscoped" + assert json.loads(described.body)["server_id"] is None + completed = await _complete_page(response, vendor=vendor) + assert vendor.calls == [] + code = parse_qs(urlparse(completed.headers["location"]).query)["code"][0] assert "resource_server_id" not in _sealed_wire_json(code, GATEWAY_AUTH_CODE_PREFIX, "gateway_authorization_code") token_response = await _redeem(code, client_id) payload = json.loads(token_response.body) @@ -979,19 +1056,150 @@ async def test_unscoped_resources_leave_flow_and_token_byte_identical(resource, @pytest.mark.asyncio async def test_scoped_authorize_delegate_server_stays_unscoped(): """A delegate-auth oauth2 server is outside the gateway-managed set (its keyless flow is - upstream PKCE via the relay), so a resource naming it never scopes the gateway flow.""" + upstream PKCE via the relay), so a resource naming it never scopes the gateway flow and + never narrows the connect page to it.""" from unittest.mock import patch client_id = (await _register([REDIRECT_URI]))["client_id"] with patch(_MANAGER_PATCH) as manager: manager.get_mcp_server_by_name.return_value = _scoped_mcp_server(delegate_auth_to_upstream=True) response = _scoped_authorize(client_id, SCOPED_RESOURCE) - assert "/ui/connect" in response.headers["location"] + location = urlparse(response.headers["location"]) + assert location.path == "/ui/connect" + assert set(parse_qs(location.query)) == {"connect_flow"} + assert json.loads((await _describe_page(response)).body)["server_id"] is None code = await _finish_connect_page(response) token_response = await _redeem(code, client_id) assert _opened_principal(json.loads(token_response.body)).resource_server_id is None +@pytest.mark.asyncio +async def test_m2m_scoped_flow_mints_without_a_user_credential(): + """A client-credentials server is already authorized by its gateway service credential, so + a resource-scoped flow finishes without consulting the per-user vault.""" + from unittest.mock import patch + + client_id = (await _register([REDIRECT_URI]))["client_id"] + m2m = _scoped_mcp_server(oauth2_flow="client_credentials") + with patch(_MANAGER_PATCH) as manager: + manager.get_mcp_server_by_name.return_value = m2m + response = _scoped_authorize(client_id, SCOPED_RESOURCE) + vendor = _VendorCredential("unavailable") + described = await _describe_page(response, scoped_server=m2m, vendor=vendor) + assert json.loads(described.body)["state"] == "m2m" + assert json.loads(described.body)["connected"] is True + assert vendor.calls == [] + completed = await _complete_page(response, scoped_server=m2m, vendor=vendor) + assert completed.status_code == 303 + assert vendor.calls == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("oauth2_flow", ["authorization_code", "client_credentials"]) +async def test_unreachable_scoped_flow_cannot_describe_or_finish(oauth2_flow): + from unittest.mock import patch + + client_id = (await _register([REDIRECT_URI]))["client_id"] + server = _scoped_mcp_server(oauth2_flow=oauth2_flow) + with patch(_MANAGER_PATCH) as manager: + manager.get_mcp_server_by_name.return_value = server + response = _scoped_authorize(client_id, SCOPED_RESOURCE) + reachable = _ServerReachability(False) + vendor = _VendorCredential("present") + described = await _describe_page(response, scoped_server=server, reachable=reachable, vendor=vendor) + assert json.loads(described.body) == { + "state": "stale", + "client_origin": "https://claude.ai", + "server_id": None, + "server_name": None, + "connected": None, + } + cache = DualCache() + refused = await _complete_page(response, scoped_server=server, reachable=reachable, vendor=vendor, cache=cache) + assert refused.status_code == 400 + assert vendor.calls == [] + assert reachable.calls == [("u1", "github-id"), ("u1", "github-id")] + completed = await _complete_page(response, scoped_server=server, vendor=vendor, cache=cache) + assert completed.status_code == 303 + + +@pytest.mark.asyncio +async def test_stale_scoped_flow_remains_distinct_from_unscoped(): + """A server removed after authorize stays a stale scoped flow, so the page cannot offer a + broader unscoped grant or report a misleading Finish action.""" + from unittest.mock import patch + + client_id = (await _register([REDIRECT_URI]))["client_id"] + with patch(_MANAGER_PATCH) as manager: + manager.get_mcp_server_by_name.return_value = _scoped_mcp_server() + response = _scoped_authorize(client_id, SCOPED_RESOURCE) + described = await _describe_page(response, scoped_server=None) + assert json.loads(described.body)["state"] == "stale" + assert json.loads(described.body)["connected"] is None + stale = await _complete_page(response, scoped_server=None) + assert stale.status_code == 400 + assert json.loads(stale.body)["error_description"] == "the requested MCP server is no longer available" + + +@pytest.mark.asyncio +async def test_scoped_flow_deny_and_stale_server_never_need_the_vendor_credential(): + """Cancel is the escape hatch: a scoped user who cannot finish the vendor step still ends + the flow with access_denied and no credential lookup. A scoped server that is no longer + gateway-managed refuses to mint (nothing could serve that code) but also burns nothing.""" + from unittest.mock import patch + + client_id = (await _register([REDIRECT_URI]))["client_id"] + github = _scoped_mcp_server() + with patch(_MANAGER_PATCH) as manager: + manager.get_mcp_server_by_name.return_value = github + response = _scoped_authorize(client_id, SCOPED_RESOURCE) + cache = DualCache() + skipped_reachability = _ServerReachability(False) + stale = await _complete_page(response, scoped_server=None, reachable=skipped_reachability, cache=cache) + assert stale.status_code == 400 + assert skipped_reachability.calls == [] + assert json.loads(stale.body)["error_description"] == "the requested MCP server is no longer available" + described = await _describe_page(response, scoped_server=github, vendor=_VendorCredential("unavailable")) + assert described.status_code == 503 + vendor = _VendorCredential("absent") + deny_reachability = _ServerReachability(False) + denied = await _complete_page( + response, + scoped_server=github, + vendor=vendor, + reachable=deny_reachability, + cache=cache, + decision="deny", + ) + assert denied.status_code == 303 + assert parse_qs(urlparse(denied.headers["location"]).query)["error"] == ["access_denied"] + assert vendor.calls == [] + assert deny_reachability.calls == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "session_user_id, cookies, expected_status, expected_error", + [ + ("u1", {}, 400, "invalid_request"), + ("u1", {"mcp_connect_flow_wrong": "garbage"}, 400, "invalid_request"), + (None, None, 401, "login_required"), + ("u2", None, 403, "access_denied"), + ], +) +async def test_describe_connect_flow_refuses_exactly_like_the_finish_step( + session_user_id, cookies, expected_status, expected_error +): + """The page's read of the flow is gated the same way minting is: the HttpOnly cookie for + that handle must open and the signed-in user must be the sealed one. A lure link with a + made-up handle therefore learns nothing and starts nothing.""" + client_id = (await _register([REDIRECT_URI]))["client_id"] + response = _authorize(client_id, session_user_id="u1") + described = await _describe_page(response, session_user_id=session_user_id, cookies=cookies) + assert described.status_code == expected_status + assert json.loads(described.body)["error"] == expected_error + + @pytest.mark.asyncio async def test_token_rejects_resource_conflicting_with_sealed_scope(): """RFC 8707 section 2.2: redeeming a scoped code (or rotating a scoped refresh token) @@ -1005,7 +1213,7 @@ async def test_token_rejects_resource_conflicting_with_sealed_scope(): with patch(_MANAGER_PATCH) as manager: manager.get_mcp_server_by_name.return_value = github response = _scoped_authorize(client_id, SCOPED_RESOURCE) - code = await _finish_connect_page(response) + code = await _finish_connect_page(response, scoped_server=github) with patch(_MANAGER_PATCH) as manager: manager.get_mcp_server_by_name.return_value = linear diff --git a/ui/litellm-dashboard/src/app/chat/integrations/page.tsx b/ui/litellm-dashboard/src/app/chat/integrations/page.tsx index 30ce62d8081..a663e16c2b2 100644 --- a/ui/litellm-dashboard/src/app/chat/integrations/page.tsx +++ b/ui/litellm-dashboard/src/app/chat/integrations/page.tsx @@ -1,43 +1,19 @@ "use client"; -import { Suspense, useEffect } from "react"; -import { useRouter, useSearchParams } from "next/navigation"; +import { Suspense } from "react"; import { useChatShell } from "@/contexts/ChatShellContext"; -import MCPAppsPanel from "@/components/chat/MCPAppsPanel"; -import ConnectFlowBanner from "@/components/chat/ConnectFlowBanner"; +import ConnectFlowSurface from "@/components/chat/ConnectFlowSurface"; // useSearchParams() requires a Suspense boundary for static export. function IntegrationsPageContent() { const { accessToken, selectedMCPServers, setSelectedMCPServers } = useChatShell(); - const router = useRouter(); - const searchParams = useSearchParams(); - const oauthReturn = searchParams.get("mcpOauthReturn"); - // Set by the gateway DCR authorize when a DCR client sends the user here to - // authorize servers before finishing sign-in (see gateway_dcr_flow.py). The - // handle keys the sealed per-flow cookie; connect_client is the client origin - // for display only. connect_flow is NOT cleaned from the URL: the finish form - // needs it, and the sealed cookie (not the URL) is the security boundary. - const connectFlow = searchParams.get("connect_flow"); - const connectClient = searchParams.get("connect_client"); - - // Clean up the OAuth return param after it's been consumed — real routing means - // we no longer need it to pick a tab, but it should not linger in the address bar. - useEffect(() => { - if (oauthReturn) { - const url = new URL(window.location.href); - url.searchParams.delete("mcpOauthReturn"); - router.replace(url.pathname + url.search); - } - }, [oauthReturn, router]); return (
- {connectFlow && } -
); diff --git a/ui/litellm-dashboard/src/app/connect/page.test.tsx b/ui/litellm-dashboard/src/app/connect/page.test.tsx index 7d49a8b6a4c..6a6ba24bd87 100644 --- a/ui/litellm-dashboard/src/app/connect/page.test.tsx +++ b/ui/litellm-dashboard/src/app/connect/page.test.tsx @@ -2,101 +2,29 @@ import { afterEach, describe, expect, it, vi } from "vitest"; import { render, screen } from "@testing-library/react"; import ConnectPage from "./page"; -interface PanelProps { +interface SurfaceProps { accessToken: string; selectedServers: string[]; onChange: (servers: string[]) => void; - connectMode?: boolean; } -interface BannerProps { - flowHandle: string; - clientOrigin: string | null; -} - -const { mockReplace, mockPanel, mockBanner, state } = vi.hoisted(() => { - const state = { - oauthReturn: null as string | null, - connectFlow: null as string | null, - connectClient: null as string | null, - }; - return { - state, - mockReplace: vi.fn(), - mockPanel: vi.fn((_props: PanelProps) =>
), - mockBanner: vi.fn((_props: BannerProps) =>
), - }; -}); - -vi.mock("next/navigation", () => ({ - useRouter: () => ({ replace: mockReplace }), - useSearchParams: () => ({ - get: (key: string) => { - if (key === "mcpOauthReturn") return state.oauthReturn; - if (key === "connect_flow") return state.connectFlow; - if (key === "connect_client") return state.connectClient; - return null; - }, - }), +const { mockSurface } = vi.hoisted(() => ({ + mockSurface: vi.fn((_props: SurfaceProps) =>
), })); + vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ default: () => ({ accessToken: "token-123" }), })); -vi.mock("@/components/chat/MCPAppsPanel", () => ({ default: mockPanel })); -vi.mock("@/components/chat/ConnectFlowBanner", () => ({ default: mockBanner })); +vi.mock("@/components/chat/ConnectFlowSurface", () => ({ default: mockSurface })); describe("ConnectPage", () => { afterEach(() => { - state.oauthReturn = null; - state.connectFlow = null; - state.connectClient = null; - mockReplace.mockClear(); - mockPanel.mockClear(); - mockBanner.mockClear(); + mockSurface.mockClear(); }); - it("renders the MCP connect panel with the user's access token", () => { + it("renders the gateway connect surface with the user's access token and an empty selection", () => { render(); - expect(screen.getByTestId("mcp-apps-panel")).toBeInTheDocument(); - expect(mockPanel.mock.calls[0][0]).toMatchObject({ accessToken: "token-123", selectedServers: [] }); - }); - - it("strips the mcpOauthReturn param from the URL after an OAuth return", () => { - state.oauthReturn = "apps"; - window.history.replaceState({}, "", "/connect?mcpOauthReturn=apps"); - render(); - expect(mockReplace).toHaveBeenCalledWith("/connect"); - }); - - it("does not rewrite the URL when there is no OAuth return param", () => { - render(); - expect(mockReplace).not.toHaveBeenCalled(); - }); - - it("mounts the gateway connect banner and puts the panel in connect mode for a DCR flow", () => { - state.connectFlow = "flow-handle-123"; - state.connectClient = "https://claude.ai"; - render(); - expect(screen.getByTestId("connect-flow-banner")).toBeInTheDocument(); - expect(mockBanner.mock.calls[0][0]).toMatchObject({ - flowHandle: "flow-handle-123", - clientOrigin: "https://claude.ai", - }); - expect(mockPanel.mock.calls[0][0].connectMode).toBe(true); - }); - - it("shows no connect banner and leaves connect mode off for a plain visit", () => { - render(); - expect(screen.queryByTestId("connect-flow-banner")).not.toBeInTheDocument(); - expect(mockBanner).not.toHaveBeenCalled(); - expect(mockPanel.mock.calls[0][0].connectMode).toBe(false); - }); - - it("keeps the connect flow handle in the URL while stripping the OAuth return param", () => { - state.oauthReturn = "apps"; - state.connectFlow = "flow-handle-123"; - window.history.replaceState({}, "", "/connect?connect_flow=flow-handle-123&mcpOauthReturn=apps"); - render(); - expect(mockReplace).toHaveBeenCalledWith("/connect?connect_flow=flow-handle-123"); + expect(screen.getByTestId("connect-flow-surface")).toBeInTheDocument(); + expect(mockSurface.mock.calls[0][0]).toMatchObject({ accessToken: "token-123", selectedServers: [] }); }); }); diff --git a/ui/litellm-dashboard/src/app/connect/page.tsx b/ui/litellm-dashboard/src/app/connect/page.tsx index 3f0c269e86b..652c044197b 100644 --- a/ui/litellm-dashboard/src/app/connect/page.tsx +++ b/ui/litellm-dashboard/src/app/connect/page.tsx @@ -1,36 +1,19 @@ "use client"; -import { Suspense, useEffect, useState } from "react"; -import { useRouter, useSearchParams } from "next/navigation"; +import { Suspense, useState } from "react"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; -import MCPAppsPanel from "@/components/chat/MCPAppsPanel"; -import ConnectFlowBanner from "@/components/chat/ConnectFlowBanner"; +import ConnectFlowSurface from "@/components/chat/ConnectFlowSurface"; function ConnectPageContent() { const { accessToken } = useAuthorized(); const [selectedServers, setSelectedServers] = useState([]); - const router = useRouter(); - const searchParams = useSearchParams(); - const oauthReturn = searchParams.get("mcpOauthReturn"); - const connectFlow = searchParams.get("connect_flow"); - const connectClient = searchParams.get("connect_client"); - - useEffect(() => { - if (oauthReturn) { - const url = new URL(window.location.href); - url.searchParams.delete("mcpOauthReturn"); - router.replace(url.pathname + url.search); - } - }, [oauthReturn, router]); return (
- {connectFlow && } -
); diff --git a/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.test.tsx b/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.test.tsx index b3cd6e229af..5caf15d1fce 100644 --- a/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.test.tsx +++ b/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.test.tsx @@ -1,59 +1,61 @@ import { afterEach, describe, expect, it, vi } from "vitest"; import { render, screen } from "@testing-library/react"; +import type { ConnectFlowStatus } from "@/components/networking"; import ConnectFlowBanner, { isLoopbackOrigin } from "./ConnectFlowBanner"; vi.mock("@/components/networking", () => ({ getProxyBaseUrl: () => "https://gateway.example.com", })); +vi.mock("@/hooks/useUserMcpOAuthFlow", () => ({ + useUserMcpOAuthFlow: () => ({ startOAuthFlow: vi.fn(), status: "idle" }), +})); + afterEach(() => { vi.restoreAllMocks(); - sessionStorage.clear(); }); +const unscoped = (client_origin: string): ConnectFlowStatus => ({ + state: "unscoped", + client_origin, + server_id: null, + server_name: null, + connected: null, +}); + +const renderBanner = (clientOrigin: string) => + render( + , + ); + describe("ConnectFlowBanner", () => { - it("posts the flow handle to the proxy /authorize/complete as a full-page form", () => { - const { container } = render(); + it("posts only the flow handle to the proxy /authorize/complete as a full-page form", () => { + const { container } = renderBanner("https://claude.ai"); const form = container.querySelector("form")!; expect(form).toHaveAttribute("method", "POST"); expect(form).toHaveAttribute("action", "https://gateway.example.com/authorize/complete"); - - const hidden = form.querySelector('input[name="flow"]') as HTMLInputElement; - expect(hidden.value).toBe("flow-handle-123"); - // No token, code, or secret is ever placed in the form; the sealed cookie carries them. + expect(screen.getByDisplayValue("flow-handle-123")).toHaveAttribute("name", "flow"); expect(form.innerHTML).not.toContain("token"); - }); - - it("shows the client origin so the user knows what they are connecting to", () => { - render(); - expect(screen.getAllByText(/claude\.ai/).length).toBeGreaterThan(0); expect(screen.getByRole("button", { name: /finish connecting/i })).toBeInTheDocument(); + expect(screen.queryByRole("button", { name: "Cancel" })).not.toBeInTheDocument(); }); - it("falls back to a generic label when the client origin is unknown", () => { - render(); - expect(screen.getAllByText(/the application/).length).toBeGreaterThan(0); - }); - - it("offers manual delivery for a loopback client, posted only when checked", () => { - const { container } = render( - , - ); - - const checkbox = container.querySelector('input[type="checkbox"][name="delivery"]') as HTMLInputElement; - expect(checkbox).not.toBeNull(); + it("offers manual delivery only for a loopback client, posted only when checked", () => { + const loopback = renderBanner("http://localhost:3118"); + const checkbox = loopback.container.querySelector('input[type="checkbox"][name="delivery"]') as HTMLInputElement; expect(checkbox.value).toBe("manual"); expect(checkbox.checked).toBe(false); - expect(screen.getByText(/remote or SSH machine/i)).toBeInTheDocument(); - }); + loopback.unmount(); - it("does not offer manual delivery for a routable client origin or an unknown one", () => { - const routable = render(); + const routable = renderBanner("https://claude.ai"); expect(routable.container.querySelector('input[name="delivery"]')).toBeNull(); - - const unknown = render(); - expect(unknown.container.querySelector('input[name="delivery"]')).toBeNull(); }); it("classifies loopback origins like the server does", () => { @@ -70,12 +72,9 @@ describe("ConnectFlowBanner", () => { }); it("does NOT complete the flow on pagehide (completion requires the explicit button)", () => { - // Security regression: an attacker could lure a signed-in victim to their own client's - // authorize URL; the victim merely closing the tab must NOT deliver a victim-bound code. - // Completion is a deliberate button press, never a side effect of leaving the page. const beaconMock = vi.fn(() => true); vi.stubGlobal("navigator", { ...navigator, sendBeacon: beaconMock }); - render(); + renderBanner("https://claude.ai"); window.dispatchEvent(new Event("pagehide")); diff --git a/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.tsx b/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.tsx index cea42f916f8..0d6e708f734 100644 --- a/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.tsx +++ b/ui/litellm-dashboard/src/components/chat/ConnectFlowBanner.tsx @@ -2,30 +2,18 @@ import React from "react"; import { CheckCircle } from "lucide-react"; -import { getProxyBaseUrl } from "@/components/networking"; +import { getProxyBaseUrl, ConnectFlowStatus } from "@/components/networking"; +import { OAuth2ConnectButton } from "@/components/chat/MCPAppsPanel"; interface Props { flowHandle: string; - clientOrigin: string | null; + flow?: ConnectFlowStatus; + accessToken: string; + onConnected: () => void; + failed: boolean; } -/** - * The interlude shown when a DCR client (Claude Desktop, MCP Inspector) sends the user - * through the gateway sign-in and lands them on the apps grid to authorize servers. The - * grid below authorizes individual servers into the per-user vault; this banner is the - * finish step that returns the user to the client. - * - * Finishing requires the explicit "Finish connecting" button: a native form POST to the proxy's - * /authorize/complete, which mints the gateway authorization code and 303-redirects to the DCR - * client's own redirect URI (the full-page navigation carries the HttpOnly per-flow cookie and - * follows the cross-origin redirect to the client's loopback). - * - * The button press IS the consent gate and must not be bypassed. An earlier version auto-finished - * on tab close via navigator.sendBeacon; that let an attacker who lured a signed-in victim to their - * own client's authorize URL harvest a victim-bound code the moment the victim closed the tab - * (no click). Merely visiting the authorize URL is attacker-inducible, so completion has to be a - * deliberate user action, not a side effect of leaving the page. - */ +/** Finish remains an explicit POST because a cross-site navigation must never mint a code. */ export function isLoopbackOrigin(origin: string | null): boolean { if (!origin) return false; try { @@ -36,10 +24,44 @@ export function isLoopbackOrigin(origin: string | null): boolean { } } -const ConnectFlowBanner: React.FC = ({ flowHandle, clientOrigin }) => { +const copyFor = (flow: ConnectFlowStatus | undefined, failed: boolean): readonly [string, string] => { + const clientLabel = flow?.client_origin ?? "the application"; + const serverLabel = flow?.server_name ?? "the requested MCP server"; + if (failed || flow === undefined || flow.state === "stale") { + return [ + "The connection cannot continue", + `The gateway could not validate this connection. Cancel to return to ${clientLabel}.`, + ]; + } + if (flow.state === "unscoped") { + return [ + `Connect your MCP servers to ${clientLabel}`, + `Authorize the servers you want to use below, then click Finish connecting to return to ${clientLabel}.`, + ]; + } + if (flow.state === "interactive" && !flow.connected) { + return [ + `Allow ${clientLabel} to use ${serverLabel}`, + `Authorize ${serverLabel} below to continue, or cancel to send ${clientLabel} away.`, + ]; + } + return [ + `Allow ${clientLabel} to use ${serverLabel}`, + `Click Finish connecting to give ${clientLabel} access to ${serverLabel} as you.`, + ]; +}; + +const ConnectFlowBanner: React.FC = ({ flowHandle, flow, accessToken, onConnected, failed }) => { const action = `${getProxyBaseUrl()}/authorize/complete`; - const clientLabel = clientOrigin ?? "the application"; - const loopbackClient = isLoopbackOrigin(clientOrigin); + const state = failed || flow === undefined ? "stale" : flow.state; + const canFinish = state === "unscoped" || (state !== "stale" && flow?.connected === true); + const canCancel = state !== "unscoped"; + const loopbackClient = isLoopbackOrigin(flow?.client_origin ?? null); + const vendorServer = + state === "interactive" && flow?.connected === false && flow.server_id !== null + ? { server_id: flow.server_id, server_name: flow.server_name } + : null; + const copy = copyFor(flow, failed); return (
@@ -47,27 +69,48 @@ const ConnectFlowBanner: React.FC = ({ flowHandle, clientOrigin }) => {
-

Connect your MCP servers to {clientLabel}

-

- Authorize the servers you want to use below, then click Finish connecting to return to {clientLabel}. -

+

{copy[0]}

+

{copy[1]}

-
- - - {loopbackClient && ( - +
+ {vendorServer !== null && ( + )} - +
+ + {canFinish && ( + + )} + {canCancel && ( + + )} + {loopbackClient && ( + + )} +
+
); diff --git a/ui/litellm-dashboard/src/components/chat/ConnectFlowSurface.test.tsx b/ui/litellm-dashboard/src/components/chat/ConnectFlowSurface.test.tsx new file mode 100644 index 00000000000..cf59dcfc368 --- /dev/null +++ b/ui/litellm-dashboard/src/components/chat/ConnectFlowSurface.test.tsx @@ -0,0 +1,118 @@ +import { afterEach, describe, expect, it, vi } from "vitest"; +import { act, render, screen, waitFor } from "@testing-library/react"; +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import ConnectFlowSurface from "./ConnectFlowSurface"; +import { fetchConnectFlow } from "@/components/networking"; + +const { startOAuthFlow, state, onSuccess } = vi.hoisted(() => ({ + startOAuthFlow: vi.fn(), + onSuccess: { current: undefined as (() => void) | undefined }, + state: { oauthReturn: null as string | null, connectFlow: null as string | null }, +})); + +vi.mock("next/navigation", () => ({ + useRouter: () => ({ replace: vi.fn() }), + useSearchParams: () => ({ + get: (key: string) => ({ mcpOauthReturn: state.oauthReturn, connect_flow: state.connectFlow })[key] ?? null, + }), +})); +vi.mock("@/components/networking", async (importOriginal) => ({ + ...(await importOriginal()), + fetchConnectFlow: vi.fn(), + getProxyBaseUrl: () => "https://gateway.example.com", +})); +vi.mock("@/components/chat/MCPAppsPanel", async (importOriginal) => ({ + ...(await importOriginal()), + default: () =>
, +})); +vi.mock("@/hooks/useUserMcpOAuthFlow", () => ({ + useUserMcpOAuthFlow: ({ onSuccess: success }: { onSuccess: () => void }) => { + onSuccess.current = success; + return { startOAuthFlow, status: "idle" }; + }, +})); + +const flow = (state: "unscoped" | "interactive" | "m2m" | "stale", connected: boolean | null = null) => ({ + state, + client_origin: "https://claude.ai", + server_id: state === "interactive" || state === "m2m" ? "s-design" : null, + server_name: state === "interactive" || state === "m2m" ? "design_tool" : null, + connected, +}); + +const renderSurface = () => + render( + + + , + ); + +afterEach(() => { + state.oauthReturn = null; + state.connectFlow = null; + onSuccess.current = undefined; + sessionStorage.clear(); + vi.clearAllMocks(); +}); + +describe("ConnectFlowSurface", () => { + it.each([ + { result: flow("unscoped"), grid: true, finish: true, cancel: false, oauthStarts: 0 }, + { result: flow("interactive", false), grid: false, finish: false, cancel: true, oauthStarts: 1 }, + { result: flow("interactive", true), grid: false, finish: true, cancel: true, oauthStarts: 0 }, + { result: flow("m2m", true), grid: false, finish: true, cancel: true, oauthStarts: 0 }, + { result: flow("stale"), grid: false, finish: false, cancel: true, oauthStarts: 0 }, + ])( + "renders $result.state without widening its action surface", + async ({ result, grid, finish, cancel, oauthStarts }) => { + state.connectFlow = "flow-handle-123"; + vi.mocked(fetchConnectFlow).mockResolvedValue(result); + renderSurface(); + + await screen.findByRole("button", { name: /finish connecting|cancel|connect/i }); + await waitFor(() => expect(startOAuthFlow).toHaveBeenCalledTimes(oauthStarts)); + expect(screen.queryByTestId("mcp-apps-panel") !== null).toBe(grid); + expect(screen.queryByRole("button", { name: /finish connecting/i }) !== null).toBe(finish); + expect(screen.queryByRole("button", { name: "Cancel" }) !== null).toBe(cancel); + }, + ); + + it("keeps the grid and Finish hidden until the gateway accepts a handle", () => { + state.connectFlow = "flow-handle-123"; + vi.mocked(fetchConnectFlow).mockReturnValue(new Promise(() => {})); + renderSurface(); + + expect(screen.queryByTestId("mcp-apps-panel")).not.toBeInTheDocument(); + expect(screen.queryByRole("button", { name: /finish connecting/i })).not.toBeInTheDocument(); + expect(screen.getByRole("button", { name: "Cancel" })).toHaveAttribute("value", "deny"); + }); + + it("keeps the grid and Finish hidden when flow validation fails", async () => { + state.connectFlow = "invalid-handle"; + vi.mocked(fetchConnectFlow).mockRejectedValue(new Error("invalid flow")); + renderSurface(); + + await screen.findByRole("button", { name: "Cancel" }); + expect(screen.queryByTestId("mcp-apps-panel")).not.toBeInTheDocument(); + expect(screen.queryByRole("button", { name: /finish connecting/i })).not.toBeInTheDocument(); + }); + + it("refetches the sealed flow after the vendor connection completes", async () => { + state.connectFlow = "flow-handle-123"; + vi.mocked(fetchConnectFlow) + .mockResolvedValueOnce(flow("interactive", false)) + .mockResolvedValueOnce(flow("interactive", true)); + renderSurface(); + + await waitFor(() => expect(startOAuthFlow).toHaveBeenCalledOnce()); + await act(async () => onSuccess.current?.()); + + await screen.findByRole("button", { name: /finish connecting/i }); + }); + + it("renders the ordinary panel without a flow handle", () => { + renderSurface(); + expect(fetchConnectFlow).not.toHaveBeenCalled(); + expect(screen.getByTestId("mcp-apps-panel")).toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/chat/ConnectFlowSurface.tsx b/ui/litellm-dashboard/src/components/chat/ConnectFlowSurface.tsx new file mode 100644 index 00000000000..33a61fceacf --- /dev/null +++ b/ui/litellm-dashboard/src/components/chat/ConnectFlowSurface.tsx @@ -0,0 +1,59 @@ +"use client"; + +import React, { useEffect } from "react"; +import { useRouter, useSearchParams } from "next/navigation"; +import { useQuery } from "@tanstack/react-query"; +import MCPAppsPanel from "@/components/chat/MCPAppsPanel"; +import ConnectFlowBanner from "@/components/chat/ConnectFlowBanner"; +import { fetchConnectFlow } from "@/components/networking"; + +interface Props { + accessToken: string; + selectedServers: string[]; + onChange: (servers: string[]) => void; +} + +/** Renders the sealed gateway connect flow without trusting URL context. */ +const ConnectFlowSurface: React.FC = ({ accessToken, selectedServers, onChange }) => { + const router = useRouter(); + const searchParams = useSearchParams(); + const oauthReturn = searchParams.get("mcpOauthReturn"); + const connectFlow = searchParams.get("connect_flow"); + + useEffect(() => { + if (oauthReturn) { + const url = new URL(window.location.href); + url.searchParams.delete("mcpOauthReturn"); + router.replace(url.pathname + url.search); + } + }, [oauthReturn, router]); + + const flowQuery = { + queryKey: ["gateway-connect-flow", connectFlow], + queryFn: () => fetchConnectFlow(connectFlow!), + enabled: !!connectFlow, + retry: false, + }; + const { data: flow, isError, refetch } = useQuery(flowQuery); + + if (connectFlow === null) { + return ; + } + + return ( + <> + + {flow?.state === "unscoped" && ( + + )} + + ); +}; + +export default ConnectFlowSurface; diff --git a/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.test.tsx b/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.test.tsx index e8795c36bc6..c9609405676 100644 --- a/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.test.tsx +++ b/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.test.tsx @@ -88,6 +88,13 @@ describe("MCPAppsPanel logos", () => { }); const connectServers = [ + { + server_id: "s-m2m", + server_name: "service_tool", + auth_type: "oauth2", + oauth2_flow: "client_credentials", + connected_app_reachable: true, + }, { server_id: "s-reach", server_name: "reachable_srv", @@ -124,6 +131,8 @@ describe("MCPAppsPanel connected-app reachability (LIT-4861)", () => { expect(vi.mocked(fetchMCPServers)).toHaveBeenCalledWith("tok", undefined, true); expect(screen.queryByText("unreachable_srv")).not.toBeInTheDocument(); expect(screen.getByText("Connected (1)")).toBeInTheDocument(); + expect(screen.getByText("service_tool")).toBeInTheDocument(); + expect(screen.queryByText("Connect", { exact: true })).not.toBeInTheDocument(); const toolCountFetchedIds = vi.mocked(listMCPTools).mock.calls.map((call) => call[1]); expect(toolCountFetchedIds).toContain("s-reach"); expect(toolCountFetchedIds).not.toContain("s-unreach"); diff --git a/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.tsx b/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.tsx index 38a095c067d..1fee8923e94 100644 --- a/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.tsx +++ b/ui/litellm-dashboard/src/components/chat/MCPAppsPanel.tsx @@ -13,23 +13,32 @@ import { getMCPOAuthUserCredentialStatus, listMCPTools, } from "../networking"; -import { AUTH_TYPE, MCPServer, MCPTool, handleTransport, isUnsupportedOnGatewayConnect } from "../mcp_tools/types"; +import { + getMcpOAuthMode, + MCPServer, + MCPTool, + handleTransport, + isUnsupportedOnGatewayConnect, +} from "../mcp_tools/types"; import { Logo } from "@/components/molecules/logo/Logo"; import { toast } from "@/lib/toast"; import { useUserMcpOAuthFlow } from "@/hooks/useUserMcpOAuthFlow"; +import { getSecureItem, setSecureItem } from "@/utils/secureStorage"; interface OAuth2ConnectButtonProps { - server: MCPServer; + server: Pick; accessToken: string; onConnect: (serverId: string) => void; variant?: "badge" | "button"; + autoStartKey?: string | null; } -const OAuth2ConnectButton: React.FC = ({ +export const OAuth2ConnectButton: React.FC = ({ server, accessToken, onConnect, variant = "badge", + autoStartKey = null, }) => { const name = server.server_name ?? server.alias ?? server.server_id; const { startOAuthFlow, status } = useUserMcpOAuthFlow({ @@ -39,6 +48,12 @@ const OAuth2ConnectButton: React.FC = ({ onSuccess: useCallback(() => onConnect(server.server_id), [onConnect, server.server_id]), }); + useEffect(() => { + if (autoStartKey === null || status !== "idle" || getSecureItem(autoStartKey) !== null) return; + setSecureItem(autoStartKey, "1"); + startOAuthFlow(); + }, [autoStartKey, status, startOAuthFlow]); + const loading = status === "authorizing" || status === "exchanging"; if (variant === "button") { @@ -190,7 +205,7 @@ const MCPAppsPanel: React.FC = ({ accessToken, selectedServers, onChange, if (!isCurrentLoad()) return; const list: MCPServer[] = Array.isArray(serverData) ? serverData : serverData?.data ?? []; const reachable = connectMode ? list.filter((s) => s.connected_app_reachable !== false) : list; - const oauthServers = reachable.filter((s) => s.auth_type === AUTH_TYPE.OAUTH2); + const oauthServers = reachable.filter((s) => getMcpOAuthMode(s) === "authorization_code"); commitServers(reachable); setOauthChecking(new Set(oauthServers.map((s) => s.server_id))); setLoading(false); @@ -274,7 +289,10 @@ const MCPAppsPanel: React.FC = ({ accessToken, selectedServers, onChange, {unavailabilityLabel} ); } - if (server.auth_type === AUTH_TYPE.OAUTH2) { + if (getMcpOAuthMode(server) === "m2m") { + return ; + } + if (getMcpOAuthMode(server) === "authorization_code") { if (oauthConnected.has(server.server_id)) { return ; } @@ -339,7 +357,10 @@ const MCPAppsPanel: React.FC = ({ accessToken, selectedServers, onChange, if (unavailabilityLabel !== null) { return {unavailabilityLabel}; } - if (detailServer.auth_type !== AUTH_TYPE.OAUTH2) { + if (getMcpOAuthMode(detailServer) === "m2m") { + return Authorized; + } + if (getMcpOAuthMode(detailServer) !== "authorization_code") { return ( +
+
+ + 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) => ( +
+ +