add safety adapters, configuration handling, server + clients

llama_toolchain/safety/api/__init__.py

@@ -3,3 +3,6 @@
 #
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
+
+from .datatypes import *  # noqa
+from .endpoints import *  # noqa

llama_toolchain/safety/api/config.py

@@ -1,25 +0,0 @@
-# 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, Optional
-
-from pydantic import BaseModel
-
-
-class LlamaGuardShieldConfig(BaseModel):
-    model_dir: str
-    excluded_categories: List[str]
-    disable_input_check: bool = False
-    disable_output_check: bool = False
-
-
-class PromptGuardShieldConfig(BaseModel):
-    model_dir: str
-
-
-class SafetyConfig(BaseModel):
-    llama_guard_shield: Optional[LlamaGuardShieldConfig] = None
-    prompt_guard_shield: Optional[PromptGuardShieldConfig] = None

llama_toolchain/safety/api/endpoints.py

@@ -15,14 +15,19 @@ from pyopenapi import webmethod
 
 @json_schema_type
 class RunShieldRequest(BaseModel):
-    shield_type: ShieldType
     messages: List[Message]
+    shields: List[ShieldDefinition]
+
+
+@json_schema_type
+class RunShieldResponse(BaseModel):
+    responses: List[ShieldResponse]
 
 
 class Safety(Protocol):
 
-    @webmethod(route="/safety/run_shield")
-    async def run_shield(
+    @webmethod(route="/safety/run_shields")
+    async def run_shields(
         self,
         request: RunShieldRequest,
-    ) -> ShieldResponse: ...
+    ) -> RunShieldResponse: ...