From b5d130fe2ad3dd29eefff437ba0ea1cd5da3af0e Mon Sep 17 00:00:00 2001 From: Dinesh Yeduguru Date: Thu, 7 Nov 2024 21:04:08 -0800 Subject: [PATCH] use correct shield impl in meta ref --- .../providers/inline/safety/meta_reference/safety.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/llama_stack/providers/inline/safety/meta_reference/safety.py b/llama_stack/providers/inline/safety/meta_reference/safety.py index ce7eaf591..f73408166 100644 --- a/llama_stack/providers/inline/safety/meta_reference/safety.py +++ b/llama_stack/providers/inline/safety/meta_reference/safety.py @@ -64,11 +64,14 @@ class MetaReferenceSafetyImpl(Safety, ShieldsProtocolPrivate): # TODO: we can refactor ShieldBase, etc. to be inline with the API types res = await shield_impl.run(messages) violation = None - if res.is_violation and shield.on_violation_action != OnViolationAction.IGNORE: + if ( + res.is_violation + and shield_impl.on_violation_action != OnViolationAction.IGNORE + ): violation = SafetyViolation( violation_level=( ViolationLevel.ERROR - if shield.on_violation_action == OnViolationAction.RAISE + if shield_impl.on_violation_action == OnViolationAction.RAISE else ViolationLevel.WARN ), user_message=res.violation_return_message,