init

llama_stack/apis/resource.py

@@ -0,0 +1,31 @@
+# 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")
+
+    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/safety.py

@@ -38,15 +38,10 @@ class RunShieldResponse(BaseModel):
     violation: Optional[SafetyViolation] = None
 
 
-class ShieldStore(Protocol):
-    async def get_shield(self, identifier: str) -> ShieldDef: ...
-
-
 @runtime_checkable
 class Safety(Protocol):
-    shield_store: ShieldStore
 
     @webmethod(route="/safety/run_shield")
     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: ...

llama_stack/apis/shields/client.py

@@ -26,16 +26,16 @@ 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: Shield) -> None:
         async with httpx.AsyncClient() as client:
             response = await client.post(
                 f"{self.base_url}/shields/register",
@@ -46,7 +46,7 @@ class ShieldsClient(Shields):
             )
             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 +61,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,22 @@ 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: Shield) -> None: ...