Update configure.py to use multiple routing keys for safety

This commit is contained in:
Ashwin Bharambe 2024-09-30 16:36:53 -07:00
parent 0996ffb3b3
commit 86834ee6c2
4 changed files with 11 additions and 71 deletions

View file

@ -117,18 +117,18 @@ def configure_api_providers(
if api_str == "safety":
# TODO: add support for other safety providers, and simplify safety provider config
if p == "meta-reference":
for shield_type in MetaReferenceShieldType:
routing_entries.append(
RoutableProviderConfig(
routing_key=shield_type.value,
provider_id=p,
config=cfg.dict(),
)
routing_entries.append(
RoutableProviderConfig(
routing_key=[s.value for s in MetaReferenceShieldType],
provider_id=p,
config=cfg.dict(),
)
)
else:
cprint(
f"[WARN] Interactive configuration of safety provider {p} is not supported, please manually configure safety shields types in routing_table of run.yaml",
f"[WARN] Interactive configuration of safety provider {p} is not supported. Please look for `{routing_key}` in run.yaml and replace it appropriately.",
"yellow",
attrs=["bold"],
)
routing_entries.append(
RoutableProviderConfig(

View file

@ -42,22 +42,7 @@ routing_table:
config:
llama_guard_shield: null
prompt_guard_shield: null
routing_key: llama_guard
- provider_id: meta-reference
config:
llama_guard_shield: null
prompt_guard_shield: null
routing_key: code_scanner_guard
- provider_id: meta-reference
config:
llama_guard_shield: null
prompt_guard_shield: null
routing_key: injection_shield
- provider_id: meta-reference
config:
llama_guard_shield: null
prompt_guard_shield: null
routing_key: jailbreak_shield
routing_key: ["llama_guard", "code_scanner_guard", "injection_shield", "jailbreak_shield"]
memory:
- provider_id: meta-reference
config: {}

View file

@ -45,22 +45,7 @@ routing_table:
config:
llama_guard_shield: null
prompt_guard_shield: null
routing_key: llama_guard
- provider_id: meta-reference
config:
llama_guard_shield: null
prompt_guard_shield: null
routing_key: code_scanner_guard
- provider_id: meta-reference
config:
llama_guard_shield: null
prompt_guard_shield: null
routing_key: injection_shield
- provider_id: meta-reference
config:
llama_guard_shield: null
prompt_guard_shield: null
routing_key: jailbreak_shield
routing_key: ["llama_guard", "code_scanner_guard", "injection_shield", "jailbreak_shield"]
memory:
- provider_id: meta-reference
config: {}

View file

@ -50,37 +50,7 @@ routing_table:
disable_output_check: false
prompt_guard_shield:
model: Prompt-Guard-86M
routing_key: llama_guard
- provider_id: meta-reference
config:
llama_guard_shield:
model: Llama-Guard-3-8B
excluded_categories: []
disable_input_check: false
disable_output_check: false
prompt_guard_shield:
model: Prompt-Guard-86M
routing_key: code_scanner_guard
- provider_id: meta-reference
config:
llama_guard_shield:
model: Llama-Guard-3-8B
excluded_categories: []
disable_input_check: false
disable_output_check: false
prompt_guard_shield:
model: Prompt-Guard-86M
routing_key: injection_shield
- provider_id: meta-reference
config:
llama_guard_shield:
model: Llama-Guard-3-8B
excluded_categories: []
disable_input_check: false
disable_output_check: false
prompt_guard_shield:
model: Prompt-Guard-86M
routing_key: jailbreak_shield
routing_key: ["llama_guard", "code_scanner_guard", "injection_shield", "jailbreak_shield"]
memory:
- provider_id: meta-reference
config: {}