improved registration flow

llama_stack/providers/datatypes.py

@@ -16,7 +16,7 @@ from llama_stack.apis.eval_tasks import EvalTaskDef
 from llama_stack.apis.memory_banks import MemoryBankDef
 from llama_stack.apis.models import ModelDef
 from llama_stack.apis.scoring_functions import ScoringFnDef
-from llama_stack.apis.shields import Shield
+from llama_stack.apis.shields import Shield, ShieldType
 
 
 @json_schema_type
@@ -51,6 +51,8 @@ class ModelsProtocolPrivate(Protocol):
 class ShieldsProtocolPrivate(Protocol):
     async def register_shield(self, shield: Shield) -> None: ...
 
+    async def supported_shield_types(self) -> List[ShieldType]: ...
+
 
 class MemoryBanksProtocolPrivate(Protocol):
     async def list_memory_banks(self) -> List[MemoryBankDef]: ...

llama_stack/providers/inline/safety/meta_reference/safety.py

@@ -21,6 +21,7 @@ from .prompt_guard import InjectionShield, JailbreakShield, PromptGuardShield
 
 
 PROMPT_GUARD_MODEL = "Prompt-Guard-86M"
+SUPPORTED_SHIELDS = [ShieldType.llama_guard, ShieldType.prompt_guard]
 
 
 class MetaReferenceSafetyImpl(Safety, ShieldsProtocolPrivate):
@@ -46,6 +47,9 @@ class MetaReferenceSafetyImpl(Safety, ShieldsProtocolPrivate):
         if shield.shield_type not in self.available_shields:
             raise ValueError(f"Shield type {shield.shield_type} not supported")
 
+    async def supported_shield_types(self) -> List[ShieldType]:
+        return SUPPORTED_SHIELDS
+
     async def run_shield(
         self,
         shield_id: str,

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

@@ -21,7 +21,7 @@ logger = logging.getLogger(__name__)
 
 
 BEDROCK_SUPPORTED_SHIELDS = [
-    ShieldType.generic_content_shield.value,
+    ShieldType.generic_content_shield,
 ]
 
 
@@ -53,6 +53,9 @@ class BedrockSafetyAdapter(Safety, ShieldsProtocolPrivate):
                 f"Shield {shield.identifier} with version {shield.params['guardrailVersion']} not found in Bedrock"
             )
 
+    async def supported_shield_types(self) -> List[ShieldType]:
+        return BEDROCK_SUPPORTED_SHIELDS
+
     async def run_shield(
         self, shield_id: str, messages: List[Message], params: Dict[str, Any] = None
     ) -> RunShieldResponse:

llama_stack/providers/tests/safety/fixtures.py

@@ -7,7 +7,7 @@
 import pytest
 import pytest_asyncio
 
-from llama_stack.apis.shields import Shield, ShieldType
+from llama_stack.apis.shields import ShieldType
 
 from llama_stack.distribution.datatypes import Api, Provider
 from llama_stack.providers.inline.safety.meta_reference import (
@@ -95,10 +95,10 @@ async def safety_stack(inference_model, safety_model, request):
     shields_impl = impls[Api.shields]
 
     # Register the appropriate shield based on provider type
-    provider_id = safety_fixture.providers[0].provider_id
     provider_type = safety_fixture.providers[0].provider_type
 
     shield_config = {}
+    shield_type = ShieldType.llama_guard
     identifier = "llama_guard"
     if provider_type == "meta-reference":
         shield_config["model"] = safety_model
@@ -107,12 +107,11 @@ async def safety_stack(inference_model, safety_model, request):
     elif provider_type == "remote::bedrock":
         identifier = get_env_or_fail("BEDROCK_GUARDRAIL_IDENTIFIER")
         shield_config["guardrailVersion"] = get_env_or_fail("BEDROCK_GUARDRAIL_VERSION")
+        shield_type = ShieldType.generic_content_shield
 
-    # Create shield
-    shield = Shield(
-        identifier=identifier,
-        shield_type=ShieldType.llama_guard,
-        provider_id=provider_id,
+    shield = await shields_impl.register_shield(
+        shield_id=identifier,
+        shield_type=shield_type,
         params=shield_config,
     )
 

llama_stack/providers/tests/safety/test_safety.py

@@ -19,9 +19,15 @@ from llama_stack.distribution.datatypes import *  # noqa: F403
 
 class TestSafety:
     @pytest.mark.asyncio
-    async def test_shield_list(self, safety_stack):
+    async def test_new_shield(self, safety_stack):
         _, shields_impl, shield = safety_stack
-        await shields_impl.register_shield(shield)
+        assert shield is not None
+        assert shield.provider_resource_identifier == shield.identifier
+        assert shield.provider_id is not None
+
+    @pytest.mark.asyncio
+    async def test_shield_list(self, safety_stack):
+        _, shields_impl, _ = safety_stack
         response = await shields_impl.list_shields()
         assert isinstance(response, list)
         assert len(response) >= 1
@@ -32,9 +38,7 @@ class TestSafety:
 
     @pytest.mark.asyncio
     async def test_run_shield(self, safety_stack):
-        safety_impl, shields_impl, shield = safety_stack
-
-        await shields_impl.register_shield(shield)
+        safety_impl, _, shield = safety_stack
 
         response = await safety_impl.run_shield(
             shield_id=shield.identifier,