Resource oriented design for shields (#399)

* init

* working bedrock tests

* bedrock test for inference fixes

* use env vars for bedrock guardrail vars

* add register in meta reference

* use correct shield impl in meta ref

* dont add together fixture

* right naming

* minor updates

* improved registration flow

* address feedback

---------

Co-authored-by: Dinesh Yeduguru <dineshyv@fb.com>

llama_stack/apis/resource.py

@@ -0,0 +1,38 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+from enum import Enum
+
+from llama_models.schema_utils import json_schema_type
+from pydantic import BaseModel, Field
+
+
+@json_schema_type
+class ResourceType(Enum):
+    model = "model"
+    shield = "shield"
+    memory_bank = "memory_bank"
+    dataset = "dataset"
+    scoring_function = "scoring_function"
+
+
+class Resource(BaseModel):
+    """Base class for all Llama Stack resources"""
+
+    identifier: str = Field(
+        description="Unique identifier for this resource in llama stack"
+    )
+
+    provider_resource_id: str = Field(
+        description="Unique identifier for this resource in the provider",
+        default=None,
+    )
+
+    provider_id: str = Field(description="ID of the provider that owns this resource")
+
+    type: ResourceType = Field(
+        description="Type of resource (e.g. 'model', 'shield', 'memory_bank', etc.)"
+    )

llama_stack/apis/safety/client.py

@@ -41,13 +41,13 @@ class SafetyClient(Safety):
         pass
 
     async def run_shield(
-        self, shield_type: str, messages: List[Message]
+        self, shield_id: str, messages: List[Message]
     ) -> RunShieldResponse:
         async with httpx.AsyncClient() as client:
             response = await client.post(
                 f"{self.base_url}/safety/run_shield",
                 json=dict(
-                    shield_type=shield_type,
+                    shield_id=shield_id,
                     messages=[encodable_dict(m) for m in messages],
                 ),
                 headers={
@@ -80,7 +80,7 @@ async def run_main(host: str, port: int, image_path: str = None):
         )
         cprint(f"User>{message.content}", "green")
         response = await client.run_shield(
-            shield_type="llama_guard",
+            shield_id="llama_guard",
             messages=[message],
         )
         print(response)
@@ -91,7 +91,7 @@ async def run_main(host: str, port: int, image_path: str = None):
     ]:
         cprint(f"User>{message.content}", "green")
         response = await client.run_shield(
-            shield_type="llama_guard",
+            shield_id="llama_guard",
             messages=[message],
         )
         print(response)

llama_stack/apis/safety/safety.py

@@ -39,7 +39,7 @@ class RunShieldResponse(BaseModel):
 
 
 class ShieldStore(Protocol):
-    async def get_shield(self, identifier: str) -> ShieldDef: ...
+    async def get_shield(self, identifier: str) -> Shield: ...
 
 
 @runtime_checkable
@@ -48,5 +48,8 @@ class Safety(Protocol):
 
     @webmethod(route="/safety/run_shield")
     async def run_shield(
-        self, identifier: str, messages: List[Message], params: Dict[str, Any] = None
+        self,
+        shield_id: str,
+        messages: List[Message],
+        params: Dict[str, Any] = None,
     ) -> RunShieldResponse: ...

llama_stack/apis/shields/client.py

@@ -5,7 +5,6 @@
 # the root directory of this source tree.
 
 import asyncio
-import json
 
 from typing import List, Optional
 
@@ -26,27 +25,38 @@ class ShieldsClient(Shields):
     async def shutdown(self) -> None:
         pass
 
-    async def list_shields(self) -> List[ShieldDefWithProvider]:
+    async def list_shields(self) -> List[Shield]:
         async with httpx.AsyncClient() as client:
             response = await client.get(
                 f"{self.base_url}/shields/list",
                 headers={"Content-Type": "application/json"},
             )
             response.raise_for_status()
-            return [ShieldDefWithProvider(**x) for x in response.json()]
+            return [Shield(**x) for x in response.json()]
 
-    async def register_shield(self, shield: ShieldDefWithProvider) -> None:
+    async def register_shield(
+        self,
+        shield_id: str,
+        shield_type: ShieldType,
+        provider_shield_id: Optional[str],
+        provider_id: Optional[str],
+        params: Optional[Dict[str, Any]],
+    ) -> None:
         async with httpx.AsyncClient() as client:
             response = await client.post(
                 f"{self.base_url}/shields/register",
                 json={
-                    "shield": json.loads(shield.json()),
+                    "shield_id": shield_id,
+                    "shield_type": shield_type,
+                    "provider_shield_id": provider_shield_id,
+                    "provider_id": provider_id,
+                    "params": params,
                 },
                 headers={"Content-Type": "application/json"},
             )
             response.raise_for_status()
 
-    async def get_shield(self, shield_type: str) -> Optional[ShieldDefWithProvider]:
+    async def get_shield(self, shield_type: str) -> Optional[Shield]:
         async with httpx.AsyncClient() as client:
             response = await client.get(
                 f"{self.base_url}/shields/get",
@@ -61,7 +71,7 @@ class ShieldsClient(Shields):
             if j is None:
                 return None
 
-            return ShieldDefWithProvider(**j)
+            return Shield(**j)
 
 
 async def run_main(host: str, port: int, stream: bool):

llama_stack/apis/shields/shields.py

@@ -8,7 +8,8 @@ from enum import Enum
 from typing import Any, Dict, List, Literal, Optional, Protocol, runtime_checkable
 
 from llama_models.schema_utils import json_schema_type, webmethod
-from pydantic import BaseModel, Field
+
+from llama_stack.apis.resource import Resource, ResourceType
 
 
 @json_schema_type
@@ -19,34 +20,29 @@ class ShieldType(Enum):
     prompt_guard = "prompt_guard"
 
 
-class ShieldDef(BaseModel):
-    identifier: str = Field(
-        description="A unique identifier for the shield type",
-    )
-    shield_type: str = Field(
-        description="The type of shield this is; the value is one of the ShieldType enum"
-    )
-    params: Dict[str, Any] = Field(
-        default_factory=dict,
-        description="Any additional parameters needed for this shield",
-    )
-
-
 @json_schema_type
-class ShieldDefWithProvider(ShieldDef):
-    type: Literal["shield"] = "shield"
-    provider_id: str = Field(
-        description="The provider ID for this shield type",
-    )
+class Shield(Resource):
+    """A safety shield resource that can be used to check content"""
+
+    type: Literal[ResourceType.shield.value] = ResourceType.shield.value
+    shield_type: ShieldType
+    params: Dict[str, Any] = {}
 
 
 @runtime_checkable
 class Shields(Protocol):
     @webmethod(route="/shields/list", method="GET")
-    async def list_shields(self) -> List[ShieldDefWithProvider]: ...
+    async def list_shields(self) -> List[Shield]: ...
 
     @webmethod(route="/shields/get", method="GET")
-    async def get_shield(self, identifier: str) -> Optional[ShieldDefWithProvider]: ...
+    async def get_shield(self, identifier: str) -> Optional[Shield]: ...
 
     @webmethod(route="/shields/register", method="POST")
-    async def register_shield(self, shield: ShieldDefWithProvider) -> None: ...
+    async def register_shield(
+        self,
+        shield_id: str,
+        shield_type: ShieldType,
+        provider_shield_id: Optional[str] = None,
+        provider_id: Optional[str] = None,
+        params: Optional[Dict[str, Any]] = None,
+    ) -> Shield: ...