improved registration flow

This commit is contained in:
Dinesh Yeduguru 2024-11-08 11:07:41 -08:00
parent 39f0c5f544
commit 0eaca98229
10 changed files with 95 additions and 26 deletions

View file

@ -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,