safety API cleanup part 1

Sample adapter implementation for Bedrock implementation of Guardrails

llama_stack/providers/adapters/safety/__init__.py

@@ -0,0 +1,5 @@
+# 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.

llama_stack/providers/adapters/safety/bedrock/__init__.py

@@ -0,0 +1,17 @@
+# 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 llama_stack.distribution.datatypes import RemoteProviderConfig
+
+from .config import BedrockSafetyRequestProviderData  # noqa: F403
+
+
+async def get_adapter_impl(config: RemoteProviderConfig, _deps):
+    from .bedrock import BedrockSafetyAdapter
+
+    impl = BedrockSafetyAdapter(config.url)
+    await impl.initialize()
+    return impl

llama_stack/providers/adapters/safety/bedrock/bedrock.py

@@ -0,0 +1,52 @@
+# 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 typing import List
+
+from llama_stack.apis.safety import *  # noqa: F403
+from llama_stack.providers.utils import get_request_provider_data
+
+from .config import BedrockSafetyRequestProviderData
+
+
+class BedrockSafetyAdapter(Safety):
+    def __init__(self, url: str) -> None:
+        self.url = url
+        pass
+
+    async def initialize(self) -> None:
+        pass
+
+    async def shutdown(self) -> None:
+        pass
+
+    async def run_shield(
+        self,
+        shield: ShieldType,
+        messages: List[Message],
+    ) -> RunShieldResponse:
+        # clients will set api_keys by doing something like:
+        #
+        # client = llama_stack.LlamaStack()
+        # await client.safety.run_shield(
+        #   shield_type="aws_guardrail_type",
+        #   messages=[ ... ],
+        #   x_llamastack_provider_data={
+        #     "aws_api_key": "..."
+        #   }
+        # )
+        #
+        # This information will arrive at the LlamaStack server via a HTTP Header.
+        #
+        # The server will then provide you a type-checked version of this provider data
+        # automagically by extracting it from the header and validating it with the
+        # BedrockSafetyRequestProviderData class you will need to register in the provider
+        # registry.
+        #
+        provider_data: BedrockSafetyRequestProviderData = get_request_provider_data()
+        # use `aws_api_key` to pass to the AWS servers in whichever form
+
+        raise NotImplementedError()

llama_stack/providers/adapters/safety/bedrock/config.py

@@ -0,0 +1,12 @@
+# 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 pydantic import BaseModel
+
+
+class BedrockSafetyRequestProviderData(BaseModel):
+    aws_api_key: str
+    # other AWS specific keys you may need

llama_stack/providers/impls/meta_reference/safety/safety.py

@@ -23,13 +23,6 @@ from .shields import (
 )
 
 
-def resolve_and_get_path(model_name: str) -> str:
-    model = resolve_model(model_name)
-    assert model is not None, f"Could not resolve model {model_name}"
-    model_dir = model_local_dir(model.descriptor())
-    return model_dir
-
-
 class MetaReferenceSafetyImpl(Safety):
     def __init__(self, config: SafetyConfig) -> None:
         self.config = config
@@ -50,16 +43,17 @@ class MetaReferenceSafetyImpl(Safety):
             model_dir = resolve_and_get_path(shield_cfg.model)
             _ = PromptGuardShield.instance(model_dir)
 
-    async def run_shields(
+    async def run_shield(
         self,
+        shield_type: ShieldType,
         messages: List[Message],
-        shields: List[ShieldDefinition],
     ) -> RunShieldResponse:
-        shields = [shield_config_to_shield(c, self.config) for c in shields]
+        assert shield_type in [
+            "llama_guard",
+            "prompt_guard",
+        ], f"Unknown shield {shield_type}"
 
-        responses = await asyncio.gather(*[shield.run(messages) for shield in shields])
-
-        return RunShieldResponse(responses=responses)
+        raise NotImplementedError()
 
 
 def shield_type_equals(a: ShieldType, b: ShieldType):

llama_stack/providers/registry/safety.py

@@ -6,7 +6,12 @@
 
 from typing import List
 
-from llama_stack.distribution.datatypes import Api, InlineProviderSpec, ProviderSpec
+from llama_stack.distribution.datatypes import (
+    Api,
+    InlineProviderSpec,
+    ProviderSpec,
+    remote_provider_spec,
+)
 
 
 def available_providers() -> List[ProviderSpec]:
@@ -23,4 +28,15 @@ def available_providers() -> List[ProviderSpec]:
             module="llama_stack.providers.impls.meta_reference.safety",
             config_class="llama_stack.providers.impls.meta_reference.safety.SafetyConfig",
         ),
+        remote_provider_spec(
+            api=Api.safety,
+            adapter=AdapterSpec(
+                adapter_id="bedrock",
+                pip_packages=[
+                    "aws-sdk",
+                ],
+                module="llama_stack.providers.adapters.safety.bedrock",
+                header_extractor="llama_stack.providers.adapters.safety.bedrock.BedrockSafetyRequestProviderData",
+            ),
+        ),
     ]