Merge pull request #38984 from BerriAI/litellm_fix_search_results_with_guardrails

fix: attach vector store search_results when a guardrail is registered

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
 

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(

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

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