diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index c8e5610b4a9..9bb613654e4 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -831,10 +831,10 @@ class CustomGuardrail(CustomLogger): # should run guardrail litellm_guardrails: Final = request_data.get("guardrails") if litellm_guardrails is None or not isinstance(litellm_guardrails, list): - return response + return None if self.should_run_guardrail(data=request_data, event_type=GuardrailEventHooks.post_call) is not True: - return response + return None # CHECK IF GUARDRAIL REJECTS THE REQUEST result: Final = await self.async_post_call_success_hook( @@ -850,7 +850,7 @@ class CustomGuardrail(CustomLogger): ) if not self._is_valid_response_type(result): - return response + return None return result diff --git a/litellm/utils.py b/litellm/utils.py index 239673a20d9..e9d3b263a91 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -1284,16 +1284,18 @@ async def async_post_call_success_deployment_hook( except ValueError: typed_call_type = None # unknown call type + modified_response = response + CustomLogger: Final = _get_cached_custom_logger() for callback in litellm.callbacks: if isinstance(callback, CustomLogger): result = await callback.async_post_call_success_deployment_hook( - request_data, cast(LLMResponseTypes, response), typed_call_type + request_data, cast(LLMResponseTypes, modified_response), typed_call_type ) if result is not None: - return result + modified_response = result - return response + return modified_response async def async_post_call_failure_deployment_hook( diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/test_litellm/integrations/test_custom_guardrail.py index 7d70b9a8862..6edc2c9bf77 100644 --- a/tests/test_litellm/integrations/test_custom_guardrail.py +++ b/tests/test_litellm/integrations/test_custom_guardrail.py @@ -1091,8 +1091,8 @@ class TestCustomGuardrailPassthroughSupport: call_type=CallTypes.allm_passthrough_route, ) - # When result is None, should return the original response - assert result == mock_response + # None means the guardrail did not modify the response (LIT-5863 contract) + assert result is None @pytest.mark.asyncio async def test_async_post_call_success_deployment_hook_with_none_call_type(self): @@ -1120,8 +1120,8 @@ class TestCustomGuardrailPassthroughSupport: call_type=None, ) - # Should return the original response when result is None - assert result == mock_response + # None means the guardrail did not modify the response (LIT-5863 contract) + assert result is None def test_is_valid_response_type_with_none(self): """ @@ -2436,3 +2436,73 @@ class TestLoggingOnlyApplyGuardrail: assert guardrail.calls == [("request", ["hello there"]), ("response", ["general kenobi"])] entries = logging_obj.model_call_details["standard_logging_object"]["guardrail_information"] assert [e["guardrail_status"] for e in entries] == ["success", "success"] + + +class TestCustomGuardrailPostCallSuccessDeploymentHook: + """Regression tests for LIT-5863: this hook answering the unmodified response instead of + None made the utils.py dispatcher treat the guardrail as having modified the response, + which starved every later callback in litellm.callbacks (notably the lazily-appended + VectorStorePreCallHook that attaches provider_specific_fields["search_results"]).""" + + @pytest.mark.asyncio + async def test_returns_none_when_request_has_no_guardrails(self): + from litellm.types.utils import ModelResponse + + guardrail = CustomGuardrail(guardrail_name="test-guardrail") + response = ModelResponse() + + assert ( + await guardrail.async_post_call_success_deployment_hook( + request_data={}, response=response, call_type=CallTypes.acompletion + ) + is None + ) + assert ( + await guardrail.async_post_call_success_deployment_hook( + request_data={"guardrails": "not-a-list"}, response=response, call_type=CallTypes.acompletion + ) + is None + ) + + @pytest.mark.asyncio + async def test_returns_none_when_guardrail_should_not_run(self): + from litellm.types.guardrails import GuardrailEventHooks + from litellm.types.utils import ModelResponse + + guardrail = CustomGuardrail( + guardrail_name="test-guardrail", + event_hook=GuardrailEventHooks.pre_call, + ) + response = ModelResponse() + + result = await guardrail.async_post_call_success_deployment_hook( + request_data={"guardrails": ["test-guardrail"]}, + response=response, + call_type=CallTypes.acompletion, + ) + + assert result is None + + @pytest.mark.asyncio + async def test_returns_modified_response_when_guardrail_runs(self): + from litellm.types.guardrails import GuardrailEventHooks + from litellm.types.utils import ModelResponse + + replacement = ModelResponse() + + class ReplacingGuardrail(CustomGuardrail): + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + return replacement + + guardrail = ReplacingGuardrail( + guardrail_name="test-guardrail", + event_hook=GuardrailEventHooks.post_call, + ) + + result = await guardrail.async_post_call_success_deployment_hook( + request_data={"guardrails": ["test-guardrail"]}, + response=ModelResponse(), + call_type=CallTypes.acompletion, + ) + + assert result is replacement diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 99c69d6bc31..5bfbbd5b50e 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -40,6 +40,7 @@ from litellm.utils import ( _is_streaming_request, _snapshot_exception_for_hook, async_post_call_failure_deployment_hook, + async_post_call_success_deployment_hook, client, get_llm_provider, get_non_default_completion_params, @@ -5808,6 +5809,90 @@ class TestHuggingFaceConfigFetch: assert request_timeout["read"] == HF_CONFIG_FETCH_TIMEOUT_SECONDS +@pytest.mark.asyncio +async def test_success_deployment_hook_chains_past_callback_returning_response( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Regression (LIT-5863): the dispatcher must run every callback, chaining each non-None + result into the next call, instead of returning at the first callback answering non-None. + A guardrail answering with the unmodified response used to starve every callback after it.""" + from litellm.types.utils import ModelResponse + + original = ModelResponse() + replacement = ModelResponse() + + class PassthroughLogger(CustomLogger): + async def async_post_call_success_deployment_hook(self, request_data, response, call_type): + return response + + class ReplacingLogger(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.seen: list = [] + + async def async_post_call_success_deployment_hook(self, request_data, response, call_type): + self.seen.append(response) + return replacement + + class ObservingLogger(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.seen: list = [] + + async def async_post_call_success_deployment_hook(self, request_data, response, call_type): + self.seen.append(response) + return None + + replacer = ReplacingLogger() + observer = ObservingLogger() + monkeypatch.setattr(litellm, "callbacks", [PassthroughLogger(), replacer, observer]) + + result = await async_post_call_success_deployment_hook( + request_data={}, response=original, call_type=CallTypes.acompletion + ) + + assert replacer.seen == [original] + assert observer.seen == [replacement] + assert result is replacement + + +@pytest.mark.asyncio +async def test_registered_guardrail_does_not_starve_vector_store_search_results( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Regression (LIT-5863): with any guardrail registered ahead of the lazily-appended + VectorStorePreCallHook, /v1/chat/completions responses lost + provider_specific_fields["search_results"] because the guardrail answered the unmodified + response and the dispatcher stopped there.""" + from types import SimpleNamespace + + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import ( + VectorStorePreCallHook, + ) + from litellm.types.utils import ModelResponse + + search_results: Final = [{"search_query": "coolant", "data": [{"content": [{"text": "Cryoline-9", "type": "text"}]}]}] + logging_obj = SimpleNamespace(model_call_details={"search_results": search_results}) + response = ModelResponse(choices=[{"message": {"role": "assistant", "content": "Cryoline-9"}}]) + + monkeypatch.setattr( + litellm, + "callbacks", + [CustomGuardrail(guardrail_name="dummy-guardrail"), VectorStorePreCallHook()], + ) + + result = await async_post_call_success_deployment_hook( + request_data={"litellm_logging_obj": logging_obj}, + response=response, + call_type=CallTypes.acompletion, + ) + + provider_fields = result.choices[0].message.provider_specific_fields + assert provider_fields is not None + assert provider_fields["search_results"] == search_results + + class TestIsVisionExplicitlyDisabled: """github_copilot and chatgpt run an OAuth device flow inside get_llm_provider; the explicit-disable lookup must adopt the declared prefix instead of resolving it, exactly