Merge ac6034ed1e into cbe89d2bdd

llama_stack/providers/datatypes.py

@@ -62,6 +62,8 @@ class ModelsProtocolPrivate(Protocol):
 class ShieldsProtocolPrivate(Protocol):
     async def register_shield(self, shield: Shield) -> None: ...
 
+    async def unregister_shield(self, identifier: str) -> None: ...
+
 
 class VectorDBsProtocolPrivate(Protocol):
     async def register_vector_db(self, vector_db: VectorDB) -> None: ...

llama_stack/providers/inline/safety/llama_guard/llama_guard.py

@@ -150,6 +150,11 @@ class LlamaGuardSafetyImpl(Safety, ShieldsProtocolPrivate):
         # The model will be validated during runtime when making inference calls
         pass
 
+    async def unregister_shield(self, identifier: str) -> None:
+        # LlamaGuard doesn't need to do anything special for unregistration
+        # The routing table handles the removal from the registry
+        pass
+
     async def run_shield(
         self,
         shield_id: str,

llama_stack/providers/inline/safety/prompt_guard/prompt_guard.py

@@ -46,6 +46,9 @@ class PromptGuardSafetyImpl(Safety, ShieldsProtocolPrivate):
         if shield.provider_resource_id != PROMPT_GUARD_MODEL:
             raise ValueError(f"Only {PROMPT_GUARD_MODEL} is supported for Prompt Guard. ")
 
+    async def unregister_shield(self, identifier: str) -> None:
+        pass
+
     async def run_shield(
         self,
         shield_id: str,

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

@@ -52,6 +52,9 @@ class BedrockSafetyAdapter(Safety, ShieldsProtocolPrivate):
                 f"Shield {shield.provider_resource_id} with version {shield.params['guardrailVersion']} not found in Bedrock"
             )
 
+    async def unregister_shield(self, identifier: str) -> None:
+        pass
+
     async def run_shield(
         self, shield_id: str, messages: list[Message], params: dict[str, Any] = None
     ) -> RunShieldResponse:

llama_stack/providers/remote/safety/nvidia/nvidia.py

@@ -40,6 +40,9 @@ class NVIDIASafetyAdapter(Safety, ShieldsProtocolPrivate):
         if not shield.provider_resource_id:
             raise ValueError("Shield model not provided.")
 
+    async def unregister_shield(self, identifier: str) -> None:
+        pass
+
     async def run_shield(
         self, shield_id: str, messages: list[Message], params: dict[str, Any] | None = None
     ) -> RunShieldResponse:

llama_stack/providers/remote/safety/sambanova/sambanova.py

@@ -68,6 +68,9 @@ class SambaNovaSafetyAdapter(Safety, ShieldsProtocolPrivate, NeedsRequestProvide
         ):
             logger.warning(f"Shield {shield.provider_resource_id} not available in {list_models_url}")
 
+    async def unregister_shield(self, identifier: str) -> None:
+        pass
+
     async def run_shield(
         self, shield_id: str, messages: list[Message], params: dict[str, Any] | None = None
     ) -> RunShieldResponse: