init

llama_stack/providers/remote/safety/bedrock/bedrock.py

@@ -40,33 +40,12 @@ class BedrockSafetyAdapter(Safety, ShieldsProtocolPrivate):
     async def shutdown(self) -> None:
         pass
 
-    async def register_shield(self, shield: ShieldDef) -> None:
+    async def register_shield(self, shield: Shield) -> None:
         raise ValueError("Registering dynamic shields is not supported")
 
-    async def list_shields(self) -> List[ShieldDef]:
-        response = self.bedrock_client.list_guardrails()
-        shields = []
-        for guardrail in response["guardrails"]:
-            # populate the shield def with the guardrail id and version
-            shield_def = ShieldDef(
-                identifier=guardrail["id"],
-                shield_type=ShieldType.generic_content_shield.value,
-                params={
-                    "guardrailIdentifier": guardrail["id"],
-                    "guardrailVersion": guardrail["version"],
-                },
-            )
-            self.registered_shields.append(shield_def)
-            shields.append(shield_def)
-        return shields
-
     async def run_shield(
-        self, identifier: str, messages: List[Message], params: Dict[str, Any] = None
+        self, shield: Shield, messages: List[Message], params: Dict[str, Any] = None
     ) -> RunShieldResponse:
-        shield_def = await self.shield_store.get_shield(identifier)
-        if not shield_def:
-            raise ValueError(f"Unknown shield {identifier}")
-
         """This is the implementation for the bedrock guardrails. The input to the guardrails is to be of this format
         ```content = [
             {
@@ -81,7 +60,7 @@ class BedrockSafetyAdapter(Safety, ShieldsProtocolPrivate):
         They contain content, role . For now we will extract the content and default the "qualifiers": ["query"]
         """
 
-        shield_params = shield_def.params
+        shield_params = shield.params
         logger.debug(f"run_shield::{shield_params}::messages={messages}")
 
         # - convert the messages into format Bedrock expects
@@ -93,7 +72,7 @@ class BedrockSafetyAdapter(Safety, ShieldsProtocolPrivate):
         )
 
         response = self.bedrock_runtime_client.apply_guardrail(
-            guardrailIdentifier=shield_params["guardrailIdentifier"],
+            guardrailIdentifier=shield.identifier,
             guardrailVersion=shield_params["guardrailVersion"],
             source="OUTPUT",  # or 'INPUT' depending on your use case
             content=content_messages,

llama_stack/providers/remote/safety/sample/sample.py

@@ -14,7 +14,7 @@ class SampleSafetyImpl(Safety):
     def __init__(self, config: SampleConfig):
         self.config = config
 
-    async def register_shield(self, shield: ShieldDef) -> None:
+    async def register_shield(self, shield: Shield) -> None:
         # these are the safety shields the Llama Stack will use to route requests to this provider
         # perform validation here if necessary
         pass