mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-10 22:41:41 +00:00
fix(guardrails): handle positional args for request_data and input_type in log_guardrail_information
When apply_guardrail is called with positional arguments, the decorator now reads request_data from args[2] and input_type from args[3] as fallbacks when they are not present in kwargs. Also adds a test asserting the positional input_type path works correctly. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
parent
479ac6ea95
commit
3438955a64
2 changed files with 53 additions and 6 deletions
|
|
@ -888,12 +888,18 @@ def log_guardrail_information(func):
|
|||
async def async_wrapper(*args, **kwargs):
|
||||
start_time = datetime.now() # Move start_time inside the wrapper
|
||||
self: CustomGuardrail = args[0]
|
||||
request_data: dict = kwargs.get("data") or kwargs.get("request_data") or {}
|
||||
# apply_guardrail signature: (self, inputs, request_data, input_type, ...)
|
||||
# Support both keyword and positional callers for request_data (args[2])
|
||||
request_data: dict = (
|
||||
kwargs.get("data")
|
||||
or kwargs.get("request_data")
|
||||
or (args[2] if len(args) > 2 and isinstance(args[2], dict) else {})
|
||||
)
|
||||
event_type = _infer_event_type_from_function_name(func.__name__)
|
||||
|
||||
# For apply_guardrail, infer event type from input_type kwarg
|
||||
# For apply_guardrail, infer event type from input_type (kwarg or positional arg[3])
|
||||
if event_type is None and func.__name__ == "apply_guardrail":
|
||||
_input_type = kwargs.get("input_type")
|
||||
_input_type = kwargs.get("input_type") or (args[3] if len(args) > 3 else None)
|
||||
if _input_type == "request":
|
||||
event_type = GuardrailEventHooks.pre_call
|
||||
elif _input_type == "response":
|
||||
|
|
@ -929,12 +935,18 @@ def log_guardrail_information(func):
|
|||
def sync_wrapper(*args, **kwargs):
|
||||
start_time = datetime.now() # Move start_time inside the wrapper
|
||||
self: CustomGuardrail = args[0]
|
||||
request_data: dict = kwargs.get("data") or kwargs.get("request_data") or {}
|
||||
# apply_guardrail signature: (self, inputs, request_data, input_type, ...)
|
||||
# Support both keyword and positional callers for request_data (args[2])
|
||||
request_data: dict = (
|
||||
kwargs.get("data")
|
||||
or kwargs.get("request_data")
|
||||
or (args[2] if len(args) > 2 and isinstance(args[2], dict) else {})
|
||||
)
|
||||
event_type = _infer_event_type_from_function_name(func.__name__)
|
||||
|
||||
# For apply_guardrail, infer event type from input_type kwarg
|
||||
# For apply_guardrail, infer event type from input_type (kwarg or positional arg[3])
|
||||
if event_type is None and func.__name__ == "apply_guardrail":
|
||||
_input_type = kwargs.get("input_type")
|
||||
_input_type = kwargs.get("input_type") or (args[3] if len(args) > 3 else None)
|
||||
if _input_type == "request":
|
||||
event_type = GuardrailEventHooks.pre_call
|
||||
elif _input_type == "response":
|
||||
|
|
|
|||
|
|
@ -1021,3 +1021,38 @@ class TestLogGuardrailInformationApplyGuardrailEventType:
|
|||
f"guardrail_mode must be a string, got {type(guardrail_mode)}: {guardrail_mode}"
|
||||
)
|
||||
assert guardrail_mode == GuardrailEventHooks.post_call
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_positional_input_type_logs_correct_event_type(self):
|
||||
"""
|
||||
When apply_guardrail is called with input_type passed positionally (args[3]),
|
||||
the decorator must still infer the correct event_type. This guards against
|
||||
regressions where only kwargs.get("input_type") is checked.
|
||||
"""
|
||||
mode = Mode(tags={}, default="pre_call")
|
||||
|
||||
class _TestGuardrail(CustomGuardrail):
|
||||
@log_guardrail_information
|
||||
async def apply_guardrail(
|
||||
self,
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
request_data: dict,
|
||||
input_type,
|
||||
logging_obj=None,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
return inputs
|
||||
|
||||
guardrail = _TestGuardrail(guardrail_name="test", event_hook=mode)
|
||||
request_data: dict = {"metadata": {}}
|
||||
inputs = GenericGuardrailAPIInputs(texts=["hello"])
|
||||
|
||||
# Pass input_type positionally — it lands in args[3], not kwargs
|
||||
await guardrail.apply_guardrail(inputs, request_data, "request")
|
||||
|
||||
slg_list = request_data["metadata"]["standard_logging_guardrail_information"]
|
||||
assert len(slg_list) == 1
|
||||
guardrail_mode = slg_list[0]["guardrail_mode"]
|
||||
assert isinstance(guardrail_mode, str), (
|
||||
f"guardrail_mode must be a string, got {type(guardrail_mode)}: {guardrail_mode}"
|
||||
)
|
||||
assert guardrail_mode == GuardrailEventHooks.pre_call
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue