From 8fc2d212a29388da9efffcf9cc39c72c0f8fb246 Mon Sep 17 00:00:00 2001 From: Xi Yan Date: Wed, 6 Nov 2024 16:30:47 -0800 Subject: [PATCH] fix safety signature mismatch (#388) * fix safety sig * shield_type->identifier --- .../providers/inline/meta_reference/safety/safety.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/llama_stack/providers/inline/meta_reference/safety/safety.py b/llama_stack/providers/inline/meta_reference/safety/safety.py index 28c78b65c..2d0db7624 100644 --- a/llama_stack/providers/inline/meta_reference/safety/safety.py +++ b/llama_stack/providers/inline/meta_reference/safety/safety.py @@ -57,13 +57,13 @@ class MetaReferenceSafetyImpl(Safety, ShieldsProtocolPrivate): async def run_shield( self, - shield_type: str, + identifier: str, messages: List[Message], params: Dict[str, Any] = None, ) -> RunShieldResponse: - shield_def = await self.shield_store.get_shield(shield_type) + shield_def = await self.shield_store.get_shield(identifier) if not shield_def: - raise ValueError(f"Unknown shield {shield_type}") + raise ValueError(f"Unknown shield {identifier}") shield = self.get_shield_impl(shield_def)