mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-08 03:08:45 +00:00
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:
parent
f079e4061b
commit
5404a7a7c2
4 changed files with 167 additions and 10 deletions
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue