fix(guardrails): stop registered guardrails starving vector store search_results

Any registered guardrail made provider_specific_fields.search_results
vanish from /v1/chat/completions vector store responses, even when the
guardrail never ran. Two defects combined:

- CustomGuardrail.async_post_call_success_deployment_hook returned the
  response instead of None when it did not run, claiming a modification
  it never made
- the async_post_call_success_deployment_hook dispatcher in utils.py
  returned at the first non-None callback result, so the lazily appended
  VectorStorePreCallHook never got a chance to attach search_results

The hook now returns None when it does not run, and the dispatcher
chains non-None results through the remaining callbacks, matching the
pre-call dispatcher's behavior
This commit is contained in:
mateo-berri 2026-08-31 11:16:19 -07:00
parent f079e4061b
commit 5404a7a7c2
4 changed files with 167 additions and 10 deletions

View file

@ -828,10 +828,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(
@ -847,7 +847,7 @@ class CustomGuardrail(CustomLogger):
)
if not self._is_valid_response_type(result):
return response
return None
return result

View file

@ -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(

View file

@ -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):
"""
@ -2237,3 +2237,73 @@ class TestRecordsOwnGuardrailInformation:
)
assert _guardrail_entries(request_data) == []
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

View file

@ -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_api_key,
get_llm_provider,
@ -5765,3 +5766,87 @@ class TestHuggingFaceConfigFetch:
assert _get_max_position_embeddings("some-org/some-model") == 512
request_timeout = hf_config_route.calls.last.request.extensions["timeout"]
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